Add ZeroGPU virtual try-on demo

README.md

@@ -1,14 +1,126 @@
 ---
-title: Tryon Demo
-emoji: 🖼
-colorFrom: purple
-colorTo: red
+title: Virtual Try-On ZeroGPU Demo
+emoji: VTON
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.5.1
+sdk_version: 5.0.0
+python_version: 3.10
 app_file: app.py
 pinned: false
-license: apache-2.0
-short_description: A demo for our unified virtual tryon model
+license: other
+short_description: ZeroGPU xlarge Gradio demo for image-based virtual try-on.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Virtual Try-On ZeroGPU Demo
+
+This Space is a Gradio SDK deployment wrapper for a virtual try-on model.
+
+## Inputs
+
+- One person image
+- One or more garment reference images
+- Garment category
+- Primary garment index
+- Seed, inference steps, guidance scale, and auto-crop controls
+
+## Output
+
+- One generated try-on image
+
+## ZeroGPU
+
+The inference callback is decorated with:
+
+```python
+@spaces.GPU(size="xlarge", duration=180)
+```
+
+Use ZeroGPU `xlarge` because the model is expected to need at least 60 GB of VRAM.
+
+## Required Space Variables
+
+Set these in the Space settings:
+
+```text
+MODEL_ID=your-org/your-vton-model
+```
+
+If the model repository is private, add this as a Space Secret:
+
+```text
+HF_TOKEN=hf_xxxxxxxxxxxxxxxxx
+```
+
+Optional variables:
+
+```text
+MODEL_REVISION=main
+MODEL_PIPELINE_MODULE=pipeline
+MODEL_PIPELINE_CLASS=VirtualTryOnPipeline
+MODEL_SUPPORTS_MULTI_GARMENT=false
+MODEL_TORCH_DTYPE=float16
+HF_HOME=/data/.huggingface
+```
+
+## Expected Model Repository Interface
+
+This wrapper downloads the model repository with `huggingface_hub.snapshot_download`.
+By default, it expects the model repo to contain:
+
+```text
+pipeline.py
+```
+
+and inside it:
+
+```python
+class VirtualTryOnPipeline:
+    @classmethod
+    def from_pretrained(cls, model_dir, torch_dtype=None):
+        ...
+
+    def to(self, device):
+        ...
+
+    def __call__(
+        self,
+        person_image,
+        garment_image=None,
+        garment_images=None,
+        garment_type="upper_body",
+        num_inference_steps=30,
+        guidance_scale=2.5,
+        generator=None,
+        seed=42,
+        auto_crop=True,
+    ):
+        ...
+```
+
+The return value may be:
+
+- `PIL.Image.Image`
+- a list or tuple whose first element is a PIL image
+- an object with `.images[0]`
+- a dict containing `image`, `images`, or `result`
+
+If your current virtual try-on model only supports a single garment image, keep:
+
+```text
+MODEL_SUPPORTS_MULTI_GARMENT=false
+```
+
+The app will use `primary_garment_index` to select exactly one garment image from the uploaded references. It does not pretend to perform multi-reference fusion.
+
+## Deployment Steps
+
+1. Create a new Hugging Face Space with SDK `Gradio` and hardware `ZeroGPU`.
+2. Copy these files into the Space repository.
+3. Add any model-specific Python dependencies to `requirements.txt`.
+4. Set `MODEL_ID` as a Space Variable.
+5. If the model is private, set `HF_TOKEN` as a Space Secret.
+6. Push the Space repository.
+7. Open the Space and run a test with one person image and one garment image.
+
+Do not commit large model weights to this Space repository. Keep weights in a Hugging Face Model Repo.

app.py

@@ -1,154 +1,153 @@
-import gradio as gr
-import numpy as np
+from __future__ import annotations
+
 import random
+import traceback
+
+import gradio as gr
+import spaces
+from PIL import Image
 
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
+from inference import run_tryon
+from preprocess import load_garment_images, resize_for_demo
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+def _validate_inputs(person_image: Image.Image | None, garment_files) -> list[Image.Image]:
+    if person_image is None:
+        raise gr.Error("Please upload a person image.")
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+    garment_images = load_garment_images(garment_files)
+    if not garment_images:
+        raise gr.Error("Please upload at least one garment reference image.")
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+    return garment_images
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
+@spaces.GPU(size="xlarge", duration=180)
+def predict(
+    person_image: Image.Image,
+    garment_files,
+    garment_type: str,
+    primary_garment_index: int,
+    seed: int,
+    randomize_seed: bool,
+    num_inference_steps: int,
+    guidance_scale: float,
+    auto_crop: bool,
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+    try:
+        garment_images = _validate_inputs(person_image, garment_files)
+
+        if randomize_seed:
+            seed = random.randint(0, 2**31 - 1)
+
+        person = resize_for_demo(person_image, max_side=1024)
+        garments = [resize_for_demo(img, max_side=1024) for img in garment_images]
+
+        result, used_index, status = run_tryon(
+            person_image=person,
+            garment_images=garments,
+            garment_type=garment_type,
+            primary_garment_index=int(primary_garment_index),
+            seed=int(seed),
+            num_inference_steps=int(num_inference_steps),
+            guidance_scale=float(guidance_scale),
+            auto_crop=bool(auto_crop),
+        )
+
+        return result, int(seed), used_index, status
+
+    except gr.Error:
+        raise
+    except Exception as exc:
+        traceback.print_exc()
+        raise gr.Error(f"Inference failed: {exc}") from exc
+
+
+with gr.Blocks(title="Virtual Try-On ZeroGPU Demo") as demo:
+    gr.Markdown(
+        """
+        # Virtual Try-On ZeroGPU Demo
+
+        Upload one person image and one or more garment reference images.
+        If the underlying model only supports one garment image, the selected
+        primary garment index is used as the single reference image.
+        """
+    )
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+    with gr.Row():
+        with gr.Column(scale=1):
+            person_image = gr.Image(
+                label="Person Image",
+                type="pil",
+                sources=["upload", "clipboard"],
+            )
 
-        result = gr.Image(label="Result", show_label=False)
+            garment_files = gr.File(
+                label="Garment Reference Images",
+                file_count="multiple",
+                file_types=["image"],
+                type="filepath",
+            )
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+            garment_type = gr.Radio(
+                choices=["upper_body", "lower_body", "dress"],
+                value="upper_body",
+                label="Garment Type",
             )
 
-            seed = gr.Slider(
-                label="Seed",
+            primary_garment_index = gr.Slider(
                 minimum=0,
-                maximum=MAX_SEED,
-                step=1,
+                maximum=9,
                 value=0,
+                step=1,
+                label="Primary Garment Index",
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Number(value=42, precision=0, label="Seed")
+                randomize_seed = gr.Checkbox(value=True, label="Randomize Seed")
+                num_inference_steps = gr.Slider(
+                    minimum=10,
+                    maximum=80,
+                    value=30,
+                    step=1,
+                    label="Inference Steps",
                 )
-
-            with gr.Row():
                 guidance_scale = gr.Slider(
-                    label="Guidance scale",
                     minimum=0.0,
-                    maximum=10.0,
+                    maximum=15.0,
+                    value=2.5,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    label="Guidance Scale",
                 )
+                auto_crop = gr.Checkbox(value=True, label="Auto Crop Person Image")
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
+            run_button = gr.Button("Generate Try-On", variant="primary")
+
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Try-On Result", type="pil")
+            used_seed = gr.Number(label="Used Seed", precision=0)
+            used_garment_index = gr.Number(label="Used Garment Index", precision=0)
+            status = gr.Textbox(label="Status", interactive=False)
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+    run_button.click(
+        fn=predict,
         inputs=[
-            prompt,
-            negative_prompt,
+            person_image,
+            garment_files,
+            garment_type,
+            primary_garment_index,
             seed,
             randomize_seed,
-            width,
-            height,
-            guidance_scale,
             num_inference_steps,
+            guidance_scale,
+            auto_crop,
         ],
-        outputs=[result, seed],
+        outputs=[output_image, used_seed, used_garment_index, status],
+        api_name="tryon",
     )
 
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(max_size=20).launch()
+

inference.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import inspect
+import os
+import random
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+from preprocess import select_primary_garment
+
+
+_MODEL: Any | None = None
+_MODEL_DIR: Path | None = None
+
+
+def _get_bool_env(name: str, default: bool = False) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _get_torch_dtype() -> torch.dtype:
+    value = os.getenv("MODEL_TORCH_DTYPE", "float16").strip().lower()
+    if value in {"bfloat16", "bf16"}:
+        return torch.bfloat16
+    if value in {"float32", "fp32"}:
+        return torch.float32
+    return torch.float16
+
+
+def _device() -> torch.device:
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "CUDA is not available. Make sure the Space hardware is set to ZeroGPU "
+            "and that inference is called from a function decorated with @spaces.GPU."
+        )
+    return torch.device("cuda")
+
+
+def set_seed(seed: int) -> torch.Generator:
+    seed = int(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    return torch.Generator(device="cuda").manual_seed(seed)
+
+
+def _download_model_repo() -> Path:
+    model_id = os.getenv("MODEL_ID")
+    if not model_id:
+        raise RuntimeError("MODEL_ID is not set. Add it as a Hugging Face Space Variable.")
+
+    token = os.getenv("HF_TOKEN")
+    revision = os.getenv("MODEL_REVISION") or None
+    local_dir = os.getenv("MODEL_LOCAL_DIR") or None
+
+    kwargs = {
+        "repo_id": model_id,
+        "revision": revision,
+        "token": token,
+    }
+    if local_dir:
+        kwargs["local_dir"] = local_dir
+        kwargs["local_dir_use_symlinks"] = False
+
+    path = snapshot_download(**kwargs)
+    return Path(path)
+
+
+def _load_module_from_file(module_path: Path) -> ModuleType:
+    module_name = f"vton_model_{module_path.stem}"
+    spec = importlib.util.spec_from_file_location(module_name, str(module_path))
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not import module from {module_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_pipeline_module(model_dir: Path) -> ModuleType:
+    module_name = os.getenv("MODEL_PIPELINE_MODULE", "pipeline")
+
+    file_candidate = model_dir / f"{module_name.replace('.', '/')}.py"
+    if file_candidate.exists():
+        return _load_module_from_file(file_candidate)
+
+    if str(model_dir) not in sys.path:
+        sys.path.insert(0, str(model_dir))
+
+    return importlib.import_module(module_name)
+
+
+def _instantiate_pipeline(module: ModuleType, model_dir: Path) -> Any:
+    class_name = os.getenv("MODEL_PIPELINE_CLASS", "VirtualTryOnPipeline")
+    if not hasattr(module, class_name):
+        raise AttributeError(
+            f"Model pipeline module '{module.__name__}' does not define '{class_name}'. "
+            "Set MODEL_PIPELINE_CLASS or provide a compatible pipeline class."
+        )
+
+    cls = getattr(module, class_name)
+    dtype = _get_torch_dtype()
+
+    if hasattr(cls, "from_pretrained") and callable(getattr(cls, "from_pretrained")):
+        try:
+            return cls.from_pretrained(str(model_dir), torch_dtype=dtype, token=os.getenv("HF_TOKEN"))
+        except TypeError:
+            try:
+                return cls.from_pretrained(str(model_dir), torch_dtype=dtype)
+            except TypeError:
+                return cls.from_pretrained(str(model_dir))
+
+    try:
+        return cls(str(model_dir), torch_dtype=dtype)
+    except TypeError:
+        return cls(str(model_dir))
+
+
+def load_model() -> Any:
+    """Download and load the virtual try-on model once per Space process."""
+    global _MODEL, _MODEL_DIR
+
+    if _MODEL is not None:
+        return _MODEL
+
+    _MODEL_DIR = _download_model_repo()
+    module = _load_pipeline_module(_MODEL_DIR)
+    model = _instantiate_pipeline(module, _MODEL_DIR)
+
+    device = _device()
+    if hasattr(model, "to") and callable(getattr(model, "to")):
+        model = model.to(device)
+    if hasattr(model, "eval") and callable(getattr(model, "eval")):
+        model.eval()
+
+    _MODEL = model
+    return _MODEL
+
+
+def _call_model(
+    model: Any,
+    *,
+    person_image: Image.Image,
+    garment_image: Image.Image,
+    garment_images: list[Image.Image],
+    garment_type: str,
+    seed: int,
+    num_inference_steps: int,
+    guidance_scale: float,
+    auto_crop: bool,
+    generator: torch.Generator,
+) -> Any:
+    supports_multi = _get_bool_env("MODEL_SUPPORTS_MULTI_GARMENT", False)
+
+    kwargs = {
+        "person_image": person_image,
+        "garment_type": garment_type,
+        "num_inference_steps": int(num_inference_steps),
+        "guidance_scale": float(guidance_scale),
+        "generator": generator,
+        "seed": int(seed),
+        "auto_crop": bool(auto_crop),
+    }
+
+    if supports_multi:
+        kwargs["garment_images"] = garment_images
+    else:
+        kwargs["garment_image"] = garment_image
+
+    call = model.run_tryon if hasattr(model, "run_tryon") else model
+    signature = inspect.signature(call)
+    accepted_kwargs = {
+        key: value
+        for key, value in kwargs.items()
+        if key in signature.parameters or any(p.kind == p.VAR_KEYWORD for p in signature.parameters.values())
+    }
+    return call(**accepted_kwargs)
+
+
+def _extract_image(result: Any) -> Image.Image:
+    if isinstance(result, Image.Image):
+        return result.convert("RGB")
+
+    if isinstance(result, dict):
+        for key in ("image", "result", "output"):
+            if key in result:
+                return _extract_image(result[key])
+        if "images" in result:
+            return _extract_image(result["images"])
+
+    if hasattr(result, "images"):
+        return _extract_image(result.images)
+
+    if isinstance(result, (list, tuple)) and result:
+        return _extract_image(result[0])
+
+    raise TypeError(
+        "The model output could not be converted to a PIL image. Return a PIL image, "
+        "a list of PIL images, an object with .images, or a dict with image/images/result."
+    )
+
+
+@torch.inference_mode()
+def run_tryon(
+    person_image: Image.Image,
+    garment_images: list[Image.Image],
+    garment_type: str = "upper_body",
+    primary_garment_index: int = 0,
+    seed: int = 42,
+    num_inference_steps: int = 30,
+    guidance_scale: float = 2.5,
+    auto_crop: bool = True,
+) -> tuple[Image.Image, int, str]:
+    if person_image is None:
+        raise ValueError("person_image is required.")
+    if not garment_images:
+        raise ValueError("At least one garment reference image is required.")
+
+    primary_garment, used_index, all_garments = select_primary_garment(
+        garment_images,
+        primary_garment_index,
+    )
+    generator = set_seed(seed)
+    model = load_model()
+
+    result = _call_model(
+        model,
+        person_image=person_image.convert("RGB"),
+        garment_image=primary_garment.convert("RGB"),
+        garment_images=[img.convert("RGB") for img in all_garments],
+        garment_type=garment_type,
+        seed=seed,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        auto_crop=auto_crop,
+        generator=generator,
+    )
+
+    image = _extract_image(result)
+    mode = "multi-reference" if _get_bool_env("MODEL_SUPPORTS_MULTI_GARMENT", False) else "single-reference"
+    status = f"Success. Used garment index {used_index} in {mode} mode."
+    return image, used_index, status
+
+
+def clear_model_cache() -> None:
+    global _MODEL, _MODEL_DIR
+    _MODEL = None
+    _MODEL_DIR = None
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()

packages.txt

@@ -0,0 +1,5 @@
+ffmpeg
+libgl1
+libglib2.0-0
+git
+

preprocess.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable
+
+from PIL import Image, ImageOps
+
+
+def load_garment_images(file_paths) -> list[Image.Image]:
+    """Load Gradio File outputs into RGB PIL images."""
+    if file_paths is None:
+        return []
+
+    if isinstance(file_paths, (str, Path)):
+        file_paths = [file_paths]
+
+    images: list[Image.Image] = []
+    for item in file_paths:
+        if item is None:
+            continue
+
+        path = item
+        if not isinstance(path, (str, Path)) and hasattr(path, "name"):
+            path = path.name
+
+        p = Path(path)
+        if not p.exists():
+            continue
+
+        with Image.open(p) as img:
+            images.append(ImageOps.exif_transpose(img).convert("RGB"))
+
+    return images
+
+
+def normalize_pil_image(image: Image.Image | None) -> Image.Image | None:
+    if image is None:
+        return None
+    return ImageOps.exif_transpose(image).convert("RGB")
+
+
+def resize_for_demo(image: Image.Image, max_side: int = 1024) -> Image.Image:
+    """Downscale very large uploads while preserving aspect ratio."""
+    image = normalize_pil_image(image)
+    if image is None:
+        raise ValueError("image is required")
+
+    width, height = image.size
+    longest = max(width, height)
+    if longest <= max_side:
+        return image
+
+    scale = max_side / float(longest)
+    new_size = (max(1, int(width * scale)), max(1, int(height * scale)))
+    return image.resize(new_size, Image.Resampling.LANCZOS)
+
+
+def clamp_primary_index(primary_garment_index: int, garment_count: int) -> int:
+    if garment_count <= 0:
+        raise ValueError("At least one garment reference image is required.")
+    return max(0, min(int(primary_garment_index), garment_count - 1))
+
+
+def select_primary_garment(
+    garment_images: Iterable[Image.Image],
+    primary_garment_index: int,
+) -> tuple[Image.Image, int, list[Image.Image]]:
+    images = list(garment_images)
+    index = clamp_primary_index(primary_garment_index, len(images))
+    return images[index], index, images
+

requirements.txt

@@ -1,6 +1,16 @@
+gradio>=5.0.0
+spaces
+huggingface_hub>=0.27.0
+torch>=2.4.0
+torchvision
 accelerate
 diffusers
-invisible_watermark
-torch
 transformers
-xformers
+safetensors
+pillow
+opencv-python-headless
+numpy
+scipy
+einops
+peft
+