import gradio as gr
import torch
import os
from hunyuan3d.pipeline import Hunyuan3D_Pipeline

# --- CONFIGURATION ---
# Path to your specific mini model
VAE_PATH = "hunyuan3d-vae-v2-mini-turbo/model.fp16.safetensor"
# Path to the main model (You still need the main weights!)
# If you don't have them, it will try to download the standard ones
MAIN_MODEL_ID = "Tencent/Hunyuan3D-1"

print(f"Initializing Hunyuan3D on CPU...")
print(f"Target VAE: {VAE_PATH}")

# 1. Check Hardware
# Free Tier = CPU. We MUST use float32 for CPU, even if the model is fp16.
device = "cpu"
dtype = torch.float32 

# 2. Load Pipeline
try:
    # We load the standard pipeline first
    pipe = Hunyuan3D_Pipeline.from_pretrained(
        MAIN_MODEL_ID,
        device=device,
        use_safetensors=True
    )
    
    # 3. SWAP THE VAE (The "Mini Turbo" part)
    if os.path.exists(VAE_PATH):
        print("Found Mini-Turbo VAE! Swapping weights...")
        # Load the custom safetensor
        from safetensors.torch import load_file
        vae_state_dict = load_file(VAE_PATH)
        
        # Convert to float32 for CPU compatibility
        vae_state_dict = {k: v.to(dtype) for k, v in vae_state_dict.items()}
        
        # Load into the pipeline's VAE
        pipe.vae.load_state_dict(vae_state_dict, strict=False)
        print("Swap complete. Running with Mini-Turbo VAE.")
    else:
        print(f"Warning: Could not find {VAE_PATH}. Using standard VAE.")

    # Convert entire pipeline to Float32 for CPU safety
    pipe.to(device, dtype)

except Exception as e:
    print(f"Error loading model: {e}")
    pipe = None

def generate_3d(prompt, steps):
    if pipe is None:
        return None
    
    print(f"Generating: {prompt}")
    
    # Run Inference
    # On CPU, this will be slow (2-5 mins)
    output = pipe(
        prompt=prompt, 
        num_inference_steps=steps, 
        guidance_scale=7.5
    )
    
    # Save the mesh
    output_path = "output.obj"
    output.save(output_path)
    
    return output_path

# --- GRADIO UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 🧊 Hunyuan3D (Mini VAE Config)")
    gr.Markdown(f"Running on **{device.upper()}**. Using VAE: `{VAE_PATH if os.path.exists(VAE_PATH) else 'Default'}`")
    
    with gr.Row():
        with gr.Column():
            txt_input = gr.Textbox(label="Prompt", value="A cute isometric house")
            steps = gr.Slider(15, 50, value=20, label="Steps (Keep low for CPU)")
            btn = gr.Button("Generate 3D")
        
        with gr.Column():
            model_out = gr.Model3D(label="3D Result", clear_color=[0,0,0,0])

    btn.click(generate_3d, inputs=[txt_input, steps], outputs=model_out)

demo.queue().launch()