Merge pull request #58 from some1galaxy/tensor-only-interpolation-upscale-pipeline

app.py

@@ -109,7 +109,7 @@ def _gpu_duration(
     return duration
 
 
-def _apply_interpolation(raw_frames, frame_factor, progress):
+def _apply_interpolation(raw_frames, frame_factor, upscale_output, progress):
     if frame_factor <= 1:
         return list(raw_frames), FIXED_FPS
 
@@ -119,8 +119,13 @@ def _apply_interpolation(raw_frames, frame_factor, progress):
     print_stage_start("interpolation")
     t0 = _time.perf_counter()
     try:
+        # When upscale will run right after, hand it GPU tensors directly instead of round-
+        # tripping through CPU numpy in between (both stages run in the same @spaces.GPU call).
         interpolated = interpolate_frames(
-            raw_frames, multiplier=int(frame_factor), progress_callback=_report_interpolation_progress
+            raw_frames,
+            multiplier=int(frame_factor),
+            progress_callback=_report_interpolation_progress,
+            as_tensor=upscale_output,
         )
         print_stage_done("interpolation", _time.perf_counter() - t0)
         return interpolated, FIXED_FPS * frame_factor
@@ -131,6 +136,16 @@ def _apply_interpolation(raw_frames, frame_factor, progress):
         return list(raw_frames), FIXED_FPS
 
 
+def _frames_to_numpy(frames):
+    # Fallback-path conversion only: interpolation may have handed back GPU tensors (C,H,W)
+    # expecting upscale to consume them next (see _apply_interpolation's as_tensor). If upscale
+    # then fails or wasn't requested, callers downstream (export_to_video, logging) need plain
+    # (H,W,C) numpy — this is a no-op when frames are already numpy.
+    if frames and isinstance(frames[0], torch.Tensor):
+        return [f.permute(1, 2, 0).float().cpu().numpy() for f in frames]
+    return frames
+
+
 def _apply_upscale(frames, progress):
     from postprocess.upscale import upscale_frames
 
@@ -196,11 +211,17 @@ def run_inference(
     raw_frames = result.frames[0]  # (T, H, W, C) float32 in [0, 1]
 
     frame_factor = frame_multiplier // FIXED_FPS
-    final_frames, final_fps = _apply_interpolation(raw_frames, frame_factor, progress)
+    final_frames, final_fps = _apply_interpolation(raw_frames, frame_factor, upscale_output, progress)
+    # Confirms which handoff mode actually ran: frame_type=Tensor (GPU, no CPU round trip) when
+    # interpolation->upscale both ran, vs. ndarray otherwise. See issue #53.
+    print_frames_info("interpolation_output", final_frames, final_fps)
 
     upscaled_frames = _apply_upscale(final_frames, progress) if upscale_output else None
 
     if upscaled_frames is None:
+        # Covers both upscale_output=False (already numpy, no-op) and upscale_output=True but
+        # upscale failing (final_frames may be GPU tensors from the as_tensor interpolation path).
+        final_frames = _frames_to_numpy(final_frames)
         print_infer_done(_time.perf_counter() - t_start, final_fps, len(final_frames))
         return final_frames, final_fps
 

postprocess/interpolation.py

@@ -60,15 +60,21 @@ def interpolate_frames(
     multiplier: int,
     scale: float = 1.0,
     progress_callback: ProgressCallback | None = None,
-) -> list[np.ndarray]:
+    as_tensor: bool = False,
+) -> list[np.ndarray] | list[torch.Tensor]:
     """Interpolate frames with RIFE to `multiplier`x the input frame rate.
 
     Args:
         frames_np: (T, H, W, C) array or list of (H, W, C) arrays, float32 in [0, 1].
         multiplier: 2, 4, or 8. Values < 2 return the frames unchanged (as a list).
         progress_callback: progress_callback(done, total) fired after each input-frame gap.
+        as_tensor: if True, skip the GPU->CPU conversion and return (C, H, W) fp16 GPU
+            tensors instead of (H, W, C) float32 numpy arrays — for callers (e.g. upscale)
+            that consume the result on GPU right after, so frames never round-trip to CPU
+            in between.
     Returns:
-        List of (H, W, C) float32 numpy arrays in [0, 1].
+        List of (H, W, C) float32 numpy arrays in [0, 1], or (if as_tensor) list of
+        (C, H, W) fp16 GPU tensors in [0, 1].
     """
     if multiplier < 2:
         return list(frames_np) if isinstance(frames_np, np.ndarray) else frames_np
@@ -88,11 +94,18 @@ def interpolate_frames(
         t = t.permute(2, 0, 1).unsqueeze(0)
         return F.pad(t, padding).half()
 
-    def from_tensor(tensor: torch.Tensor) -> np.ndarray:
+    def from_tensor_np(tensor: torch.Tensor) -> np.ndarray:
         t = tensor[0, :, :height, :width]
         t = t.permute(1, 2, 0)
         return t.float().cpu().numpy()
 
+    def from_tensor_gpu(tensor: torch.Tensor) -> torch.Tensor:
+        # .contiguous() detaches from the padded working tensor's storage (a plain crop is
+        # just a view into it) so the padded buffer isn't kept alive for the batch's lifetime.
+        return tensor[0, :, :height, :width].contiguous()
+
+    from_tensor = from_tensor_gpu if as_tensor else from_tensor_np
+
     def make_inference(i0: torch.Tensor, i1: torch.Tensor, n: int) -> list[torch.Tensor]:
         if model.version >= 3.9:
             return [model.inference(i0, i1, (i + 1) / (n + 1), scale) for i in range(n)]

postprocess/upscale/upscale.py

@@ -146,19 +146,29 @@ def _tile_process(model: torch.nn.Module, img: torch.Tensor) -> torch.Tensor:
     return output
 
 
-def _frames_to_tensor(frames: list[np.ndarray]) -> torch.Tensor:
+def _frame_hw(frame: np.ndarray | torch.Tensor) -> tuple[int, int]:
+    # numpy frames are (H, W, C); GPU tensor frames (from RIFE's as_tensor=True path) are (C, H, W).
+    return tuple(frame.shape[-2:]) if isinstance(frame, torch.Tensor) else tuple(frame.shape[:2])
+
+
+def _frames_to_tensor(frames: list[np.ndarray] | list[torch.Tensor]) -> torch.Tensor:
+    if isinstance(frames[0], torch.Tensor):
+        # Already (C, H, W) fp16 GPU tensors (RIFE's as_tensor=True output) — just batch them,
+        # no host round trip.
+        return torch.stack(frames, dim=0).to(device=device, dtype=dtype)
     array = np.stack(frames)
     return torch.from_numpy(array).permute(0, 3, 1, 2).to(device=device, dtype=dtype)
 
 
 def upscale_frames(
-    frames: list[np.ndarray], progress_callback: ProgressCallback | None = None
+    frames: list[np.ndarray] | list[torch.Tensor], 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 either numpy float32 arrays in `[0,1]`, shape `(H,W,C)`, or (C,H,W) fp16 GPU
+    tensors — the same representations `_apply_interpolation` can produce, so callers don't
+    need to round-trip through `PIL.Image`, or through the CPU at all, 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
@@ -172,10 +182,10 @@ def upscale_frames(
 
     i = 0
     while i < len(frames):
-        h, w = frames[i].shape[:2]
+        h, w = _frame_hw(frames[i])
         batch = [frames[i]]
         i += 1
-        while i < len(frames) and len(batch) < FRAME_BATCH_SIZE and frames[i].shape[:2] == (h, w):
+        while i < len(frames) and len(batch) < FRAME_BATCH_SIZE and _frame_hw(frames[i]) == (h, w):
             batch.append(frames[i])
             i += 1