Skip to content

Models

mlx_diffuser.modeling.ModelMixin

Bases: Module, Generic[C]

Base class for all mlx-diffuser networks (generic over its config type).

Source code in src/mlx_diffuser/modeling.py 
 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
class ModelMixin(nn.Module, Generic[C]):
    """Base class for all mlx-diffuser networks (generic over its config type)."""

    #: The Config subclass this model is constructed from.
    config_class: type[C]

    config: C

    def save_pretrained(self, save_directory: str | Path, push_to_hub: str | None = None) -> Path:
        """Write ``config.json`` + ``model.safetensors`` into ``save_directory``.

        If ``push_to_hub`` is a repo id, the folder is uploaded afterwards.
        """
        save_directory = Path(save_directory)
        save_directory.mkdir(parents=True, exist_ok=True)
        self.config.save(save_directory)
        self.save_weights(str(save_directory / WEIGHTS_NAME))
        if push_to_hub:
            push_folder(save_directory, push_to_hub)
        return save_directory

    @classmethod
    def from_pretrained(
        cls: type[M],
        path_or_repo_id: str | Path,
        *,
        dtype: str | mx.Dtype | None = None,
        quantize: int | None = None,
        quant_group_size: int = 64,
        strict: bool = True,
        revision: str | None = None,
        **config_overrides: Any,
    ) -> M:
        """Load a model from a local directory or a Hub repo id.

        Args:
            path_or_repo_id: local dir (with ``config.json`` + weights) or Hub id.
            dtype: cast floating-point params to this dtype after loading.
            quantize: if set (2/3/4/6/8), weight-quantize after loading.
            quant_group_size: group size for quantization.
            strict: require the checkpoint to match the model's params exactly.
            revision: Hub revision (branch/tag/commit) when downloading.
            **config_overrides: override individual config fields at load time.
        """
        local = resolve(path_or_repo_id, revision=revision)
        config = cls.config_class.load(local)
        if config_overrides:
            config = config.replace(**config_overrides)

        model = cls(config)  # type: ignore[call-arg]  # subclasses take a config
        model.load_weights(str(local / WEIGHTS_NAME), strict=strict)

        resolved_dtype = as_dtype(dtype)
        if resolved_dtype is not None:
            model.set_dtype(resolved_dtype)
        if quantize is not None:
            quantize_module(model, bits=quantize, group_size=quant_group_size)

        mx.eval(model.parameters())
        model.eval()
        return model

    def num_parameters(self, trainable_only: bool = False) -> int:
        """Total number of parameter elements in the model."""
        params = self.trainable_parameters() if trainable_only else self.parameters()
        return sum(v.size for _, v in tree_flatten(params))  # type: ignore[union-attr, str-unpack]

    def __repr__(self) -> str:  # pragma: no cover - cosmetic
        n = self.num_parameters()
        return f"{type(self).__name__}({n / 1e6:.1f}M params)"

from_pretrained(path_or_repo_id, *, dtype=None, quantize=None, quant_group_size=64, strict=True, revision=None, **config_overrides) classmethod

Load a model from a local directory or a Hub repo id.

Parameters:

Name Type Description Default
path_or_repo_id str | Path

local dir (with config.json + weights) or Hub id.

 required
dtype str | Dtype | None

cast floating-point params to this dtype after loading.

 None
quantize int | None

if set (2/3/4/6/8), weight-quantize after loading.

 None
quant_group_size int

group size for quantization.

 64
strict bool

require the checkpoint to match the model's params exactly.

 True
revision str | None

Hub revision (branch/tag/commit) when downloading.

 None
**config_overrides Any

override individual config fields at load time.

 {}
Source code in src/mlx_diffuser/modeling.py 
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
@classmethod
def from_pretrained(
    cls: type[M],
    path_or_repo_id: str | Path,
    *,
    dtype: str | mx.Dtype | None = None,
    quantize: int | None = None,
    quant_group_size: int = 64,
    strict: bool = True,
    revision: str | None = None,
    **config_overrides: Any,
) -> M:
    """Load a model from a local directory or a Hub repo id.

    Args:
        path_or_repo_id: local dir (with ``config.json`` + weights) or Hub id.
        dtype: cast floating-point params to this dtype after loading.
        quantize: if set (2/3/4/6/8), weight-quantize after loading.
        quant_group_size: group size for quantization.
        strict: require the checkpoint to match the model's params exactly.
        revision: Hub revision (branch/tag/commit) when downloading.
        **config_overrides: override individual config fields at load time.
    """
    local = resolve(path_or_repo_id, revision=revision)
    config = cls.config_class.load(local)
    if config_overrides:
        config = config.replace(**config_overrides)

    model = cls(config)  # type: ignore[call-arg]  # subclasses take a config
    model.load_weights(str(local / WEIGHTS_NAME), strict=strict)

    resolved_dtype = as_dtype(dtype)
    if resolved_dtype is not None:
        model.set_dtype(resolved_dtype)
    if quantize is not None:
        quantize_module(model, bits=quantize, group_size=quant_group_size)

    mx.eval(model.parameters())
    model.eval()
    return model

num_parameters(trainable_only=False)

Total number of parameter elements in the model.

Source code in src/mlx_diffuser/modeling.py 
96
97
98
99
def num_parameters(self, trainable_only: bool = False) -> int:
    """Total number of parameter elements in the model."""
    params = self.trainable_parameters() if trainable_only else self.parameters()
    return sum(v.size for _, v in tree_flatten(params))  # type: ignore[union-attr, str-unpack]

save_pretrained(save_directory, push_to_hub=None)

Write config.json + model.safetensors into save_directory.

If push_to_hub is a repo id, the folder is uploaded afterwards.

Source code in src/mlx_diffuser/modeling.py 
42
43
44
45
46
47
48
49
50
51
52
53
def save_pretrained(self, save_directory: str | Path, push_to_hub: str | None = None) -> Path:
    """Write ``config.json`` + ``model.safetensors`` into ``save_directory``.

    If ``push_to_hub`` is a repo id, the folder is uploaded afterwards.
    """
    save_directory = Path(save_directory)
    save_directory.mkdir(parents=True, exist_ok=True)
    self.config.save(save_directory)
    self.save_weights(str(save_directory / WEIGHTS_NAME))
    if push_to_hub:
        push_folder(save_directory, push_to_hub)
    return save_directory

mlx_diffuser.models.DiT

Bases: ModelMixin[DiTConfig]

Source code in src/mlx_diffuser/models/dit.py 
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
class DiT(ModelMixin[DiTConfig]):
    config_class = DiTConfig

    def __init__(self, config: DiTConfig):
        super().__init__()
        self.config = config
        if config.hidden_size % 4 != 0:
            raise ValueError("hidden_size must be divisible by 4 (2D positional embedding).")

        self.patch_embed = PatchEmbed(config.in_channels, config.hidden_size, config.patch_size)
        self.t_embed = TimestepEmbedder(config.hidden_size)
        self.y_embed = (
            LabelEmbedder(config.num_classes, config.hidden_size, config.class_dropout_prob)
            if config.num_classes > 0
            else None
        )
        self.blocks = [
            DiTBlock(config.hidden_size, config.num_heads, config.mlp_ratio)
            for _ in range(config.depth)
        ]
        patch_dim = config.patch_size**2 * config.resolved_out_channels
        self.final = FinalLayer(config.hidden_size, patch_dim)

    def _unpatchify(self, x: mx.array, h: int, w: int) -> mx.array:
        p = self.config.patch_size
        c = self.config.resolved_out_channels
        b = x.shape[0]
        x = x.reshape(b, h, w, p, p, c)
        x = x.transpose(0, 1, 3, 2, 4, 5)  # (b, h, p, w, p, c)
        return x.reshape(b, h * p, w * p, c)

    def __call__(
        self,
        x: mx.array,
        t: mx.array,
        y: mx.array | None = None,
        *,
        training: bool = False,
        key: mx.array | None = None,
    ) -> mx.array:
        tokens, (h, w) = self.patch_embed(x)
        pos = get_2d_sincos_pos_embed(self.config.hidden_size, h, w)
        tokens = tokens + pos[None].astype(tokens.dtype)

        c = self.t_embed(t)
        if self.y_embed is not None:
            if y is None:
                raise ValueError("This DiT is class-conditional; pass labels `y`.")
            c = c + self.y_embed(y, training=training, key=key)

        for block in self.blocks:
            tokens = block(tokens, c)

        tokens = self.final(tokens, c)
        return self._unpatchify(tokens, h, w)

mlx_diffuser.models.DiTConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/dit.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@dataclasses.dataclass
class DiTConfig(Config):
    in_channels: int = 4
    out_channels: int | None = None
    patch_size: int = 2
    hidden_size: int = 384
    depth: int = 12
    num_heads: int = 6
    mlp_ratio: float = 4.0
    num_classes: int = 0  # 0 => unconditional (no label embedding)
    class_dropout_prob: float = 0.1

    @property
    def resolved_out_channels(self) -> int:
        return self.out_channels if self.out_channels is not None else self.in_channels

mlx_diffuser.models.UNet2D

Bases: ModelMixin[UNet2DConfig]

Source code in src/mlx_diffuser/models/unet2d.py 
 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
class UNet2D(ModelMixin[UNet2DConfig]):
    config_class = UNet2DConfig

    def __init__(self, config: UNet2DConfig):
        super().__init__()
        self.config = config
        boc = list(config.block_out_channels)
        n = len(boc)
        lpb = config.layers_per_block
        groups = config.norm_groups
        heads = config.num_heads
        ctx = config.cross_attention_dim
        attn_levels = (
            list(config.attention_levels)
            if config.attention_levels is not None
            else [i > 0 for i in range(n)]
        )

        self.time_embed = TimestepEmbedder(config.time_embed_dim)
        self.conv_in = nn.Conv2d(config.in_channels, boc[0], 3, padding=1)

        # --- down path ---
        self.down_resnets: list[ResnetBlock2D] = []
        self.down_attns: list[nn.Module] = []
        self.downsamplers: list[nn.Module] = []
        skip_channels = [boc[0]]
        prev = boc[0]
        for i, ch in enumerate(boc):
            for _ in range(lpb):
                self.down_resnets.append(ResnetBlock2D(prev, ch, config.time_embed_dim, groups))
                self.down_attns.append(
                    SpatialAttention(ch, heads, ctx, groups) if attn_levels[i] else Identity()
                )
                prev = ch
                skip_channels.append(ch)
            if i < n - 1:
                self.downsamplers.append(Downsample2D(ch))
                skip_channels.append(ch)
            else:
                self.downsamplers.append(Identity())

        # --- mid ---
        mid_ch = boc[-1]
        self.mid_resnet1 = ResnetBlock2D(mid_ch, mid_ch, config.time_embed_dim, groups)
        self.mid_attn = SpatialAttention(mid_ch, heads, ctx, groups)
        self.mid_resnet2 = ResnetBlock2D(mid_ch, mid_ch, config.time_embed_dim, groups)

        # --- up path ---
        self.up_resnets: list[ResnetBlock2D] = []
        self.up_attns: list[nn.Module] = []
        self.upsamplers: list[nn.Module] = []
        for i, ch in enumerate(reversed(boc)):
            level = n - 1 - i
            for _ in range(lpb + 1):
                skip = skip_channels.pop()
                self.up_resnets.append(
                    ResnetBlock2D(prev + skip, ch, config.time_embed_dim, groups)
                )
                self.up_attns.append(
                    SpatialAttention(ch, heads, ctx, groups) if attn_levels[level] else Identity()
                )
                prev = ch
            self.upsamplers.append(Upsample2D(ch) if i < n - 1 else Identity())

        self.norm_out = nn.GroupNorm(valid_groups(boc[0], groups), boc[0], pytorch_compatible=True)
        self.conv_out = nn.Conv2d(boc[0], config.out_channels, 3, padding=1)

        self._n_levels = n
        self._lpb = lpb

    def __call__(self, x: mx.array, t: mx.array, context: mx.array | None = None) -> mx.array:
        temb = self.time_embed(t)
        x = self.conv_in(x)

        skips = [x]
        ri = 0
        for i in range(self._n_levels):
            for _ in range(self._lpb):
                x = self.down_resnets[ri](x, temb)
                x = self.down_attns[ri](x, context)
                skips.append(x)
                ri += 1
            if i < self._n_levels - 1:
                x = self.downsamplers[i](x)
                skips.append(x)

        x = self.mid_resnet1(x, temb)
        x = self.mid_attn(x, context)
        x = self.mid_resnet2(x, temb)

        ri = 0
        for i in range(self._n_levels):
            for _ in range(self._lpb + 1):
                x = mx.concatenate([x, skips.pop()], axis=-1)
                x = self.up_resnets[ri](x, temb)
                x = self.up_attns[ri](x, context)
                ri += 1
            x = self.upsamplers[i](x)

        x = self.conv_out(nn.silu(self.norm_out(x)))
        return x

mlx_diffuser.models.UNet2DConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/unet2d.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
@dataclasses.dataclass
class UNet2DConfig(Config):
    in_channels: int = 4
    out_channels: int = 4
    block_out_channels: tuple[int, ...] = (224, 448, 672, 896)
    layers_per_block: int = 2
    attention_levels: tuple[bool, ...] | None = None  # per level; default: all but first
    num_heads: int = 8
    cross_attention_dim: int | None = None
    norm_groups: int = 32

    @property
    def time_embed_dim(self) -> int:
        return self.block_out_channels[0] * 4

mlx_diffuser.models.AutoencoderKL

Bases: ModelMixin[AutoencoderKLConfig]

Source code in src/mlx_diffuser/models/autoencoder_kl.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
class AutoencoderKL(ModelMixin[AutoencoderKLConfig]):
    config_class = AutoencoderKLConfig

    def __init__(self, config: AutoencoderKLConfig):
        super().__init__()
        self.config = config
        self.encoder = _Encoder(config)
        self.decoder = _Decoder(config)

    @property
    def scaling_factor(self) -> float:
        return self.config.scaling_factor

    def encode(self, x: mx.array) -> DiagonalGaussian:
        return DiagonalGaussian(self.encoder(x))

    def decode(self, z: mx.array) -> mx.array:
        return self.decoder(z)

    def __call__(
        self, x: mx.array, key: mx.array | None = None
    ) -> tuple[mx.array, DiagonalGaussian]:
        posterior = self.encode(x)
        return self.decode(posterior.sample(key)), posterior

mlx_diffuser.models.AutoencoderKLConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/autoencoder_kl.py 
44
45
46
47
48
49
50
51
@dataclasses.dataclass
class AutoencoderKLConfig(Config):
    in_channels: int = 3
    latent_channels: int = 4
    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
    layers_per_block: int = 2
    norm_groups: int = 32
    scaling_factor: float = 0.18215

mlx_diffuser.models.VideoDiT

Bases: ModelMixin[VideoDiTConfig]

Source code in src/mlx_diffuser/models/video_dit.py 
 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
class VideoDiT(ModelMixin[VideoDiTConfig]):
    config_class = VideoDiTConfig

    def __init__(self, config: VideoDiTConfig):
        super().__init__()
        self.config = config
        if config.hidden_size % config.num_heads != 0:
            raise ValueError("hidden_size must be divisible by num_heads.")
        if config.head_dim % 2 != 0:
            raise ValueError("head_dim (hidden_size / num_heads) must be even for RoPE.")

        self.patch_embed = PatchEmbed3D(config.in_channels, config.hidden_size, config.patch_size)
        self.t_embed = TimestepEmbedder(config.hidden_size)
        self.blocks = [
            VideoDiTBlock(
                config.hidden_size,
                config.num_heads,
                config.mlp_ratio,
                cross_attn_dim=config.cross_attn_dim,
                qk_norm=config.qk_norm,
            )
            for _ in range(config.depth)
        ]
        pt, ph, pw = config.patch_size
        patch_dim = pt * ph * pw * config.resolved_out_channels
        self.final = FinalLayer(config.hidden_size, patch_dim)

    def _unpatchify(self, x: mx.array, t: int, h: int, w: int) -> mx.array:
        pt, ph, pw = self.config.patch_size
        c = self.config.resolved_out_channels
        b = x.shape[0]
        x = x.reshape(b, t, h, w, pt, ph, pw, c)
        x = x.transpose(0, 1, 4, 2, 5, 3, 6, 7)  # (b, t, pt, h, ph, w, pw, c)
        return x.reshape(b, t * pt, h * ph, w * pw, c)

    def __call__(
        self,
        x: mx.array,
        t: mx.array,
        context: mx.array | None = None,
        *,
        training: bool = False,
    ) -> mx.array:
        """Predict the flow/noise target for video latents ``x`` at timestep ``t``.

        Args:
            x: video latents ``(B, T, H, W, C)``.
            t: per-sample timestep / sigma ``(B,)``.
            context: per-token text embeddings ``(B, L, cross_attn_dim)`` for
                cross-attention, or ``None`` for unconditional.
        """
        tokens, (gt, gh, gw) = self.patch_embed(x)
        cos, sin = rope_3d_freqs(self.config.head_dim, gt, gh, gw, self.config.rope_theta)
        rope = (cos[None, None].astype(tokens.dtype), sin[None, None].astype(tokens.dtype))

        c = self.t_embed(t)
        for block in self.blocks:
            tokens = block(tokens, c, rope, context=context)

        tokens = self.final(tokens, c)
        return self._unpatchify(tokens, gt, gh, gw)

__call__(x, t, context=None, *, training=False)

Predict the flow/noise target for video latents x at timestep t.

Parameters:

Name Type Description Default
x array

video latents (B, T, H, W, C).

 required
t array

per-sample timestep / sigma (B,).

 required
context array | None

per-token text embeddings (B, L, cross_attn_dim) for cross-attention, or None for unconditional.

 None
Source code in src/mlx_diffuser/models/video_dit.py 
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
def __call__(
    self,
    x: mx.array,
    t: mx.array,
    context: mx.array | None = None,
    *,
    training: bool = False,
) -> mx.array:
    """Predict the flow/noise target for video latents ``x`` at timestep ``t``.

    Args:
        x: video latents ``(B, T, H, W, C)``.
        t: per-sample timestep / sigma ``(B,)``.
        context: per-token text embeddings ``(B, L, cross_attn_dim)`` for
            cross-attention, or ``None`` for unconditional.
    """
    tokens, (gt, gh, gw) = self.patch_embed(x)
    cos, sin = rope_3d_freqs(self.config.head_dim, gt, gh, gw, self.config.rope_theta)
    rope = (cos[None, None].astype(tokens.dtype), sin[None, None].astype(tokens.dtype))

    c = self.t_embed(t)
    for block in self.blocks:
        tokens = block(tokens, c, rope, context=context)

    tokens = self.final(tokens, c)
    return self._unpatchify(tokens, gt, gh, gw)

mlx_diffuser.models.VideoDiTConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/video_dit.py 
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
@dataclasses.dataclass
class VideoDiTConfig(Config):
    in_channels: int = 16  # latent channels of the video VAE
    out_channels: int | None = None
    patch_size: tuple[int, int, int] = (1, 2, 2)  # (time, height, width)
    hidden_size: int = 1024
    depth: int = 24
    num_heads: int = 16
    mlp_ratio: float = 4.0
    cross_attn_dim: int | None = 4096  # text-embedding dim (T5); None => unconditional
    rope_theta: float = 10000.0
    qk_norm: bool = True

    @property
    def resolved_out_channels(self) -> int:
        return self.out_channels if self.out_channels is not None else self.in_channels

    @property
    def head_dim(self) -> int:
        return self.hidden_size // self.num_heads

    # --- presets (architecture shapes of well-known video models) ---------
    @classmethod
    def wan_t2v_1_3b(cls) -> VideoDiTConfig:
        """WAN 2.1 text-to-video 1.3B shape (16-channel latents, umT5 conditioning)."""
        return cls(
            in_channels=16,
            patch_size=(1, 2, 2),
            hidden_size=1536,
            depth=30,
            num_heads=12,
            mlp_ratio=8960 / 1536,
            cross_attn_dim=4096,
        )

    @classmethod
    def wan_t2v_14b(cls) -> VideoDiTConfig:
        """WAN 2.1 text-to-video 14B shape."""
        return cls(
            in_channels=16,
            patch_size=(1, 2, 2),
            hidden_size=5120,
            depth=40,
            num_heads=40,
            mlp_ratio=13824 / 5120,
            cross_attn_dim=4096,
        )

    @classmethod
    def ltx_video(cls) -> VideoDiTConfig:
        """LTX-Video shape (128-channel highly-compressed latents, T5-XXL conditioning)."""
        return cls(
            in_channels=128,
            patch_size=(1, 1, 1),
            hidden_size=2048,
            depth=28,
            num_heads=16,
            mlp_ratio=4.0,
            cross_attn_dim=4096,
        )

ltx_video() classmethod

LTX-Video shape (128-channel highly-compressed latents, T5-XXL conditioning).

Source code in src/mlx_diffuser/models/video_dit.py 
79
80
81
82
83
84
85
86
87
88
89
90
@classmethod
def ltx_video(cls) -> VideoDiTConfig:
    """LTX-Video shape (128-channel highly-compressed latents, T5-XXL conditioning)."""
    return cls(
        in_channels=128,
        patch_size=(1, 1, 1),
        hidden_size=2048,
        depth=28,
        num_heads=16,
        mlp_ratio=4.0,
        cross_attn_dim=4096,
    )

wan_t2v_14b() classmethod

WAN 2.1 text-to-video 14B shape.

Source code in src/mlx_diffuser/models/video_dit.py 
66
67
68
69
70
71
72
73
74
75
76
77
@classmethod
def wan_t2v_14b(cls) -> VideoDiTConfig:
    """WAN 2.1 text-to-video 14B shape."""
    return cls(
        in_channels=16,
        patch_size=(1, 2, 2),
        hidden_size=5120,
        depth=40,
        num_heads=40,
        mlp_ratio=13824 / 5120,
        cross_attn_dim=4096,
    )

wan_t2v_1_3b() classmethod

WAN 2.1 text-to-video 1.3B shape (16-channel latents, umT5 conditioning).

Source code in src/mlx_diffuser/models/video_dit.py 
53
54
55
56
57
58
59
60
61
62
63
64
@classmethod
def wan_t2v_1_3b(cls) -> VideoDiTConfig:
    """WAN 2.1 text-to-video 1.3B shape (16-channel latents, umT5 conditioning)."""
    return cls(
        in_channels=16,
        patch_size=(1, 2, 2),
        hidden_size=1536,
        depth=30,
        num_heads=12,
        mlp_ratio=8960 / 1536,
        cross_attn_dim=4096,
    )

mlx_diffuser.models.AutoencoderKLVideo

Bases: ModelMixin[AutoencoderKLVideoConfig]

Source code in src/mlx_diffuser/models/autoencoder_kl_video.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
class AutoencoderKLVideo(ModelMixin[AutoencoderKLVideoConfig]):
    config_class = AutoencoderKLVideoConfig

    def __init__(self, config: AutoencoderKLVideoConfig):
        super().__init__()
        self.config = config
        self.encoder = _Encoder(config)
        self.decoder = _Decoder(config)

    @property
    def scaling_factor(self) -> float:
        return self.config.scaling_factor

    def encode(self, x: mx.array) -> DiagonalGaussian:
        return DiagonalGaussian(self.encoder(x))

    def decode(self, z: mx.array) -> mx.array:
        return self.decoder(z)

    def __call__(
        self, x: mx.array, key: mx.array | None = None
    ) -> tuple[mx.array, DiagonalGaussian]:
        posterior = self.encode(x)
        return self.decode(posterior.sample(key)), posterior

mlx_diffuser.models.AutoencoderKLVideoConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/autoencoder_kl_video.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@dataclasses.dataclass
class AutoencoderKLVideoConfig(Config):
    in_channels: int = 3
    latent_channels: int = 16
    block_out_channels: tuple[int, ...] = (64, 128, 128)
    layers_per_block: int = 1
    temporal_compression: int = 4  # power of two, <= number of spatial downsamples
    norm_groups: int = 32
    scaling_factor: float = 1.0

    @property
    def num_downsamples(self) -> int:
        return len(self.block_out_channels) - 1

    @property
    def spatial_compression(self) -> int:
        return 2**self.num_downsamples

    @property
    def temporal_levels(self) -> int:
        return int(round(math.log2(self.temporal_compression)))

WAN 2.1 (weight-compatible ports)

mlx_diffuser.models.WanTransformer3DModel

Bases: ModelMixin[WanTransformerConfig]

WAN 2.1 text-to-video diffusion transformer. Channels-last (B, T, H, W, C).

Source code in src/mlx_diffuser/models/wan_transformer.py 
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
class WanTransformer3DModel(ModelMixin[WanTransformerConfig]):
    """WAN 2.1 text-to-video diffusion transformer. Channels-last ``(B, T, H, W, C)``."""

    config_class = WanTransformerConfig

    def __init__(self, config: WanTransformerConfig):
        super().__init__()
        self.config = config
        d = config.inner_dim
        self.rope = WanRotaryPosEmbed(config.attention_head_dim, config.theta)
        self.patch_embedding = nn.Conv3d(
            config.in_channels, d, kernel_size=config.patch_size, stride=config.patch_size
        )
        self.condition_embedder = WanConditionEmbedder(config)
        self.blocks = [WanTransformerBlock(config) for _ in range(config.num_layers)]
        self.norm_out = nn.LayerNorm(d, eps=config.eps, affine=False)
        self.proj_out = nn.Linear(d, config.out_channels * math.prod(config.patch_size))
        self.scale_shift_table = mx.zeros((1, 2, d))

    def __call__(
        self,
        x: mx.array,
        timestep: mx.array,
        context: mx.array,
        cache: FirstBlockCache | None = None,
    ) -> mx.array:
        """Predict the flow target for latents ``x`` ``(B, T, H, W, C)`` at ``timestep``.

        ``context`` is per-token text embeddings ``(B, L, text_dim)`` (umT5). When a
        :class:`~mlx_diffuser.caching.FirstBlockCache` is supplied, the bulk of the
        transformer blocks may be skipped on steps where the output barely changes.
        """
        b, t, h, w, _ = x.shape
        pt, ph, pw = self.config.patch_size
        gt, gh, gw = t // pt, h // ph, w // pw

        tokens = self.patch_embedding(x)  # (B, gt, gh, gw, d)
        tokens = tokens.reshape(b, gt * gh * gw, -1)
        rope = self.rope(gt, gh, gw)

        temb, timestep_proj, context = self.condition_embedder(timestep, context)
        timestep_proj = timestep_proj.reshape(b, 6, -1)

        tokens = self._run_blocks(tokens, context, timestep_proj, rope, cache)

        mod = self.scale_shift_table + temb.astype(mx.float32)[:, None]  # (B, 2, d)
        shift, scale = mod[:, 0:1], mod[:, 1:2]
        tokens = _layernorm_f32(tokens, self.config.eps) * (1 + scale) + shift
        tokens = self.proj_out(tokens.astype(x.dtype))

        # unpatchify -> (B, T, H, W, C_out)
        c = self.config.out_channels
        tokens = tokens.reshape(b, gt, gh, gw, pt, ph, pw, c)
        tokens = tokens.transpose(0, 1, 4, 2, 5, 3, 6, 7)
        return tokens.reshape(b, gt * pt, gh * ph, gw * pw, c)

    def _run_blocks(self, tokens, context, timestep_proj, rope, cache):
        """Run the transformer blocks, optionally reusing the First-Block Cache.

        On a cache hit we run only block 0 and add the residual cached from the last
        full step; otherwise we run every block and refresh the cached residual.
        """
        if cache is None:
            for block in self.blocks:
                tokens = block(tokens, context, timestep_proj, rope)
            return tokens

        first = self.blocks[0](tokens, context, timestep_proj, rope)
        # FBCache signal is the *first block's contribution* (first - input), not its
        # output: the output is dominated by the slowly-drifting input, which would
        # mask real change and skip the wrong steps.
        if cache.should_reuse(first - tokens):
            return first + cache.residual
        hidden = first
        for block in self.blocks[1:]:
            hidden = block(hidden, context, timestep_proj, rope)
        cache.residual = hidden - first
        return hidden

__call__(x, timestep, context, cache=None)

Predict the flow target for latents x (B, T, H, W, C) at timestep.

context is per-token text embeddings (B, L, text_dim) (umT5). When a :class:~mlx_diffuser.caching.FirstBlockCache is supplied, the bulk of the transformer blocks may be skipped on steps where the output barely changes.

Source code in src/mlx_diffuser/models/wan_transformer.py 
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def __call__(
    self,
    x: mx.array,
    timestep: mx.array,
    context: mx.array,
    cache: FirstBlockCache | None = None,
) -> mx.array:
    """Predict the flow target for latents ``x`` ``(B, T, H, W, C)`` at ``timestep``.

    ``context`` is per-token text embeddings ``(B, L, text_dim)`` (umT5). When a
    :class:`~mlx_diffuser.caching.FirstBlockCache` is supplied, the bulk of the
    transformer blocks may be skipped on steps where the output barely changes.
    """
    b, t, h, w, _ = x.shape
    pt, ph, pw = self.config.patch_size
    gt, gh, gw = t // pt, h // ph, w // pw

    tokens = self.patch_embedding(x)  # (B, gt, gh, gw, d)
    tokens = tokens.reshape(b, gt * gh * gw, -1)
    rope = self.rope(gt, gh, gw)

    temb, timestep_proj, context = self.condition_embedder(timestep, context)
    timestep_proj = timestep_proj.reshape(b, 6, -1)

    tokens = self._run_blocks(tokens, context, timestep_proj, rope, cache)

    mod = self.scale_shift_table + temb.astype(mx.float32)[:, None]  # (B, 2, d)
    shift, scale = mod[:, 0:1], mod[:, 1:2]
    tokens = _layernorm_f32(tokens, self.config.eps) * (1 + scale) + shift
    tokens = self.proj_out(tokens.astype(x.dtype))

    # unpatchify -> (B, T, H, W, C_out)
    c = self.config.out_channels
    tokens = tokens.reshape(b, gt, gh, gw, pt, ph, pw, c)
    tokens = tokens.transpose(0, 1, 4, 2, 5, 3, 6, 7)
    return tokens.reshape(b, gt * pt, gh * ph, gw * pw, c)

mlx_diffuser.models.WanTransformerConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/wan_transformer.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@dataclasses.dataclass
class WanTransformerConfig(Config):
    patch_size: tuple[int, int, int] = (1, 2, 2)
    num_attention_heads: int = 12
    attention_head_dim: int = 128
    in_channels: int = 16
    out_channels: int = 16
    text_dim: int = 4096
    freq_dim: int = 256
    ffn_dim: int = 8960
    num_layers: int = 30
    cross_attn_norm: bool = True
    eps: float = 1e-6
    rope_max_seq_len: int = 1024
    theta: float = 10000.0

    @property
    def inner_dim(self) -> int:
        return self.num_attention_heads * self.attention_head_dim

    @classmethod
    def wan_t2v_1_3b(cls) -> WanTransformerConfig:
        return cls()  # defaults match the 1.3B model

mlx_diffuser.models.AutoencoderKLWan

Bases: ModelMixin[AutoencoderKLWanConfig]

WAN 2.1 causal 3D VAE. Channels-last (B, T, H, W, C).

Source code in src/mlx_diffuser/models/autoencoder_kl_wan.py 
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
class AutoencoderKLWan(ModelMixin[AutoencoderKLWanConfig]):
    """WAN 2.1 causal 3D VAE. Channels-last ``(B, T, H, W, C)``."""

    config_class = AutoencoderKLWanConfig

    def __init__(self, config: AutoencoderKLWanConfig):
        super().__init__()
        self.config = config
        self.encoder = WanEncoder3d(config)
        self.quant_conv = WanCausalConv3d(config.z_dim * 2, config.z_dim * 2, 1)
        self.post_quant_conv = WanCausalConv3d(config.z_dim, config.z_dim, 1)
        self.decoder = WanDecoder3d(config)
        self._enc_conv_num = _count_causal_convs(self.encoder)
        self._dec_conv_num = _count_causal_convs(self.decoder)

    # --- latent normalization helpers (used by the diffusion pipeline) -------
    @property
    def latents_mean(self) -> mx.array:
        return mx.array(self.config.latents_mean).reshape(1, 1, 1, 1, -1)

    @property
    def latents_std(self) -> mx.array:
        return mx.array(self.config.latents_std).reshape(1, 1, 1, 1, -1)

    def normalize_latents(self, z: mx.array) -> mx.array:
        """Pixel-latent -> normalized latent space the transformer operates in."""
        return (z - self.latents_mean) / self.latents_std

    def denormalize_latents(self, z: mx.array) -> mx.array:
        return z * self.latents_std + self.latents_mean

    # --- encode / decode -----------------------------------------------------
    def encode(self, x: mx.array) -> DiagonalGaussian:
        """Encode ``(B, T, H, W, C)`` video into a latent ``DiagonalGaussian``.

        The clip is processed causally in temporal chunks (frame 0 alone, then
        groups of 4), accumulating ``1 + (T-1)//4`` latent frames.
        """
        num_frames = x.shape[1]
        feat_cache = [None] * self._enc_conv_num
        feat_idx = [0]
        iters = 1 + (num_frames - 1) // 4
        outs = []
        for i in range(iters):
            feat_idx[0] = 0
            chunk = x[:, :1] if i == 0 else x[:, 1 + 4 * (i - 1) : 1 + 4 * i]
            outs.append(self.encoder(chunk, feat_cache, feat_idx))
        out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=1)
        return DiagonalGaussian(self.quant_conv(out))

    def decode(self, z: mx.array) -> mx.array:
        """Decode a latent ``(B, T, H, W, C)`` back to pixels in ``[-1, 1]``."""
        num_frames = z.shape[1]
        x = self.post_quant_conv(z)
        feat_cache = [None] * self._dec_conv_num
        feat_idx = [0]
        outs = []
        for i in range(num_frames):
            feat_idx[0] = 0
            outs.append(self.decoder(x[:, i : i + 1], feat_cache, feat_idx))
        out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=1)
        return mx.clip(out, -1.0, 1.0)

    def __call__(
        self, x: mx.array, key: mx.array | None = None
    ) -> tuple[mx.array, DiagonalGaussian]:
        posterior = self.encode(x)
        return self.decode(posterior.mode()), posterior

decode(z)

Decode a latent (B, T, H, W, C) back to pixels in [-1, 1].

Source code in src/mlx_diffuser/models/autoencoder_kl_wan.py 
416
417
418
419
420
421
422
423
424
425
426
427
def decode(self, z: mx.array) -> mx.array:
    """Decode a latent ``(B, T, H, W, C)`` back to pixels in ``[-1, 1]``."""
    num_frames = z.shape[1]
    x = self.post_quant_conv(z)
    feat_cache = [None] * self._dec_conv_num
    feat_idx = [0]
    outs = []
    for i in range(num_frames):
        feat_idx[0] = 0
        outs.append(self.decoder(x[:, i : i + 1], feat_cache, feat_idx))
    out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=1)
    return mx.clip(out, -1.0, 1.0)

encode(x)

Encode (B, T, H, W, C) video into a latent DiagonalGaussian.

The clip is processed causally in temporal chunks (frame 0 alone, then groups of 4), accumulating 1 + (T-1)//4 latent frames.

Source code in src/mlx_diffuser/models/autoencoder_kl_wan.py 
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def encode(self, x: mx.array) -> DiagonalGaussian:
    """Encode ``(B, T, H, W, C)`` video into a latent ``DiagonalGaussian``.

    The clip is processed causally in temporal chunks (frame 0 alone, then
    groups of 4), accumulating ``1 + (T-1)//4`` latent frames.
    """
    num_frames = x.shape[1]
    feat_cache = [None] * self._enc_conv_num
    feat_idx = [0]
    iters = 1 + (num_frames - 1) // 4
    outs = []
    for i in range(iters):
        feat_idx[0] = 0
        chunk = x[:, :1] if i == 0 else x[:, 1 + 4 * (i - 1) : 1 + 4 * i]
        outs.append(self.encoder(chunk, feat_cache, feat_idx))
    out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=1)
    return DiagonalGaussian(self.quant_conv(out))

normalize_latents(z)

Pixel-latent -> normalized latent space the transformer operates in.

Source code in src/mlx_diffuser/models/autoencoder_kl_wan.py 
390
391
392
def normalize_latents(self, z: mx.array) -> mx.array:
    """Pixel-latent -> normalized latent space the transformer operates in."""
    return (z - self.latents_mean) / self.latents_std

mlx_diffuser.models.AutoencoderKLWanConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/autoencoder_kl_wan.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
@dataclasses.dataclass
class AutoencoderKLWanConfig(Config):
    base_dim: int = 96
    z_dim: int = 16
    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
    num_res_blocks: int = 2
    attn_scales: tuple[float, ...] = ()
    temperal_downsample: tuple[bool, ...] = (False, True, True)
    in_channels: int = 3
    out_channels: int = 3
    latents_mean: tuple[float, ...] = tuple(_LATENTS_MEAN)
    latents_std: tuple[float, ...] = tuple(_LATENTS_STD)

    @property
    def temperal_upsample(self) -> tuple[bool, ...]:
        return tuple(reversed(self.temperal_downsample))

mlx_diffuser.models.umt5.UMT5EncoderModel

Bases: ModelMixin[UMT5Config]

umT5 text encoder. Produces per-token embeddings for WAN cross-attention.

Source code in src/mlx_diffuser/models/umt5.py 
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
class UMT5EncoderModel(ModelMixin[UMT5Config]):
    """umT5 text encoder. Produces per-token embeddings for WAN cross-attention."""

    config_class = UMT5Config

    def __init__(self, config: UMT5Config):
        super().__init__()
        self.config = config
        self.shared = nn.Embedding(config.vocab_size, config.d_model)
        self.encoder = _Stack(config)

    def __call__(self, input_ids: mx.array, attention_mask: mx.array | None = None) -> mx.array:
        """``input_ids`` ``(B, L)`` -> embeddings ``(B, L, d_model)``.

        ``attention_mask`` ``(B, L)`` with 1 for real tokens, 0 for padding.
        """
        # Run activations in fp32: T5/umT5 activations routinely exceed the fp16
        # range, so a low-precision compute path overflows to inf/NaN on some
        # prompts. The (quantized) weights stay small; only the activations are fp32.
        x = self.shared(input_ids).astype(mx.float32)
        mask = None
        if attention_mask is not None:
            neg = (1.0 - attention_mask.astype(mx.float32)) * -1e9
            mask = neg[:, None, None, :]  # (B, 1, 1, L) additive
        return self.encoder(x, mask)

__call__(input_ids, attention_mask=None)

input_ids (B, L) -> embeddings (B, L, d_model).

attention_mask (B, L) with 1 for real tokens, 0 for padding.

Source code in src/mlx_diffuser/models/umt5.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def __call__(self, input_ids: mx.array, attention_mask: mx.array | None = None) -> mx.array:
    """``input_ids`` ``(B, L)`` -> embeddings ``(B, L, d_model)``.

    ``attention_mask`` ``(B, L)`` with 1 for real tokens, 0 for padding.
    """
    # Run activations in fp32: T5/umT5 activations routinely exceed the fp16
    # range, so a low-precision compute path overflows to inf/NaN on some
    # prompts. The (quantized) weights stay small; only the activations are fp32.
    x = self.shared(input_ids).astype(mx.float32)
    mask = None
    if attention_mask is not None:
        neg = (1.0 - attention_mask.astype(mx.float32)) * -1e9
        mask = neg[:, None, None, :]  # (B, 1, 1, L) additive
    return self.encoder(x, mask)

Stable Diffusion XL (weight-compatible ports)

mlx_diffuser.models.SDXLUNet

Bases: ModelMixin[SDXLUNetConfig]

SDXL base UNet. Channels-last (B, H, W, C).

Source code in src/mlx_diffuser/models/unet_sdxl.py 
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
class SDXLUNet(ModelMixin[SDXLUNetConfig]):
    """SDXL base UNet. Channels-last ``(B, H, W, C)``."""

    config_class = SDXLUNetConfig

    def __init__(self, config: SDXLUNetConfig):
        super().__init__()
        self.config = config
        c = config
        boc = list(c.block_out_channels)
        n = len(boc)
        temb = c.time_embed_dim
        g = c.norm_groups

        self.conv_in = nn.Conv2d(c.in_channels, boc[0], 3, padding=1)
        self.time_embedding = _TimeMLP(boc[0], temb)
        self.add_embedding = _TimeMLP(c.projection_class_embeddings_input_dim, temb)

        self.down_blocks: list = []
        out = boc[0]
        for i in range(n):
            in_ch, out = out, boc[i]
            cross = c.down_block_types[i].startswith("CrossAttn")
            down_cls = _CrossAttnDownBlock if cross else _DownBlock
            self.down_blocks.append(
                down_cls(
                    in_ch,
                    out,
                    c.layers_per_block,
                    temb,
                    g,
                    add_down=i < n - 1,
                    heads=c.num_attention_heads[i],
                    depth=c.transformer_layers_per_block[i],
                    cross_dim=c.cross_attention_dim,
                )
            )

        self.mid_block = _MidBlock(
            boc[-1],
            temb,
            g,
            c.num_attention_heads[-1],
            c.transformer_layers_per_block[-1],
            c.cross_attention_dim,
        )

        rev_boc = list(reversed(boc))
        rev_heads = list(reversed(c.num_attention_heads))
        rev_depth = list(reversed(c.transformer_layers_per_block))
        self.up_blocks: list = []
        prev_out = boc[-1]
        for i in range(n):
            out = rev_boc[i]
            in_ch = rev_boc[min(i + 1, n - 1)]  # skip channel = symmetric down level input
            cross = c.up_block_types[i].startswith("CrossAttn")
            up_cls = _CrossAttnUpBlock if cross else _UpBlock
            self.up_blocks.append(
                up_cls(
                    in_ch,
                    out,
                    prev_out,
                    c.layers_per_block + 1,
                    temb,
                    g,
                    add_up=i < n - 1,
                    heads=rev_heads[i],
                    depth=rev_depth[i],
                    cross_dim=c.cross_attention_dim,
                )
            )
            prev_out = out

        self.conv_norm_out = nn.GroupNorm(g, boc[0], pytorch_compatible=True)
        self.conv_out = nn.Conv2d(boc[0], c.out_channels, 3, padding=1)

    def __call__(
        self,
        sample: mx.array,
        timestep: mx.array,
        encoder_hidden_states: mx.array,
        text_embeds: mx.array,
        time_ids: mx.array,
        cache: DeepCache | None = None,
    ) -> mx.array:
        """Predict noise for latents ``sample`` ``(B, H, W, C)``.

        ``encoder_hidden_states`` is the ``(B, L, 2048)`` concatenated CLIP context;
        ``text_embeds`` the ``(B, 1280)`` pooled bigG embedding; ``time_ids`` the
        ``(B, 6)`` SDXL micro-conditioning (sizes / crop). With a ``DeepCache``, the
        deep blocks may be skipped on most steps.
        """
        b = sample.shape[0]
        emb = self.time_embedding(
            _timesteps(mx.broadcast_to(timestep, (b,)), self.config.block_out_channels[0])
        )
        tid = _timesteps(time_ids.reshape(-1), self.config.addition_time_embed_dim).reshape(b, -1)
        emb = emb + self.add_embedding(mx.concatenate([text_embeds, tid], axis=-1))

        ehs = encoder_hidden_states
        if cache is not None and cache.should_reuse():
            return self._cached_forward(sample, emb, ehs, cache)

        x = self.conv_in(sample)
        res_samples = [x]
        for block in self.down_blocks:
            x, res = block(x, emb, ehs)
            res_samples += res

        x = self.mid_block(x, emb, ehs)

        for block in self.up_blocks[:-1]:
            take = len(block.resnets)
            x = block(x, res_samples[-take:], emb, ehs)
            res_samples = res_samples[:-take]

        if cache is not None:
            cache.deep = x  # the input to the shallow up block — reused on cached steps

        last = self.up_blocks[-1]
        x = last(x, res_samples[-len(last.resnets) :], emb, ehs)
        return self.conv_out(nn.silu(self.conv_norm_out(x)))

    def _cached_forward(self, sample, emb, ehs, cache: DeepCache) -> mx.array:
        """DeepCache step: run only conv_in + first down block + last up block."""
        x = self.conv_in(sample)
        shallow = [x]
        x, res0 = self.down_blocks[0](x, emb, ehs)
        shallow += res0
        last = self.up_blocks[-1]
        skips = shallow[: len(last.resnets)]  # conv_in + first down block's resnets
        x = last(cache.deep, skips, emb, ehs)
        return self.conv_out(nn.silu(self.conv_norm_out(x)))

__call__(sample, timestep, encoder_hidden_states, text_embeds, time_ids, cache=None)

Predict noise for latents sample (B, H, W, C).

encoder_hidden_states is the (B, L, 2048) concatenated CLIP context; text_embeds the (B, 1280) pooled bigG embedding; time_ids the (B, 6) SDXL micro-conditioning (sizes / crop). With a DeepCache, the deep blocks may be skipped on most steps.

Source code in src/mlx_diffuser/models/unet_sdxl.py 
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
def __call__(
    self,
    sample: mx.array,
    timestep: mx.array,
    encoder_hidden_states: mx.array,
    text_embeds: mx.array,
    time_ids: mx.array,
    cache: DeepCache | None = None,
) -> mx.array:
    """Predict noise for latents ``sample`` ``(B, H, W, C)``.

    ``encoder_hidden_states`` is the ``(B, L, 2048)`` concatenated CLIP context;
    ``text_embeds`` the ``(B, 1280)`` pooled bigG embedding; ``time_ids`` the
    ``(B, 6)`` SDXL micro-conditioning (sizes / crop). With a ``DeepCache``, the
    deep blocks may be skipped on most steps.
    """
    b = sample.shape[0]
    emb = self.time_embedding(
        _timesteps(mx.broadcast_to(timestep, (b,)), self.config.block_out_channels[0])
    )
    tid = _timesteps(time_ids.reshape(-1), self.config.addition_time_embed_dim).reshape(b, -1)
    emb = emb + self.add_embedding(mx.concatenate([text_embeds, tid], axis=-1))

    ehs = encoder_hidden_states
    if cache is not None and cache.should_reuse():
        return self._cached_forward(sample, emb, ehs, cache)

    x = self.conv_in(sample)
    res_samples = [x]
    for block in self.down_blocks:
        x, res = block(x, emb, ehs)
        res_samples += res

    x = self.mid_block(x, emb, ehs)

    for block in self.up_blocks[:-1]:
        take = len(block.resnets)
        x = block(x, res_samples[-take:], emb, ehs)
        res_samples = res_samples[:-take]

    if cache is not None:
        cache.deep = x  # the input to the shallow up block — reused on cached steps

    last = self.up_blocks[-1]
    x = last(x, res_samples[-len(last.resnets) :], emb, ehs)
    return self.conv_out(nn.silu(self.conv_norm_out(x)))

mlx_diffuser.models.SDXLUNetConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/unet_sdxl.py 
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
@dataclasses.dataclass
class SDXLUNetConfig(Config):
    in_channels: int = 4
    out_channels: int = 4
    block_out_channels: tuple[int, ...] = (320, 640, 1280)
    down_block_types: tuple[str, ...] = (
        "DownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
    )
    up_block_types: tuple[str, ...] = (
        "CrossAttnUpBlock2D",
        "CrossAttnUpBlock2D",
        "UpBlock2D",
    )
    layers_per_block: int = 2
    transformer_layers_per_block: tuple[int, ...] = (1, 2, 10)
    num_attention_heads: tuple[int, ...] = (5, 10, 20)
    cross_attention_dim: int = 2048
    addition_time_embed_dim: int = 256
    projection_class_embeddings_input_dim: int = 2816
    norm_groups: int = 32

    @property
    def time_embed_dim(self) -> int:
        return self.block_out_channels[0] * 4

mlx_diffuser.models.AutoencoderKLSD

Bases: ModelMixin[AutoencoderKLSDConfig]

SD / SDXL VAE. Channels-last (B, H, W, C).

Source code in src/mlx_diffuser/models/autoencoder_kl_sd.py 
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
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
class AutoencoderKLSD(ModelMixin[AutoencoderKLSDConfig]):
    """SD / SDXL VAE. Channels-last ``(B, H, W, C)``."""

    config_class = AutoencoderKLSDConfig

    def __init__(self, config: AutoencoderKLSDConfig):
        super().__init__()
        self.config = config
        self.encoder = _Encoder(config)
        self.decoder = _Decoder(config)
        self.quant_conv: nn.Conv2d | None = None
        self.post_quant_conv: nn.Conv2d | None = None
        if config.use_quant_conv:  # FLUX / SD3 VAEs skip the quant convs
            self.quant_conv = nn.Conv2d(2 * config.latent_channels, 2 * config.latent_channels, 1)
            self.post_quant_conv = nn.Conv2d(config.latent_channels, config.latent_channels, 1)

    @property
    def scaling_factor(self) -> float:
        return self.config.scaling_factor

    @property
    def shift_factor(self) -> float:
        return self.config.shift_factor

    def encode(self, x: mx.array) -> DiagonalGaussian:
        moments = self.encoder(x)
        if self.quant_conv is not None:
            moments = self.quant_conv(moments)
        return DiagonalGaussian(moments)

    def decode(
        self, z: mx.array, *, tile: bool = False, tile_latent: int = 64, overlap_latent: int = 16
    ) -> mx.array:
        """Decode latents to pixels. With ``tile=True``, decode in overlapping tiles.

        Tiling bounds peak memory for large latents: the latent is split into
        ``tile_latent``-sized spatial tiles (stride ``tile_latent - overlap_latent``),
        each decoded independently, then feather-blended back together.
        """
        if self.post_quant_conv is not None:
            z = self.post_quant_conv(z)
        h, w = z.shape[1], z.shape[2]
        if not tile or (h <= tile_latent and w <= tile_latent):
            return self.decoder(z)
        return self._tiled(z, tile_latent, overlap_latent)

    def _tiled(self, z: mx.array, tile: int, overlap: int) -> mx.array:
        h, w = z.shape[1], z.shape[2]
        stride = tile - overlap
        ys = sorted({*range(0, max(h - tile, 0) + 1, stride), max(h - tile, 0)})
        xs = sorted({*range(0, max(w - tile, 0) + 1, stride), max(w - tile, 0)})
        scale = 2 ** (len(self.config.block_out_channels) - 1)  # VAE spatial upsampling factor
        out = mx.zeros((z.shape[0], h * scale, w * scale, self.config.out_channels))
        wsum = mx.zeros((1, h * scale, w * scale, 1))
        for y in ys:
            for x in xs:
                dec = self.decoder(z[:, y : y + tile, x : x + tile])  # (B, tile*8, tile*8, C)
                th, tw = dec.shape[1], dec.shape[2]
                win = (
                    _feather(th, overlap * scale)[:, None] * _feather(tw, overlap * scale)[None]
                )[None, :, :, None]
                py, px = y * scale, x * scale
                out[:, py : py + th, px : px + tw] += dec * win
                wsum[:, py : py + th, px : px + tw] += win
        return out / wsum

    def __call__(
        self, x: mx.array, key: mx.array | None = None
    ) -> tuple[mx.array, DiagonalGaussian]:
        posterior = self.encode(x)
        return self.decode(posterior.sample(key)), posterior

decode(z, *, tile=False, tile_latent=64, overlap_latent=16)

Decode latents to pixels. With tile=True, decode in overlapping tiles.

Tiling bounds peak memory for large latents: the latent is split into tile_latent-sized spatial tiles (stride tile_latent - overlap_latent), each decoded independently, then feather-blended back together.

Source code in src/mlx_diffuser/models/autoencoder_kl_sd.py 
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
def decode(
    self, z: mx.array, *, tile: bool = False, tile_latent: int = 64, overlap_latent: int = 16
) -> mx.array:
    """Decode latents to pixels. With ``tile=True``, decode in overlapping tiles.

    Tiling bounds peak memory for large latents: the latent is split into
    ``tile_latent``-sized spatial tiles (stride ``tile_latent - overlap_latent``),
    each decoded independently, then feather-blended back together.
    """
    if self.post_quant_conv is not None:
        z = self.post_quant_conv(z)
    h, w = z.shape[1], z.shape[2]
    if not tile or (h <= tile_latent and w <= tile_latent):
        return self.decoder(z)
    return self._tiled(z, tile_latent, overlap_latent)

mlx_diffuser.models.AutoencoderKLSDConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/autoencoder_kl_sd.py 
23
24
25
26
27
28
29
30
31
32
33
@dataclasses.dataclass
class AutoencoderKLSDConfig(Config):
    in_channels: int = 3
    out_channels: int = 3
    latent_channels: int = 4
    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
    layers_per_block: int = 2
    norm_groups: int = 32
    scaling_factor: float = 0.13025  # SDXL
    shift_factor: float = 0.0  # FLUX/SD3 shift latents before scaling: (z - shift) * scale
    use_quant_conv: bool = True  # SD/SDXL have quant convs; FLUX/SD3 VAEs do not

mlx_diffuser.models.CLIPTextModel

Bases: ModelMixin[CLIPTextConfig]

CLIP text encoder. Returns all hidden states + the projected pooled embedding.

Source code in src/mlx_diffuser/models/clip_text.py 
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
class CLIPTextModel(ModelMixin[CLIPTextConfig]):
    """CLIP text encoder. Returns all hidden states + the projected pooled embedding."""

    config_class = CLIPTextConfig

    def __init__(self, config: CLIPTextConfig):
        super().__init__()
        self.config = config
        self.text_model = _CLIPTextTransformer(config)
        self.text_projection = (
            nn.Linear(config.hidden_size, config.projection_dim, bias=False)
            if config.with_projection
            else None
        )

    def __call__(self, input_ids: mx.array) -> tuple[list[mx.array], mx.array]:
        """``(B, L) int -> (hidden_states, pooled)``.

        ``hidden_states`` has length ``num_layers + 1`` (embeddings then each layer
        output); SDXL reads ``[-2]``. ``pooled`` is the final-norm hidden state at the
        EOS token (the highest token id), projected by ``text_projection`` when present
        (the bigG encoder). SDXL uses the pooled output only from that encoder.
        """
        x = self.text_model.embeddings(input_ids)
        seq = input_ids.shape[-1]
        mask = mx.triu(mx.full((seq, seq), -mx.inf, dtype=x.dtype), k=1)

        hidden_states = [x]
        for layer in self.text_model.encoder.layers:
            x = layer(x, mask)
            hidden_states.append(x)

        last = self.text_model.final_layer_norm(x)
        eos = mx.argmax(input_ids, axis=-1)
        pooled = last[mx.arange(input_ids.shape[0]), eos]
        if self.text_projection is not None:
            pooled = self.text_projection(pooled)
        return hidden_states, pooled

__call__(input_ids)

(B, L) int -> (hidden_states, pooled).

hidden_states has length num_layers + 1 (embeddings then each layer output); SDXL reads [-2]. pooled is the final-norm hidden state at the EOS token (the highest token id), projected by text_projection when present (the bigG encoder). SDXL uses the pooled output only from that encoder.

Source code in src/mlx_diffuser/models/clip_text.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
def __call__(self, input_ids: mx.array) -> tuple[list[mx.array], mx.array]:
    """``(B, L) int -> (hidden_states, pooled)``.

    ``hidden_states`` has length ``num_layers + 1`` (embeddings then each layer
    output); SDXL reads ``[-2]``. ``pooled`` is the final-norm hidden state at the
    EOS token (the highest token id), projected by ``text_projection`` when present
    (the bigG encoder). SDXL uses the pooled output only from that encoder.
    """
    x = self.text_model.embeddings(input_ids)
    seq = input_ids.shape[-1]
    mask = mx.triu(mx.full((seq, seq), -mx.inf, dtype=x.dtype), k=1)

    hidden_states = [x]
    for layer in self.text_model.encoder.layers:
        x = layer(x, mask)
        hidden_states.append(x)

    last = self.text_model.final_layer_norm(x)
    eos = mx.argmax(input_ids, axis=-1)
    pooled = last[mx.arange(input_ids.shape[0]), eos]
    if self.text_projection is not None:
        pooled = self.text_projection(pooled)
    return hidden_states, pooled

mlx_diffuser.models.CLIPTextConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/clip_text.py 
27
28
29
30
31
32
33
34
35
36
37
38
@dataclasses.dataclass
class CLIPTextConfig(Config):
    vocab_size: int = 49408
    hidden_size: int = 768
    intermediate_size: int = 3072
    num_hidden_layers: int = 12
    num_attention_heads: int = 12
    max_position_embeddings: int = 77
    hidden_act: str = "quick_gelu"  # "gelu" for the bigG encoder
    projection_dim: int = 768
    layer_norm_eps: float = 1e-5
    with_projection: bool = False  # True for the bigG encoder (CLIPTextModelWithProjection)

FLUX.1 (weight-compatible ports)

mlx_diffuser.models.FluxTransformer2DModel

Bases: ModelMixin[FluxConfig]

FLUX.1 text-to-image diffusion transformer (MMDiT). Channels-last latents.

Source code in src/mlx_diffuser/models/flux_transformer.py 
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
class FluxTransformer2DModel(ModelMixin[FluxConfig]):
    """FLUX.1 text-to-image diffusion transformer (MMDiT). Channels-last latents."""

    config_class = FluxConfig

    def __init__(self, config: FluxConfig):
        super().__init__()
        self.config = config
        d = config.inner_dim
        self.time_text_embed = CombinedTimestepTextEmbed(config)
        self.context_embedder = nn.Linear(config.joint_attention_dim, d)
        self.x_embedder = nn.Linear(config.in_channels, d)
        self.transformer_blocks = [FluxTransformerBlock(config) for _ in range(config.num_layers)]
        self.single_transformer_blocks = [
            FluxSingleTransformerBlock(config) for _ in range(config.num_single_layers)
        ]
        self.norm_out = _AdaLNContinuous(d)
        self.proj_out = nn.Linear(d, config.patch_size**2 * config.out_ch)

    def __call__(
        self,
        hidden_states: mx.array,
        timestep: mx.array,
        encoder_hidden_states: mx.array,
        pooled_projections: mx.array,
        img_ids: mx.array,
        txt_ids: mx.array,
        guidance: mx.array | None = None,
        cache: FirstBlockCache | None = None,
    ) -> mx.array:
        """Predict the flow velocity for packed latents ``hidden_states`` ``(B, L_img, 64)``.

        ``timestep`` is the flow time in ``[0, 1]`` (scaled by 1000 internally, matching
        diffusers). ``encoder_hidden_states`` are T5 tokens ``(B, L_txt, 4096)``;
        ``pooled_projections`` is the CLIP pooled embed ``(B, 768)``. ``img_ids`` /
        ``txt_ids`` are the ``(L, 3)`` rotary position ids. ``guidance`` (FLUX-dev only)
        is the distilled guidance scale.
        """
        hidden = self.x_embedder(hidden_states)
        timestep = timestep * 1000
        guidance = guidance * 1000 if guidance is not None else None
        temb = self.time_text_embed(timestep, pooled_projections, guidance)
        encoder = self.context_embedder(encoder_hidden_states)

        ids = mx.concatenate([txt_ids, img_ids], axis=0)
        rope = flux_rope(ids, self.config.axes_dims_rope)

        encoder, hidden = self._run_blocks(hidden, encoder, temb, rope, cache)

        hidden = self.norm_out(hidden, temb)
        return self.proj_out(hidden)

    def _run_blocks(self, hidden, encoder, temb, rope, cache):
        """Run double then single blocks, optionally reusing the First-Block Cache.

        The cache signal is the first double block's contribution to the image stream;
        on a hit we add the cached residual of every remaining block instead of running
        them. Text-stream updates from the skipped blocks do not affect the output (only
        the image stream is read out), so caching the image residual alone is exact.
        """
        if cache is None or not cache.enabled:
            for block in self.transformer_blocks:
                encoder, hidden = block(hidden, encoder, temb, rope)
            for block in self.single_transformer_blocks:
                encoder, hidden = block(hidden, encoder, temb, rope)
            return encoder, hidden

        first_in = hidden
        encoder, hidden = self.transformer_blocks[0](hidden, encoder, temb, rope)
        if cache.should_reuse(hidden - first_in):
            return encoder, first_in + cache.residual
        after_first = hidden
        for block in self.transformer_blocks[1:]:
            encoder, hidden = block(hidden, encoder, temb, rope)
        for block in self.single_transformer_blocks:
            encoder, hidden = block(hidden, encoder, temb, rope)
        cache.residual = hidden - after_first
        return encoder, hidden

__call__(hidden_states, timestep, encoder_hidden_states, pooled_projections, img_ids, txt_ids, guidance=None, cache=None)

Predict the flow velocity for packed latents hidden_states (B, L_img, 64).

timestep is the flow time in [0, 1] (scaled by 1000 internally, matching diffusers). encoder_hidden_states are T5 tokens (B, L_txt, 4096); pooled_projections is the CLIP pooled embed (B, 768). img_ids / txt_ids are the (L, 3) rotary position ids. guidance (FLUX-dev only) is the distilled guidance scale.

Source code in src/mlx_diffuser/models/flux_transformer.py 
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
def __call__(
    self,
    hidden_states: mx.array,
    timestep: mx.array,
    encoder_hidden_states: mx.array,
    pooled_projections: mx.array,
    img_ids: mx.array,
    txt_ids: mx.array,
    guidance: mx.array | None = None,
    cache: FirstBlockCache | None = None,
) -> mx.array:
    """Predict the flow velocity for packed latents ``hidden_states`` ``(B, L_img, 64)``.

    ``timestep`` is the flow time in ``[0, 1]`` (scaled by 1000 internally, matching
    diffusers). ``encoder_hidden_states`` are T5 tokens ``(B, L_txt, 4096)``;
    ``pooled_projections`` is the CLIP pooled embed ``(B, 768)``. ``img_ids`` /
    ``txt_ids`` are the ``(L, 3)`` rotary position ids. ``guidance`` (FLUX-dev only)
    is the distilled guidance scale.
    """
    hidden = self.x_embedder(hidden_states)
    timestep = timestep * 1000
    guidance = guidance * 1000 if guidance is not None else None
    temb = self.time_text_embed(timestep, pooled_projections, guidance)
    encoder = self.context_embedder(encoder_hidden_states)

    ids = mx.concatenate([txt_ids, img_ids], axis=0)
    rope = flux_rope(ids, self.config.axes_dims_rope)

    encoder, hidden = self._run_blocks(hidden, encoder, temb, rope, cache)

    hidden = self.norm_out(hidden, temb)
    return self.proj_out(hidden)

mlx_diffuser.models.FluxConfig dataclass

Bases: Config

Source code in src/mlx_diffuser/models/flux_transformer.py 
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
@dataclasses.dataclass
class FluxConfig(Config):
    patch_size: int = 1
    in_channels: int = 64
    out_channels: int | None = None
    num_layers: int = 19
    num_single_layers: int = 38
    attention_head_dim: int = 128
    num_attention_heads: int = 24
    joint_attention_dim: int = 4096
    pooled_projection_dim: int = 768
    guidance_embeds: bool = False
    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)

    @property
    def inner_dim(self) -> int:
        return self.num_attention_heads * self.attention_head_dim

    @property
    def out_ch(self) -> int:
        return self.out_channels if self.out_channels is not None else self.in_channels

    @classmethod
    def flux_schnell(cls) -> FluxConfig:
        return cls(guidance_embeds=False)

    @classmethod
    def flux_dev(cls) -> FluxConfig:
        return cls(guidance_embeds=True)

mlx_diffuser.models.T5EncoderModel

Bases: ModelMixin[T5Config]

T5 v1.1 text encoder. Produces per-token embeddings for FLUX joint attention.

Source code in src/mlx_diffuser/models/t5.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class T5EncoderModel(ModelMixin[T5Config]):
    """T5 v1.1 text encoder. Produces per-token embeddings for FLUX joint attention."""

    config_class = T5Config

    def __init__(self, config: T5Config):
        super().__init__()
        self.config = config
        self.shared = nn.Embedding(config.vocab_size, config.d_model)
        self.encoder = _Stack(config)

    def __call__(self, input_ids: mx.array, attention_mask: mx.array | None = None) -> mx.array:
        """``input_ids`` ``(B, L)`` -> embeddings ``(B, L, d_model)``.

        Activations run in fp32 (T5 routinely exceeds the fp16 range); only the
        (optionally quantized) weights stay small.
        """
        x = self.shared(input_ids).astype(mx.float32)
        mask = None
        if attention_mask is not None:
            neg = (1.0 - attention_mask.astype(mx.float32)) * -1e9
            mask = neg[:, None, None, :]  # (B, 1, 1, L) additive
        return self.encoder(x, mask)

__call__(input_ids, attention_mask=None)

input_ids (B, L) -> embeddings (B, L, d_model).

Activations run in fp32 (T5 routinely exceeds the fp16 range); only the (optionally quantized) weights stay small.

Source code in src/mlx_diffuser/models/t5.py 
148
149
150
151
152
153
154
155
156
157
158
159
def __call__(self, input_ids: mx.array, attention_mask: mx.array | None = None) -> mx.array:
    """``input_ids`` ``(B, L)`` -> embeddings ``(B, L, d_model)``.

    Activations run in fp32 (T5 routinely exceeds the fp16 range); only the
    (optionally quantized) weights stay small.
    """
    x = self.shared(input_ids).astype(mx.float32)
    mask = None
    if attention_mask is not None:
        neg = (1.0 - attention_mask.astype(mx.float32)) * -1e9
        mask = neg[:, None, None, :]  # (B, 1, 1, L) additive
    return self.encoder(x, mask)

mlx_diffuser.models.T5Config dataclass

Bases: Config

Source code in src/mlx_diffuser/models/t5.py 
27
28
29
30
31
32
33
34
35
36
37
@dataclasses.dataclass
class T5Config(Config):
    vocab_size: int = 32128
    d_model: int = 4096
    d_kv: int = 64
    d_ff: int = 10240
    num_layers: int = 24
    num_heads: int = 64
    relative_attention_num_buckets: int = 32
    relative_attention_max_distance: int = 128
    layer_norm_epsilon: float = 1e-6