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 = """
๐ CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models
ยท This demo and our weights are only for Non-commercial Use.
ยท Thanks to ZeroGPU for providing A100 for our HuggingFace Space.
ยท SafetyChecker is set to filter NSFW content, but it may block normal results too. Please adjust the `seed` for normal outcomes.
"""
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(
'Two ways to provide Mask:
1. Upload the person image and use the `๐๏ธ` above to draw the Mask (higher priority)
2. Select the `Try-On Cloth Type` to generate automatically '
)
cloth_type = gr.Radio(
label="Try-On Cloth Type",
choices=["upper", "lower", "overall"],
value="upper",
)
submit_flux = gr.Button("Submit")
gr.Markdown(
'!!! Click only Once, Wait for Delay !!!'
)
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(
'Two ways to provide Mask:
1. Upload the person image and use the `๐๏ธ` above to draw the Mask (higher priority)
2. Select the `Try-On Cloth Type` to generate automatically '
)
cloth_type = gr.Radio(
label="Try-On Cloth Type",
choices=["upper", "lower", "overall"],
value="upper",
)
submit_flux = gr.Button("Submit")
gr.Markdown(
'!!! Click only Once, Wait for Delay !!!'
)
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()