Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 import gradio as gr
 from rdkit import Chem
 from rdkit.Chem import Draw
@@ -15,6 +14,7 @@ import glob
 import csv
 from datetime import datetime
 import json
+# import spaces
 
 from evaluator import Evaluator
 # from loader import load_graph_decoder
@@ -90,6 +90,8 @@ def random_properties():
 def load_model(model_choice):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = load_graph_decoder(path=model_choice)
+    model.to(device)
+    print('in load_model', torch.device("cuda" if torch.cuda.is_available() else "cpu"))
     return (model, device)
 
 # Create a flagged folder if it doesn't exist
@@ -118,7 +120,7 @@ def save_interesting_log(smiles, properties, suggested_properties):
 
 # @spaces.GPU(duration=60)
 def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
-    print('in generate graph', torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    print('in generate_graph', torch.device("cuda" if torch.cuda.is_available() else "cpu"))
     model, device = model_state
     
     properties = [CH4, CO2, H2, N2, O2]
@@ -135,54 +137,53 @@ def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_ti
     num_nodes = None if num_nodes == 0 else num_nodes
     
     for _ in range(repeating_time):
-        try:
-            model.to(device)
-            generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
-
-            # Create GIF if img_list is available
-            gif_path = None
-            if img_list and len(img_list) > 0:
-                imgs = [np.array(pil_img) for pil_img in img_list]
-                imgs.extend([imgs[-1]] * 10)
-                gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
-                imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
-            
-            if generated_molecule is not None:
-                mol = Chem.MolFromSmiles(generated_molecule)
-                if mol is not None:
-                    standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
-                    is_novel = standardized_smiles not in knwon_smiles['smiles'].values
-                    novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
-                    img = Draw.MolToImage(mol)
-                    
-                    # Evaluate the generated molecule
-                    suggested_properties = {}
-                    for prop, evaluator in evaluators.items():
-                        suggested_properties[prop] = evaluator([standardized_smiles])[0]
-                    
-                    suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
-                    
-                    return (
-                        f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
-                        f"**{nan_message}**\n\n"
-                        f"**{novelty_status}**\n\n"
-                        f"**Suggested Properties:**\n{suggested_properties_text}",
-                        img,
-                        gif_path,
-                        properties,  # Add this
-                        suggested_properties  # Add this
-                    )
-            else:
+        # try:
+        generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
+
+        # Create GIF if img_list is available
+        gif_path = None
+        if img_list and len(img_list) > 0:
+            imgs = [np.array(pil_img) for pil_img in img_list]
+            imgs.extend([imgs[-1]] * 10)
+            gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
+            imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
+        
+        if generated_molecule is not None:
+            mol = Chem.MolFromSmiles(generated_molecule)
+            if mol is not None:
+                standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
+                is_novel = standardized_smiles not in knwon_smiles['smiles'].values
+                novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
+                img = Draw.MolToImage(mol)
+                
+                # Evaluate the generated molecule
+                suggested_properties = {}
+                for prop, evaluator in evaluators.items():
+                    suggested_properties[prop] = evaluator([standardized_smiles])[0]
+                
+                suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
+                
                 return (
-                    f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
-                    None,
+                    f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
+                    f"**{nan_message}**\n\n"
+                    f"**{novelty_status}**\n\n"
+                    f"**Suggested Properties:**\n{suggested_properties_text}",
+                    img,
                     gif_path,
-                    properties,
-                    None,
+                    properties,  # Add this
+                    suggested_properties  # Add this
                 )
-        except Exception as e:
-            print(f"Error in generation: {e}")
-            continue
+        else:
+            return (
+                f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
+                None,
+                gif_path,
+                properties,
+                None,
+            )
+        # except Exception as e:
+        #     print(f"Error in generation: {e}")
+        #     continue
     
     return f"**Generation failed:** Could not generate a valid molecule after {repeating_time} attempts.\n\n**{nan_message}**", None, None