#!/usr/bin/env python3
"""
Generate 3D mesh from a 2D image using Hunyuan3D-2 via HuggingFace Spaces.
Outputs GLB (with color) and STL (geometry only).

Usage:
    python3 generate_3d_from_image.py /path/to/image.jpg [caption]
"""
import sys, os

def generate(image_path, caption=None, out_dir="/tmp"):
    from gradio_client import Client, handle_file
    import trimesh

    c = Client("tencent/Hunyuan3D-2")
    img = handle_file(image_path)

    job = c.submit(
        image=img,
        caption=caption,
        mv_image_front=None, mv_image_back=None,
        mv_image_left=None, mv_image_right=None,
        steps=30,
        guidance_scale=5.0,
        seed=1234,
        octree_resolution=256,
        check_box_rembg=True,
        num_chunks=8000,
        randomize_seed=True,
        api_name="/shape_generation"
    )
    print("Job submitted, waiting...")
    result = job.result(timeout=300)

    glb_path = result[0]["value"]
    base = os.path.splitext(os.path.basename(image_path))[0]
    out_glb = os.path.join(out_dir, f"{base}.glb")
    out_stl = os.path.join(out_dir, f"{base}.stl")

    # Copy GLB
    import shutil
    shutil.copy(glb_path, out_glb)

    # Convert to STL
    mesh = trimesh.load(glb_path, force='mesh')
    mesh.export(out_stl)

    print(f"GLB: {out_glb}  ({os.path.getsize(out_glb)/1024/1024:.1f} MB)")
    print(f"STL: {out_stl}  ({os.path.getsize(out_stl)/1024/1024:.1f} MB)")
    print(f"Faces: {len(mesh.faces)}  Vertices: {len(mesh.vertices)}")
    print(f"Watertight: {mesh.is_watertight}")
    return out_glb, out_stl

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 generate_3d_from_image.py <image.jpg> [caption]")
        sys.exit(1)
    generate(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None)