Update app.py

app.py

@@ -10,6 +10,7 @@ import random
 os.system("git clone https://github.com/luost26/diffusion-point-cloud")
 sys.path.append("diffusion-point-cloud")
 
+#Codes reference : https://github.com/luost26/diffusion-point-cloud
 
 from models.vae_gaussian import *
 from models.vae_flow import *
@@ -44,7 +45,10 @@ def normalize_point_clouds(pcs,mode):
         pcs[i] = pc
     return pcs
 
-def predict(Seed,ckpt):
+    
+
+
+def predict(Seed,ckpt,truncate_std):
   if Seed==None:
     Seed=777 
   seed_all(Seed)
@@ -59,14 +63,14 @@ def predict(Seed,ckpt):
   gen_pcs = []
   with torch.no_grad():
       z = torch.randn([1, ckpt['args'].latent_dim]).to(device)
-      x = model.sample(z, 2048, flexibility=ckpt['args'].flexibility)
+      x = model.sample(z, 2048, flexibility=ckpt['args'].flexibility,truncate_std=truncate_std))
       gen_pcs.append(x.detach().cpu())
   gen_pcs = torch.cat(gen_pcs, dim=0)[:1]
   gen_pcs = normalize_point_clouds(gen_pcs, mode="shape_bbox")
 
   return gen_pcs[0]
 
-def generate(seed,value):
+def generate(seed,value,truncate_std):
     if value=="Airplane":
       ckpt=ckpt_airplane
     elif value=="Chair":
@@ -75,7 +79,7 @@ def generate(seed,value):
        ckpt=ckpt_airplane
 
     colors=(238, 75, 43)
-    points=predict(seed,ckpt)
+    points=predict(seed,ckpt,truncate_std)
     num_points=points.shape[0]
 
 
@@ -99,9 +103,18 @@ def generate(seed,value):
     
 markdown=f'''
   # Diffusion Probabilistic Models for 3D Point Cloud Generation
+
+  
   [The space demo for the CVPR 2021 paper "Diffusion Probabilistic Models for 3D Point Cloud Generation".](https://arxiv.org/abs/2103.01458)
+  
   [For the official implementation.](https://github.com/luost26/diffusion-point-cloud)
 
+  ### Future Work based on interest
+  - Adding new models for new type objects
+  - New Customization 
+  
+  
+
   It is running on {device}
   
 
@@ -113,6 +126,7 @@ with gr.Blocks() as demo:
         with gr.Row():
             seed = gr.Slider( minimum=0, maximum=2**16,label='Seed')
             value=gr.Dropdown(choices=["Airplane","Chair"],label="Choose Model Type")
+            truncate_std = gr.Slider( minimum=1, maximum=2,label='Truncate Std')
 
         btn = gr.Button(value="Generate")
         point_cloud = gr.Plot()