Upload app.py with huggingface_hub

app.py

@@ -52,23 +52,15 @@ atexit.register(cleanup_temp_files)
 def random_properties():
     return known_labels[all_properties].sample(1).values.tolist()[0]
 
-# def load_model(model_choice):
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    # return (model, device)
-
-model_all = load_graph_decoder(path='model_all')
-model_labeled = load_graph_decoder(path='model_labeled')
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print('device')
-
+def load_model(model_choice):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = load_graph_decoder(path=model_choice)
+    return (model, device)
 
 # Create a flagged folder if it doesn't exist
 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")
@@ -89,8 +81,10 @@ def save_interesting_log(smiles, properties, suggested_properties):
         }
         writer.writerow(log_data)
 
-# def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model, num_chain_steps, fps):    
-def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, num_chain_steps, fps):    
+@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
+    
     properties = [CH4, CO2, H2, N2, O2]
     
     def is_nan_like(x):
@@ -112,8 +106,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()
-            # print('Before generation, move model to', device)
-            # model.to(device)
+
+            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
@@ -186,12 +180,8 @@ 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, num_chain_steps, fps):
-# def on_generate(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
-    # model = model_state
-    # result = generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model, num_chain_steps, fps)
-    result = generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, num_chain_steps, fps)
+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
     if result[0].startswith("**Generated polymer SMILES:**"):
         smiles = result[0].split("**Generated polymer SMILES:** `")[1].split("`")[0]
@@ -286,8 +276,7 @@ with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
         ```
         """)
 
-    # model_state = gr.State(value=model_labeled)
-    # model_state = model_labeled
+    model_state = gr.State(lambda: load_model("model_all"))
 
     with gr.Row():
         CH4_input = gr.Slider(0, property_ranges['CH4'][1], value=2.5, label=f"CH₄ (Barrier) [0-{property_ranges['CH4'][1]:.1f}]")
@@ -316,17 +305,13 @@ with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
         feedback_btn = gr.Button("🌟 I think this polymer is interesting!", visible=True, interactive=False)
         feedback_result = gr.Textbox(label="Feedback Result", visible=False)
     
-    # def switch_model(choice):
-    #     # Convert display name back to internal name
-    #     internal_name = next(key for key, value in model_name_mapping.items() if value == choice)
-    #     if internal_name == 'model_labeled':
-    #         return model_labeled
-    #     elif internal_name == 'model_all':
-    #         return model_all
-    #     else:
-    #         raise ValueError('Not support model', internal_name)
-
-    # model_choice.change(switch_model, inputs=[model_choice], outputs=[model_state])
+    # Add model switching functionality
+    def switch_model(choice):
+        # Convert display name back to internal name
+        internal_name = next(key for key, value in model_name_mapping.items() if value == choice)
+        return load_model(internal_name)
+
+    model_choice.change(switch_model, inputs=[model_choice], outputs=[model_state])
 
     # Hidden components to store generation data
     hidden_smiles = gr.Textbox(visible=False)
@@ -341,8 +326,7 @@ with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
 
     generate_btn.click(
         on_generate,
-        # inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps],
-        inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale, num_nodes, repeating_time, num_chain_steps, fps],
+        inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps],
         outputs=[result_text, result_image, result_gif, hidden_smiles, hidden_properties, hidden_suggested_properties, feedback_btn]
     )