Spaces:
Runtime error
Runtime error
| 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 | |
| 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''' | |
| <div style="text-align: center; max-width: 650px; margin: 0 auto;"> | |
| <div> | |
| <img class="logo" src="file/images/mirage.png" alt="Mirage Logo" | |
| style="margin: auto; max-width: 7rem;"> | |
| <br /> | |
| <h1 style="font-weight: 900; font-size: 2.5rem;"> | |
| Shap-E Web UI | |
| </h1> | |
| </div> | |
| <h3 style="font-weight: 900; font-size: 1.5rem;"> | |
| If the Queue is Too Long, <a href="https://app.mirageml.com/generate" style="text-decoration: underline" target="_blank">Try it on Mirage</a>! | |
| </h3> | |
| <br /> | |
| <p style="margin-bottom: 10px; font-size: 94%"> | |
| Generate 3D Assets in 1 minute with a prompt or image! | |
| Based on the <a href="https://github.com/openai/shap-e">Shap-E</a> implementation | |
| </p> | |
| <br /> | |
| <p>There's only one step left before you can train your model: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a <b>T4 GPU</b> to it (via the Settings tab)</a> 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.</p> | |
| </div> | |
| ''') | |
| else: | |
| top_description = gr.HTML(f''' | |
| <div style="text-align: center; max-width: 650px; margin: 0 auto;"> | |
| <div> | |
| <img class="logo" src="file/images/mirage.png" alt="Mirage Logo" | |
| style="margin: auto; max-width: 7rem;"> | |
| <br /> | |
| <h1 style="font-weight: 900; font-size: 2.5rem;"> | |
| Shap-E Web UI | |
| </h1> | |
| </div> | |
| <h3 style="font-weight: 900; font-size: 1.5rem;"> | |
| If the Queue is Too Long, <a href="https://app.mirageml.com/generate" style="text-decoration: underline" target="_blank">Try it on Mirage</a>! | |
| </h3> | |
| <br /> | |
| <p style="margin-bottom: 10px; font-size: 94%"> | |
| Generate 3D Assets in 1 minute with a prompt or image! | |
| Based on the <a href="https://github.com/openai/shap-e">Shap-E</a> implementation | |
| </p> | |
| </div> | |
| ''') | |
| 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) | |