Upload app.py with huggingface_hub

app.py

@@ -6,24 +6,18 @@ import random
 from rdkit import Chem
 from rdkit.Chem import Draw
 
-class RandomPolymerGenerator(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.fc1 = nn.Linear(5, 64)
-        self.fc2 = nn.Linear(64, 128)
-        self.fc3 = nn.Linear(128, 256)
-        self.fc4 = nn.Linear(256, 100)  # Output size set to 100 for simplicity
-
-    def forward(self, x):
-        x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = torch.relu(self.fc3(x))
-        return torch.sigmoid(self.fc4(x))
-
-def load_graph_decoder():
-    model = RandomPolymerGenerator()
+from graph_decoder.diffusion_model import GraphDiT
+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()
     return model
 
+
 ATOM_SYMBOLS = ['C', 'N', 'O', 'H']
 
 def generate_random_smiles(length=10):