Upload 2 files

app (1).py

@@ -0,0 +1,181 @@
+"""
+app.py — OOTDiffusion Hugging Face Space
+Place this file in the ROOT of your Space repo.
+
+Your Space structure should look like:
+  OOODdiffusion/
+  ├── app.py                  ← this file (root level)
+  ├── requirements.txt        ← root level
+  ├── README.md               ← root level
+  └── OOTDiffusion-main/      ← the uploaded zip contents
+      ├── ootd/
+      ├── run/
+      ├── preprocess/
+      ├── checkpoints/
+      └── ...
+"""
+
+import sys
+import os
+
+# ── Path setup ────────────────────────────────────────────────────────────────
+ROOT_DIR  = os.path.dirname(os.path.abspath(__file__))
+
+# Support both flat layout and nested OOTDiffusion-main/ layout
+OOTD_DIR = ROOT_DIR
+for candidate in ["OOTDiffusion-main", "OOTDiffusion"]:
+    candidate_path = os.path.join(ROOT_DIR, candidate)
+    if os.path.isdir(candidate_path):
+        OOTD_DIR = candidate_path
+        break
+
+RUN_DIR = os.path.join(OOTD_DIR, "run")
+
+sys.path.insert(0, OOTD_DIR)
+sys.path.insert(0, RUN_DIR)
+
+print(f"[OOTDiffusion] ROOT_DIR : {ROOT_DIR}")
+print(f"[OOTDiffusion] OOTD_DIR : {OOTD_DIR}")
+
+import torch
+import numpy as np
+import gradio as gr
+from PIL import Image
+
+# ── Device ────────────────────────────────────────────────────────────────────
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"[OOTDiffusion] Device: {DEVICE}")
+
+# ── Lazy-load models ──────────────────────────────────────────────────────────
+_pipe_hd = None
+_pipe_dc = None
+
+
+def load_pipeline(model_type: str):
+    global _pipe_hd, _pipe_dc
+    if model_type == "hd":
+        if _pipe_hd is None:
+            from ootd.inference_ootd_hd import OOTDiffusionHD
+            print("[OOTDiffusion] Loading HD pipeline …")
+            _pipe_hd = OOTDiffusionHD(OOTD_DIR)
+        return _pipe_hd
+    else:
+        if _pipe_dc is None:
+            from ootd.inference_ootd_dc import OOTDiffusionDC
+            print("[OOTDiffusion] Loading DC pipeline …")
+            _pipe_dc = OOTDiffusionDC(OOTD_DIR)
+        return _pipe_dc
+
+
+# ── Category mapping ──────────────────────────────────────────────────────────
+CATEGORY_MAP = {
+    "Upper-body": 0,
+    "Lower-body":  1,
+    "Dress":       2,
+}
+
+
+# ── Inference ─────────────────────────────────────────────────────────────────
+def run_tryon(model_image, cloth_image, model_type, category_label,
+              n_samples, n_steps, guidance_scale, seed):
+
+    if model_image is None:
+        raise gr.Error("Please upload a model (person) image.")
+    if cloth_image is None:
+        raise gr.Error("Please upload a garment image.")
+
+    if isinstance(model_image, np.ndarray):
+        model_image = Image.fromarray(model_image)
+    if isinstance(cloth_image, np.ndarray):
+        cloth_image = Image.fromarray(cloth_image)
+
+    model_image = model_image.convert("RGB")
+    cloth_image = cloth_image.convert("RGB")
+
+    category_idx = CATEGORY_MAP[category_label]
+
+    try:
+        pipe = load_pipeline(model_type)
+    except Exception as e:
+        raise gr.Error(
+            f"Failed to load model: {e}\n"
+            "Make sure OOTDiffusion-main/ folder with ootd/ and checkpoints/ is present."
+        )
+
+    try:
+        result = pipe(
+            model_type=model_type,
+            category=category_idx,
+            image_garm=cloth_image,
+            image_vton=model_image,
+            mask=None,
+            image_ori=model_image,
+            num_samples=int(n_samples),
+            num_steps=int(n_steps),
+            guidance_scale=float(guidance_scale),
+            seed=int(seed),
+        )
+    except Exception as e:
+        raise gr.Error(f"Inference failed: {e}")
+
+    if isinstance(result, (list, tuple)):
+        return result
+    return [result]
+
+
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+with gr.Blocks(title="OOTDiffusion Virtual Try-On", theme=gr.themes.Soft()) as demo:
+
+    gr.Markdown("""
+    # 👗 OOTDiffusion — Virtual Try-On
+    **[AAAI 2025]** Upload a *person photo* and a *garment image*, then click **Run Try-On**.
+    > ⚠️ Non-commercial use only (CC-BY-NC-SA-4.0)
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            model_img = gr.Image(label="👤 Model Image (person)", type="pil", height=380)
+            cloth_img = gr.Image(label="👕 Garment Image", type="pil", height=380)
+
+        with gr.Column(scale=1):
+            model_type = gr.Radio(
+                choices=["hd", "dc"], value="hd",
+                label="Model Type",
+                info="hd = half-body (VITON-HD)  |  dc = full-body (Dress Code)"
+            )
+            category = gr.Dropdown(
+                choices=list(CATEGORY_MAP.keys()), value="Upper-body",
+                label="Garment Category",
+                info="Only matters when Model Type = dc"
+            )
+            n_samples = gr.Slider(1, 4, step=1, value=1, label="Number of Samples")
+            n_steps   = gr.Slider(10, 40, step=5, value=20, label="Denoising Steps",
+                                  info="More steps = better quality, slower")
+            guidance  = gr.Slider(1.0, 5.0, step=0.5, value=2.0, label="Guidance Scale")
+            seed      = gr.Number(value=42, label="Seed  (-1 = random)", precision=0)
+            run_btn   = gr.Button("🚀 Run Try-On", variant="primary", size="lg")
+
+        with gr.Column(scale=1):
+            output_gallery = gr.Gallery(
+                label="✨ Try-On Results",
+                columns=2, height=500, object_fit="contain"
+            )
+
+    gr.Markdown("""
+    ### 💡 Tips
+    - **HD model** — best for upper-body garments on half-body photos
+    - **DC model** — supports upper / lower / dress on full-body photos
+    - Steps **30–40** give noticeably better quality
+    - **Seed = -1** gives a different result each run
+    """)
+
+    run_btn.click(
+        fn=run_tryon,
+        inputs=[model_img, cloth_img, model_type, category,
+                n_samples, n_steps, guidance, seed],
+        outputs=output_gallery,
+    )
+
+# ── Launch ────────────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    demo.launch()

requirements (2).txt

@@ -0,0 +1,14 @@
+torch
+torchvision
+torchaudio
+numpy
+Pillow
+diffusers==0.24.0
+transformers==4.36.2
+accelerate==0.25.0
+omegaconf
+einops
+opencv-python-headless
+scikit-image
+huggingface_hub
+onnxruntime