app.py

import os
from flask import Flask, request, jsonify
from PIL import Image
from io import BytesIO
import base64
import torch
import requests
import numpy as np
import uuid
from transformers import (
    CLIPImageProcessor,
    CLIPVisionModelWithProjection,
    CLIPTextModel,
    CLIPTextModelWithProjection,
    AutoTokenizer
)
from diffusers import DDPMScheduler, AutoencoderKL
from utils_mask import get_mask_location
from torchvision import transforms
import apply_net
from preprocess.humanparsing.run_parsing import Parsing
from preprocess.openpose.run_openpose import OpenPose
from detectron2.data.detection_utils import convert_PIL_to_numpy, _apply_exif_orientation
from torchvision.transforms.functional import to_pil_image

app = Flask(__name__)

# Variables globales pour stocker les modèles
models_loaded = False

def load_models():
    global unet, tokenizer_one, tokenizer_two, noise_scheduler, text_encoder_one, text_encoder_two
    global image_encoder, vae, UNet_Encoder, parsing_model, openpose_model, pipe
    
    if not models_loaded:
        base_path = 'yisol/IDM-VTON'
        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", 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)
        
        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])])
        
        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
        )
        pipe.unet_encoder = UNet_Encoder
        
        global models_loaded
        models_loaded = True

def pil_to_binary_mask(pil_image, threshold=0):
    np_image = np.array(pil_image.convert("L"))  # Convert to grayscale directly
    binary_mask = np_image > threshold
    mask = np.uint8(binary_mask * 255)
    return Image.fromarray(mask)

def get_image_from_url(url):
    try:
        response = requests.get(url)
        response.raise_for_status()
        return Image.open(BytesIO(response.content))
    except Exception as e:
        logging.error(f"Error fetching image from URL: {e}")
        raise

def decode_image_from_base64(base64_str):
    try:
        img_data = base64.b64decode(base64_str)
        return Image.open(BytesIO(img_data))
    except Exception as e:
        logging.error(f"Error decoding image: {e}")
        raise

def encode_image_to_base64(img):
    try:
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")
    except Exception as e:
        logging.error(f"Error encoding image: {e}")
        raise

def save_image(img):
    unique_name = f"{uuid.uuid4()}.webp"
    img.save(unique_name, format="WEBP", lossless=True)
    return unique_name

def clear_gpu_memory():
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

@spaces.GPU
def start_tryon(human_dict, garment_image, garment_description, use_auto_mask, use_auto_crop, denoise_steps, seed, category='upper_body'):
    device = "cuda"
    openpose_model.preprocessor.body_estimation.model.to(device)
    pipe.to(device)
    pipe.unet_encoder.to(device)

    garment_image = garment_image.convert("RGB").resize((768, 1024))
    human_image_orig = human_dict["background"].convert("RGB")

    if use_auto_crop:
        width, height = human_image_orig.size
        target_width = int(min(width, height * (3 / 4)))
        target_height = int(min(height, width * (4 / 3)))
        left, top = (width - target_width) / 2, (height - target_height) / 2
        right, bottom = (width + target_width) / 2, (height + target_height) / 2
        cropped_img = human_image_orig.crop((left, top, right, bottom)).resize((768, 1024))
    else:
        cropped_img = human_image_orig.resize((768, 1024))

    if use_auto_mask:
        keypoints = openpose_model(cropped_img.resize((384, 512)))
        model_parse, _ = parsing_model(cropped_img.resize((384, 512)))
        mask, mask_gray = get_mask_location('hd', category, model_parse, keypoints)
        mask = mask.resize((768, 1024))
    else:
        mask = pil_to_binary_mask(human_dict['layers'][0].convert("RGB").resize((768, 1024)))
    
    mask_gray = (1 - transforms.ToTensor()(mask)) * transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])(cropped_img)
    mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)

    human_image_arg = _apply_exif_orientation(cropped_img.resize((384, 512)))
    human_image_arg = convert_PIL_to_numpy(human_image_arg, format="BGR")

    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_image = args.func(args, human_image_arg)
    pose_image = Image.fromarray(pose_image[:, :, ::-1]).resize((768, 1024))

    with torch.no_grad(), torch.cuda.amp.autocast():
        prompt = "model is wearing " + garment_description
        negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
        prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
            prompt, num_images_per_prompt=1, do_classifier_free_guidance=True, negative_prompt=negative_prompt
        )

        prompt_c = "a photo of " + garment_description
        negative_prompt_c = "monochrome, lowres, bad anatomy, worst quality, low quality"
        prompt_embeds_c, _, _, _ = pipe.encode_prompt(
            prompt_c, num_images_per_prompt=1, do_classifier_free_guidance=False, negative_prompt=negative_prompt_c
        )

        pose_image = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])(pose_image).unsqueeze(0).to(device, torch.float16)
        garment_tensor = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])(garment_image).unsqueeze(0).to(device, torch.float16)

        images = pipe(
            prompt_embeds=prompt_embeds.to(device, torch.float16),
            negative_prompt_embeds=negative_prompt_embeds.to(device, torch.float16),
            pose_image=pose_image,
            garment_image=garment_tensor,
            mask_image=mask_gray.to(device, torch.float16),
            generator=torch.Generator(device).manual_seed(seed),
            num_inference_steps=denoise_steps
        ).images

    if images:
        output_image = images[0]
        output_base64 = encode_image_to_base64(output_image)
        mask_image = mask
        mask_base64 = encode_image_to_base64(mask_image)
        return output_image, mask_image
    else:
        raise ValueError("Failed to generate image")


# Route pour récupérer l'image générée
@app.route('/api/get_image/<image_id>', methods=['GET'])
def get_image(image_id):
    # Construire le chemin complet de l'image
    image_path = image_id  # Assurez-vous que le nom de fichier correspond à celui que vous avez utilisé lors de la sauvegarde

    # Renvoyer l'image
    try:
        return send_file(image_path, mimetype='image/webp')
    except FileNotFoundError:
        return jsonify({'error': 'Image not found'}), 404

@app.route('/tryon', methods=['POST'])
def tryon_handler():
    try:
        data = request.json
        human_image = decode_image_from_base64(data['human_image'])
        garment_image = decode_image_from_base64(data['garment_image'])
        description = data.get('description')
        use_auto_mask = data.get('use_auto_mask', True)
        use_auto_crop = data.get('use_auto_crop', False)
        denoise_steps = int(data.get('denoise_steps', 30))
        seed = int(data.get('seed', 42))
        category = data.get('category', 'upper_body')
        
        human_dict = {
            'background': human_image,
            '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, category
        )

        output_base64 = encode_image_to_base64(output_image)
        mask_base64 = encode_image_to_base64(mask_image)

        return jsonify({
            'output_image': output_base64,
            'mask_image': mask_base64
        })
    except Exception as e:
        logging.error(f"Error in tryon_handler: {e}")
        return jsonify({'error': str(e)}), 500

if __name__ == "__main__":
    load_models()  # Charge les modèles au démarrage
    app.run(host='0.0.0.0', port=7860)