from __future__ import annotations import argparse import importlib from .config import ColorCorrectionConfig, DecoderConfig, InferenceConfig, StreamConfig from .console import print_banner, print_configuration, print_stage from .loading import ComponentFactories, ModelPaths from .pipeline import build_inference_pipeline from .request import InferenceRequest def _load_factory(value: str | None): if value is None: return None module_name, symbol = value.split(":", 1) return getattr(importlib.import_module(module_name), symbol) def parse_args(): parser = argparse.ArgumentParser( description="FlashVSR+ stream video super-resolution inference", allow_abbrev=False, ) parser.add_argument("--input", required=True) parser.add_argument("--output", required=True) parser.add_argument("--dit", default=None) parser.add_argument("--student", default=None) parser.add_argument("--lq_projection", default=None) parser.add_argument("--context", default=None) parser.add_argument("--vae", default=None) parser.add_argument("--tiny_decoder", default=None) parser.add_argument("--decoder", choices=["vae", "tiny"], default="tiny") parser.add_argument( "--decode-every", "--decode_every", default="auto", help="auto loads the full input; a positive multiple of 4 enables incremental reading and decoding", ) parser.add_argument( "--tiny_decoder_parallel_mode", choices=["auto", "serial", "parallel"], default="auto", help="TinyDecoder mode: auto, serial, or full temporal parallel", ) parser.add_argument( "--tiny_decoder_channels_last", action=argparse.BooleanOptionalAction, default=True, help="use channels-last BF16 TinyDecoder convolutions (default: true)", ) parser.add_argument("--is_sparse", action="store_true") parser.add_argument("--is_hybrid_sliding", action="store_true") parser.add_argument( "--attn_kernel", choices=[ "auto", "sdpa", "block_sparse", "triton_block_sparse", "flash_attn2", "flash_attn3", ], default="auto", ) parser.add_argument("--block_size", type=int, default=1) parser.add_argument("--sparse_topk", type=int, default=None) parser.add_argument("--sparse_topk_ratio", type=float, default=2.5) parser.add_argument("--sparse_local_blocks", type=int, default=None) parser.add_argument("--sparse_local_range", type=int, default=None) parser.add_argument( "--sparse_window_size", nargs=3, type=int, default=(2, 8, 8), ) parser.add_argument( "--sparse_kv_ratio", type=float, default=3.0, help="stream KV temporal-group count; every group keeps the full spatial grid (default: 3.0)", ) parser.add_argument( "--stream_cache_offload_layers", nargs="*", type=int, default=(), help="optional DiT layer indices whose stream KV cache is kept on CPU between chunks", ) parser.add_argument( "--stream_cache_offload_mode", choices=["auto", "never", "always"], default="auto", help="stream KV residency policy: auto by resolution/VRAM, or force never/always", ) parser.add_argument( "--stream_cache_pin_memory", action=argparse.BooleanOptionalAction, default=True, help="use pinned host buffers for stream KV offload (default: true)", ) parser.add_argument("--replacement_layers", nargs="*", type=int, default=()) parser.add_argument("--sliding_local_range", type=int, default=3) parser.add_argument("--sliding_temporal_block_size", type=int, default=2) parser.add_argument( "--sliding_temporal_mode", choices=["local", "causal", "short_causal"], default="local", ) parser.add_argument( "--sliding_spatial_mode", choices=["local", "base"], default="local", ) parser.add_argument("--sliding_cache_temporal_windows", type=int, default=1) parser.add_argument("--lq_projection_layers", type=int, default=1) parser.add_argument("--chunk_latents", type=int, default=4) # Accepted only so older launch scripts keep working. Stream scheduling is # now defined in latent space by --chunk_latents. parser.add_argument("--chunk_frames", type=int, default=None, help=argparse.SUPPRESS) parser.add_argument( "--temporal_compression", type=int, choices=[1, 2, 4], default=4, help="effective temporal compression exposed to the input/output timeline (default: 4)", ) parser.add_argument( "--temporal_merge", choices=["gaussian", "first"], default="gaussian", help="how repeated decoded frames are collapsed for temporal compression 1 or 2", ) parser.add_argument( "--temporal_gaussian_sigma", type=float, default=1.0, help="Gaussian sigma used by --temporal_merge gaussian", ) parser.add_argument("--height", type=int, default=None) parser.add_argument("--width", type=int, default=None) parser.add_argument("--max_frames", type=int, default=None) parser.add_argument("--fps", type=float, default=None) parser.add_argument("--steps", type=int, default=1) parser.add_argument("--sigma_shift", type=float, default=5.0) parser.add_argument("--seed", type=int, default=2026) parser.add_argument("--device", default="cuda:0") parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16") parser.add_argument( "--color-correction", "--color_correction", action=argparse.BooleanOptionalAction, default=False, help="match output low-frequency color to the LQ input (default: disabled)", ) parser.add_argument( "--color-correction-levels", "--color_correction_levels", type=int, default=5 ) parser.add_argument( "--color-correction-chunk-size", "--color_correction_chunk_size", type=int, default=16, ) parser.add_argument("--denoiser_factory", default=None, help="Denoiser factory in module:function format") parser.add_argument("--vae_factory", default=None, help="VAE factory in module:function format") parser.add_argument("--tiny_factory", default=None, help="Tiny decoder factory in module:function format") return parser.parse_args() def main() -> None: args = parse_args() config = InferenceConfig( is_sparse=args.is_sparse, is_hybrid_sliding=args.is_hybrid_sliding, attn_kernel=args.attn_kernel, block_size=args.block_size, sparse_topk=args.sparse_topk, sparse_topk_ratio=args.sparse_topk_ratio, sparse_local_blocks=args.sparse_local_blocks, sparse_local_range=args.sparse_local_range, sparse_window_size=tuple(args.sparse_window_size), sparse_kv_ratio=args.sparse_kv_ratio, stream_cache_offload_layers=tuple(args.stream_cache_offload_layers), stream_cache_offload_mode=args.stream_cache_offload_mode, stream_cache_pin_memory=args.stream_cache_pin_memory, replacement_layers=tuple(args.replacement_layers), sliding_local_range=args.sliding_local_range, sliding_temporal_block_size=args.sliding_temporal_block_size, sliding_temporal_mode=args.sliding_temporal_mode, sliding_spatial_mode=args.sliding_spatial_mode, sliding_cache_temporal_windows=args.sliding_cache_temporal_windows, lq_projection_layers=args.lq_projection_layers, decoder=DecoderConfig( kind=args.decoder, decode_every=args.decode_every, tiny_decoder_parallel_mode=args.tiny_decoder_parallel_mode, tiny_decoder_channels_last=args.tiny_decoder_channels_last, ), color_correction=ColorCorrectionConfig( enabled=args.color_correction, levels=args.color_correction_levels, chunk_size=args.color_correction_chunk_size, ), streaming=StreamConfig(args.chunk_latents), temporal_compression=args.temporal_compression, temporal_merge=args.temporal_merge, temporal_gaussian_sigma=args.temporal_gaussian_sigma, num_inference_steps=args.steps, sigma_shift=args.sigma_shift, seed=args.seed, device=args.device, dtype=args.dtype, ) config.validate() paths = ModelPaths( dit=args.dit, student=args.student, lq_projection=args.lq_projection, context=args.context, vae=args.vae, tiny_decoder=args.tiny_decoder, ) factories = ComponentFactories( denoiser=_load_factory(args.denoiser_factory), vae=_load_factory(args.vae_factory), tiny_decoder=_load_factory(args.tiny_factory), ) print_banner() print_configuration(config, args) print_stage("Loading model weights...") pipeline = build_inference_pipeline(config, paths, factories=factories) print_stage("Model loading complete", done=True) print_stage(f"Starting inference: {args.input}") result = pipeline.infer( InferenceRequest( input_path=args.input, output_path=args.output, height=args.height, width=args.width, max_frames=args.max_frames, fps=args.fps, ) ) print_stage( f"Inference complete · {result.frames} frames · Saved to {result.output_path}", done=True, ) if __name__ == "__main__": main()