Spaces:
Running on Zero
Running on Zero
File size: 6,699 Bytes
ee69ab7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from .config import LQProjectionConfig
class CausalConv3d(nn.Conv3d):
"""Causal convolution matching the cache behavior used by FlashVSR 1.1."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
pt, ph, pw = self.padding
self._causal_padding = (pw, pw, ph, ph, 2 * pt, 0)
self.padding = (0, 0, 0)
def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None) -> torch.Tensor:
padding = list(self._causal_padding)
if cache is not None and padding[4] > 0:
x = torch.cat([cache.to(device=x.device, dtype=x.dtype), x], dim=2)
padding[4] = max(0, padding[4] - cache.shape[2])
x = F.pad(x, padding, mode="replicate")
return super().forward(x)
class ChannelRMSNorm(nn.Module):
def __init__(self, channels: int, eps: float = 1e-6):
super().__init__()
# Keep the FlashVSR 1.1 checkpoint shape: [C,1,1,1].
self.weight = nn.Parameter(torch.ones(channels, 1, 1, 1))
self.scale = math.sqrt(channels)
def forward(self, x: torch.Tensor) -> torch.Tensor:
value = F.normalize(x, dim=1) * self.scale
return value * self.weight.unsqueeze(0)
class PixelUnshuffle3D(nn.Module):
def __init__(self, spatial_factor: int):
super().__init__()
self.spatial_factor = int(spatial_factor)
def forward(self, x: torch.Tensor) -> torch.Tensor:
factor = self.spatial_factor
batch, channels, frames, height, width = x.shape
if height % factor or width % factor:
raise ValueError(f"LQ video spatial size {(height, width)} must be divisible by {factor}")
x = x.view(batch, channels, frames, height // factor, factor, width // factor, factor)
return x.permute(0, 1, 4, 6, 2, 3, 5).reshape(
batch, channels * factor * factor, frames, height // factor, width // factor
)
class LQProjection(nn.Module):
"""Independent replacement for the original Buffer_LQ4x_Proj."""
def __init__(self, config: LQProjectionConfig, output_dim: int):
super().__init__()
self.config = config
factor = config.spatial_factor
self.pixel_shuffle = PixelUnshuffle3D(factor)
self.conv1 = CausalConv3d(
config.input_channels * factor * factor,
config.hidden_dim,
kernel_size=(4, 3, 3),
stride=(2, 1, 1),
padding=(1, 1, 1),
)
self.norm1 = ChannelRMSNorm(config.hidden_dim)
self.conv2 = CausalConv3d(
config.hidden_dim,
config.hidden_dim_2,
kernel_size=(4, 3, 3),
stride=(2, 1, 1),
padding=(1, 1, 1),
)
self.norm2 = ChannelRMSNorm(config.hidden_dim_2)
self.act = nn.SiLU()
self.linear_layers = nn.ModuleList(
nn.Linear(config.hidden_dim_2, output_dim) for _ in range(config.num_injection_layers)
)
if config.zero_init:
for layer in self.linear_layers:
nn.init.zeros_(layer.weight)
nn.init.zeros_(layer.bias)
self.clear_cache()
def clear_cache(self) -> None:
self._stream_cache1: torch.Tensor | None = None
self._stream_cache2: torch.Tensor | None = None
self._stream_clip_index = 0
def stream_forward(self, video_clip: torch.Tensor) -> list[torch.Tensor] | None:
"""Process one FlashVSR 1.1 LQ chunk while preserving causal conv state.
The first call is the conv2 warm-up and intentionally returns ``None``.
Every subsequent complete 4-frame call returns one latent-token frame.
"""
if video_clip.ndim != 5:
raise ValueError(f"LQProjection expects [B,C,F,H,W], got {video_clip.shape}")
if self._stream_clip_index == 0:
if video_clip.shape[2] != 1:
raise ValueError(
"the first LQ stream chunk must contain exactly one frame"
)
first = video_clip[:, :, :1].repeat(1, 1, 3, 1, 1)
video_clip = torch.cat([first, video_clip], dim=2)
elif video_clip.shape[2] != 4:
raise ValueError(
f"later LQ stream chunks must contain 4 frames, got {video_clip.shape[2]}"
)
x = self.pixel_shuffle(video_clip)
next_cache1 = x[:, :, -2:].clone()
x = self.act(self.norm1(self.conv1(x, self._stream_cache1)))
self._stream_cache1 = next_cache1
next_cache2 = x[:, :, -2:].clone()
if self._stream_clip_index == 0:
self._stream_cache2 = next_cache2
self._stream_clip_index += 1
return None
x = self.act(self.norm2(self.conv2(x, self._stream_cache2)))
self._stream_cache2 = next_cache2
self._stream_clip_index += 1
tokens = x.permute(0, 2, 3, 4, 1).flatten(1, 3)
return [layer(tokens) for layer in self.linear_layers]
def forward(self, video: torch.Tensor) -> list[torch.Tensor]:
if video.ndim != 5:
raise ValueError(f"LQProjection expects [B,C,F,H,W], got {video.shape}")
if self.config.pad_last_frames:
video = torch.cat(
[video, video[:, :, -1:].repeat(1, 1, self.config.pad_last_frames, 1, 1)], dim=2
)
iterations = 1 + (video.shape[2] - 1) // 4
video = torch.cat([video[:, :, :1].repeat(1, 1, 3, 1, 1), video], dim=2)
cache1 = None
cache2 = None
output_chunks = []
for chunk_index in range(iterations):
chunk = video[:, :, chunk_index * 4 : (chunk_index + 1) * 4]
if chunk.shape[2] != 4:
raise ValueError(
f"FlashVSR 1.1 LQ projection requires complete 4-frame chunks, got {chunk.shape[2]}"
)
x = self.pixel_shuffle(chunk)
next_cache1 = x[:, :, -2:].clone()
x = self.act(self.norm1(self.conv1(x, cache1)))
cache1 = next_cache1
next_cache2 = x[:, :, -2:].clone()
if chunk_index == 0:
cache2 = next_cache2
continue
x = self.act(self.norm2(self.conv2(x, cache2)))
cache2 = next_cache2
output_chunks.append(x)
if not output_chunks:
raise ValueError("LQ video is too short to produce projection tokens")
x = torch.cat(output_chunks, dim=2)
tokens = x.permute(0, 2, 3, 4, 1).flatten(1, 3)
return [layer(tokens) for layer in self.linear_layers]
|