Merge pull request #57 from some1galaxy/remove-upscale-pil-roundtrip

Remove redundant numpy->PIL->numpy round trip between interpolation and upscale

app.py

@@ -9,7 +9,6 @@ import numpy as np
 import spaces
 import torch
 from diffusers.utils.export_utils import export_to_video
-from PIL import Image
 
 from logging_utils import (
     LogUploader,
@@ -142,10 +141,7 @@ def _apply_upscale(frames, progress):
     print_stage_start("upscaling")
     t0 = _time.perf_counter()
     try:
-        pil_frames = [
-            Image.fromarray((np.asarray(frame) * 255.0).round().astype(np.uint8)) for frame in frames
-        ]
-        upscaled = upscale_frames(pil_frames, progress_callback=_report_upscale_progress)
+        upscaled = upscale_frames(frames, progress_callback=_report_upscale_progress)
         print_stage_done("upscaling", _time.perf_counter() - t0)
         return upscaled
     except Exception as e:

postprocess/upscale/upscale.py

@@ -2,9 +2,10 @@
 
 Inference structure (pre-pad, tile, merge, crop) is adapted from xinntao/Real-ESRGAN's `RealESRGANer`
 (BSD-3-Clause License, see `LICENSES/REAL-ESRGAN-LICENSE`, `realesrgan/utils.py`), trimmed to the single code path
-this Space needs: one fixed compact model, RGB `PIL.Image` in and out (no alpha channel, no 16-bit, no cv2/BGR
-round-trip, no `dni` model blending, no CLI prefetch/IO-queue threads — none of those apply to frames coming
-straight out of the Wan 2.2 decode step).
+this Space needs: one fixed compact model, RGB frames in (numpy float32 `[0,1]`, `(H,W,C)` — the same
+representation `_apply_interpolation` already produces, so no PIL round trip is needed at the boundary) and
+`PIL.Image` out (no alpha channel, no 16-bit, no cv2/BGR round-trip, no `dni` model blending, no CLI
+prefetch/IO-queue threads — none of those apply to frames coming straight out of the Wan 2.2 decode step).
 
 The weights are `4xLSDIRCompact` (Philip Hofmann, CC BY 4.0, see `LICENSES/4xLSDIRCompact-LICENSE`) rather than
 Real-ESRGAN's own `realesr-general-x4v3`: same `SRVGGNetCompact` architecture (just `num_conv=16` instead of 32,
@@ -145,16 +146,20 @@ def _tile_process(model: torch.nn.Module, img: torch.Tensor) -> torch.Tensor:
     return output
 
 
-def _frames_to_tensor(frames: list[Image.Image]) -> torch.Tensor:
-    array = np.stack([np.asarray(f.convert("RGB"), dtype=np.float32) / 255.0 for f in frames])
+def _frames_to_tensor(frames: list[np.ndarray]) -> torch.Tensor:
+    array = np.stack(frames)
     return torch.from_numpy(array).permute(0, 3, 1, 2).to(device=device, dtype=dtype)
 
 
 def upscale_frames(
-    frames: list[Image.Image], progress_callback: ProgressCallback | None = None
+    frames: list[np.ndarray], progress_callback: ProgressCallback | None = None
 ) -> list[Image.Image]:
     """Upscale every frame 4x with `4xLSDIRCompact`, tiled to bound peak memory.
 
+    frames are numpy float32 arrays in `[0,1]`, shape `(H,W,C)` — the same representation
+    `_apply_interpolation` already produces, so callers don't need to round-trip through
+    `PIL.Image` just to cross this boundary.
+
     Frames are processed in batches (of up to FRAME_BATCH_SIZE, split wherever frame size
     changes) rather than one at a time, so the model sees a real batch dimension instead of
     batch=1 on every call. progress_callback(done, total), if given, fires after each frame
@@ -167,10 +172,10 @@ def upscale_frames(
 
     i = 0
     while i < len(frames):
-        w, h = frames[i].size
+        h, w = frames[i].shape[:2]
         batch = [frames[i]]
         i += 1
-        while i < len(frames) and len(batch) < FRAME_BATCH_SIZE and frames[i].size == (w, h):
+        while i < len(frames) and len(batch) < FRAME_BATCH_SIZE and frames[i].shape[:2] == (h, w):
             batch.append(frames[i])
             i += 1