Upload app.py with huggingface_hub

app.py

@@ -61,6 +61,9 @@ def load_model(model_choice):
 flagged_folder = "flagged"
 os.makedirs(flagged_folder, exist_ok=True)
 
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# model.to(device)
+
 def save_interesting_log(smiles, properties, suggested_properties):
     """Save interesting polymer data to a CSV file."""
     log_file = os.path.join(flagged_folder, "log.csv")
@@ -81,7 +84,6 @@ def save_interesting_log(smiles, properties, suggested_properties):
         }
         writer.writerow(log_data)
 
-@spaces.GPU(duration=60)
 def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
     model, device = model_state
     
@@ -106,10 +108,8 @@ def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_ti
                 # print('Before generation, move model to', device)
                 # return generated_molecule, img_list
             # generated_molecule, img_list = generate_func()
-
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            # print('Before generation, move model to', device)
             model.to(device)
-            print('Before generation, move model to', device)
             generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
 
             # Create GIF if img_list is available
@@ -182,6 +182,7 @@ def numpy_to_python(obj):
     else:
         return obj
         
+@spaces.GPU(duration=60)
 def on_generate(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
     result = generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps)
     # Check if the generation was successful