diff --git a/FlashVSR/LQ_proj_model.py b/FlashVSR/LQ_proj_model.py new file mode 100644 index 0000000..215fd87 --- /dev/null +++ b/FlashVSR/LQ_proj_model.py @@ -0,0 +1,157 @@ +from einops import rearrange + +import torch +import torch.nn as nn +import torch.nn.functional as F + +CACHE_T = 2 + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding, mode='replicate') + + return super().forward(x) + +class PixelShuffle3d(nn.Module): + def __init__(self, ff, hh, ww): + super().__init__() + self.ff = ff + self.hh = hh + self.ww = ww + + def forward(self, x): + # x: (B, C, F, H, W) + return rearrange(x, + 'b c (f ff) (h hh) (w ww) -> b (c ff hh ww) f h w', + ff=self.ff, hh=self.hh, ww=self.ww) + +class Buffer_LQ4x_Proj(nn.Module): + + def __init__(self, in_dim, out_dim, layer_num=30): + super().__init__() + self.ff = 1 + self.hh = 16 + self.ww = 16 + self.hidden_dim1 = 2048 + self.hidden_dim2 = 3072 + self.layer_num = layer_num + + self.pixel_shuffle = PixelShuffle3d(self.ff, self.hh, self.ww) + + self.conv1 = CausalConv3d(in_dim*self.ff*self.hh*self.ww, self.hidden_dim1, (4, 3, 3), stride=(2, 1, 1), padding=(1, 1, 1)) # f -> f/2 h -> h w -> w + self.norm1 = RMS_norm(self.hidden_dim1, images=False) + self.act1 = nn.SiLU() + + self.conv2 = CausalConv3d(self.hidden_dim1, self.hidden_dim2, (4, 3, 3), stride=(2, 1, 1), padding=(1, 1, 1)) # f -> f/2 h -> h w -> w + self.norm2 = RMS_norm(self.hidden_dim2, images=False) + self.act2 = nn.SiLU() + + self.linear_layers = nn.ModuleList([nn.Linear(self.hidden_dim2, out_dim) for _ in range(layer_num)]) + + self.clip_idx = 0 + + def forward(self, video): + self.clear_cache() + # x: (B, C, F, H, W) + + t = video.shape[2] + iter_ = 1 + (t - 1) // 4 + first_frame = video[:, :, :1, :, :].repeat(1, 1, 3, 1, 1) + video = torch.cat([first_frame, video], dim=2) + + out_x = [] + for i in range(iter_): + x = self.pixel_shuffle(video[:,:,i*4:(i+1)*4,:,:]) + cache1_x = x[:, :, -CACHE_T:, :, :].clone() + self.cache['conv1'] = cache1_x + x = self.conv1(x, self.cache['conv1']) + x = self.norm1(x) + x = self.act1(x) + cache2_x = x[:, :, -CACHE_T:, :, :].clone() + self.cache['conv2'] = cache2_x + if i == 0: + continue + x = self.conv2(x, self.cache['conv2']) + x = self.norm2(x) + x = self.act2(x) + out_x.append(x) + out_x = torch.cat(out_x, dim = 2) + + out_x = rearrange(out_x, 'b c f h w -> b (f h w) c') + outputs = [] + for i in range(self.layer_num): + outputs.append(self.linear_layers[i](out_x)) + self.clear_cache() + return outputs + + def clear_cache(self): + self.cache = {} + self.cache['conv1'] = None + self.cache['conv2'] = None + self.clip_idx = 0 + + def stream_forward(self, video_clip): + if self.clip_idx == 0: + # self.clear_cache() + first_frame = video_clip[:, :, :1, :, :].repeat(1, 1, 3, 1, 1) + video_clip = torch.cat([first_frame, video_clip], dim=2) + x = self.pixel_shuffle(video_clip) + cache1_x = x[:, :, -CACHE_T:, :, :].clone() + self.cache['conv1'] = cache1_x + x = self.conv1(x, self.cache['conv1']) + x = self.norm1(x) + x = self.act1(x) + cache2_x = x[:, :, -CACHE_T:, :, :].clone() + self.cache['conv2'] = cache2_x + self.clip_idx += 1 + return None + else: + x = self.pixel_shuffle(video_clip) + cache1_x = x[:, :, -CACHE_T:, :, :].clone() + self.cache['conv1'] = cache1_x + x = self.conv1(x, self.cache['conv1']) + x = self.norm1(x) + x = self.act1(x) + cache2_x = x[:, :, -CACHE_T:, :, :].clone() + self.cache['conv2'] = cache2_x + x = self.conv2(x, self.cache['conv2']) + x = self.norm2(x) + x = self.act2(x) + out_x = rearrange(x, 'b c f h w -> b (f h w) c') + outputs = [] + for i in range(self.layer_num): + outputs.append(self.linear_layers[i](out_x)) + self.clip_idx += 1 + return outputs diff --git a/FlashVSR/TCDecoder.py b/FlashVSR/TCDecoder.py new file mode 100644 index 0000000..3e50e01 --- /dev/null +++ b/FlashVSR/TCDecoder.py @@ -0,0 +1,261 @@ +""" +Tiny AutoEncoder for Hunyuan Video (Decoder-only, pruned) +- Encoder removed +- Transplant/widening helpers removed +- Deepening (IdentityConv2d+ReLU) is now built into the decoder structure itself +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm.auto import tqdm +from collections import namedtuple +from einops import rearrange +import torch.nn.init as init + +DecoderResult = namedtuple("DecoderResult", ("frame", "memory")) +TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) + +# ---------------------------- +# Utility / building blocks +# ---------------------------- + +class IdentityConv2d(nn.Conv2d): + """Same-shape Conv2d initialized to identity (Dirac).""" + def __init__(self, C, kernel_size=3, bias=False): + pad = kernel_size // 2 + super().__init__(C, C, kernel_size, padding=pad, bias=bias) + with torch.no_grad(): + init.dirac_(self.weight) + if self.bias is not None: + self.bias.zero_() + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + +class MemBlock(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential( + conv(n_in * 2, n_out), nn.ReLU(inplace=True), + conv(n_out, n_out), nn.ReLU(inplace=True), + conv(n_out, n_out) + ) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + def forward(self, x, past): + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + +class TPool(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f*stride, n_f, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + +class TGrow(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + x = self.conv(x) + return x.reshape(-1, C, H, W) + +class PixelShuffle3d(nn.Module): + def __init__(self, ff, hh, ww): + super().__init__() + self.ff = ff + self.hh = hh + self.ww = ww + def forward(self, x): + # x: (B, C, F, H, W) + B, C, F, H, W = x.shape + if F % self.ff != 0: + first_frame = x[:, :, 0:1, :, :].repeat(1, 1, self.ff - F % self.ff, 1, 1) + x = torch.cat([first_frame, x], dim=2) + return rearrange( + x, + 'b c (f ff) (h hh) (w ww) -> b (c ff hh ww) f h w', + ff=self.ff, hh=self.hh, ww=self.ww + ).transpose(1, 2) + +# ---------------------------- +# Generic NTCHW graph executor (kept; used by decoder) +# ---------------------------- + +def apply_model_with_memblocks(model, x, parallel, show_progress_bar, mem=None): + """ + Apply a sequential model with memblocks to the given input. + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - parallel: if True, parallelize over timesteps (fast but uses O(T) memory) + if False, each timestep will be processed sequentially (slow but uses O(1) memory) + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor" + N, T, C, H, W = x.shape + if parallel: + x = x.reshape(N*T, C, H, W) + for b in tqdm(model, disable=not show_progress_bar): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T = NT // N + _x = x.reshape(N, T, C, H, W) + mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape) + x = b(x, mem) + else: + x = b(x) + NT, C, H, W = x.shape + T = NT // N + x = x.view(N, T, C, H, W) + else: + out = [] + work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))] + progress_bar = tqdm(range(T), disable=not show_progress_bar) + while work_queue: + xt, i = work_queue.pop(0) + if i == 0: + progress_bar.update(1) + if i == len(model): + out.append(xt) + else: + b = model[i] + if isinstance(b, MemBlock): + if mem[i] is None: + xt_new = b(xt, xt * 0) + mem[i] = xt + else: + xt_new = b(xt, mem[i]) + mem[i].copy_(xt) + work_queue.insert(0, TWorkItem(xt_new, i+1)) + elif isinstance(b, TPool): + if mem[i] is None: + mem[i] = [] + mem[i].append(xt) + if len(mem[i]) > b.stride: + raise ValueError("TPool internal state invalid.") + elif len(mem[i]) == b.stride: + N_, C_, H_, W_ = xt.shape + xt = b(torch.cat(mem[i], 1).view(N_*b.stride, C_, H_, W_)) + mem[i] = [] + work_queue.insert(0, TWorkItem(xt, i+1)) + elif isinstance(b, TGrow): + xt = b(xt) + NT, C_, H_, W_ = xt.shape + for xt_next in reversed(xt.view(N, b.stride*C_, H_, W_).chunk(b.stride, 1)): + work_queue.insert(0, TWorkItem(xt_next, i+1)) + else: + xt = b(xt) + work_queue.insert(0, TWorkItem(xt, i+1)) + progress_bar.close() + x = torch.stack(out, 1) + return x, mem + +# ---------------------------- +# Decoder-only TAEHV +# ---------------------------- + +class TAEHV(nn.Module): + image_channels = 3 + def __init__( + self, + decoder_time_upscale=(True, True), + decoder_space_upscale=(True, True, True), + channels = [256, 128, 64, 64], + latent_channels = 16, + dtype=torch.float32 + ): + """Initialize TAEHV (decoder-only) with built-in deepening after every ReLU. + Deepening config: how_many_each=1, k=3 (fixed as requested). + """ + super().__init__() + self.dtype = dtype + self.latent_channels = latent_channels + n_f = channels + self.frames_to_trim = 2**sum(decoder_time_upscale) - 1 + + # Build the decoder "skeleton" + base_decoder = nn.Sequential( + Clamp(), conv(self.latent_channels, n_f[0]), nn.ReLU(inplace=True), + + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), + TGrow(n_f[0], 1), + conv(n_f[0], n_f[1], bias=False), + + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), + TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), + conv(n_f[1], n_f[2], bias=False), + + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), + TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), + conv(n_f[2], n_f[3], bias=False), + + nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), + ) + + # Inline deepening: insert (IdentityConv2d(k=3) + ReLU) after every ReLU + self.decoder = self._apply_identity_deepen(base_decoder, how_many_each=1, k=3) + + self.pixel_shuffle = PixelShuffle3d(4, 8, 8) + + # Initialize decoder mem state + self.clean_mem() + + @staticmethod + def _apply_identity_deepen(decoder: nn.Sequential, how_many_each=1, k=3) -> nn.Sequential: + """Return a new Sequential where every nn.ReLU is followed by how_many_each*(IdentityConv2d(k)+ReLU).""" + new_layers = [] + for b in decoder: + new_layers.append(b) + if isinstance(b, nn.ReLU): + # Deduce channel count from preceding layer + C = None + if len(new_layers) >= 2 and isinstance(new_layers[-2], nn.Conv2d): + C = new_layers[-2].out_channels + elif len(new_layers) >= 2 and isinstance(new_layers[-2], MemBlock): + C = new_layers[-2].conv[-1].out_channels + if C is not None: + for _ in range(how_many_each): + new_layers.append(IdentityConv2d(C, kernel_size=k, bias=False)) + new_layers.append(nn.ReLU(inplace=True)) + return nn.Sequential(*new_layers) + + def decode_video(self, x, parallel=False, show_progress_bar=False, cond=None): + """Decode a sequence of frames from latents. + x: NTCHW latent tensor; returns NTCHW RGB in ~[0, 1]. + """ + trim_flag = self.mem[-8] is None # keeps original relative check + + if cond is not None: + shuffled = self.pixel_shuffle(cond) + x = torch.cat([shuffled[:, :x.shape[1]], x], dim=2) + + x, self.mem = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar, mem=self.mem) + self.clean_mem() + + if trim_flag: + return x[:, self.frames_to_trim:] + + return x + + def clean_mem(self): + self.mem = [None] * len(self.decoder) + + +def build_tcdecoder(new_channels = [512, 256, 128, 128], device="cuda", dtype=torch.bfloat16, new_latent_channels=None): + big = TAEHV(channels=new_channels, latent_channels=new_latent_channels, dtype=dtype).to(device).to(dtype) + return big diff --git a/FlashVSR/flashvsr_nodes.py b/FlashVSR/flashvsr_nodes.py new file mode 100644 index 0000000..81f289d --- /dev/null +++ b/FlashVSR/flashvsr_nodes.py @@ -0,0 +1,69 @@ +import folder_paths +import torch + +from comfy.utils import load_torch_file +import comfy.model_management as mm +device = mm.get_torch_device() +offload_device = mm.unet_offload_device() + +class WanVideoAddFlashVSRInput: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "embeds": ("WANVIDIMAGE_EMBEDS",), + "images": ("IMAGE", {"tooltip": "Low-res video frames to enhance"}), + } + } + + RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",) + RETURN_NAMES = ("image_embeds",) + FUNCTION = "add" + CATEGORY = "WanVideoWrapper" + + def add(self, embeds, images): + updated = dict(embeds) + updated["flashvsr_LQ_images"] = images + return (updated,) + + +class WanVideoFlashVSRDecoderLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}), + }, + "optional": { + "precision": (["fp16", "fp32", "bf16"], + {"default": "bf16"} + ), + } + } + + RETURN_TYPES = ("WANVAE",) + RETURN_NAMES = ("vae", ) + FUNCTION = "loadmodel" + CATEGORY = "WanVideoWrapper" + DESCRIPTION = "Loads Wan VAE model from 'ComfyUI/models/vae'" + + def loadmodel(self, model_name, precision): + from .TCDecoder import build_tcdecoder + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] + model_path = folder_paths.get_full_path("vae", model_name) + sd = load_torch_file(model_path, safe_load=True) + + TCDecoder = build_tcdecoder(new_channels=[512, 256, 128, 128], new_latent_channels=16+768, dtype=dtype) + TCDecoder.load_state_dict(sd, strict=True) + TCDecoder.to(dtype) + + return (TCDecoder,) + +NODE_CLASS_MAPPINGS = { + "WanVideoAddFlashVSRInput": WanVideoAddFlashVSRInput, + "WanVideoFlashVSRDecoderLoader": WanVideoFlashVSRDecoderLoader, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WanVideoAddFlashVSRInput": "WanVideo Add FlashVSR Input", + "WanVideoFlashVSRDecoderLoader": "WanVideo FlashVSR Decoder Loader", +} diff --git a/__init__.py b/__init__.py index 94012c1..d12d003 100644 --- a/__init__.py +++ b/__init__.py @@ -24,6 +24,7 @@ from .nodes_utility import NODE_CLASS_MAPPINGS as UTILITY_NODE_CLASS_MAPPINGS, N from .cache_methods.nodes_cache import NODE_CLASS_MAPPINGS as NODE_CACHE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as NODE_CACHE_DISPLAY_NAME_MAPPINGS from .nodes_deprecated import NODE_CLASS_MAPPINGS as DEPRECATED_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as DEPRECATED_NODE_DISPLAY_NAME_MAPPINGS from .s2v.nodes import NODE_CLASS_MAPPINGS as S2V_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as S2V_NODE_DISPLAY_NAME_MAPPINGS +from .FlashVSR.flashvsr_nodes import NODE_CLASS_MAPPINGS as FLASHVSR_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FLASHVSR_NODE_DISPLAY_NAME_MAPPINGS try: from .qwen.qwen import NODE_CLASS_MAPPINGS as QWEN_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as QWEN_NODE_DISPLAY_NAME_MAPPINGS @@ -96,6 +97,7 @@ NODE_CLASS_MAPPINGS.update(HUMO_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(SAMPLER_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(LYNX_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(OVI_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(FLASHVSR_NODE_CLASS_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(RECAM_MASTER_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(UNIANIMATE_NODE_DISPLAY_NAME_MAPPINGS) @@ -118,5 +120,6 @@ NODE_DISPLAY_NAME_MAPPINGS.update(HUMO_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(SAMPLER_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(LYNX_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(OVI_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(FLASHVSR_NODE_DISPLAY_NAME_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/example_workflows/wanvideo_1_3B_FlashVSR_upscale_example.json b/example_workflows/wanvideo_1_3B_FlashVSR_upscale_example.json new file mode 100644 index 0000000..7f21f46 --- /dev/null +++ b/example_workflows/wanvideo_1_3B_FlashVSR_upscale_example.json @@ -0,0 +1,1169 @@ +{ + "id": "206247b6-9fec-4ed2-8927-e4f388c674d4", + "revision": 0, + "last_node_id": 127, + "last_link_id": 215, + "nodes": [ + { + "id": 35, + "type": "WanVideoTorchCompileSettings", + "pos": [ + -664.2244873046875, + -1274.048828125 + ], + "size": [ + 390.5999755859375, + 202 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "torch_compile_args", + "type": "WANCOMPILEARGS", + "slot_index": 0, + "links": [] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50", + "Node name for S&R": "WanVideoTorchCompileSettings" + }, + "widgets_values": [ + "inductor", + false, + "default", + false, + 64, + true, + 128 + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 104, + "type": "WanVideoTextEncode", + "pos": [ + 22.614208221435547, + -727.9273071289062 + ], + "size": [ + 525.8681640625, + 384.3753356933594 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "t5", + "shape": 7, + "type": "WANTEXTENCODER", + "link": 162 + }, + { + "name": "model_to_offload", + "shape": 7, + "type": "WANVIDEOMODEL", + "link": 163 + } + ], + "outputs": [ + { + "name": "text_embeds", + "type": "WANVIDEOTEXTEMBEDS", + "slot_index": 0, + "links": [ + 166 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50", + "Node name for S&R": "WanVideoTextEncode" + }, + "widgets_values": [ + "people talking", + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards\"", + true, + true, + "gpu" + ], + "color": "#332922", + "bgcolor": "#593930" + }, + { + "id": 11, + "type": "LoadWanVideoT5TextEncoder", + "pos": [ + -443.7521667480469, + -535.9060668945312 + ], + "size": [ + 377.1661376953125, + 130 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "wan_t5_model", + "type": "WANTEXTENCODER", + "slot_index": 0, + "links": [ + 162 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50", + "Node name for S&R": "LoadWanVideoT5TextEncoder" + }, + "widgets_values": [ + "umt5-xxl-enc-bf16.safetensors", + "bf16", + "offload_device", + "disabled" + ], + "color": "#332922", + "bgcolor": "#593930" + }, + { + "id": 123, + "type": "VHS_LoadVideo", + "pos": [ + -179.3857421875, + -165.7681427001953 + ], + "size": [ + 247.455078125, + 551.455078125 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 206 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": [] + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "ver": "0edce8ef7ce173ac97a3ed3d6f4636029d1a4530", + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "example0.mp4", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "AnimateDiff", + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "example0.mp4", + "type": "input", + "format": "video/mp4", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1 + } + } + } + }, + { + "id": 116, + "type": "ImageResizeKJv2", + "pos": [ + 131.98231506347656, + -163.7329559326172 + ], + "size": [ + 270, + 336 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 206 + }, + { + "name": "mask", + "shape": 7, + "type": "MASK", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 186, + 210 + ] + }, + { + "name": "width", + "type": "INT", + "links": [ + 189 + ] + }, + { + "name": "height", + "type": "INT", + "links": [ + 190 + ] + }, + { + "name": "mask", + "type": "MASK", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-kjnodes", + "ver": "3fcd22f2fe2be69c3229f192362b91888277cbcb", + "Node name for S&R": "ImageResizeKJv2" + }, + "widgets_values": [ + 1024, + 1024, + "lanczos", + "crop", + "0, 0, 0", + "center", + 2, + "cpu", + "Output: 85 x 1024 x 1024 | 1020.00MB" + ] + }, + { + "id": 78, + "type": "WanVideoEmptyEmbeds", + "pos": [ + 518.80322265625, + -107.40808868408203 + ], + "size": [ + 272.431640625, + 126 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "control_embeds", + "shape": 7, + "type": "WANVIDIMAGE_EMBEDS", + "link": null + }, + { + "name": "extra_latents", + "shape": 7, + "type": "LATENT", + "link": null + }, + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": 189 + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": 190 + } + ], + "outputs": [ + { + "name": "image_embeds", + "type": "WANVIDIMAGE_EMBEDS", + "links": [ + 181 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "6bc53b771d5d2af316801cb69e2ee10dbf7d18b1", + "Node name for S&R": "WanVideoEmptyEmbeds" + }, + "widgets_values": [ + 720, + 720, + 85 + ] + }, + { + "id": 114, + "type": "WanVideoAddFlashVSRInput", + "pos": [ + 837.4269409179688, + -203.31698608398438 + ], + "size": [ + 234.45761108398438, + 46 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "embeds", + "type": "WANVIDIMAGE_EMBEDS", + "link": 181 + }, + { + "name": "images", + "type": "IMAGE", + "link": 186 + } + ], + "outputs": [ + { + "name": "image_embeds", + "type": "WANVIDIMAGE_EMBEDS", + "links": [ + 182 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "3d42bf62cedf4f760edadba9e9c55869105d0bba", + "Node name for S&R": "WanVideoAddFlashVSRInput" + }, + "widgets_values": [] + }, + { + "id": 117, + "type": "ImageConcatMulti", + "pos": [ + 1645.5211181640625, + -942.0350952148438 + ], + "size": [ + 270, + 150 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "image_1", + "type": "IMAGE", + "link": 210 + }, + { + "name": "image_2", + "shape": 7, + "type": "IMAGE", + "link": 215 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 193 + ] + } + ], + "properties": { + "cnr_id": "comfyui-kjnodes", + "ver": "3fcd22f2fe2be69c3229f192362b91888277cbcb" + }, + "widgets_values": [ + 2, + "right", + true, + null + ] + }, + { + "id": 121, + "type": "WanVideoDecode", + "pos": [ + 1233.328125, + -938.7442626953125 + ], + "size": [ + 315, + 198 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "WANVAE", + "link": 213 + }, + { + "name": "samples", + "type": "LATENT", + "link": 196 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 214 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50", + "Node name for S&R": "WanVideoDecode" + }, + "widgets_values": [ + false, + 272, + 272, + 144, + 128, + "default" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 126, + "type": "GetImageSizeAndCount", + "pos": [ + 1618.2845458984375, + -665.3915405273438 + ], + "size": [ + 190.86483764648438, + 86 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 214 + } + ], + "outputs": [ + { + "name": "image", + "type": "IMAGE", + "links": [ + 215 + ] + }, + { + "label": "1024 width", + "name": "width", + "type": "INT", + "links": null + }, + { + "label": "1024 height", + "name": "height", + "type": "INT", + "links": null + }, + { + "label": "81 count", + "name": "count", + "type": "INT", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-kjnodes", + "ver": "3fcd22f2fe2be69c3229f192362b91888277cbcb", + "Node name for S&R": "GetImageSizeAndCount" + } + }, + { + "id": 27, + "type": "WanVideoSampler", + "pos": [ + 1221.549560546875, + -644.4454956054688 + ], + "size": [ + 315, + 999 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "WANVIDEOMODEL", + "link": 144 + }, + { + "name": "image_embeds", + "type": "WANVIDIMAGE_EMBEDS", + "link": 182 + }, + { + "name": "text_embeds", + "shape": 7, + "type": "WANVIDEOTEXTEMBEDS", + "link": 166 + }, + { + "name": "samples", + "shape": 7, + "type": "LATENT", + "link": null + }, + { + "name": "feta_args", + "shape": 7, + "type": "FETAARGS", + "link": null + }, + { + "name": "context_options", + "shape": 7, + "type": "WANVIDCONTEXT", + "link": null + }, + { + "name": "cache_args", + "shape": 7, + "type": "CACHEARGS", + "link": null + }, + { + "name": "flowedit_args", + "shape": 7, + "type": "FLOWEDITARGS", + "link": null + }, + { + "name": "slg_args", + "shape": 7, + "type": "SLGARGS", + "link": null + }, + { + "name": "loop_args", + "shape": 7, + "type": "LOOPARGS", + "link": null + }, + { + "name": "experimental_args", + "shape": 7, + "type": "EXPERIMENTALARGS", + "link": null + }, + { + "name": "sigmas", + "shape": 7, + "type": "SIGMAS", + "link": null + }, + { + "name": "unianimate_poses", + "shape": 7, + "type": "UNIANIMATE_POSE", + "link": null + }, + { + "name": "fantasytalking_embeds", + "shape": 7, + "type": "FANTASYTALKING_EMBEDS", + "link": null + }, + { + "name": "uni3c_embeds", + "shape": 7, + "type": "UNI3C_EMBEDS", + "link": null + }, + { + "name": "multitalk_embeds", + "shape": 7, + "type": "MULTITALK_EMBEDS", + "link": null + }, + { + "name": "freeinit_args", + "shape": 7, + "type": "FREEINITARGS", + "link": null + } + ], + "outputs": [ + { + "name": "samples", + "type": "LATENT", + "slot_index": 0, + "links": [ + 196 + ] + }, + { + "name": "denoised_samples", + "type": "LATENT", + "links": null + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50", + "Node name for S&R": "WanVideoSampler" + }, + "widgets_values": [ + 1, + 1, + 5, + 74, + "fixed", + true, + "euler", + 0, + 1, + false, + "comfy", + 0, + -1, + false + ] + }, + { + "id": 22, + "type": "WanVideoModelLoader", + "pos": [ + -72.96549987792969, + -1239.5343017578125 + ], + "size": [ + 624.1026611328125, + 338 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "compile_args", + "shape": 7, + "type": "WANCOMPILEARGS", + "link": null + }, + { + "name": "block_swap_args", + "shape": 7, + "type": "BLOCKSWAPARGS", + "link": null + }, + { + "name": "lora", + "shape": 7, + "type": "WANVIDLORA", + "link": null + }, + { + "name": "vram_management_args", + "shape": 7, + "type": "VRAM_MANAGEMENTARGS", + "link": null + }, + { + "name": "extra_model", + "shape": 7, + "type": "VACEPATH", + "link": 180 + }, + { + "name": "fantasytalking_model", + "shape": 7, + "type": "FANTASYTALKINGMODEL", + "link": null + }, + { + "name": "multitalk_model", + "shape": 7, + "type": "MULTITALKMODEL", + "link": null + }, + { + "name": "fantasyportrait_model", + "shape": 7, + "type": "FANTASYPORTRAITMODEL", + "link": null + }, + { + "name": "vace_model", + "shape": 7, + "type": "VACEPATH", + "link": null + } + ], + "outputs": [ + { + "name": "model", + "type": "WANVIDEOMODEL", + "slot_index": 0, + "links": [ + 144, + 163 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50", + "Node name for S&R": "WanVideoModelLoader" + }, + "widgets_values": [ + "WanVideo\\FlashVSR\\Wan2_1-T2V-1_3B_FlashVSR_fp32.safetensors", + "fp16", + "disabled", + "offload_device", + "sdpa", + "default" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 119, + "type": "WanVideoFlashVSRDecoderLoader", + "pos": [ + -670.894775390625, + -824.1834716796875 + ], + "size": [ + 360.68438720703125, + 82 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "vae", + "type": "WANVAE", + "links": [ + 213 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "3d42bf62cedf4f760edadba9e9c55869105d0bba", + "Node name for S&R": "WanVideoFlashVSRDecoderLoader" + }, + "widgets_values": [ + "Wan2_1_FlashVSR_TCDecoder_fp32.safetensors", + "fp16" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 113, + "type": "WanVideoExtraModelSelect", + "pos": [ + -677.6390380859375, + -939.74267578125 + ], + "size": [ + 507.62115478515625, + 58 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "prev_model", + "shape": 7, + "type": "VACEPATH", + "link": null + } + ], + "outputs": [ + { + "name": "extra_model", + "type": "VACEPATH", + "links": [ + 180 + ] + } + ], + "properties": { + "cnr_id": "ComfyUI-WanVideoWrapper", + "ver": "3d42bf62cedf4f760edadba9e9c55869105d0bba", + "Node name for S&R": "WanVideoExtraModelSelect" + }, + "widgets_values": [ + "WanVideo\\FlashVSR\\Wan2_1_FlashVSR_PQ_proj_model_bf16.safetensors" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 84, + "type": "MarkdownNote", + "pos": [ + -1130.9693603515625, + -908.9747314453125 + ], + "size": [ + 408.80950927734375, + 131.6868133544922 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Model links:\n\n[https://huggingface.co/Kijai/WanVideo_comfy/tree/main/FlashVSR](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/FlashVSR)\n\nText encoder:\n\n[https://huggingface.co/Kijai/WanVideo_comfy/blob/main/umt5-xxl-enc-bf16.safetensors](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/umt5-xxl-enc-bf16.safetensors)\n" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 127, + "type": "Note", + "pos": [ + 630.6744384765625, + -1230.7738037109375 + ], + "size": [ + 316.588623046875, + 203.8870391845703 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: WIP", + "properties": {}, + "widgets_values": [ + "This is a barebones implementation of FlashVSR:\n\nhttps://huggingface.co/JunhaoZhuang/FlashVSR\n\nNone of the sparse attention and streaming stuff is implemented, however it still gives results like this and may be useful" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 30, + "type": "VHS_VideoCombine", + "pos": [ + 1978.0262451171875, + -921.9905395507812 + ], + "size": [ + 1369.7001953125, + 1022.85009765625 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 193 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "ver": "0a75c7958fe320efcb052f1d9f8451fd20c730a8", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WanVideoWrapper_FlashVSR", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "WanVideoWrapper_EchoShot_00020.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 16, + "workflow": "WanVideoWrapper_EchoShot_00020.png", + "fullpath": "N:\\AI\\ComfyUI\\temp\\WanVideoWrapper_EchoShot_00020.mp4" + } + } + } + } + ], + "links": [ + [ + 144, + 22, + 0, + 27, + 0, + "WANVIDEOMODEL" + ], + [ + 162, + 11, + 0, + 104, + 0, + "WANTEXTENCODER" + ], + [ + 163, + 22, + 0, + 104, + 1, + "WANVIDEOMODEL" + ], + [ + 166, + 104, + 0, + 27, + 2, + "WANVIDEOTEXTEMBEDS" + ], + [ + 180, + 113, + 0, + 22, + 4, + "VACEPATH" + ], + [ + 181, + 78, + 0, + 114, + 0, + "WANVIDIMAGE_EMBEDS" + ], + [ + 182, + 114, + 0, + 27, + 1, + "WANVIDIMAGE_EMBEDS" + ], + [ + 186, + 116, + 0, + 114, + 1, + "IMAGE" + ], + [ + 189, + 116, + 1, + 78, + 2, + "INT" + ], + [ + 190, + 116, + 2, + 78, + 3, + "INT" + ], + [ + 193, + 117, + 0, + 30, + 0, + "IMAGE" + ], + [ + 196, + 27, + 0, + 121, + 1, + "LATENT" + ], + [ + 206, + 123, + 0, + 116, + 0, + "IMAGE" + ], + [ + 210, + 116, + 0, + 117, + 0, + "IMAGE" + ], + [ + 213, + 119, + 0, + 121, + 0, + "WANVAE" + ], + [ + 214, + 121, + 0, + 126, + 0, + "IMAGE" + ], + [ + 215, + 126, + 0, + 117, + 1, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.611590904484194, + "offset": [ + 1778.97094924203, + 1358.0539389275748 + ] + }, + "frontendVersion": "1.28.1", + "node_versions": { + "ComfyUI-WanVideoWrapper": "5a2383621a05825d0d0437781afcb8552d9590fd", + "comfy-core": "0.3.26", + "ComfyUI-KJNodes": "a5bd3c86c8ed6b83c55c2d0e7a59515b15a0137f", + "ComfyUI-VideoHelperSuite": "0a75c7958fe320efcb052f1d9f8451fd20c730a8" + }, + "VHS_latentpreview": true, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py index b608f1f..acb70e2 100644 --- a/nodes.py +++ b/nodes.py @@ -1997,6 +1997,8 @@ class WanVideoDecode: drop_last = samples.get("drop_last", False) is_looped = samples.get("looped", False) + flashvsr_LQ_images = samples.get("flashvsr_LQ_images", None) + vae.to(device) latents = latents.to(device = device, dtype = vae.dtype) @@ -2009,7 +2011,7 @@ class WanVideoDecode: latents = latents[:, :, :-1] if type(vae).__name__ == "TAEHV": - images = vae.decode_video(latents.permute(0, 2, 1, 3, 4))[0].permute(1, 0, 2, 3) + images = vae.decode_video(latents.permute(0, 2, 1, 3, 4), cond=flashvsr_LQ_images.to(vae.dtype))[0].permute(1, 0, 2, 3) images = torch.clamp(images, 0.0, 1.0) images = images.permute(1, 2, 3, 0).cpu().float() return (images,) diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 1b0e4de..8683325 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -1424,6 +1424,12 @@ class WanVideoModelLoader: sd.update(extra_sd) del extra_sd + # FlashVSR + if "LQ_proj_in.norm1.gamma" in sd: + log.info("FlashVSR model detected, patching model...") + from .FlashVSR.LQ_proj_model import Buffer_LQ4x_Proj + transformer.LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1) + # Additional cond latents if "add_conv_in.weight" in sd: def zero_module(module): diff --git a/nodes_sampler.py b/nodes_sampler.py index c0651aa..d41a27a 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -829,8 +829,17 @@ class WanVideoSampler: if extra_channel_latents is not None: extra_channel_latents = extra_channel_latents[0].to(noise) + # FlashVSR + flashvsr_LQ_latent = None + flashvsr_LQ_images = image_embeds.get("flashvsr_LQ_images", None) + if flashvsr_LQ_images is not None: + LQ_images = flashvsr_LQ_images.unsqueeze(0).movedim(-1, 1).to(device, dtype) * 2 - 1 + flashvsr_LQ_latent = transformer.LQ_proj_in(LQ_images) + log.info(f"flashvsr_LQ_latent: {flashvsr_LQ_latent[0].shape}") + noise = noise[:, :-1] + seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1]) + latent = noise - print("Latent shape:", latent.shape, "Latent dtype:", latent.dtype, "Latent device:", latent.device) #controlnet controlnet_latents = controlnet = None @@ -1335,6 +1344,7 @@ class WanVideoSampler: "x_ovi": [latent_model_input_ovi.to(z)] if latent_model_input_ovi is not None else None, # Audio latent model input for Ovi "seq_len_ovi": seq_len_ovi, # Audio latent model sequence length for Ovi "ovi_negative_text_embeds": ovi_negative_text_embeds, # Audio latent model negative text embeds for Ovi + "flashvsr_LQ_latent": flashvsr_LQ_latent, # FlashVSR LQ latent for upsampling } batch_size = 1 @@ -2603,7 +2613,10 @@ class WanVideoSampler: self.cache_state = [None, None] - if ref_latent is not None: + noise = torch.randn(16, latent_window_size + 1, lat_h, lat_w, dtype=torch.float32, device=torch.device("cpu"), generator=seed_g).to(device) + seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1]) + + if current_ref_images is not None: if offload: offload_transformer(transformer) offloaded = True @@ -2631,12 +2644,14 @@ class WanVideoSampler: else: temporal_ref_latents = temporal_ref_latents[:, :msk.shape[1]] - temporal_ref_latents = torch.cat([msk, temporal_ref_latents], dim=0) # 4+C T H W - image_cond_in = torch.cat([ref_latent.to(device), temporal_ref_latents], dim=1) # 4+C T+trefs H W - del temporal_ref_latents, msk, bg_image_slice - - noise = torch.randn(16, latent_window_size + 1, lat_h, lat_w, dtype=torch.float32, device=torch.device("cpu"), generator=seed_g).to(device) - seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1]) + if ref_latent is not None: + temporal_ref_latents = torch.cat([msk, temporal_ref_latents], dim=0) # 4+C T H W + image_cond_in = torch.cat([ref_latent.to(device), temporal_ref_latents], dim=1) # 4+C T+trefs H W + del temporal_ref_latents, msk, bg_image_slice + else: + image_cond_in = torch.cat([torch.tile(torch.zeros_like(noise[:1]), [4, 1, 1, 1]), torch.zeros_like(noise)], dim=0).to(device) + else: + image_cond_in = torch.cat([torch.tile(torch.zeros_like(noise[:1]), [4, 1, 1, 1]), torch.zeros_like(noise)], dim=0).to(device) pose_input_slice = None if pose_images is not None: @@ -2978,6 +2993,7 @@ class WanVideoSampler: "original_image": original_image.cpu() if original_image is not None else None, "cache_states": cache_states, "latent_ovi_audio": latent_ovi.unsqueeze(0).transpose(1, 2).cpu() if latent_ovi is not None else None, + "flashvsr_LQ_images": LQ_images, },{ "samples": callback_latent.unsqueeze(0).cpu() if callback is not None else None, }) diff --git a/taehv/taehv.py b/taehv/taehv.py index 7e2e93d..c0a8fc8 100644 --- a/taehv/taehv.py +++ b/taehv/taehv.py @@ -211,7 +211,7 @@ class TAEHV(nn.Module): if self.patch_size > 1: x = F.pixel_unshuffle(x, self.patch_size) return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) - def decode_video(self, x, parallel=False, show_progress_bar=True): + def decode_video(self, x, parallel=False, show_progress_bar=True, **kwargs): """Decode a sequence of frames. Args: diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 8b8a95f..22f2439 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -2153,6 +2153,7 @@ class WanModel(torch.nn.Module): wananim_pose_strength=1.0, wananim_face_strength=1.0, lynx_embeds=None, x_ovi=None, seq_len_ovi=None, ovi_negative_text_embeds=None, + flashvsr_LQ_latent=None, ): r""" Forward pass through the diffusion model @@ -2807,6 +2808,9 @@ class WanModel(torch.nn.Module): lynx_ref_feature = lynx_ref_buffer.get(block_idx, None) else: lynx_ref_feature = None + # FlashVSR + if flashvsr_LQ_latent is not None and b < len(flashvsr_LQ_latent): + x += flashvsr_LQ_latent[b].to(x) # Prefetch blocks if enabled if self.prefetch_blocks > 0: for prefetch_offset in range(1, self.prefetch_blocks + 1):