EcoTry/IDM-VTON

1 2import spaces3import gradio as gr4from PIL import Image5from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline6from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref7from src.unet_hacked_tryon import UNet2DConditionModel8from transformers import (9    CLIPImageProcessor,10    CLIPVisionModelWithProjection,11    CLIPTextModel,12    CLIPTextModelWithProjection,13)14from diffusers import DDPMScheduler,AutoencoderKL15from typing import List16 17import torch18import os19from transformers import AutoTokenizer20import numpy as np21from utils_mask import get_mask_location22from torchvision import transforms23import apply_net24from preprocess.humanparsing.run_parsing import Parsing25from preprocess.openpose.run_openpose import OpenPose26from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation27from torchvision.transforms.functional import to_pil_image28 29 30def pil_to_binary_mask(pil_image, threshold=0):31    np_image = np.array(pil_image)32    grayscale_image = Image.fromarray(np_image).convert("L")33    binary_mask = np.array(grayscale_image) > threshold34    mask = np.zeros(binary_mask.shape, dtype=np.uint8)35    for i in range(binary_mask.shape[0]):36        for j in range(binary_mask.shape[1]):37            if binary_mask[i,j] == True :38                mask[i,j] = 139    mask = (mask*255).astype(np.uint8)40    output_mask = Image.fromarray(mask)41    return output_mask42 43 44base_path = 'yisol/IDM-VTON'45example_path = os.path.join(os.path.dirname(__file__), 'example')46 47unet = UNet2DConditionModel.from_pretrained(48    base_path,49    subfolder="unet",50    torch_dtype=torch.float16,51)52unet.requires_grad_(False)53tokenizer_one = AutoTokenizer.from_pretrained(54    base_path,55    subfolder="tokenizer",56    revision=None,57    use_fast=False,58)59tokenizer_two = AutoTokenizer.from_pretrained(60    base_path,61    subfolder="tokenizer_2",62    revision=None,63    use_fast=False,64)65noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")66 67text_encoder_one = CLIPTextModel.from_pretrained(68    base_path,69    subfolder="text_encoder",70    torch_dtype=torch.float16,71)72text_encoder_two = CLIPTextModelWithProjection.from_pretrained(73    base_path,74    subfolder="text_encoder_2",75    torch_dtype=torch.float16,76)77image_encoder = CLIPVisionModelWithProjection.from_pretrained(78    base_path,79    subfolder="image_encoder",80    torch_dtype=torch.float16,81    )82vae = AutoencoderKL.from_pretrained(base_path,83                                    subfolder="vae",84                                    torch_dtype=torch.float16,85)86 87# "stabilityai/stable-diffusion-xl-base-1.0",88UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(89    base_path,90    subfolder="unet_encoder",91    torch_dtype=torch.float16,92)93 94parsing_model = Parsing(0)95openpose_model = OpenPose(0)96 97UNet_Encoder.requires_grad_(False)98image_encoder.requires_grad_(False)99vae.requires_grad_(False)100unet.requires_grad_(False)101text_encoder_one.requires_grad_(False)102text_encoder_two.requires_grad_(False)103tensor_transfrom = transforms.Compose(104            [105                transforms.ToTensor(),106                transforms.Normalize([0.5], [0.5]),107            ]108    )109 110pipe = TryonPipeline.from_pretrained(111        base_path,112        unet=unet,113        vae=vae,114        feature_extractor= CLIPImageProcessor(),115        text_encoder = text_encoder_one,116        text_encoder_2 = text_encoder_two,117        tokenizer = tokenizer_one,118        tokenizer_2 = tokenizer_two,119        scheduler = noise_scheduler,120        image_encoder=image_encoder,121        torch_dtype=torch.float16,122)123pipe.unet_encoder = UNet_Encoder124 125@spaces.GPU126def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed):127    device = "cuda"128    129    openpose_model.preprocessor.body_estimation.model.to(device)130    pipe.to(device)131    pipe.unet_encoder.to(device)132 133    garm_img= garm_img.convert("RGB").resize((768,1024))134    human_img_orig = dict["background"].convert("RGB")    135    136    if is_checked_crop:137        width, height = human_img_orig.size138        target_width = int(min(width, height * (3 / 4)))139        target_height = int(min(height, width * (4 / 3)))140        left = (width - target_width) / 2141        top = (height - target_height) / 2142        right = (width + target_width) / 2143        bottom = (height + target_height) / 2144        cropped_img = human_img_orig.crop((left, top, right, bottom))145        crop_size = cropped_img.size146        human_img = cropped_img.resize((768,1024))147    else:148        human_img = human_img_orig.resize((768,1024))149 150 151    if is_checked:152        keypoints = openpose_model(human_img.resize((384,512)))153        model_parse, _ = parsing_model(human_img.resize((384,512)))154        mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints)155        mask = mask.resize((768,1024))156    else:157        mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))158        # mask = transforms.ToTensor()(mask)159        # mask = mask.unsqueeze(0)160    mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img)161    mask_gray = to_pil_image((mask_gray+1.0)/2.0)162 163 164    human_img_arg = _apply_exif_orientation(human_img.resize((384,512)))165    human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")166     167    168 169    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'))170    # verbosity = getattr(args, "verbosity", None)171    pose_img = args.func(args,human_img_arg)    172    pose_img = pose_img[:,:,::-1]    173    pose_img = Image.fromarray(pose_img).resize((768,1024))174    175    with torch.no_grad():176        # Extract the images177        with torch.cuda.amp.autocast():178            with torch.no_grad():179                prompt = "model is wearing " + garment_des180                negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"181                with torch.inference_mode():182                    (183                        prompt_embeds,184                        negative_prompt_embeds,185                        pooled_prompt_embeds,186                        negative_pooled_prompt_embeds,187                    ) = pipe.encode_prompt(188                        prompt,189                        num_images_per_prompt=1,190                        do_classifier_free_guidance=True,191                        negative_prompt=negative_prompt,192                    )193                                    194                    prompt = "a photo of " + garment_des195                    negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"196                    if not isinstance(prompt, List):197                        prompt = [prompt] * 1198                    if not isinstance(negative_prompt, List):199                        negative_prompt = [negative_prompt] * 1200                    with torch.inference_mode():201                        (202                            prompt_embeds_c,203                            _,204                            _,205                            _,206                        ) = pipe.encode_prompt(207                            prompt,208                            num_images_per_prompt=1,209                            do_classifier_free_guidance=False,210                            negative_prompt=negative_prompt,211                        )212 213 214 215                    pose_img =  tensor_transfrom(pose_img).unsqueeze(0).to(device,torch.float16)216                    garm_tensor =  tensor_transfrom(garm_img).unsqueeze(0).to(device,torch.float16)217                    generator = torch.Generator(device).manual_seed(seed) if seed is not None else None218                    images = pipe(219                        prompt_embeds=prompt_embeds.to(device,torch.float16),220                        negative_prompt_embeds=negative_prompt_embeds.to(device,torch.float16),221                        pooled_prompt_embeds=pooled_prompt_embeds.to(device,torch.float16),222                        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(device,torch.float16),223                        num_inference_steps=denoise_steps,224                        generator=generator,225                        strength = 1.0,226                        pose_img = pose_img.to(device,torch.float16),227                        text_embeds_cloth=prompt_embeds_c.to(device,torch.float16),228                        cloth = garm_tensor.to(device,torch.float16),229                        mask_image=mask,230                        image=human_img, 231                        height=1024,232                        width=768,233                        ip_adapter_image = garm_img.resize((768,1024)),234                        guidance_scale=2.0,235                    )[0]236 237    if is_checked_crop:238        out_img = images[0].resize(crop_size)        239        human_img_orig.paste(out_img, (int(left), int(top)))    240        return human_img_orig, mask_gray241    else:242        return images[0], mask_gray243    # return images[0], mask_gray244 245garm_list = os.listdir(os.path.join(example_path,"cloth"))246garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]247 248human_list = os.listdir(os.path.join(example_path,"human"))249human_list_path = [os.path.join(example_path,"human",human) for human in human_list]250 251human_ex_list = []252for ex_human in human_list_path:253    ex_dict= {}254    ex_dict['background'] = ex_human255    ex_dict['layers'] = None256    ex_dict['composite'] = None257    human_ex_list.append(ex_dict)258 259##default human260 261 262image_blocks = gr.Blocks().queue()263with image_blocks as demo:264    gr.Markdown("## IDM-VTON 👕👔👚")265    gr.Markdown("Virtual Try-on with your image and garment image. Check out the [source codes](https://github.com/yisol/IDM-VTON) and the [model](https://huggingface.co/yisol/IDM-VTON)")266    with gr.Row():267        with gr.Column():268            imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)269            with gr.Row():270                is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)271            with gr.Row():272                is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)273 274            example = gr.Examples(275                inputs=imgs,276                examples_per_page=10,277                examples=human_ex_list278            )279 280        with gr.Column():281            garm_img = gr.Image(label="Garment", sources='upload', type="pil")282            with gr.Row(elem_id="prompt-container"):283                with gr.Row():284                    prompt = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")285            example = gr.Examples(286                inputs=garm_img,287                examples_per_page=8,288                examples=garm_list_path)289        with gr.Column():290            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)291            masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)292        with gr.Column():293            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)294            image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False)295 296 297 298 299    with gr.Column():300        try_button = gr.Button(value="Try-on")301        with gr.Accordion(label="Advanced Settings", open=False):302            with gr.Row():303                denoise_steps = gr.Number(label="Denoising Steps", minimum=20, maximum=40, value=30, step=1)304                seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)305 306 307 308    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')309            310image_blocks.launch()311 312