if mian

app.py

@@ -19,6 +19,12 @@ image_size = 112
 batch_size = 32
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+models_info = [{"name": "ae_model_tf_2024-03-05_00-35-21.pth", "feature_dim": 32},
+        {"name": "autoencoder-epoch=09-train_loss=1.00.ckpt", "feature_dim": 64},
+        {"name": "autoencoder-epoch=29-train_loss=1.01.ckpt", "feature_dim": 64},
+        {"name": "autoencoder-epoch=49-train_loss=1.01.ckpt", "feature_dim": 64}]
+models = []
+
 @spaces.GPU
 def load_model(model_path, feature_dim):
     model = AutoencoderModule(feature_dim=feature_dim)
@@ -141,21 +147,6 @@ def setup(model_info, input_image=None):
         fig = get_heatmaps(0, image_size // 2, image_size // 2, input_image)
         return fig
 
-
-models_info = [{"name": "ae_model_tf_2024-03-05_00-35-21.pth", "feature_dim": 32},
-        {"name": "autoencoder-epoch=09-train_loss=1.00.ckpt", "feature_dim": 64},
-        {"name": "autoencoder-epoch=29-train_loss=1.01.ckpt", "feature_dim": 64},
-        {"name": "autoencoder-epoch=49-train_loss=1.01.ckpt", "feature_dim": 64}]
-models = []
-for model_info in models_info:
-    model_name = model_info["name"]
-    feature_dim = model_info["feature_dim"]
-    model_path = f"checkpoints/{model_name}"
-    model = load_model(model_path, feature_dim)
-    models.append(model)
-    
-setup(models_info[0])
-
 with gr.Blocks() as demo:
     # title
     gr.Markdown("# TripletGeoEncoder Feature Map Visualization")
@@ -204,7 +195,18 @@ with gr.Blocks() as demo:
     ],
         inputs=[input_image],
     )
-            
+
+
+if __name__ == "__main__":    
+    for model_info in models_info:
+        model_name = model_info["name"]
+        feature_dim = model_info["feature_dim"]
+        model_path = f"checkpoints/{model_name}"
+        model = load_model(model_path, feature_dim)
+        models.append(model)
+        
+    setup(models_info[0])
+
     # JavaScriptコードをロード
     demo.launch()