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
import zipfile
import io

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load models
transmitter = load_model("transmitter", device=device)
xm = load_model("text300M", device=device)
diffusion = diffusion_from_config(load_config("diffusion"))

def generate_3d(prompt):
    with torch.no_grad():
        latents = sample_latents(
            batch_size=1,
            model=xm,
            diffusion=diffusion,
            guidance_scale=3.0,
            model_kwargs=dict(texts=[prompt]),
            progress=True,
            clip_denoised=True,
            use_fp16=True,
            device=device,
        )
        mesh = decode_latent_mesh(xm, latents[0]).tri_mesh()

        # Save zip
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w") as z:
            z.writestr("mesh.ply", mesh.to_ply())
        buf.seek(0)
        return buf

demo = gr.Interface(
    fn=generate_3d,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.File(label="Generated 3D Model (.zip)")
)

demo.launch()