app.py

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 create_pan_cameras, decode_latent_images, gif_widget
from shap_e.util.image_util import load_image

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'))

batch_size = 1
guidance_scale = 9

def get_latents(prompt):
    return sample_latents(
    batch_size=batch_size,
    model=model,
    diffusion=diffusion,
    guidance_scale=guidance_scale,
    model_kwargs=dict(texts=[prompt] * batch_size),
    progress=True,
    clip_denoised=False,
    use_fp16=True,
    use_karras=True,
    karras_steps=64,
    sigma_min=1e-3,
    sigma_max=160,
    s_churn=0,
    )

render_mode = 'nerf' 
size = 128 

cameras = create_pan_cameras(size, device)

def get_gif(prompt):
    for i, latent in enumerate(get_latents(prompt)):
      images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)
      return gif_widget(images)
        
iface = gr.Interface(fn = get_gif, inputs = "text", outputs=["html" , "text"] ,title = 'LatentVerse')
iface.queue().launch(inline = False)