app.py

import argparse
import os
import gc
import psutil
import threading
from pathlib import Path
import shutil
import time
import glob
from datetime import datetime
os.environ['CUDA_HOME'] = '/usr/local/cuda'
os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
from datetime import datetime
import cv2
import gradio as gr
import spaces
import numpy as np
import torch
from diffusers.image_processor import VaeImageProcessor
from huggingface_hub import snapshot_download
from PIL import Image

torch.jit.script = lambda f: f
from model.cloth_masker import AutoMasker, vis_mask
from model.pipeline import CatVTONPipeline, CatVTONPix2PixPipeline
from model.flux.pipeline_flux_tryon import FluxTryOnPipeline
from utils import init_weight_dtype, resize_and_crop, resize_and_padding


def parse_args():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument(
        "--base_model_path",
        type=str,
        default="booksforcharlie/stable-diffusion-inpainting",
        help=(
            "The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub."
        ),
    )
    parser.add_argument(
        "--p2p_base_model_path",
        type=str,
        default="timbrooks/instruct-pix2pix",
        help=(
            "The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub."
        ),
    )
    parser.add_argument(
        "--resume_path",
        type=str,
        default="zhengchong/CatVTON",
        help=(
            "The Path to the checkpoint of trained tryon model."
        ),
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="resource/demo/output",
        help="The output directory where the model predictions will be written.",
    )

    parser.add_argument(
        "--width",
        type=int,
        default=768,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    parser.add_argument(
        "--height",
        type=int,
        default=1024,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    parser.add_argument(
        "--repaint",
        action="store_true",
        help="Whether to repaint the result image with the original background."
    )
    parser.add_argument(
        "--allow_tf32",
        action="store_true",
        default=True,
        help=(
            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
        ),
    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default="bf16",
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
        ),
    )

    args = parser.parse_args()
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
    if env_local_rank != -1 and env_local_rank != args.local_rank:
        args.local_rank = env_local_rank

    return args


def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols

    w, h = imgs[0].size
    grid = Image.new("RGB", size=(cols * w, rows * h))

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid


args = parse_args()
OUTPUT_DIR = "generated_images"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Mask-based CatVTON
catvton_repo = "zhengchong/CatVTON"
repo_path = snapshot_download(repo_id=catvton_repo)
# Pipeline
pipeline = CatVTONPipeline(
    base_ckpt=args.base_model_path,
    attn_ckpt=repo_path,
    attn_ckpt_version="mix",
    weight_dtype=init_weight_dtype(args.mixed_precision),
    use_tf32=args.allow_tf32,
    device='cuda'
)
# AutoMasker
mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
automasker = AutoMasker(
    densepose_ckpt=os.path.join(repo_path, "DensePose"),
    schp_ckpt=os.path.join(repo_path, "SCHP"),
    device='cuda',
)

# Flux-based CatVTON
access_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
flux_repo = "black-forest-labs/FLUX.1-Fill-dev"
pipeline_flux = FluxTryOnPipeline.from_pretrained(flux_repo, use_auth_token=access_token)
pipeline_flux.load_lora_weights(
    os.path.join(repo_path, "flux-lora"),
    weight_name='pytorch_lora_weights.safetensors'
)
pipeline_flux.to("cuda", init_weight_dtype(args.mixed_precision))

def save_generated_image(image, frame_no):
    """Save generated image with timestamp and model name"""
    filename = f"{frame_no}_frame.png"
    filepath = os.path.join(OUTPUT_DIR, filename)
    image.save(filepath)
    return filepath
    
def print_image_info(img):
    # Basic attributes
    info = {
       
        "Format": img.format,
        "Mode": img.mode,
        "Size": img.size,
        "Width": img.width,
        "Height": img.height,
        "DPI": img.info.get('dpi', "N/A"),
        "Is Animated": getattr(img, "is_animated", False),
        "Frames": getattr(img, "n_frames", 1)
    }

    print("----- Image Information -----")
    for key, value in info.items():
        print(f"{key}: {value}")

def extract_frames(video_path):
    if not os.path.exists(video_path):
        print("Video file does not exist:", video_path)
        return None
    # Open the video file
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video file {video_path}")
        return []
    frames = []
    success, frame = cap.read()
    print(f"cap read status {success}")
    while success:
        print("getting frame")
        # Convert frame from BGR (OpenCV default) to RGB
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # Convert the numpy array (frame) to a PIL Image
        pil_frame = Image.fromarray(frame_rgb)
        frames.append(pil_frame)
        success, frame = cap.read()
    cap.release()
    return frames

#process_video_frames
@spaces.GPU(duration=175)
def process_video_frames(
        video,
        cloth_image,
        cloth_type,
        num_inference_steps,
        guidance_scale,
        seed,
        show_type
):
    """
    Process each frame of the video through the flux pipeline
    
    Args:
        video (str): Path to the input video file
        cloth_image (str): Path to the cloth image
        ... (other parameters from original function)
    
    Returns:
        list: Processed frames
    """
    # Extract frames from video
    frames = extract_frames(video)
    
    processed_frames = []
    print(f"processed_frames {len(frames)}")
    for index, person_image in enumerate(frames):
        result_image = proc_function_vidfl(
            person_image,
            cloth_image,
        cloth_type,
        num_inference_steps,
        guidance_scale,
        seed,
        show_type
        )
        print_image_info(result_image)
        save_generated_image(result_image,index)
        gallery_images = update_gallery()
        processed_frames.append(result_image)
        print("YEILEDING process_video_frames")
        yield result_image,gallery_images
    gallery_images = update_gallery()
    yield processed_frames, gallery_images


@spaces.GPU(duration=175)
def proc_function_vidfl(
        person_image,
        cloth_image,
        cloth_type,
        num_inference_steps,
        guidance_scale,
        seed,
        show_type
):

    print_image_info(person_image)
    # Set random seed
    generator = None
    if seed != -1:
        generator = torch.Generator(device='cuda').manual_seed(seed)

    # Process input images
    #person_image = Image.open(person_image).convert("RGB")
    cloth_image = Image.open(cloth_image).convert("RGB")

    # Adjust image sizes
    person_image = resize_and_crop(person_image, (args.width, args.height))
    cloth_image = resize_and_padding(cloth_image, (args.width, args.height))

    # Process mask

    mask = automasker(
        person_image,
        cloth_type
    )['mask']
    mask = mask_processor.blur(mask, blur_factor=9)

    # Inference
    result_image = pipeline_flux(
        image=person_image,
        condition_image=cloth_image,
        mask_image=mask,
        width=args.width,
        height=args.height,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator
    ).images[0]
    print("YEILEDING proc_function_vidfl")
    return result_image
        
@spaces.GPU(duration=175)
def submit_function_flux(
        person_image,
        cloth_image,
        cloth_type,
        num_inference_steps,
        guidance_scale,
        seed,
        show_type
):
    # Process image editor input
    person_image, mask = person_image["background"], person_image["layers"][0]
    mask = Image.open(mask).convert("L")
    if len(np.unique(np.array(mask))) == 1:
        mask = None
    else:
        mask = np.array(mask)
        mask[mask > 0] = 255
        mask = Image.fromarray(mask)

    # Set random seed
    generator = None
    if seed != -1:
        generator = torch.Generator(device='cuda').manual_seed(seed)

    # Process input images
    person_image = Image.open(person_image).convert("RGB")
    cloth_image = Image.open(cloth_image).convert("RGB")

    # Adjust image sizes
    person_image = resize_and_crop(person_image, (args.width, args.height))
    cloth_image = resize_and_padding(cloth_image, (args.width, args.height))

    # Process mask
    if mask is not None:
        mask = resize_and_crop(mask, (args.width, args.height))
    else:
        mask = automasker(
            person_image,
            cloth_type
        )['mask']
    mask = mask_processor.blur(mask, blur_factor=9)

    # Inference
    result_image = pipeline_flux(
        image=person_image,
        condition_image=cloth_image,
        mask_image=mask,
        width=args.width,
        height=args.height,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator
    ).images[0]

    # Post-processing
    masked_person = vis_mask(person_image, mask)

    # Return result based on show type
    if show_type == "result only":
        return result_image
    else:
        width, height = person_image.size
        if show_type == "input & result":
            condition_width = width // 2
            conditions = image_grid([person_image, cloth_image], 2, 1)
        else:
            condition_width = width // 3
            conditions = image_grid([person_image, masked_person, cloth_image], 3, 1)

        conditions = conditions.resize((condition_width, height), Image.NEAREST)
        new_result_image = Image.new("RGB", (width + condition_width + 5, height))
        new_result_image.paste(conditions, (0, 0))
        new_result_image.paste(result_image, (condition_width + 5, 0))
        gallery_images = update_gallery()
        return new_result_image, gallery_images


def person_example_fn(image_path):
    return image_path

def get_generated_images():
    """Get list of generated images with their details"""
    files = glob.glob(os.path.join(OUTPUT_DIR, "*.png"))
    files.sort(key=os.path.getctime, reverse=True)  # Sort by creation time
    return [
        {
            "path": f,
            "name": os.path.basename(f),
            "date": datetime.fromtimestamp(os.path.getctime(f)).strftime("%Y-%m-%d %H:%M:%S"),
            "size": f"{os.path.getsize(f) / 1024:.1f} KB"
        }
        for f in files
    ]

def update_gallery():
        """Update the file gallery"""
        files = get_generated_images()
        return [
            (f["path"], f"{f['name']}\n{f['date']}")
            for f in files
        ]
    
HEADER = """
<h1 style="text-align: center;"> 🐈 CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models </h1>

<br>
· This demo and our weights are only for Non-commercial Use. <br>
· Thanks to <a href="https://huggingface.co/zero-gpu-explorers">ZeroGPU</a> for providing A100 for our <a href="https://huggingface.co/spaces/zhengchong/CatVTON">HuggingFace Space</a>. <br>
· SafetyChecker is set to filter NSFW content, but it may block normal results too. Please adjust the <span>`seed`</span> for normal outcomes.<br> 
"""


def app_gradio():
    with gr.Blocks(title="CatVTON") as demo:
        gr.Markdown(HEADER)

        with gr.Tab("Mask-based & Flux.1 Fill Dev"):
            with gr.Row():
                with gr.Column(scale=1, min_width=350):
                    with gr.Row():
                        image_path_flux = gr.Image(
                            type="filepath",
                            interactive=True,
                            visible=False,
                        )
                        person_image_flux = gr.ImageEditor(
                            interactive=True, label="Person Image", type="filepath"
                        )

                    with gr.Row():
                        with gr.Column(scale=1, min_width=230):
                            cloth_image_flux = gr.Image(
                                interactive=True, label="Condition Image", type="filepath"
                            )
                        with gr.Column(scale=1, min_width=120):
                            gr.Markdown(
                                '<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `🖌️` above to draw the Mask (higher priority)<br>2. Select the `Try-On Cloth Type` to generate automatically </span>'
                            )
                            cloth_type = gr.Radio(
                                label="Try-On Cloth Type",
                                choices=["upper", "lower", "overall"],
                                value="upper",
                            )

                    submit_flux = gr.Button("Submit")
                    gr.Markdown(
                        '<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>'
                    )

                    with gr.Accordion("Advanced Options", open=False):
                        num_inference_steps_flux = gr.Slider(
                            label="Inference Step", minimum=10, maximum=100, step=5, value=50
                        )
                        # Guidence Scale
                        guidance_scale_flux = gr.Slider(
                            label="CFG Strenth", minimum=0.0, maximum=50, step=0.5, value=30
                        )
                        # Random Seed
                        seed_flux = gr.Slider(
                            label="Seed", minimum=-1, maximum=10000, step=1, value=42
                        )
                        show_type = gr.Radio(
                            label="Show Type",
                            choices=["result only", "input & result", "input & mask & result"],
                            value="input & mask & result",
                        )

                with gr.Column(scale=2, min_width=500):
                    result_image_flux = gr.Image(interactive=False, label="Result")
                    with gr.Row():
                        # Photo Examples
                        root_path = "resource/demo/example"
                        with gr.Column():
                            gal_output = gr.Gallery(label="Processed Frames")
                        

                image_path_flux.change(
                    person_example_fn, inputs=image_path_flux, outputs=person_image_flux
                )

                submit_flux.click(
                    submit_function_flux,
                    [person_image_flux, cloth_image_flux, cloth_type, num_inference_steps_flux, guidance_scale_flux,
                     seed_flux, show_type],
                    [result_image_flux,gal_output]
                )

        with gr.Tab("Video Flux"):
            with gr.Row():
                with gr.Column(scale=1, min_width=350):
                    with gr.Row():
                        image_path_vidflux = gr.Image(
                            type="filepath",
                            interactive=True,
                            visible=False,
                        )
                        person_image_vidflux = gr.Video(
                            
                        )

                    with gr.Row():
                        with gr.Column(scale=1, min_width=230):
                            cloth_image_vidflux = gr.Image(
                                interactive=True, label="Condition Image", type="filepath"
                            )
                        with gr.Column(scale=1, min_width=120):
                            gr.Markdown(
                                '<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `🖌️` above to draw the Mask (higher priority)<br>2. Select the `Try-On Cloth Type` to generate automatically </span>'
                            )
                            cloth_type = gr.Radio(
                                label="Try-On Cloth Type",
                                choices=["upper", "lower", "overall"],
                                value="upper",
                            )

                    submit_flux = gr.Button("Submit")
                    gr.Markdown(
                        '<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>'
                    )

                    with gr.Accordion("Advanced Options", open=False):
                        num_inference_steps_vidflux = gr.Slider(
                            label="Inference Step", minimum=10, maximum=100, step=5, value=50
                        )
                        # Guidence Scale
                        guidance_scale_vidflux = gr.Slider(
                            label="CFG Strenth", minimum=0.0, maximum=50, step=0.5, value=30
                        )
                        # Random Seed
                        seed_vidflux = gr.Slider(
                            label="Seed", minimum=-1, maximum=10000, step=1, value=42
                        )
                        show_type = gr.Radio(
                            label="Show Type",
                            choices=["result only", "input & result", "input & mask & result"],
                            value="input & mask & result",
                        )

                with gr.Column(scale=2, min_width=500):
                    result_image_vidflux = gr.Image(interactive=False, label="Result")
                    with gr.Row():
                        # Photo Examples
                        root_path = "resource/demo/example"
                        with gr.Column():
                            gal_output = gr.Gallery(
                        label="Generated Images",
                        show_label=True,
                        elem_id="gal_output",
                        columns=3,
                        height=800,
                        visible=True
                    )
                            refresh_button = gr.Button("Refresh Gallery")
                        

                image_path_vidflux.change(
                    person_example_fn, inputs=image_path_vidflux, outputs=person_image_vidflux
                )

                refresh_button.click(
        fn=update_gallery,
        inputs=[],
        outputs=[gal_output],
    )

                submit_flux.click(
                    process_video_frames,
                    [person_image_vidflux, cloth_image_vidflux, cloth_type, num_inference_steps_vidflux, guidance_scale_vidflux,
                     seed_vidflux, show_type],
                    [result_image_vidflux,gal_output]
                )


    demo.queue().launch(share=True, show_error=True)


if __name__ == "__main__":
    app_gradio()