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 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)) 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=1, color=colors) ) ], layout=dict( scene=dict( xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False) ) ) ) return fig demo = gr.Interface( fn=inference, inputs="text", outputs=gr.Plot(), examples=[["a red motorcycle"]], 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." ) demo.launch(debug=True)