# Hunyuan3D-2.1 MLX — Forward Pass (Stage 2: Texture Synthesis) End-to-end data flow from a reference image + 3D mesh to a textured GLB. All tensors are in MLX (NHWC), all weights live on Apple Silicon's unified memory. ``` REFERENCE IMAGE MESH (.glb) (PIL, 512×512) (Trimesh, native UVs) │ │ │ ▼ │ ┌─────────────────────┐ │ │ MeshRenderMLX │ │ │ (Metal rasterizer) │ │ └─────────┬───────────┘ │ │ │ 6 azim × 1 elev (+ top/bot) │ │ │ ┌──────────────┼─────────────────┐ │ ▼ ▼ ▼ │ normal_maps position_maps alpha (face IDs) │ (6, H, W, 3) (6, H, W, 3) used for view sel. │ │ │ │ ▼ ▼ │ ┌───────────────────────────────┐ │ │ VAE.encode (×3, frozen) │ │ │ → z_normal, z_pos (6, h, w, 4)│ │ │ + raw position_maps (no VAE) │ │ └──────────────┬────────────────┘ │ │ │ ▼ │ ┌─────────────────────────────────┐ │ │ calc_multires_voxel_idxs │ │ │ → {seq_len: voxel_indices, │ │ │ voxel_resolution} │ │ │ for each level [4096,1024,256,64]│ │ └──────────────┬──────────────────┘ ▼ │ ┌─────────────────────┐ │ │ DINOv2 (frozen) │ │ │ + ImageProjModel │ │ │ → dino_proj │ │ │ (1, N_tok, 1024) │ │ └──────────┬──────────┘ │ │ │ │ ┌───────────────┘ │ │ ┌──────────▼───▼──────────────────────────────┐ │ extract_reference_features (run ONCE) │ │ → unet_dual (4-ch, vanilla SD2.1) │ │ forward(ref_latent, t=0, text_ref) │ │ capture_dict[block_id] = norm_pre_attn1 │ │ → ref_features dict (16 transformer blocks) │ └──────────────┬───────────────────────────────┘ │ ┌───────┴───────┐ │ LEARNED │ │ TEXT TOKENS │ │ (77, 1024) │ │ per material │ │ + ref + neg │ └───────┬───────┘ │ ▼ ┌───────────────────────────────────────────────────────────────────┐ │ DENOISING LOOP (15 steps, UniPC, v_prediction) │ │ │ │ latents ~ N(0,1) shape (n_pbr*n_views, h, w, 4) = (12, h, w, 4)│ │ │ │ for t in scheduler.timesteps: │ │ for chunk_view in chunks(n_views=6): │ │ │ │ unet_in = concat(latents, z_normal, z_pos) → (12, h, w, 12)│ │ │ │ ┌─────────── 3 CFG passes (PT-faithful) ───────────┐ │ │ │ pred_uncond = UNet(unet_in, text=NEG) │ │ │ │ no DINO, no ref_features │ │ │ │ pred_ref = UNet(unet_in, text=POS, │ │ │ │ ref_features=ref_features) │ │ │ │ no DINO yet │ │ │ │ pred_full = UNet(unet_in, text=POS, │ │ │ │ ref_features=ref_features, │ │ │ │ dino_features=dino_proj, │ │ │ │ position_voxel_indices=...) │ │ │ └───────────────────────────────────────────────────┘ │ │ │ │ view_scale = cam_mapping(azim) per view ∈ [1, 2, ..., 5] │ │ guided = pred_uncond │ │ + g · vs · (pred_ref - pred_uncond) │ │ + g · vs · (pred_full - pred_ref) │ │ │ │ latents = scheduler.step(noise_guided, t, latents) │ └───────────────────────────────────────────────────────────────────┘ │ ▼ ┌────────────────────────────┐ │ VAE.decode (per view) │ │ → albedo (6, 512, 512, 3) │ │ mr (6, 512, 512, 3) │ └─────────────┬───────────────┘ │ ▼ ┌────────────────────────────┐ │ RealESRGAN x4 (MLX) │ │ → 512² → 2048² per view │ └─────────────┬───────────────┘ │ ▼ ┌────────────────────────────────────────┐ │ back_project (per view) │ │ for each view: project_texture, │ │ project_cos_map, boundary │ │ cos = view_weight * (cos**bake_exp) │ └─────────────┬──────────────────────────┘ │ ▼ ┌────────────────────────────────────────┐ │ fast_bake_texture (mode="weighted") │ │ cosine-weighted blend across all views │ │ (matches PT MeshRender.fast_bake_texture)│ │ → atlas (2048, 2048, 3) + trust mask │ └─────────────┬──────────────────────────┘ │ ▼ ┌────────────────────────────────────────┐ │ uv_inpaint (mesh-aware) │ │ 1. mesh_vertex_inpaint (Python port │ │ of meshVerticeInpaint): │ │ - seed vtx_color from texels │ │ - propagate via 3D adjacency │ │ - face barycentric raster +4 px │ │ conservative margin │ │ 2. cv2.INPAINT_NS radius=3 (PT parity) │ │ 3. EDT nearest-fill post-pass │ │ (MLX-specific: pads UV gutters so │ │ 3D viewers doing bilinear sampling │ │ across island edges don't pull in │ │ atlas background → no black lines) │ └─────────────┬──────────────────────────┘ │ ▼ ┌────────────────────────────┐ │ set_texture(atlas) │ │ set_texture_mr(atlas_mr) │ │ save_mesh(.obj, downsample=False) │ │ + trimesh export → .glb │ └────────────────────────────┘ ``` ## UNet (per CFG pass) — internal data flow ``` sample (B*n_pbr*n_views, h, w, 12) │ ▼ conv_in → (B*n_pbr*n_views, h, w, 320) │ ├─────────────────────── DownBlocks (3 cross-attn + 1 plain) ──────┐ │ │ │ CrossAttnDownBlock2D × 3: │ │ for each ResNet+Transformer pair: │ │ hidden = ResNet(hidden, temb) │ │ hidden = Transformer2DModel(hidden, text, **kwargs) │ │ norm + proj_in │ │ for each transformer_block: │ │ ┌──────── BasicTransformerBlock ────────┐ │ │ │ norm_hs = norm1(hidden) │ │ │ │ │ │ │ │ ─ Step 1: SELF-ATTN (or MDA) ─ │ │ │ │ if MDA enabled (n_views > 0): │ │ │ │ split albedo/MR by index │ │ │ │ albedo: attn1(norm_hs[albedo]) │ │ │ │ mr: attn1.processor.{q,k,v,out}_mr│ │ │ │ else: attn_out = attn1(norm_hs) │ │ │ │ hidden += attn_out │ │ │ │ │ │ │ │ ─ Step 2: REFERENCE-ATTN ─ │ │ │ │ if ref_features and _block_id: │ │ │ │ query_albedo = norm_hs[:,0,...] │ │ │ │ ref_ctx = ref_features[block_id] │ │ │ │ Q,K shared, V per material: │ │ │ │ albedo: attn_refview.to_v │ │ │ │ mr: processor.to_v_mr │ │ │ │ sdpa(Q, K, V_per_mat) │ │ │ │ hidden += per-material output │ │ │ │ │ │ │ │ ─ Step 3: MULTIVIEW-ATTN (RoPE 3D) ─ │ │ │ │ if n_views > 1: │ │ │ │ mv_input = norm_hs.reshape(B*n_pbr, n_views*L, C) │ │ │ Q,K = rotate(Q,K, voxel_indices, voxel_res) │ │ │ mv_out = sdpa(Q, K, V) (flash-style) │ │ │ hidden += mv_out │ │ │ │ │ │ │ │ ─ Step 4: CROSS-ATTN (text) ─ │ │ │ │ hidden += attn2(norm2(hidden), text)│ │ │ │ │ │ │ │ ─ Step 5: DINO CROSS-ATTN ─ │ │ │ │ if dino_features: │ │ │ │ hidden += attn_dino(norm_dino(h), │ │ │ │ dino_features) │ │ │ │ │ │ │ │ ─ Step 6: FFN ─ │ │ │ │ hidden += ff(norm3(hidden)) │ │ │ └────────────────────────────────────────┘ │ │ proj_out + residual │ │ output_states.append(hidden) │ │ Downsample2D │ │ │ │ DownBlock2D × 1 (no attention): │ │ ResNet × 2 │ │ │ ▼ │ mid_block (UNetMidBlock2DCrossAttn) │ ResNet → Transformer (same 6-step block) → ResNet │ │ │ ▼ │ ├─────────────────── UpBlocks (1 plain + 3 cross-attn) ◀────────────┘ │ UpBlock2D × 1 (concat skip + ResNet × 3 + Upsample) │ CrossAttnUpBlock2D × 3 (concat skip + ResNet+Transformer × 3 + Upsample) │ ▼ conv_norm_out + SiLU + conv_out → (B*n_pbr*n_views, h, w, 4) noise prediction ``` ## Porting principle: match PT config exactly During development we drifted from the PT reference config (introduced `fast_bake_texture(mode="face_wta")`, `uv_inpaint(method="edt")`, increased `max_num_view`, halved `render_size`/`texture_size`). Each deviation was added to "work around" an artifact — and each one silently masked a still-unfixed port bug upstream. Re-aligning the defaults with PT after the port bugs were resolved sharpened every view, eliminated UV seam cracks, and matched PT quality. **Rule**: the first working version must match the reference config exactly. Only deviate when there is a target-framework constraint documented in code (e.g. `texture_size=2048` because the MLX Metal rasterizer lacks tiling and times out on the command buffer at PT's 4096). Every other config-level "knob" is almost certainly a port bug waiting to be found. ## Critical invariants (learned the hard way) Two classes of subtle bugs produced visibly degraded output for us; both follow the same pattern — **a norm tensor is computed once and every attention path in that step must reuse it, never recompute against the updated hidden_states**: 1. **norm1 invariant**: in each transformer block, `norm1(hidden_states)` is computed once. Self-attn (or MDA), reference-attn, and multiview- attn all consume that single `norm_hs`. The block's hidden_states accumulates additive updates from each attention but `norm_hs` itself is NEVER recomputed inside the block. 2. **norm2 invariant**: same for `norm2`. Computed once. Text cross-attn AND DINO cross-attn both consume the pre-attn2 `norm2_hs`. Recomputing after attn2 feeds DINO a signal off by attn2's residual, which weakens the reference/style conditioning downstream. Other PT-parity traps to watch for when porting: - **VAE encode/decode scale**: `vae.encode()` multiplies by `scaling_factor` internally; `vae.decode()` divides by it. Callers must NOT scale again. (We doubled it once and conditioning latents ended up 30x too small.) - **`position_voxel_indices` lookup key**: indexed by the *multiview* sequence length `n_views * L`, not per-view `L` — matches `multivew_hidden_states.shape[1]` in PT's modules.py. - **Dual-stream reference UNet**: the checkpoint ships BOTH a main `unet.*` (2.5D) AND a vanilla `unet_dual.*` (plain SD 2.1) for reference feature extraction. Must load both. - **3-pass CFG DINO/ref routing**: uncond has no DINO/no ref; ref has ref only; full has both. Passing DINO to both ref AND full collapses the `full - ref` guidance term. - **DINOv2**: run at native 518x518 with fp32 weights — pos_embed interpolation to 224 amplifies ~450x through 40 blocks. - **Per-material ref-attn V/out**: `attn_refview.processor.to_v_mr` / `to_out_mr` exist in the checkpoint and are required for non-albedo materials to get reference conditioning. - **Diffusers `attention_head_dim`**: is a misnomer — it's the number of heads, not per-head dim. `head_dim = channels // attention_head_dim`. ## Tensor shapes for default config (resolution=512, 6 views, 2 PBR materials) | Stage | Tensor | Shape | Notes | |-------|--------|-------|-------| | Conditioning | normal_maps (PIL) | 6 × (512, 512, 3) | RGB float [0,1] | | Conditioning | position_maps (PIL) | 6 × (512, 512, 3) | RGB float [0,1] | | Conditioning | reference (PIL) | (512, 512, 3) | white background | | VAE encoded | z_normal | (6, 64, 64, 4) | scaled by 0.18215 | | VAE encoded | z_pos | (6, 64, 64, 4) | scaled by 0.18215 | | Voxel idx | level 0 | (1, 6×4096, 3), vox_res=512 | for L=4096 multiview attn | | Voxel idx | level 1 | (1, 6×1024, 3), vox_res=256 | for L=1024 | | Voxel idx | level 2 | (1, 6×256, 3), vox_res=128 | for L=256 | | Voxel idx | level 3 | (1, 6×64, 3), vox_res=64 | for L=64 | | DINOv2 | dino_proj | (1, 4 tokens, 1024) | from ImageProjModel | | Ref features | per block | dict[16 blocks] → (1, 4096, C) at L0 | from unet_dual | | Latent init | latents | (12, 64, 64, 4) | n_pbr × n_views | | UNet input | unet_in | (12, 64, 64, 12) | latent + z_normal + z_pos | | UNet output | noise_pred | (12, 64, 64, 4) | per-view, per-material | | Decoded | albedo[i] | (512, 512, 3) | uint8 RGB | | Super-res | albedo[i] | (2048, 2048, 3) | RealESRGAN x4 | | Bake atlas | albedo atlas | (2048, 2048, 3) | float [0,1] | ## Key references in the codebase | Component | File | |-----------|------| | Pipeline orchestration | `hy3dpaint/textureGenPipeline_mlx.py` | | Inference loop, CFG, voxel calc | `hy3dpaint/hunyuanpaintpbr_mlx/inference.py` | | Model loading + ref-feature extraction | `hy3dpaint/hunyuanpaintpbr_mlx/load_model.py` | | UNet construction | `hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py` | | BasicTransformerBlock + 6-step forward | `hy3dpaint/hunyuanpaintpbr_mlx/unet/blocks_mlx.py` | | Per-material V/out, RoPE, MDA processors | `hy3dpaint/hunyuanpaintpbr_mlx/unet/attn_processor_mlx.py` | | 3D voxel index computation | `hy3dpaint/hunyuanpaintpbr_mlx/unet/voxel_indices.py` | | VAE | `hy3dpaint/hunyuanpaintpbr_mlx/vae_mlx.py` | | DINOv2 + ImageProjModel | `hy3dpaint/hunyuanpaintpbr_mlx/dino_mlx.py` | | UniPC scheduler (v_prediction) | `hy3dpaint/hunyuanpaintpbr_mlx/scheduler_mlx.py` | | Metal rasterizer + bake | `hy3dpaint/DifferentiableRenderer/mesh_render_mlx.py` | | Mesh-aware UV inpaint | `hy3dpaint/DifferentiableRenderer/mesh_inpaint_py.py` | | RealESRGAN super-res | `hy3dpaint/utils/image_super_utils_mlx.py` |