|
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', ''), |
|
) |
|
|
|
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) |
|
|
