import torch
import gradio as gr
from shap_e.diffusion.sample import sample_latents
from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
from shap_e.models.download import load_model, load_config
from shap_e.util.notebooks import decode_latent_mesh
import trimesh
import os

# Initialize models
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

xm = load_model('transmitter', device=device)
model = load_model('text300M', device=device)
diffusion = diffusion_from_config(load_config('diffusion'))

def generate_3d(prompt, guidance_scale):
    if not prompt:
        return None
    
    # Generate latents
    latents = sample_latents(
        batch_size=1,
        model=model,
        diffusion=diffusion,
        guidance_scale=guidance_scale,
        model_kwargs=dict(texts=[prompt]),
        progress=True,
        clip_denoised=True,
        use_fp16=torch.cuda.is_available(),
        use_karras=True,
        karras_steps=64,
        sigma_min=1e-3,
        sigma_max=160,
        s_churn=0,
    )

    # Decode to mesh
    mesh = decode_latent_mesh(xm, latents[0]).triangulate()
    
    # Save as GLB
    temp_obj = "temp.obj"
    with open(temp_obj, "w") as f:
        mesh.write_obj(f)
        
    t_mesh = trimesh.load(temp_obj)
    glb_path = "output.glb"
    t_mesh.export(glb_path)
    
    return glb_path

# UI Design
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧊 Prompt2Shape: AI 3D Generator")
    gr.Markdown("Transform your text descriptions into 3D models instantly.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Text Prompt", placeholder="a cute cat, a futuristic car...", lines=3)
            guidance = gr.Slider(minimum=1, maximum=20, value=15, label="Guidance Scale")
            btn = gr.Button("Generate 3D Model", variant="primary")
            
        with gr.Column():
            output = gr.Model3D(label="3D Preview")

    btn.click(fn=generate_3d, inputs=[prompt, guidance], outputs=output)
    
    gr.Markdown("""
    ### 🚀 Architecture Highlights
    - **Semantic Layer**: CLIP-based prompt understanding.
    - **Diffusion Layer**: Latent diffusion sampling.
    - **Geometry Layer**: Mesh extraction and refinement.
    """)

if __name__ == "__main__":
    demo.launch()