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 argparse | |
| from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config | |
| from point_e.diffusion.sampler import PointCloudSampler | |
| from point_e.models.download import load_checkpoint | |
| from point_e.models.configs import MODEL_CONFIGS, model_from_config | |
| from point_e.util.plotting import plot_point_cloud | |
| from point_e.util.pc_to_mesh import marching_cubes_mesh | |
| from diffusers import StableDiffusionPipeline | |
| 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} | |
| .duplicate-button img{margin: 0} | |
| ''' | |
| def set_state(s): | |
| print(s) | |
| global state | |
| state = s | |
| def get_state(): | |
| return state | |
| def load_img2mesh_model(model_name): | |
| set_state(f'Creating img2mesh model {model_name}...') | |
| i2m_name = model_name | |
| i2m_model = model_from_config(MODEL_CONFIGS[i2m_name], device) | |
| i2m_model.eval() | |
| base_diffusion_i2m = diffusion_from_config(DIFFUSION_CONFIGS[i2m_name]) | |
| set_state(f'Downloading img2mesh checkpoint {model_name}...') | |
| i2m_model.load_state_dict(load_checkpoint(i2m_name, device)) | |
| return i2m_model, base_diffusion_i2m | |
| def get_sampler(model_name, txt2obj, guidance_scale): | |
| if txt2obj: | |
| set_state('Creating txt2mesh model...') | |
| t2m_name = 'base40M-textvec' | |
| t2m_model = model_from_config(MODEL_CONFIGS[t2m_name], device) | |
| t2m_model.eval() | |
| base_diffusion_t2m = diffusion_from_config(DIFFUSION_CONFIGS[t2m_name]) | |
| set_state('Downloading txt2mesh checkpoint...') | |
| t2m_model.load_state_dict(load_checkpoint(t2m_name, device)) | |
| else: | |
| i2m_model, base_diffusion_i2m = load_img2mesh_model(model_name) | |
| set_state('Creating upsample model...') | |
| upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device) | |
| upsampler_model.eval() | |
| upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample']) | |
| set_state('Downloading upsampler checkpoint...') | |
| upsampler_model.load_state_dict(load_checkpoint('upsample', device)) | |
| return PointCloudSampler( | |
| device=device, | |
| models=[t2m_model if txt2obj else i2m_model, upsampler_model], | |
| diffusions=[base_diffusion_t2m if txt2obj else base_diffusion_i2m, upsampler_diffusion], | |
| num_points=[1024, 4096 - 1024], | |
| aux_channels=['R', 'G', 'B'], | |
| guidance_scale=[guidance_scale, 0.0 if txt2obj else guidance_scale], | |
| model_kwargs_key_filter=('texts', '') if txt2obj else ("*",) | |
| ) | |
| def generate_txt2img(prompt): | |
| pipe = StableDiffusionPipeline.from_pretrained("point_e_model_cache/stable-diffusion-2-1", torch_dtype=torch.float16) | |
| pipe = pipe.to("cuda") | |
| image = pipe(prompt).images[0] | |
| return image | |
| def generate_3D(input, model_name='base1B', guidance_scale=3.0, grid_size=128): | |
| set_state('Entered generate function...') | |
| # try: | |
| # input = Image.fromarray(input) | |
| # except: | |
| # img = generate_txt2img(input) | |
| # img.save('/tmp/img.png') | |
| # input = Image.open('/tmp/img.png') | |
| if isinstance(input, Image.Image): | |
| input = prepare_img(input) | |
| # if input is a string, it's a text prompt | |
| sampler = get_sampler(model_name, txt2obj=True if isinstance(input, str) else False, guidance_scale=guidance_scale) | |
| # Produce a sample from the model. | |
| set_state('Sampling...') | |
| samples = None | |
| kw_args = dict(texts=[input]) if isinstance(input, str) else dict(images=[input]) | |
| for x in sampler.sample_batch_progressive(batch_size=1, model_kwargs=kw_args): | |
| samples = x | |
| set_state('Converting to point cloud...') | |
| pc = sampler.output_to_point_clouds(samples)[0] | |
| set_state('Converting to mesh...') | |
| save_ply(pc, '/tmp/mesh.ply', grid_size) | |
| set_state('') | |
| return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), create_gif(pc), 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) | |
| def create_gif(pc): | |
| fig = plt.figure(facecolor='black', figsize=(4, 4)) | |
| ax = fig.add_subplot(111, projection='3d', facecolor='black') | |
| fixed_bounds=((-0.75, -0.75, -0.75),(0.75, 0.75, 0.75)) | |
| # Create an empty list to store the frames | |
| frames = [] | |
| # Create a loop to generate the frames for the GIF | |
| for angle in range(0, 360, 4): | |
| # Clear the plot and plot the point cloud | |
| ax.clear() | |
| color_args = np.stack( | |
| [pc.channels["R"], pc.channels["G"], pc.channels["B"]], axis=-1 | |
| ) | |
| c = pc.coords | |
| ax.scatter(c[:, 0], c[:, 1], c[:, 2], c=color_args) | |
| # Set the viewpoint for the plot | |
| ax.view_init(elev=10, azim=angle) | |
| # Turn off the axis labels and ticks | |
| ax.axis('off') | |
| ax.set_xlim3d(fixed_bounds[0][0], fixed_bounds[1][0]) | |
| ax.set_ylim3d(fixed_bounds[0][1], fixed_bounds[1][1]) | |
| ax.set_zlim3d(fixed_bounds[0][2], fixed_bounds[1][2]) | |
| # Draw the figure to update the image data | |
| fig.canvas.draw() | |
| # Save the plot as a frame for the GIF | |
| frame = np.array(fig.canvas.renderer.buffer_rgba()) | |
| w, h = frame.shape[0], frame.shape[1] | |
| i = int(round((h - int(h*0.6)) / 2.)) | |
| frame = frame[i:i + int(h*0.6),i:i + int(h*0.6)] | |
| frames.append(frame) | |
| # Save the GIF using imageio | |
| imageio.mimsave('/tmp/pointcloud.mp4', frames, fps=30) | |
| return '/tmp/pointcloud.mp4' | |
| 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;"> | |
| Point-E Web UI | |
| </h1> | |
| <br /> | |
| <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/point-e?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a> | |
| </div> | |
| <br /> | |
| <p style="margin-bottom: 10px; font-size: 94%"> | |
| Generate 3D Assets in 2 minutes with a prompt or image! | |
| Based on the <a href="https://github.com/openai/point-e">Point-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;"> | |
| Point-E Web UI | |
| </h1> | |
| <br /> | |
| <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/point-e?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a> | |
| </div> | |
| <br /> | |
| <p style="margin-bottom: 10px; font-size: 94%"> | |
| Generate 3D Assets in 2 minutes with a prompt or image! | |
| Based on the <a href="https://github.com/openai/point-e">Point-E</a> implementation | |
| </p> | |
| </div> | |
| ''') | |
| with gr.Row(): | |
| with gr.Column(): | |
| 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.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.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.Video(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=[ | |
| ["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, model, scale], outputs=[model_3d, model_gif, file_out]) | |
| text_button.click(fn=generate_3D, inputs=[prompt, model, scale], outputs=[model_3d, model_gif, file_out]) | |
| block.launch(show_api=False) | |