Update app.py

app.py

@@ -285,34 +285,26 @@ def tryon():
 
 @spaces.GPU 
 def generate_mask(human_img, categorie='upper_body'):
+    device = "cuda"
+    openpose_model.preprocessor.body_estimation.model.to(device)
+    pipe.to(device)
+
     try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
         # Redimensionner l'image pour le modèle
-        human_img_resized = human_img.resize((384, 512))
-        
-        # Déplacer les modèles sur le bon appareil
-        openpose_model.to(device)
-        parsing_model.to(device)
-        
-        # Convertir l'image en tenseur et déplacer sur l'appareil
-        input_tensor = transforms.ToTensor()(human_img_resized).unsqueeze(0).to(device)
-        
-        # Génération des points clés et du masque
-        keypoints = openpose_model(input_tensor)
-        model_parse, _ = parsing_model(input_tensor)
+        human_img_resized = human_img.convert("RGB").resize((384, 512))
+
+        # Générer les points clés et le masque
+        keypoints = openpose_model(human_img_resized)
+        model_parse, _ = parsing_model(human_img_resized)
         mask, _ = get_mask_location('hd', categorie, model_parse, keypoints)
-        
-        # Déplacement du masque sur le CPU pour traitement avec PIL
-        mask = mask.cpu()
-        
+
         # Redimensionner le masque à la taille d'origine de l'image
-        mask_resized = transforms.ToPILImage()(mask.squeeze(0)).resize(human_img.size)
+        mask_resized = mask.resize(human_img.size)
         
         return mask_resized
     except Exception as e:
         logging.error(f"Error generating mask: {e}")
-        raise e  # Renvoyer l'exception pour que l'API puisse la gérer
+        raise e
 
    
 @app.route('/generate_mask', methods=['POST'])
@@ -324,7 +316,7 @@ def generate_mask_api():
         categorie = data.get('categorie', 'upper_body')
         
         # Décodage de l'image à partir de base64
-        human_img = decode_image_from_base64(base64_image).convert("RGB")
+        human_img = decode_image_from_base64(base64_image)
         
         # Appeler la fonction pour générer le masque
         mask_resized = generate_mask(human_img, categorie)