Merge remote-tracking branch 'origin/develop' into add-mypy-strict-typecheck

# Conflicts:
# app.py
# postprocess/interpolation.py
# postprocess/upscale/upscale.py

app.py

@@ -5,7 +5,7 @@ import random
 import tempfile
 import threading
 import time as _time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import gradio as gr
 import numpy as np
@@ -42,9 +42,12 @@ from model.pipeline import (
     estimate_generation_seconds,
     get_num_frames,
     load_pipeline,
+    resize_and_crop_to_match,
     resize_image,
 )
+from postprocess.interpolation import ensure_weights_downloaded as ensure_rife_weights_downloaded
 from postprocess.interpolation import interpolate_frames
+from postprocess.upscale.upscale import ensure_weights_downloaded as ensure_upscale_weights_downloaded
 
 MAX_SEED = np.iinfo(np.int32).max
 FRAME_MULTIPLIER_CHOICES = [FIXED_FPS, FIXED_FPS * 2, FIXED_FPS * 4, FIXED_FPS * 8]
@@ -61,6 +64,12 @@ print_startup_env()
 
 pipe = load_pipeline()
 
+# RIFE/upscale weight fetches are CPU/network-only (no CUDA needed) — done here, at process
+# startup, so a cold container doesn't pay for them out of the metered @spaces.GPU allocation
+# on its first interpolation/upscale request.
+ensure_rife_weights_downloaded()
+ensure_upscale_weights_downloaded()
+
 # Debug logging (LOG-2): gated entirely by LOG_HF_TOKEN/LOG_DATASET_REPO — a silent no-op when
 # either is unset. Deliberately not the deploy HF_TOKEN: this uses its own write-scoped token so a
 # leak has a much smaller blast radius (append-only on one dataset repo, not the whole Space repo).
@@ -68,6 +77,9 @@ _log_uploader = LogUploader(
     token=os.environ.get("LOG_HF_TOKEN"),
     repo_id=os.environ.get("LOG_DATASET_REPO"),
     max_bytes=int(float(os.environ.get("LOG_STORAGE_CAP_GB", "10")) * 1024**3),
+    # Hub's commit endpoint rejects pushes once a directory holds >10000 files; stay well under
+    # that per-directory cap (each logged stem adds one file to data/, images/, and videos/).
+    max_files=int(os.environ.get("LOG_MAX_FILES", "8000")),
     batch_interval=int(os.environ.get("LOG_BATCH_INTERVAL", "60")),
 )
 
@@ -85,6 +97,7 @@ _COMMUNITY_MD = (
 
 def _gpu_duration(
     resized_image: Image.Image,
+    processed_last_image: Image.Image | None,
     prompt: str,
     negative_prompt: str,
     steps: int,
@@ -114,8 +127,8 @@ def _gpu_duration(
 
 
 def _apply_interpolation(
-    raw_frames: FrameArray, frame_factor: int, progress: gr.Progress
-) -> tuple[list[FrameArray], int]:
+    raw_frames: FrameArray, frame_factor: int, upscale_output: bool, progress: gr.Progress
+) -> tuple[list[FrameArray] | list[torch.Tensor], int]:
     if frame_factor <= 1:
         return list(raw_frames), FIXED_FPS
 
@@ -125,8 +138,13 @@ def _apply_interpolation(
     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
@@ -137,7 +155,19 @@ def _apply_interpolation(
         return list(raw_frames), FIXED_FPS
 
 
-def _apply_upscale(frames: list[FrameArray], progress: gr.Progress) -> list[Image.Image] | None:
+def _frames_to_numpy(frames: list[FrameArray] | list[torch.Tensor]) -> list[FrameArray]:
+    # 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  # type: ignore[return-value]  # narrowed by the isinstance check above
+
+
+def _apply_upscale(
+    frames: list[FrameArray] | list[torch.Tensor], progress: gr.Progress
+) -> list[Image.Image] | None:
     from postprocess.upscale import upscale_frames
 
     def _report_upscale_progress(done: int, total: int) -> None:
@@ -147,10 +177,7 @@ def _apply_upscale(frames: list[FrameArray], progress: gr.Progress) -> list[Imag
     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:
@@ -163,6 +190,7 @@ def _apply_upscale(frames: list[FrameArray], progress: gr.Progress) -> list[Imag
 @spaces.GPU(duration=_gpu_duration)  # type: ignore[untyped-decorator]  # spaces ships no type stubs
 def run_inference(
     resized_image: Image.Image,
+    processed_last_image: Image.Image | None,
     prompt: str,
     negative_prompt: str,
     steps: int,
@@ -172,15 +200,22 @@ def run_inference(
     frame_multiplier: int,
     upscale_output: bool,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
-) -> tuple[Frames, int, list[FrameArray] | None, int]:
+) -> tuple[Frames, int]:
     print_infer_start(prompt, negative_prompt, seed, steps, guidance_scale, frame_multiplier, upscale_output)
     t_start = _time.perf_counter()
 
+    def _report_generation_progress(
+        pipe_: Any, step_index: int, timestep: int, callback_kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        progress(0.7 * (step_index + 1) / int(steps), desc=f"Generating ({step_index + 1}/{int(steps)} steps)...")
+        return callback_kwargs
+
     print_stage_start("generation")
     t0 = _time.perf_counter()
     try:
         result = pipe(
             image=resized_image,
+            last_image=processed_last_image,
             prompt=prompt,
             negative_prompt=negative_prompt,
             height=resized_image.height,
@@ -190,6 +225,7 @@ def run_inference(
             num_inference_steps=int(steps),
             generator=torch.Generator(device="cuda").manual_seed(seed),
             output_type="np",
+            callback_on_step_end=_report_generation_progress,
         )
     except Exception as e:
         print_stage_error("generation", e)
@@ -198,23 +234,27 @@ 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, None, final_fps
+        return final_frames, final_fps
 
-    # LOG-4: the debug log records the pre-upscale (post-interpolation) result, not the final
-    # upscaled video and not the raw un-interpolated base output.
     print_infer_done(_time.perf_counter() - t_start, final_fps, len(upscaled_frames))
-    return upscaled_frames, final_fps, final_frames, final_fps
+    return upscaled_frames, final_fps
 
 
 def _spawn_log(
     input_image: Image.Image,
-    output_frames: list[FrameArray] | None,
+    output_frames: Frames | None,
     output_fps: int | None,
     prompt: str,
     negative_prompt: str,
@@ -247,6 +287,7 @@ def _spawn_log(
 
 def generate_video(
     input_image: Image.Image | None,
+    last_image: Image.Image | None,
     prompt: str,
     negative_prompt: str,
     duration_seconds: float,
@@ -262,6 +303,7 @@ def generate_video(
         raise gr.Error("Please upload an input image.")
 
     resized_image = resize_image(input_image)
+    processed_last_image = resize_and_crop_to_match(last_image, resized_image) if last_image is not None else None
     num_frames = get_num_frames(duration_seconds)
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
@@ -271,8 +313,9 @@ def generate_video(
 
     t0 = _time.perf_counter()
     try:
-        final_frames, final_fps, pre_upscale_frames, pre_upscale_fps = run_inference(
+        final_frames, final_fps = run_inference(
             resized_image,
+            processed_last_image,
             prompt,
             negative_prompt,
             steps,
@@ -299,7 +342,6 @@ def generate_video(
     duration = _time.perf_counter() - t0
 
     print_frames_info("final_frames", final_frames, final_fps)
-    print_frames_info("pre_upscale_frames", pre_upscale_frames, pre_upscale_fps)
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
@@ -318,9 +360,8 @@ def generate_video(
         # run_inference — a video was still produced, so this still counts as success=True.
         # interpolation_enabled/upscale_output reflect what the user requested, not whether an
         # internal fallback fired.
-        log_frames = pre_upscale_frames if pre_upscale_frames is not None else final_frames
         _spawn_log(
-            resized_image, log_frames, pre_upscale_fps, prompt, negative_prompt, current_seed,
+            resized_image, final_frames, final_fps, prompt, negative_prompt, current_seed,
             steps, guidance_scale, interpolation_enabled, interpolation_multiplier, upscale_output,
             resized_image.width, resized_image.height, duration, True,
         )
@@ -355,6 +396,10 @@ with gr.Blocks() as demo:
         "- Service availability and response times may vary as this Space runs on shared GPU infrastructure.\n"
         "- When using this Space, please comply with the "
         "[Hugging Face Content Policy](https://huggingface.co/content-policy).\n"
+        "- This Space and its outputs are provided **\"as is\"**, without warranties of any kind. You are "
+        "solely responsible for the content you generate and for how you use or share it; the operator "
+        "accepts no liability for any loss, damage, or claim arising from use of this Space or its "
+        "outputs.\n"
         f"- To monitor application performance and improve quality, input data (image/prompt) and generated "
         f"outputs are logged solely for debugging purposes and retained in a private Hugging Face dataset "
         f"accessible only to {_OPERATOR_MD}, automatically pruned once the dataset reaches its configured "
@@ -369,11 +414,11 @@ with gr.Blocks() as demo:
         gr.Markdown(
             f"**Space:** High Quality Video Generation  \n"
             f"**Operator:** {_OPERATOR_MD}  \n"
-            f"**Last updated:** 2026-08-30\n"
+            f"**Last updated:** 2026-09-18\n"
             "\n"
             "#### What data is collected\n"
             "When you submit a request, the following is logged: your uploaded input image, the generated "
-            "output video (pre-upscale, post-interpolation), your prompt and negative prompt, seed, inference "
+            "output video, your prompt and negative prompt, seed, inference "
             "steps, guidance scale, frame-interpolation and upscaling settings, output resolution/fps/duration, "
             "generation duration, and success/error status. No user account, IP address, or session data is "
             "collected.\n"
@@ -428,6 +473,11 @@ with gr.Blocks() as demo:
                 "interpolation, within the same GPU allocation, and adds to the GPU quota used.",
             )
             with gr.Accordion("Advanced Settings", open=False):
+                last_image_component = gr.Image(
+                    type="pil",
+                    label="Last Image (Optional) — end frame to animate towards",
+                    sources=["upload", "clipboard"],
+                )
                 negative_prompt_input = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=3)
                 steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps")
                 guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale")
@@ -443,6 +493,7 @@ with gr.Blocks() as demo:
         fn=generate_video,
         inputs=[
             input_image_component,
+            last_image_component,
             prompt_input,
             negative_prompt_input,
             duration_input,
@@ -458,11 +509,7 @@ with gr.Blocks() as demo:
 
     preview_inputs = [input_image_component, duration_input, frame_multiplier_input, upscale_checkbox]
     for component in preview_inputs:
-        # `component` collapses to the base Component type across this mixed list, and gradio's
-        # stubs only declare .change on the concrete subclasses, not the base.
-        component.change(  # type: ignore[attr-defined]
-            fn=preview_effect, inputs=preview_inputs, outputs=effect_preview
-        )
+        component.change(fn=preview_effect, inputs=preview_inputs, outputs=effect_preview)
     demo.load(fn=preview_effect, inputs=preview_inputs, outputs=effect_preview)
 
 if __name__ == "__main__":

logging_utils.py

@@ -3,8 +3,11 @@ Face Hub dataset repo (LOG-2 through LOG-9).
 
 Modeled on the FireRed-Image-Edit-1.0-Fast reference project's logging_utils.py/inference.py
 pipeline, adapted for video: logged media is referenced by path rather than embedded as bytes
-(LOG-5), the logged output video is the pre-upscale result (LOG-4), and retention is a
-total-storage-size cap rather than a file-count cap (LOG-7).
+(LOG-5), the logged output video is the same final result served to the user, and retention
+prunes on both a total-storage-size cap and a per-directory file-count cap (LOG-7) — the latter exists
+because the Hub commit endpoint rejects pushes once a directory (data/, images/, videos/)
+holds more than 10000 files, which a size-only cap doesn't bound since data/ and images/
+files are tiny compared to videos/.
 """
 
 import json
@@ -271,19 +274,26 @@ def _group_by_stem(existing: list[tuple[str, int]]) -> dict[str, dict[str, Any]]
     return groups
 
 
-def _build_delete_ops_by_size(existing: list[tuple[str, int]], new_batch_size: int,
-                               max_bytes: int) -> list[CommitOperationDelete]:
+def _build_delete_ops(existing: list[tuple[str, int]], new_batch: list[tuple[str, str]],
+                       max_bytes: int, max_files: int) -> list[CommitOperationDelete]:
     groups = _group_by_stem(existing)
-    total = sum(g["size"] for g in groups.values()) + new_batch_size
-    if max_bytes <= 0 or total <= max_bytes:
+    new_stems = {_stem_of(p) for p, _ in new_batch}
+    total_bytes = sum(g["size"] for g in groups.values()) + sum(os.path.getsize(local) for _, local in new_batch)
+    total_files = len(groups) + len(new_stems)
+
+    def _over_cap() -> bool:
+        return (max_bytes > 0 and total_bytes > max_bytes) or (max_files > 0 and total_files > max_files)
+
+    if not _over_cap():
         return []
     ops: list[CommitOperationDelete] = []
     for stem in sorted(groups):
-        if total <= max_bytes:
+        if not _over_cap():
             break
         group = groups[stem]
         ops.extend(CommitOperationDelete(path_in_repo=p) for p in group["paths"])
-        total -= group["size"]
+        total_bytes -= group["size"]
+        total_files -= 1
     return ops
 
 
@@ -313,10 +323,12 @@ def _delete_temp_files(paths: list[str]) -> None:
 
 
 class LogUploader:
-    def __init__(self, token: str | None, repo_id: str | None, max_bytes: int, batch_interval: int = 60) -> None:
+    def __init__(self, token: str | None, repo_id: str | None, max_bytes: int, max_files: int,
+                 batch_interval: int = 60) -> None:
         self._token = token
         self._repo_id = repo_id
         self._max_bytes = max_bytes
+        self._max_files = max_files
         self._batch_interval = batch_interval
         self._pending: list[tuple[str, str]] = []
         self._lock = threading.Lock()
@@ -412,9 +424,8 @@ class LogUploader:
         api = HfApi(token=self._token)
         api.create_repo(repo_id=self._repo_id, repo_type="dataset", private=True, exist_ok=True)
         existing = _list_existing_files_with_sizes(api, self._repo_id)
-        new_batch_size = sum(os.path.getsize(local) for _, local in batch)
         add_ops = [CommitOperationAdd(path_in_repo=p, path_or_fileobj=local) for p, local in batch]
-        del_ops = _build_delete_ops_by_size(existing, new_batch_size, self._max_bytes)
+        del_ops = _build_delete_ops(existing, batch, self._max_bytes, self._max_files)
         api.create_commit(
             repo_id=self._repo_id, repo_type="dataset",
             operations=[*add_ops, *del_ops],

model/pipeline.py

@@ -1,14 +1,21 @@
 """WAMU_v3 image-to-video pipeline loading and sizing helpers.
 
-Baseline (non-AOT) loading per the SRS's "baseline first, AOT after" sequencing: FR-6's
-AOT-compiled-transformer requirement is deferred until AOT-package compatibility with WAMU_v3
-is confirmed (open item C-2, tracked in issue #3). Weight quantization (below) is independent
-of AOT compilation and is applied now: the full bf16 pipe packs to ~69GB, which made ZeroGPU's
-per-call GPU init (loading the packed weights onto the GPU) exceed the requested call duration.
+Weight quantization (below) shrinks the full bf16 pipe from ~69GB, which made ZeroGPU's per-call
+GPU init (loading the packed weights onto the GPU) exceed the requested call duration. The
+transformers additionally get an AOT-compiled forward (FR-6) via `spaces.aoti_load()`, matching
+the reference project: this is what makes `Float8DynamicActivationFloat8WeightConfig` (fp8x8
+matmuls, both weights and activations quantized) safe to run. AOTInductor pre-plans a single
+memory pool for every op in the compiled graph, including that config's per-layer
+activation-quantize casts; run the same config eagerly (no AOT), as an earlier version of this
+file did, and every one of those casts is an unplanned CUDA allocator request instead — that
+churn reliably crashed ZeroGPU with an NVML assert (#35), fixed at the time by downgrading to
+weight-only fp8 (no runtime activation quantization) as a stopgap. See #36 for the AOT
+investigation (package-reuse feasibility) that unblocked switching back.
 """
 
 from __future__ import annotations
 
+import spaces
 import torch
 import torch._dynamo
 from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
@@ -17,6 +24,15 @@ from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, Int8
 
 MODEL_ID = "thornmaze/WAMU_v3_WAN2.2_I2V_LIGHTNING"
 
+# Precompiled AOT package for the transformer, from the reference project (zerogpu-aoti's
+# wan2-2-fp8da-aoti-faster Space). Confirmed reusable for WAMU_v3 without recompiling (#36):
+# WAMU_v3's transformer/transformer_2 share the exact same config (layer count, head dims, etc.)
+# as the Wan2.2-I2V-A14B base model this package was compiled against — `spaces.aoti_load()`
+# applies compiled kernels to an already-loaded model rather than embedding weights, so
+# compatibility hinges on that shape/architecture match, not on WAMU_v3's fine-tuned weight
+# values. "sm120" in the name matches ZeroGPU's current hardware (RTX Pro 6000 Blackwell).
+AOTI_TRANSFORMER_REPO = "thornmaze/WanTransformer3DModel-sm120-cu130-raa"
+
 MAX_DIM = 832
 MIN_DIM = 480
 SQUARE_DIM = 640
@@ -41,7 +57,6 @@ def load_pipeline() -> WanImageToVideoPipeline:
     pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
 
     # Quantize to shrink the ~69GB bf16 pipe (drops per-call ZeroGPU init time and peak memory).
-    # Independent of AOT compilation (FR-6/C-2, deferred) — same recipe as the reference project.
     quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
     torch._dynamo.reset()
     quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
@@ -49,6 +64,18 @@ def load_pipeline() -> WanImageToVideoPipeline:
     quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
     torch._dynamo.reset()
 
+    # AOT-compile the transformers (FR-6) — see module docstring and #36 for why this must come
+    # after the dynamic-activation quantize_() calls above, not before: the compiled graph is
+    # built from (and expects) that config's op layout, matching how the reference project's own
+    # usage example sequences it.
+    #
+    # Not yet validated against this repo's pinned torch==2.11.0: the package was compiled with
+    # torch==2.9.1+cu130, and AOTInductor artifacts can be ABI-sensitive across torch minor
+    # versions. Needs real-GPU testing on the dev Space — if loading fails here, the fallback is
+    # either pinning torch back to 2.9.1 or recompiling this package against 2.11.0.
+    spaces.aoti_load(module=pipe.transformer, repo_id=AOTI_TRANSFORMER_REPO)
+    spaces.aoti_load(module=pipe.transformer_2, repo_id=AOTI_TRANSFORMER_REPO)
+
     return pipe
 
 
@@ -86,6 +113,18 @@ def resize_image(image: Image.Image) -> Image.Image:
     return image_to_resize.resize((final_w, final_h), Image.Resampling.LANCZOS)
 
 
+def resize_and_crop_to_match(target_image: Image.Image, reference_image: Image.Image) -> Image.Image:
+    """Resize/crop `target_image` to exactly match `reference_image`'s size (the pipeline requires
+    the first and last frame to share dimensions)."""
+    ref_width, ref_height = reference_image.size
+    target_width, target_height = target_image.size
+    scale = max(ref_width / target_width, ref_height / target_height)
+    new_width, new_height = int(target_width * scale), int(target_height * scale)
+    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+    left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
+    return resized.crop((left, top, left + ref_width, top + ref_height))
+
+
 def get_num_frames(duration_seconds: float) -> int:
     raw = int(round(duration_seconds * FIXED_FPS))
     raw = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, raw))
@@ -99,6 +138,10 @@ def estimate_generation_seconds(
 
     Exact limits are still an open item (SRS Section 4, #3) pending empirical testing on the
     dev Space; this heuristic is a starting point, not a final calibration.
+
+    Kept conservative even now that the transformers are AOT-compiled (#36): the reference
+    project's numbers below assume a validated AOT setup, and ours hasn't been measured on real
+    GPU yet (see load_pipeline's AOT caveat) — recalibrate once that's confirmed working.
     """
     base_frames_height_width = 81 * 832 * 624
     base_step_duration = 5.0
@@ -108,9 +151,8 @@ def estimate_generation_seconds(
     if guidance_scale > 1:
         gen_time *= 2.4
 
-    # The reference project's "15" was calibrated with AOT-compiled transformers (FR-6, deferred
-    # here per C-2/issue #3). Without AOT, per-call ZeroGPU init (loading the quantized-but-not-
-    # compiled pipe onto the GPU) plus eager-mode execution both take longer, so pad generously
-    # until this can be measured empirically on the dev Space.
+    # Padded well above the reference project's flat overhead (its pipe loads a validated AOT
+    # setup already proven on its own GPU quota) until ours is confirmed working and measured
+    # empirically on the dev Space (#36).
     init_overhead = 60
     return init_overhead + gen_time

postprocess/interpolation.py

@@ -6,6 +6,10 @@ Model code/weights come from `thornmaze/RIFE`, a packaged build of hzwer/Practic
 `from model.warplayer import warp` (and, at module load time, from `model.loss` and
 `model.pytorch_msssim`) expecting a sibling `model` package — see the vendored copies in
 model/warplayer.py, model/loss.py, model/pytorch_msssim/.
+
+The zip download/unzip itself needs no GPU, so `ensure_weights_downloaded()` is called from
+app.py at process startup (like `load_pipeline()`), before any `@spaces.GPU` call — otherwise
+a cold container pays for that fetch out of the metered GPU allocation on its first request.
 """
 
 from __future__ import annotations
@@ -13,7 +17,7 @@ from __future__ import annotations
 import subprocess
 from functools import lru_cache
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, cast
 
 import numpy as np
 import torch
@@ -37,7 +41,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ProgressCallback = Callable[[int, int], None]
 
 
-def _ensure_weights_downloaded() -> None:
+def ensure_weights_downloaded() -> None:
+    """Download and unpack the RIFE weights if not already present. CPU/network only — no CUDA
+    needed — so this is safe (and intended) to call outside any @spaces.GPU allocation."""
     if RIFE_DIR.exists():
         return
     subprocess.run(["wget", "-q", RIFE_ZIP_URL, "-O", str(RIFE_ZIP_PATH)], check=True)
@@ -46,7 +52,7 @@ def _ensure_weights_downloaded() -> None:
 
 @lru_cache(maxsize=1)
 def _load_model() -> Any:
-    _ensure_weights_downloaded()
+    ensure_weights_downloaded()
     from train_log.RIFE_HDv3 import Model  # only importable after the zip is unpacked
 
     model = Model()
@@ -63,15 +69,21 @@ def interpolate_frames(
     multiplier: int,
     scale: float = 1.0,
     progress_callback: ProgressCallback | None = None,
-) -> list[FrameArray]:
+    as_tensor: bool = False,
+) -> list[FrameArray] | 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
@@ -91,11 +103,19 @@ def interpolate_frames(
         t = t.permute(2, 0, 1).unsqueeze(0)
         return F.pad(t, padding).half()
 
-    def from_tensor(tensor: torch.Tensor) -> FrameArray:
+    def from_tensor_np(tensor: torch.Tensor) -> FrameArray:
         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: Callable[[torch.Tensor], FrameArray | torch.Tensor]
+    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)]
@@ -108,7 +128,7 @@ def interpolate_frames(
             return [*first_half, middle, *second_half]
         return [*first_half, *second_half]
 
-    output_frames: list[FrameArray] = []
+    output_frames: list[FrameArray | torch.Tensor] = []
     next_tensor = to_tensor(frames_np[0])
     total_steps = total_frames - 1
 
@@ -126,4 +146,6 @@ def interpolate_frames(
     del current_tensor, next_tensor
     torch.cuda.empty_cache()
 
-    return output_frames
+    if as_tensor:
+        return cast("list[torch.Tensor]", output_frames)
+    return cast("list[FrameArray]", output_frames)

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,
@@ -16,6 +17,10 @@ flattens real skin/fabric/hair texture into a smooth, CG-like look on frames tha
 Runs inside the caller's `@spaces.GPU` allocation (see app.py's `run_inference`), same as RIFE interpolation
 (postprocess/interpolation.py) — moved off the Space's shared CPU per issue #11 to avoid CPU contention across
 concurrent visitors. Falls back to CPU/fp32 automatically when no CUDA device is present (e.g. local runs).
+
+The weight download itself needs no GPU, so `ensure_weights_downloaded()` is called from app.py at process
+startup (like `load_pipeline()`), before any `@spaces.GPU` call — otherwise a cold container pays for that
+fetch out of the metered GPU allocation on its first upscale request.
 """
 
 from __future__ import annotations
@@ -23,7 +28,7 @@ from __future__ import annotations
 import math
 from functools import lru_cache
 from pathlib import Path
-from typing import Callable
+from typing import TYPE_CHECKING, Callable, cast
 
 import numpy as np
 import torch
@@ -34,34 +39,71 @@ from torch.hub import download_url_to_file, get_dir
 
 from .srvgg_arch import SRVGGNetCompact
 
+if TYPE_CHECKING:
+    # numpy.typing needs numpy>=1.20; requirements.txt allows down to 1.16, so this must stay out
+    # of the runtime import path.
+    import numpy.typing as npt
+
+    FrameArray = npt.NDArray[np.float32]
+else:
+    FrameArray = np.ndarray
+
 WEIGHTS_URL = "https://huggingface.co/Phips/4xLSDIRCompact/resolve/main/4xLSDIRCompact.safetensors"
 SCALE = 4
 NUM_CONV = 16
-TILE_SIZE = 256
+# Larger tiles mean fewer, bigger model() calls (less pad/crop overhead per pixel, better GPU
+# occupancy) at the cost of more peak memory per call. 512 is still conservative for a 16GB+
+# Spaces GPU (T4/L4/A10G) with such a compact network — bump further if profiling on the dev
+# Space shows headroom.
+TILE_SIZE = 512
 TILE_PAD = 10
+# Frames are grouped into batches (of equal size) before tiling, and same-shaped tile patches
+# — across both the frame batch and tile positions — are run through the model together instead
+# of one at a time, since SRVGGNetCompact is small enough that a single 256x256-ish tile leaves
+# most of the GPU idle. Conservative defaults for a 16GB+ Spaces GPU (T4/L4/A10G); bump if
+# profiling on the dev Space shows headroom.
+FRAME_BATCH_SIZE = 4
+MAX_TILE_BATCH = 16
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.float16 if device.type == "cuda" else torch.float32
 
 ProgressCallback = Callable[[int, int], None]
 
 
-def _download_weights() -> Path:
+def _weights_path() -> Path:
     cache_dir = Path(get_dir()) / "checkpoints"
     cache_dir.mkdir(parents=True, exist_ok=True)
-    weights_path = cache_dir / Path(WEIGHTS_URL).name
+    return cache_dir / Path(WEIGHTS_URL).name
+
+
+def ensure_weights_downloaded() -> None:
+    """Download the 4xLSDIRCompact weights if not already cached. CPU/network only — no CUDA
+    needed — so this is safe (and intended) to call outside any @spaces.GPU allocation."""
+    weights_path = _weights_path()
     if not weights_path.exists():
         download_url_to_file(WEIGHTS_URL, str(weights_path))
-    return weights_path
 
 
 @lru_cache(maxsize=1)
-def _load_model() -> SRVGGNetCompact:
-    model = SRVGGNetCompact(  # type: ignore[no-untyped-call]  # vendored, untyped (see srvgg_arch.py)
+def _load_model() -> torch.nn.Module:
+    model: torch.nn.Module = SRVGGNetCompact(  # type: ignore[no-untyped-call]  # vendored, untyped (see srvgg_arch.py)
         num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=NUM_CONV, upscale=SCALE, act_type="prelu"
     )
-    state_dict = load_file(_download_weights())
+    ensure_weights_downloaded()
+    state_dict = load_file(_weights_path())
     model.load_state_dict(state_dict, strict=True)
     model.to(device=device, dtype=dtype).eval()
+    if device.type == "cuda":
+        try:
+            # dynamic=True avoids a fresh recompile for every distinct tile shape (interior
+            # tiles vs. the last row/column, clipped against the image edge) — there are only a
+            # handful of those per video, but recompiling per-shape would otherwise eat into the
+            # first frame's latency. Guarded: torch.compile can fail to initialize (e.g. missing
+            # or incompatible triton on the target GPU image), in which case the eager model is
+            # used as-is.
+            model = cast(torch.nn.Module, torch.compile(model, dynamic=True))
+        except Exception as e:
+            print(f"[upscale] torch.compile failed to initialize, using eager model: {e!r}", flush=True)
     return model
 
 
@@ -70,51 +112,101 @@ def _pre_pad(tensor: torch.Tensor, pad: int) -> torch.Tensor:
 
 
 @torch.no_grad()
-def _tile_process(model: SRVGGNetCompact, img: torch.Tensor) -> torch.Tensor:
+def _tile_process(model: torch.nn.Module, img: torch.Tensor) -> torch.Tensor:
     batch, channel, height, width = img.shape
     output = img.new_zeros((batch, channel, height * SCALE, width * SCALE))
     tiles_x = math.ceil(width / TILE_SIZE)
     tiles_y = math.ceil(height / TILE_SIZE)
 
-    for y in range(tiles_y):
-        for x in range(tiles_x):
-            in_x0, in_x1 = x * TILE_SIZE, min((x + 1) * TILE_SIZE, width)
-            in_y0, in_y1 = y * TILE_SIZE, min((y + 1) * TILE_SIZE, height)
-            pad_x0, pad_x1 = max(in_x0 - TILE_PAD, 0), min(in_x1 + TILE_PAD, width)
-            pad_y0, pad_y1 = max(in_y0 - TILE_PAD, 0), min(in_y1 + TILE_PAD, height)
-
-            tile_out = model(img[:, :, pad_y0:pad_y1, pad_x0:pad_x1])
-
-            trim_x0, trim_y0 = (in_x0 - pad_x0) * SCALE, (in_y0 - pad_y0) * SCALE
-            trim_x1 = trim_x0 + (in_x1 - in_x0) * SCALE
-            trim_y1 = trim_y0 + (in_y1 - in_y0) * SCALE
-            output[:, :, in_y0 * SCALE:in_y1 * SCALE, in_x0 * SCALE:in_x1 * SCALE] = (
-                tile_out[:, :, trim_y0:trim_y1, trim_x0:trim_x1]
+    # Every tile job needs its own input patch (pad_*) and output placement (in_*/trim_*), but
+    # patches share the same H/W across the batch dim and across interior tile positions — only
+    # the last row/column of tiles is smaller, clipped against the image edge. Group jobs by
+    # patch shape so each model() call is a single batched forward pass over same-shaped patches
+    # (chunked to MAX_TILE_BATCH to bound memory) instead of one tile at a time.
+    jobs_by_shape: dict[tuple[int, int], list[tuple[int, int, int, int, int, int, int, int, int]]] = {}
+    for b in range(batch):
+        for y in range(tiles_y):
+            for x in range(tiles_x):
+                in_x0, in_x1 = x * TILE_SIZE, min((x + 1) * TILE_SIZE, width)
+                in_y0, in_y1 = y * TILE_SIZE, min((y + 1) * TILE_SIZE, height)
+                pad_x0, pad_x1 = max(in_x0 - TILE_PAD, 0), min(in_x1 + TILE_PAD, width)
+                pad_y0, pad_y1 = max(in_y0 - TILE_PAD, 0), min(in_y1 + TILE_PAD, height)
+                shape = (pad_y1 - pad_y0, pad_x1 - pad_x0)
+                jobs_by_shape.setdefault(shape, []).append(
+                    (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1)
+                )
+
+    for jobs in jobs_by_shape.values():
+        for chunk_start in range(0, len(jobs), MAX_TILE_BATCH):
+            chunk = jobs[chunk_start:chunk_start + MAX_TILE_BATCH]
+            patch = torch.cat(
+                [img[b:b + 1, :, pad_y0:pad_y1, pad_x0:pad_x1] for b, _, _, _, _, pad_x0, pad_y0, pad_x1, pad_y1 in chunk],
+                dim=0,
             )
+            tile_out = model(patch)
+
+            for i, (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1) in enumerate(chunk):
+                trim_x0, trim_y0 = (in_x0 - pad_x0) * SCALE, (in_y0 - pad_y0) * SCALE
+                trim_x1 = trim_x0 + (in_x1 - in_x0) * SCALE
+                trim_y1 = trim_y0 + (in_y1 - in_y0) * SCALE
+                output[b:b + 1, :, in_y0 * SCALE:in_y1 * SCALE, in_x0 * SCALE:in_x1 * SCALE] = (
+                    tile_out[i:i + 1, :, trim_y0:trim_y1, trim_x0:trim_x1]
+                )
     return output
 
 
+def _frame_hw(frame: FrameArray | 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).
+    h, w = frame.shape[-2:] if isinstance(frame, torch.Tensor) else frame.shape[:2]
+    return int(h), int(w)
+
+
+def _frames_to_tensor(frames: list[FrameArray | 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(cast("list[torch.Tensor]", frames), dim=0).to(device=device, dtype=dtype)
+    array = np.stack(cast("list[FrameArray]", 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[FrameArray] | list[torch.Tensor], progress_callback: ProgressCallback | None = None
 ) -> list[Image.Image]:
     """Upscale every frame 4x with `4xLSDIRCompact`, tiled to bound peak memory.
 
-    progress_callback(done, total), if given, fires after each frame finishes.
+    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
+    batch=1 on every call. progress_callback(done, total), if given, fires after each frame
+    finishes.
     """
     if not frames:
         return []
     model = _load_model()
-    out_frames = []
-    for frame in frames:
-        w, h = frame.size
-        array = np.asarray(frame.convert("RGB"), dtype=np.float32) / 255.0
-        tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=dtype)
+    out_frames: list[Image.Image] = []
+
+    i = 0
+    while i < len(frames):
+        h, w = _frame_hw(frames[i])
+        batch: list[FrameArray | torch.Tensor] = [frames[i]]
+        i += 1
+        while i < len(frames) and len(batch) < FRAME_BATCH_SIZE and _frame_hw(frames[i]) == (h, w):
+            batch.append(frames[i])
+            i += 1
+
+        tensor = _frames_to_tensor(batch)
         padded = _pre_pad(tensor, TILE_PAD)
         upscaled = _tile_process(model, padded)
         # crop the pre-pad border (scaled) back off
         upscaled = upscaled[:, :, : h * SCALE, : w * SCALE]
-        result = upscaled.squeeze(0).clamp(0, 1).permute(1, 2, 0).float().cpu().numpy()
-        out_frames.append(Image.fromarray((result * 255.0).round().astype(np.uint8)))
-        if progress_callback is not None:
-            progress_callback(len(out_frames), len(frames))
+        results = upscaled.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy()
+        for result in results:
+            out_frames.append(Image.fromarray((result * 255.0).round().astype(np.uint8)))
+            if progress_callback is not None:
+                progress_callback(len(out_frames), len(frames))
     return out_frames