Update app.py

app.py

@@ -56,18 +56,6 @@ 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)
-    #### test
-    # from graph_decoder.diffusion_model import GraphDiT
-    
-    # model_config_path = f"model_labeled/config.yaml"
-    # data_info_path = f"model_labeled/data.meta.json"
-    # model = GraphDiT(
-    #     model_config_path=model_config_path,
-    #     data_info_path=data_info_path,
-    #     # model_dtype=torch.float16,
-    #     model_dtype=torch.float32,
-    # )
-    ### test
     return (model, device)
 
 # Create a flagged folder if it doesn't exist
@@ -112,55 +100,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, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
-        # generated_molecule = 'C'
-        # img_list = []
-        # 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()])
-                
+        try:
+            model.to(device)
+            generated_molecule, img_list = model.generate(properties, 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:
                 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,
+                    f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
+                    None,
                     gif_path,
-                    properties,  # Add this
-                    suggested_properties  # Add this
+                    properties,
+                    None,
                 )
-        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
+        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