Upload app.py with huggingface_hub

app.py

@@ -1,6 +1,5 @@
-import gradio as gr
 import spaces
-
+import gradio as gr
 from rdkit import Chem
 from rdkit.Chem import Draw
 import numpy as np
@@ -16,45 +15,9 @@ import glob
 import csv
 from datetime import datetime
 import json
-# import spaces
 
 from evaluator import Evaluator
-# from loader import load_graph_decoder
-
-### loader
-from graph_decoder.diffusion_model import GraphDiT
-def count_parameters(model):
-    r"""
-    Returns the number of trainable parameters and number of all parameters in the model.
-    """
-    trainable_params, all_param = 0, 0
-    for param in model.parameters():
-        num_params = param.numel()
-        all_param += num_params
-        if param.requires_grad:
-            trainable_params += num_params
-
-    return trainable_params, all_param
-
-def load_graph_decoder(path='model_labeled'):
-    model_config_path = f"{path}/config.yaml"
-    data_info_path = f"{path}/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,
-    )
-    model.init_model(path)
-    model.disable_grads()
-
-    trainable_params, all_param = count_parameters(model)
-    param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
-        path, trainable_params, all_param, 100 * trainable_params / all_param
-    )
-    print(param_stats)
-    return model
+from loader import load_graph_decoder
 
 # Load the CSV data
 known_labels = pd.read_csv('data/known_labels.csv')
@@ -92,8 +55,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)
-    # 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
@@ -122,7 +83,6 @@ 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"))
     model, device = model_state
     
     properties = [CH4, CO2, H2, N2, O2]
@@ -139,53 +99,61 @@ 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:
-        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()])
-                
+        try:
+            # def generate_func():
+                # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+                # model.to(device)
+                # print('Before generation, move model to', device)
+                # return generated_molecule, img_list
+            # generated_molecule, img_list = generate_func()
+
+            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:
                 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