Update app.py

app.py

@@ -1,10 +1,10 @@
 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
 import base64
-import io 
+import io
 import logging
 import gradio as gr
 import numpy as np
@@ -33,52 +33,48 @@ from torchvision.transforms.functional import to_pil_image
 
 app = Flask(__name__)
 
-# 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')
+base_path = 'yisol/IDM-VTON'
 
+# Téléchargez les fichiers nécessaires via huggingface_hub
+def download_model_files(base_path, filename):
+    return hf_hub_download(repo_id=base_path, filename=filename)
+
+# Téléchargement et chargement des fichiers de modèle
+unet_path = download_model_files(base_path, "unet/pytorch_model.bin")
 unet = UNet2DConditionModel.from_pretrained(
-    pretrained_model_name_or_path=os.path.dirname(unet_path),
-    torch_dtype=torch.float16,
-    force_download=False
+    unet_path,
+    torch_dtype=torch.float16
 )
 unet.requires_grad_(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')
+tokenizer_one_path = download_model_files(base_path, "tokenizer/config.json")
+tokenizer_one = AutoTokenizer.from_pretrained(tokenizer_one_path, use_fast=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)
+tokenizer_two_path = download_model_files(base_path, "tokenizer_2/config.json")
+tokenizer_two = AutoTokenizer.from_pretrained(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))
+noise_scheduler_path = download_model_files(base_path, "scheduler/scheduler_config.json")
+noise_scheduler = DDPMScheduler.from_pretrained(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_one_path = download_model_files(base_path, "text_encoder/pytorch_model.bin")
+text_encoder_one = CLIPTextModel.from_pretrained(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)
+text_encoder_two_path = download_model_files(base_path, "text_encoder_2/pytorch_model.bin")
+text_encoder_two = CLIPTextModelWithProjection.from_pretrained(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)
+image_encoder_path = download_model_files(base_path, "image_encoder/pytorch_model.bin")
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(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)
+vae_path = download_model_files(base_path, "vae/pytorch_model.bin")
+vae = AutoencoderKL.from_pretrained(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)
+unet_encoder_path = download_model_files(base_path, "unet_encoder/pytorch_model.bin")
+UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(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
+# Définir les modèles comme non entraînables
 UNet_Encoder.requires_grad_(False)
 image_encoder.requires_grad_(False)
 vae.requires_grad_(False)
@@ -86,17 +82,13 @@ 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]),
-    ]
-)
+# Autres transformations et pipeline
+tensor_transfrom = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize([0.5], [0.5]),
+])
 
-# Configuration du pipeline Tryon
 pipe = TryonPipeline.from_pretrained(
-    'yisol/IDM-VTON',
     unet=unet,
     vae=vae,
     feature_extractor=CLIPImageProcessor(),
@@ -106,11 +98,8 @@ pipe = TryonPipeline.from_pretrained(
     tokenizer_2=tokenizer_two,
     scheduler=noise_scheduler,
     image_encoder=image_encoder,
-    torch_dtype=torch.float16,
-    force_download=False
+    torch_dtype=torch.float16
 )
-
-# Ajout du UNet Encoder dans le pipeline
 pipe.unet_encoder = UNet_Encoder
 
 def pil_to_binary_mask(pil_image, threshold=0):
@@ -259,7 +248,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)))
@@ -295,7 +284,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)