Update app.py

app.py

@@ -18,7 +18,7 @@ from rdkit import Chem
 from rdkit.Chem import Draw
 
 from evaluator import Evaluator
-from loader import load_graph_decoder
+# from loader import load_graph_decoder
 
 # Load the CSV data
 known_labels = pd.read_csv('data/known_labels.csv')
@@ -55,7 +55,19 @@ 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 = 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