Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import requests
 from flask import Flask, request, jsonify,send_file
 from PIL import Image
 from io import BytesIO
@@ -33,70 +32,107 @@ 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)
-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)
+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
+)
 
-# Préparation du pipeline Tryon
-pipe = TryonPipeline.from_pretrained(
+UNet_Encoder = UNet2DConditionModel_ref.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,
+    subfolder="unet_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]),
-])
+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]),
+            ]
+    )
+
+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
+)
+pipe.unet_encoder = UNet_Encoder
 
 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)
-    mask[binary_mask] = 1
-    return Image.fromarray((mask * 255).astype(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
 
 def get_image_from_url(url):
     try:
@@ -121,7 +157,8 @@ def encode_image_to_base64(img):
     try:
         buffered = BytesIO()
         img.save(buffered, format="PNG")
-        return base64.b64encode(buffered.getvalue()).decode("utf-8")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        return img_str
     except Exception as e:
         logging.error(f"Error encoding image: {e}")
         raise
@@ -236,16 +273,17 @@ 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 , mask
+        return images[0], mask_gray
 
         
 @app.route('/tryon-v2', methods=['POST'])
 def tryon_v2():
 
-   data = request.json
+    data = request.json
     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)
 
@@ -256,34 +294,21 @@ 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
     }
-
-    try:
-        # Utiliser torch.no_grad() pour économiser de la mémoire
-        with torch.no_grad():
-            output_image, mask_image, mask = start_tryon(human_dict, garment_image, description, use_auto_mask, use_auto_crop, denoise_steps, seed, categorie)
-        
-        # Vider la mémoire GPU après traitement
-        clear_gpu_memory()
-
-        return jsonify({
-            'image_id': save_image(output_image),
-            'mask_gray_id': save_image(mask_image),
-            'mask_id': save_image(mask)
-        })
-
-    except Exception as e:
-        logging.error(f"Error during tryon process: {e}")
-        return jsonify({'error': 'An error occurred during tryon process.'}), 500
+    output_image, mask_image = 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)
+    })
 
 def clear_gpu_memory():
     torch.cuda.empty_cache()