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
390 lines
16 KiB
Python
390 lines
16 KiB
Python
"""Full inference for HunyuanPaintPBR with CFG + multi-view.
|
|
|
|
Reproduces the PyTorch pipeline's denoising loop:
|
|
- 6-view generation (albedo + metallic-roughness)
|
|
- Classifier-free guidance (triple-batch: uncond / ref / full)
|
|
- Learned material text embeddings
|
|
- DINO feature conditioning
|
|
- View-dependent guidance scaling
|
|
"""
|
|
|
|
from typing import Dict, List, Optional
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from .dino_mlx import preprocess_for_dino
|
|
from .load_model import HunyuanPaintModelMLX, extract_reference_features
|
|
|
|
|
|
def _compute_dino_features_hf(model, ref_u8: np.ndarray) -> mx.array:
|
|
"""Run HF DINOv2-giant once on the reference image.
|
|
|
|
Uses fp16 on CPU (1.2GB vs 4.4GB fp32) to avoid OOM with the paint
|
|
UNet + dual UNet + VAE + super-res all resident. Frees the model
|
|
after feature extraction.
|
|
|
|
Falls back to pure MLX DINO if HF transformers isn't available.
|
|
"""
|
|
try:
|
|
import torch
|
|
import gc as _gc
|
|
from transformers import AutoImageProcessor, AutoModel
|
|
from PIL import Image as PILImage
|
|
except Exception:
|
|
from .dino_mlx import preprocess_for_dino
|
|
return model.dino(preprocess_for_dino(ref_u8))
|
|
|
|
proc = AutoImageProcessor.from_pretrained("facebook/dinov2-giant")
|
|
pt_dino = AutoModel.from_pretrained(
|
|
"facebook/dinov2-giant", torch_dtype=torch.float16,
|
|
)
|
|
for p in pt_dino.parameters():
|
|
p.requires_grad_(False)
|
|
pt_in = proc(images=PILImage.fromarray(ref_u8), return_tensors="pt")
|
|
with torch.no_grad():
|
|
pt_feat = pt_dino(pt_in.pixel_values.half())[0]
|
|
out = mx.array(pt_feat.float().numpy())
|
|
del proc, pt_dino, pt_in, pt_feat
|
|
_gc.collect()
|
|
return out
|
|
|
|
|
|
def _cam_mapping(azim: float) -> float:
|
|
"""View-dependent guidance scale (per the original pipeline)."""
|
|
azim = azim % 360
|
|
if 0 <= azim < 90:
|
|
return azim / 90.0 + 1.0
|
|
elif 90 <= azim < 330:
|
|
return 2.0
|
|
else:
|
|
return -azim / 90.0 + 5.0
|
|
|
|
|
|
def generate_multiview_pbr(
|
|
model: HunyuanPaintModelMLX,
|
|
normal_maps: List[np.ndarray],
|
|
position_maps: List[np.ndarray],
|
|
reference_image: np.ndarray,
|
|
camera_azims: Optional[List[float]] = None,
|
|
num_inference_steps: int = 15,
|
|
guidance_scale: float = 3.0,
|
|
view_size: int = 512,
|
|
seed: int = 0,
|
|
) -> Dict[str, List[np.ndarray]]:
|
|
"""Generate multiview PBR textures.
|
|
|
|
Args:
|
|
model: Loaded HunyuanPaintModelMLX.
|
|
normal_maps: List of N (H, W, 3) float32 normal renders in [0, 1].
|
|
position_maps: List of N (H, W, 3) float32 position renders in [0, 1].
|
|
reference_image: (H, W, 3) uint8 or float32 reference image.
|
|
camera_azims: Azimuth angles for each view (for view-dependent guidance).
|
|
num_inference_steps: Number of denoising steps.
|
|
guidance_scale: CFG scale.
|
|
view_size: Generation resolution.
|
|
seed: Random seed.
|
|
|
|
Returns:
|
|
dict with 'albedo' and 'mr' keys, each a list of N (H, W, 3) uint8 arrays.
|
|
"""
|
|
mx.random.seed(seed)
|
|
|
|
n_views = len(normal_maps)
|
|
n_pbr = 2 # albedo + mr
|
|
latent_size = view_size // 8
|
|
scaling_factor = 0.18215
|
|
|
|
if camera_azims is None:
|
|
camera_azims = [0.0] * n_views
|
|
|
|
# ------------------------------------------------------------------
|
|
# 1. Encode conditions
|
|
# ------------------------------------------------------------------
|
|
print(" Encoding conditions...")
|
|
|
|
def encode_images(images):
|
|
"""Encode list of images to latent space."""
|
|
arrays = []
|
|
for img in images:
|
|
if img.dtype == np.uint8:
|
|
img = img.astype(np.float32) / 255.0
|
|
from PIL import Image as PILImage
|
|
pil = PILImage.fromarray((img * 255).astype(np.uint8)).resize(
|
|
(view_size, view_size)
|
|
)
|
|
arr = np.array(pil).astype(np.float32) / 255.0
|
|
arrays.append(arr)
|
|
batch = mx.array(np.stack(arrays)) # (N, H, W, 3)
|
|
# vae.encode already multiplies by self.scaling_factor internally;
|
|
# do NOT re-multiply (was double-scaling latents to ~0.033 of the
|
|
# expected magnitude, making ref_latents and z_normal/z_pos
|
|
# essentially invisible to the UNet → washed-out diffusion output).
|
|
return model.vae.encode(batch * 2 - 1)
|
|
|
|
z_normal = encode_images(normal_maps) # (N, h, w, 4)
|
|
z_pos = encode_images(position_maps) # (N, h, w, 4)
|
|
mx.synchronize()
|
|
|
|
# Multi-resolution voxel indices for 3D RoPE in multiview attention.
|
|
# Built from the raw position maps (NOT the VAE-encoded latents).
|
|
from .unet.voxel_indices import calc_multires_voxel_idxs
|
|
pos_resized = []
|
|
for pm in position_maps:
|
|
if pm.dtype == np.uint8:
|
|
pm = pm.astype(np.float32) / 255.0
|
|
from PIL import Image as PILImage
|
|
pil = PILImage.fromarray((pm * 255).astype(np.uint8)).resize(
|
|
(view_size, view_size)
|
|
)
|
|
pos_resized.append(np.array(pil).astype(np.float32) / 255.0)
|
|
# (N, H, W, 3) -> (1, N, 3, H, W)
|
|
pos_np = np.stack(pos_resized).transpose(0, 3, 1, 2)[None]
|
|
voxel_indices = calc_multires_voxel_idxs(pos_np)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 2. DINO features
|
|
# ------------------------------------------------------------------
|
|
print(" Extracting DINO features...")
|
|
if reference_image.dtype != np.uint8:
|
|
ref_u8 = (reference_image * 255).astype(np.uint8)
|
|
else:
|
|
ref_u8 = reference_image
|
|
# Pure MLX DINO at native 518 (1370 tokens). Pool spatial grid from
|
|
# 37x37 down to 16x16 so the final token count (257) matches what
|
|
# HF DINOv2 produces at 224 input — the distribution attn_dino's
|
|
# to_k / to_v were trained against. HF DINO bridge OOMs 32 GB in
|
|
# both fp32 and fp16.
|
|
dino_in = preprocess_for_dino(ref_u8)
|
|
dino_feat = model.dino(dino_in) # (1, 1370, 1536)
|
|
# Separate CLS + patch grid, avg-pool 37x37 -> 16x16
|
|
import math as _math
|
|
N_full = dino_feat.shape[1] - 1
|
|
g_old = int(round(_math.sqrt(N_full)))
|
|
if g_old * g_old == N_full and g_old > 16:
|
|
cls = dino_feat[:, :1, :]
|
|
patch = dino_feat[:, 1:, :].reshape(1, g_old, g_old, -1)
|
|
# Downsample to 16x16 via adaptive average pool (bilinear in MLX)
|
|
up = nn.Upsample(scale_factor=(16.0 / g_old, 16.0 / g_old),
|
|
mode="linear")
|
|
patch_d = up(patch).reshape(1, 16 * 16, -1)
|
|
dino_feat = mx.concatenate([cls, patch_d], axis=1)
|
|
dino_proj = model.image_proj(dino_feat)
|
|
mx.synchronize()
|
|
|
|
# ------------------------------------------------------------------
|
|
# 3. Text embeddings for CFG
|
|
# ------------------------------------------------------------------
|
|
# Per-material learned tokens
|
|
text_albedo = model.learned_text_clip["albedo"] # (77, 1024)
|
|
text_mr = model.learned_text_clip["mr"]
|
|
# PT (pipeline.py:275) sets ``negative_prompt_embeds = prompt_embeds``
|
|
# when use_learned_text_clip is True — the uncond batch reuses the
|
|
# SAME learned material tokens, not zeros. CFG signal then comes
|
|
# from ref_features/DINO presence, not from text. Using zeros here
|
|
# was pushing the uncond prediction off-distribution and weakening
|
|
# the guidance.
|
|
|
|
# ------------------------------------------------------------------
|
|
# 2b. Reference features — capture norm_hidden_states from each
|
|
# transformer block by running a single forward pass on the
|
|
# reference image latent. Matches PyTorch mode "w".
|
|
# ------------------------------------------------------------------
|
|
print(" Extracting reference features...")
|
|
ref_latent = encode_images([reference_image]) # (1, h, w, 4)
|
|
text_ref = model.learned_text_clip["ref"] # PT uses learned_text_clip_ref
|
|
ref_text_batch = text_ref[None] # (1, 77, 1024)
|
|
# Use the dual-stream reference UNet when available — its weights are
|
|
# different from the main UNet and are specifically trained to extract
|
|
# ref-attention features. Falling back to the main UNet (as we did) gives
|
|
# noticeably worse reference grounding.
|
|
ref_unet = getattr(model, "unet_dual", None) or model.unet
|
|
ref_features = extract_reference_features(
|
|
ref_unet, ref_latent, ref_text_batch
|
|
)
|
|
mx.synchronize()
|
|
|
|
# Build per-sample text: (N_pbr * N_views, 77, 1024)
|
|
# Order: [albedo_v0, albedo_v1, ..., mr_v0, mr_v1, ...]
|
|
text_full = mx.concatenate([
|
|
mx.broadcast_to(text_albedo[None], (n_views, 77, 1024)),
|
|
mx.broadcast_to(text_mr[None], (n_views, 77, 1024)),
|
|
], axis=0) # (N_pbr*N_views, 77, 1024) = (12, 77, 1024)
|
|
|
|
# Negative and positive text batches are identical (PT pipeline.py:275).
|
|
text_neg_batch = text_full
|
|
|
|
# CFG triple: [negative, ref(=full), full]
|
|
text_cfg = mx.concatenate([text_neg_batch, text_full, text_full], axis=0)
|
|
# (36, 77, 1024)
|
|
|
|
# DINO features for CFG: replicate for all 36 samples
|
|
# But set to zero for negative (uncond) batch
|
|
dino_for_cfg = mx.broadcast_to(
|
|
dino_proj, (n_pbr * n_views, dino_proj.shape[1], dino_proj.shape[2])
|
|
)
|
|
dino_zeros = mx.zeros_like(dino_for_cfg)
|
|
dino_cfg = mx.concatenate([dino_zeros, dino_for_cfg, dino_for_cfg], axis=0)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 4. View-dependent guidance scale
|
|
# ------------------------------------------------------------------
|
|
view_scales = [_cam_mapping(a) for a in camera_azims]
|
|
# Repeat for materials: [albedo_scales, mr_scales]
|
|
view_scale_flat = np.array(view_scales * n_pbr, dtype=np.float32)
|
|
view_scale = mx.array(view_scale_flat)[:, None, None, None] # (12, 1, 1, 1)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 5. Prepare condition latents for concat
|
|
# ------------------------------------------------------------------
|
|
# Replicate conditions for N_pbr materials
|
|
z_normal_rep = mx.concatenate([z_normal] * n_pbr, axis=0) # (12, h, w, 4)
|
|
z_pos_rep = mx.concatenate([z_pos] * n_pbr, axis=0) # (12, h, w, 4)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 6. Denoising loop
|
|
# ------------------------------------------------------------------
|
|
latents = mx.random.normal((n_pbr * n_views, latent_size, latent_size, 4))
|
|
model.scheduler.set_timesteps(num_inference_steps)
|
|
|
|
print(f" Denoising ({num_inference_steps} steps, {n_views} views, CFG={guidance_scale})...")
|
|
|
|
for i, t in enumerate(model.scheduler.timesteps):
|
|
t_int = int(t)
|
|
|
|
# Process all n_views in a single forward so the multiview
|
|
# attention path inside BasicTransformerBlock fires (it guards on
|
|
# n_views > 1). Combined with the 3D RoPE positional encoding
|
|
# from calc_multires_voxel_idxs and the norm_hs-reuse fix, the
|
|
# cross-view consistency this provides is a big part of matching
|
|
# PT's output.
|
|
chunk_size = n_views
|
|
t_arr = mx.array([t_int])
|
|
noise_guided = mx.zeros_like(latents)
|
|
|
|
for v_start in range(0, n_views, chunk_size):
|
|
v_end = min(v_start + chunk_size, n_views)
|
|
n_chunk = v_end - v_start
|
|
|
|
# Gather chunk indices for both materials
|
|
# Order: [albedo_v_start..v_end, mr_v_start..v_end]
|
|
# Maps into full latent array ordered as
|
|
# [albedo_0..N-1, mr_0..N-1]
|
|
chunk_idx = list(range(v_start, v_end)) + [
|
|
n_views + j for j in range(v_start, v_end)
|
|
]
|
|
lat_chunk = latents[mx.array(chunk_idx)] # (2*n_chunk, h, w, 4)
|
|
zn_chunk = z_normal_rep[mx.array(chunk_idx)]
|
|
zp_chunk = z_pos_rep[mx.array(chunk_idx)]
|
|
unet_in = mx.concatenate([lat_chunk, zn_chunk, zp_chunk], axis=-1)
|
|
|
|
n_c = len(chunk_idx)
|
|
text_neg_c = text_neg_batch[:n_c]
|
|
text_full_c = text_full[mx.array(chunk_idx)]
|
|
dino_c = dino_for_cfg[:n_c]
|
|
|
|
# 3 CFG passes (sequential to save memory).
|
|
# Mirrors the PyTorch pipeline (hunyuanpaintpbr/pipeline.py:300+,
|
|
# hunyuanpaintpbr/pipeline.py:685):
|
|
# uncond -> ref_scale=0, dino=zero, prompt=neg
|
|
# ref -> ref_scale=1, dino=zero, prompt=pos
|
|
# full -> ref_scale=1, dino=full, prompt=pos
|
|
# Then: noise = uncond + g*v*(ref - uncond) + g*v*(full - ref).
|
|
# Earlier we passed DINO to BOTH ref and full, making
|
|
# (full - ref) ~= 0 and dropping half of the guidance.
|
|
|
|
# Slice voxel indices to the current chunk's views
|
|
chunk_view_ids = list(range(v_start, v_end))
|
|
chunk_voxel = {}
|
|
for seq_len, vd in voxel_indices.items():
|
|
vi_full = vd["voxel_indices"] # (1, n_views * GxG, 3)
|
|
gxg = vi_full.shape[1] // n_views
|
|
# Pick this chunk's views
|
|
vi_per_view = vi_full.reshape(1, n_views, gxg, 3)
|
|
vi_chunk = vi_per_view[:, mx.array(chunk_view_ids), :, :]
|
|
vi_chunk = vi_chunk.reshape(1, len(chunk_view_ids) * gxg, 3)
|
|
chunk_voxel[len(chunk_view_ids) * gxg] = {
|
|
"voxel_indices": vi_chunk,
|
|
"voxel_resolution": vd["voxel_resolution"],
|
|
}
|
|
|
|
# uncond: no DINO, no ref features, neg text
|
|
pred_uncond = model.unet(
|
|
unet_in, t_arr, text_neg_c,
|
|
n_views=n_chunk, n_pbr=n_pbr,
|
|
position_voxel_indices=chunk_voxel,
|
|
)
|
|
mx.synchronize()
|
|
|
|
# ref: ref_features ON, DINO OFF, pos text
|
|
ctx_ref = dict(
|
|
n_views=n_chunk, n_pbr=n_pbr,
|
|
position_voxel_indices=chunk_voxel,
|
|
)
|
|
if ref_features is not None:
|
|
ctx_ref["ref_features"] = ref_features
|
|
pred_ref = model.unet(unet_in, t_arr, text_full_c, **ctx_ref)
|
|
mx.synchronize()
|
|
|
|
# full: ref_features ON, DINO ON, pos text
|
|
ctx_full = dict(
|
|
n_views=n_chunk, n_pbr=n_pbr,
|
|
dino_features=dino_c,
|
|
position_voxel_indices=chunk_voxel,
|
|
)
|
|
if ref_features is not None:
|
|
ctx_full["ref_features"] = ref_features
|
|
pred_full = model.unet(unet_in, t_arr, text_full_c, **ctx_full)
|
|
mx.synchronize()
|
|
|
|
vs = view_scale[mx.array(chunk_idx)]
|
|
guided = (
|
|
pred_uncond
|
|
+ guidance_scale * vs * (pred_ref - pred_uncond)
|
|
+ guidance_scale * vs * (pred_full - pred_ref)
|
|
)
|
|
|
|
# Scatter chunk results back to correct positions in the
|
|
# full latent array. chunk_idx maps each chunk sample to
|
|
# its position in the [albedo_0..N-1, mr_0..N-1] layout.
|
|
noise_guided[mx.array(chunk_idx)] = guided
|
|
|
|
# Scheduler step
|
|
latents = model.scheduler.step(noise_guided, t_int, latents)
|
|
mx.synchronize()
|
|
|
|
if i % 5 == 0 or i == num_inference_steps - 1:
|
|
r = float(latents.max() - latents.min())
|
|
print(f" step {i}/{num_inference_steps}: t={t_int}, range={r:.1f}")
|
|
|
|
# ------------------------------------------------------------------
|
|
# 7. Decode (one image at a time to avoid OOM)
|
|
# ------------------------------------------------------------------
|
|
print(" Decoding...")
|
|
all_images = []
|
|
n_total = n_pbr * n_views
|
|
for idx in range(n_total):
|
|
# latents already at scaling_factor magnitude; vae.decode handles
|
|
# the divide internally — pass z as-is.
|
|
z_i = latents[idx : idx + 1]
|
|
dec_i = model.vae.decode(z_i)
|
|
img_i = mx.clip((dec_i + 1) / 2, 0, 1)
|
|
mx.synchronize()
|
|
all_images.append(np.array(img_i[0]))
|
|
if (idx + 1) % n_views == 0:
|
|
print(f" decoded {idx + 1}/{n_total}")
|
|
|
|
result = {
|
|
"albedo": [
|
|
(all_images[i] * 255).astype(np.uint8)
|
|
for i in range(n_views)
|
|
],
|
|
"mr": [
|
|
(all_images[n_views + i] * 255).astype(np.uint8)
|
|
for i in range(n_views)
|
|
],
|
|
}
|
|
return result
|