Pipelines

mlx_diffuser.pipelines.DiffusionPipeline

Source code in src/mlx_diffuser/pipelines/base.py

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

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

@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

@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

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

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

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

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

@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

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

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

    def release_text_encoder_models(self) -> None:
        """Release both CLIP encoders after prompt encoding to reduce peak memory."""
        self.text_encoder = None  # type: ignore[assignment]
        self.text_encoder_2 = None  # type: ignore[assignment]
        gc.collect()
        mx.clear_cache()

    # --- generation ----------------------------------------------------------
    def __call__(
        self,
        prompt: str,
        *,
        negative_prompt: str = "",
        image=None,
        strength: float = 0.8,
        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,
        release_text_encoders: bool = False,
        progress: bool = True,
    ) -> mx.array:
        """Generate an image ``(1, height, width, 3)`` in ``[-1, 1]``.

        Pass ``image`` (a path, PIL image, or HWC/BHWC array) for image-to-image;
        ``strength`` controls how far the result may move from the input (0 < strength
        <= 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.
        ``release_text_encoders`` reclaims both CLIP encoders before denoising.
        """
        if height % 8 or width % 8:
            raise ValueError("height and width must be multiples of 8.")
        if image is not None and not 0.0 < strength <= 1.0:
            raise ValueError("strength must be in the interval (0, 1].")

        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]
        if release_text_encoders:
            self.release_text_encoder_models()
        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
        if image is None:
            latents = mx.random.normal(
                (1, height // 8, width // 8, self.vae.config.latent_channels),
                key=mx.random.key(seed),
            )
            latents = latents * self.scheduler.init_noise_sigma
        else:
            input_image = prepare_image(image, height=height, width=width, dtype=mx.float32)
            if input_image.shape[0] != 1:
                raise ValueError("SDXL image-to-image currently accepts exactly one input image.")
            latent_key, noise_key = mx.random.split(mx.random.key(seed))
            clean_latents = (
                self.vae.encode(input_image).sample(latent_key) * self.vae.scaling_factor
            )
            noise = mx.random.normal(clean_latents.shape, key=noise_key)
            denoise_steps = max(1, min(num_inference_steps, int(num_inference_steps * strength)))
            begin_index = num_inference_steps - denoise_steps
            assert self.scheduler.sigmas is not None
            latents = self.scheduler.add_noise_sigma(
                clean_latents, noise, self.scheduler.sigmas[begin_index]
            )
            self.scheduler.set_begin_index(begin_index)
            steps = steps[begin_index:]
            mx.eval(latents)

        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='', image=None, strength=0.8, height=1024, width=1024, num_inference_steps=30, guidance_scale=5.0, seed=0, cache_interval=1, tile_vae=False, release_text_encoders=False, progress=True)

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

Pass image (a path, PIL image, or HWC/BHWC array) for image-to-image; strength controls how far the result may move from the input (0 < strength <= 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. release_text_encoders reclaims both CLIP encoders before denoising.

Source code in src/mlx_diffuser/pipelines/sdxl.py

def __call__(
    self,
    prompt: str,
    *,
    negative_prompt: str = "",
    image=None,
    strength: float = 0.8,
    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,
    release_text_encoders: bool = False,
    progress: bool = True,
) -> mx.array:
    """Generate an image ``(1, height, width, 3)`` in ``[-1, 1]``.

    Pass ``image`` (a path, PIL image, or HWC/BHWC array) for image-to-image;
    ``strength`` controls how far the result may move from the input (0 < strength
    <= 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.
    ``release_text_encoders`` reclaims both CLIP encoders before denoising.
    """
    if height % 8 or width % 8:
        raise ValueError("height and width must be multiples of 8.")
    if image is not None and not 0.0 < strength <= 1.0:
        raise ValueError("strength must be in the interval (0, 1].")

    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]
    if release_text_encoders:
        self.release_text_encoder_models()
    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
    if image is None:
        latents = mx.random.normal(
            (1, height // 8, width // 8, self.vae.config.latent_channels),
            key=mx.random.key(seed),
        )
        latents = latents * self.scheduler.init_noise_sigma
    else:
        input_image = prepare_image(image, height=height, width=width, dtype=mx.float32)
        if input_image.shape[0] != 1:
            raise ValueError("SDXL image-to-image currently accepts exactly one input image.")
        latent_key, noise_key = mx.random.split(mx.random.key(seed))
        clean_latents = (
            self.vae.encode(input_image).sample(latent_key) * self.vae.scaling_factor
        )
        noise = mx.random.normal(clean_latents.shape, key=noise_key)
        denoise_steps = max(1, min(num_inference_steps, int(num_inference_steps * strength)))
        begin_index = num_inference_steps - denoise_steps
        assert self.scheduler.sigmas is not None
        latents = self.scheduler.add_noise_sigma(
            clean_latents, noise, self.scheduler.sigmas[begin_index]
        )
        self.scheduler.set_begin_index(begin_index)
        steps = steps[begin_index:]
        mx.eval(latents)

    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

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

@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)

release_text_encoder_models()

Release both CLIP encoders after prompt encoding to reduce peak memory.

Source code in src/mlx_diffuser/pipelines/sdxl.py

def release_text_encoder_models(self) -> None:
    """Release both CLIP encoders after prompt encoding to reduce peak memory."""
    self.text_encoder = None  # type: ignore[assignment]
    self.text_encoder_2 = None  # type: ignore[assignment]
    gc.collect()
    mx.clear_cache()

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

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

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

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

@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

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()

mlx_diffuser.pipelines.LTX2Pipeline

LTX-2.3 text-to-video-with-audio (channels-last MLX tensors).

Calling the pipeline returns (video, audio): frames (1, F, H, W, 3) in [-1, 1] and a 48 kHz stereo waveform (2, samples) in [-1, 1].

Source code in src/mlx_diffuser/pipelines/ltx2.py

class LTX2Pipeline:
    """LTX-2.3 text-to-video-with-audio (channels-last MLX tensors).

    Calling the pipeline returns ``(video, audio)``: frames ``(1, F, H, W, 3)``
    in [-1, 1] and a 48 kHz stereo waveform ``(2, samples)`` in [-1, 1].
    """

    def __init__(self, folder: str | Path):
        self.folder = Path(folder)
        required = (
            "transformer",
            "connectors",
            "vae_decoder",
            "audio_decoder",
            "vocoder",
            "text_encoder",
            "tokenizer",
        )
        for sub in required:
            if not (self.folder / sub).exists():
                raise FileNotFoundError(
                    f"{self.folder / sub} not found — run the LTX-2.3 conversion first "
                    "(mlx-diffuser generate --model ltx-2.3 --download; it only fetches "
                    "components that are missing)."
                )
        self.text_encoder: Gemma3TextEncoder | None = None
        self.connectors: LTX2TextConnectors | None = None
        self.transformer: LTX2Transformer3DModel | None = None
        self.vae: LTX2VideoDecoder | None = None
        self.audio_decoder: LTX2AudioDecoder | None = None
        self.vocoder: LTX2Vocoder | None = None
        self._tokenizer = None

    @classmethod
    def from_converted(cls, folder: str | Path) -> LTX2Pipeline:
        """Open a converted LTX-2.3 checkpoint folder (components load lazily)."""
        return cls(folder)

    # --- staged component management ------------------------------------------
    def _release(self, *names: str) -> None:
        for name in names:
            setattr(self, name, None)
        gc.collect()
        mx.clear_cache()

    def _load_tokenizer(self):
        if self._tokenizer is None:
            try:
                from transformers import AutoTokenizer
            except ImportError as exc:  # pragma: no cover
                raise ImportError(
                    "The LTX-2 pipeline needs `transformers` for tokenization."
                ) from exc
            self._tokenizer = AutoTokenizer.from_pretrained(str(self.folder / "tokenizer"))
            self._tokenizer.padding_side = "left"
            if self._tokenizer.pad_token is None:
                self._tokenizer.pad_token = self._tokenizer.eos_token
        return self._tokenizer

    # --- text encoding ----------------------------------------------------------
    def encode_prompt(self, prompt: str) -> tuple[mx.array, mx.array]:
        """Encode ``prompt`` into per-modality text streams (video, audio).

        Loads the Gemma encoder + connectors on first use; call
        :meth:`release_text_stack` afterwards to reclaim ~8 GB.
        """
        tokenizer = self._load_tokenizer()
        if self.text_encoder is None:
            self.text_encoder = Gemma3TextEncoder.from_pretrained(self.folder / "text_encoder")
        if self.connectors is None:
            self.connectors = LTX2TextConnectors.from_pretrained(self.folder / "connectors")

        enc = tokenizer(
            prompt.strip(),
            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"))
        hidden = self.text_encoder(ids, mask)  # (1, L, D, 49)
        video_text, audio_text = self.connectors(hidden, mask)
        mx.eval(video_text, audio_text)
        return video_text, audio_text

    def release_text_stack(self) -> None:
        self._release("text_encoder", "connectors")

    # --- generation ---------------------------------------------------------------
    def __call__(
        self,
        prompt: str,
        *,
        negative_prompt: str = "",
        height: int = 512,
        width: int = 768,
        num_frames: int = 121,
        frame_rate: float = 24.0,
        guidance_scale: float = 1.0,
        seed: int = 0,
        cache_threshold: float = 0.0,
        release_stages: bool = True,
        progress: bool = True,
    ) -> tuple[mx.array, mx.array]:
        """Generate ``(video, audio)`` for ``prompt``.

        ``video`` is ``(1, num_frames, height, width, 3)`` in [-1, 1]; ``audio``
        is a 48 kHz stereo waveform ``(2, samples)`` in [-1, 1] covering the
        same duration. ``height``/``width`` must be multiples of 32 and
        ``num_frames`` must be ``1 + 8*k`` (the VAE's compression grid). The
        distilled model runs its fixed 8-step schedule at ``guidance_scale=1``
        (no CFG); values > 1 add a classifier-free pass against
        ``negative_prompt``.

        ``release_stages`` frees each component when the next stage begins —
        essential on 16 GB machines, only worth disabling for repeated calls on
        big-memory hosts.
        """
        if height % 32 or width % 32:
            raise ValueError("height and width must be multiples of 32.")
        if (num_frames - 1) % 8 != 0:
            raise ValueError("num_frames must be 1 + a multiple of 8 (e.g. 57, 97, 121).")

        use_cfg = guidance_scale > 1.0

        # Stage 1: text encoding (Gemma-3-12B + connectors), then free ~8 GB.
        video_text, audio_text = self.encode_prompt(prompt)
        if use_cfg:
            neg_video_text, neg_audio_text = self.encode_prompt(negative_prompt)
            video_text = mx.concatenate([video_text, neg_video_text], axis=0)
            audio_text = mx.concatenate([audio_text, neg_audio_text], axis=0)
        if release_stages:
            self.release_text_stack()

        # Stage 2: joint audio-video denoising with the 22B transformer.
        if self.transformer is None:
            self.transformer = LTX2Transformer3DModel.from_pretrained(self.folder / "transformer")
        cfg = self.transformer.config

        f = (num_frames - 1) // cfg.vae_scale_factors[0] + 1
        h = height // cfg.vae_scale_factors[1]
        w = width // cfg.vae_scale_factors[2]
        audio_frames = round(num_frames / frame_rate * AUDIO_LATENTS_PER_SECOND)

        key = mx.random.key(seed)
        k_video, k_audio = mx.random.split(key)
        latents = mx.random.normal((1, f * h * w, cfg.in_channels), key=k_video)
        audio_latents = mx.random.normal(
            (1, audio_frames, cfg.audio_in_channels),
            key=k_audio,  # 8 ch x 16 mel bins packed
        )

        video_coords, audio_coords = self.transformer.prepare_coords(
            video_text.shape[0], f, h, w, audio_frames, frame_rate
        )

        sigmas = np.array(DISTILLED_SIGMAS + [0.0], dtype=np.float32)
        cache = FirstBlockCache(cache_threshold) if cache_threshold > 0 else None
        bar = _DenoiseProgress(len(sigmas) - 1, enabled=progress)
        for i in range(len(sigmas) - 1):
            sigma, sigma_next = float(sigmas[i]), float(sigmas[i + 1])
            t = mx.full((video_text.shape[0],), sigma * 1000.0)
            if use_cfg:
                x_in = mx.concatenate([latents, latents], axis=0)
                a_in = mx.concatenate([audio_latents, audio_latents], axis=0)
            else:
                x_in, a_in = latents, audio_latents
            v_video, v_audio = self.transformer(
                x_in.astype(video_text.dtype),
                a_in.astype(audio_text.dtype),
                video_text,
                audio_text,
                t,
                video_coords,
                audio_coords,
                cache=cache,
            )
            v_video = v_video.astype(mx.float32)
            v_audio = v_audio.astype(mx.float32)
            if use_cfg:
                v_video = _cfg_x0(v_video, latents, sigma, guidance_scale)
                v_audio = _cfg_x0(v_audio, audio_latents, sigma, guidance_scale)
            latents = latents + (sigma_next - sigma) * v_video
            audio_latents = audio_latents + (sigma_next - sigma) * v_audio
            mx.eval(latents, audio_latents)
            bar.update(cache)
        bar.close()
        if release_stages:
            self._release("transformer")

        # Stage 3: audio decode (mel spectrogram -> 48 kHz stereo waveform).
        # Small models, but run in float32: bf16 accumulation across the
        # vocoder's 100+ sequential convolutions audibly degrades the spectrum.
        if self.audio_decoder is None:
            self.audio_decoder = LTX2AudioDecoder.from_pretrained(
                self.folder / "audio_decoder", dtype=mx.float32
            )
        if self.vocoder is None:
            self.vocoder = LTX2Vocoder.from_pretrained(self.folder / "vocoder", dtype=mx.float32)
        mel = self.audio_decoder.decode(audio_latents.astype(mx.float32))
        audio = self.vocoder(mel)[0].transpose(1, 0)  # (2, samples) at 48 kHz
        mx.eval(audio)
        if release_stages:
            self._release("audio_decoder", "vocoder")

        # Stage 4: unpack + denormalize + decode the video latents.
        if self.vae is None:
            self.vae = LTX2VideoDecoder.from_pretrained(
                self.folder / "vae_decoder", dtype=mx.bfloat16
            )
        z = latents.reshape(1, f, h, w, cfg.in_channels)
        z = self.vae.denormalize_latents(z.astype(mx.float32))
        video = self.vae.decode(z.astype(mx.bfloat16))
        video = mx.clip(video.astype(mx.float32), -1.0, 1.0)
        mx.eval(video)
        if release_stages:
            self._release("vae")
        return video, audio

__call__(prompt, *, negative_prompt='', height=512, width=768, num_frames=121, frame_rate=24.0, guidance_scale=1.0, seed=0, cache_threshold=0.0, release_stages=True, progress=True)

Generate (video, audio) for prompt.

video is (1, num_frames, height, width, 3) in [-1, 1]; audio is a 48 kHz stereo waveform (2, samples) in [-1, 1] covering the same duration. height/width must be multiples of 32 and num_frames must be 1 + 8*k (the VAE's compression grid). The distilled model runs its fixed 8-step schedule at guidance_scale=1 (no CFG); values > 1 add a classifier-free pass against negative_prompt.

release_stages frees each component when the next stage begins — essential on 16 GB machines, only worth disabling for repeated calls on big-memory hosts.

Source code in src/mlx_diffuser/pipelines/ltx2.py

def __call__(
    self,
    prompt: str,
    *,
    negative_prompt: str = "",
    height: int = 512,
    width: int = 768,
    num_frames: int = 121,
    frame_rate: float = 24.0,
    guidance_scale: float = 1.0,
    seed: int = 0,
    cache_threshold: float = 0.0,
    release_stages: bool = True,
    progress: bool = True,
) -> tuple[mx.array, mx.array]:
    """Generate ``(video, audio)`` for ``prompt``.

    ``video`` is ``(1, num_frames, height, width, 3)`` in [-1, 1]; ``audio``
    is a 48 kHz stereo waveform ``(2, samples)`` in [-1, 1] covering the
    same duration. ``height``/``width`` must be multiples of 32 and
    ``num_frames`` must be ``1 + 8*k`` (the VAE's compression grid). The
    distilled model runs its fixed 8-step schedule at ``guidance_scale=1``
    (no CFG); values > 1 add a classifier-free pass against
    ``negative_prompt``.

    ``release_stages`` frees each component when the next stage begins —
    essential on 16 GB machines, only worth disabling for repeated calls on
    big-memory hosts.
    """
    if height % 32 or width % 32:
        raise ValueError("height and width must be multiples of 32.")
    if (num_frames - 1) % 8 != 0:
        raise ValueError("num_frames must be 1 + a multiple of 8 (e.g. 57, 97, 121).")

    use_cfg = guidance_scale > 1.0

    # Stage 1: text encoding (Gemma-3-12B + connectors), then free ~8 GB.
    video_text, audio_text = self.encode_prompt(prompt)
    if use_cfg:
        neg_video_text, neg_audio_text = self.encode_prompt(negative_prompt)
        video_text = mx.concatenate([video_text, neg_video_text], axis=0)
        audio_text = mx.concatenate([audio_text, neg_audio_text], axis=0)
    if release_stages:
        self.release_text_stack()

    # Stage 2: joint audio-video denoising with the 22B transformer.
    if self.transformer is None:
        self.transformer = LTX2Transformer3DModel.from_pretrained(self.folder / "transformer")
    cfg = self.transformer.config

    f = (num_frames - 1) // cfg.vae_scale_factors[0] + 1
    h = height // cfg.vae_scale_factors[1]
    w = width // cfg.vae_scale_factors[2]
    audio_frames = round(num_frames / frame_rate * AUDIO_LATENTS_PER_SECOND)

    key = mx.random.key(seed)
    k_video, k_audio = mx.random.split(key)
    latents = mx.random.normal((1, f * h * w, cfg.in_channels), key=k_video)
    audio_latents = mx.random.normal(
        (1, audio_frames, cfg.audio_in_channels),
        key=k_audio,  # 8 ch x 16 mel bins packed
    )

    video_coords, audio_coords = self.transformer.prepare_coords(
        video_text.shape[0], f, h, w, audio_frames, frame_rate
    )

    sigmas = np.array(DISTILLED_SIGMAS + [0.0], dtype=np.float32)
    cache = FirstBlockCache(cache_threshold) if cache_threshold > 0 else None
    bar = _DenoiseProgress(len(sigmas) - 1, enabled=progress)
    for i in range(len(sigmas) - 1):
        sigma, sigma_next = float(sigmas[i]), float(sigmas[i + 1])
        t = mx.full((video_text.shape[0],), sigma * 1000.0)
        if use_cfg:
            x_in = mx.concatenate([latents, latents], axis=0)
            a_in = mx.concatenate([audio_latents, audio_latents], axis=0)
        else:
            x_in, a_in = latents, audio_latents
        v_video, v_audio = self.transformer(
            x_in.astype(video_text.dtype),
            a_in.astype(audio_text.dtype),
            video_text,
            audio_text,
            t,
            video_coords,
            audio_coords,
            cache=cache,
        )
        v_video = v_video.astype(mx.float32)
        v_audio = v_audio.astype(mx.float32)
        if use_cfg:
            v_video = _cfg_x0(v_video, latents, sigma, guidance_scale)
            v_audio = _cfg_x0(v_audio, audio_latents, sigma, guidance_scale)
        latents = latents + (sigma_next - sigma) * v_video
        audio_latents = audio_latents + (sigma_next - sigma) * v_audio
        mx.eval(latents, audio_latents)
        bar.update(cache)
    bar.close()
    if release_stages:
        self._release("transformer")

    # Stage 3: audio decode (mel spectrogram -> 48 kHz stereo waveform).
    # Small models, but run in float32: bf16 accumulation across the
    # vocoder's 100+ sequential convolutions audibly degrades the spectrum.
    if self.audio_decoder is None:
        self.audio_decoder = LTX2AudioDecoder.from_pretrained(
            self.folder / "audio_decoder", dtype=mx.float32
        )
    if self.vocoder is None:
        self.vocoder = LTX2Vocoder.from_pretrained(self.folder / "vocoder", dtype=mx.float32)
    mel = self.audio_decoder.decode(audio_latents.astype(mx.float32))
    audio = self.vocoder(mel)[0].transpose(1, 0)  # (2, samples) at 48 kHz
    mx.eval(audio)
    if release_stages:
        self._release("audio_decoder", "vocoder")

    # Stage 4: unpack + denormalize + decode the video latents.
    if self.vae is None:
        self.vae = LTX2VideoDecoder.from_pretrained(
            self.folder / "vae_decoder", dtype=mx.bfloat16
        )
    z = latents.reshape(1, f, h, w, cfg.in_channels)
    z = self.vae.denormalize_latents(z.astype(mx.float32))
    video = self.vae.decode(z.astype(mx.bfloat16))
    video = mx.clip(video.astype(mx.float32), -1.0, 1.0)
    mx.eval(video)
    if release_stages:
        self._release("vae")
    return video, audio

encode_prompt(prompt)

Encode prompt into per-modality text streams (video, audio).

Loads the Gemma encoder + connectors on first use; call :meth:release_text_stack afterwards to reclaim ~8 GB.

Source code in src/mlx_diffuser/pipelines/ltx2.py

def encode_prompt(self, prompt: str) -> tuple[mx.array, mx.array]:
    """Encode ``prompt`` into per-modality text streams (video, audio).

    Loads the Gemma encoder + connectors on first use; call
    :meth:`release_text_stack` afterwards to reclaim ~8 GB.
    """
    tokenizer = self._load_tokenizer()
    if self.text_encoder is None:
        self.text_encoder = Gemma3TextEncoder.from_pretrained(self.folder / "text_encoder")
    if self.connectors is None:
        self.connectors = LTX2TextConnectors.from_pretrained(self.folder / "connectors")

    enc = tokenizer(
        prompt.strip(),
        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"))
    hidden = self.text_encoder(ids, mask)  # (1, L, D, 49)
    video_text, audio_text = self.connectors(hidden, mask)
    mx.eval(video_text, audio_text)
    return video_text, audio_text

from_converted(folder) classmethod

Open a converted LTX-2.3 checkpoint folder (components load lazily).

Source code in src/mlx_diffuser/pipelines/ltx2.py

@classmethod
def from_converted(cls, folder: str | Path) -> LTX2Pipeline:
    """Open a converted LTX-2.3 checkpoint folder (components load lazily)."""
    return cls(folder)

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

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

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

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

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

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