Upload app.py with huggingface_hub

app.py

@@ -52,10 +52,12 @@ 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")
-    model = load_graph_decoder(path=model_choice)
-    return (model, device)
+# def load_model(model_choice):
+#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model_all = load_graph_decoder(path='model_all')
+model_labeled = load_graph_decoder(path='model_labeled')
+    # return (model, device)
 
 # Create a flagged folder if it doesn't exist
 flagged_folder = "flagged"
@@ -279,7 +281,7 @@ with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
         ```
         """)
 
-    model_state = gr.State(lambda: load_model("model_all"))
+    model_state = gr.State(model_labeled)
 
     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}]")
@@ -312,7 +314,12 @@ with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
     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)
+        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])