Upload 4 files

app.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""
+Manga Light Colorizer - Gradio app (local ONNX inference)
+
+Runs entirely inside a Hugging Face Space (or locally) using onnxruntime.
+No external API is called: the v6 generator + SAM encoder ONNX models are
+loaded from the local `models/` folder and run directly.
+
+Models (auto-detected, relative to this script):
+    standalone/models/v6_generator.onnx
+    standalone/models/v6_sam_encoder.onnx
+
+Launch:
+    python app.py
+"""
+
+import sys
+import time
+from pathlib import Path
+
+import cv2
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+try:
+    import onnxruntime as ort
+except ImportError:
+    print("Error: onnxruntime not installed. Install with: pip install onnxruntime")
+    sys.exit(1)
+
+print(f"[startup] Python {sys.version}", flush=True)
+print(f"[startup] gradio version: {gr.__version__}", flush=True)
+print(f"[startup] onnxruntime version: {ort.__version__}", flush=True)
+
+
+# ============================================================================
+# CONFIG
+# ============================================================================
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+GENERATOR_PATH = SCRIPT_DIR / "models" / "v6_generator.onnx"
+SAM_PATH = SCRIPT_DIR / "models" / "v6_sam_encoder.onnx"
+EXAMPLES_DIR = SCRIPT_DIR / "input"
+
+INFER_SIZE_OPTIONS = [512, 768, 1024]
+DEFAULT_INFER_SIZE = 768
+
+
+# ============================================================================
+# CORE ONNX INFERENCE (ported from inference.py)
+# ============================================================================
+
+def denormalize_rgb(rgb_norm: np.ndarray) -> np.ndarray:
+    """[-1, 1] -> [0, 255] uint8."""
+    return np.clip((rgb_norm + 1.0) * 127.5, 0, 255).astype(np.uint8)
+
+
+def extract_sam_features_onnx(sam_session: ort.InferenceSession, L_bw_norm: np.ndarray):
+    """
+    Extract SAM features via ONNX. WD14 is intentionally DISABLED (zeros).
+
+    Args:
+        sam_session: ONNX Runtime session for SAM encoder
+        L_bw_norm: (H, W) grayscale in [-1, 1]
+
+    Returns:
+        sam_level0, sam_level1, wd14_embedding (all numpy)
+    """
+    L_01 = (L_bw_norm + 1.0) / 2.0  # [-1,1] -> [0,1]
+    L_1024 = cv2.resize(L_01, (1024, 1024), interpolation=cv2.INTER_LINEAR)
+    rgb_sam = np.stack([L_1024, L_1024, L_1024], axis=0)[np.newaxis].astype(np.float32)
+
+    sam_out = sam_session.run(None, {"rgb_input": rgb_sam})
+    sam_level0 = sam_out[0]  # (1, 256, 64, 64)
+    sam_level1 = sam_out[1]  # (1, 256, 32, 32)
+
+    wd14_embedding = np.zeros((1, 1024), dtype=np.float32)
+    return sam_level0, sam_level1, wd14_embedding
+
+
+def colorize_onnx(
+    session: ort.InferenceSession,
+    L_bw: np.ndarray,
+    sam_level0: np.ndarray,
+    sam_level1: np.ndarray,
+    wd14_embedding: np.ndarray,
+) -> np.ndarray:
+    """Run generator ONNX inference. Returns RGB (H, W, 3) in [0, 255]."""
+    L_norm = (L_bw.astype(np.float32) / 127.5) - 1.0
+    L_tensor = L_norm[np.newaxis, np.newaxis, :, :]  # (1, 1, H, W)
+
+    ort_inputs = {
+        "L_bw": L_tensor,
+        "sam_level0": sam_level0,
+        "sam_level1": sam_level1,
+        "wd14_embedding": wd14_embedding,
+    }
+
+    rgb_pred = session.run(None, ort_inputs)[0]  # (1, 3, H, W)
+    rgb_pred = rgb_pred[0].transpose(1, 2, 0)  # (H, W, 3)
+    return denormalize_rgb(rgb_pred)
+
+
+# ============================================================================
+# MODEL LOADING (once, at startup)
+# ============================================================================
+
+def load_sessions():
+    """Load generator (+ optional SAM) ONNX sessions. Prefers CUDA if available."""
+    available = ort.get_available_providers()
+    if "CUDAExecutionProvider" in available:
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    else:
+        providers = ["CPUExecutionProvider"]
+
+    if not GENERATOR_PATH.exists():
+        raise FileNotFoundError(f"Generator ONNX not found: {GENERATOR_PATH}")
+
+    print(f"[startup] Loading generator: {GENERATOR_PATH}", flush=True)
+    session = ort.InferenceSession(str(GENERATOR_PATH), providers=providers)
+    print(f"[startup] Generator provider: {session.get_providers()[0]}", flush=True)
+
+    sam_session = None
+    if SAM_PATH.exists():
+        print(f"[startup] Loading SAM encoder: {SAM_PATH}", flush=True)
+        sam_session = ort.InferenceSession(str(SAM_PATH), providers=providers)
+        print("[startup] SAM encoder loaded", flush=True)
+    else:
+        print("[startup] SAM encoder NOT found -> using zeros", flush=True)
+
+    return session, sam_session
+
+
+SESSION, SAM_SESSION = load_sessions()
+HAS_SAM = SAM_SESSION is not None
+
+
+# ============================================================================
+# GRADIO INFERENCE HANDLER
+# ============================================================================
+
+def colorize_image(input_image: Image.Image, infer_size: int):
+    """
+    Colorize a grayscale manga image using local ONNX models.
+
+    Args:
+        input_image: PIL Image (any mode).
+        infer_size: Square inference resolution.
+
+    Returns:
+        (colorized PIL Image or None, status message).
+    """
+    if input_image is None:
+        return None, "⚠️ Please upload an image first."
+
+    t_start = time.time()
+
+    # PIL -> grayscale numpy
+    gray = np.array(input_image.convert("L"))
+    orig_H, orig_W = gray.shape
+
+    infer_size = int(infer_size)
+
+    # Always resize input to infer_size for inference
+    L_bw = cv2.resize(gray, (infer_size, infer_size), interpolation=cv2.INTER_AREA)
+    H_in, W_in = L_bw.shape
+    L_norm = (L_bw.astype(np.float32) / 127.5) - 1.0
+
+    if HAS_SAM:
+        sam_level0, sam_level1, wd14_embedding = extract_sam_features_onnx(SAM_SESSION, L_norm)
+    else:
+        sam_level0 = np.zeros((1, 256, H_in // 16, W_in // 16), dtype=np.float32)
+        sam_level1 = np.zeros((1, 256, H_in // 32, W_in // 32), dtype=np.float32)
+        wd14_embedding = np.zeros((1, 1024), dtype=np.float32)
+
+    rgb_output = colorize_onnx(SESSION, L_bw, sam_level0, sam_level1, wd14_embedding)
+    # (infer_size, infer_size, 3) -> back to original input resolution
+    rgb_output = cv2.resize(rgb_output, (orig_W, orig_H), interpolation=cv2.INTER_LANCZOS4)
+
+    result = Image.fromarray(rgb_output)
+    elapsed = time.time() - t_start
+    status = (
+        f"✅ Colorization complete! "
+        f"({orig_W}×{orig_H} px, infer {infer_size}×{infer_size}, {elapsed:.2f}s)"
+    )
+    return result, status
+
+
+# ============================================================================
+# GRADIO UI
+# ============================================================================
+
+def collect_examples():
+    """Build example list from the input/ folder."""
+    examples = []
+    if EXAMPLES_DIR.is_dir():
+        for ext in ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp"):
+            for f in sorted(EXAMPLES_DIR.glob(ext)):
+                examples.append([str(f), DEFAULT_INFER_SIZE])
+    return examples
+
+
+def build_interface() -> gr.Blocks:
+    with gr.Blocks(
+        title="Manga Light Colorizer",
+        theme=gr.themes.Soft(),
+    ) as demo:
+        gr.Markdown(
+            """
+            # 🎨 Manga Light Colorizer
+            Upload a black-and-white manga image and let the AI bring it to life in color.
+
+            > Runs **fully locally** with ONNX Runtime — no external API call.
+            > The model was trained at **512×512**; the further the inference resolution
+            > differs from 512, the less faithful the colors may be.
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(label="Input Image", type="pil")
+                infer_size = gr.Radio(
+                    choices=INFER_SIZE_OPTIONS,
+                    value=DEFAULT_INFER_SIZE,
+                    label="Inference Resolution",
+                    info=(
+                        "Square resolution used for inference. Output is resized back "
+                        "to the original input resolution. 512 = best color fidelity."
+                    ),
+                )
+                colorize_btn = gr.Button("🎨 Colorize", variant="primary", size="lg")
+
+            with gr.Column(scale=1):
+                output_image = gr.Image(
+                    label="Colorized Output",
+                    type="pil",
+                    interactive=False,
+                )
+                status_text = gr.Textbox(label="Status", interactive=False, lines=2)
+
+        colorize_btn.click(
+            fn=colorize_image,
+            inputs=[input_image, infer_size],
+            outputs=[output_image, status_text],
+        )
+
+        examples = collect_examples()
+        if examples:
+            gr.Examples(
+                examples=examples,
+                inputs=[input_image, infer_size],
+                outputs=[output_image, status_text],
+                fn=colorize_image,
+                cache_examples=False,
+            )
+
+        gr.Markdown(
+            """
+            ---
+            ### 📝 Notes
+            - Supported input formats: **JPEG, PNG, WebP, BMP**.
+            - Inference runs locally via **ONNX Runtime** (CUDA if available, else CPU).
+            - Pipeline: `grayscale → resize → SAM encoder → generator → resize to original`.
+            """
+        )
+
+    return demo
+
+
+print("[startup] calling build_interface()...", flush=True)
+demo = build_interface()
+print(f"[startup] demo object created: {demo}", flush=True)
+
+if __name__ == "__main__":
+    print("[startup] running as __main__, calling demo.launch()", flush=True)
+    demo.launch()

packages.txt

@@ -0,0 +1 @@
+python3-opencv

requirements-cuda.txt

@@ -3,3 +3,7 @@
 onnxruntime-gpu>=1.16.0
 numpy>=1.24.0
 opencv-python>=4.8.0
+
+# Gradio app (app.py) — local ONNX inference on Hugging Face Spaces
+gradio>=6.0.0
+Pillow>=10.0.0

requirements.txt

@@ -3,3 +3,8 @@
 onnxruntime>=1.16.0
 numpy>=1.24.0
 opencv-python>=4.8.0
+
+# Gradio app (app.py) — local ONNX inference on Hugging Face Spaces
+# On HF Spaces, OpenCV also needs the Debian package python3-opencv (see packages.txt).
+gradio>=6.0.0
+Pillow>=10.0.0