Update app.py

app.py

@@ -4,6 +4,8 @@ import plotly.graph_objects as go
 import sys
 import torch
 from huggingface_hub import hf_hub_download
+import numpy as np
+import random
 
 os.system("git clone https://github.com/luost26/diffusion-point-cloud")
 sys.path.append("diffusion-point-cloud")
@@ -17,8 +19,13 @@ chair=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename=
 
 device='cuda' if torch.cuda.is_available() else 'cpu'
 
-ckpt_airplane = torch.load(airplane)
-ckpt_chair = torch.load(chair)
+ckpt_airplane = torch.load(airplane,map_location=torch.device(device))
+ckpt_chair = torch.load(chair,map_location=torch.device(device))
+
+def seed_all(seed):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
 
 def normalize_point_clouds(pcs,mode):
     if mode is None:
@@ -67,7 +74,6 @@ def generate(seed,value):
     else :
        ckpt=ckpt_airplane
 
-    print(value)
     colors=(238, 75, 43)
     points=predict(seed,ckpt)
     num_points=points.shape[0]
@@ -90,14 +96,13 @@ def generate(seed,value):
         )
     )
     return fig
+    
 markdown=f'''
   # Diffusion Probabilistic Models for 3D Point Cloud Generation
-
   [[The Paper](https://arxiv.org/abs/2103.01458)] [[Original Code](https://github.com/luost26/diffusion-point-cloud)]
-
   The space demo for our CVPR 2021 paper "Diffusion Probabilistic Models for 3D Point Cloud Generation".
 
-
+  It is running on {device}
 '''
 with gr.Blocks() as demo:
     with gr.Column():