Skip to content

Pipelines

mlx_diffuser.pipelines.DiffusionPipeline

Source code in src/mlx_diffuser/pipelines/base.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
class DiffusionPipeline:
    #: Names of attributes that are persisted as components.
    _component_names: tuple[str, ...] = ()

    def __init__(self, **components: Any):
        for name, value in components.items():
            setattr(self, name, value)

    # --- persistence ------------------------------------------------------
    def save_pretrained(self, save_directory: str | Path) -> Path:
        save_directory = Path(save_directory)
        save_directory.mkdir(parents=True, exist_ok=True)
        index: dict[str, Any] = {"_class_name": type(self).__name__, "_components": {}}
        for name in self._component_names:
            comp = getattr(self, name)
            comp.save_pretrained(save_directory / name)
            index["_components"][name] = type(comp).__name__
        (save_directory / MODEL_INDEX_NAME).write_text(json.dumps(index, indent=2, sort_keys=True))
        return save_directory

    @classmethod
    def from_pretrained(
        cls,
        path_or_repo_id: str | Path,
        *,
        dtype: str | mx.Dtype | None = None,
        quantize: int | None = None,
        revision: str | None = None,
    ) -> DiffusionPipeline:
        local = resolve(path_or_repo_id, revision=revision)
        index = json.loads((local / MODEL_INDEX_NAME).read_text())

        target: type[DiffusionPipeline] = cls
        if cls is DiffusionPipeline:  # dispatch to the concrete pipeline
            resolved = PIPELINE_REGISTRY.get(index["_class_name"])
            if resolved is None:
                raise ValueError(f"Unknown pipeline {index['_class_name']!r}. Is it registered?")
            target = resolved

        components = {
            name: _load_component(class_name, local / name, dtype, quantize)
            for name, class_name in index["_components"].items()
        }
        return target(**components)

    # --- helpers for subclasses ------------------------------------------
    @staticmethod
    def _resolve_key(seed: int | None, key: mx.array | None) -> mx.array:
        if key is not None:
            return key
        return mx.random.key(seed if seed is not None else 0)

    @staticmethod
    def classifier_free_guidance(cond: mx.array, uncond: mx.array, scale: float) -> mx.array:
        return uncond + scale * (cond - uncond)

    def denoising_loop(
        self,
        scheduler: Scheduler,
        latents: mx.array,
        predict: Callable[[mx.array, mx.array], mx.array],
        key: mx.array,
        progress: bool = False,
    ) -> mx.array:
        """Run the reverse process, calling ``predict(scaled_latents, t)`` per step.

        ``predict`` returns the (already guidance-combined) model output. Evaluation
        is forced once per step to keep the lazy graph bounded.
        """
        steps = scheduler.timesteps
        assert steps is not None, "call scheduler.set_timesteps() before sampling"
        if progress:
            from ..utils import get_logger

            get_logger().info("sampling %d steps", len(steps))
        for t in steps:
            scaled = scheduler.scale_model_input(latents, t)
            model_output = predict(scaled, t)
            key, step_key = mx.random.split(key)
            latents = scheduler.step(model_output, t, latents, key=step_key)
            mx.eval(latents)
        return latents

denoising_loop(scheduler, latents, predict, key, progress=False)

Run the reverse process, calling predict(scaled_latents, t) per step.

predict returns the (already guidance-combined) model output. Evaluation is forced once per step to keep the lazy graph bounded.

Source code in src/mlx_diffuser/pipelines/base.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def denoising_loop(
    self,
    scheduler: Scheduler,
    latents: mx.array,
    predict: Callable[[mx.array, mx.array], mx.array],
    key: mx.array,
    progress: bool = False,
) -> mx.array:
    """Run the reverse process, calling ``predict(scaled_latents, t)`` per step.

    ``predict`` returns the (already guidance-combined) model output. Evaluation
    is forced once per step to keep the lazy graph bounded.
    """
    steps = scheduler.timesteps
    assert steps is not None, "call scheduler.set_timesteps() before sampling"
    if progress:
        from ..utils import get_logger

        get_logger().info("sampling %d steps", len(steps))
    for t in steps:
        scaled = scheduler.scale_model_input(latents, t)
        model_output = predict(scaled, t)
        key, step_key = mx.random.split(key)
        latents = scheduler.step(model_output, t, latents, key=step_key)
        mx.eval(latents)
    return latents

mlx_diffuser.pipelines.ClassConditionalPipeline

Bases: DiffusionPipeline

Source code in src/mlx_diffuser/pipelines/class_conditional.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
@register_pipeline
class ClassConditionalPipeline(DiffusionPipeline):
    _component_names = ("model", "scheduler")

    def __init__(self, model: DiT, scheduler: Scheduler):
        super().__init__(model=model, scheduler=scheduler)
        self.model = model
        self.scheduler = scheduler

    def __call__(
        self,
        class_labels: list[int] | mx.array,
        *,
        sample_size: int = 32,
        num_inference_steps: int = 50,
        guidance_scale: float = 4.0,
        seed: int | None = None,
        key: mx.array | None = None,
        compile: bool = True,
        progress: bool = False,
    ) -> mx.array:
        labels = mx.array(class_labels) if not isinstance(class_labels, mx.array) else class_labels
        b = labels.shape[0]
        c = self.model.config.in_channels
        key = self._resolve_key(seed, key)

        key, noise_key = mx.random.split(key)
        latents = mx.random.normal((b, sample_size, sample_size, c), key=noise_key)

        self.scheduler.set_timesteps(num_inference_steps)
        use_cfg = guidance_scale and guidance_scale > 1.0
        null = mx.full((b,), self.model.config.num_classes)
        # Per-step shapes are constant, so a single compiled graph is reused.
        model = compile_model(self.model) if compile else self.model

        def predict(scaled: mx.array, t: mx.array) -> mx.array:
            tb = mx.ones((b,)) * t
            if use_cfg:
                model_in = mx.concatenate([scaled, scaled], axis=0)
                y = mx.concatenate([labels, null], axis=0)
                tt = mx.concatenate([tb, tb], axis=0)
                out = model(model_in, tt, y)
                cond, uncond = out[:b], out[b:]
                return self.classifier_free_guidance(cond, uncond, guidance_scale)
            return model(scaled, tb, labels)

        latents = self.denoising_loop(self.scheduler, latents, predict, key, progress=progress)
        return latents

mlx_diffuser.pipelines.TextToVideoPipeline

Bases: DiffusionPipeline

Source code in src/mlx_diffuser/pipelines/text_to_video.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
@register_pipeline
class TextToVideoPipeline(DiffusionPipeline):
    _component_names = ("transformer", "vae", "scheduler")

    def __init__(self, transformer: VideoDiT, vae: AutoencoderKLVideo, scheduler: Scheduler):
        super().__init__(transformer=transformer, vae=vae, scheduler=scheduler)
        self.transformer = transformer
        self.vae = vae
        self.scheduler = scheduler

    def _latent_shape(self, batch: int, num_frames: int, height: int, width: int) -> tuple:
        c = self.vae.config
        sp, tc = c.spatial_compression, c.temporal_compression
        if num_frames % tc or height % sp or width % sp:
            raise ValueError(
                f"num_frames must be divisible by {tc} and height/width by {sp} "
                f"(got {num_frames}, {height}, {width})."
            )
        return (
            batch,
            num_frames // tc,
            height // sp,
            width // sp,
            self.transformer.config.in_channels,
        )

    def __call__(
        self,
        prompt_embeds: mx.array,
        *,
        negative_embeds: mx.array | None = None,
        num_frames: int = 16,
        height: int = 256,
        width: int = 256,
        num_inference_steps: int = 50,
        guidance_scale: float = 5.0,
        seed: int | None = None,
        key: mx.array | None = None,
        compile: bool = True,
        decode: bool = True,
        progress: bool = False,
    ) -> mx.array:
        """Generate video(s) from text embeddings.

        Args:
            prompt_embeds: ``(B, L, D)`` per-token text embeddings.
            negative_embeds: ``(B, L, D)`` for the unconditional branch (CFG);
                defaults to zeros when guidance is enabled.
            num_frames: number of output frames.
            height: output height in pixels.
            width: output width in pixels.
            decode: if ``False``, return raw latents instead of pixels.

        Returns:
            ``(B, T, H, W, C)`` video in roughly ``[-1, 1]``, or latents if
            ``decode=False``.
        """
        b = prompt_embeds.shape[0]
        key = self._resolve_key(seed, key)
        key, noise_key = mx.random.split(key)
        latents = mx.random.normal(self._latent_shape(b, num_frames, height, width), key=noise_key)

        self.scheduler.set_timesteps(num_inference_steps)
        use_cfg = guidance_scale and guidance_scale > 1.0
        if use_cfg and negative_embeds is None:
            negative_embeds = mx.zeros_like(prompt_embeds)

        model = compile_model(self.transformer) if compile else self.transformer

        def predict(scaled: mx.array, t: mx.array) -> mx.array:
            tb = mx.ones((b,)) * t
            if use_cfg:
                assert negative_embeds is not None  # set above when CFG is enabled
                x2 = mx.concatenate([scaled, scaled], axis=0)
                t2 = mx.concatenate([tb, tb], axis=0)
                ctx = mx.concatenate([prompt_embeds, negative_embeds], axis=0)
                out = model(x2, t2, ctx)
                cond, uncond = out[:b], out[b:]
                return self.classifier_free_guidance(cond, uncond, guidance_scale)
            return model(scaled, tb, prompt_embeds)

        latents = self.denoising_loop(self.scheduler, latents, predict, key, progress=progress)
        if not decode:
            return latents
        video = self.vae.decode(latents / self.vae.scaling_factor)
        mx.eval(video)
        return video

__call__(prompt_embeds, *, negative_embeds=None, num_frames=16, height=256, width=256, num_inference_steps=50, guidance_scale=5.0, seed=None, key=None, compile=True, decode=True, progress=False)

Generate video(s) from text embeddings.

Parameters:

Name Type Description Default
prompt_embeds array

(B, L, D) per-token text embeddings.

 required
negative_embeds array | None

(B, L, D) for the unconditional branch (CFG); defaults to zeros when guidance is enabled.

 None
num_frames int

number of output frames.

 16
height int

output height in pixels.

 256
width int

output width in pixels.

 256
decode bool

if False, return raw latents instead of pixels.

 True

Returns:

Type Description
array

(B, T, H, W, C) video in roughly [-1, 1], or latents if
array

decode=False.
Source code in src/mlx_diffuser/pipelines/text_to_video.py 
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def __call__(
    self,
    prompt_embeds: mx.array,
    *,
    negative_embeds: mx.array | None = None,
    num_frames: int = 16,
    height: int = 256,
    width: int = 256,
    num_inference_steps: int = 50,
    guidance_scale: float = 5.0,
    seed: int | None = None,
    key: mx.array | None = None,
    compile: bool = True,
    decode: bool = True,
    progress: bool = False,
) -> mx.array:
    """Generate video(s) from text embeddings.

    Args:
        prompt_embeds: ``(B, L, D)`` per-token text embeddings.
        negative_embeds: ``(B, L, D)`` for the unconditional branch (CFG);
            defaults to zeros when guidance is enabled.
        num_frames: number of output frames.
        height: output height in pixels.
        width: output width in pixels.
        decode: if ``False``, return raw latents instead of pixels.

    Returns:
        ``(B, T, H, W, C)`` video in roughly ``[-1, 1]``, or latents if
        ``decode=False``.
    """
    b = prompt_embeds.shape[0]
    key = self._resolve_key(seed, key)
    key, noise_key = mx.random.split(key)
    latents = mx.random.normal(self._latent_shape(b, num_frames, height, width), key=noise_key)

    self.scheduler.set_timesteps(num_inference_steps)
    use_cfg = guidance_scale and guidance_scale > 1.0
    if use_cfg and negative_embeds is None:
        negative_embeds = mx.zeros_like(prompt_embeds)

    model = compile_model(self.transformer) if compile else self.transformer

    def predict(scaled: mx.array, t: mx.array) -> mx.array:
        tb = mx.ones((b,)) * t
        if use_cfg:
            assert negative_embeds is not None  # set above when CFG is enabled
            x2 = mx.concatenate([scaled, scaled], axis=0)
            t2 = mx.concatenate([tb, tb], axis=0)
            ctx = mx.concatenate([prompt_embeds, negative_embeds], axis=0)
            out = model(x2, t2, ctx)
            cond, uncond = out[:b], out[b:]
            return self.classifier_free_guidance(cond, uncond, guidance_scale)
        return model(scaled, tb, prompt_embeds)

    latents = self.denoising_loop(self.scheduler, latents, predict, key, progress=progress)
    if not decode:
        return latents
    video = self.vae.decode(latents / self.vae.scaling_factor)
    mx.eval(video)
    return video

mlx_diffuser.pipelines.WanPipeline

WAN 2.1 text-to-video pipeline (channels-last (B, T, H, W, C)).

Source code in src/mlx_diffuser/pipelines/wan.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
class WanPipeline:
    """WAN 2.1 text-to-video pipeline (channels-last ``(B, T, H, W, C)``)."""

    def __init__(
        self,
        transformer: WanTransformer3DModel,
        vae: AutoencoderKLWan,
        text_encoder: UMT5EncoderModel | None,
        tokenizer,
        scheduler: FlowMatchEulerScheduler,
    ):
        self.transformer = transformer
        self.vae = vae
        self.text_encoder = text_encoder
        self.tokenizer = tokenizer
        self.scheduler = scheduler

    # --- loading -------------------------------------------------------------
    @classmethod
    def from_diffusers(
        cls,
        folder: str | Path,
        *,
        text_encoder: str | Path | None = None,
        quantize_text: int | None = 4,
        quantize_transformer: int | None = None,
        transformer_dtype: mx.Dtype = mx.bfloat16,
        shift: float = 3.0,
    ) -> WanPipeline:
        """Load + convert an official ``Wan2.1-T2V-*-Diffusers`` folder into MLX.

        ``quantize_text`` keeps the 5.6B umT5 encoder small (4-bit ≈ 3 GB);
        ``quantize_transformer`` optionally quantizes the DiT (defaults to bf16).
        ``text_encoder`` overrides where the umT5 weights come from — a folder or a
        single ``.safetensors`` file (e.g. an fp16 community checkpoint, half the
        size of the fp32 default); falls back to ``folder/text_encoder``.
        """
        from ..converters import get_converter

        folder = Path(folder)
        te_source = Path(text_encoder) if text_encoder is not None else folder / "text_encoder"
        # Convert the large text encoder first (quantizing while little else is
        # resident), then the transformer, to keep the peak conversion memory down.
        vae = cast(AutoencoderKLWan, get_converter("AutoencoderKLWan").convert(folder / "vae"))
        text_encoder_model = cast(
            UMT5EncoderModel,
            get_converter("UMT5EncoderModel").convert(te_source, quantize=quantize_text),
        )
        transformer = cast(
            WanTransformer3DModel,
            get_converter("WanTransformer3DModel").convert(
                folder / "transformer", dtype=transformer_dtype, quantize=quantize_transformer
            ),
        )
        tokenizer = _load_tokenizer(folder / "tokenizer")
        scheduler = FlowMatchEulerScheduler(FlowMatchConfig(shift=shift))
        return cls(transformer, vae, text_encoder_model, tokenizer, scheduler)

    # --- text encoding -------------------------------------------------------
    def encode_text(self, prompt: str) -> mx.array:
        """Tokenize and encode ``prompt`` into ``(1, L, text_dim)`` embeddings."""
        if self.text_encoder is None:
            raise RuntimeError("Text encoder was released; re-load the pipeline to encode.")
        enc = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=MAX_PROMPT_TOKENS,
            truncation=True,
            return_tensors="np",
        )
        ids = mx.array(enc["input_ids"].astype("int32"))
        mask = mx.array(enc["attention_mask"].astype("int32"))
        embeds = self.text_encoder(ids, mask)
        # Zero the padded positions so the transformer's cross-attention sees the
        # same fixed token budget it was trained on. Without the padding the text
        # attention is far too concentrated and the video collapses to a flat frame.
        embeds = embeds * mask[..., None].astype(embeds.dtype)
        mx.eval(embeds)
        return embeds

    def release_text_encoder(self) -> None:
        """Free the text encoder (call after encoding to reclaim memory)."""
        self.text_encoder = None
        gc.collect()
        mx.clear_cache()

    # --- generation ----------------------------------------------------------
    def __call__(
        self,
        prompt: str,
        *,
        negative_prompt: str = "",
        num_frames: int = 17,
        height: int = 256,
        width: int = 256,
        num_inference_steps: int = 40,
        guidance_scale: float = 5.0,
        seed: int = 0,
        cache_threshold: float = 0.0,
        release_text_encoder: bool = True,
        progress: bool = True,
    ) -> mx.array:
        """Generate a video ``(1, num_frames, height, width, 3)`` in ``[-1, 1]``.

        ``num_frames`` must satisfy ``(num_frames - 1) % 4 == 0`` and ``height`` /
        ``width`` must be multiples of 8 (the VAE's spatial compression).

        ``cache_threshold`` enables First-Block Cache (``0`` = off, exact). On the
        1.3B model, ``0.1`` ≈ 1.5x and ``0.2`` ≈ 2.2x faster with no visible quality
        change; ``>= 0.3`` starts to degrade.
        """
        if (num_frames - 1) % 4 != 0:
            raise ValueError("num_frames must be 1 + a multiple of 4 (e.g. 17, 33, 49, 81).")
        if height % 8 or width % 8:
            raise ValueError("height and width must be multiples of 8.")

        use_cfg = bool(guidance_scale) and guidance_scale > 1.0
        prompt_embeds = self.encode_text(prompt)
        # Batch the conditional + unconditional CFG passes into one forward.
        if use_cfg:
            context = mx.concatenate([prompt_embeds, self.encode_text(negative_prompt)], axis=0)
        else:
            context = prompt_embeds
        if release_text_encoder:
            self.release_text_encoder()

        latent_frames = (num_frames - 1) // 4 + 1
        lh, lw = height // 8, width // 8
        c = self.transformer.config.in_channels
        latents = mx.random.normal((1, latent_frames, lh, lw, c), key=mx.random.key(seed))

        cache = FirstBlockCache(cache_threshold) if cache_threshold > 0 else None
        self.scheduler.set_timesteps(num_inference_steps)
        steps = self.scheduler.timesteps
        assert steps is not None
        bar = _DenoiseProgress(len(steps), enabled=progress)
        for t in steps:
            ts = t * 1000.0
            if use_cfg:
                x2 = mx.concatenate([latents, latents], axis=0)
                out = self.transformer(x2, mx.broadcast_to(ts, (2,)), context, cache=cache)
                cond, uncond = out[:1], out[1:]
                v = uncond + guidance_scale * (cond - uncond)
            else:
                v = self.transformer(latents, mx.broadcast_to(ts, (1,)), context, cache=cache)
            latents = self.scheduler.step(v, t, latents)
            mx.eval(latents)
            bar.update(cache)
        bar.close()

        z = self.vae.denormalize_latents(latents).astype(mx.float32)
        video = self.vae.decode(z)
        mx.eval(video)
        return video

__call__(prompt, *, negative_prompt='', num_frames=17, height=256, width=256, num_inference_steps=40, guidance_scale=5.0, seed=0, cache_threshold=0.0, release_text_encoder=True, progress=True)

Generate a video (1, num_frames, height, width, 3) in [-1, 1].

num_frames must satisfy (num_frames - 1) % 4 == 0 and height / width must be multiples of 8 (the VAE's spatial compression).

cache_threshold enables First-Block Cache (0 = off, exact). On the 1.3B model, 0.1 ≈ 1.5x and 0.2 ≈ 2.2x faster with no visible quality change; >= 0.3 starts to degrade.

Source code in src/mlx_diffuser/pipelines/wan.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def __call__(
    self,
    prompt: str,
    *,
    negative_prompt: str = "",
    num_frames: int = 17,
    height: int = 256,
    width: int = 256,
    num_inference_steps: int = 40,
    guidance_scale: float = 5.0,
    seed: int = 0,
    cache_threshold: float = 0.0,
    release_text_encoder: bool = True,
    progress: bool = True,
) -> mx.array:
    """Generate a video ``(1, num_frames, height, width, 3)`` in ``[-1, 1]``.

    ``num_frames`` must satisfy ``(num_frames - 1) % 4 == 0`` and ``height`` /
    ``width`` must be multiples of 8 (the VAE's spatial compression).

    ``cache_threshold`` enables First-Block Cache (``0`` = off, exact). On the
    1.3B model, ``0.1`` ≈ 1.5x and ``0.2`` ≈ 2.2x faster with no visible quality
    change; ``>= 0.3`` starts to degrade.
    """
    if (num_frames - 1) % 4 != 0:
        raise ValueError("num_frames must be 1 + a multiple of 4 (e.g. 17, 33, 49, 81).")
    if height % 8 or width % 8:
        raise ValueError("height and width must be multiples of 8.")

    use_cfg = bool(guidance_scale) and guidance_scale > 1.0
    prompt_embeds = self.encode_text(prompt)
    # Batch the conditional + unconditional CFG passes into one forward.
    if use_cfg:
        context = mx.concatenate([prompt_embeds, self.encode_text(negative_prompt)], axis=0)
    else:
        context = prompt_embeds
    if release_text_encoder:
        self.release_text_encoder()

    latent_frames = (num_frames - 1) // 4 + 1
    lh, lw = height // 8, width // 8
    c = self.transformer.config.in_channels
    latents = mx.random.normal((1, latent_frames, lh, lw, c), key=mx.random.key(seed))

    cache = FirstBlockCache(cache_threshold) if cache_threshold > 0 else None
    self.scheduler.set_timesteps(num_inference_steps)
    steps = self.scheduler.timesteps
    assert steps is not None
    bar = _DenoiseProgress(len(steps), enabled=progress)
    for t in steps:
        ts = t * 1000.0
        if use_cfg:
            x2 = mx.concatenate([latents, latents], axis=0)
            out = self.transformer(x2, mx.broadcast_to(ts, (2,)), context, cache=cache)
            cond, uncond = out[:1], out[1:]
            v = uncond + guidance_scale * (cond - uncond)
        else:
            v = self.transformer(latents, mx.broadcast_to(ts, (1,)), context, cache=cache)
        latents = self.scheduler.step(v, t, latents)
        mx.eval(latents)
        bar.update(cache)
    bar.close()

    z = self.vae.denormalize_latents(latents).astype(mx.float32)
    video = self.vae.decode(z)
    mx.eval(video)
    return video

encode_text(prompt)

Tokenize and encode prompt into (1, L, text_dim) embeddings.

Source code in src/mlx_diffuser/pipelines/wan.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def encode_text(self, prompt: str) -> mx.array:
    """Tokenize and encode ``prompt`` into ``(1, L, text_dim)`` embeddings."""
    if self.text_encoder is None:
        raise RuntimeError("Text encoder was released; re-load the pipeline to encode.")
    enc = self.tokenizer(
        prompt,
        padding="max_length",
        max_length=MAX_PROMPT_TOKENS,
        truncation=True,
        return_tensors="np",
    )
    ids = mx.array(enc["input_ids"].astype("int32"))
    mask = mx.array(enc["attention_mask"].astype("int32"))
    embeds = self.text_encoder(ids, mask)
    # Zero the padded positions so the transformer's cross-attention sees the
    # same fixed token budget it was trained on. Without the padding the text
    # attention is far too concentrated and the video collapses to a flat frame.
    embeds = embeds * mask[..., None].astype(embeds.dtype)
    mx.eval(embeds)
    return embeds

from_diffusers(folder, *, text_encoder=None, quantize_text=4, quantize_transformer=None, transformer_dtype=mx.bfloat16, shift=3.0) classmethod

Load + convert an official Wan2.1-T2V-*-Diffusers folder into MLX.

quantize_text keeps the 5.6B umT5 encoder small (4-bit ≈ 3 GB); quantize_transformer optionally quantizes the DiT (defaults to bf16). text_encoder overrides where the umT5 weights come from — a folder or a single .safetensors file (e.g. an fp16 community checkpoint, half the size of the fp32 default); falls back to folder/text_encoder.

Source code in src/mlx_diffuser/pipelines/wan.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
@classmethod
def from_diffusers(
    cls,
    folder: str | Path,
    *,
    text_encoder: str | Path | None = None,
    quantize_text: int | None = 4,
    quantize_transformer: int | None = None,
    transformer_dtype: mx.Dtype = mx.bfloat16,
    shift: float = 3.0,
) -> WanPipeline:
    """Load + convert an official ``Wan2.1-T2V-*-Diffusers`` folder into MLX.

    ``quantize_text`` keeps the 5.6B umT5 encoder small (4-bit ≈ 3 GB);
    ``quantize_transformer`` optionally quantizes the DiT (defaults to bf16).
    ``text_encoder`` overrides where the umT5 weights come from — a folder or a
    single ``.safetensors`` file (e.g. an fp16 community checkpoint, half the
    size of the fp32 default); falls back to ``folder/text_encoder``.
    """
    from ..converters import get_converter

    folder = Path(folder)
    te_source = Path(text_encoder) if text_encoder is not None else folder / "text_encoder"
    # Convert the large text encoder first (quantizing while little else is
    # resident), then the transformer, to keep the peak conversion memory down.
    vae = cast(AutoencoderKLWan, get_converter("AutoencoderKLWan").convert(folder / "vae"))
    text_encoder_model = cast(
        UMT5EncoderModel,
        get_converter("UMT5EncoderModel").convert(te_source, quantize=quantize_text),
    )
    transformer = cast(
        WanTransformer3DModel,
        get_converter("WanTransformer3DModel").convert(
            folder / "transformer", dtype=transformer_dtype, quantize=quantize_transformer
        ),
    )
    tokenizer = _load_tokenizer(folder / "tokenizer")
    scheduler = FlowMatchEulerScheduler(FlowMatchConfig(shift=shift))
    return cls(transformer, vae, text_encoder_model, tokenizer, scheduler)

release_text_encoder()

Free the text encoder (call after encoding to reclaim memory).

Source code in src/mlx_diffuser/pipelines/wan.py 
113
114
115
116
117
def release_text_encoder(self) -> None:
    """Free the text encoder (call after encoding to reclaim memory)."""
    self.text_encoder = None
    gc.collect()
    mx.clear_cache()

mlx_diffuser.pipelines.StableDiffusionXLPipeline

SDXL base text-to-image pipeline (channels-last (B, H, W, C)).

Source code in src/mlx_diffuser/pipelines/sdxl.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
class StableDiffusionXLPipeline:
    """SDXL base text-to-image pipeline (channels-last ``(B, H, W, C)``)."""

    def __init__(
        self,
        unet: SDXLUNet,
        vae: AutoencoderKLSD,
        text_encoder: CLIPTextModel,
        text_encoder_2: CLIPTextModel,
        tokenizer,
        tokenizer_2,
        scheduler: EulerDiscreteScheduler,
    ):
        self.unet = unet
        self.vae = vae
        self.text_encoder = text_encoder
        self.text_encoder_2 = text_encoder_2
        self.tokenizer = tokenizer
        self.tokenizer_2 = tokenizer_2
        self.scheduler = scheduler

    # --- loading -------------------------------------------------------------
    @classmethod
    def from_diffusers(
        cls,
        folder: str | Path,
        *,
        dtype: mx.Dtype = mx.float16,
        quantize_unet: int | None = None,
    ) -> StableDiffusionXLPipeline:
        """Load + convert a ``stable-diffusion-xl-base`` folder into MLX.

        The UNet and text encoders run in ``dtype`` (fp16 by default, SDXL's native
        precision); the VAE runs in fp32 because SDXL's VAE overflows fp16.
        ``quantize_unet`` (4/8) weight-quantizes the UNet to save memory.
        """
        from ..converters import get_converter

        folder = Path(folder)
        unet = cast(
            SDXLUNet,
            get_converter("UNet2DConditionModel").convert(
                folder / "unet", dtype=dtype, quantize=quantize_unet
            ),
        )
        vae = cast(
            AutoencoderKLSD,
            get_converter("AutoencoderKL").convert(folder / "vae", dtype=mx.float32),
        )
        te = cast(
            CLIPTextModel,
            get_converter("CLIPTextModel").convert(folder / "text_encoder", dtype=dtype),
        )
        te2 = cast(
            CLIPTextModel,
            get_converter("CLIPTextModelWithProjection").convert(
                folder / "text_encoder_2", dtype=dtype
            ),
        )
        tok = _load_tokenizer(folder / "tokenizer")
        tok2 = _load_tokenizer(folder / "tokenizer_2")
        scheduler = EulerDiscreteScheduler(
            EulerConfig(
                beta_schedule="scaled_linear",
                beta_start=0.00085,
                beta_end=0.012,
                prediction_type="epsilon",
                timestep_spacing="leading",
                steps_offset=1,
            )
        )
        return cls(unet, vae, te, te2, tok, tok2, scheduler)

    # --- text encoding -------------------------------------------------------
    def encode_prompt(self, prompt: str) -> tuple[mx.array, mx.array]:
        """Tokenize + encode with both CLIPs.

        Returns ``(prompt_embeds (1, 77, 2048), pooled (1, 1280))``: the two encoders'
        penultimate hidden states concatenated, and the bigG projected pooled embed.
        """

        def ids(tok) -> mx.array:
            enc = tok(
                prompt,
                padding="max_length",
                max_length=MAX_TOKENS,
                truncation=True,
                return_tensors="np",
            )
            return mx.array(enc["input_ids"].astype("int32"))

        hs1, _ = self.text_encoder(ids(self.tokenizer))
        hs2, pooled = self.text_encoder_2(ids(self.tokenizer_2))
        embeds = mx.concatenate([hs1[-2], hs2[-2]], axis=-1)  # (1, 77, 768+1280)
        return embeds, pooled

    # --- generation ----------------------------------------------------------
    def __call__(
        self,
        prompt: str,
        *,
        negative_prompt: str = "",
        height: int = 1024,
        width: int = 1024,
        num_inference_steps: int = 30,
        guidance_scale: float = 5.0,
        seed: int = 0,
        cache_interval: int = 1,
        tile_vae: bool = False,
        progress: bool = True,
    ) -> mx.array:
        """Generate an image ``(1, height, width, 3)`` in ``[-1, 1]``.

        ``height`` / ``width`` must be multiples of 8. ``cache_interval`` enables
        DeepCache (1 = off/exact; 2 ≈ 1.5-1.8x by skipping the deep UNet blocks on every
        other step). ``tile_vae`` decodes the VAE in tiles to bound memory at high res.
        """
        if height % 8 or width % 8:
            raise ValueError("height and width must be multiples of 8.")

        use_cfg = guidance_scale > 1.0
        pe, pooled = self.encode_prompt(prompt)
        if use_cfg:
            npe, npooled = self.encode_prompt(negative_prompt)
            context = mx.concatenate([pe, npe], axis=0)
            text_embeds = mx.concatenate([pooled, npooled], axis=0)
        else:
            context, text_embeds = pe, pooled
        n = context.shape[0]
        time_ids = mx.broadcast_to(
            mx.array([[float(height), float(width), 0.0, 0.0, float(height), float(width)]]),
            (n, 6),
        ).astype(context.dtype)

        self.scheduler.set_timesteps(num_inference_steps)
        steps = self.scheduler.timesteps
        assert steps is not None
        latents = mx.random.normal((1, height // 8, width // 8, 4), key=mx.random.key(seed))
        latents = latents * self.scheduler.init_noise_sigma

        cache = DeepCache(cache_interval) if cache_interval > 1 else None
        for i, t in enumerate(steps):
            scaled = self.scheduler.scale_model_input(latents, t)
            tt = mx.broadcast_to(t, (n,))
            if use_cfg:
                x2 = mx.concatenate([scaled, scaled], axis=0)
                out = self.unet(x2, tt, context, text_embeds, time_ids, cache=cache)
                cond, uncond = out[:1], out[1:]
                noise = uncond + guidance_scale * (cond - uncond)
            else:
                noise = self.unet(scaled, tt, context, text_embeds, time_ids, cache=cache)
            latents = self.scheduler.step(noise, t, latents)
            mx.eval(latents)
            if progress:
                print(f"  step {i + 1}/{len(steps)}", end="\r")
        if progress:
            print()

        image = self.vae.decode(latents.astype(mx.float32) / self.vae.scaling_factor, tile=tile_vae)
        mx.eval(image)
        return image

__call__(prompt, *, negative_prompt='', height=1024, width=1024, num_inference_steps=30, guidance_scale=5.0, seed=0, cache_interval=1, tile_vae=False, progress=True)

Generate an image (1, height, width, 3) in [-1, 1].

height / width must be multiples of 8. cache_interval enables DeepCache (1 = off/exact; 2 ≈ 1.5-1.8x by skipping the deep UNet blocks on every other step). tile_vae decodes the VAE in tiles to bound memory at high res.

Source code in src/mlx_diffuser/pipelines/sdxl.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def __call__(
    self,
    prompt: str,
    *,
    negative_prompt: str = "",
    height: int = 1024,
    width: int = 1024,
    num_inference_steps: int = 30,
    guidance_scale: float = 5.0,
    seed: int = 0,
    cache_interval: int = 1,
    tile_vae: bool = False,
    progress: bool = True,
) -> mx.array:
    """Generate an image ``(1, height, width, 3)`` in ``[-1, 1]``.

    ``height`` / ``width`` must be multiples of 8. ``cache_interval`` enables
    DeepCache (1 = off/exact; 2 ≈ 1.5-1.8x by skipping the deep UNet blocks on every
    other step). ``tile_vae`` decodes the VAE in tiles to bound memory at high res.
    """
    if height % 8 or width % 8:
        raise ValueError("height and width must be multiples of 8.")

    use_cfg = guidance_scale > 1.0
    pe, pooled = self.encode_prompt(prompt)
    if use_cfg:
        npe, npooled = self.encode_prompt(negative_prompt)
        context = mx.concatenate([pe, npe], axis=0)
        text_embeds = mx.concatenate([pooled, npooled], axis=0)
    else:
        context, text_embeds = pe, pooled
    n = context.shape[0]
    time_ids = mx.broadcast_to(
        mx.array([[float(height), float(width), 0.0, 0.0, float(height), float(width)]]),
        (n, 6),
    ).astype(context.dtype)

    self.scheduler.set_timesteps(num_inference_steps)
    steps = self.scheduler.timesteps
    assert steps is not None
    latents = mx.random.normal((1, height // 8, width // 8, 4), key=mx.random.key(seed))
    latents = latents * self.scheduler.init_noise_sigma

    cache = DeepCache(cache_interval) if cache_interval > 1 else None
    for i, t in enumerate(steps):
        scaled = self.scheduler.scale_model_input(latents, t)
        tt = mx.broadcast_to(t, (n,))
        if use_cfg:
            x2 = mx.concatenate([scaled, scaled], axis=0)
            out = self.unet(x2, tt, context, text_embeds, time_ids, cache=cache)
            cond, uncond = out[:1], out[1:]
            noise = uncond + guidance_scale * (cond - uncond)
        else:
            noise = self.unet(scaled, tt, context, text_embeds, time_ids, cache=cache)
        latents = self.scheduler.step(noise, t, latents)
        mx.eval(latents)
        if progress:
            print(f"  step {i + 1}/{len(steps)}", end="\r")
    if progress:
        print()

    image = self.vae.decode(latents.astype(mx.float32) / self.vae.scaling_factor, tile=tile_vae)
    mx.eval(image)
    return image

encode_prompt(prompt)

Tokenize + encode with both CLIPs.

Returns (prompt_embeds (1, 77, 2048), pooled (1, 1280)): the two encoders' penultimate hidden states concatenated, and the bigG projected pooled embed.

Source code in src/mlx_diffuser/pipelines/sdxl.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def encode_prompt(self, prompt: str) -> tuple[mx.array, mx.array]:
    """Tokenize + encode with both CLIPs.

    Returns ``(prompt_embeds (1, 77, 2048), pooled (1, 1280))``: the two encoders'
    penultimate hidden states concatenated, and the bigG projected pooled embed.
    """

    def ids(tok) -> mx.array:
        enc = tok(
            prompt,
            padding="max_length",
            max_length=MAX_TOKENS,
            truncation=True,
            return_tensors="np",
        )
        return mx.array(enc["input_ids"].astype("int32"))

    hs1, _ = self.text_encoder(ids(self.tokenizer))
    hs2, pooled = self.text_encoder_2(ids(self.tokenizer_2))
    embeds = mx.concatenate([hs1[-2], hs2[-2]], axis=-1)  # (1, 77, 768+1280)
    return embeds, pooled

from_diffusers(folder, *, dtype=mx.float16, quantize_unet=None) classmethod

Load + convert a stable-diffusion-xl-base folder into MLX.

The UNet and text encoders run in dtype (fp16 by default, SDXL's native precision); the VAE runs in fp32 because SDXL's VAE overflows fp16. quantize_unet (4/8) weight-quantizes the UNet to save memory.

Source code in src/mlx_diffuser/pipelines/sdxl.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@classmethod
def from_diffusers(
    cls,
    folder: str | Path,
    *,
    dtype: mx.Dtype = mx.float16,
    quantize_unet: int | None = None,
) -> StableDiffusionXLPipeline:
    """Load + convert a ``stable-diffusion-xl-base`` folder into MLX.

    The UNet and text encoders run in ``dtype`` (fp16 by default, SDXL's native
    precision); the VAE runs in fp32 because SDXL's VAE overflows fp16.
    ``quantize_unet`` (4/8) weight-quantizes the UNet to save memory.
    """
    from ..converters import get_converter

    folder = Path(folder)
    unet = cast(
        SDXLUNet,
        get_converter("UNet2DConditionModel").convert(
            folder / "unet", dtype=dtype, quantize=quantize_unet
        ),
    )
    vae = cast(
        AutoencoderKLSD,
        get_converter("AutoencoderKL").convert(folder / "vae", dtype=mx.float32),
    )
    te = cast(
        CLIPTextModel,
        get_converter("CLIPTextModel").convert(folder / "text_encoder", dtype=dtype),
    )
    te2 = cast(
        CLIPTextModel,
        get_converter("CLIPTextModelWithProjection").convert(
            folder / "text_encoder_2", dtype=dtype
        ),
    )
    tok = _load_tokenizer(folder / "tokenizer")
    tok2 = _load_tokenizer(folder / "tokenizer_2")
    scheduler = EulerDiscreteScheduler(
        EulerConfig(
            beta_schedule="scaled_linear",
            beta_start=0.00085,
            beta_end=0.012,
            prediction_type="epsilon",
            timestep_spacing="leading",
            steps_offset=1,
        )
    )
    return cls(unet, vae, te, te2, tok, tok2, scheduler)

mlx_diffuser.pipelines.FluxPipeline

FLUX.1 text-to-image pipeline (channels-last (B, H, W, C)).

Source code in src/mlx_diffuser/pipelines/flux.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class FluxPipeline:
    """FLUX.1 text-to-image pipeline (channels-last ``(B, H, W, C)``)."""

    def __init__(
        self,
        transformer: FluxTransformer2DModel,
        vae: AutoencoderKLSD,
        text_encoder: CLIPTextModel,
        text_encoder_2: T5EncoderModel,
        tokenizer,
        tokenizer_2,
        scheduler: FlowMatchEulerScheduler,
        shift_config: dict,
    ):
        self.transformer = transformer
        self.vae = vae
        self.text_encoder = text_encoder
        self.text_encoder_2 = text_encoder_2
        self.tokenizer = tokenizer
        self.tokenizer_2 = tokenizer_2
        self.scheduler = scheduler
        self.shift_config = shift_config

    # --- loading -------------------------------------------------------------
    @classmethod
    def from_diffusers(
        cls,
        folder: str | Path,
        *,
        dtype: mx.Dtype = mx.bfloat16,
        quantize_transformer: int | None = 4,
        quantize_t5: int | None = 4,
    ) -> FluxPipeline:
        """Load + convert a ``FLUX.1-schnell`` / ``FLUX.1-dev`` folder into MLX.

        The transformer and CLIP run in ``dtype`` (bf16); the VAE runs in fp32. The
        12B transformer and the T5 encoder default to 4-bit weights so the pipeline
        fits in ~10 GB; pass ``quantize_transformer=None`` / ``quantize_t5=None`` for
        full precision (needs far more memory).
        """
        from ..converters import get_converter

        folder = Path(folder)
        transformer = cast(
            FluxTransformer2DModel,
            get_converter("FluxTransformer2DModel").convert(
                folder / "transformer", dtype=dtype, quantize=quantize_transformer
            ),
        )
        vae = cast(
            AutoencoderKLSD,
            get_converter("AutoencoderKL").convert(folder / "vae", dtype=mx.float32),
        )
        te = cast(
            CLIPTextModel,
            get_converter("CLIPTextModel").convert(folder / "text_encoder", dtype=dtype),
        )
        te2 = cast(
            T5EncoderModel,
            get_converter("T5EncoderModel").convert(
                folder / "text_encoder_2", dtype=dtype, quantize=quantize_t5
            ),
        )
        tok = _load_tokenizer(folder / "tokenizer")
        tok2 = _load_tokenizer(folder / "tokenizer_2")
        shift_config = _load_scheduler_config(folder / "scheduler")
        scheduler = FlowMatchEulerScheduler(FlowMatchConfig(prediction_type="velocity"))
        return cls(transformer, vae, te, te2, tok, tok2, scheduler, shift_config)

    def release_text_encoders(self) -> None:
        """Free the CLIP + T5 encoders after encoding (they are ~half the footprint)."""
        self.text_encoder = None  # type: ignore[assignment]
        self.text_encoder_2 = None  # type: ignore[assignment]
        gc.collect()
        mx.clear_cache()

    # --- text encoding -------------------------------------------------------
    def encode_prompt(self, prompt: str, max_sequence_length: int) -> tuple[mx.array, mx.array]:
        """Encode ``prompt`` with CLIP (pooled) and T5 (sequence).

        Returns ``(t5_embeds (1, L, 4096), pooled (1, 768))``: FLUX conditions joint
        attention on the T5 token sequence and modulation on CLIP's pooled embedding.
        """
        clip_ids = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=CLIP_MAX_TOKENS,
            truncation=True,
            return_tensors="np",
        )["input_ids"].astype("int32")
        _, pooled = self.text_encoder(mx.array(clip_ids))

        t5_ids = self.tokenizer_2(
            prompt,
            padding="max_length",
            max_length=max_sequence_length,
            truncation=True,
            return_tensors="np",
        )["input_ids"].astype("int32")
        t5_embeds = self.text_encoder_2(mx.array(t5_ids))
        return t5_embeds, pooled

    # --- generation ----------------------------------------------------------
    def __call__(
        self,
        prompt: str,
        *,
        height: int = 1024,
        width: int = 1024,
        num_inference_steps: int = 4,
        guidance_scale: float = 0.0,
        max_sequence_length: int = 256,
        seed: int = 0,
        cache_threshold: float = 0.0,
        release_text_encoders: bool = False,
        tile_vae: bool = False,
        progress: bool = True,
    ) -> mx.array:
        """Generate an image ``(1, height, width, 3)`` in ``[-1, 1]``.

        ``height`` / ``width`` must be multiples of 16 (8× VAE × 2× patch packing).
        For schnell use the defaults (4 steps, ``guidance_scale=0``); for dev use ~50
        steps and ``guidance_scale≈3.5``. ``cache_threshold`` (>0) enables First-Block
        caching for a speed-up. ``release_text_encoders`` frees CLIP + T5 right after
        encoding to lower peak memory before the denoising loop. ``tile_vae`` decodes the
        final latent in tiles: at 1024px the fp32 VAE decode is the memory high-water mark
        (~18 GB), and tiling brings it under 16 GB so it fits a 16 GB Mac without swapping.
        """
        if height % 16 or width % 16:
            raise ValueError("height and width must be multiples of 16.")

        t5_embeds, pooled = self.encode_prompt(prompt, max_sequence_length)
        if release_text_encoders:
            self.release_text_encoders()

        lh, lw = height // 8, width // 8  # VAE latent grid
        ph, pw = lh // 2, lw // 2  # packed token grid
        c = self.vae.config.latent_channels
        latents = mx.random.normal((1, lh, lw, c), key=mx.random.key(seed))
        latents = _pack_latents(latents)  # (1, ph*pw, c*4)

        img_ids = _prepare_image_ids(ph, pw)
        txt_ids = mx.zeros((t5_embeds.shape[1], 3))

        sigmas = self._sigma_schedule(num_inference_steps, ph * pw)
        self.scheduler.set_sigmas(sigmas)
        steps = self.scheduler.timesteps
        assert steps is not None

        guidance = None
        if self.transformer.config.guidance_embeds:
            guidance = mx.full((1,), guidance_scale, dtype=mx.float32)

        cache = FirstBlockCache(cache_threshold) if cache_threshold > 0 else None
        bar = _DenoiseProgress(len(steps), enabled=progress)
        for t in steps:
            v = self.transformer(
                latents,
                mx.broadcast_to(t, (1,)),
                t5_embeds,
                pooled,
                img_ids,
                txt_ids,
                guidance=guidance,
                cache=cache,
            )
            latents = self.scheduler.step(v, t, latents)
            mx.eval(latents)
            bar.update(cache)
        bar.close()

        latents = _unpack_latents(latents, ph, pw, c)  # (1, lh, lw, c)
        z = latents.astype(mx.float32) / self.vae.scaling_factor + self.vae.shift_factor
        image = self.vae.decode(z, tile=tile_vae, tile_latent=48, overlap_latent=8)
        mx.eval(image)
        return image

    def _sigma_schedule(self, num_steps: int, image_seq_len: int) -> mx.array:
        """FLUX sigma schedule: ``linspace(1, 1/N, N)`` with resolution-dependent shift."""
        sc = self.shift_config
        sigmas = mx.linspace(1.0, 1.0 / num_steps, num_steps)
        if sc.get("use_dynamic_shifting", False):
            m = (sc["max_shift"] - sc["base_shift"]) / (
                sc["max_image_seq_len"] - sc["base_image_seq_len"]
            )
            mu = sc["base_shift"] + m * (image_seq_len - sc["base_image_seq_len"])
            shift = math.exp(mu)
            sigmas = shift / (shift + (1.0 / sigmas - 1.0))  # exponential time shift
        else:
            s = sc.get("shift", 1.0)
            if s != 1.0:
                sigmas = s * sigmas / (1.0 + (s - 1.0) * sigmas)
        return sigmas

__call__(prompt, *, height=1024, width=1024, num_inference_steps=4, guidance_scale=0.0, max_sequence_length=256, seed=0, cache_threshold=0.0, release_text_encoders=False, tile_vae=False, progress=True)

Generate an image (1, height, width, 3) in [-1, 1].

height / width must be multiples of 16 (8× VAE × 2× patch packing). For schnell use the defaults (4 steps, guidance_scale=0); for dev use ~50 steps and guidance_scale≈3.5. cache_threshold (>0) enables First-Block caching for a speed-up. release_text_encoders frees CLIP + T5 right after encoding to lower peak memory before the denoising loop. tile_vae decodes the final latent in tiles: at 1024px the fp32 VAE decode is the memory high-water mark (~18 GB), and tiling brings it under 16 GB so it fits a 16 GB Mac without swapping.

Source code in src/mlx_diffuser/pipelines/flux.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def __call__(
    self,
    prompt: str,
    *,
    height: int = 1024,
    width: int = 1024,
    num_inference_steps: int = 4,
    guidance_scale: float = 0.0,
    max_sequence_length: int = 256,
    seed: int = 0,
    cache_threshold: float = 0.0,
    release_text_encoders: bool = False,
    tile_vae: bool = False,
    progress: bool = True,
) -> mx.array:
    """Generate an image ``(1, height, width, 3)`` in ``[-1, 1]``.

    ``height`` / ``width`` must be multiples of 16 (8× VAE × 2× patch packing).
    For schnell use the defaults (4 steps, ``guidance_scale=0``); for dev use ~50
    steps and ``guidance_scale≈3.5``. ``cache_threshold`` (>0) enables First-Block
    caching for a speed-up. ``release_text_encoders`` frees CLIP + T5 right after
    encoding to lower peak memory before the denoising loop. ``tile_vae`` decodes the
    final latent in tiles: at 1024px the fp32 VAE decode is the memory high-water mark
    (~18 GB), and tiling brings it under 16 GB so it fits a 16 GB Mac without swapping.
    """
    if height % 16 or width % 16:
        raise ValueError("height and width must be multiples of 16.")

    t5_embeds, pooled = self.encode_prompt(prompt, max_sequence_length)
    if release_text_encoders:
        self.release_text_encoders()

    lh, lw = height // 8, width // 8  # VAE latent grid
    ph, pw = lh // 2, lw // 2  # packed token grid
    c = self.vae.config.latent_channels
    latents = mx.random.normal((1, lh, lw, c), key=mx.random.key(seed))
    latents = _pack_latents(latents)  # (1, ph*pw, c*4)

    img_ids = _prepare_image_ids(ph, pw)
    txt_ids = mx.zeros((t5_embeds.shape[1], 3))

    sigmas = self._sigma_schedule(num_inference_steps, ph * pw)
    self.scheduler.set_sigmas(sigmas)
    steps = self.scheduler.timesteps
    assert steps is not None

    guidance = None
    if self.transformer.config.guidance_embeds:
        guidance = mx.full((1,), guidance_scale, dtype=mx.float32)

    cache = FirstBlockCache(cache_threshold) if cache_threshold > 0 else None
    bar = _DenoiseProgress(len(steps), enabled=progress)
    for t in steps:
        v = self.transformer(
            latents,
            mx.broadcast_to(t, (1,)),
            t5_embeds,
            pooled,
            img_ids,
            txt_ids,
            guidance=guidance,
            cache=cache,
        )
        latents = self.scheduler.step(v, t, latents)
        mx.eval(latents)
        bar.update(cache)
    bar.close()

    latents = _unpack_latents(latents, ph, pw, c)  # (1, lh, lw, c)
    z = latents.astype(mx.float32) / self.vae.scaling_factor + self.vae.shift_factor
    image = self.vae.decode(z, tile=tile_vae, tile_latent=48, overlap_latent=8)
    mx.eval(image)
    return image

encode_prompt(prompt, max_sequence_length)

Encode prompt with CLIP (pooled) and T5 (sequence).

Returns (t5_embeds (1, L, 4096), pooled (1, 768)): FLUX conditions joint attention on the T5 token sequence and modulation on CLIP's pooled embedding.

Source code in src/mlx_diffuser/pipelines/flux.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def encode_prompt(self, prompt: str, max_sequence_length: int) -> tuple[mx.array, mx.array]:
    """Encode ``prompt`` with CLIP (pooled) and T5 (sequence).

    Returns ``(t5_embeds (1, L, 4096), pooled (1, 768))``: FLUX conditions joint
    attention on the T5 token sequence and modulation on CLIP's pooled embedding.
    """
    clip_ids = self.tokenizer(
        prompt,
        padding="max_length",
        max_length=CLIP_MAX_TOKENS,
        truncation=True,
        return_tensors="np",
    )["input_ids"].astype("int32")
    _, pooled = self.text_encoder(mx.array(clip_ids))

    t5_ids = self.tokenizer_2(
        prompt,
        padding="max_length",
        max_length=max_sequence_length,
        truncation=True,
        return_tensors="np",
    )["input_ids"].astype("int32")
    t5_embeds = self.text_encoder_2(mx.array(t5_ids))
    return t5_embeds, pooled

from_diffusers(folder, *, dtype=mx.bfloat16, quantize_transformer=4, quantize_t5=4) classmethod

Load + convert a FLUX.1-schnell / FLUX.1-dev folder into MLX.

The transformer and CLIP run in dtype (bf16); the VAE runs in fp32. The 12B transformer and the T5 encoder default to 4-bit weights so the pipeline fits in ~10 GB; pass quantize_transformer=None / quantize_t5=None for full precision (needs far more memory).

Source code in src/mlx_diffuser/pipelines/flux.py 
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
@classmethod
def from_diffusers(
    cls,
    folder: str | Path,
    *,
    dtype: mx.Dtype = mx.bfloat16,
    quantize_transformer: int | None = 4,
    quantize_t5: int | None = 4,
) -> FluxPipeline:
    """Load + convert a ``FLUX.1-schnell`` / ``FLUX.1-dev`` folder into MLX.

    The transformer and CLIP run in ``dtype`` (bf16); the VAE runs in fp32. The
    12B transformer and the T5 encoder default to 4-bit weights so the pipeline
    fits in ~10 GB; pass ``quantize_transformer=None`` / ``quantize_t5=None`` for
    full precision (needs far more memory).
    """
    from ..converters import get_converter

    folder = Path(folder)
    transformer = cast(
        FluxTransformer2DModel,
        get_converter("FluxTransformer2DModel").convert(
            folder / "transformer", dtype=dtype, quantize=quantize_transformer
        ),
    )
    vae = cast(
        AutoencoderKLSD,
        get_converter("AutoencoderKL").convert(folder / "vae", dtype=mx.float32),
    )
    te = cast(
        CLIPTextModel,
        get_converter("CLIPTextModel").convert(folder / "text_encoder", dtype=dtype),
    )
    te2 = cast(
        T5EncoderModel,
        get_converter("T5EncoderModel").convert(
            folder / "text_encoder_2", dtype=dtype, quantize=quantize_t5
        ),
    )
    tok = _load_tokenizer(folder / "tokenizer")
    tok2 = _load_tokenizer(folder / "tokenizer_2")
    shift_config = _load_scheduler_config(folder / "scheduler")
    scheduler = FlowMatchEulerScheduler(FlowMatchConfig(prediction_type="velocity"))
    return cls(transformer, vae, te, te2, tok, tok2, scheduler, shift_config)

release_text_encoders()

Free the CLIP + T5 encoders after encoding (they are ~half the footprint).

Source code in src/mlx_diffuser/pipelines/flux.py 
105
106
107
108
109
110
def release_text_encoders(self) -> None:
    """Free the CLIP + T5 encoders after encoding (they are ~half the footprint)."""
    self.text_encoder = None  # type: ignore[assignment]
    self.text_encoder_2 = None  # type: ignore[assignment]
    gc.collect()
    mx.clear_cache()

Inference caching

mlx_diffuser.caching.FirstBlockCache

Decide, per step, whether to reuse the cached transformer residual.

Parameters:

Name Type Description Default
threshold float

accumulated relative first-block change that triggers a full recompute. 0 disables caching (always full / exact). On WAN 2.1 1.3B (256px, 25 steps) 0.1 ≈ 1.5x and 0.2 ≈ 2.2x with no visible quality change (the sample differs but stays sharp and coherent); >= 0.3 starts to degrade. Higher = faster, lower fidelity.

 0.0
Source code in src/mlx_diffuser/caching.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
class FirstBlockCache:
    """Decide, per step, whether to reuse the cached transformer residual.

    Args:
        threshold: accumulated relative first-block change that triggers a full
            recompute. ``0`` disables caching (always full / exact). On WAN 2.1 1.3B
            (256px, 25 steps) ``0.1`` ≈ 1.5x and ``0.2`` ≈ 2.2x with no visible quality
            change (the sample differs but stays sharp and coherent); ``>= 0.3``
            starts to degrade. Higher = faster, lower fidelity.
    """

    def __init__(self, threshold: float = 0.0):
        self.threshold = threshold
        self.prev_first: mx.array | None = None
        self.residual: mx.array | None = None
        self._accum = 0.0
        self.steps = 0
        self.skipped = 0

    def reset(self) -> None:
        """Clear state before a new generation."""
        self.prev_first = None
        self.residual = None
        self._accum = 0.0
        self.steps = 0
        self.skipped = 0

    @property
    def enabled(self) -> bool:
        return self.threshold > 0.0

    def should_reuse(self, first_residual: mx.array) -> bool:
        """Given this step's first-block *contribution*, return whether to reuse.

        Accumulates the relative L1 change of the first block's residual (its output
        minus its input); while the running total stays under ``threshold`` we reuse,
        and we force a recompute (resetting the accumulator) once it crosses. The
        first step, and any step before a residual exists, always recomputes.
        """
        self.steps += 1
        if not self.enabled or self.prev_first is None or self.residual is None:
            self.prev_first = first_residual
            return False
        rel = mx.mean(mx.abs(first_residual - self.prev_first)) / (
            mx.mean(mx.abs(self.prev_first)) + 1e-8
        )
        self.prev_first = first_residual
        self._accum += float(rel)  # one tiny scalar sync per step (~nothing vs a block)
        if self._accum < self.threshold:
            self.skipped += 1
            return True
        self._accum = 0.0
        return False

reset()

Clear state before a new generation.

Source code in src/mlx_diffuser/caching.py 
76
77
78
79
80
81
82
def reset(self) -> None:
    """Clear state before a new generation."""
    self.prev_first = None
    self.residual = None
    self._accum = 0.0
    self.steps = 0
    self.skipped = 0

should_reuse(first_residual)

Given this step's first-block contribution, return whether to reuse.

Accumulates the relative L1 change of the first block's residual (its output minus its input); while the running total stays under threshold we reuse, and we force a recompute (resetting the accumulator) once it crosses. The first step, and any step before a residual exists, always recomputes.

Source code in src/mlx_diffuser/caching.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def should_reuse(self, first_residual: mx.array) -> bool:
    """Given this step's first-block *contribution*, return whether to reuse.

    Accumulates the relative L1 change of the first block's residual (its output
    minus its input); while the running total stays under ``threshold`` we reuse,
    and we force a recompute (resetting the accumulator) once it crosses. The
    first step, and any step before a residual exists, always recomputes.
    """
    self.steps += 1
    if not self.enabled or self.prev_first is None or self.residual is None:
        self.prev_first = first_residual
        return False
    rel = mx.mean(mx.abs(first_residual - self.prev_first)) / (
        mx.mean(mx.abs(self.prev_first)) + 1e-8
    )
    self.prev_first = first_residual
    self._accum += float(rel)  # one tiny scalar sync per step (~nothing vs a block)
    if self._accum < self.threshold:
        self.skipped += 1
        return True
    self._accum = 0.0
    return False

mlx_diffuser.caching.DeepCache

DeepCache for U-Net diffusion: skip the deep blocks on most steps.

A U-Net's deep (bottleneck) features change slowly across denoising steps, while the shallow blocks carry the high-frequency detail. DeepCache recomputes the full network only every interval steps; in between it runs just the shallowest down/up blocks and reuses the cached deep feature — skipping the most expensive levels (for SDXL, the 1280-channel / 10-transformer-layer blocks).

interval = 1 disables caching (every step full). interval = 2 caches every other step (~1.5-1.8x). The first step is always full (no cache yet).

Source code in src/mlx_diffuser/caching.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class DeepCache:
    """DeepCache for U-Net diffusion: skip the deep blocks on most steps.

    A U-Net's deep (bottleneck) features change slowly across denoising steps, while
    the shallow blocks carry the high-frequency detail. DeepCache recomputes the full
    network only every ``interval`` steps; in between it runs just the shallowest
    down/up blocks and reuses the cached deep feature — skipping the most expensive
    levels (for SDXL, the 1280-channel / 10-transformer-layer blocks).

    ``interval = 1`` disables caching (every step full). ``interval = 2`` caches every
    other step (~1.5-1.8x). The first step is always full (no cache yet).
    """

    def __init__(self, interval: int = 1):
        self.interval = interval
        self.deep: mx.array | None = None
        self.steps = 0
        self.skipped = 0

    @property
    def enabled(self) -> bool:
        return self.interval > 1

    def reset(self) -> None:
        self.deep = None
        self.steps = 0
        self.skipped = 0

    def should_reuse(self) -> bool:
        """Return whether this step should reuse the cached deep feature."""
        reuse = self.enabled and self.deep is not None and (self.steps % self.interval != 0)
        self.steps += 1
        if reuse:
            self.skipped += 1
        return reuse

should_reuse()

Return whether this step should reuse the cached deep feature.

Source code in src/mlx_diffuser/caching.py 
48
49
50
51
52
53
54
def should_reuse(self) -> bool:
    """Return whether this step should reuse the cached deep feature."""
    reuse = self.enabled and self.deep is not None and (self.steps % self.interval != 0)
    self.steps += 1
    if reuse:
        self.skipped += 1
    return reuse