import os import gradio as gr from PIL import Image import torch import matplotlib.pyplot as plt import imageio import numpy as np import math import argparse import tempfile import torch import base64 import io import os from typing import Union 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, decode_latent_mesh from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera from shap_e.models.transmitter.base import Transmitter, VectorDecoder from shap_e.util.collections import AttrDict import trimesh state = "" device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') css = ''' .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important} .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important} #component-4, #component-3, #component-10{min-height: 0} ''' def set_state(s): print(s) global state state = s def get_state(): return state def to_video(frames: list[Image.Image], fps: int = 5) -> str: out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps) for frame in frames: writer.append_data(np.asarray(frame)) writer.close() return out_file.name def generate_3D(input, grid_size=64): set_state('Entered generate function...') # if input is a string, it's a text prompt xm = load_model('transmitter', device=device) diffusion = diffusion_from_config(load_config('diffusion')) batch_size = 4 if isinstance(input, np.ndarray): input = Image.fromarray(input) if isinstance(input, Image.Image): input = prepare_img(input) model = load_model('image300M', device=device) guidance_scale = 3.0 model_kwargs = dict(images=[input] * batch_size) else: model = load_model('text300M', device=device) guidance_scale = 15.0 model_kwargs = dict(texts=[input] * batch_size) print(input) latents = sample_latents( batch_size=batch_size, model=model, diffusion=diffusion, guidance_scale=guidance_scale, model_kwargs=model_kwargs, progress=True, clip_denoised=True, use_fp16=True, use_karras=True, karras_steps=64, sigma_min=1e-3, sigma_max=160, s_churn=0, ) render_mode = 'stf' # you can change this to 'stf' size = grid_size # this is the size of the renders; higher values take longer to render. cameras = create_pan_cameras(size, device) with open(f'/tmp/mesh.ply', 'wb') as f: decode_latent_mesh(xm, latents[0]).tri_mesh().write_ply(f) set_state('Converting to point cloud...') # pc = sampler.output_to_point_clouds(samples)[0] set_state('Converting to mesh...') # save_ply(pc, 'output/mesh.ply', grid_size) set_state('') images = decode_latent_images(xm, latents[0], cameras, rendering_mode=render_mode) return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), to_video(images), gr.update(value=['/tmp/mesh.glb', '/tmp/mesh.ply'], visible=True) def prepare_img(img): w, h = img.size if w > h: img = img.crop((w - h) / 2, 0, w - (w - h) / 2, h) else: img = img.crop((0, (h - w) / 2, w, h - (h - w) / 2)) # resize to 256x256 img = img.resize((256, 256)) return img def ply_to_glb(ply_file, glb_file): mesh = trimesh.load(ply_file) # Save the mesh as a glb file using Trimesh mesh.export(glb_file, file_type='glb') return glb_file # def save_ply(pc, file_name, grid_size): # set_state('Creating SDF model...') # sdf_name = 'sdf' # sdf_model = model_from_config(MODEL_CONFIGS[sdf_name], device) # sdf_model.eval() # set_state('Loading SDF model...') # sdf_model.load_state_dict(load_checkpoint(sdf_name, device)) # # Produce a mesh (with vertex colors) # mesh = marching_cubes_mesh( # pc=pc, # model=sdf_model, # batch_size=4096, # grid_size=grid_size, # increase to 128 for resolution used in evals # progress=True, # ) # # Write the mesh to a PLY file to import into some other program. # with open(file_name, 'wb') as f: # mesh.write_ply(f) block = gr.Blocks().queue(max_size=250, concurrency_count=6) with block: with gr.Box(): if(not torch.cuda.is_available()): top_description = gr.HTML(f'''

Shap-E Web UI

If the Queue is Too Long, Try it on Mirage!


Generate 3D Assets in 1 minute with a prompt or image! Based on the Shap-E implementation


There's only one step left before you can train your model: attribute a T4 GPU to it (via the Settings tab) and run the training below. Other GPUs are not compatible for now. You will be billed by the minute from when you activate the GPU until when it is turned it off.

''') else: top_description = gr.HTML(f'''

Shap-E Web UI

If the Queue is Too Long, Try it on Mirage!


Generate 3D Assets in 1 minute with a prompt or image! Based on the Shap-E implementation

''') with gr.Row(): with gr.Column(): with gr.Tab("Text to 3D"): gr.Markdown("Uses Stable Diffusion to create an image from the prompt.") prompt = gr.Textbox(label="Prompt", placeholder="A HD photo of a Corgi") text_button = gr.Button(label="Generate") with gr.Tab("Image to 3D"): gr.Markdown("Best results with images of objects on an empty background.") input_image = gr.Image(label="Image") img_button = gr.Button(label="Generate") # with gr.Accordion("Advanced options", open=False): # model = gr.Radio(["base40M", "base300M", "base1B"], label="Model", value="base1B") # scale = gr.Slider( # label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.1 # ) with gr.Column(): model_gif = gr.Model3D(label="3D Model GIF") # btn_pc_to_obj = gr.Button(value="Convert to OBJ", visible=False) model_3d = gr.Model3D(value=None) file_out = gr.File(label="Files", visible=False) if torch.cuda.is_available(): gr.Examples( examples=[ ["a shark"], ["an avocado"], ], inputs=[prompt], outputs=[model_3d, model_gif, file_out], fn=generate_3D, cache_examples=True ) gr.Examples( examples=[ ["images/pumpkin.png"], ["images/fantasy_world.png"], ], inputs=[input_image], outputs=[model_3d, model_gif, file_out], fn=generate_3D, cache_examples=True ) img_button.click(fn=generate_3D, inputs=[input_image], outputs=[model_3d, model_gif, file_out]) text_button.click(fn=generate_3D, inputs=[prompt], outputs=[model_3d, model_gif, file_out]) block.launch(show_api=False)