Hunyuan3D-2.2-mrp-MLX/hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py

"""Stable Diffusion 2.1 UNet for MLX (NHWC layout).

UNet2DConditionModelMLX: standard SD 2.1 UNet with optional 2.5D attention
modules attached externally via load_model._enhance_unet.
"""

from typing import Dict, Optional, Tuple

import os

import mlx.core as mx
import mlx.nn as nn
import numpy as np

from .blocks_mlx import (
    CrossAttnDownBlock2D,
    CrossAttnUpBlock2D,
    DownBlock2D,
    ResnetBlock2D,
    TimestepEmbedding,
    UNetMidBlock2DCrossAttn,
    UpBlock2D,
    get_timestep_embedding,
)


class UNet2DConditionModelMLX(nn.Module):
    """Stable Diffusion 2.1 UNet in MLX (NHWC layout).

    Architecture:
        in_channels=4 (or 12 for concat conditioning)
        block_out_channels=(320, 640, 1280, 1280)
        cross_attention_dim=1024 (SD 2.1)
        down: CrossAttn, CrossAttn, CrossAttn, DownOnly
        mid: CrossAttn
        up: UpOnly, CrossAttn, CrossAttn, CrossAttn
    """

    def __init__(
        self,
        in_channels: int = 4,
        out_channels: int = 4,
        block_out_channels: tuple = (320, 640, 1280, 1280),
        cross_attention_dim: int = 1024,
        attention_head_dim: tuple | int = (5, 10, 20, 20),
        layers_per_block: int = 2,
        transformer_layers_per_block: int = 1,
    ):
        super().__init__()
        self.in_channels = in_channels
        time_embed_dim = block_out_channels[0] * 4  # 1280

        # Normalize attention_head_dim to a list
        if isinstance(attention_head_dim, int):
            attention_head_dim = [attention_head_dim] * len(block_out_channels)
        self._attention_head_dims = list(attention_head_dim)

        # Timestep embedding
        self.time_proj_dim = block_out_channels[0]  # 320
        self.time_embedding = TimestepEmbedding(
            block_out_channels[0], time_embed_dim
        )

        # Input conv
        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], 3, padding=1)

        # Down blocks
        self.down_blocks = []
        output_channel = block_out_channels[0]
        for i, ch in enumerate(block_out_channels):
            is_last = i == len(block_out_channels) - 1
            input_channel = output_channel
            output_channel = ch

            if i < len(block_out_channels) - 1:
                # Diffusers convention: attention_head_dim is actually num_heads.
                # The true per-head dim is ch // num_heads.
                num_heads = self._attention_head_dims[i]
                self.down_blocks.append(
                    CrossAttnDownBlock2D(
                        in_channels=input_channel,
                        out_channels=output_channel,
                        temb_channels=time_embed_dim,
                        num_layers=layers_per_block,
                        num_attention_heads=num_heads,
                        cross_attention_dim=cross_attention_dim,
                        add_downsample=not is_last,
                        transformer_layers_per_block=transformer_layers_per_block,
                    )
                )
            else:
                # Last block: no attention
                self.down_blocks.append(
                    DownBlock2D(
                        in_channels=input_channel,
                        out_channels=output_channel,
                        temb_channels=time_embed_dim,
                        num_layers=layers_per_block,
                        add_downsample=False,
                    )
                )

        # Mid block
        mid_channels = block_out_channels[-1]
        num_mid_heads = self._attention_head_dims[-1]
        self.mid_block = UNetMidBlock2DCrossAttn(
            in_channels=mid_channels,
            temb_channels=time_embed_dim,
            num_attention_heads=num_mid_heads,
            cross_attention_dim=cross_attention_dim,
            transformer_layers_per_block=transformer_layers_per_block,
        )

        # Up blocks (reversed)
        reversed_channels = list(reversed(block_out_channels))
        self.up_blocks = []
        output_channel = reversed_channels[0]
        for i, ch in enumerate(reversed_channels):
            is_last = i == len(reversed_channels) - 1
            prev_output_channel = output_channel
            output_channel = ch
            input_channel = reversed_channels[min(i + 1, len(reversed_channels) - 1)]

            if i > 0:
                # CrossAttn — use reversed num_heads to match down blocks.
                # Diffusers convention: attention_head_dim == num_heads.
                rev_idx = len(block_out_channels) - 1 - i
                num_heads = self._attention_head_dims[max(rev_idx, 0)]
                self.up_blocks.append(
                    CrossAttnUpBlock2D(
                        in_channels=input_channel,
                        prev_output_channel=prev_output_channel,
                        out_channels=output_channel,
                        temb_channels=time_embed_dim,
                        num_layers=layers_per_block + 1,
                        num_attention_heads=num_heads,
                        cross_attention_dim=cross_attention_dim,
                        add_upsample=not is_last,
                        transformer_layers_per_block=transformer_layers_per_block,
                    )
                )
            else:
                # First up block: no attention
                self.up_blocks.append(
                    UpBlock2D(
                        in_channels=input_channel,
                        prev_output_channel=prev_output_channel,
                        out_channels=output_channel,
                        temb_channels=time_embed_dim,
                        num_layers=layers_per_block + 1,
                        add_upsample=True,
                    )
                )

        # Output
        self.conv_norm_out = nn.GroupNorm(32, block_out_channels[0], pytorch_compatible=True)
        self.conv_act = nn.SiLU()
        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)

    def __call__(
        self,
        sample: mx.array,
        timestep: mx.array,
        encoder_hidden_states: Optional[mx.array] = None,
        **kwargs,
    ) -> mx.array:
        # HY3D_MLX_COMPILE=1 routes through mx.compile (fixed-shape denoise
        # loop -> one trace per chunk-shape/ctx-structure; corridorkey
        # precedent 4c660df). Default off = numerics identical to upstream.
        if os.environ.get("HY3D_MLX_COMPILE", "0") == "1":
            fn = getattr(self, "_compiled_forward", None)
            if fn is None:
                fn = mx.compile(self._forward)
                self._compiled_forward = fn
            return fn(sample, timestep, encoder_hidden_states, **kwargs)
        return self._forward(sample, timestep, encoder_hidden_states, **kwargs)

    def _forward(
        self,
        sample: mx.array,
        timestep: mx.array,
        encoder_hidden_states: Optional[mx.array] = None,
        **kwargs,
    ) -> mx.array:
        """
        Args:
            sample: (B, H, W, C_in) noisy latent.
            timestep: (B,) or scalar timestep.
            encoder_hidden_states: (B, L, C_text) text embeddings.
            **kwargs: Extra context passed to transformer blocks
                      (e.g. dino_features for 2.5D attention).

        Returns:
            (B, H, W, C_out) noise prediction.
        """
        # Timestep embedding
        if not isinstance(timestep, mx.array):
            timestep = mx.array([timestep])
        if timestep.ndim == 0:
            timestep = mx.expand_dims(timestep, 0)
        t_emb = get_timestep_embedding(timestep, self.time_proj_dim)
        emb = self.time_embedding(t_emb)  # (B, time_embed_dim)

        # Broadcast for batch
        if emb.shape[0] == 1 and sample.shape[0] > 1:
            emb = mx.broadcast_to(emb, (sample.shape[0],) + emb.shape[1:])

        # Input conv
        sample = self.conv_in(sample)

        # Down
        down_block_res_samples = [sample]
        for block in self.down_blocks:
            if hasattr(block, "has_cross_attention") and block.has_cross_attention:
                sample, res = block(sample, emb, encoder_hidden_states, **kwargs)
            else:
                sample, res = block(sample, emb)
            down_block_res_samples.extend(res)

        # Mid
        sample = self.mid_block(sample, emb, encoder_hidden_states, **kwargs)

        # Up
        for block in self.up_blocks:
            n_res = len(block.resnets)
            res_samples = down_block_res_samples[-n_res:]
            down_block_res_samples = down_block_res_samples[:-n_res]

            if hasattr(block, "has_cross_attention") and block.has_cross_attention:
                sample = block(sample, emb, res_samples, encoder_hidden_states, **kwargs)
            else:
                sample = block(sample, emb, res_samples)

        # Output
        sample = self.conv_act(self.conv_norm_out(sample))
        sample = self.conv_out(sample)
        return sample