gpu

app.py

@@ -93,8 +93,6 @@ model_index = 0
 # ヒートマップの生成関数
 @spaces.GPU
 def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
-    global model, mean_vector_list
-    
     if type(uploaded_image) == str:
         uploaded_image = Image.open(uploaded_image)
     if type(source_num) == str:
@@ -104,14 +102,14 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
     if type(y_coords) == str:
         y_coords = int(y_coords)
 
-    dec5, _ = model(x)
+    dec5, _ = models[model_index](x)
     feature_map = dec5
     # アップロード画像の前処理
     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).to(device)
-    target_feature_map, _ = model(uploaded_image)
+    target_feature_map, _ = models[model_index](uploaded_image)
     img = torch.cat((x, uploaded_image))
     feature_map = torch.cat((feature_map, target_feature_map))