Update app.py

app.py

@@ -283,18 +283,13 @@ def tryon():
         'mask_image': mask_base64
     })
 
-@spaces.GPU    
-@app.route('/generate_mask', methods=['POST'])
-def generate_mask():
-    data = request.json
-    base64_image = data.get('image')
-    categorie = data.get('categorie', 'upper_body')
-    
-    # Décodage de l'image
-    human_img = decode_image_from_base64(base64_image).convert("RGB")
-    human_img_resized = human_img.resize((384, 512))
-    
+@spaces.GPU 
+def generate_mask(human_img, categorie='upper_body'):
+    device = "cuda"
     try:
+        # Redimensionnement de l'image pour le modèle
+        human_img_resized = human_img.resize((384, 512))
+        
         # Génération des points clés et du masque
         keypoints = openpose_model(human_img_resized)
         model_parse, _ = parsing_model(human_img_resized)
@@ -303,12 +298,32 @@ def generate_mask():
         # Redimensionnement du masque à la taille d'origine de l'image
         mask_resized = mask.resize(human_img.size)
         
-        # Encodage du masque en base64
+        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
+
+   
+@app.route('/generate_mask', methods=['POST'])
+def generate_mask_api():
+    try:
+        # Récupérer les données de l'image à partir de la requête
+        data = request.json
+        base64_image = data.get('image')
+        categorie = data.get('categorie', 'upper_body')
+        
+        # Décodage de l'image à partir de base64
+        human_img = decode_image_from_base64(base64_image).convert("RGB")
+        
+        # Appeler la fonction pour générer le masque
+        mask_resized = generate_mask(human_img, categorie)
+        
+        # Encodage du masque en base64 pour la réponse
         mask_base64 = encode_image_to_base64(mask_resized)
         
         return jsonify({
             'mask_image': mask_base64
-        })
+        }), 200
     except Exception as e:
         logging.error(f"Error generating mask: {e}")
         return jsonify({'error': str(e)}), 500