Upload app.py with huggingface_hub

app.py

@@ -3,18 +3,18 @@ import gradio as gr
 import torch
 from rdkit import Chem
 from rdkit.Chem import Draw
-# from graph_decoder.diffusion_model import GraphDiT
+from graph_decoder.diffusion_model import GraphDiT
 
 # Load the model
 def load_graph_decoder(path='model_labeled'):
-    # model = GraphDiT(
-    #     model_config_path=f"{path}/config.yaml",
-    #     data_info_path=f"{path}/data.meta.json",
-    #     model_dtype=torch.float32,
-    # )
-    # model.init_model(path)
-    # model.disable_grads()
-    model = None
+    model = GraphDiT(
+        model_config_path=f"{path}/config.yaml",
+        data_info_path=f"{path}/data.meta.json",
+        model_dtype=torch.float32,
+    )
+    model.init_model(path)
+    model.disable_grads()
+    # model = None
     return model
 
 model = load_graph_decoder()
@@ -25,8 +25,8 @@ def generate_polymer(CH4, CO2, H2, N2, O2, guidance_scale):
     properties = [CH4, CO2, H2, N2, O2]
     
     try:
-        print('enter generate polymer')
         model.to(device)
+        print('enter function')
         generated_molecule, _ = model.generate(properties, device=device, guide_scale=guidance_scale)
         
         if generated_molecule is not None: