diff --git a/docs/diffusion/api/cli.md b/docs/diffusion/api/cli.md index 85928b03e..b35cbf60a 100644 --- a/docs/diffusion/api/cli.md +++ b/docs/diffusion/api/cli.md @@ -41,6 +41,34 @@ The SGLang-diffusion CLI provides a quick way to access the inference pipeline f - `--fps {FPS}`: Frames per second for the saved output, if this is a video-generation task +**Frame Interpolation** (video only) + +Frame interpolation is a post-processing step that synthesizes new frames +between each pair of consecutive generated frames, producing smoother +motion without re-running the diffusion model. The `--frame-interpolation-exp` +flag controls how many rounds of interpolation to apply: each round inserts one +new frame into every gap between adjacent frames, so the output frame count +follows the formula **(N − 1) × 2^exp + 1** (e.g. 5 original frames with +`exp=1` → 4 gaps × 1 new frame + 5 originals = **9** frames; with `exp=2` → +**17** frames). + +- `--enable-frame-interpolation`: Enable frame interpolation. Model weights are downloaded automatically on first use. +- `--frame-interpolation-exp {EXP}`: Interpolation exponent — `1` = 2× temporal resolution, `2` = 4×, etc. (default: `1`) +- `--frame-interpolation-scale {SCALE}`: RIFE inference scale; use `0.5` for high-resolution inputs to save memory (default: `1.0`) +- `--frame-interpolation-model-path {PATH}`: Local directory or HuggingFace repo ID containing RIFE `flownet.pkl` weights (default: `elfgum/RIFE-4.22.lite`, downloaded automatically) + +Example — generate a 5-frame video and interpolate to 9 frames ((5 − 1) × 2¹ + 1 = 9): + +```bash +sglang generate \ + --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --prompt "A dog running through a park" \ + --num-frames 5 \ + --enable-frame-interpolation \ + --frame-interpolation-exp 1 \ + --save-output +``` + **Output Options** - `--output-path {PATH}`: Directory to save the generated video diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 746d38bdd..1b411101d 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -100,6 +100,14 @@ class SamplingParams: output_quality: str | None = "default" output_compression: int | None = None + # Frame interpolation + enable_frame_interpolation: bool = False + frame_interpolation_exp: int = 1 # 1=2x, 2=4x + frame_interpolation_scale: float = 1.0 # RIFE inference scale (0.5 for high-res) + frame_interpolation_model_path: str | None = ( + None # local dir or HF repo ID with flownet.pkl (default: elfgum/RIFE-4.22.lite) + ) + # Batch info num_outputs_per_prompt: int = 1 seed: int = 42 @@ -817,6 +825,30 @@ class SamplingParams: default=SamplingParams.enable_sequence_shard, help="Enable sequence dimension shard with sequence parallelism.", ) + parser.add_argument( + "--enable-frame-interpolation", + action="store_true", + help="Enable post-generation frame interpolation using RIFE 4.22.lite.", + ) + parser.add_argument( + "--frame-interpolation-exp", + type=int, + default=SamplingParams.frame_interpolation_exp, + help="Frame interpolation exponent: 1=2x, 2=4x (default: 1).", + ) + parser.add_argument( + "--frame-interpolation-scale", + type=float, + default=SamplingParams.frame_interpolation_scale, + help="RIFE inference scale factor (default: 1.0; use 0.5 for high-res).", + ) + parser.add_argument( + "--frame-interpolation-model-path", + type=str, + default=SamplingParams.frame_interpolation_model_path, + help="Local directory or HuggingFace repo ID containing RIFE flownet.pkl weights " + "(default: elfgum/RIFE-4.22.lite, downloaded automatically).", + ) return parser @classmethod diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 7e03c3bf4..ced687e90 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -243,6 +243,10 @@ class DiffGenerator: audios_out=audios_out, frames_out=frames_out, output_compression=req.output_compression, + enable_frame_interpolation=req.enable_frame_interpolation, + frame_interpolation_exp=req.frame_interpolation_exp, + frame_interpolation_scale=req.frame_interpolation_scale, + frame_interpolation_model_path=req.frame_interpolation_model_path, ) for idx in range(len(samples_out)): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 30a60b35a..b2ab9623b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -141,6 +141,10 @@ async def forward_to_scheduler( lambda _idx: output_file_path, audio=response.audio, audio_sample_rate=response.audio_sample_rate, + enable_frame_interpolation=sp.enable_frame_interpolation, + frame_interpolation_exp=sp.frame_interpolation_exp, + frame_interpolation_scale=sp.frame_interpolation_scale, + frame_interpolation_model_path=sp.frame_interpolation_model_path, ) if hasattr(response, "model_dump"): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index 770713190..9d0e89904 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -90,6 +90,11 @@ class VideoGenerationsRequest(BaseModel): ) negative_prompt: Optional[str] = None enable_teacache: Optional[bool] = False + # Frame interpolation + enable_frame_interpolation: Optional[bool] = False + frame_interpolation_exp: Optional[int] = 1 # 1=2×, 2=4× + frame_interpolation_scale: Optional[float] = 1.0 + frame_interpolation_model_path: Optional[str] = None output_quality: Optional[str] = "default" output_compression: Optional[int] = None output_path: Optional[str] = None diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 0f8ea4397..5dc6195f9 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -259,6 +259,10 @@ async def process_generation_batch( audio=result.audio, audio_sample_rate=result.audio_sample_rate, output_compression=batch.output_compression, + enable_frame_interpolation=batch.enable_frame_interpolation, + frame_interpolation_exp=batch.frame_interpolation_exp, + frame_interpolation_scale=batch.frame_interpolation_scale, + frame_interpolation_model_path=batch.frame_interpolation_model_path, ) total_time = time.perf_counter() - total_start_time diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index d98f0a5e2..f92d48c30 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -69,6 +69,10 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque guidance_scale_2=request.guidance_scale_2, negative_prompt=request.negative_prompt, enable_teacache=request.enable_teacache, + enable_frame_interpolation=request.enable_frame_interpolation, + frame_interpolation_exp=request.frame_interpolation_exp, + frame_interpolation_scale=request.frame_interpolation_scale, + frame_interpolation_model_path=request.frame_interpolation_model_path, output_path=request.output_path, output_compression=request.output_compression, output_quality=request.output_quality, @@ -158,6 +162,10 @@ async def create_video( guidance_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), enable_teacache: Optional[bool] = Form(False), + enable_frame_interpolation: Optional[bool] = Form(False), + frame_interpolation_exp: Optional[int] = Form(1), + frame_interpolation_scale: Optional[float] = Form(1.0), + frame_interpolation_model_path: Optional[str] = Form(None), output_quality: Optional[str] = Form("default"), output_compression: Optional[int] = Form(None), extra_body: Optional[str] = Form(None), @@ -211,6 +219,10 @@ async def create_video( negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, + enable_frame_interpolation=enable_frame_interpolation, + frame_interpolation_exp=frame_interpolation_exp, + frame_interpolation_scale=frame_interpolation_scale, + frame_interpolation_model_path=frame_interpolation_model_path, output_compression=output_compression, output_quality=output_quality, **( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index f0c980e44..95bc98ef7 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -341,6 +341,10 @@ def save_outputs( audios_out: Optional[list[Any]] = None, frames_out: Optional[list[Any]] = None, output_compression: Optional[int] = None, + enable_frame_interpolation: bool = False, + frame_interpolation_exp: int = 1, + frame_interpolation_scale: float = 1.0, + frame_interpolation_model_path: Optional[str] = None, ) -> list[str]: """Save outputs to files and return the list of file paths.""" output_paths: list[str] = [] @@ -349,6 +353,7 @@ def save_outputs( sample = output if data_type == DataType.VIDEO: sample = attach_audio_to_video_sample(sample, audio, idx) + frames = post_process_sample( sample, data_type, @@ -357,7 +362,12 @@ def save_outputs( save_file_path, audio_sample_rate=audio_sample_rate, output_compression=output_compression, + enable_frame_interpolation=enable_frame_interpolation, + frame_interpolation_exp=frame_interpolation_exp, + frame_interpolation_scale=frame_interpolation_scale, + frame_interpolation_model_path=frame_interpolation_model_path, ) + if samples_out is not None: samples_out.append(sample) if audios_out is not None: @@ -384,14 +394,19 @@ def post_process_sample( save_file_path: Optional[str] = None, audio_sample_rate: Optional[int] = None, output_compression: Optional[int] = None, + enable_frame_interpolation: bool = False, + frame_interpolation_exp: int = 1, + frame_interpolation_scale: float = 1.0, + frame_interpolation_model_path: Optional[str] = None, ): """ - Process sample output and save video if necessary + Process sample output, optionally interpolate video frames, and save. """ audio = None if isinstance(sample, (tuple, list)) and len(sample) == 2: sample, audio = sample + # 1. Convert tensor / array to list of uint8 HWC frames frames = None if isinstance(sample, torch.Tensor): if sample.dim() == 3: @@ -424,7 +439,21 @@ def post_process_sample( arr = (np.clip(arr, 0.0, 1.0) * 255.0).astype(np.uint8) frames = list(arr) - # 2. Save outputs if requested + # 2. Frame interpolation (video only) + if enable_frame_interpolation and data_type == DataType.VIDEO and len(frames) > 1: + from sglang.multimodal_gen.runtime.postprocess import ( + interpolate_video_frames, + ) + + frames, multiplier = interpolate_video_frames( + frames, + exp=frame_interpolation_exp, + scale=frame_interpolation_scale, + model_path=frame_interpolation_model_path, + ) + fps = fps * multiplier + + # 3. Save outputs if requested if save_output: if save_file_path: os.makedirs(os.path.dirname(save_file_path), exist_ok=True) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 6f1d5f904..d1e618d9d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -264,6 +264,10 @@ class GPUWorker: audio=output_batch.audio, audio_sample_rate=output_batch.audio_sample_rate, output_compression=req.output_compression, + enable_frame_interpolation=req.enable_frame_interpolation, + frame_interpolation_exp=req.frame_interpolation_exp, + frame_interpolation_scale=req.frame_interpolation_scale, + frame_interpolation_model_path=req.frame_interpolation_model_path, ) output_batch.output_file_paths = output_paths output_batch.output = None diff --git a/python/sglang/multimodal_gen/runtime/postprocess/__init__.py b/python/sglang/multimodal_gen/runtime/postprocess/__init__.py new file mode 100644 index 000000000..f70951a2d --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/postprocess/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Frame interpolation support for SGLang diffusion pipelines.""" + +from sglang.multimodal_gen.runtime.postprocess.rife_interpolator import ( + FrameInterpolator, + interpolate_video_frames, +) + +__all__ = ["FrameInterpolator", "interpolate_video_frames"] diff --git a/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py b/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py new file mode 100644 index 000000000..1d722e065 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py @@ -0,0 +1,483 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +RIFE 4.22.lite frame interpolation for SGLang diffusion pipelines. + +RIFE model code is vendored and adapted from: + - https://github.com/hzwer/ECCV2022-RIFE (MIT License) + - https://github.com/hzwer/Practical-RIFE (MIT License) + Copyright (c) 2021 Zhewei Huang + +The FrameInterpolator wrapper and integration code are original work. +""" + +import os +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Default HuggingFace repo for RIFE 4.22.lite weights +_DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite" + +# Module-level cache: model_path -> Model instance +_MODEL_CACHE: dict[str, "Model"] = {} + + +# --------------------------------------------------------------------------- +# Vendored RIFE 4.22.lite model code +# (IFBlock, IFNet_HDv3 backbone, Model wrapper) +# --------------------------------------------------------------------------- + + +def warp(tenInput: torch.Tensor, tenFlow: torch.Tensor) -> torch.Tensor: + """Warp tenInput by tenFlow using grid_sample.""" + # Build base grid for the current size + tenHorizontal = ( + torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=tenFlow.device) + .view(1, 1, 1, tenFlow.shape[3]) + .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) + ) + tenVertical = ( + torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=tenFlow.device) + .view(1, 1, tenFlow.shape[2], 1) + .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) + ) + tenGrid = torch.cat([tenHorizontal, tenVertical], dim=1) + + tenFlow = torch.cat( + [ + tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), + tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0), + ], + dim=1, + ) + + grid = (tenGrid + tenFlow).permute(0, 2, 3, 1) + return F.grid_sample( + input=tenInput, + grid=grid, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + +def _conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + """Conv2d + LeakyReLU helper (matches RIFE 4.22 conv()).""" + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +class ResConv(nn.Module): + """Residual convolution block with learnable beta scaling (RIFE 4.22).""" + + def __init__(self, c: int, dilation: int = 1): + super().__init__() + self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.conv(x) * self.beta + x) + + +class IFBlock(nn.Module): + """Single-scale optical flow + mask + feature block (RIFE 4.22).""" + + def __init__(self, in_planes: int, c: int = 64): + super().__init__() + self.conv0 = nn.Sequential( + _conv(in_planes, c // 2, 3, 2, 1), + _conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1), + nn.PixelShuffle(2), + ) + + def forward( + self, + x: torch.Tensor, + flow: Optional[torch.Tensor] = None, + scale: float = 1.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = F.interpolate( + x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False + ) + if flow is not None: + flow = ( + F.interpolate( + flow, + scale_factor=1.0 / scale, + mode="bilinear", + align_corners=False, + ) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + feat = self.conv0(x) + feat = self.convblock(feat) + tmp = self.lastconv(feat) + tmp = F.interpolate( + tmp, scale_factor=scale, mode="bilinear", align_corners=False + ) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + feat = tmp[:, 5:] + return flow, mask, feat + + +class Head(nn.Module): + """Feature encoder producing 4-channel features at full resolution (RIFE 4.22).""" + + def __init__(self): + super().__init__() + self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1) + self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x0 = self.cnn0(x) + x = self.relu(x0) + x1 = self.cnn1(x) + x = self.relu(x1) + x2 = self.cnn2(x) + x = self.relu(x2) + x3 = self.cnn3(x) + return x3 + + +class IFNet(nn.Module): + """4-scale IFNet optical flow network (RIFE 4.22 backbone).""" + + def __init__(self): + super().__init__() + self.block0 = IFBlock(7 + 8, c=192) + self.block1 = IFBlock(8 + 4 + 8 + 8, c=128) + self.block2 = IFBlock(8 + 4 + 8 + 8, c=64) + self.block3 = IFBlock(8 + 4 + 8 + 8, c=32) + self.encode = Head() + + def forward( + self, + x: torch.Tensor, + timestep: float = 0.5, + scale_list: Optional[list] = None, + ) -> tuple[list, torch.Tensor, list]: + if scale_list is None: + scale_list = [8, 4, 2, 1] + + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + + if not torch.is_tensor(timestep): + timestep = (x[:, :1].clone() * 0 + 1) * timestep + else: + timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) + + f0 = self.encode(img0[:, :3]) + f1 = self.encode(img1[:, :3]) + + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + mask = None + + block = [self.block0, self.block1, self.block2, self.block3] + for i in range(4): + if flow is None: + flow, mask, feat = block[i]( + torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), + None, + scale=scale_list[i], + ) + else: + wf0 = warp(f0, flow[:, :2]) + wf1 = warp(f1, flow[:, 2:4]) + fd, m0, feat = block[i]( + torch.cat( + ( + warped_img0[:, :3], + warped_img1[:, :3], + wf0, + wf1, + timestep, + mask, + feat, + ), + 1, + ), + flow, + scale=scale_list[i], + ) + mask = m0 + flow = flow + fd + + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged.append((warped_img0, warped_img1)) + + mask = torch.sigmoid(mask) + merged[3] = warped_img0 * mask + warped_img1 * (1 - mask) + + return flow_list, mask_list[3], merged + + +class Model: + """Wraps IFNet, provides load_model() and inference() API.""" + + def __init__(self): + self.flownet = IFNet() + self.device_type: str = "cpu" + + def eval(self) -> "Model": + self.flownet.eval() + return self + + def device(self) -> torch.device: + return next(self.flownet.parameters()).device + + def load_model(self, path: str, strip_module_prefix: bool = True) -> None: + """Load weights from {path}/flownet.pkl. + + Args: + path: Directory containing ``flownet.pkl``. + strip_module_prefix: If True, strip the ``module.`` prefix that + ``DataParallel`` / ``DistributedDataParallel`` adds to keys. + """ + flownet_path = os.path.join(path, "flownet.pkl") + if not os.path.isfile(flownet_path): + raise FileNotFoundError( + f"RIFE weight file not found: {flownet_path}\n" + "Expected layout: /flownet.pkl" + ) + + def convert(param): + if strip_module_prefix: + return { + k.replace("module.", ""): v + for k, v in param.items() + if "module." in k + } + else: + return {k: v for k, v in param.items() if "module." not in k} + + state = torch.load(flownet_path, map_location="cpu", weights_only=False) + self.flownet.load_state_dict(convert(state), strict=False) + logger.info("Loaded RIFE weights from %s", flownet_path) + + def inference( + self, + img0: torch.Tensor, + img1: torch.Tensor, + scale: float = 1.0, + timestep: float = 0.5, + ) -> torch.Tensor: + """Interpolate a single intermediate frame between img0 and img1.""" + n, c, h, w = img0.shape + + # Pad to multiples of 32 so that RIFE's downsample/upsample round-trips + # preserve spatial dimensions exactly. + ph = ((h - 1) // 32 + 1) * 32 + pw = ((w - 1) // 32 + 1) * 32 + pad = (0, pw - w, 0, ph - h) + img0 = F.pad(img0, pad) + img1 = F.pad(img1, pad) + + imgs = torch.cat((img0, img1), 1) + scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale] + with torch.no_grad(): + flow_list, mask, merged = self.flownet( + imgs, + timestep=timestep, + scale_list=scale_list, + ) + + # Crop back to original resolution + return merged[3][:, :, :h, :w] + + +# --------------------------------------------------------------------------- +# FrameInterpolator public class +# --------------------------------------------------------------------------- + + +class FrameInterpolator: + """ + Lazy-loaded RIFE 4.22.lite frame interpolator. + + Weights are loaded on first call to `.interpolate()` and cached globally + per model_path to avoid reloading across requests. + """ + + def __init__(self, model_path: Optional[str] = None): + self._model_path = model_path + self._resolved_path: Optional[str] = None + + def _ensure_model_loaded(self) -> Model: + """Load RIFE model weights. + + Accepts a local directory **or** a HuggingFace repo ID. When *None* + (the default) the weights are downloaded (and cached) automatically + from ``elfgum/RIFE-4.22.lite`` via ``maybe_download_model()``. + """ + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + maybe_download_model, + ) + + model_path = self._model_path or _DEFAULT_RIFE_HF_REPO + + # Resolve: local path pass-through, HF repo ID → download & cache + model_path = maybe_download_model(model_path) + + self._resolved_path = model_path + + if model_path in _MODEL_CACHE: + return _MODEL_CACHE[model_path] + + device = current_platform.get_local_torch_device() + model = Model() + model.load_model(model_path, strip_module_prefix=True) + model.eval() + model.flownet = model.flownet.to(device) + _MODEL_CACHE[model_path] = model + logger.info("RIFE model loaded on device: %s", device) + return model + + @staticmethod + def _frame_to_tensor(frame: np.ndarray, device: torch.device) -> torch.Tensor: + """Convert uint8 HWC numpy frame to float32 CHW tensor on device.""" + t = torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0 + return t.to(device) + + @staticmethod + def _tensor_to_frame(t: torch.Tensor) -> np.ndarray: + """Convert float32 CHW tensor (batch=1) to uint8 HWC numpy frame.""" + arr = t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy() + return (arr * 255.0).astype(np.uint8) + + def _make_inference( + self, model: Model, I0: torch.Tensor, I1: torch.Tensor, n: int, scale: float + ) -> list[torch.Tensor]: + """ + Recursively generate n-1 intermediate frames between I0 and I1. + + Returns a list of intermediate frame tensors (not including I0 or I1). + """ + if n == 1: + return [model.inference(I0, I1, scale=scale)] + mid = model.inference(I0, I1, scale=scale) + return ( + self._make_inference(model, I0, mid, n // 2, scale) + + [mid] + + self._make_inference(model, mid, I1, n // 2, scale) + ) + + def interpolate( + self, + frames: list[np.ndarray], + exp: int = 1, + scale: float = 1.0, + ) -> tuple[list[np.ndarray], int]: + """ + Interpolate frames using RIFE. + + Args: + frames: List of uint8 numpy arrays with shape [H, W, 3]. + exp: Exponent for interpolation factor. 1 → 2×, 2 → 4×. + scale: RIFE inference scale. Use 0.5 for high-resolution inputs. + + Returns: + (interpolated_frames, multiplier) where multiplier = 2**exp. + """ + if len(frames) < 2: + logger.warning( + "Frame interpolation requires at least 2 frames; returning input unchanged." + ) + return frames, 1 + + model = self._ensure_model_loaded() + device = model.device() + + n_intermediate = 2**exp // 2 # intermediates per adjacent pair + + result: list[np.ndarray] = [] + for i in range(len(frames) - 1): + I0 = self._frame_to_tensor(frames[i], device) + I1 = self._frame_to_tensor(frames[i + 1], device) + + intermediate_tensors = self._make_inference( + model, I0, I1, n_intermediate, scale + ) + + result.append(frames[i]) + for t in intermediate_tensors: + result.append(self._tensor_to_frame(t)) + + result.append(frames[-1]) + multiplier = 2**exp + return result, multiplier + + +# --------------------------------------------------------------------------- +# Module-level convenience function +# --------------------------------------------------------------------------- + + +def interpolate_video_frames( + frames: list[np.ndarray], + exp: int = 1, + scale: float = 1.0, + model_path: Optional[str] = None, +) -> tuple[list[np.ndarray], int]: + """ + Convenience wrapper around FrameInterpolator. + + Args: + frames: List of uint8 HWC numpy frames. + exp: Interpolation exponent (1=2×, 2=4×). + scale: RIFE inference scale (default 1.0; use 0.5 for high-res). + model_path: Local directory or HuggingFace repo ID containing + ``flownet.pkl``. *None* → default ``elfgum/RIFE-4.22.lite``. + + Returns: + (interpolated_frames, multiplier) + """ + interpolator = FrameInterpolator(model_path=model_path) + return interpolator.interpolate(frames, exp=exp, scale=scale) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index f1e062208..08abfd296 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -1210,7 +1210,7 @@ }, "fastwan2_2_ti2v_5b": { "stages_ms": { - "InputValidationStage": 300.00, + "InputValidationStage": 300.0, "TextEncodingStage": 843.86, "TimestepPreparationStage": 58.66, "LatentPreparationStage": 28.55, @@ -1303,7 +1303,7 @@ "39": 1599.78 }, "expected_e2e_ms": 123182.9887, - "expected_avg_denoise_ms": 2831.00, + "expected_avg_denoise_ms": 2831.0, "expected_median_denoise_ms": 1600.09 }, "turbo_wan2_2_i2v_a14b_2gpu": { @@ -1320,7 +1320,7 @@ "denoise_step_ms": { "0": 3042.56, "1": 485.88, - "2": 721.20, + "2": 721.2, "3": 475.58 }, "expected_e2e_ms": 11605.97, @@ -1919,6 +1919,72 @@ "expected_e2e_ms": 2103.05, "expected_avg_denoise_ms": 173.83, "expected_median_denoise_ms": 178.08 + }, + "wan2_1_t2v_1.3b_frame_interp_2x": { + "stages_ms": { + "InputValidationStage": 0.03, + "TextEncodingStage": 1155.78, + "LatentPreparationStage": 0.12, + "TimestepPreparationStage": 2.17, + "DenoisingStage": 4977.09, + "DecodingStage": 98.46, + "per_frame_generation": 1271.97 + }, + "denoise_step_ms": { + "0": 85.33, + "1": 76.41, + "2": 100.5, + "3": 99.28, + "4": 100.14, + "5": 100.33, + "6": 100.17, + "7": 100.35, + "8": 100.14, + "9": 100.32, + "10": 100.18, + "11": 100.37, + "12": 100.17, + "13": 100.34, + "14": 100.21, + "15": 100.35, + "16": 100.15, + "17": 100.32, + "18": 100.23, + "19": 100.35, + "20": 100.19, + "21": 100.57, + "22": 100.27, + "23": 100.24, + "24": 100.27, + "25": 100.27, + "26": 100.27, + "27": 100.29, + "28": 100.25, + "29": 100.25, + "30": 100.3, + "31": 100.3, + "32": 100.23, + "33": 100.25, + "34": 100.26, + "35": 100.27, + "36": 100.24, + "37": 100.26, + "38": 100.26, + "39": 100.29, + "40": 100.31, + "41": 100.31, + "42": 100.31, + "43": 100.26, + "44": 100.26, + "45": 100.27, + "46": 100.3, + "47": 100.27, + "48": 100.23, + "49": 99.65 + }, + "expected_e2e_ms": 6359.87, + "expected_avg_denoise_ms": 99.47, + "expected_median_denoise_ms": 100.27 } } } diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index c137e0b76..094487b0f 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -37,6 +37,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import ( from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack from sglang.multimodal_gen.test.test_utils import ( get_expected_image_format, + get_video_frame_count, is_image_url, prepare_perf_log, validate_image, @@ -663,6 +664,7 @@ def get_generate_fn( seconds: int | None = None, input_reference: Any | None = None, extra_body: dict[Any] | None = None, + expected_frame_count: int | None = None, ) -> str: """ Create a video job via /v1/videos, poll until completion, @@ -746,6 +748,13 @@ def get_generate_fn( tmp_path, expected_filename, expected_width, expected_height ) + if expected_frame_count is not None: + actual_count = get_video_frame_count(tmp_path) + assert actual_count == expected_frame_count, ( + f"{case_id}: frame count mismatch after interpolation — " + f"expected {expected_frame_count}, got {actual_count}" + ) + upload_file_to_slack( case_id=case_id, model=model_path, @@ -976,6 +985,20 @@ def get_generate_fn( extra_body = {} if sampling_params.enable_teacache: extra_body["enable_teacache"] = True + if sampling_params.num_frames: + extra_body["num_frames"] = sampling_params.num_frames + if sampling_params.enable_frame_interpolation: + extra_body["enable_frame_interpolation"] = True + extra_body["frame_interpolation_exp"] = ( + sampling_params.frame_interpolation_exp + ) + + # Compute expected output frame count for validation + expected_frame_count = None + if sampling_params.enable_frame_interpolation and sampling_params.num_frames: + n = sampling_params.num_frames + exp = sampling_params.frame_interpolation_exp + expected_frame_count = (n - 1) * (2**exp) + 1 return _create_and_download_video( client, @@ -985,6 +1008,7 @@ def get_generate_fn( size=output_size, seconds=video_seconds, extra_body=extra_body if extra_body else None, + expected_frame_count=expected_frame_count, ) def generate_image_to_video(case_id, client) -> tuple[str, bytes]: diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index d57c1ae34..f8c980b2f 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -218,6 +218,10 @@ class DiffusionSamplingParams: # TeaCache acceleration enable_teacache: bool = False + # Frame interpolation + enable_frame_interpolation: bool = False + frame_interpolation_exp: int = 1 # 1 = 2×, 2 = 4× + @dataclass(frozen=True) class DiffusionTestCase: @@ -494,6 +498,22 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ enable_teacache=True, ), ), + # Frame interpolation correctness (2× / exp=1) + # Uses the same 1.3B model already in the suite; + DiffusionTestCase( + "wan2_1_t2v_1.3b_frame_interp_2x", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + modality="video", + custom_validator="video", + ), + DiffusionSamplingParams( + prompt=T2V_PROMPT, + num_frames=5, + enable_frame_interpolation=True, + frame_interpolation_exp=1, + ), + ), # LoRA test case for single transformer + merge/unmerge API test # Note: Uses dynamic_lora_path instead of lora_path to test LayerwiseOffload + set_lora interaction # Server starts WITHOUT LoRA, then set_lora is called after startup (Wan models auto-enable layerwise offload) diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 921ec5519..b4efafc65 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -337,6 +337,22 @@ def get_video_dimensions(file_path: str) -> tuple[int, int]: cap.release() +def get_video_frame_count(file_path: str) -> int: + """Return the number of frames in a video file using OpenCV.""" + cap = cv2.VideoCapture(file_path) + try: + count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + if count > 0: + return count + # Fallback: count frames manually + n = 0 + while cap.read()[0]: + n += 1 + return n + finally: + cap.release() + + def validate_video_file( file_path: str, expected_filename: str,