Refactor code structure for improved readability and maintainability

workers/colorize/Dockerfile

@@ -0,0 +1,24 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_NO_CACHE_DIR=1 \
+    PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        python3 python3-pip python3-dev \
+        libgl1 libglib2.0-0 \
+        git ca-certificates \
+    && rm -rf /var/lib/apt/lists/* \
+    && ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+COPY workers/colorize/requirements.txt .
+RUN pip3 install --upgrade pip && pip3 install -r requirements.txt
+# Install runpod LAST so modelscope's dependency pins can't break `import runpod`
+# at worker startup (the cause of "worker never becomes ready, no logs").
+RUN pip3 install runpod==1.7.9
+
+COPY workers/colorize/handler.py .
+
+CMD ["python3", "-u", "handler.py"]

workers/colorize/README.md

@@ -0,0 +1,33 @@
+# DDColor Colorization Worker (#8)
+
+RunPod Serverless worker that colorizes black-and-white images using the
+ModelScope DDColor pipeline (`damo/cv_ddcolor_image-colorization`).
+
+## Contract
+
+Input:
+```json
+{ "image": "<base64 (optionally data:...;base64, prefixed)>", "input_size": 512 }
+```
+
+Output:
+```json
+{ "image": "<base64 PNG, same dimensions as input>" }
+```
+
+On error: `{ "error": "..." }`.
+
+## Deploy
+
+- Dockerfile path: `/workers/colorize/Dockerfile`
+- Build context: repo root (`docker build -f workers/colorize/Dockerfile .`)
+- GPU: medium (RTX 3080 / 3090, ~10-12 GB VRAM)
+- App env var to set: `RUNPOD_COLORIZE_URL` (point your app at this endpoint)
+
+## Notes
+
+- Model weights are cached to the network volume via `MODELSCOPE_CACHE`
+  (`/runpod-volume/modelscope`) when `/runpod-volume` is mounted, so the first
+  request downloads and subsequent cold starts reuse the cached weights.
+- `input_size` controls the DDColor internal working resolution (default 512).
+  The output is always resized back to the exact input dimensions.

workers/colorize/handler.py

@@ -0,0 +1,80 @@
+_V = "/runpod-volume"
+import os
+if os.path.isdir(_V):
+    os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
+    os.environ.setdefault("MODELSCOPE_CACHE", os.path.join(_V, "modelscope"))
+
+import base64
+import runpod
+
+_pipeline = None
+
+
+def _get_pipeline():
+    global _pipeline
+    if _pipeline is None:
+        from modelscope.pipelines import pipeline
+        from modelscope.utils.constant import Tasks
+        _pipeline = pipeline(
+            Tasks.image_colorization,
+            model="damo/cv_ddcolor_image-colorization",
+        )
+    return _pipeline
+
+
+def _decode_image(b64):
+    if "," in b64:
+        b64 = b64.split(",", 1)[1]
+    return base64.b64decode(b64)
+
+
+def handler(event):
+    try:
+        import numpy as np
+        import cv2
+
+        inp = (event or {}).get("input") or {}
+        img_b64 = inp.get("image")
+        if not img_b64:
+            return {"error": "ValueError: missing 'image' in input"}
+
+        try:
+            input_size = int(inp.get("input_size", 512))
+        except (TypeError, ValueError):
+            input_size = 512
+
+        raw = _decode_image(img_b64)
+        arr = np.frombuffer(raw, dtype=np.uint8)
+        # Decode to BGR (modelscope DDColor expects BGR numpy array).
+        bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+        if bgr is None:
+            return {"error": "ValueError: could not decode input image"}
+
+        orig_h, orig_w = bgr.shape[:2]
+
+        pipe = _get_pipeline()
+
+        from modelscope.outputs import OutputKeys
+        result = pipe(bgr, input_size=input_size)
+        out_bgr = result[OutputKeys.OUTPUT_IMG]
+        out_bgr = np.asarray(out_bgr)
+
+        # Ensure output matches the input dimensions exactly.
+        if out_bgr.shape[0] != orig_h or out_bgr.shape[1] != orig_w:
+            out_bgr = cv2.resize(
+                out_bgr, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4
+            )
+
+        # BGR -> RGB, then encode PNG. cv2.imencode expects BGR, so re-convert.
+        out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
+        ok, buf = cv2.imencode(".png", cv2.cvtColor(out_rgb, cv2.COLOR_RGB2BGR))
+        if not ok:
+            return {"error": "RuntimeError: PNG encoding failed"}
+
+        out_b64 = base64.b64encode(buf.tobytes()).decode("utf-8")
+        return {"image": out_b64}
+    except Exception as exc:
+        return {"error": f"{type(exc).__name__}: {exc}"}
+
+
+runpod.serverless.start({"handler": handler})

workers/colorize/requirements.txt

@@ -0,0 +1,9 @@
+# NOTE: runpod is installed LAST in the Dockerfile (after modelscope), so
+# modelscope's aggressive dependency pins can't break `import runpod` at startup.
+torch==2.1.2
+torchvision==0.16.2
+modelscope==1.9.5
+opencv-python-headless==4.9.0.80
+pillow==10.2.0
+numpy==1.24.4
+timm==0.9.12

workers/colorize/test.json

@@ -0,0 +1 @@
+{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}