MirageML/point-e
37
1import os2import gradio as gr3from PIL import Image4import torch5import matplotlib.pyplot as plt6import imageio7import numpy as np8import argparse9 10from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config11from point_e.diffusion.sampler import PointCloudSampler12from point_e.models.download import load_checkpoint13from point_e.models.configs import MODEL_CONFIGS, model_from_config14from point_e.util.plotting import plot_point_cloud15from point_e.util.pc_to_mesh import marching_cubes_mesh16 17from diffusers import StableDiffusionPipeline18 19import trimesh20 21 22state = ""23device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')24 25css = '''26 .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}27 .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}28 #component-4, #component-3, #component-10{min-height: 0}29 .duplicate-button img{margin: 0}30'''31 32def set_state(s):33 print(s)34 global state35 state = s36 37def get_state():38 return state39 40def load_img2mesh_model(model_name):41 set_state(f'Creating img2mesh model {model_name}...')42 i2m_name = model_name43 i2m_model = model_from_config(MODEL_CONFIGS[i2m_name], device)44 i2m_model.eval()45 base_diffusion_i2m = diffusion_from_config(DIFFUSION_CONFIGS[i2m_name])46 47 set_state(f'Downloading img2mesh checkpoint {model_name}...')48 i2m_model.load_state_dict(load_checkpoint(i2m_name, device))49 50 return i2m_model, base_diffusion_i2m51 52 53 54def get_sampler(model_name, txt2obj, guidance_scale):55 if txt2obj:56 set_state('Creating txt2mesh model...')57 t2m_name = 'base40M-textvec'58 t2m_model = model_from_config(MODEL_CONFIGS[t2m_name], device)59 t2m_model.eval()60 base_diffusion_t2m = diffusion_from_config(DIFFUSION_CONFIGS[t2m_name])61 62 set_state('Downloading txt2mesh checkpoint...')63 t2m_model.load_state_dict(load_checkpoint(t2m_name, device))64 else:65 i2m_model, base_diffusion_i2m = load_img2mesh_model(model_name)66 67 set_state('Creating upsample model...')68 upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)69 upsampler_model.eval()70 upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])71 72 set_state('Downloading upsampler checkpoint...')73 upsampler_model.load_state_dict(load_checkpoint('upsample', device))74 75 return PointCloudSampler(76 device=device,77 models=[t2m_model if txt2obj else i2m_model, upsampler_model],78 diffusions=[base_diffusion_t2m if txt2obj else base_diffusion_i2m, upsampler_diffusion],79 num_points=[1024, 4096 - 1024],80 aux_channels=['R', 'G', 'B'],81 guidance_scale=[guidance_scale, 0.0 if txt2obj else guidance_scale],82 model_kwargs_key_filter=('texts', '') if txt2obj else ("*",)83 )84 85def generate_txt2img(prompt):86 pipe = StableDiffusionPipeline.from_pretrained("point_e_model_cache/stable-diffusion-2-1", torch_dtype=torch.float16)87 pipe = pipe.to("cuda")88 image = pipe(prompt).images[0]89 90 return image91 92def generate_3D(input, model_name='base1B', guidance_scale=3.0, grid_size=128):93 set_state('Entered generate function...')94 95 # try:96 # input = Image.fromarray(input)97 # except:98 # img = generate_txt2img(input)99 # img.save('/tmp/img.png')100 # input = Image.open('/tmp/img.png')101 102 if isinstance(input, Image.Image):103 input = prepare_img(input)104 105 # if input is a string, it's a text prompt106 sampler = get_sampler(model_name, txt2obj=True if isinstance(input, str) else False, guidance_scale=guidance_scale)107 108 # Produce a sample from the model.109 set_state('Sampling...')110 samples = None111 kw_args = dict(texts=[input]) if isinstance(input, str) else dict(images=[input])112 for x in sampler.sample_batch_progressive(batch_size=1, model_kwargs=kw_args):113 samples = x114 115 set_state('Converting to point cloud...')116 pc = sampler.output_to_point_clouds(samples)[0]117 118 set_state('Converting to mesh...')119 save_ply(pc, '/tmp/mesh.ply', grid_size)120 121 set_state('')122 123 return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), create_gif(pc), gr.update(value=['/tmp/mesh.glb', '/tmp/mesh.ply'], visible=True)124 125def prepare_img(img):126 127 w, h = img.size128 if w > h:129 img = img.crop((w - h) / 2, 0, w - (w - h) / 2, h)130 else:131 img = img.crop((0, (h - w) / 2, w, h - (h - w) / 2))132 133 # resize to 256x256134 img = img.resize((256, 256))135 136 return img137 138 139def ply_to_glb(ply_file, glb_file):140 mesh = trimesh.load(ply_file)141 142 # Save the mesh as a glb file using Trimesh143 mesh.export(glb_file, file_type='glb')144 145 return glb_file146 147def save_ply(pc, file_name, grid_size):148 set_state('Creating SDF model...')149 sdf_name = 'sdf'150 sdf_model = model_from_config(MODEL_CONFIGS[sdf_name], device)151 sdf_model.eval()152 153 set_state('Loading SDF model...')154 sdf_model.load_state_dict(load_checkpoint(sdf_name, device))155 156 # Produce a mesh (with vertex colors)157 mesh = marching_cubes_mesh(158 pc=pc,159 model=sdf_model,160 batch_size=4096,161 grid_size=grid_size, # increase to 128 for resolution used in evals162 progress=True,163 )164 165 # Write the mesh to a PLY file to import into some other program.166 with open(file_name, 'wb') as f:167 mesh.write_ply(f)168 169def create_gif(pc):170 fig = plt.figure(facecolor='black', figsize=(4, 4))171 ax = fig.add_subplot(111, projection='3d', facecolor='black')172 fixed_bounds=((-0.75, -0.75, -0.75),(0.75, 0.75, 0.75))173 174 # Create an empty list to store the frames175 frames = []176 177 # Create a loop to generate the frames for the GIF178 for angle in range(0, 360, 4):179 # Clear the plot and plot the point cloud180 ax.clear()181 color_args = np.stack(182 [pc.channels["R"], pc.channels["G"], pc.channels["B"]], axis=-1183 )184 c = pc.coords185 186 187 ax.scatter(c[:, 0], c[:, 1], c[:, 2], c=color_args)188 189 # Set the viewpoint for the plot190 ax.view_init(elev=10, azim=angle)191 192 # Turn off the axis labels and ticks193 ax.axis('off')194 ax.set_xlim3d(fixed_bounds[0][0], fixed_bounds[1][0])195 ax.set_ylim3d(fixed_bounds[0][1], fixed_bounds[1][1])196 ax.set_zlim3d(fixed_bounds[0][2], fixed_bounds[1][2])197 198 # Draw the figure to update the image data199 fig.canvas.draw()200 201 # Save the plot as a frame for the GIF202 frame = np.array(fig.canvas.renderer.buffer_rgba())203 w, h = frame.shape[0], frame.shape[1]204 i = int(round((h - int(h*0.6)) / 2.))205 frame = frame[i:i + int(h*0.6),i:i + int(h*0.6)]206 frames.append(frame)207 208 # Save the GIF using imageio209 imageio.mimsave('/tmp/pointcloud.mp4', frames, fps=30)210 return '/tmp/pointcloud.mp4'211 212block = gr.Blocks().queue(max_size=250, concurrency_count=6)213with block:214 with gr.Box():215 if(not torch.cuda.is_available()):216 top_description = gr.HTML(f'''217 <div style="text-align: center; max-width: 650px; margin: 0 auto;">218 <div>219 <img class="logo" src="file/images/mirage.png" alt="Mirage Logo"220 style="margin: auto; max-width: 7rem;">221 <br />222 <h1 style="font-weight: 900; font-size: 2.5rem;">223 Point-E Web UI224 </h1>225 <br />226 <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>227 </div>228 <br />229 <p style="margin-bottom: 10px; font-size: 94%">230 Generate 3D Assets in 2 minutes with a prompt or image!231 Based on the <a href="https://github.com/openai/point-e">Point-E</a> implementation232 </p>233 <br />234 <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>235 </div>236 ''')237 else:238 top_description = gr.HTML(f'''239 <div style="text-align: center; max-width: 650px; margin: 0 auto;">240 <div>241 <img class="logo" src="file/images/mirage.png" alt="Mirage Logo"242 style="margin: auto; max-width: 7rem;">243 <br />244 <h1 style="font-weight: 900; font-size: 2.5rem;">245 Point-E Web UI246 </h1>247 <br />248 <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>249 </div>250 <br />251 <p style="margin-bottom: 10px; font-size: 94%">252 Generate 3D Assets in 2 minutes with a prompt or image!253 Based on the <a href="https://github.com/openai/point-e">Point-E</a> implementation254 </p>255 </div>256 ''')257 with gr.Row():258 with gr.Column():259 with gr.Tab("Image to 3D"):260 gr.Markdown("Best results with images of objects on an empty background.")261 input_image = gr.Image(label="Image")262 img_button = gr.Button(label="Generate")263 264 with gr.Tab("Text to 3D"):265 gr.Markdown("Uses Stable Diffusion to create an image from the prompt.")266 prompt = gr.Textbox(label="Prompt", placeholder="A HD photo of a Corgi")267 text_button = gr.Button(label="Generate")268 269 with gr.Accordion("Advanced options", open=False):270 model = gr.Radio(["base40M", "base300M", "base1B"], label="Model", value="base1B")271 scale = gr.Slider(272 label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.1273 )274 275 with gr.Column():276 model_gif = gr.Video(label="3D Model GIF")277 # btn_pc_to_obj = gr.Button(value="Convert to OBJ", visible=False)278 model_3d = gr.Model3D(value=None)279 file_out = gr.File(label="Files", visible=False)280 281 if torch.cuda.is_available():282 gr.Examples(283 examples=[284 ["images/pumpkin.png"],285 ["images/fantasy_world.png"],286 ],287 inputs=[input_image],288 outputs=[model_3d, model_gif, file_out],289 fn=generate_3D,290 cache_examples=True291 )292 293 img_button.click(fn=generate_3D, inputs=[input_image, model, scale], outputs=[model_3d, model_gif, file_out])294 text_button.click(fn=generate_3D, inputs=[prompt, model, scale], outputs=[model_3d, model_gif, file_out])295 296block.launch(show_api=False)297