import os, sys, time
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
import torch
from PIL import Image
from hy3dgen.rembg import BackgroundRemover
from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline

OCT   = int(sys.argv[1]) if len(sys.argv) > 1 else 256
STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 50
G     = float(sys.argv[3]) if len(sys.argv) > 3 else 5.0
TAG   = sys.argv[4] if len(sys.argv) > 4 else "mv"

P = "/Users/danicosta/Desktop/Atapuerca/3D/prep/multiview"
views = {
    "front": f"{P}/00_front.jpg",
    "left":  f"{P}/270_left_prof.jpg",
    "back":  f"{P}/180_back.jpg",
    "right": f"{P}/075_right_prof.jpg",
}
print(f"[mv] octree={OCT} steps={STEPS} guidance={G} vistas={list(views)}")
rb = BackgroundRemover()
imgs = {}
for k, p in views.items():
    imgs[k] = rb(Image.open(p).convert("RGB"))

pipe = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
    'tencent/Hunyuan3D-2mv', subfolder='hunyuan3d-dit-v2-mv', variant='fp16',
    device='mps', dtype=torch.float32)

t0 = time.time()
mesh = pipe(image=imgs, num_inference_steps=STEPS, octree_resolution=OCT,
            num_chunks=20000, guidance_scale=G,
            generator=torch.Generator("cpu").manual_seed(123), output_type='trimesh')[0]
print(f"[mv] {time.time()-t0:.0f}s verts={len(mesh.vertices)} faces={len(mesh.faces)}")
out = f"/Users/danicosta/Desktop/Atapuerca/3D/export/neandertal_head_{TAG}.glb"
mesh.export(out)
mesh.export(out.replace('.glb', '.obj'))
print("[mv] EXPORTADO", out)