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
114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
"""End-to-end paint inference on the real mermaid mesh + astronaut reference.
|
|
|
|
Validates that the full MLX pipeline produces coherent textures after the
|
|
attention-head-dim fix (commit message for context). Outputs 6-view albedo
|
|
and MR images to /tmp/full/e2e_fixed/.
|
|
|
|
Usage:
|
|
cd /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint
|
|
python3 tests/test_e2e_paint.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
MESH_PATH = os.path.join(
|
|
os.path.dirname(__file__), "..", "assets", "case_2", "mesh.glb"
|
|
)
|
|
REF_IMAGE_PATH = "/tmp/full/ComfyUI_temp_cpzbr_00007_.png"
|
|
# HF repo ID (auto-downloaded via huggingface_hub) or local path override
|
|
# via HUNYUAN3D_MLX_WEIGHTS_DIR env var.
|
|
WEIGHTS_SOURCE = "dgrauet/hunyuan3d-2.1-mlx"
|
|
OUT_DIR = "/tmp/full/e2e_fixed"
|
|
|
|
VIEW_SIZE = 512
|
|
N_VIEWS = 6
|
|
NUM_STEPS = 15
|
|
CFG = 3.0
|
|
AZIMS = [0, 60, 120, 180, 240, 300]
|
|
|
|
|
|
def main() -> None:
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
|
|
# --- Render conditions ---
|
|
print(f"[1/4] Rendering {N_VIEWS} views at {VIEW_SIZE}px from {MESH_PATH}")
|
|
import trimesh
|
|
from DifferentiableRenderer.mesh_render_mlx import MeshRenderMLX
|
|
|
|
mesh = trimesh.load(MESH_PATH, force="mesh")
|
|
renderer = MeshRenderMLX(default_resolution=VIEW_SIZE, shader_type="face")
|
|
renderer.load_mesh(mesh)
|
|
|
|
normal_maps = []
|
|
position_maps = []
|
|
for i, az in enumerate(AZIMS):
|
|
n = renderer.render_normal(0, az, bg_color=(0.5, 0.5, 0.5))
|
|
p = renderer.render_position(0, az, bg_color=(0.5, 0.5, 0.5))
|
|
normal_maps.append(np.asarray(n, dtype=np.float32))
|
|
position_maps.append(np.asarray(p, dtype=np.float32))
|
|
Image.fromarray((n * 255).astype(np.uint8)).save(
|
|
os.path.join(OUT_DIR, f"cond_normal_{i}_az{az}.png")
|
|
)
|
|
Image.fromarray((p * 255).astype(np.uint8)).save(
|
|
os.path.join(OUT_DIR, f"cond_position_{i}_az{az}.png")
|
|
)
|
|
print(f" rendered {len(normal_maps)} normal + position maps")
|
|
|
|
# --- Reference image ---
|
|
ref = np.array(Image.open(REF_IMAGE_PATH).convert("RGB"))
|
|
print(f"[2/4] Reference image: {ref.shape}, dtype={ref.dtype}")
|
|
|
|
# --- Load model ---
|
|
print("[3/4] Loading HunyuanPaintPBR MLX weights...")
|
|
from hunyuanpaintpbr_mlx.load_model import HunyuanPaintModelMLX
|
|
|
|
t0 = time.time()
|
|
model = HunyuanPaintModelMLX.from_pretrained(WEIGHTS_SOURCE)
|
|
print(f" loaded in {time.time() - t0:.1f}s")
|
|
|
|
# --- Generate ---
|
|
print(f"[4/4] Running denoising ({NUM_STEPS} steps, CFG={CFG})...")
|
|
from hunyuanpaintpbr_mlx.inference import generate_multiview_pbr
|
|
|
|
t0 = time.time()
|
|
result = generate_multiview_pbr(
|
|
model=model,
|
|
normal_maps=normal_maps,
|
|
position_maps=position_maps,
|
|
reference_image=ref,
|
|
camera_azims=AZIMS,
|
|
num_inference_steps=NUM_STEPS,
|
|
guidance_scale=CFG,
|
|
view_size=VIEW_SIZE,
|
|
seed=42,
|
|
)
|
|
print(f" denoise+decode in {time.time() - t0:.1f}s")
|
|
|
|
# --- Save ---
|
|
for i, (al, mr) in enumerate(zip(result["albedo"], result["mr"])):
|
|
Image.fromarray(al).save(
|
|
os.path.join(OUT_DIR, f"albedo_{i}_az{AZIMS[i]}.png")
|
|
)
|
|
Image.fromarray(mr).save(
|
|
os.path.join(OUT_DIR, f"mr_{i}_az{AZIMS[i]}.png")
|
|
)
|
|
print(f"\nOutputs written to {OUT_DIR}")
|
|
|
|
# Report rough stats on the first albedo view to catch cyan-blob regressions
|
|
a0 = result["albedo"][0].astype(np.float32) / 255.0
|
|
mean_rgb = a0.reshape(-1, 3).mean(axis=0)
|
|
std_rgb = a0.reshape(-1, 3).std(axis=0)
|
|
print(f" albedo[0] mean RGB: {mean_rgb.round(3).tolist()} "
|
|
f"std: {std_rgb.round(3).tolist()}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|