first commit

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+.DS_Store
+.gradio/
+.hf_cache/
+.cache/
+venv/
+.venv/

README.md

@@ -0,0 +1,38 @@
+---
+title: RealRestorer Demo
+emoji: 🖼️
+colorFrom: teal
+colorTo: orange
+sdk: gradio
+app_file: app.py
+pinned: false
+python_version: 3.10
+---
+
+# RealRestorer Hugging Face Space
+
+This Space runs the public `RealRestorer/RealRestorer` model with a Gradio interface for single-image restoration.
+
+It is designed as a lightweight Space repo:
+
+- the model weights are loaded from Hugging Face
+- the custom RealRestorer-enabled `diffusers` build is installed from the main GitHub repo
+- the UI layout follows the existing internal Gradio demo, but is simplified for Spaces deployment
+
+## Notes
+
+- A GPU Space is strongly recommended. The released model is not practical on CPU.
+- The default model repo is `RealRestorer/RealRestorer`.
+- You can override the model repo with the `REALRESTORER_MODEL_REPO` Space environment variable.
+
+## Local Run
+
+```bash
+python -m pip install -r requirements.txt
+python app.py
+```
+
+## Space Variables
+
+- `REALRESTORER_MODEL_REPO`: optional model repo id, default `RealRestorer/RealRestorer`
+- `HF_TOKEN`: optional token if you want to load a private model repo

app.py

@@ -0,0 +1,414 @@
+from __future__ import annotations
+
+import base64
+import io
+import os
+import random
+import threading
+import time
+import traceback
+from typing import Optional
+
+import gradio as gr
+import torch
+from PIL import Image
+
+MODEL_REPO_ID = os.environ.get("REALRESTORER_MODEL_REPO", "RealRestorer/RealRestorer")
+HF_TOKEN = os.environ.get("HF_TOKEN") or None
+
+TASK_PRESETS = {
+    "Low-light Enhancement": "Please restore this low-quality image, recovering its normal brightness and clarity.",
+    "Deblurring": "Please deblur the image and make it sharper.",
+    "Deraining": "Please remove the rain from the image and restore its clarity.",
+    "Compression Artifact Removal": "Please restore the image clarity and artifacts.",
+    "Deflare": "Please remove the lens flare and glare from the image.",
+    "Demoire": "Please remove the moire patterns from the image.",
+    "Dehazing": "Please dehaze the image.",
+    "Denoising": "Please remove noise from the image.",
+    "Reflection Removal": "Please remove the reflection from the image.",
+}
+
+DEFAULT_PRESET = "Reflection Removal"
+DEFAULT_STATUS = (
+    "Model not loaded yet. Upload an image, pick a preset or edit the prompt, "
+    "then click Run Inference."
+)
+DEFAULT_SLIDER = """
+<div style="text-align:center; padding:96px 24px; color:#7b8794; font-size:1.05rem;">
+  The interactive before/after slider will appear here after the first run.
+</div>
+"""
+
+CUSTOM_CSS = """
+:root {
+    --rr-ink: #102a43;
+    --rr-muted: #5c6b7a;
+    --rr-panel: rgba(255, 255, 255, 0.9);
+    --rr-edge: rgba(16, 42, 67, 0.12);
+    --rr-accent: #0f766e;
+    --rr-accent-2: #f59e0b;
+}
+
+body {
+    background:
+        radial-gradient(circle at top left, rgba(15, 118, 110, 0.12), transparent 30%),
+        radial-gradient(circle at top right, rgba(245, 158, 11, 0.12), transparent 28%),
+        linear-gradient(180deg, #f7fafc 0%, #eef2f7 100%);
+}
+
+.gradio-container {
+    max-width: 1380px !important;
+    margin: 0 auto !important;
+}
+
+.rr-hero {
+    padding: 20px 8px 10px;
+}
+
+.rr-hero h1 {
+    margin: 0;
+    text-align: center;
+    font-size: 2.55rem;
+    line-height: 1.05;
+    letter-spacing: -0.04em;
+    color: var(--rr-ink);
+}
+
+.rr-hero p {
+    margin: 12px auto 0;
+    max-width: 820px;
+    text-align: center;
+    color: var(--rr-muted);
+    font-size: 1rem;
+    line-height: 1.6;
+}
+
+.rr-shell {
+    border: 1px solid var(--rr-edge);
+    border-radius: 24px;
+    background: rgba(255, 255, 255, 0.7);
+    backdrop-filter: blur(16px);
+    box-shadow: 0 18px 48px rgba(15, 23, 42, 0.08);
+    padding: 18px;
+}
+
+.rr-note {
+    border-left: 4px solid var(--rr-accent);
+    padding: 12px 14px;
+    border-radius: 12px;
+    background: rgba(15, 118, 110, 0.08);
+    color: var(--rr-ink);
+    margin-bottom: 8px;
+}
+
+.rr-foot {
+    text-align: center;
+    color: var(--rr-muted);
+    font-size: 0.9rem;
+    padding: 10px 0 2px;
+}
+
+#run-btn {
+    background: linear-gradient(135deg, var(--rr-accent) 0%, #0b8a7b 55%, var(--rr-accent-2) 100%) !important;
+    color: white !important;
+    border: none !important;
+    font-weight: 700 !important;
+}
+"""
+
+PIPELINE = None
+PIPELINE_LOCK = threading.Lock()
+INFERENCE_LOCK = threading.Lock()
+PIPELINE_DEVICE = "cpu"
+PIPELINE_DTYPE = "float32"
+
+
+def _pick_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def _pick_dtype(device: str) -> tuple[torch.dtype, str]:
+    if device == "cuda":
+        if torch.cuda.is_bf16_supported():
+            return torch.bfloat16, "bfloat16"
+        return torch.float16, "float16"
+    return torch.float32, "float32"
+
+
+def _load_pipeline():
+    global PIPELINE, PIPELINE_DEVICE, PIPELINE_DTYPE
+
+    if PIPELINE is not None:
+        return PIPELINE
+
+    with PIPELINE_LOCK:
+        if PIPELINE is not None:
+            return PIPELINE
+
+        from diffusers import RealRestorerPipeline
+
+        device = _pick_device()
+        torch_dtype, dtype_name = _pick_dtype(device)
+        pipe = RealRestorerPipeline.from_pretrained(
+            MODEL_REPO_ID,
+            torch_dtype=torch_dtype,
+            token=HF_TOKEN,
+        )
+        if device == "cuda":
+            pipe.enable_model_cpu_offload(device=device)
+        else:
+            pipe.to(device)
+
+        PIPELINE = pipe
+        PIPELINE_DEVICE = device
+        PIPELINE_DTYPE = dtype_name
+        return PIPELINE
+
+
+def _resolve_seed(seed: float | int | None) -> int:
+    if seed is None:
+        return random.randint(0, 2**31 - 1)
+    seed_value = int(seed)
+    if seed_value < 0:
+        return random.randint(0, 2**31 - 1)
+    return seed_value
+
+
+def _placeholder_slider(message: str) -> str:
+    return (
+        "<div style='text-align:center; padding:72px 24px; color:#7b8794; font-size:1rem;'>"
+        f"{message}"
+        "</div>"
+    )
+
+
+def _pil_to_data_url(image: Image.Image) -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{encoded}"
+
+
+def _build_slider_html(before_image: Image.Image, after_image: Image.Image) -> str:
+    if before_image.size != after_image.size:
+        before_image = before_image.resize(after_image.size, Image.LANCZOS)
+
+    before_url = _pil_to_data_url(before_image)
+    after_url = _pil_to_data_url(after_image)
+    width, height = after_image.size
+    slider_id = f"rr_slider_{int(time.time() * 1000)}"
+    on_input = (
+        f"var p=this.value;"
+        f"document.getElementById('{slider_id}_top').style.clipPath='inset(0 '+(100-p)+'% 0 0)';"
+        f"document.getElementById('{slider_id}_line').style.left=p+'%';"
+    )
+
+    return f"""
+    <div style="display:flex; justify-content:space-between; align-items:end; gap:12px; margin-bottom:14px; color:#102a43;">
+      <div>
+        <div style="font-size:1.12rem; font-weight:700;">Before / After Comparison</div>
+        <div style="font-size:0.92rem; color:#5c6b7a;">Drag the slider to inspect restored details.</div>
+      </div>
+      <div style="font-size:0.92rem; color:#5c6b7a;">{width} x {height}px</div>
+    </div>
+    <div style="position:relative; width:100%; aspect-ratio:{width}/{height}; overflow:hidden; border-radius:18px; border:1px solid rgba(16, 42, 67, 0.12); box-shadow:0 18px 36px rgba(15, 23, 42, 0.12); background:white;">
+      <img src="{after_url}" style="position:absolute; inset:0; width:100%; height:100%; object-fit:cover;" draggable="false" />
+      <img id="{slider_id}_top" src="{before_url}" style="position:absolute; inset:0; width:100%; height:100%; object-fit:cover; clip-path:inset(0 50% 0 0);" draggable="false" />
+      <div id="{slider_id}_line" style="position:absolute; top:0; left:50%; width:3px; height:100%; background:white; box-shadow:0 0 10px rgba(0,0,0,0.35); transform:translateX(-50%);">
+        <div style="position:absolute; top:50%; left:50%; transform:translate(-50%, -50%); width:44px; height:44px; border-radius:999px; background:white; display:flex; align-items:center; justify-content:center; box-shadow:0 10px 24px rgba(15, 23, 42, 0.18); color:#102a43; font-weight:700;">
+          &#8596;
+        </div>
+      </div>
+    </div>
+    <div style="padding-top:18px;">
+      <input type="range" min="0" max="100" value="50" oninput="{on_input}" style="width:100%; cursor:ew-resize;" />
+      <div style="display:flex; justify-content:space-between; color:#486581; font-size:0.92rem; padding-top:6px;">
+        <span>Original</span>
+        <span>Restored</span>
+      </div>
+    </div>
+    """
+
+
+def _on_preset_change(preset_name: str) -> str:
+    return TASK_PRESETS.get(preset_name, "")
+
+
+def run_inference(
+    image: Optional[Image.Image],
+    preset_name: str,
+    prompt: str,
+    steps: float,
+    guidance_scale: float,
+    size_level: float,
+    seed: float,
+):
+    if image is None:
+        return None, "Please upload an input image first.", _placeholder_slider("Please upload an image to start.")
+
+    source_image = image.convert("RGB")
+    final_prompt = prompt.strip() or TASK_PRESETS.get(preset_name, TASK_PRESETS[DEFAULT_PRESET])
+    final_seed = _resolve_seed(seed)
+
+    try:
+        pipeline = _load_pipeline()
+        start_time = time.time()
+        with INFERENCE_LOCK:
+            output = pipeline(
+                image=source_image,
+                prompt=final_prompt,
+                num_inference_steps=int(steps),
+                guidance_scale=float(guidance_scale),
+                size_level=int(size_level),
+                seed=final_seed,
+            ).images[0]
+        elapsed = time.time() - start_time
+    except Exception as exc:
+        traceback.print_exc()
+        error_text = (
+            f"Inference failed: {exc}\n"
+            f"Model repo: {MODEL_REPO_ID}"
+        )
+        return None, error_text, _placeholder_slider("Inference failed. Check the Space logs for details.")
+
+    status = (
+        f"Done in {elapsed:.2f}s\n"
+        f"Model: {MODEL_REPO_ID}\n"
+        f"Device: {PIPELINE_DEVICE}\n"
+        f"Dtype: {PIPELINE_DTYPE}\n"
+        f"Steps: {int(steps)} | Guidance: {float(guidance_scale):.1f} | Size level: {int(size_level)}\n"
+        f"Seed: {final_seed}"
+    )
+    slider_html = _build_slider_html(source_image, output)
+    return output, status, slider_html
+
+
+def build_demo() -> gr.Blocks:
+    with gr.Blocks(css=CUSTOM_CSS, title="RealRestorer Demo") as demo:
+        gr.HTML(
+            """
+            <div class="rr-hero">
+              <h1>RealRestorer Image Restoration</h1>
+              <p>
+                Restore real-world degraded photos with the released RealRestorer model.
+                Upload one image, choose a restoration preset, and compare the result with the
+                interactive before/after slider.
+              </p>
+            </div>
+            """
+        )
+
+        with gr.Group(elem_classes=["rr-shell"]):
+            gr.HTML(
+                """
+                <div class="rr-note">
+                  This Space loads the public <b>RealRestorer/RealRestorer</b> model.
+                  A GPU Space is recommended because the released checkpoint is heavy.
+                </div>
+                """
+            )
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    input_image = gr.Image(
+                        label="Input Image",
+                        type="pil",
+                        height=360,
+                    )
+
+                with gr.Column(scale=1):
+                    preset_dropdown = gr.Dropdown(
+                        choices=list(TASK_PRESETS.keys()),
+                        value=DEFAULT_PRESET,
+                        label="Preset",
+                    )
+                    prompt_box = gr.Textbox(
+                        label="Instruction",
+                        value=TASK_PRESETS[DEFAULT_PRESET],
+                        lines=4,
+                    )
+
+                    with gr.Accordion("Advanced Settings", open=False):
+                        steps_slider = gr.Slider(
+                            minimum=12,
+                            maximum=40,
+                            value=28,
+                            step=1,
+                            label="Inference Steps",
+                        )
+                        guidance_slider = gr.Slider(
+                            minimum=1.0,
+                            maximum=6.0,
+                            value=3.0,
+                            step=0.1,
+                            label="Guidance Scale",
+                        )
+                        size_level_slider = gr.Slider(
+                            minimum=512,
+                            maximum=1280,
+                            value=1024,
+                            step=64,
+                            label="Resize Target",
+                        )
+                        seed_box = gr.Number(
+                            label="Seed (-1 for random)",
+                            value=-1,
+                            precision=0,
+                        )
+
+                    run_button = gr.Button("Run Inference", variant="primary", elem_id="run-btn")
+                    status_box = gr.Textbox(
+                        label="Status",
+                        value=DEFAULT_STATUS,
+                        lines=6,
+                        interactive=False,
+                    )
+
+            gr.Markdown("---")
+
+            with gr.Tabs():
+                with gr.Tab("Interactive Compare"):
+                    slider_html = gr.HTML(DEFAULT_SLIDER)
+                with gr.Tab("Restored Output"):
+                    output_image = gr.Image(label="Output Image", type="pil", interactive=False)
+
+            gr.HTML(
+                """
+                <div class="rr-foot">
+                  Prompts are fully editable. Presets are only shortcuts for common restoration tasks.
+                </div>
+                """
+            )
+
+        preset_dropdown.change(
+            fn=_on_preset_change,
+            inputs=[preset_dropdown],
+            outputs=[prompt_box],
+        )
+        run_button.click(
+            fn=run_inference,
+            inputs=[
+                input_image,
+                preset_dropdown,
+                prompt_box,
+                steps_slider,
+                guidance_slider,
+                size_level_slider,
+                seed_box,
+            ],
+            outputs=[output_image, status_box, slider_html],
+        )
+
+    return demo
+
+
+demo = build_demo()
+demo.queue(max_size=8, default_concurrency_limit=1)
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,13 @@
+accelerate>=1.0.0
+einops>=0.8.0
+gradio>=5.0.0,<6.0.0
+huggingface_hub>=0.36.0
+opencv-python>=4.10.0
+Pillow>=10.0.0
+safetensors>=0.5.0
+sentencepiece>=0.2.0
+timm>=1.0.0
+torch>=2.5.0
+torchvision>=0.20.0
+transformers>=4.57.0
+git+https://github.com/yfyang007/RealRestorer.git@main#subdirectory=diffusers