gpu

app.py

@@ -35,22 +35,6 @@ def load_model(model_path, feature_dim):
     print(f"{model_path} loaded successfully.")
     return model
 
-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 = []
-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)
-
 def load_data(img_dir="resources/trainB/", image_size=112, batch_size=256):
     filenames = load_filenames(img_dir)
     train_X = filenames[:1000]
@@ -84,6 +68,25 @@ def load_keypoints(img_dir="resources/trainB/", image_size=112, batch_size=32):
     print("Keypoints loaded successfully.")
     return test_imgs, points
 
+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 = []
+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)
+
+x = load_data()
+test_imgs, points = load_keypoints()
+
 # ヒートマップの生成関数
 @spaces.GPU
 def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
@@ -131,8 +134,9 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
     plt.close(fig)
     return fig
 
+@spaces.GPU
 def setup(model_info, input_image=None):
-    global model, x, test_imgs, points, mean_vector_list
+    global model, mean_vector_list
     # str -> dictに変換
     if type(model_info) == str:
         model_info = eval(model_info)
@@ -140,9 +144,6 @@ def setup(model_info, input_image=None):
     index = models_info.index(model_info)
     model = models[index]
 
-    x = load_data()
-    test_imgs, points = load_keypoints()
-    
     feature_map, _ = model(test_imgs)
     mean_vector_list = utils.get_mean_vector(feature_map, points)