Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import requests
 from flask import Flask, request, jsonify,send_file
 from PIL import Image
 from io import BytesIO
@@ -32,107 +33,70 @@ from torchvision.transforms.functional import to_pil_image
 
 app = Flask(__name__)
 
+# Chemins de base pour les modèles
 base_path = 'yisol/IDM-VTON'
-example_path = os.path.join(os.path.dirname(__file__), 'example')
 
+# Chargement des modèles
 unet = UNet2DConditionModel.from_pretrained(
     base_path,
     subfolder="unet",
     torch_dtype=torch.float16,
     force_download=False
 )
-unet.requires_grad_(False)
 tokenizer_one = AutoTokenizer.from_pretrained(
     base_path,
     subfolder="tokenizer",
-    revision=None,
     use_fast=False,
     force_download=False
 )
 tokenizer_two = AutoTokenizer.from_pretrained(
     base_path,
     subfolder="tokenizer_2",
-    revision=None,
     use_fast=False,
     force_download=False
 )
 noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")
-
-text_encoder_one = CLIPTextModel.from_pretrained(
-    base_path,
-    subfolder="text_encoder",
-    torch_dtype=torch.float16,
-    force_download=False
-)
-text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
-    base_path,
-    subfolder="text_encoder_2",
-    torch_dtype=torch.float16,
-    force_download=False
-)
-image_encoder = CLIPVisionModelWithProjection.from_pretrained(
-    base_path,
-    subfolder="image_encoder",
-    torch_dtype=torch.float16,
-    force_download=False
-)
-vae = AutoencoderKL.from_pretrained(base_path,
-                                    subfolder="vae",
-                                    torch_dtype=torch.float16,
-                                    force_download=False
-)
-
-UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
-    base_path,
-    subfolder="unet_encoder",
-    torch_dtype=torch.float16,
-    force_download=False
-)
+text_encoder_one = CLIPTextModel.from_pretrained(base_path, subfolder="text_encoder", torch_dtype=torch.float16)
+text_encoder_two = CLIPTextModelWithProjection.from_pretrained(base_path, subfolder="text_encoder_2", torch_dtype=torch.float16)
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(base_path, subfolder="image_encoder", torch_dtype=torch.float16)
+vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae", torch_dtype=torch.float16)
+UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(base_path, subfolder="unet_encoder", torch_dtype=torch.float16)
 
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
-UNet_Encoder.requires_grad_(False)
-image_encoder.requires_grad_(False)
-vae.requires_grad_(False)
-unet.requires_grad_(False)
-text_encoder_one.requires_grad_(False)
-text_encoder_two.requires_grad_(False)
-tensor_transfrom = transforms.Compose(
-            [
-                transforms.ToTensor(),
-                transforms.Normalize([0.5], [0.5]),
-            ]
-    )
-
+# Préparation du pipeline Tryon
 pipe = TryonPipeline.from_pretrained(
-        base_path,
-        unet=unet,
-        vae=vae,
-        feature_extractor= CLIPImageProcessor(),
-        text_encoder = text_encoder_one,
-        text_encoder_2 = text_encoder_two,
-        tokenizer = tokenizer_one,
-        tokenizer_2 = tokenizer_two,
-        scheduler = noise_scheduler,
-        image_encoder=image_encoder,
-        torch_dtype=torch.float16,
-        force_download=False
+    base_path,
+    unet=unet,
+    vae=vae,
+    feature_extractor=CLIPImageProcessor(),
+    text_encoder=text_encoder_one,
+    text_encoder_2=text_encoder_two,
+    tokenizer=tokenizer_one,
+    tokenizer_2=tokenizer_two,
+    scheduler=noise_scheduler,
+    image_encoder=image_encoder,
+    torch_dtype=torch.float16,
+    force_download=False
 )
 pipe.unet_encoder = UNet_Encoder
 
+# Utilisation des transformations d'images
+tensor_transfrom = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize([0.5], [0.5]),
+])
+
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
     grayscale_image = Image.fromarray(np_image).convert("L")
     binary_mask = np.array(grayscale_image) > threshold
     mask = np.zeros(binary_mask.shape, dtype=np.uint8)
-    for i in range(binary_mask.shape[0]):
-        for j in range(binary_mask.shape[1]):
-            if binary_mask[i, j]:
-                mask[i, j] = 1
-    mask = (mask * 255).astype(np.uint8)
-    output_mask = Image.fromarray(mask)
-    return output_mask
+    mask[binary_mask] = 1
+    return Image.fromarray((mask * 255).astype(np.uint8))
+    
+
 
 def get_image_from_url(url):
     try:
@@ -157,8 +121,7 @@ def encode_image_to_base64(img):
     try:
         buffered = BytesIO()
         img.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return img_str
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
     except Exception as e:
         logging.error(f"Error encoding image: {e}")
         raise
@@ -198,7 +161,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
         mask, mask_gray = get_mask_location('hd', categorie , model_parse, keypoints)
         mask = mask.resize((768, 1024))
     else:
-        mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
+        mask = dict['layers'][0].convert("RGB").resize((768, 1024))#pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
     mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
     mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)
 
@@ -273,7 +236,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
         human_img_orig.paste(out_img, (int(left), int(top)))
         return human_img_orig, mask_gray
     else:
-        return images[0], mask_gray
+        return images[0], mask_gray , mask
 
         
 @app.route('/tryon-v2', methods=['POST'])
@@ -283,7 +246,6 @@ def tryon_v2():
     human_image_data = data['human_image']
     garment_image_data = data['garment_image']
 
-    # Process images (base64 ou URL)
     human_image = process_image(human_image_data)
     garment_image = process_image(garment_image_data)
 
@@ -294,20 +256,22 @@ def tryon_v2():
     seed = int(data.get('seed', random.randint(0, 9999999)))
     categorie = data.get('categorie', 'upper_body')
 
-        # Vérifie si 'mask_image' est présent dans les données
     mask_image = None
     if 'mask_image' in data:
         mask_image_data = data['mask_image']
         mask_image = process_image(mask_image_data)
-    
+
     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_image, mask_image , mask = start_tryon(human_dict, garment_image, description, use_auto_mask, use_auto_crop, denoise_steps, seed, categorie)
     return jsonify({
-        'image_id': save_image(output_image)
+        'image_id': save_image(output_image),
+        'mask_gray_id' : save_image(mask_image),
+        'mask_id' : save_image(mask)
     })
 
 def clear_gpu_memory():
@@ -350,6 +314,15 @@ def tryon():
     })
 
 
+# Route index
+@app.route('/', methods=['GET'])
+def get_image():
+
+    # Renvoyer l'image
+    try:
+        return 'Welcome to IDM VTON API'
+    except FileNotFoundError:
+        return jsonify({'error': 'Image not found'}), 404
 
 # Route pour récupérer l'image générée
 @app.route('/api/get_image/<image_id>', methods=['GET'])