diff --git a/frontend/src/components/SettingsPanel.tsx b/frontend/src/components/SettingsPanel.tsx index b21d592bc..0929ed3f1 100644 --- a/frontend/src/components/SettingsPanel.tsx +++ b/frontend/src/components/SettingsPanel.tsx @@ -87,6 +87,9 @@ interface SettingsPanelProps { onVaceEnabledChange?: (enabled: boolean) => void; vaceContextScale?: number; onVaceContextScaleChange?: (scale: number) => void; + // RIFE settings + rifeEnabled?: boolean; + onRifeEnabledChange?: (enabled: boolean) => void; } export function SettingsPanel({ @@ -126,6 +129,8 @@ export function SettingsPanel({ onVaceEnabledChange, vaceContextScale = 1.0, onVaceContextScaleChange, + rifeEnabled = true, + onRifeEnabledChange, }: SettingsPanelProps) { // Local slider state management hooks const noiseScaleSlider = useLocalSliderValue(noiseScale, onNoiseScaleChange); @@ -389,6 +394,27 @@ export function SettingsPanel({ )} + {/* RIFE Toggle */} +
+
+ + {})} + variant="outline" + size="sm" + className="h-7" + disabled={isStreaming || isLoading} + > + {rifeEnabled ? "ON" : "OFF"} + +
+
+ {currentPipeline?.supportsLoRA && (
0 ? fps.toFixed(1) : "N/A"; const bitrateValue = formatBitrate(bitrate); + const originalFPSValue = + originalFPS !== null && originalFPS !== undefined && originalFPS > 0 + ? originalFPS.toFixed(1) + : null; + const interpolatedFPSValue = + interpolatedFPS !== null && + interpolatedFPS !== undefined && + interpolatedFPS > 0 + ? interpolatedFPS.toFixed(1) + : null; + + // Show detailed FPS (Original + Interpolated) only when RIFE is enabled (interpolatedFPS exists) + // Otherwise show regular FPS + const showDetailedFPS = interpolatedFPSValue !== null; return (
+ {showDetailedFPS ? ( + <> + {originalFPSValue !== null && ( + + )} + {interpolatedFPSValue !== null && ( + + )} + + ) : ( + )}
diff --git a/frontend/src/hooks/useStreamState.ts b/frontend/src/hooks/useStreamState.ts index 5fa290775..a12d4d352 100644 --- a/frontend/src/hooks/useStreamState.ts +++ b/frontend/src/hooks/useStreamState.ts @@ -148,6 +148,7 @@ export function useStreamState() { paused: false, loraMergeStrategy: "permanent_merge", inputMode: initialDefaults.inputMode, + rifeEnabled: true, // RIFE enabled by default }); const [promptData, setPromptData] = useState({ diff --git a/frontend/src/hooks/useWebRTC.ts b/frontend/src/hooks/useWebRTC.ts index 01f985706..af377d748 100644 --- a/frontend/src/hooks/useWebRTC.ts +++ b/frontend/src/hooks/useWebRTC.ts @@ -19,6 +19,7 @@ interface InitialParameters { kv_cache_attention_bias?: number; vace_ref_images?: string[]; vace_context_scale?: number; + rife_enabled?: boolean; } interface UseWebRTCOptions { @@ -38,6 +39,8 @@ export function useWebRTC(options?: UseWebRTCOptions) { useState("new"); const [isConnecting, setIsConnecting] = useState(false); const [isStreaming, setIsStreaming] = useState(false); + const [originalFPS, setOriginalFPS] = useState(null); + const [interpolatedFPS, setInterpolatedFPS] = useState(null); const peerConnectionRef = useRef(null); const dataChannelRef = useRef(null); @@ -115,11 +118,28 @@ export function useWebRTC(options?: UseWebRTCOptions) { peerConnectionRef.current.close(); peerConnectionRef.current = null; } + // Reset FPS data + setOriginalFPS(null); + setInterpolatedFPS(null); + // Notify parent component if (options?.onStreamStop) { options.onStreamStop(); } } + + // Handle FPS update notification from backend + if (data.type === "fps_update") { + if (typeof data.original_fps === "number") { + setOriginalFPS(data.original_fps); + } + if (typeof data.interpolated_fps === "number") { + setInterpolatedFPS(data.interpolated_fps); + } else { + // Clear interpolated FPS if not provided (RIFE disabled) + setInterpolatedFPS(null); + } + } } catch (error) { console.error("Failed to parse data channel message:", error); } @@ -327,6 +347,7 @@ export function useWebRTC(options?: UseWebRTCOptions) { spout_receiver?: { enabled: boolean; name: string }; vace_ref_images?: string[]; vace_context_scale?: number; + rife_enabled?: boolean; }) => { if ( dataChannelRef.current && @@ -373,6 +394,10 @@ export function useWebRTC(options?: UseWebRTCOptions) { // Clear any queued ICE candidates queuedCandidatesRef.current = []; + // Reset FPS data + setOriginalFPS(null); + setInterpolatedFPS(null); + setRemoteStream(null); setConnectionState("new"); setIsStreaming(false); @@ -397,5 +422,7 @@ export function useWebRTC(options?: UseWebRTCOptions) { stopStream, updateVideoTrack, sendParameterUpdate, + originalFPS, + interpolatedFPS, }; } diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index 50fa5b811..b71b9d5b2 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -25,6 +25,7 @@ export interface WebRTCOfferRequest { kv_cache_attention_bias?: number; vace_ref_images?: string[]; vace_context_scale?: number; + rife_enabled?: boolean; }; } @@ -424,4 +425,4 @@ export const getPipelineSchemas = const result = await response.json(); return result; - }; + }; \ No newline at end of file diff --git a/frontend/src/pages/StreamPage.tsx b/frontend/src/pages/StreamPage.tsx index 0de1ef2e6..db02fb4ec 100644 --- a/frontend/src/pages/StreamPage.tsx +++ b/frontend/src/pages/StreamPage.tsx @@ -157,6 +157,8 @@ export function StreamPage() { stopStream, updateVideoTrack, sendParameterUpdate, + originalFPS, + interpolatedFPS, } = useWebRTC(); // Computed loading state - true when downloading models, loading pipeline, or connecting WebRTC @@ -299,6 +301,8 @@ export function StreamPage() { inputMode: modeToUse, denoisingSteps: defaults.denoisingSteps, resolution, + // RIFE is enabled by default for all pipelines + rifeEnabled: settings.rifeEnabled ?? true, noiseScale: defaults.noiseScale, noiseController: defaults.noiseController, loras: [], // Clear LoRA controls when switching pipelines @@ -487,6 +491,16 @@ export function StreamPage() { // Note: This setting requires pipeline reload, so we don't send parameter update here }; + const handleRifeEnabledChange = (enabled: boolean) => { + updateSettings({ rifeEnabled: enabled }); + // Send RIFE enabled update to backend if streaming + if (isStreaming) { + sendParameterUpdate({ + rife_enabled: enabled, + }); + } + }; + const handleRefImagesChange = (images: string[]) => { updateSettings({ refImages: images }); }; @@ -749,6 +763,9 @@ export function StreamPage() { loadParams = { ...loadParams, ...vaceParams }; } + // Add RIFE parameter + loadParams.rife_enabled = settings.rifeEnabled ?? true; + console.log( `Loading ${pipelineIdToUse} with resolution ${resolution.width}x${resolution.height}`, loadParams @@ -791,6 +808,7 @@ export function StreamPage() { spout_receiver?: { enabled: boolean; name: string }; vace_ref_images?: string[]; vace_context_scale?: number; + rife_enabled?: boolean; } = { // Signal the intended input mode to the backend so it doesn't // briefly fall back to text mode before video frames arrive @@ -841,6 +859,9 @@ export function StreamPage() { initialParameters.spout_receiver = settings.spoutReceiver; } + // RIFE interpolation + initialParameters.rife_enabled = settings.rifeEnabled ?? true; + // Reset paused state when starting a fresh stream updateSettings({ paused: false }); @@ -1117,12 +1138,19 @@ export function StreamPage() { onVaceEnabledChange={handleVaceEnabledChange} vaceContextScale={settings.vaceContextScale ?? 1.0} onVaceContextScaleChange={handleVaceContextScaleChange} + rifeEnabled={settings.rifeEnabled ?? true} + onRifeEnabledChange={handleRifeEnabledChange} />
{/* Status Bar */} - + {/* Download Dialog */} {pipelineNeedsModels && ( diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index ba373389a..24e048aa3 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -73,6 +73,8 @@ export interface SettingsState { vaceEnabled?: boolean; refImages?: string[]; vaceContextScale?: number; + // RIFE-specific settings + rifeEnabled?: boolean; } export interface PipelineInfo { diff --git a/src/scope/core/rife/IFNet_HDv3.py b/src/scope/core/rife/IFNet_HDv3.py new file mode 100644 index 000000000..4aa620c6d --- /dev/null +++ b/src/scope/core/rife/IFNet_HDv3.py @@ -0,0 +1,193 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from scope.core.rife.warplayer import warp + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + 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), + ) + + +def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=False, + ), + nn.BatchNorm2d(out_planes), + nn.LeakyReLU(0.2, True), + ) + + +class Head(nn.Module): + 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, feat=False): + 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) + if feat: + return [x0, x1, x2, x3] + return x3 + + +class ResConv(nn.Module): + def __init__(self, c, dilation=1): + super().__init__() + self.conv = nn.Conv2d( + c, c, 3, 1, dilation, dilation=dilation, groups=1, bias=True + ) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x): + return self.relu(self.conv(x) * self.beta + x) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=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, flow=None, scale=1): + 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 IFNet(nn.Module): + 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, + timestep=0.5, + scale_list=(8, 4, 2, 1), + training=False, + fastmode=True, + ensemble=False, + ): + if training is False: + 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], + ) + if ensemble: + print("warning: ensemble is not supported since RIFEv4.21") + 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], + ) + if ensemble: + print("warning: ensemble is not supported since RIFEv4.21") + else: + 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)) + if not fastmode: + print("contextnet is removed") + return flow_list, mask_list[3], merged diff --git a/src/scope/core/rife/RIFE_HDv3.py b/src/scope/core/rife/RIFE_HDv3.py new file mode 100644 index 000000000..9009fc69d --- /dev/null +++ b/src/scope/core/rife/RIFE_HDv3.py @@ -0,0 +1,102 @@ +import itertools + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW + +from scope.core.rife.IFNet_HDv3 import IFNet +from scope.core.rife.loss import EPE, SOBEL + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class Model(nn.Module): + def __init__(self, local_rank=-1): + super().__init__() + self.flownet = IFNet() + self.device() + self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4) + self.epe = EPE() + # self.vgg = VGGPerceptualLoss().to(device) + self.sobel = SOBEL() + if local_rank != -1: + self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) + + def train(self, mode: bool = True): + self.flownet.train(mode) + return super().train(mode) + + def eval(self): + self.flownet.eval() + return super().eval() + + def device(self): + self.flownet.to(device) + + def load_model(self, path, rank=0): + def convert(param): + if rank == -1: + return {k.replace("module.", ""): v for k, v in param.items()} + else: + return param + + if rank > 0: + return + + raw = ( + torch.load(f"{path}/flownet.pkl") + if torch.cuda.is_available() + else torch.load(f"{path}/flownet.pkl", map_location="cpu") + ) + state_dict = convert(raw) + + # Filter out unexpected keys (teacher.*, caltime.*, etc.) to avoid strict load errors + model_state = self.flownet.state_dict() + filtered = {k: v for k, v in state_dict.items() if k in model_state} + unexpected = [k for k in state_dict.keys() if k not in model_state] + if unexpected: + print(f"[RIFE] Ignoring {len(unexpected)} unexpected keys: {unexpected[:3]}...") + + self.flownet.load_state_dict(filtered, strict=False) + + def save_model(self, path, rank=0): + if rank == 0: + torch.save(self.flownet.state_dict(), f"{path}/flownet.pkl") + + def inference(self, img0, img1, timestep=0.5, scale=1.0): + imgs = torch.cat((img0, img1), 1) + scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale] + flow, mask, merged = self.flownet(imgs, timestep=timestep, scale_list=scale_list) + return merged[3] + + def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): + for param_group in self.optimG.param_groups: + param_group["lr"] = learning_rate + img0 = imgs[:, :3] + img1 = imgs[:, 3:] + if training: + self.train() + else: + self.eval() + scale = [8, 4, 2, 1] + flow, mask, merged = self.flownet( + torch.cat((imgs, gt), 1), timestep=0.5, scale_list=scale, training=training + ) + loss_l1 = (merged[3] - gt).abs().mean() + loss_smooth = self.sobel(flow[3], flow[3] * 0).mean() + # loss_vgg = self.vgg(merged[2], gt) + if training: + self.optimG.zero_grad() + loss_G = loss_l1 + loss_smooth * 0.1 + loss_G.backward() + self.optimG.step() + else: + flow_teacher = flow[2] + return merged[3], { + "mask": mask, + "flow": flow[3][:, :2], + "loss_l1": loss_l1, + "loss_smooth": loss_smooth, + } diff --git a/src/scope/core/rife/__init__.py b/src/scope/core/rife/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/scope/core/rife/loss.py b/src/scope/core/rife/loss.py new file mode 100644 index 000000000..72e5de6af --- /dev/null +++ b/src/scope/core/rife/loss.py @@ -0,0 +1,128 @@ +import torch +import numpy as np +import torch.nn as nn +import torch.nn.functional as F +import torchvision.models as models + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class EPE(nn.Module): + def __init__(self): + super(EPE, self).__init__() + + def forward(self, flow, gt, loss_mask): + loss_map = (flow - gt.detach()) ** 2 + loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5 + return (loss_map * loss_mask) + + +class Ternary(nn.Module): + def __init__(self): + super(Ternary, self).__init__() + patch_size = 7 + out_channels = patch_size * patch_size + self.w = np.eye(out_channels).reshape( + (patch_size, patch_size, 1, out_channels)) + self.w = np.transpose(self.w, (3, 2, 0, 1)) + self.w = torch.tensor(self.w).float().to(device) + + def transform(self, img): + patches = F.conv2d(img, self.w, padding=3, bias=None) + transf = patches - img + transf_norm = transf / torch.sqrt(0.81 + transf**2) + return transf_norm + + def rgb2gray(self, rgb): + r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :] + gray = 0.2989 * r + 0.5870 * g + 0.1140 * b + return gray + + def hamming(self, t1, t2): + dist = (t1 - t2) ** 2 + dist_norm = torch.mean(dist / (0.1 + dist), 1, True) + return dist_norm + + def valid_mask(self, t, padding): + n, _, h, w = t.size() + inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t) + mask = F.pad(inner, [padding] * 4) + return mask + + def forward(self, img0, img1): + img0 = self.transform(self.rgb2gray(img0)) + img1 = self.transform(self.rgb2gray(img1)) + return self.hamming(img0, img1) * self.valid_mask(img0, 1) + + +class SOBEL(nn.Module): + def __init__(self): + super(SOBEL, self).__init__() + self.kernelX = torch.tensor([ + [1, 0, -1], + [2, 0, -2], + [1, 0, -1], + ]).float() + self.kernelY = self.kernelX.clone().T + self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device) + self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device) + + def forward(self, pred, gt): + N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3] + img_stack = torch.cat( + [pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0) + sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1) + sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1) + pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:] + pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:] + + L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y) + loss = (L1X+L1Y) + return loss + +class MeanShift(nn.Conv2d): + def __init__(self, data_mean, data_std, data_range=1, norm=True): + c = len(data_mean) + super(MeanShift, self).__init__(c, c, kernel_size=1) + std = torch.Tensor(data_std) + self.weight.data = torch.eye(c).view(c, c, 1, 1) + if norm: + self.weight.data.div_(std.view(c, 1, 1, 1)) + self.bias.data = -1 * data_range * torch.Tensor(data_mean) + self.bias.data.div_(std) + else: + self.weight.data.mul_(std.view(c, 1, 1, 1)) + self.bias.data = data_range * torch.Tensor(data_mean) + self.requires_grad = False + +class VGGPerceptualLoss(torch.nn.Module): + def __init__(self, rank=0): + super(VGGPerceptualLoss, self).__init__() + blocks = [] + pretrained = True + self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features + self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda() + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X, Y, indices=None): + X = self.normalize(X) + Y = self.normalize(Y) + indices = [2, 7, 12, 21, 30] + weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5] + k = 0 + loss = 0 + for i in range(indices[-1]): + X = self.vgg_pretrained_features[i](X) + Y = self.vgg_pretrained_features[i](Y) + if (i+1) in indices: + loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1 + k += 1 + return loss + +if __name__ == '__main__': + img0 = torch.zeros(3, 3, 256, 256).float().to(device) + img1 = torch.tensor(np.random.normal( + 0, 1, (3, 3, 256, 256))).float().to(device) + ternary_loss = Ternary() + print(ternary_loss(img0, img1).shape) diff --git a/src/scope/core/rife/warplayer.py b/src/scope/core/rife/warplayer.py new file mode 100644 index 000000000..21b0b904c --- /dev/null +++ b/src/scope/core/rife/warplayer.py @@ -0,0 +1,22 @@ +import torch +import torch.nn as nn + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +backwarp_tenGrid = {} + + +def warp(tenInput, tenFlow): + k = (str(tenFlow.device), str(tenFlow.size())) + if k not in backwarp_tenGrid: + tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=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=device).view( + 1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) + backwarp_tenGrid[k] = torch.cat( + [tenHorizontal, tenVertical], 1).to(device) + + tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), + tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1) + + g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1) + return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True) diff --git a/src/scope/server/frame_processor.py b/src/scope/server/frame_processor.py index ca58d0c24..789033681 100644 --- a/src/scope/server/frame_processor.py +++ b/src/scope/server/frame_processor.py @@ -9,6 +9,7 @@ from aiortc.mediastreams import VideoFrame from .pipeline_manager import PipelineManager, PipelineNotAvailableException +from .rife_interpolation import RIFEInterpolator logger = logging.getLogger(__name__) @@ -73,11 +74,19 @@ def __init__( self.processing_time_per_frame = deque( maxlen=2 ) # Keep last 2 processing_time/num_frames values for averaging + self.original_fps_data = deque( + maxlen=2 + ) # Keep last 2 original frame generation FPS values + self.interpolated_fps_data = deque( + maxlen=2 + ) # Keep last 2 interpolated FPS values self.last_fps_update = time.time() self.fps_update_interval = 0.5 # Update FPS every 0.5 seconds self.min_fps = MIN_FPS self.max_fps = MAX_FPS self.current_pipeline_fps = DEFAULT_FPS + self.current_original_fps = DEFAULT_FPS # Original frame generation FPS + self.current_interpolated_fps = None # After interpolation FPS (None when RIFE disabled) self.fps_lock = threading.Lock() # Lock for thread-safe FPS updates # Input FPS tracking variables @@ -108,6 +117,12 @@ def __init__( # This determines whether we wait for video frames or generate immediately. self._video_mode = (initial_parameters or {}).get("input_mode") == "video" + # RIFE interpolation + rife_enabled = (initial_parameters or {}).get("rife_enabled", True) + self.rife_interpolator = RIFEInterpolator( + enabled=rife_enabled, device=torch.device("cuda" if torch.cuda.is_available() else "cpu") + ) + def start(self): if self.running: return @@ -229,6 +244,16 @@ def get_current_pipeline_fps(self) -> float: with self.fps_lock: return self.current_pipeline_fps + def get_original_fps(self) -> float: + """Get the original frame generation FPS (before interpolation)""" + with self.fps_lock: + return self.current_original_fps + + def get_interpolated_fps(self) -> float | None: + """Get the interpolated FPS (after RIFE interpolation), returns None if RIFE is disabled""" + with self.fps_lock: + return self.current_interpolated_fps + def get_output_fps(self) -> float: """Get the output FPS that frames should be sent at. @@ -239,6 +264,10 @@ def get_output_fps(self) -> float: input_fps = self._get_input_fps() pipeline_fps = self.get_current_pipeline_fps() + # RIFE interpolation doubles the frame count, so double the output rate + if self.rife_interpolator.enabled: + pipeline_fps = pipeline_fps * 2 + if input_fps is None: return pipeline_fps @@ -325,6 +354,39 @@ def _calculate_pipeline_fps(self, start_time: float, num_frames: int): with self.fps_lock: self.current_pipeline_fps = estimated_fps + # Update original and interpolated FPS if we have data + if len(self.original_fps_data) >= 1: + avg_original_fps = sum(self.original_fps_data) / len(self.original_fps_data) + avg_original_fps = max(self.min_fps, min(self.max_fps, avg_original_fps)) + with self.fps_lock: + self.current_original_fps = avg_original_fps + + # Only calculate interpolated FPS if RIFE is enabled + if self.rife_interpolator.enabled and len(self.interpolated_fps_data) >= 1: + avg_interpolated_fps = sum(self.interpolated_fps_data) / len(self.interpolated_fps_data) + avg_interpolated_fps = max(self.min_fps, min(self.max_fps, avg_interpolated_fps)) + with self.fps_lock: + self.current_interpolated_fps = avg_interpolated_fps + else: + # If RIFE is disabled or no interpolated data, set to None (will not be sent to frontend) + with self.fps_lock: + self.current_interpolated_fps = None + + # Send FPS update to frontend via notification callback + if self.notification_callback: + with self.fps_lock: + fps_data = { + "type": "fps_update", + "original_fps": self.current_original_fps, + } + # Only include interpolated_fps if RIFE is enabled and we have data + if self.current_interpolated_fps is not None: + fps_data["interpolated_fps"] = self.current_interpolated_fps + try: + self.notification_callback(fps_data) + except Exception as e: + logger.debug(f"Failed to send FPS update notification: {e}") + self.last_fps_update = current_time def _get_pipeline_dimensions(self) -> tuple[int, int]: @@ -351,6 +413,24 @@ def update_parameters(self, parameters: dict[str, Any]): spout_config = parameters.pop("spout_receiver") self._update_spout_receiver(spout_config) + # Handle RIFE interpolation settings + if "rife_enabled" in parameters: + rife_enabled = parameters.pop("rife_enabled") + try: + self.rife_interpolator.set_enabled(rife_enabled) + logger.info(f"RIFE interpolation {'enabled' if rife_enabled else 'disabled'}") + # Clear interpolated FPS data when RIFE is disabled + if not rife_enabled: + with self.fps_lock: + self.interpolated_fps_data.clear() + self.current_interpolated_fps = None + except Exception as e: + logger.error(f"Failed to {'enable' if rife_enabled else 'disable'} RIFE interpolation: {e}") + # If enabling fails, keep RIFE disabled + if rife_enabled: + self.rife_interpolator.enabled = False + raise + # Put new parameters in queue (replace any pending update) try: # Add new update @@ -749,11 +829,16 @@ def process_chunk(self): self.parameters.pop("transition", None) processing_time = time.time() - start_time - num_frames = output.shape[0] + original_num_frames = output.shape[0] logger.debug( - f"Processed pipeline in {processing_time:.4f}s, {num_frames} frames" + f"Processed pipeline in {processing_time:.4f}s, {original_num_frames} frames" ) + # Calculate original frame generation FPS (before interpolation) + if processing_time > 0 and original_num_frames > 0: + original_fps = original_num_frames / processing_time + self.original_fps_data.append(original_fps) + # Normalize to [0, 255] and convert to uint8 output = ( (output * 255.0) @@ -764,6 +849,26 @@ def process_chunk(self): .cpu() ) + # Apply RIFE interpolation if enabled + rife_interpolation_time = None + if self.rife_interpolator.enabled: + rife_start_time = time.time() + output = self.rife_interpolator.interpolate(output) + rife_interpolation_time = time.time() - rife_start_time + interpolated_num_frames = output.shape[0] + logger.debug( + f"RIFE interpolation: {original_num_frames} frames -> {interpolated_num_frames} frames in {rife_interpolation_time:.4f}s" + ) + num_frames = interpolated_num_frames + + # Calculate interpolated FPS (total time including interpolation) + total_processing_time = processing_time + rife_interpolation_time + if total_processing_time > 0 and interpolated_num_frames > 0: + interpolated_fps = interpolated_num_frames / total_processing_time + self.interpolated_fps_data.append(interpolated_fps) + else: + num_frames = original_num_frames + # Resize output queue to meet target max size target_output_queue_max_size = num_frames * OUTPUT_QUEUE_MAX_SIZE_FACTOR if self.output_queue.maxsize < target_output_queue_max_size: @@ -864,5 +969,12 @@ def _is_recoverable(error: Exception) -> bool: """ if isinstance(error, torch.cuda.OutOfMemoryError): return False + + # RIFE-related errors are not recoverable - RIFE must be properly installed + if isinstance(error, RuntimeError): + error_msg = str(error) + if "RIFE" in error_msg and ("not available" in error_msg or "not found" in error_msg or "Failed to load" in error_msg): + return False + # Add more non-recoverable error types here as needed return True diff --git a/src/scope/server/pipeline_manager.py b/src/scope/server/pipeline_manager.py index ef8873ec7..ee492b821 100644 --- a/src/scope/server/pipeline_manager.py +++ b/src/scope/server/pipeline_manager.py @@ -6,6 +6,7 @@ import os import threading from enum import Enum +from pathlib import Path from typing import Any import torch @@ -233,11 +234,25 @@ def _configure_vace(self, config: dict, load_params: dict | None = None) -> None Adds vace_path to config and optionally extracts VACE-specific parameters from load_params (ref_images, vace_context_scale). + If VACE model file doesn't exist, VACE will be automatically disabled + by setting vace_path to None. + Args: config: Pipeline configuration dict to modify load_params: Optional load parameters containing VACE settings """ - config["vace_path"] = self._get_vace_checkpoint_path() + vace_path = self._get_vace_checkpoint_path() + + # Check if VACE model file exists + if not Path(vace_path).exists(): + logger.warning( + f"VACE model file not found at {vace_path}. " + f"VACE will be disabled. To enable VACE, download the model first." + ) + config["vace_path"] = None + return + + config["vace_path"] = vace_path logger.debug(f"_configure_vace: Using VACE checkpoint at {config['vace_path']}") # Extract VACE-specific parameters from load_params if present diff --git a/src/scope/server/rife_interpolation.py b/src/scope/server/rife_interpolation.py new file mode 100644 index 000000000..8bb2f2749 --- /dev/null +++ b/src/scope/server/rife_interpolation.py @@ -0,0 +1,354 @@ +"""RIFE (Real-Time Intermediate Flow Estimation) HDv3 frame interpolation module. + +This module provides frame interpolation functionality using RIFE HDv3 to double +the frame rate of video output from the pipeline. + +The RIFE HDv3 model code is integrated into scope.core.rife, and model weights +should be placed in ~/.daydream-scope/models/RIFE/flownet.pkl. +""" + +import logging +from pathlib import Path +from typing import Optional + +import torch +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +# Try to import RIFE model +RIFE_AVAILABLE = False +RIFE_MODEL_CLASS = None + +# Import RIFE HDv3 model from our codebase +try: + from scope.core.rife.RIFE_HDv3 import Model as RIFEModel + RIFE_MODEL_CLASS = RIFEModel + RIFE_AVAILABLE = True + logger.info("RIFE HDv3 model found and imported") +except ImportError as e: + RIFE_AVAILABLE = False + logger.debug(f"RIFE HDv3 import failed: {e}") + +class RIFEInterpolator: + """RIFE HDv3-based frame interpolator. + + This class handles frame interpolation using RIFE HDv3 to generate intermediate + frames between consecutive frames, effectively doubling the frame rate. + + Attributes: + enabled: Whether interpolation is enabled + device: Device to run interpolation on + model: RIFE HDv3 model instance (if available) + model_path: Path to RIFE HDv3 model weights directory + """ + + def __init__( + self, + enabled: bool = False, + device: Optional[torch.device] = None, + model_path: Optional[str] = None, + ): + """Initialize RIFE interpolator. + + Args: + enabled: Whether interpolation is enabled + device: Device to run interpolation on (defaults to CUDA if available, else CPU) + model_path: Optional path to RIFE model weights file + """ + self.enabled = enabled + self.device = device or (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")) + self.model = None + self.model_path = model_path + + if enabled: + if not RIFE_AVAILABLE: + raise ImportError( + "RIFE interpolation requested but RIFE HDv3 is not available. " + "Please install RIFE HDv3 from https://github.com/hzwer/arXiv2020-RIFE. " + "See docs/rife.md for installation instructions." + ) + + try: + self._load_model() + if self.model is None: + raise RuntimeError( + "RIFE HDv3 model weights not found. " + "Please download RIFE HDv3 model weights and place flownet.pkl in " + "arXiv2020-RIFE/train_log/ directory. " + "See docs/rife.md for installation instructions." + ) + logger.info("RIFE HDv3 model loaded successfully") + except Exception as e: + logger.error(f"Failed to load RIFE model: {e}") + raise RuntimeError( + f"Failed to load RIFE model: {e}. " + "Please ensure RIFE is properly installed. " + "See docs/rife.md for installation instructions." + ) from e + + def _load_model(self): + """Load RIFE model.""" + if not RIFE_AVAILABLE or RIFE_MODEL_CLASS is None: + return + + try: + # Initialize RIFE model + self.model = RIFE_MODEL_CLASS() + + # Find model weights directory (RIFE expects a directory containing flownet.pkl) + model_dir_path = None + if self.model_path: + # If path is a file, get its parent directory + model_path_obj = Path(self.model_path) + if model_path_obj.is_file(): + model_dir_path = str(model_path_obj.parent) + elif model_path_obj.is_dir(): + model_dir_path = str(model_path_obj) + else: + # Try default model directories + from .models_config import get_models_dir + + # Get project root directory (where pyproject.toml is located) + # Start from this file's directory and walk up to find project root + current_file = Path(__file__).resolve() + project_root = None + for parent in current_file.parents: + if (parent / "pyproject.toml").exists(): + project_root = parent + break + + default_model_dirs = [] + # First priority: project root weights folder + if project_root: + weights_dir = project_root / "weights" / "RIFE" + default_model_dirs.append(weights_dir) + + # Second priority: configured models directory + default_model_dirs.append(get_models_dir() / "RIFE") + + # Third priority: default home directory + default_model_dirs.append(Path.home() / ".daydream-scope" / "models" / "RIFE") + + for model_dir in default_model_dirs: + if model_dir.exists() and (model_dir / "flownet.pkl").exists(): + model_dir_path = str(model_dir) + break + + if model_dir_path: + # Load model (RIFE uses -1 for latest version) + # RIFE's load_model expects a directory path, not a file path + self.model.load_model(model_dir_path, -1) + self.model.eval() + self.model.device() # Move model to device + logger.info(f"Loaded RIFE HDv3 weights from {model_dir_path}/flownet.pkl") + + # Compile model for faster inference (PyTorch 2.0+) + if hasattr(torch, "compile"): + try: + self.model.flownet = torch.compile( + self.model.flownet, + mode="reduce-overhead", # Best for inference + fullgraph=False, # Allow graph breaks for compatibility + ) + logger.info("RIFE flownet compiled with torch.compile()") + except Exception as e: + logger.warning(f"torch.compile() failed, using eager mode: {e}") + else: + raise FileNotFoundError( + "RIFE HDv3 model weights (flownet.pkl) not found. " + "Please download RIFE HDv3 model from https://github.com/hzwer/arXiv2020-RIFE " + "and place flownet.pkl in ~/.daydream-scope/models/RIFE/ directory. " + "See docs/rife.md for installation instructions." + ) + except FileNotFoundError: + # Re-raise FileNotFoundError for model weights + raise + except Exception as e: + logger.error(f"Error loading RIFE HDv3 model: {e}", exc_info=True) + raise RuntimeError( + f"Failed to load RIFE HDv3 model: {e}. " + "See docs/rife.md for installation instructions." + ) from e + + def interpolate(self, frames: torch.Tensor) -> torch.Tensor: + """Interpolate frames to double the frame rate. + + Args: + frames: Input frames tensor of shape [T, H, W, C] with values in [0, 255] (uint8) + + Returns: + Interpolated frames tensor of shape [T*2-1, H, W, C] with values in [0, 255] (uint8) + + Raises: + RuntimeError: If RIFE is enabled but model is not available or loaded + """ + if not self.enabled: + return frames + + if frames.shape[0] < 2: + # Can't interpolate with less than 2 frames + return frames + + if not RIFE_AVAILABLE or self.model is None: + raise RuntimeError( + "RIFE interpolation is enabled but RIFE HDv3 model is not available. " + "Please ensure RIFE HDv3 is properly installed and model weights are loaded. " + "See docs/rife.md for installation instructions." + ) + + # Convert to float32 for processing + frames_float = frames.float() + + # Use RIFE for interpolation + return self._rife_interpolate(frames_float) + + def _rife_interpolate(self, frames: torch.Tensor) -> torch.Tensor: + """Use RIFE model for interpolation. + + Optimized implementation with: + - BF16 mixed precision for faster tensor core inference + - Batched processing of all frame pairs in single forward pass + - Pre-computed padding applied once to all frames + - Minimal CPU-GPU transfers (single transfer at end) + + Args: + frames: Input frames tensor of shape [T, H, W, C] with values in [0, 255] + + Returns: + Interpolated frames tensor of shape [T*2-1, H, W, C] with values in [0, 255] (uint8) + """ + num_frames = frames.shape[0] + if num_frames < 2: + return frames.clamp(0, 255).to(torch.uint8) + + T, H, W, C = frames.shape + + # Convert from [T, H, W, C] to [T, C, H, W] and normalize to [0, 1] + # Move to GPU once at the start + frames_chw = (frames.permute(0, 3, 1, 2) / 255.0).to(self.device).contiguous() + + # Calculate padding - RIFE requires dimensions to be multiples of 32 + tmp = 32 # For scale=1.0 + ph = ((H - 1) // tmp + 1) * tmp + pw = ((W - 1) // tmp + 1) * tmp + padding = (0, pw - W, 0, ph - H) + + # Pre-compute padding for all frames at once (optimization) + frames_padded = F.pad(frames_chw, padding) # [T, C, pH, pW] + + with torch.no_grad(): + # Use BF16 mixed precision for faster inference on modern GPUs + autocast_dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32 + with torch.amp.autocast(device_type=self.device.type, dtype=autocast_dtype): + # Batch all frame pairs: frames[0:T-1] and frames[1:T] + frames1_padded = frames_padded[:-1] # [T-1, C, pH, pW] + frames2_padded = frames_padded[1:] # [T-1, C, pH, pW] + + # Batched inference - process all pairs at once + mid_frames_padded = self.model.inference( + frames1_padded, frames2_padded, scale=1.0 + ) # [T-1, C, pH, pW] + + # Remove padding from interpolated frames (stay on GPU) + mid_frames = mid_frames_padded[:, :, :H, :W].float() # [T-1, C, H, W] + + # Get original frames without padding (stay on GPU) + original_frames = frames_padded[:, :, :H, :W] # [T, C, H, W] + + # Interleave original and interpolated frames on GPU + # Result: [orig[0], mid[0], orig[1], mid[1], ..., orig[T-2], mid[T-2], orig[T-1]] + result_frames = torch.zeros( + (T * 2 - 1, C, H, W), dtype=torch.float32, device=self.device + ) + result_frames[0::2] = original_frames # Original frames at even indices + result_frames[1::2] = mid_frames # Interpolated frames at odd indices + + # Scale to [0, 255] and transfer to CPU once at the end + result_chw = (result_frames * 255.0).cpu() + + # Convert back to [T*2-1, H, W, C] + result = result_chw.permute(0, 2, 3, 1).contiguous() + + # Clamp to valid range and convert to uint8 + result = result.clamp(0.0, 255.0).to(torch.uint8) + + return result + + def set_enabled(self, enabled: bool): + """Enable or disable interpolation. + + Args: + enabled: Whether to enable interpolation + + Raises: + RuntimeError: If enabling RIFE but model is not available + """ + if enabled: + if not RIFE_AVAILABLE: + raise RuntimeError( + "RIFE interpolation cannot be enabled: RIFE HDv3 is not available. " + "Please install RIFE HDv3 from https://github.com/hzwer/arXiv2020-RIFE. " + "See docs/rife.md for installation instructions." + ) + + if self.model is None: + # Try to load model if not already loaded + try: + self._load_model() + if self.model is None: + raise RuntimeError( + "RIFE HDv3 model weights not found. " + "Please download RIFE HDv3 model weights. " + "See docs/rife.md for installation instructions." + ) + except Exception as e: + raise RuntimeError( + f"Failed to load RIFE HDv3 model when enabling: {e}. " + "See docs/rife.md for installation instructions." + ) from e + + self.enabled = enabled + + +def is_rife_available() -> bool: + """Check if RIFE is available for use. + + Returns: + True if RIFE is available, False otherwise + """ + return RIFE_AVAILABLE + + +def get_rife_model_path() -> Optional[Path]: + """Get the default RIFE HDv3 model directory path. + + Returns: + Path to RIFE HDv3 model directory (containing flownet.pkl) if found, None otherwise + """ + from .models_config import get_models_dir + + # Get project root directory + current_file = Path(__file__).resolve() + project_root = None + for parent in current_file.parents: + if (parent / "pyproject.toml").exists(): + project_root = parent + break + + default_dirs = [] + # First priority: project root weights folder + if project_root: + weights_dir = project_root / "weights" / "RIFE" + default_dirs.append(weights_dir) + + # Second priority: configured models directory + default_dirs.append(get_models_dir() / "RIFE") + + # Third priority: default home directory + default_dirs.append(Path.home() / ".daydream-scope" / "models" / "RIFE") + + for model_dir in default_dirs: + if model_dir.exists() and (model_dir / "flownet.pkl").exists(): + return model_dir + return None diff --git a/src/scope/server/schema.py b/src/scope/server/schema.py index 6cb6ed28e..e4639c404 100644 --- a/src/scope/server/schema.py +++ b/src/scope/server/schema.py @@ -113,6 +113,10 @@ class Parameters(BaseModel): ge=0.0, le=2.0, ) + rife_enabled: bool | None = Field( + default=None, + description="Enable RIFE (Real-Time Intermediate Flow Estimation) frame interpolation to double the frame rate of the output video. This increases smoothness but may add latency.", + ) class SpoutConfig(BaseModel): @@ -345,6 +349,10 @@ class LongLiveLoadParams(LoRAEnabledLoadParams): default=True, description="Enable VACE (Video All-In-One Creation and Editing) support for reference image conditioning and structural guidance. When enabled, incoming video in V2V mode is routed to VACE for conditioning. When disabled, V2V uses faster regular encoding.", ) + rife_enabled: bool = Field( + default=True, + description="Enable RIFE (Real-Time Intermediate Flow Estimation) frame interpolation to double the frame rate of the output video. This increases smoothness but may add latency.", + ) class KreaRealtimeVideoLoadParams(LoRAEnabledLoadParams): diff --git a/src/scope/server/webrtc.py b/src/scope/server/webrtc.py index d0d2c458d..d485cfb4e 100644 --- a/src/scope/server/webrtc.py +++ b/src/scope/server/webrtc.py @@ -105,7 +105,11 @@ def _send_message_threadsafe(self, message: dict): def send_sync(): try: self.data_channel.send(message_str) - logger.info(f"Sent notification to frontend: {message}") + # Use debug level for frequent updates like fps_update + if message.get("type") == "fps_update": + logger.debug(f"Sent notification to frontend: {message}") + else: + logger.info(f"Sent notification to frontend: {message}") except Exception as e: logger.error(f"Failed to send notification: {e}") diff --git a/weights/flownet.pkl b/weights/flownet.pkl new file mode 100644 index 000000000..d4e22122b Binary files /dev/null and b/weights/flownet.pkl differ