Update app.py

app.py

@@ -1,4 +1,3 @@
-from PIL import Image
 import torch
 from tqdm.auto import tqdm
 
@@ -10,22 +9,45 @@ from point_e.util.plotting import plot_point_cloud
 import streamlit as st
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 st.write('creating base model...')
-base_name = 'base40M' # use base300M or base1B for better results
+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])
+
 st.write('creating upsample model...')
 upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)
 upsampler_model.eval()
 upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])
+
 st.write('downloading base checkpoint...')
 base_model.load_state_dict(load_checkpoint(base_name, device))
+
 st.write('downloading upsampler checkpoint...')
 upsampler_model.load_state_dict(load_checkpoint('upsample', device))
 
 
+# Define Sampler
+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
+)
+
+
+# Load an image to condition on.
+prompt = st.sidebar.text_input("Prompt")
+
 
+# Produce a sample from the model.
+samples = None
+for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=[prompt]))):
+    samples = x