Update app.py

app.py

@@ -1,6 +1,6 @@
 
 import os
-from flask import Flask, request, jsonify,send_file
+from flask import Flask, request, jsonify
 from PIL import Image
 from io import BytesIO
 import torch
@@ -34,95 +34,84 @@ from torchvision.transforms.functional import to_pil_image
 
 app = Flask(__name__)
 
-#base_path = 'yisol/IDM-VTON'
-base_path = huggingface_hub.hf_hub_download("yisol/IDM-VTON")
-
-example_path = os.path.join(os.path.dirname(__file__), 'example')
+# Téléchargement du modèle UNet
+unet_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='unet', filename='pytorch_model.bin')
+unet_config_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='unet', filename='config.json')
 
 unet = UNet2DConditionModel.from_pretrained(
-    base_path,
-    subfolder="unet",
+    pretrained_model_name_or_path=os.path.dirname(unet_path),
     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
-)
+# Téléchargement des tokenizers
+tokenizer_one_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='tokenizer', filename='tokenizer.json')
+tokenizer_two_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='tokenizer_2', filename='tokenizer.json')
 
-UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
-    base_path,
-    subfolder="unet_encoder",
-    torch_dtype=torch.float16,
-    force_download=False
-)
+tokenizer_one = AutoTokenizer.from_pretrained(os.path.dirname(tokenizer_one_path), use_fast=False)
+tokenizer_two = AutoTokenizer.from_pretrained(os.path.dirname(tokenizer_two_path), use_fast=False)
+
+# Téléchargement du noise scheduler
+noise_scheduler_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='scheduler', filename='scheduler.json')
+noise_scheduler = DDPMScheduler.from_pretrained(os.path.dirname(noise_scheduler_path))
+
+# Téléchargement et chargement des text encoders
+text_encoder_one_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='text_encoder', filename='pytorch_model.bin')
+text_encoder_one = CLIPTextModel.from_pretrained(os.path.dirname(text_encoder_one_path), torch_dtype=torch.float16)
+
+text_encoder_two_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='text_encoder_2', filename='pytorch_model.bin')
+text_encoder_two = CLIPTextModelWithProjection.from_pretrained(os.path.dirname(text_encoder_two_path), torch_dtype=torch.float16)
+
+# Téléchargement et chargement de l'image encoder
+image_encoder_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='image_encoder', filename='pytorch_model.bin')
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(os.path.dirname(image_encoder_path), torch_dtype=torch.float16)
+
+# Téléchargement et chargement du VAE
+vae_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='vae', filename='pytorch_model.bin')
+vae = AutoencoderKL.from_pretrained(os.path.dirname(vae_path), torch_dtype=torch.float16)
 
+# Téléchargement et chargement de l'UNet Encoder
+unet_encoder_path = hf_hub_download(repo_id='yisol/IDM-VTON', subfolder='unet_encoder', filename='pytorch_model.bin')
+UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(os.path.dirname(unet_encoder_path), torch_dtype=torch.float16)
+
+# Initialisation des autres modèles (parsing et openpose)
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
+# Désactivation des gradients
 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)
+
+# Transformation des tensors
 tensor_transfrom = transforms.Compose(
-            [
-                transforms.ToTensor(),
-                transforms.Normalize([0.5], [0.5]),
-            ]
-    )
+    [
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+)
 
+# Configuration 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
+    'yisol/IDM-VTON',
+    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
 )
+
+# Ajout du UNet Encoder dans le pipeline
 pipe.unet_encoder = UNet_Encoder
 
 def pil_to_binary_mask(pil_image, threshold=0):
@@ -271,7 +260,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
                     ip_adapter_image=garm_img.resize((768, 1024)),
                     guidance_scale=2.0,
                 )[0]
-
+    del prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, generator, images
     if is_checked_crop:
         out_img = images[0].resize(crop_size)
         human_img_orig.paste(out_img, (int(left), int(top)))
@@ -307,7 +296,7 @@ def tryon():
         'layers': [human_image] if not use_auto_mask else None,
         'composite': None
     }
-    #clear_gpu_memory()
+    clear_gpu_memory()
 
     output_image, mask_image = start_tryon(human_dict, garment_image, description, use_auto_mask, use_auto_crop, denoise_steps, seed , categorie)