Spaces:

bhagu75
/

comfyui-zero-gpu

Sleeping

Final CatVTON ZeroGPU app

60f2208 verified about 1 month ago

8.45 kB

import os
import sys
import subprocess
from pathlib import Path

import gradio as gr
import spaces
import torch
from PIL import Image
from diffusers.image_processor import VaeImageProcessor
from huggingface_hub import snapshot_download

# ============================================================
# CatVTON ZeroGPU — stable Gradio app
# Based on the user's ComfyUI workflow:
#
# LoadCatVTONPipeline_AMD
# Person image       -> target_image
# Garment image      -> refer_image
# Clothing mask      -> mask_image
# ImageInvert        -> represented by the mask polarity note below
# CatVTON_AMD        -> result
#
# IMPORTANT:
# The uploaded UI mask should be:
#   WHITE = clothing / replace area
#   BLACK = background / keep area
# Therefore this version does NOT invert the uploaded mask.
# ============================================================

ROOT = Path("/tmp/catvton_src")
ROOT.mkdir(parents=True, exist_ok=True)

BASE_MODEL = os.getenv(
    "CATVTON_BASE_MODEL",
    "booksforcharlie/stable-diffusion-inpainting",
)
CATVTON_MODEL = os.getenv(
    "CATVTON_MODEL",
    "zhengchong/CatVTON",
)

WIDTH = int(os.getenv("CATVTON_WIDTH", "768"))
HEIGHT = int(os.getenv("CATVTON_HEIGHT", "1024"))


def ensure_catvton_source():
    marker = ROOT / "model" / "pipeline.py"

    if marker.exists():
        return

    if (ROOT / ".git").exists():
        subprocess.run(
            ["git", "-C", str(ROOT), "pull", "--ff-only"],
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        return

    subprocess.run(
        [
            "git",
            "clone",
            "--depth",
            "1",
            "--branch",
            "edited",
            "https://github.com/Zheng-Chong/CatVTON.git",
            str(ROOT),
        ],
        check=True,
    )


ensure_catvton_source()
sys.path.insert(0, str(ROOT))

from model.pipeline import CatVTONPipeline
from utils import resize_and_crop, resize_and_padding


# ============================================================
# MODEL
# ============================================================

print("Downloading CatVTON checkpoint...")
catvton_repo = snapshot_download(repo_id=CATVTON_MODEL)

print("Loading CatVTON pipeline...")

pipe = CatVTONPipeline(
    base_ckpt=BASE_MODEL,
    attn_ckpt=catvton_repo,
    attn_ckpt_version="mix",
    weight_dtype=torch.float16,
    use_tf32=True,
    device="cuda",
    # The upstream safety checker can produce false positives on
    # normal fully-clothed fashion images. The user's successful
    # test required this to be disabled.
    skip_safety_check=True,
)

mask_processor = VaeImageProcessor(
    vae_scale_factor=8,
    do_normalize=False,
    do_binarize=True,
    do_convert_grayscale=True,
)


# ============================================================
# IMAGE HELPERS
# ============================================================

def load_rgb(value):
    if value is None:
        return None

    if isinstance(value, Image.Image):
        return value.convert("RGB")

    return Image.open(value).convert("RGB")


def load_mask(value):
    if value is None:
        return None

    if isinstance(value, Image.Image):
        return value.convert("L")

    return Image.open(value).convert("L")


def prepare_inputs(person_file, garment_file, mask_file):
    person = load_rgb(person_file)
    garment = load_rgb(garment_file)
    mask = load_mask(mask_file)

    # Same target/reference preprocessing used by the working app.
    person = resize_and_crop(person, (WIDTH, HEIGHT))
    garment = resize_and_padding(garment, (WIDTH, HEIGHT))

    # IMPORTANT:
    # The UI mask is already:
    #   white = clothing/replacement region
    #   black = preserved region
    #
    # Do not invert it again.
    mask = resize_and_crop(mask, (WIDTH, HEIGHT))

    # Light mask edge smoothing for cleaner garment boundaries.
    mask = mask_processor.blur(mask, blur_factor=9)

    return person, garment, mask


# ============================================================
# GENERATION
# ============================================================

@spaces.GPU(duration=120)
def generate_outfit(person_file, garment_file, mask_file, steps, cfg, seed):
    if person_file is None:
        raise gr.Error("Please upload the person image.")

    if garment_file is None:
        raise gr.Error("Please upload the garment/clothing image.")

    if mask_file is None:
        raise gr.Error(
            "Please upload the clothing mask. "
            "White = clothing area, Black = background."
        )

    person, garment, mask = prepare_inputs(
        person_file,
        garment_file,
        mask_file,
    )

    steps = int(steps)
    cfg = float(cfg)
    seed = int(seed)

    generator = None

    # -1 means random seed.
    if seed >= 0:
        generator = torch.Generator(device="cuda").manual_seed(seed)

    print(
        f"CatVTON inference | steps={steps} "
        f"cfg={cfg} seed={seed}"
    )

    with torch.inference_mode():
        result = pipe(
            image=person,
            condition_image=garment,
            mask=mask,
            num_inference_steps=steps,
            guidance_scale=cfg,
            height=HEIGHT,
            width=WIDTH,
            generator=generator,
        )[0]

    return result


# ============================================================
# UI
# ============================================================

DESCRIPTION = """
# 🐈 CatVTON Outfit Try-On
Upload a **person photo**, a **garment/product photo**, and a **clothing mask**.
**Mask format:** ⚪ White = replace clothing area · ⚫ Black = preserve area
The default settings match the tested workflow: **30 steps / CFG 2.5**.
"""

with gr.Blocks(title="CatVTON Outfit Try-On — ZeroGPU") as demo:

    gr.Markdown(DESCRIPTION)

    with gr.Row():

        # ----------------------------------------------------
        # INPUTS
        # ----------------------------------------------------
        with gr.Column():

            person = gr.Image(
                type="filepath",
                label="👤 Person Image",
            )

            garment = gr.Image(
                type="filepath",
                label="👕 Garment / Clothing Image",
            )

            mask = gr.Image(
                type="filepath",
                label="🎭 Clothing Mask",
            )

            gr.Markdown(
                """
**Mask example**
- White shirt shape = area CatVTON replaces
- Black background = area CatVTON preserves
"""
            )

        # ----------------------------------------------------
        # OUTPUT
        # ----------------------------------------------------
        with gr.Column():

            result = gr.Image(
                type="pil",
                label="✨ Generated Result",
            )

            gr.Markdown(
                "Use the image toolbar to download the generated result."
            )

    # --------------------------------------------------------
    # SETTINGS
    # --------------------------------------------------------

    with gr.Row():

        steps = gr.Slider(
            minimum=10,
            maximum=50,
            step=1,
            value=30,
            label="Inference Steps",
        )

        cfg = gr.Slider(
            minimum=1.0,
            maximum=5.0,
            step=0.1,
            value=2.5,
            label="CFG",
        )

        seed = gr.Number(
            value=-1,
            precision=0,
            label="Seed (-1 = Random)",
        )

    # --------------------------------------------------------
    # BUTTON
    # --------------------------------------------------------

    generate_button = gr.Button(
        "🚀 GENERATE OUTFIT",
        variant="primary",
        size="lg",
    )

    generate_button.click(
        fn=generate_outfit,
        inputs=[
            person,
            garment,
            mask,
            steps,
            cfg,
            seed,
        ],
        outputs=result,
    )

    gr.Markdown(
        """
### Workflow mapping
`Person → target_image`  
`Garment → refer_image`  
`White/black mask → mask_image`  
`CatVTON → Generated Result`
**Model:** `zhengchong/CatVTON`  
**Resolution:** 768 × 1024  
**Default:** 30 steps / CFG 2.5
"""
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch()

	import os
	import sys
	import subprocess
	from pathlib import Path

	import gradio as gr
	import spaces
	import torch
	from PIL import Image
	from diffusers.image_processor import VaeImageProcessor
	from huggingface_hub import snapshot_download

	# ============================================================
	# CatVTON ZeroGPU — stable Gradio app
	# Based on the user's ComfyUI workflow:
	#
	# LoadCatVTONPipeline_AMD
	# Person image -> target_image
	# Garment image -> refer_image
	# Clothing mask -> mask_image
	# ImageInvert -> represented by the mask polarity note below
	# CatVTON_AMD -> result
	#
	# IMPORTANT:
	# The uploaded UI mask should be:
	# WHITE = clothing / replace area
	# BLACK = background / keep area
	# Therefore this version does NOT invert the uploaded mask.
	# ============================================================

	ROOT = Path("/tmp/catvton_src")
	ROOT.mkdir(parents=True, exist_ok=True)

	BASE_MODEL = os.getenv(
	"CATVTON_BASE_MODEL",
	"booksforcharlie/stable-diffusion-inpainting",
	)
	CATVTON_MODEL = os.getenv(
	"CATVTON_MODEL",
	"zhengchong/CatVTON",
	)

	WIDTH = int(os.getenv("CATVTON_WIDTH", "768"))
	HEIGHT = int(os.getenv("CATVTON_HEIGHT", "1024"))


	def ensure_catvton_source():
	marker = ROOT / "model" / "pipeline.py"

	if marker.exists():
	return

	if (ROOT / ".git").exists():
	subprocess.run(
	["git", "-C", str(ROOT), "pull", "--ff-only"],
	check=False,
	stdout=subprocess.DEVNULL,
	stderr=subprocess.DEVNULL,
	)
	return

	subprocess.run(
	[
	"git",
	"clone",
	"--depth",
	"1",
	"--branch",
	"edited",
	"https://github.com/Zheng-Chong/CatVTON.git",
	str(ROOT),
	],
	check=True,
	)


	ensure_catvton_source()
	sys.path.insert(0, str(ROOT))

	from model.pipeline import CatVTONPipeline
	from utils import resize_and_crop, resize_and_padding


	# ============================================================
	# MODEL
	# ============================================================

	print("Downloading CatVTON checkpoint...")
	catvton_repo = snapshot_download(repo_id=CATVTON_MODEL)

	print("Loading CatVTON pipeline...")

	pipe = CatVTONPipeline(
	base_ckpt=BASE_MODEL,
	attn_ckpt=catvton_repo,
	attn_ckpt_version="mix",
	weight_dtype=torch.float16,
	use_tf32=True,
	device="cuda",
	# The upstream safety checker can produce false positives on
	# normal fully-clothed fashion images. The user's successful
	# test required this to be disabled.
	skip_safety_check=True,
	)

	mask_processor = VaeImageProcessor(
	vae_scale_factor=8,
	do_normalize=False,
	do_binarize=True,
	do_convert_grayscale=True,
	)


	# ============================================================
	# IMAGE HELPERS
	# ============================================================

	def load_rgb(value):
	if value is None:
	return None

	if isinstance(value, Image.Image):
	return value.convert("RGB")

	return Image.open(value).convert("RGB")


	def load_mask(value):
	if value is None:
	return None

	if isinstance(value, Image.Image):
	return value.convert("L")

	return Image.open(value).convert("L")


	def prepare_inputs(person_file, garment_file, mask_file):
	person = load_rgb(person_file)
	garment = load_rgb(garment_file)
	mask = load_mask(mask_file)

	# Same target/reference preprocessing used by the working app.
	person = resize_and_crop(person, (WIDTH, HEIGHT))
	garment = resize_and_padding(garment, (WIDTH, HEIGHT))

	# IMPORTANT:
	# The UI mask is already:
	# white = clothing/replacement region
	# black = preserved region
	#
	# Do not invert it again.
	mask = resize_and_crop(mask, (WIDTH, HEIGHT))

	# Light mask edge smoothing for cleaner garment boundaries.
	mask = mask_processor.blur(mask, blur_factor=9)

	return person, garment, mask


	# ============================================================
	# GENERATION
	# ============================================================

	@spaces.GPU(duration=120)
	def generate_outfit(person_file, garment_file, mask_file, steps, cfg, seed):
	if person_file is None:
	raise gr.Error("Please upload the person image.")

	if garment_file is None:
	raise gr.Error("Please upload the garment/clothing image.")

	if mask_file is None:
	raise gr.Error(
	"Please upload the clothing mask. "
	"White = clothing area, Black = background."
	)

	person, garment, mask = prepare_inputs(
	person_file,
	garment_file,
	mask_file,
	)

	steps = int(steps)
	cfg = float(cfg)
	seed = int(seed)

	generator = None

	# -1 means random seed.
	if seed >= 0:
	generator = torch.Generator(device="cuda").manual_seed(seed)

	print(
	f"CatVTON inference \| steps={steps} "
	f"cfg={cfg} seed={seed}"
	)

	with torch.inference_mode():
	result = pipe(
	image=person,
	condition_image=garment,
	mask=mask,
	num_inference_steps=steps,
	guidance_scale=cfg,
	height=HEIGHT,
	width=WIDTH,
	generator=generator,
	)[0]

	return result


	# ============================================================
	# UI
	# ============================================================

	DESCRIPTION = """
	# 🐈 CatVTON Outfit Try-On

	Upload a person photo, a garment/product photo, and a clothing mask.

	Mask format: ⚪ White = replace clothing area · ⚫ Black = preserve area

	The default settings match the tested workflow: 30 steps / CFG 2.5.
	"""

	with gr.Blocks(title="CatVTON Outfit Try-On — ZeroGPU") as demo:

	gr.Markdown(DESCRIPTION)

	with gr.Row():

	# ----------------------------------------------------
	# INPUTS
	# ----------------------------------------------------
	with gr.Column():

	person = gr.Image(
	type="filepath",
	label="👤 Person Image",
	)

	garment = gr.Image(
	type="filepath",
	label="👕 Garment / Clothing Image",
	)

	mask = gr.Image(
	type="filepath",
	label="🎭 Clothing Mask",
	)

	gr.Markdown(
	"""
	Mask example

	- White shirt shape = area CatVTON replaces
	- Black background = area CatVTON preserves
	"""
	)

	# ----------------------------------------------------
	# OUTPUT
	# ----------------------------------------------------
	with gr.Column():

	result = gr.Image(
	type="pil",
	label="✨ Generated Result",
	)

	gr.Markdown(
	"Use the image toolbar to download the generated result."
	)

	# --------------------------------------------------------
	# SETTINGS
	# --------------------------------------------------------

	with gr.Row():

	steps = gr.Slider(
	minimum=10,
	maximum=50,
	step=1,
	value=30,
	label="Inference Steps",
	)

	cfg = gr.Slider(
	minimum=1.0,
	maximum=5.0,
	step=0.1,
	value=2.5,
	label="CFG",
	)

	seed = gr.Number(
	value=-1,
	precision=0,
	label="Seed (-1 = Random)",
	)

	# --------------------------------------------------------
	# BUTTON
	# --------------------------------------------------------

	generate_button = gr.Button(
	"🚀 GENERATE OUTFIT",
	variant="primary",
	size="lg",
	)

	generate_button.click(
	fn=generate_outfit,
	inputs=[
	person,
	garment,
	mask,
	steps,
	cfg,
	seed,
	],
	outputs=result,
	)

	gr.Markdown(
	"""
	### Workflow mapping

	`Person → target_image`
	`Garment → refer_image`
	`White/black mask → mask_image`
	`CatVTON → Generated Result`

	Model: `zhengchong/CatVTON`
	Resolution: 768 × 1024
	Default: 30 steps / CFG 2.5
	"""
	)


	if __name__ == "__main__":
	demo.queue(max_size=8).launch()