Upload app.py with huggingface_hub

app.py

@@ -7,31 +7,16 @@ from rdkit import Chem
 from rdkit.Chem import Draw
 from graph_decoder.diffusion_model import GraphDiT
 
-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(path='model_labeled'):
-    model = GraphDiT(
+ATOM_SYMBOLS = ['C', 'N', 'O', 'H']
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+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']
+        model_dtype=torch.float32
+        )
+model.to(device)
 
 def generate_random_smiles(length=10):
     return ''.join(random.choices(ATOM_SYMBOLS, k=length))
@@ -42,15 +27,6 @@ def generate_polymer(CH4, CO2, H2, N2, O2, guidance_scale):
     
     print('in generate_polymer')
     try:
-        model = load_graph_decoder()
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model.to(device)
-        properties = properties.to(device)
-        
-        with torch.no_grad():
-            output = model(properties)
-        print('output', output)
-        
         # Generate a random SMILES string (this is a placeholder)
         generated_molecule = generate_random_smiles()