Update app.py

app.py

@@ -15,6 +15,7 @@ from models.vae_flow import *
 airplane=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_airplane.pt",revision="main")
 chair=network_pkl=hf_hub_download("SerdarHelli/diffusion-point-cloud", filename="GEN_chair.pt",revision="main")
 
+device='cuda' if torch.cuda.is_available() else 'cpu'
 
 ckpt_airplane = torch.load(airplane)
 ckpt_chair = torch.load(chair)
@@ -42,15 +43,15 @@ def predict(Seed,ckpt):
   seed_all(Seed)
 
   if ckpt['args'].model == 'gaussian':
-      model = GaussianVAE(ckpt['args']).to("cuda")
+      model = GaussianVAE(ckpt['args']).to(device)
   elif ckpt['args'].model == 'flow':
-      model = FlowVAE(ckpt['args']).to("cuda")
+      model = FlowVAE(ckpt['args']).to(device)
 
   model.load_state_dict(ckpt['state_dict'])
   # Generate Point Clouds
   gen_pcs = []
   with torch.no_grad():
-      z = torch.randn([1, ckpt['args'].latent_dim]).to("cuda")
+      z = torch.randn([1, ckpt['args'].latent_dim]).to(device)
       x = model.sample(z, 2048, flexibility=ckpt['args'].flexibility)
       gen_pcs.append(x.detach().cpu())
   gen_pcs = torch.cat(gen_pcs, dim=0)[:1]