e4cfa9d1e9
Fork of dgrauet/Hunyuan3D-2.1-mlx + our generate_e2e.py CLI, env-tunable remesh (HY3D_REMESH_FACES), and HARDWARE.md. Upstream training data (mini_trainset) and demo images stripped — inference needs none of it. Full upstream history: github.com/dgrauet/Hunyuan3D-2.1-mlx
21 KiB
21 KiB
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:
-
norm1 invariant: in each transformer block,
norm1(hidden_states)is computed once. Self-attn (or MDA), reference-attn, and multiview- attn all consume that singlenorm_hs. The block's hidden_states accumulates additive updates from each attention butnorm_hsitself is NEVER recomputed inside the block. -
norm2 invariant: same for
norm2. Computed once. Text cross-attn AND DINO cross-attn both consume the pre-attn2norm2_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 byscaling_factorinternally;vae.decode()divides by it. Callers must NOT scale again. (We doubled it once and conditioning latents ended up 30x too small.) position_voxel_indiceslookup key: indexed by the multiview sequence lengthn_views * L, not per-viewL— matchesmultivew_hidden_states.shape[1]in PT's modules.py.- Dual-stream reference UNet: the checkpoint ships BOTH a main
unet.*(2.5D) AND a vanillaunet_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 - refguidance 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_mrexist 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 |