import gradio as gr import plotly.graph_objects as go import torch from tqdm.auto import tqdm 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 import trimesh device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print('creating base model...') base_name = 'base40M-textvec' base_model = model_from_config(MODEL_CONFIGS[base_name], device) base_model.eval() base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name]) print('creating upsample model...') upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device) upsampler_model.eval() upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample']) print('downloading base checkpoint...') base_model.load_state_dict(load_checkpoint(base_name, device)) print('downloading upsampler checkpoint...') upsampler_model.load_state_dict(load_checkpoint('upsample', device)) print('creating SDF model...') name = 'sdf' sdf_model = model_from_config(MODEL_CONFIGS[name], device) sdf_model.eval() print('loading SDF model...') sdf_model.load_state_dict(load_checkpoint(name, device)) sampler = PointCloudSampler( device=device, models=[base_model, upsampler_model], diffusions=[base_diffusion, upsampler_diffusion], num_points=[1024, 4096 - 1024], aux_channels=['R', 'G', 'B'], guidance_scale=[3.0, 0.0], model_kwargs_key_filter=('texts', ''), # Do not condition the upsampler at all ) def inference(prompt): samples = None for x in sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=[prompt])): samples = x pc = sampler.output_to_point_clouds(samples)[0] pc = sampler.output_to_point_clouds(samples)[0] colors=(238, 75, 43) fig = go.Figure( data=[ go.Scatter3d( x=pc.coords[:,0], y=pc.coords[:,1], z=pc.coords[:,2], mode='markers', marker=dict( size=2, color=['rgb({},{},{})'.format(r,g,b) for r,g,b in zip(pc.channels["R"], pc.channels["G"], pc.channels["B"])], ) ) ], layout=dict( scene=dict( xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False) ) ), ) # Produce a mesh (with vertex colors) mesh = marching_cubes_mesh( pc=pc, model=sdf_model, batch_size=4096, grid_size=32, # 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("mesh.ply", 'wb') as f: mesh.write_ply(f) obj_file = '3d_model.obj' mesh = trimesh.load('mesh.ply') mesh.export(obj_file) return fig, obj_file demo = gr.Interface( fn=inference, inputs="text", outputs=[gr.Plot(),gr.Model3D(value=None)], examples=[ ["a red motorcycle"], ["a RED pumpkin"], ["a yellow rubber duck"] ], title="Point-E demo: text to 3D", description="""Generated 3D Point Cloiuds with [Point-E](https://github.com/openai/point-e/tree/main). This demo uses a small, worse quality text-to-3D model to produce 3D point clouds directly from text descriptions. Check out the [notebook](https://github.com/openai/point-e/blob/main/point_e/examples/text2pointcloud.ipynb). """ ) demo.queue(max_size=30) demo.launch(debug=True)