gpu

app.py

@@ -15,16 +15,6 @@ import utils
 
 import spaces
 
-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 = []
-
 def load_model(model_path, feature_dim):
     model = AutoencoderModule(feature_dim=feature_dim)
     state_dict = torch.load(model_path)
@@ -42,10 +32,25 @@ def load_model(model_path, feature_dim):
         model.eval()
     
     model.to(device)
-    print("Model loaded successfully.")
+    print(f"{model_path} loaded successfully.")
     return model
 
-@spaces.GPU
+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]
@@ -62,7 +67,6 @@ def load_data(img_dir="resources/trainB/", image_size=112, batch_size=256):
     print("Data loaded successfully.")
     return x
 
-@spaces.GPU
 def load_keypoints(img_dir="resources/trainB/", image_size=112, batch_size=32):
     filenames = load_filenames(img_dir)
     train_X = filenames[:1000]
@@ -98,7 +102,7 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
     if uploaded_image is not None:
         uploaded_image = utils.preprocess_uploaded_image(uploaded_image['composite'], image_size)
     else:
-        uploaded_image = torch.zeros(1, 3, image_size, image_size, device=device)
+        uploaded_image = torch.zeros(1, 3, image_size, image_size).to(device)
     target_feature_map, _ = model(uploaded_image)
     img = torch.cat((x, uploaded_image))
     feature_map = torch.cat((feature_map, target_feature_map))
@@ -135,7 +139,7 @@ def setup(model_info, input_image=None):
         
     index = models_info.index(model_info)
     model = models[index]
-    
+
     x = load_data()
     test_imgs, points = load_keypoints()
     
@@ -195,14 +199,7 @@ 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)
+if __name__ == "__main__":
         
     setup(models_info[0])