Update app.py

app.py

@@ -1,17 +1,17 @@
+
 import os
-from flask import Flask, request, jsonify, send_file
+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
 import spaces
 import uuid
 import random
-from huggingface_hub import hf_hub_download
 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
@@ -34,71 +34,91 @@ from torchvision.transforms.functional import to_pil_image
 app = Flask(__name__)
 
 base_path = 'yisol/IDM-VTON'
+example_path = os.path.join(os.path.dirname(__file__), 'example')
 
-# 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(
-    unet_path,
-    torch_dtype=torch.float16
+    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")
 
-tokenizer_one_path = download_model_files(base_path, "tokenizer/config.json")
-tokenizer_one = AutoTokenizer.from_pretrained(tokenizer_one_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)
-
-noise_scheduler_path = download_model_files(base_path, "scheduler/scheduler_config.json")
-noise_scheduler = DDPMScheduler.from_pretrained(noise_scheduler_path)
-
-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 = 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)
-
-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)
-
-vae_path = download_model_files(base_path, "vae/pytorch_model.bin")
-vae = AutoencoderKL.from_pretrained(vae_path, torch_dtype=torch.float16)
+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_path = download_model_files(base_path, "unet_encoder/pytorch_model.bin")
-UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(unet_encoder_path, torch_dtype=torch.float16)
+UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
+    base_path,
+    subfolder="unet_encoder",
+    torch_dtype=torch.float16,
+    force_download=False
+)
 
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
-# Définir les modèles comme non entraînables
 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)
-
-# Autres transformations et pipeline
-tensor_transfrom = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Normalize([0.5], [0.5]),
-])
+tensor_transfrom = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+    )
 
 pipe = TryonPipeline.from_pretrained(
-    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
+        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
 
@@ -284,7 +304,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)