"""
Gradio demo per Shap-E (text-to-3D) – Hugging Face Spaces
Autore: tu
"""

import os
import gradio as gr
import torch
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

# ---------- caricamento modelli ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
xm = load_model('transmitter', device=device)
model = load_model('text300M', device=device)
diffusion = diffusion_from_config(load_config('diffusion'))

# ---------- logica ----------
def generate(prompt: str,
             guidance: float = 15.0,
             steps: int = 64):
    """Genera la mesh e restituisce il file .ply scaricabile."""
    latents = sample_latents(
        batch_size=1,
        model=model,
        diffusion=diffusion,
        guidance_scale=guidance,
        model_kwargs=dict(texts=[prompt]),
        progress=True,
        clip_denoised=True,
        use_fp16=True,
        use_karras=True,
        karras_steps=steps,
        sigma_min=1e-3,
        sigma_max=160,
        s_churn=0,
    )

    t = decode_latent_mesh(xm, latents[0]).tri_mesh()
    out_path = "output.ply"
    with open(out_path, "wb") as f:
        t.write_ply(f)
    return out_path

# ---------- interfaccia Gradio ----------
iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(1, 30, value=15, label="Guidance scale"),
        gr.Slider(32, 128, value=64, step=16, label="Karras steps")
    ],
    outputs=gr.File(label="Scarica mesh .ply"),
    title="Shap-E Text-to-3D",
    description="Genera una mesh 3D da una descrizione testuale con Shap-E.",
    examples=[["a high–quality red sports car"],
              ["a cute low-poly cat"]],
    cache_examples=False   # vogliamo sempre generare
)

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