Upload app.py

app.py

@@ -34,13 +34,18 @@ import gradio as gr
 
 SD_INPAINT_ID = "runwayml/stable-diffusion-inpainting"
 CONTROLNET_ID = "lllyasviel/control_v11f1p_sd15_depth"   # -> hr16/ControlNet-HandRefiner-pruned for best
+TILE_CN_ID = "lllyasviel/control_v11f1e_sd15_tile"        # detail-regeneration ControlNet
+SD_BASE_ID = "runwayml/stable-diffusion-v1-5"             # base SD for img2img detail pass
 MESHGRAPHORMER_ID = "hr16/ControlNet-HandRefiner-pruned"
 MAX_SIDE = 768
+DETAIL_MAX_SIDE = 1280   # detail pass can work larger since it's tiled-friendly
 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"
+DETAIL_NEG = "blurry, soft, out of focus, jpeg artifacts, low quality, smudged, messy lines"
 
 _PIPE = None
 _MESH = None
+_DETAIL = None
 
 def _load():
     """Load on CPU at import time. Models are moved to GPU inside the @spaces.GPU call,
@@ -74,6 +79,32 @@ try:
 except Exception as _e:
     print("[load] preload deferred:", _e, flush=True)
 
+def _load_detail():
+    """Tile-ControlNet img2img pipeline for detail/lineart recovery. Loaded lazily on CPU."""
+    global _DETAIL
+    if _DETAIL is not None:
+        return
+    import time
+    from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler
+    t0 = time.time()
+    print("[load] detail pipeline (tile CN) on CPU…", flush=True)
+    tile = ControlNetModel.from_pretrained(TILE_CN_ID, torch_dtype=torch.float16)
+    pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+        SD_BASE_ID, controlnet=tile, torch_dtype=torch.float16, safety_checker=None
+    )
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    try: pipe.enable_attention_slicing()
+    except Exception as e: print("[load] attn-slicing skip:", e, flush=True)
+    try: pipe.enable_vae_tiling()
+    except Exception as e: print("[load] vae-tiling skip:", e, flush=True)
+    _DETAIL = pipe
+    print(f"[load] detail pipeline ready ({time.time()-t0:.0f}s)", flush=True)
+
+def _fit_to(img, max_side):
+    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)
+
 def _fit(img):
     w, h = img.size
     s = min(1.0, MAX_SIDE / max(w, h))
@@ -126,19 +157,66 @@ def fix_hands(image, mask_layers, prompt, strength):
         print("[fix] ERROR:\n" + traceback.format_exc(), flush=True)
         raise gr.Error(f"Fix failed: {e}")
 
-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")
+@spaces.GPU(duration=120)
+def detail_pass(image, strength, scale):
+    """Detail/lineart recovery via Tile-ControlNet img2img at low denoise.
+    Regenerates real detail and clean lines while preserving composition + style.
+    No prompt is used (per ControlNet-tile guidance) so it can't redraw the subject."""
+    import time, traceback
+    if image is None:
+        raise gr.Error("Upload an image first.")
+    try:
+        t0 = time.time()
+        _load_detail()
+        _DETAIL.to("cuda")
+        src = image["background"] if isinstance(image, dict) else image
+        src = src.convert("RGB")
+        # optionally enlarge first (Lanczos) — the model then fills in real detail at the higher res
+        scale = float(scale)
+        if scale > 1.01:
+            src = src.resize((int(src.width*scale), int(src.height*scale)), Image.LANCZOS)
+        work, (ow, oh) = _fit_to(src, DETAIL_MAX_SIDE)
+        print(f"[detail] working at {work.size}, denoise={strength}", flush=True)
+        # tile controlnet uses the image itself as the control signal
+        out = _DETAIL(
+            prompt="", negative_prompt=DETAIL_NEG,
+            image=work, control_image=work,
+            num_inference_steps=30, strength=float(strength),
+            guidance_scale=6.0, controlnet_conditioning_scale=1.0,
+        ).images[0]
+        if out.size != (ow, oh):
+            out = out.resize((ow, oh), Image.LANCZOS)
+        print(f"[detail] done, total {time.time()-t0:.0f}s", flush=True)
+        return out
+    except Exception as e:
+        print("[detail] ERROR:\n" + traceback.format_exc(), flush=True)
+        raise gr.Error(f"Detail pass failed: {e}")
+
+with gr.Blocks(title="DARKROOM", theme=gr.themes.Base()) as demo:
+    gr.Markdown("## 🎨 DARKROOM\nAI-art repair on GPU. **Fix hands** regenerates malformed hands "
+                "with correct geometry. **Add detail** uses Tile-ControlNet img2img to recover real "
+                "sharpness and clean lineart while keeping your original style.")
+    with gr.Tab("Fix hands"):
+        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")
+    with gr.Tab("Add detail"):
+        with gr.Row():
+            with gr.Column():
+                dinp = gr.Image(type="pil", label="Image to sharpen / add detail")
+                dstrength = gr.Slider(0.15, 0.6, value=0.3, step=0.05,
+                    label="Detail strength (low = safe & on-style, high = more new detail / more drift)")
+                dscale = gr.Slider(1.0, 2.0, value=1.0, step=0.5, label="Enlarge first (×)")
+                dbtn = gr.Button("Add detail", variant="primary")
+            with gr.Column():
+                dout = gr.Image(type="pil", label="Result")
+        dbtn.click(detail_pass, inputs=[dinp, dstrength, dscale], outputs=dout, api_name="detail_pass")
 
 if __name__ == "__main__":
     demo.queue().launch()