Update app.py

app.py

@@ -32,107 +32,70 @@ 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,
-    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
-)
+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)
 
-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]),
-            ]
-    )
-
+# Préparation 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
+    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
 
+# Utilisation des transformations d'images
+tensor_transfrom = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize([0.5], [0.5]),
+])
+
 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)
-    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
+    mask[binary_mask] = 1
+    return Image.fromarray((mask * 255).astype(np.uint8))
+    
+
 
 def get_image_from_url(url):
     try:
@@ -157,8 +120,7 @@ def encode_image_to_base64(img):
     try:
         buffered = BytesIO()
         img.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return img_str
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
     except Exception as e:
         logging.error(f"Error encoding image: {e}")
         raise
@@ -283,7 +245,6 @@ def tryon_v2():
     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)
 
@@ -294,18 +255,18 @@ 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
     }
-    output_image, mask_image = start_tryon(human_dict, garment_image, description, use_auto_mask, use_auto_crop, denoise_steps, seed , categorie)
+
+    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)
     })