Upload 3 files

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.11-slim
+
+RUN apt-get update && apt-get install -y \
+    libglib2.0-0 libsm6 libxext6 libxrender-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app.py .
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1", "--timeout-keep-alive", "300"]

app-5.py

@@ -0,0 +1,330 @@
+"""
+╔══════════════════════════════════════════════════════════════════╗
+║   PIXELSKITS — WORLD CLASS VIRTUAL TRY-ON API  v2.0             ║
+║                                                                  ║
+║   PIPELINE:                                                      ║
+║     1. IDM-VTON  — World #1 virtual try-on (HF servers)         ║
+║     2. Real-ESRGAN — Upscale to 720p/1080p/2K/4K/8K             ║
+║                                                                  ║
+║   Our Space = lightweight middleman (~200MB RAM)                 ║
+║   All heavy AI runs on HF servers FREE                          ║
+╚══════════════════════════════════════════════════════════════════╝
+"""
+
+import os, io, hashlib, time, logging, tempfile, base64
+from typing import Optional
+from datetime import datetime
+
+import requests as http_requests
+from fastapi import FastAPI, File, UploadFile, Header, HTTPException, Request, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import Response
+from PIL import Image
+from gradio_client import Client, handle_file
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+log = logging.getLogger("pixelskits-tryon")
+
+# ──────────────────────────────────────────────────────────
+MASTER_KEY  = os.environ.get("PIXELSKITS_API_KEY", "")
+HF_TOKEN    = os.environ.get("HF_TOKEN", "")
+RATE_LIMIT  = 30
+RATE_WINDOW = 3600
+
+# Real-ESRGAN upscaler on HF Inference API
+UPSCALE_URL = "https://api-inference.huggingface.co/models/ai-forever/Real-ESRGAN"
+
+# Resolution targets (width x height)
+RESOLUTIONS = {
+    "720p":  (1280,  720),
+    "1080p": (1920, 1080),
+    "2k":    (2560, 1440),
+    "4k":    (3840, 2160),
+    "8k":    (7680, 4320),
+}
+
+app = FastAPI(title="Pixelskits Virtual Try-On API v2", version="2.0.0", docs_url="/docs")
+app.add_middleware(CORSMiddleware, allow_origins=["*"],
+                   allow_methods=["GET", "POST"], allow_headers=["*"])
+
+_rate: dict   = {}
+_client       = None
+_ready_flag   = False
+
+
+# ══════════════════════════════════════════════════════════
+#  RATE + AUTH
+# ══════════════════════════════════════════════════════════
+
+def check_rate(ip: str):
+    now = time.time()
+    h   = hashlib.sha256(ip.encode()).hexdigest()[:16]
+    rec = _rate.get(h, {"count": 0, "start": now})
+    if now - rec["start"] > RATE_WINDOW:
+        rec = {"count": 0, "start": now}
+    rec["count"] += 1
+    _rate[h] = rec
+    if rec["count"] > RATE_LIMIT:
+        raise HTTPException(429, f"Rate limit: {RATE_LIMIT} req/hour.")
+
+def _auth(key: Optional[str]):
+    if MASTER_KEY and key != MASTER_KEY:
+        raise HTTPException(401, "Invalid API key.")
+
+def _get_ip(req: Request) -> str:
+    return req.headers.get("x-forwarded-for", req.client.host or "unknown")
+
+
+# ══════════════════════════════════════════════════════════
+#  STARTUP
+# ══════════════════════════════════════════════════════════
+
+@app.on_event("startup")
+async def startup():
+    global _client, _ready_flag
+    log.info("Connecting to IDM-VTON Space…")
+    try:
+        token    = HF_TOKEN if HF_TOKEN else None
+        _client  = Client("yisol/IDM-VTON", hf_token=token)
+        log.info("✅ IDM-VTON connected. API online.")
+    except Exception as e:
+        log.warning(f"IDM-VTON connect warning: {e} — will retry per request")
+    _ready_flag = True
+
+def _ready():
+    return _ready_flag
+
+
+# ══════════════════════════════════════════════════════════
+#  HELPERS
+# ══════════════════════════════════════════════════════════
+
+def _save_temp(data: bytes, suffix: str = ".png") -> str:
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+    tmp.write(data)
+    tmp.flush()
+    tmp.close()
+    return tmp.name
+
+def _pil_to_bytes(img: Image.Image, fmt: str = "PNG") -> bytes:
+    buf = io.BytesIO()
+    img.save(buf, format=fmt, compress_level=6 if fmt == "PNG" else None)
+    return buf.getvalue()
+
+def _validate_and_save(data: bytes, label: str) -> str:
+    try:
+        img = Image.open(io.BytesIO(data))
+        img.verify()
+    except Exception:
+        raise HTTPException(400, f"{label}: invalid image.")
+    img = Image.open(io.BytesIO(data)).convert("RGB")
+    return _save_temp(_pil_to_bytes(img))
+
+
+# ══════════════════════════════════════════════════════════
+#  STEP 1 — IDM-VTON Try-On
+# ══════════════════════════════════════════════════════════
+
+def run_tryon(person_path: str, garment_path: str,
+              category: str = "upper_body",
+              denoise_steps: int = 30,
+              seed: int = 42) -> Image.Image:
+    global _client
+
+    if _client is None:
+        log.info("Reconnecting to IDM-VTON…")
+        _client = Client("yisol/IDM-VTON", hf_token=HF_TOKEN or None)
+
+    log.info(f"TryOn: category={category} steps={denoise_steps}")
+
+    result = _client.predict(
+        dict(background=handle_file(person_path), layers=[], composite=None),
+        handle_file(garment_path),
+        "Output",
+        True,           # auto-masking
+        True,           # auto-crop
+        denoise_steps,
+        seed,
+        api_name="/tryon"
+    )
+
+    result_path = result[0]
+    if isinstance(result_path, dict):
+        result_path = result_path.get("image", result_path.get("path", ""))
+
+    log.info(f"TryOn: result at {result_path}")
+    return Image.open(result_path).convert("RGB")
+
+
+# ══════════════════════════════════════════════════════════
+#  STEP 2 — Real-ESRGAN Upscaler (HF Inference API)
+# ══════════════════════════════════════════════════════════
+
+def upscale_image(img: Image.Image, target_res: str) -> Image.Image:
+    """
+    Upscale using Real-ESRGAN on HF Inference API.
+    Then resize to exact target resolution.
+    Falls back to Lanczos resize if HF API unavailable.
+    """
+    if target_res not in RESOLUTIONS:
+        return img
+
+    target_w, target_h = RESOLUTIONS[target_res]
+    log.info(f"Upscale: target={target_res} ({target_w}×{target_h})")
+
+    # Only call Real-ESRGAN if we need meaningful upscale (>1.5x)
+    src_w, src_h = img.size
+    scale_needed = max(target_w / src_w, target_h / src_h)
+
+    if scale_needed > 1.5 and HF_TOKEN:
+        try:
+            img_bytes = _pil_to_bytes(img)
+            headers   = {"Authorization": f"Bearer {HF_TOKEN}"}
+
+            # Retry up to 3 times for model warmup
+            for attempt in range(3):
+                resp = http_requests.post(
+                    UPSCALE_URL,
+                    headers=headers,
+                    data=img_bytes,
+                    timeout=120
+                )
+                if resp.status_code == 503:
+                    wait = min(float(resp.json().get("estimated_time", 15)), 25)
+                    log.info(f"Upscaler warming up, waiting {wait}s… (attempt {attempt+1})")
+                    time.sleep(wait)
+                    continue
+                if resp.status_code == 200:
+                    upscaled = Image.open(io.BytesIO(resp.content)).convert("RGB")
+                    log.info(f"Upscale: Real-ESRGAN success → {upscaled.size}")
+                    # Final resize to exact target dimensions
+                    final = upscaled.resize((target_w, target_h), Image.LANCZOS)
+                    return final
+                log.warning(f"Upscaler returned {resp.status_code} — using Lanczos")
+                break
+
+        except Exception as e:
+            log.warning(f"Upscaler error: {e} — using Lanczos fallback")
+
+    # Lanczos fallback (still very good quality for moderate upscales)
+    log.info(f"Upscale: Lanczos resize → {target_w}×{target_h}")
+    return img.resize((target_w, target_h), Image.LANCZOS)
+
+
+# ══════════════════════════════════════════════════════════
+#  ROUTES
+# ══════════════════════════════════════════════════════════
+
+@app.get("/")
+async def root():
+    return {
+        "api":         "Pixelskits Virtual Try-On API v2",
+        "status":      "online" if _ready() else "loading",
+        "engine":      "IDM-VTON (World #1) + Real-ESRGAN Upscaler",
+        "resolutions": list(RESOLUTIONS.keys()),
+        "usage":       "POST /tryon — person_image + garment_image + resolution",
+        "category":    "upper_body | lower_body | dresses",
+    }
+
+@app.get("/health")
+async def health():
+    return {
+        "status":     "ok" if _ready() else "loading",
+        "hf_token":   "set" if HF_TOKEN else "missing",
+        "idm_client": "connected" if _client else "disconnected",
+        "ts":         datetime.utcnow().isoformat() + "Z",
+    }
+
+
+@app.post("/tryon")
+async def virtual_tryon(
+    request:       Request,
+    person_image:  UploadFile = File(...,  description="Photo of the person"),
+    garment_image: UploadFile = File(...,  description="Photo of the garment/outfit"),
+    category:      str = Form("upper_body", description="upper_body | lower_body | dresses"),
+    resolution:    str = Form("1080p",      description="720p | 1080p | 2k | 4k | 8k"),
+    denoise_steps: int = Form(30,           description="Quality: 20–40 (higher=better, slower)"),
+    seed:          int = Form(42,           description="Seed for reproducibility (-1 = random)"),
+    x_api_key:     Optional[str] = Header(None, alias="X-API-Key"),
+):
+    """
+    Virtual try-on endpoint.
+    Send person photo + garment photo → get photorealistic try-on result at your chosen resolution.
+    """
+    _auth(x_api_key)
+    check_rate(_get_ip(request))
+
+    if not _ready():
+        raise HTTPException(503, "API starting. Retry in 10s.")
+
+    # Validate inputs
+    allowed = ("image/jpeg", "image/jpg", "image/png", "image/webp")
+    if person_image.content_type not in allowed:
+        raise HTTPException(400, "person_image: use JPEG, PNG, or WebP.")
+    if garment_image.content_type not in allowed:
+        raise HTTPException(400, "garment_image: use JPEG, PNG, or WebP.")
+    if category not in ("upper_body", "lower_body", "dresses"):
+        raise HTTPException(400, "category: upper_body | lower_body | dresses")
+    if resolution not in RESOLUTIONS:
+        raise HTTPException(400, f"resolution: {' | '.join(RESOLUTIONS.keys())}")
+
+    denoise_steps = max(20, min(40, denoise_steps))
+    if seed == -1:
+        import random
+        seed = random.randint(0, 999999)
+
+    person_bytes  = await person_image.read()
+    garment_bytes = await garment_image.read()
+
+    if len(person_bytes) > 20 * 1024 * 1024:
+        raise HTTPException(413, "person_image: max 20 MB.")
+    if len(garment_bytes) > 20 * 1024 * 1024:
+        raise HTTPException(413, "garment_image: max 20 MB.")
+
+    t0 = time.time()
+    log.info(f"/tryon — category={category} resolution={resolution} steps={denoise_steps}")
+
+    person_path  = _validate_and_save(person_bytes,  "person_image")
+    garment_path = _validate_and_save(garment_bytes, "garment_image")
+
+    try:
+        # Step 1: IDM-VTON try-on
+        log.info("Step 1: IDM-VTON try-on…")
+        tryon_result = run_tryon(
+            person_path   = person_path,
+            garment_path  = garment_path,
+            category      = category,
+            denoise_steps = denoise_steps,
+            seed          = seed,
+        )
+
+        # Step 2: Upscale to chosen resolution
+        log.info(f"Step 2: Upscaling to {resolution}…")
+        final_img = upscale_image(tryon_result, resolution)
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        log.error(f"TryOn pipeline failed: {e}")
+        raise HTTPException(500, f"Try-on failed: {str(e)}. Please retry.")
+    finally:
+        import os as _os
+        for p in [person_path, garment_path]:
+            try: _os.unlink(p)
+            except: pass
+
+    elapsed  = round(time.time() - t0, 2)
+    out_w, out_h = final_img.size
+    log.info(f"/tryon done — {elapsed}s, output={out_w}×{out_h}")
+
+    return Response(
+        content    = _pil_to_bytes(final_img),
+        media_type = "image/png",
+        headers    = {
+            "X-Processing-Time": f"{elapsed}s",
+            "X-Engine":          "IDM-VTON + Real-ESRGAN",
+            "X-Category":        category,
+            "X-Resolution":      resolution,
+            "X-Output-Size":     f"{out_w}x{out_h}",
+        },
+    )

requirements-8.txt

@@ -0,0 +1,6 @@
+fastapi==0.111.0
+uvicorn==0.30.1
+python-multipart==0.0.9
+Pillow==10.3.0
+gradio_client==0.17.0
+requests==2.32.3