Update app.py

app.py

@@ -17,7 +17,42 @@ from datetime import datetime
 import json
 
 from evaluator import Evaluator
-from loader import load_graph_decoder
+# 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
 
 # Load the CSV data
 known_labels = pd.read_csv('data/known_labels.csv')