Skip to content

v0.8.3 — Stable Cascade (Würstchen)

Choose a tag to compare

@dahai80 dahai80 released this 06 Aug 18:57
a0d03fb

Stable Cascade (Würstchen) native MLX image pipeline (#370)

Native MLX port of Stable Cascade — a 3-stage Würstchen pipeline
(prior → decoder → VQGAN), from scratch under fusion_mlx/image/cascade/.
Does not wrap mflux.

Highlights

  • Unified StableCascadeUNet serves both the prior
    (switch_level=(False,)UpDownBlock2d 1×1 mapping, no spatial
    change) and the decoder (switch_level=NoneConv2d(k=2,s=2) down
    / ConvTranspose2d(k=2,s=2) up).
  • PaellaVQModel decode-only VQGAN (scale_factor=0.3764).
  • CLIP-ViT-bigG text encoder (dims=1280, 32 layers, 20 heads).
  • DDPMWuerstchenScheduler (cosine _alpha_cumprod).
  • NHWC throughout (MLX conv2d-native).

Weight-layout fix (key gotcha)

  • Conv2d/DepthwiseConv2d (out,in,k,k) → OHWI (0,2,3,1)
  • ConvTranspose2d (in,out,k,k) → OHWI (1,2,3,0)different transpose
  • nn.Linear / nn.Embedding left as-is (MLX stores (out,in) like
    PyTorch; transposing 2D weights on load is a bug)

Validation

  • Key alignment vs real stabilityai weights: prior 1550/1550,
    decoder 1726/1726, VQGAN 121/122, CLIP 517/517.
  • 24 unit tests green; real-weights E2E (512×512 / 768×768 PNG, CFG);
    server-path E2E via ImageGenEngine.
  • ruff + black clean across the repo.

Wiring

  • image_gen engine stable_cascade variant (auto-detected from
    cascade/wuerstchen model names) → /v1/images/generate.
  • model_discovery maps StableCascadePriorPipeline /
    StableCascadeDecoderPipelinetext-to-image.
  • docs: docs/cascade-image.md

Closes #370.

PyPI publish is blocked by trusted-publisher re-registration
(issue #384); Homebrew tap updated.