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

def load(path):
    im = Image.open(path)
    if im.mode != "RGBA":
        im = BackgroundRemover()(im.convert("RGB"))
    return im

pipe = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained("tencent/Hunyuan3D-2", device="mps", dtype=torch.float32)

def boxness(mesh):
    v = np.asarray(mesh.vertices); lo=v.min(0); hi=v.max(0); eps=(hi-lo).max()*0.01
    return (((np.abs(v-lo)<eps)|(np.abs(v-hi)<eps)).any(1)).mean()

IMG = "/Users/danicosta/Desktop/Atapuerca/3D/prep/single/head_crop.png"
img = load(IMG)

# (guidance, box_v, steps, octree)
configs = [
    (5.0, 1.01, 20, 96),
    (3.0, 1.01, 20, 96),
    (1.0, 1.01, 20, 96),
    (0.0, 1.01, 20, 96),
    (5.0, 1.50, 20, 96),
]
for g, bv, steps, oct in configs:
    try:
        gen = torch.Generator(device="cpu").manual_seed(42)
        m = pipe(image=img, num_inference_steps=steps, octree_resolution=oct,
                 guidance_scale=g, box_v=bv, generator=gen)[0]
        b = boxness(m)
        print(f"RESULT guidance={g} box_v={bv}: verts={len(m.vertices)} boxness={b:.3f} => {'CUBO' if b>0.5 else 'CABEZA OK'}", flush=True)
    except Exception as e:
        print(f"RESULT guidance={g} box_v={bv}: ERROR {e}", flush=True)