update

app.py

@@ -17,17 +17,19 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 os.makedirs("results", exist_ok=True)
 os.makedirs("models", exist_ok=True)
 
-subprocess.run(
-    'wget https://zenodo.org/record/7121300/files/geom_size_gnn.ckpt?download=1 -O models/geom_size_gnn.ckpt',
-    shell=True
-)
+size_gnn_path = 'models/geom_size_gnn.ckpt'
+if not os.path.exists(size_gnn_path):
+    print('Downloading SizeGNN model...')
+    link = 'https://zenodo.org/record/7121300/files/geom_size_gnn.ckpt?download=1'
+    subprocess.run(f'wget {link} -O {size_gnn_path}', shell=True)
 size_nn = SizeClassifier.load_from_checkpoint('models/geom_size_gnn.ckpt', map_location=device).eval().to(device)
 print('Loaded SizeGNN model')
 
-subprocess.run(
-    'wget https://zenodo.org/record/7121300/files/geom_difflinker.ckpt?download=1 -O models/geom_difflinker.ckpt',
-    shell=True
-)
+diffusion_path = 'models/geom_difflinker.ckpt'
+if not os.path.exists(diffusion_path):
+    print('Downloading Diffusion model...')
+    link = 'https://zenodo.org/record/7121300/files/geom_difflinker.ckpt?download=1'
+    subprocess.run(f'wget {link} -O {diffusion_path}', shell=True)
 ddpm = DDPM.load_from_checkpoint('models/geom_difflinker.ckpt', map_location=device).eval().to(device)
 print('Loaded diffusion model')
 
@@ -174,4 +176,4 @@ with demo:
         outputs=[visualization, output_files],
     )
 
-demo.launch()
+demo.launch(share=True)