app-2.py

"""
DARKROOM HandRefiner — Hugging Face ZeroGPU Space
=================================================
Standard Gradio Interface (the pattern ZeroGPU actually supports): upload an
image, optionally paint a mask, get the hands structurally fixed on a free
on-demand GPU. This is the reliable shape — the previous "custom FastAPI route"
build failed with "No @spaces.GPU function detected" because ZeroGPU only
detects GPU functions wired into a normal Gradio app.

PIPELINE: MeshGraphormer hand-mesh -> depth map -> depth ControlNet ->
Stable Diffusion inpainting (HandRefiner). Fixes only the hand region.

--------------------------------------------------------------------------
DEPLOY  (needs a HF PRO account to CREATE a ZeroGPU Space — $9/mo)
--------------------------------------------------------------------------
1. huggingface.co -> New Space -> SDK: Gradio -> Hardware: ZeroGPU
2. Upload: app.py, requirements.txt, README.md
3. Wait for build, then use the Space UI (or call it from the DARKROOM tool
   via the gradio_client endpoint shown on the Space's "View API" page).

HONEST LIMITS:
* Creating a ZeroGPU Space requires PRO. Using one is free within a daily quota
  (resets 24h after first use); each fix is a few GPU-seconds.
* GPU duration is capped (~120s max). We request 90s.
* Stock depth ControlNet is okay-not-perfect; swap CONTROLNET_ID to
  hr16/ControlNet-HandRefiner-pruned for finetuned quality.
* MeshGraphormer can't fix unreadable hands or crossed fingers.
"""

import spaces                       # must precede torch for ZeroGPU
import torch
from PIL import Image, ImageFilter
import gradio as gr

SD_INPAINT_ID = "runwayml/stable-diffusion-inpainting"
CONTROLNET_ID = "lllyasviel/control_v11f1p_sd15_depth"   # -> hr16/ControlNet-HandRefiner-pruned for best
MESHGRAPHORMER_ID = "hr16/ControlNet-HandRefiner-pruned"
MAX_SIDE = 768
DEFAULT_PROMPT = "a detailed, anatomically correct hand with five fingers, natural proportions, same art style and lighting"
NEG = "extra fingers, fused fingers, missing fingers, deformed, mutated, blurry, low quality"

_PIPE = None
_MESH = None

def _load():
    global _PIPE, _MESH
    if _PIPE is not None:
        return
    from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
    from controlnet_aux import MeshGraphormerDetector
    _MESH = MeshGraphormerDetector.from_pretrained(MESHGRAPHORMER_ID).to("cuda")
    cn = ControlNetModel.from_pretrained(CONTROLNET_ID, torch_dtype=torch.float16)
    pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
        SD_INPAINT_ID, controlnet=cn, torch_dtype=torch.float16, safety_checker=None
    )
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
    _PIPE = pipe.to("cuda")

def _fit(img):
    w, h = img.size
    s = min(1.0, MAX_SIDE / max(w, h))
    return img.resize((max(8, int(round(w*s/8))*8), max(8, int(round(h*s/8))*8)), Image.LANCZOS), (w, h)

@spaces.GPU(duration=90)
def fix_hands(image, mask_layers, prompt, strength):
    """ZeroGPU-allocated worker, wired directly into the Gradio Interface below."""
    if image is None:
        raise gr.Error("Upload an image first.")
    _load()
    init, (ow, oh) = _fit(image.convert("RGB"))
    W, H = init.size

    # optional hand-drawn mask from the ImageMask/Sketchpad component
    sent_mask = None
    if isinstance(mask_layers, dict):
        layers = mask_layers.get("layers") or []
        if layers:
            m = layers[0].convert("L").resize((W, H), Image.LANCZOS)
            if m.getbbox() is not None:
                sent_mask = m

    mg = _MESH(init)
    depth_img, auto_mask = (mg[0], (mg[1] if len(mg) > 1 else None)) if isinstance(mg, tuple) else (mg, None)
    depth_img = depth_img.convert("RGB").resize((W, H), Image.LANCZOS)
    mask_img = sent_mask or (auto_mask.convert("L").resize((W, H), Image.LANCZOS) if auto_mask else None)
    if mask_img is None:
        raise gr.Error("No hands detected. Paint a mask over the hand and try again.")

    mask_img = mask_img.filter(ImageFilter.GaussianBlur(2))
    out = _PIPE(
        prompt=prompt or DEFAULT_PROMPT, negative_prompt=NEG, image=init, mask_image=mask_img,
        control_image=depth_img, num_inference_steps=30, strength=float(strength),
        guidance_scale=7.5, controlnet_conditioning_scale=0.7,
    ).images[0]
    return out.resize((ow, oh), Image.LANCZOS)

with gr.Blocks(title="DARKROOM HandRefiner", theme=gr.themes.Base()) as demo:
    gr.Markdown("## 🖐️ DARKROOM HandRefiner\nUpload AI art with bad hands. It auto-detects hands "
                "(MeshGraphormer) and regenerates them with correct geometry. Optionally paint a mask "
                "to target a specific hand. Free GPU runs a few seconds per fix.")
    with gr.Row():
        with gr.Column():
            inp = gr.ImageMask(type="pil", label="Image (optionally paint over the bad hand)")
            prompt = gr.Textbox(value=DEFAULT_PROMPT, label="Prompt", lines=2)
            strength = gr.Slider(0.3, 1.0, value=0.75, step=0.05, label="Fix strength (denoise)")
            btn = gr.Button("Fix hands", variant="primary")
        with gr.Column():
            out = gr.Image(type="pil", label="Result")
    btn.click(fix_hands, inputs=[inp, inp, prompt, strength], outputs=out, api_name="fix_hands")

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