Update app.py

app.py

@@ -36,59 +36,19 @@ app = Flask(__name__)
 base_path = 'yisol/IDM-VTON'
 example_path = os.path.join(os.path.dirname(__file__), 'example')
 
-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
-)
+def load_models():
+    global unet, tokenizer_one, tokenizer_two, noise_scheduler, text_encoder_one, text_encoder_two, image_encoder, vae, UNet_Encoder
+    unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet", torch_dtype=torch.float16, force_download=False)
+    tokenizer_one = AutoTokenizer.from_pretrained(base_path, subfolder="tokenizer", use_fast=False, force_download=False)
+    tokenizer_two = AutoTokenizer.from_pretrained(base_path, subfolder="tokenizer_2", 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)
+
+load_models()
 
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
@@ -169,6 +129,12 @@ def save_image(img):
     img.save(unique_name, format="WEBP", lossless=True) 
     return unique_name
 
+
+def clear_gpu_memory():
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+
+
 @spaces.GPU
 def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, categorie = 'upper_body'):
     device = "cuda"
@@ -176,7 +142,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
-    garm_img = garm_img.convert("RGB").resize((768, 1024))
+    garm_img = garm_img.convert("RGB").resize((512, 768))
     human_img_orig = dict["background"].convert("RGB")
 
     if is_checked_crop:
@@ -189,9 +155,9 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
         bottom = (height + target_height) / 2
         cropped_img = human_img_orig.crop((left, top, right, bottom))
         crop_size = cropped_img.size
-        human_img = cropped_img.resize((768, 1024))
+        human_img = cropped_img.resize((512, 768))
     else:
-        human_img = human_img_orig.resize((768, 1024))
+        human_img = human_img_orig.resize((512, 768))
 
     if is_checked:
         keypoints = openpose_model(human_img.resize((384, 512)))
@@ -199,7 +165,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 = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((512, 768)))
     mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
     mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)
 
@@ -209,7 +175,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
     args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
     pose_img = args.func(args, human_img_arg)
     pose_img = pose_img[:, :, ::-1]
-    pose_img = Image.fromarray(pose_img).resize((768, 1024))
+    pose_img = Image.fromarray(pose_img).resize((512, 768))
 
     with torch.no_grad():
         with torch.cuda.amp.autocast():
@@ -265,10 +231,10 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
                     image=human_img,
                     height=1024,
                     width=768,
-                    ip_adapter_image=garm_img.resize((768, 1024)),
+                    ip_adapter_image=garm_img.resize((512, 768)),
                     guidance_scale=2.0,
                 )[0]
-
+    clear_gpu_memory()
     if is_checked_crop:
         out_img = images[0].resize(crop_size)
         human_img_orig.paste(out_img, (int(left), int(top)))
@@ -277,10 +243,6 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
         return images[0], mask_gray
 
 
-def clear_gpu_memory():
-    torch.cuda.empty_cache()
-    torch.cuda.synchronize()
-
 def process_image(image_data):
     # Vérifie si l'image est en base64 ou URL
     if image_data.startswith('http://') or image_data.startswith('https://'):
@@ -304,7 +266,6 @@ def tryon():
         'layers': [human_image] if not use_auto_mask else None,
         'composite': None
     }
-    clear_gpu_memory()
 
     output_image, mask_image = start_tryon(human_dict, garment_image, description, use_auto_mask, use_auto_crop, denoise_steps, seed , categorie)