Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify,send_file
 from PIL import Image
 from io import BytesIO
 import torch
@@ -9,6 +9,8 @@ import logging
 import gradio as gr
 import numpy as np
 import spaces
+import uuid
+import random
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
 from src.unet_hacked_tryon import UNet2DConditionModel
@@ -123,6 +125,15 @@ def pil_to_binary_mask(pil_image, threshold=0):
     output_mask = Image.fromarray(mask)
     return output_mask
 
+def get_image_from_url(url):
+    try:
+        response = requests.get(url)
+        response.raise_for_status()  # Vérifie les erreurs HTTP
+        img = Image.open(BytesIO(response.content))
+        return img
+    except Exception as e:
+        logging.error(f"Error fetching image from URL: {e}")
+        raise
 
 def decode_image_from_base64(base64_str):
     try:
@@ -143,6 +154,11 @@ def encode_image_to_base64(img):
         logging.error(f"Error encoding image: {e}")
         raise
 
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".webp"  
+    img.save(unique_name, format="WEBP", lossless=True) 
+    return unique_name
+
 @spaces.GPU
 def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, categorie = 'upper_body'):
     device = "cuda"
@@ -255,11 +271,18 @@ def clear_gpu_memory():
     torch.cuda.empty_cache()
     torch.cuda.synchronize()
 
+def process_image(image_data):
+    # Vérifie si l'image est en base64 ou URL
+    if image_data.startswith('http://') or image_data.startswith('https://'):
+        return get_image_from_url(image_data)  # Télécharge l'image depuis l'URL
+    else:
+        return decode_image_from_base64(image_data)  # Décode l'image base64
+
 @app.route('/tryon', methods=['POST'])
 def tryon():
     data = request.json
-    human_image = decode_image_from_base64(data['human_image'])
-    garment_image = decode_image_from_base64(data['garment_image'])
+    human_image = process_image(data['human_image'])
+    garment_image = process_image(data['garment_image'])
     description = data.get('description')
     use_auto_mask = data.get('use_auto_mask', True)
     use_auto_crop = data.get('use_auto_crop', False)
@@ -283,6 +306,32 @@ def tryon():
         'mask_image': mask_base64
     })
 
+@app.route('/tryon-v2', methods=['POST'])
+def tryon():
+    data = request.json
+    human_image = process_image(data['human_image'])
+    garment_image = process_image(data['garment_image'])
+    mask_image = process_image(data['mask_image'])
+    description = data.get('description')
+    use_auto_mask = data.get('use_auto_mask', True)
+    use_auto_crop = data.get('use_auto_crop', False)
+    denoise_steps = int(data.get('denoise_steps', 30))
+    seed = int(data.get('seed', random.randint(0, 9999999)))
+    categorie = data.get('categorie' , 'upper_body')
+    human_dict = {
+        'background': human_image,
+        'layers': [mask_image] if not use_auto_mask else None,
+        'composite': None
+    }
+    output_image, mask_image = start_tryon(human_dict, garment_image, description, use_auto_mask, use_auto_crop, denoise_steps, seed , categorie)
+
+    output_id = 
+    mask_base64 = encode_image_to_base64(mask_image)
+
+    return jsonify({
+        'image_id': save_image(output_image)
+    })
+
 @spaces.GPU 
 def generate_mask(human_img, categorie='upper_body'):
     device = "cuda"
@@ -316,7 +365,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)
+        human_img = process_image(base64_image)
         
         # Appeler la fonction pour générer le masque
         mask_resized = generate_mask(human_img, categorie)
@@ -331,6 +380,17 @@ def generate_mask_api():
         logging.error(f"Error generating mask: {e}")
         return jsonify({'error': str(e)}), 500
 
+# Route pour récupérer l'image générée
+@app.route('/api/get_image/<image_id>', methods=['GET'])
+def get_image(image_id):
+    # Construire le chemin complet de l'image
+    image_path = image_id  # Assurez-vous que le nom de fichier correspond à celui que vous avez utilisé lors de la sauvegarde
+
+    # Renvoyer l'image
+    try:
+        return send_file(image_path, mimetype='image/webp')
+    except FileNotFoundError:
+        return jsonify({'error': 'Image not found'}), 404
 
 if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)