"""Load HunyuanPaintPBR models from converted MLX weights. Handles the full weight loading including 2.5D attention modules that get dynamically attached to transformer blocks. """ import os from pathlib import Path from typing import Optional import mlx.core as mx import mlx.nn as nn from mlx.utils import tree_flatten from .dino_mlx import DINOv2MLX from .scheduler_mlx import UniPCMultistepSchedulerMLX from .vae_mlx import AutoencoderKLMLX from .unet.blocks_mlx import Attention from .unet.unet_mlx import UNet2DConditionModelMLX from .unet.attn_processor_mlx import ImageProjModel # --------------------------------------------------------------------------- # 2.5D module containers (match checkpoint key paths) # --------------------------------------------------------------------------- class MDAProcessor(nn.Module): """Material-dimension attention: extra Q/K/V/out for non-primary materials. Key path: attn1.processor.to_{q,k,v,out}_mr.* """ def __init__(self, dim: int): super().__init__() self.to_q_mr = nn.Linear(dim, dim, bias=False) self.to_k_mr = nn.Linear(dim, dim, bias=False) self.to_v_mr = nn.Linear(dim, dim, bias=False) self.to_out_mr = nn.Linear(dim, dim) class RefProcessor(nn.Module): """Reference attention: material-specific V and out projections. Key path: attn_refview.processor.to_{v,out}_mr.* """ def __init__(self, dim: int): super().__init__() self.to_v_mr = nn.Linear(dim, dim, bias=False) self.to_out_mr = nn.Linear(dim, dim) def _enhance_transformer_block(block, dim: int, cross_attention_dim: int): """Add 2.5D attention modules to a BasicTransformerBlock. Attaches modules with key paths matching the checkpoint: - block.attn1.processor -> MDAProcessor - block.attn_dino -> Attention (cross-attn to DINO features) - block.attn_multiview -> Attention (self-attn with RoPE) - block.attn_refview -> Attention + processor """ # Get head config from the block's existing attn1 (set by UNet with correct head_dims) heads = block.attn1.heads head_dim = block.attn1.dim_head # MDA processor on attn1 block.attn1.processor = MDAProcessor(dim) # DINO cross-attention block.attn_dino = Attention( query_dim=dim, cross_attention_dim=cross_attention_dim, heads=heads, dim_head=head_dim, ) # Multiview self-attention block.attn_multiview = Attention( query_dim=dim, heads=heads, dim_head=head_dim, ) # Reference cross-attention block.attn_refview = Attention( query_dim=dim, heads=heads, dim_head=head_dim, ) block.attn_refview.processor = RefProcessor(dim) def _assign_block_ids(unet: UNet2DConditionModelMLX) -> None: """Assign sequential ``_block_id`` to every transformer block. Used on the dual-stream reference UNet so its forward pass populates the capture dict in the same order as the main UNet's reference-attention expects (each transformer block reads ``condition_embed_dict[block_id]``). Does not attach 2.5D modules — unet_dual is plain SD 2.1. """ counter = 0 for blocks in (unet.down_blocks, [unet.mid_block], unet.up_blocks): for block in blocks: if hasattr(block, "attentions"): for attn_block in block.attentions: for tb in attn_block.transformer_blocks: tb._block_id = str(counter) counter += 1 def _enhance_unet(unet: UNet2DConditionModelMLX, cross_attention_dim: int = 1024): """Add 2.5D modules to all transformer blocks in a UNet. Also assigns a unique ``_block_id`` (str) to each transformer block so that the capture-based reference feature extraction can key features without monkey-patching ``__call__``. """ counter = 0 # Down blocks for bi, block in enumerate(unet.down_blocks): if hasattr(block, "attentions"): for ai, attn_block in enumerate(block.attentions): dim = attn_block.in_channels for ti, tb in enumerate(attn_block.transformer_blocks): _enhance_transformer_block(tb, dim, cross_attention_dim) tb._block_id = str(counter) counter += 1 # Mid block if hasattr(unet.mid_block, "attentions"): for ai, attn_block in enumerate(unet.mid_block.attentions): dim = attn_block.in_channels for ti, tb in enumerate(attn_block.transformer_blocks): _enhance_transformer_block(tb, dim, cross_attention_dim) tb._block_id = str(counter) counter += 1 # Up blocks for bi, block in enumerate(unet.up_blocks): if hasattr(block, "attentions"): for ai, attn_block in enumerate(block.attentions): dim = attn_block.in_channels for ti, tb in enumerate(attn_block.transformer_blocks): _enhance_transformer_block(tb, dim, cross_attention_dim) tb._block_id = str(counter) counter += 1 # --------------------------------------------------------------------------- # Reference feature extraction # --------------------------------------------------------------------------- def extract_reference_features( unet: "UNet2DConditionModelMLX", ref_latent: mx.array, ref_text: mx.array, condition_latents: mx.array = None, ) -> dict: """Run a forward pass through the UNet to capture transformer block outputs. Used for reference attention: the reference image is processed once, and the intermediate features at each transformer block are cached. These are then used as key/value context by attn_refview in the main pass. Instead of monkey-patching ``__call__`` (which hangs on Apple Silicon), this passes a ``_capture_dict`` through kwargs. Each ``BasicTransformerBlock`` that has a ``_block_id`` attribute will store its post-self-attention hidden states into this dict during the forward pass. Args: unet: The UNet (main or dual) to extract features from. Transformer blocks must have ``_block_id`` set (done by ``_enhance_unet``). ref_latent: (1, H, W, 4) reference image latent. ref_text: (1, 77, C) text embeddings for reference. condition_latents: (1, H, W, 8) optional normal+position latent concat. Returns: Dict mapping block id (str) to (1, L, C) feature tensors. """ # Prepare input — match the UNet's expected in_channels. expected_in = unet.in_channels if condition_latents is not None and expected_in == 12: unet_input = mx.concatenate([ref_latent, condition_latents], axis=-1) elif expected_in == 12: # Main UNet path (use_dual_stream=False in PT terms): pad with zeros pad = mx.zeros((*ref_latent.shape[:-1], 8)) unet_input = mx.concatenate([ref_latent, pad], axis=-1) else: # Dual-stream UNet path (in_channels=4): pass the latent as-is unet_input = ref_latent # Dict that BasicTransformerBlock will populate during the forward pass capture_dict: dict = {} # Forward pass at timestep 0 (frozen reference). # The _capture_dict kwarg is propagated through UNet -> blocks -> transformer # blocks via **kwargs. _ = unet( unet_input, mx.array([0]), ref_text, _capture_dict=capture_dict, _ref_n_views=1, n_views=0, # disable multiview / MDA paths during capture n_pbr=1, ) mx.synchronize() return capture_dict # --------------------------------------------------------------------------- # Full model loading # --------------------------------------------------------------------------- class HunyuanPaintModelMLX: """Container for all HunyuanPaintPBR model components. Loads weights from mlx-forge converted safetensors. """ def __init__( self, unet: UNet2DConditionModelMLX, vae: AutoencoderKLMLX, dino: DINOv2MLX, image_proj: ImageProjModel, learned_text_clip: dict, scheduler: UniPCMultistepSchedulerMLX, cross_attention_dim: int = 1024, ): self.unet = unet self.vae = vae self.dino = dino self.image_proj = image_proj self.learned_text_clip = learned_text_clip self.scheduler = scheduler self.cross_attention_dim = cross_attention_dim @staticmethod def from_pretrained( weights_source: str = "dgrauet/hunyuan3d-2.1-mlx", ) -> "HunyuanPaintModelMLX": """Load all components from converted weights. Args: weights_source: Either a HuggingFace repo ID (e.g. ``dgrauet/hunyuan3d-2.1-mlx``) which will be downloaded via ``huggingface_hub.snapshot_download`` and cached locally, or an absolute path to a local directory containing ``paint_{unet,vae,dino}.safetensors``. Environment: HUNYUAN3D_MLX_WEIGHTS_DIR — overrides ``weights_source`` if set. """ env_override = os.environ.get("HUNYUAN3D_MLX_WEIGHTS_DIR") if env_override: weights_source = env_override if os.path.isdir(weights_source): weights_dir = Path(weights_source) else: # Treat as HF repo ID. Only fetch the paint_* files we need. from huggingface_hub import snapshot_download weights_dir = Path(snapshot_download( repo_id=weights_source, allow_patterns=[ "paint_unet.safetensors", "paint_vae.safetensors", "paint_dino.safetensors", "config.json", "split_model.json", ], )) # --- VAE --- print("Loading VAE...") vae = AutoencoderKLMLX() vae_w = dict(mx.load(str(weights_dir / "paint_vae.safetensors"))) vae.load_weights(list(vae_w.items())) del vae_w # --- DINOv2 --- # Weights ship as fp16. With 40 transformer blocks the per-op # quantisation error accumulates (output diverges ~5x vs fp32 HF on # the same input — measured). Upcast to fp32 for DINO: the model is # small enough (~1.2 GB weights) that the extra memory is fine, and # the reference conditioning quality depends on clean features. print("Loading DINOv2...") dino = DINOv2MLX() dino_w = {k: v.astype(mx.float32) for k, v in mx.load(str(weights_dir / "paint_dino.safetensors")).items()} dino.load_weights(list(dino_w.items())) del dino_w # --- UNet --- print("Loading UNet...") unet = UNet2DConditionModelMLX( in_channels=12, out_channels=4, block_out_channels=(320, 640, 1280, 1280), cross_attention_dim=1024, attention_head_dim=(5, 10, 20, 20), ) # Add 2.5D modules to match checkpoint structure _enhance_unet(unet, cross_attention_dim=1024) unet_w = dict(mx.load(str(weights_dir / "paint_unet.safetensors"))) # The checkpoint contains TWO UNets: # unet.* — main 2.5D UNet that runs the denoising loop # unet_dual.* — vanilla SD 2.1 UNet that processes the reference # image once to extract per-block features for # reference-attention. Equivalent to PT's # self.unet_dual when use_dual_stream=True. stripped = {k[len("unet."):]: v for k, v in unet_w.items() if k.startswith("unet.")} stripped_dual = {k[len("unet_dual."):]: v for k, v in unet_w.items() if k.startswith("unet_dual.")} model_keys = set(k for k, _ in tree_flatten(unet.parameters())) # Handle .to_out_mr.0. -> .to_out_mr. (ModuleList artifact) normalized = {} for k, v in stripped.items(): nk = k.replace(".to_out_mr.0.", ".to_out_mr.") nk = nk.replace(".to_out_albedo.0.", ".to_out_albedo.") normalized[nk] = v matched = [(k, v) for k, v in normalized.items() if k in model_keys] unet.load_weights(matched) # The remaining keys in `stripped` (learned_text_clip_* and # image_proj_model_dino.*) belong to separate modules and are # consumed below — they are not truly "unmatched". print(f" Loaded {len(matched)}/{len(stripped)} main UNet weights" f" (remaining keys load into text embeds + DINO proj below)") # --- Learned text embeddings --- learned = {} for token in ("albedo", "mr", "ref"): key = f"learned_text_clip_{token}" if key in stripped: learned[token] = stripped[key] else: learned[token] = mx.zeros((77, 1024)) # --- ImageProjModel --- image_proj = ImageProjModel( cross_attention_dim=1024, clip_embeddings_dim=1536, clip_extra_context_tokens=4, ) proj_weights = [ (k.replace("image_proj_model_dino.", ""), v) for k, v in stripped.items() if k.startswith("image_proj_model_dino.") ] if proj_weights: image_proj.load_weights(proj_weights) # --- Dual-stream reference UNet (vanilla SD 2.1 — no 2.5D modules) --- # Takes 4-channel reference latent only (no normal/position concat # — those are reserved for the main 12-ch UNet). unet_dual = None if stripped_dual: print("Loading dual-stream reference UNet...") unet_dual = UNet2DConditionModelMLX( in_channels=4, out_channels=4, block_out_channels=(320, 640, 1280, 1280), cross_attention_dim=1024, attention_head_dim=(5, 10, 20, 20), ) # NOTE: do NOT call _enhance_unet — unet_dual is plain SD 2.1. dual_keys = set(k for k, _ in tree_flatten(unet_dual.parameters())) matched_dual = [(k, v) for k, v in stripped_dual.items() if k in dual_keys] unet_dual.load_weights(matched_dual) _assign_block_ids(unet_dual) print(f" Loaded {len(matched_dual)}/{len(stripped_dual)}" f" dual UNet weights") del unet_w, stripped, stripped_dual scheduler = UniPCMultistepSchedulerMLX() model = HunyuanPaintModelMLX( unet=unet, vae=vae, dino=dino, image_proj=image_proj, learned_text_clip=learned, scheduler=scheduler, ) model.unet_dual = unet_dual print("All components loaded.") return model