[diffusion] Postprocess: implement frame interpolation using RIFE (#19384)
This commit is contained in:
Junhao Liu
@@ -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
|
||||
|
||||
@@ -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)):
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
**(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
@@ -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: <model_path>/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)
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user