Update app.py

app.py

@@ -314,6 +314,7 @@ def tryon():
         'mask_image': mask_base64
     })
 
+@spaces.GPU
 @app.route('/get_mask', methods=['POST'])
 def get_mask():
     try:
@@ -327,20 +328,28 @@ def get_mask():
         keypoints = openpose_model(img)  # Utilise votre modèle
         model_parse, _ = parsing_model(img)  # Utilise votre modèle
 
+        # Déplacer le modèle et les images sur le même dispositif
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        img_tensor = transforms.ToTensor()(img).unsqueeze(0).to(device)  # Convertir et déplacer l'image
+
+        # Assurez-vous que le modèle est sur le même dispositif
+        parsing_model.to(device)
+
         # Obtenir le masque
-        mask, mask_gray = get_mask_location('hd', categorie , model_parse, keypoints)
-        
+        mask, mask_gray = get_mask_location('hd', categorie, model_parse, keypoints)
+
         # Convertir le masque en image (si nécessaire)
-        mask_gray = (1 - transforms.ToTensor()(mask_gray)) * tensor_transfrom(img)
+        mask_gray = (1 - transforms.ToTensor()(mask_gray)) * tensor_transform(img_tensor)
         mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)
 
         # Convertir l'image en base64 si besoin pour le retour
         img_byte_arr = io.BytesIO()
         mask_gray.save(img_byte_arr, format='PNG')
         img_byte_arr.seek(0)
-        return jsonify({'mask': img_byte_arr.getvalue().decode('latin1')})  # Utiliser une méthode appropriée pour l'encodage
+        return jsonify({'mask': img_byte_arr.getvalue().decode('latin1')})
 
     except Exception as e:
+        print(e)
         return jsonify({'error': str(e)}), 500
 
 # Route index