gpu

app.py

@@ -17,8 +17,8 @@ import spaces
 
 image_size = 112
 batch_size = 32
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-@spaces.GPU
 def load_model(model_path="checkpoints/autoencoder-epoch=49-train_loss=1.01.ckpt", feature_dim=64):
     model = AutoencoderModule(feature_dim=feature_dim)
     state_dict = torch.load(model_path)
@@ -35,10 +35,9 @@ def load_model(model_path="checkpoints/autoencoder-epoch=49-train_loss=1.01.ckpt
         model.load_state_dict(new_state_dict)
         model.eval()
     
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
     print("Model loaded successfully.")
-    return model, device
+    return model
 
 def load_data(device, img_dir="resources/trainB/", image_size=112, batch_size=256):
     filenames = load_filenames(img_dir)
@@ -121,14 +120,14 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
     return fig
 
 def setup(model_dict, input_image=None):
-    global model, device, x, test_imgs, points, mean_vector_list
+    global model, x, test_imgs, points, mean_vector_list
     # str -> dictに変換
     if type(model_dict) == str:
         model_dict = eval(model_dict)
     model_name = model_dict["name"]
     feature_dim = model_dict["feature_dim"]
     model_path = f"checkpoints/{model_name}"
-    model, device = load_model(model_path, feature_dim)
+    model = load_model(model_path, feature_dim)
     x = load_data(device)
     test_imgs, points = load_keypoints(device)
     feature_map, _ = model(test_imgs)