Update app.py

app.py

@@ -21,7 +21,7 @@ import moleculekit
 print(moleculekit.__version__)
 
 
-def update(inp, file, mode, custom_resids, clustering_threshold, distance_cutoff):
+def update(inp, file, mode, custom_resids, clustering_threshold, distance_cutoff,model, device):
     try:
         filepath = file.name
     except:
@@ -70,18 +70,9 @@ def update(inp, file, mode, custom_resids, clustering_threshold, distance_cutoff
             "Error",
             f"""<div class="text-center mt-4"> Something went wrong with the voxelization, reset custom residues and other input fiels and check error message <br> <br> <code>{e}</code></div>""",
         )
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
     voxels.to(device)
 
-    model = Model()
-    model.to(device)
-    model.load_state_dict(
-        torch.load(
-            "weights/metal_0.5A_v3_d0.2_16Abox.pth",
-            map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
-        )
-    )
-    model.eval()
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore")
         output = model(voxels)
@@ -108,7 +99,8 @@ def update(inp, file, mode, custom_resids, clustering_threshold, distance_cutoff
         threshold=distance_cutoff,
         p=clustering_threshold,
     )
-
+    del voxels
+    torch.cuda.empty_cache()
     return message, molecule(
         filepath,
         f"output/probes_{identifier}.pdb",
@@ -324,7 +316,17 @@ def set_examples(example):
     n, code, resids = example
     return [n, code, resids]
 
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = Model()
+model.to(device)
+model.load_state_dict(
+    torch.load(
+        "weights/metal_0.5A_v3_d0.2_16Abox.pth",
+        map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    )
+)
+model.eval()
+    
 metal3d = gr.Blocks()
 
 with metal3d:
@@ -389,7 +391,7 @@ with metal3d:
     mol = gr.HTML()
     btn.click(
         fn=update,
-        inputs=[inp, file, mode, custom_resids, clustering_threshold, distance_cutoff],
+        inputs=[inp, file, mode, custom_resids, clustering_threshold, distance_cutoff, model, device],
         outputs=[out, mol],
     )