swiftvr/runner.py

"""Four-stage pipelined runner: reader -> H2D -> GPU -> writer.

Overlaps host reading, host->device copy, GPU restoration and disk writing on
separate threads/CUDA streams to maximise sustained throughput.
"""

import time
import queue
import threading
import traceback
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import torch

from .streaming.chunk import ChunkType
from .io import (
    iter_video_clips_fixed_scheme,
    preprocess_clip_uint8,
    crop_spatial_padding_ntchw,
    ntchw_to_uint8_frames,
    append_chunk_to_png_dir,
    open_stream_video_writer,
)


def cuda_synchronize():
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def enable_max_fps_runtime(allow_tf32=True):
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32)
        torch.backends.cudnn.allow_tf32 = bool(allow_tf32)
    try:
        torch.set_float32_matmul_precision("high")
    except Exception:
        pass


def _event_elapsed_seconds(start_event, end_event):
    try:
        end_event.synchronize()
        return start_event.elapsed_time(end_event) / 1000.0
    except Exception:
        return 0.0


@dataclass
class _Item:
    clip_idx: int
    spec: object = None
    cpu_rgb: Optional[torch.Tensor] = None
    gpu_rgb: Optional[torch.Tensor] = None
    rgb_out_gpu: Optional[torch.Tensor] = None
    rgb_out_cpu: Optional[torch.Tensor] = None
    h2d_event: Optional[torch.cuda.Event] = None
    d2h_event: Optional[torch.cuda.Event] = None
    timings: Dict[str, float] = field(default_factory=dict)
    stop: bool = False


def _stop_item():
    return _Item(clip_idx=-1, stop=True)


def run_pipeline(
    *,
    video_path,
    final_output_path,
    png_output_dir,
    tae_stream,
    dit_stream,
    prompt_emb,
    device,
    dtype,
    total_frames: int,
    clip_len: int,
    lq_h: int,
    lq_w: int,
    out_h: int,
    out_w: int,
    pad_h: int,
    pad_w: int,
    upscale_mode: str,
    source_fps: float,
    png_save: bool,
    quality: int,
    save_format: str = "",
    ffmpeg_preset: str = "",
    queue_size: int = 3,
    png_frame_names: Optional[List[str]] = None,
    verbose: bool = True,
):
    q_read = queue.Queue(maxsize=max(1, queue_size))
    q_gpu = queue.Queue(maxsize=max(1, queue_size))
    q_write = queue.Queue(maxsize=max(1, queue_size))

    stage_errors = []
    stop_event = threading.Event()
    frames_state = {"next_idx": 0, "saved": 0}
    png_written_once = set()

    use_cuda = torch.cuda.is_available() and device.type == "cuda"
    h2d_stream = torch.cuda.Stream(device=device) if use_cuda else None
    d2h_stream = torch.cuda.Stream(device=device) if use_cuda else None

    def record_error(stage_name):
        stage_errors.append((stage_name, traceback.format_exc()))
        stop_event.set()
        for q in (q_read, q_gpu, q_write):
            try:
                q.put(_stop_item())
            except Exception:
                pass

    def reader_worker():
        try:
            clips = iter_video_clips_fixed_scheme(
                video_path, clip_len=clip_len, total_frames=total_frames, crop_h=lq_h, crop_w=lq_w)
            for spec, cpu_rgb in clips:
                if stop_event.is_set():
                    break
                try:
                    cpu_rgb = cpu_rgb.pin_memory()
                except Exception:
                    pass
                q_read.put(_Item(clip_idx=spec.clip_idx, spec=spec, cpu_rgb=cpu_rgb))
            q_read.put(_stop_item())
        except Exception:
            record_error("reader")

    def h2d_worker():
        try:
            while True:
                item = q_read.get()
                if item.stop:
                    q_gpu.put(_stop_item())
                    break
                if stop_event.is_set():
                    continue
                if use_cuda:
                    se = torch.cuda.Event(enable_timing=True)
                    ee = torch.cuda.Event(enable_timing=True)
                    with torch.cuda.stream(h2d_stream):
                        se.record(h2d_stream)
                        item.gpu_rgb = item.cpu_rgb.to(device=device, non_blocking=True)
                        ee.record(h2d_stream)
                    item.h2d_event = ee
                    item.timings["_h2d"] = (se, ee)
                else:
                    item.gpu_rgb = item.cpu_rgb.to(device=device)
                item.cpu_rgb = None
                q_gpu.put(item)
        except Exception:
            record_error("h2d")

    def _start_d2h(item):
        if item.rgb_out_gpu is None or item.rgb_out_gpu.shape[1] == 0:
            item.rgb_out_cpu = None
            return item
        try:
            cpu_buf = torch.empty(item.rgb_out_gpu.shape, dtype=item.rgb_out_gpu.dtype,
                                  device="cpu", pin_memory=True)
        except Exception:
            cpu_buf = torch.empty(item.rgb_out_gpu.shape, dtype=item.rgb_out_gpu.dtype, device="cpu")
        if use_cuda:
            se = torch.cuda.Event(enable_timing=True)
            ee = torch.cuda.Event(enable_timing=True)
            with torch.cuda.stream(d2h_stream):
                se.record(d2h_stream)
                cpu_buf.copy_(item.rgb_out_gpu, non_blocking=True)
                try:
                    item.rgb_out_gpu.record_stream(d2h_stream)
                except Exception:
                    pass
                ee.record(d2h_stream)
            item.d2h_event = ee
        else:
            cpu_buf.copy_(item.rgb_out_gpu)
        item.rgb_out_cpu = cpu_buf
        return item

    def gpu_worker():
        try:
            tae_stream.reset()
            dit_stream.reset()
            n_lat = clip_len // 4
            prev_dit_out_cpu = None

            while True:
                item = q_gpu.get()
                if item.stop:
                    break
                if stop_event.is_set():
                    continue

                if item.h2d_event is not None:
                    torch.cuda.current_stream(device=device).wait_event(item.h2d_event)
                spec = item.spec

                if use_cuda:
                    t_start = torch.cuda.Event(enable_timing=True)
                    t_end = torch.cuda.Event(enable_timing=True)
                    t_start.record(torch.cuda.current_stream(device=device))
                else:
                    t0 = time.perf_counter()

                clip_rgb = preprocess_clip_uint8(
                    item.gpu_rgb, out_h=out_h, out_w=out_w, mode=upscale_mode,
                    pad_h=pad_h, pad_w=pad_w, dtype=dtype)
                z = tae_stream.encode_chunk_fixed(clip_rgb, spec)

                if spec.ctype == ChunkType.LAST:
                    z_ntchw = dit_stream.denoise_last_chunk(
                        z, spec, prompt_emb, prev_dit_out_cpu, n_lat, device, dtype)
                else:
                    z_bcfhw = z.permute(0, 2, 1, 3, 4).contiguous()
                    z_den = dit_stream.denoise(z_bcfhw, prompt_emb)
                    z_ntchw = z_den.permute(0, 2, 1, 3, 4).contiguous()
                    prev_dit_out_cpu = z_bcfhw[:, :, -n_lat:].detach().cpu().clone()

                rgb_out = tae_stream.decode_chunk_fixed(z_ntchw, spec)
                if rgb_out is not None and rgb_out.shape[1] > 0:
                    item.rgb_out_gpu = crop_spatial_padding_ntchw(rgb_out, pad_h, pad_w).detach()
                else:
                    item.rgb_out_gpu = None

                if use_cuda:
                    t_end.record(torch.cuda.current_stream(device=device))
                    item.timings["gpu"] = _event_elapsed_seconds(t_start, t_end)
                else:
                    item.timings["gpu"] = time.perf_counter() - t0

                if verbose:
                    out_n = 0 if item.rgb_out_gpu is None else item.rgb_out_gpu.shape[1]
                    fps = out_n / item.timings["gpu"] if item.timings["gpu"] > 0 else 0.0
                    print(f"  [gpu] {spec.ctype.value:6s} clip {item.clip_idx}: "
                          f"out={out_n}f time={item.timings['gpu']:.3f}s gpu_fps={fps:.2f}")

                item.gpu_rgb = None
                q_write.put(_start_d2h(item))
                del clip_rgb, z, z_ntchw, rgb_out

            q_write.put(_stop_item())
        except Exception:
            record_error("gpu")

    def writer_worker():
        writer = None
        try:
            while True:
                item = q_write.get()
                if item.stop:
                    break
                if stop_event.is_set():
                    continue
                if item.d2h_event is not None:
                    item.d2h_event.synchronize()
                item.rgb_out_gpu = None

                n_written = n_consumed = 0
                if item.rgb_out_cpu is not None and item.rgb_out_cpu.shape[1] > 0:
                    if png_save:
                        n_written, n_consumed = append_chunk_to_png_dir(
                            item.rgb_out_cpu, png_output_dir, start_idx=frames_state["next_idx"],
                            frame_names=png_frame_names, written_once=png_written_once)
                    else:
                        if writer is None:
                            writer = open_stream_video_writer(
                                final_output_path, fps=source_fps, video_format=save_format,
                                preset=ffmpeg_preset, quality=quality)
                        frames = ntchw_to_uint8_frames(item.rgb_out_cpu)
                        if frames is not None:
                            for frame in frames:
                                writer.append_data(frame)
                            n_written = n_consumed = int(frames.shape[0])

                frames_state["next_idx"] += n_consumed
                frames_state["saved"] += n_written
        except Exception:
            record_error("writer")
        finally:
            if writer is not None:
                writer.close()

    threads = [
        threading.Thread(target=reader_worker, name="reader", daemon=True),
        threading.Thread(target=h2d_worker, name="h2d", daemon=True),
        threading.Thread(target=gpu_worker, name="gpu", daemon=True),
        threading.Thread(target=writer_worker, name="writer", daemon=True),
    ]
    t0 = time.perf_counter()
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    wall_time = time.perf_counter() - t0

    if stage_errors:
        name, err = stage_errors[0]
        raise RuntimeError(f"Pipeline stage '{name}' failed:\n{err}")

    written = frames_state["saved"] if png_save else frames_state["next_idx"]
    return written, wall_time