diff --git a/nodes.py b/nodes.py index 7864a9e..a648725 100644 --- a/nodes.py +++ b/nodes.py @@ -1263,7 +1263,6 @@ class WanVideoSampler: return { "required": { "model": ("WANVIDEOMODEL",), - "image_embeds": ("WANVIDIMAGE_EMBEDS", ), "steps": ("INT", {"default": 30, "min": 1}), "cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}), @@ -1292,6 +1291,8 @@ class WanVideoSampler: "uni3c_embeds": ("UNI3C_EMBEDS", ), "multitalk_embeds": ("MULTITALK_EMBEDS", ), "freeinit_args": ("FREEINITARGS", ), + "start_step": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1, "tooltip": "Start step for the sampling, 0 means full sampling, otherwise samples only from this step"}), + "end_step": ("INT", {"default": -1, "min": -1, "max": 10000, "step": 1, "tooltip": "End step for the sampling, -1 means full sampling, otherwise samples only until this step"}), } } @@ -1303,16 +1304,20 @@ class WanVideoSampler: def process(self, model, image_embeds, shift, steps, cfg, seed, scheduler, riflex_freq_index, text_embeds=None, force_offload=True, samples=None, feta_args=None, denoise_strength=1.0, context_options=None, cache_args=None, teacache_args=None, flowedit_args=None, batched_cfg=False, slg_args=None, rope_function="default", loop_args=None, - experimental_args=None, sigmas=None, unianimate_poses=None, fantasytalking_embeds=None, uni3c_embeds=None, multitalk_embeds=None, freeinit_args=None): + experimental_args=None, sigmas=None, unianimate_poses=None, fantasytalking_embeds=None, uni3c_embeds=None, multitalk_embeds=None, freeinit_args=None, start_step=0, end_step=-1): patcher = model model = model.model transformer = model.diffusion_model + dtype = model["dtype"] gguf = model["gguf"] control_lora = model["control_lora"] transformer_options = patcher.model_options.get("transformer_options", None) + is_5b = transformer.out_dim == 48 + vae_upscale_factor = 16 if is_5b else 8 + if len(patcher.patches) != 0 and transformer_options.get("linear_with_lora", False) is True: log.info(f"Using {len(patcher.patches)} LoRA weight patches for WanVideo model") if not gguf: @@ -1356,6 +1361,17 @@ class WanVideoSampler: sample_scheduler, timesteps = get_scheduler(scheduler, steps, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas) else: timesteps = torch.tensor([1000, 750, 500, 250], device=device) + log.info(f"sigmas: {sample_scheduler.sigmas}") + + if end_step != -1: + timesteps = timesteps[:end_step] + sample_scheduler.sigmas = sample_scheduler.sigmas[:end_step+1] + else: + timesteps = timesteps[start_step:] + sample_scheduler.sigmas = sample_scheduler.sigmas[start_step:] + + if hasattr(sample_scheduler, 'timesteps'): + sample_scheduler.timesteps = timesteps scheduler_step_args = {"generator": seed_g} step_sig = inspect.signature(sample_scheduler.step) @@ -1370,14 +1386,14 @@ class WanVideoSampler: control_latents = control_camera_latents = clip_fea = clip_fea_neg = end_image = recammaster = camera_embed = unianim_data = None vace_data = vace_context = vace_scale = None fun_or_fl2v_model = has_ref = drop_last = False - phantom_latents = None - fun_ref_image = None - - image_cond = image_embeds.get("image_embeds", None) - ATI_tracks = None + phantom_latents = fun_ref_image = ATI_tracks = None add_cond = attn_cond = attn_cond_neg = None - + + #I2V + image_cond = image_embeds.get("image_embeds", None) if image_cond is not None: + if transformer.in_dim == 16: + raise ValueError("T2V (text to video) model detected, encoded images only work with I2V (Image to video) models") log.info(f"image_cond shape: {image_cond.shape}") #ATI tracks if transformer_options is not None: @@ -1400,21 +1416,18 @@ class WanVideoSampler: add_cond_end_percent = add_cond_latents["pose_cond_end_percent"] end_image = image_embeds.get("end_image", None) - lat_h = image_embeds.get("lat_h", None) - lat_w = image_embeds.get("lat_w", None) - if lat_h is None or lat_w is None: - raise ValueError("Clip encoded image embeds must be provided for I2V (Image to Video) model") fun_or_fl2v_model = image_embeds.get("fun_or_fl2v_model", False) - noise = torch.randn( - 16, + + noise = torch.randn( #C, T, H, W + 48 if is_5b else 16, (image_embeds["num_frames"] - 1) // 4 + (2 if end_image is not None and not fun_or_fl2v_model else 1), - lat_h, - lat_w, + image_embeds["lat_h"], + image_embeds["lat_w"], dtype=torch.float32, generator=seed_g, device=torch.device("cpu")) seq_len = image_embeds["max_seq_len"] - + clip_fea = image_embeds.get("clip_context", None) if clip_fea is not None: clip_fea = clip_fea.to(dtype) @@ -1474,10 +1487,10 @@ class WanVideoSampler: }) noise = torch.randn( - target_shape[0], + 48 if is_5b else 16, target_shape[1] + 1 if has_ref else target_shape[1], - target_shape[2], - target_shape[3], + target_shape[2] // 2 if is_5b else target_shape[2], #todo make this smarter + target_shape[3] // 2 if is_5b else target_shape[3], #todo make this smarter dtype=torch.float32, device=torch.device("cpu"), generator=seed_g) @@ -1638,9 +1651,6 @@ class WanVideoSampler: context_frames = (context_options["context_frames"] - 1) // 4 + 1 context_stride = context_options["context_stride"] // 4 context_overlap = context_options["context_overlap"] // 4 - # context_vae = context_options.get("vae", None) - # if context_vae is not None: - # context_vae.to(device) context_reference_latent = context_options.get("reference_latent", None) # Get total number of prompts @@ -1680,20 +1690,26 @@ class WanVideoSampler: # vid2vid if samples is not None: + saved_generator_state = samples.get("generator_state", None) + if saved_generator_state is not None: + seed_g.set_state(saved_generator_state) input_samples = samples["samples"].squeeze(0).to(noise) if input_samples.shape[1] != noise.shape[1]: - input_samples = torch.cat([input_samples[:, :1].repeat(1, noise.shape[1] - input_samples.shape[1], 1, 1), input_samples], dim=1) - original_image = input_samples.to(device) + input_samples = torch.cat([input_samples[:, :1].repeat(1, noise.shape[1] - input_samples.shape[1], 1, 1), input_samples], dim=1) + if denoise_strength < 1.0: - latent_timestep = timesteps[:1].to(noise) - noise = noise * latent_timestep / 1000 + (1 - latent_timestep / 1000) * input_samples - + latent_timestep = timesteps[:1].to(noise) + noise = noise * latent_timestep / 1000 + (1 - latent_timestep / 1000) * input_samples + else: + noise = input_samples mask = samples.get("mask", None) if mask is not None: + original_image = input_samples.to(device) if mask.shape[2] != noise.shape[1]: mask = torch.cat([torch.zeros(1, noise.shape[0], noise.shape[1] - mask.shape[2], noise.shape[2], noise.shape[3]), mask], dim=2) - # extra latents (Pusa) + # extra latents (Pusa) and 5b + encoded_image_latents = None if (extra_latents := image_embeds.get("extra_latents", None)) is not None: encoded_image_latents = extra_latents["samples"].squeeze(0).to(noise) if (empty_latent_indices := extra_latents.get("empty_latent_indices", None)) is not None and len(empty_latent_indices) > 0: @@ -1903,6 +1919,7 @@ class WanVideoSampler: def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, idx, image_cond=None, clip_fea=None, control_latents=None, vace_data=None, unianim_data=None, audio_proj=None, control_camera_latents=None, add_cond=None, cache_state=None, context_window=None, multitalk_audio_embeds=None): + nonlocal transformer z = z.to(dtype) with torch.autocast(device_type=mm.get_autocast_device(device), dtype=dtype, enabled=("fp8" in model["quantization"])): @@ -2200,7 +2217,7 @@ class WanVideoSampler: from .latent_preview import prepare_callback #custom for tiny VAE previews callback = prepare_callback(patcher, steps) - log.info(f"Sampling {(latent_video_length-1) * 4 + 1} frames at {latent.shape[3]*8}x{latent.shape[2]*8} with {steps} steps") + log.info(f"Sampling {(latent_video_length-1) * 4 + 1} frames at {latent.shape[3]*vae_upscale_factor}x{latent.shape[2]*vae_upscale_factor} with {steps} steps") intermediate_device = device @@ -2294,7 +2311,7 @@ class WanVideoSampler: latent_model_input = latent.to(device) timestep = torch.tensor([t]).to(device) - if scheduler == "flowmatch_pusa": + if scheduler == "flowmatch_pusa" or (is_5b and encoded_image_latents is not None): timestep = timestep.unsqueeze(1).repeat(1, latent_video_length) if extra_latents is not None: if empty_latent_indices is not None and len(empty_latent_indices) > 0: @@ -2302,7 +2319,7 @@ class WanVideoSampler: non_noise_indices = [i for i in range(timestep.shape[1]) if i not in empty_latent_indices] timestep[:, non_noise_indices] = 0 else: - timestep[:,0:encoded_image_latents.shape[1]] = 0 + timestep[:,0:encoded_image_latents.shape[1]] = 0 #print(f"timestep: {timestep}") current_step_percentage = idx / len(timesteps) @@ -2821,7 +2838,7 @@ class WanVideoSampler: latent = latent.to(intermediate_device) latent = sample_scheduler.step( noise_pred[:, :orig_noise_len].unsqueeze(0) if recammaster is not None else noise_pred.unsqueeze(0), - timestep, + t, latent[:, :orig_noise_len].unsqueeze(0) if recammaster is not None else latent.unsqueeze(0), **scheduler_step_args)[0].squeeze(0) @@ -2865,7 +2882,7 @@ class WanVideoSampler: torch.cuda.reset_peak_memory_stats(device) except: pass - + print("samples out stats: mean", latent.mean().item(), "std", latent.std().item(), "min", latent.min().item(), "max", latent.max().item()) return ({ "samples": latent.unsqueeze(0).cpu(), "looped": is_looped, @@ -2873,7 +2890,7 @@ class WanVideoSampler: "has_ref": has_ref, "drop_last": drop_last, "generator_state": seed_g.get_state(), - }, ) + },) #region VideoDecode class WanVideoDecode: @@ -2960,7 +2977,7 @@ class WanVideoDecode: if is_looped: #images = images[:, warmup_latent_count * 4:] temp_latents = torch.cat([latents[:, :, -3:]] + [latents[:, :, :2]], dim=2) - temp_images = vae.decode(temp_latents, device=device, end_=(end_image is not None), tiled=enable_vae_tiling, tile_size=(tile_x//8, tile_y//8), tile_stride=(tile_stride_x//8, tile_stride_y//8))[0] + temp_images = vae.decode(temp_latents, device=device, end_=(end_image is not None), tiled=enable_vae_tiling, tile_size=(tile_x//vae.upsampling_factor, tile_y//vae.upsampling_factor), tile_stride=(tile_stride_x//vae.upsampling_factor, tile_stride_y//vae.upsampling_factor))[0] temp_images = (temp_images - temp_images.min()) / (temp_images.max() - temp_images.min()) images = torch.cat([temp_images[:, 9:].to(images), images[:, 5:]], dim=1) @@ -3054,7 +3071,7 @@ class WanVideoEncode: latents = vae.encode_video(image.permute(0, 2, 1, 3, 4), parallel=False)# B, T, C, H, W latents = latents.permute(0, 2, 1, 3, 4) else: - latents = vae.encode(image * 2.0 - 1.0, device=device, tiled=enable_vae_tiling, tile_size=(tile_x//8, tile_y//8), tile_stride=(tile_stride_x//8, tile_stride_y//8)) + latents = vae.encode(image * 2.0 - 1.0, device=device, tiled=enable_vae_tiling, tile_size=(tile_x//vae.upsampling_factor, tile_y//vae.upsampling_factor), tile_stride=(tile_stride_x//vae.upsampling_factor, tile_stride_y//vae.upsampling_factor)) vae.model.clear_cache() if latent_strength != 1.0: latents *= latent_strength diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 78be6b4..30e8c5b 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -9,7 +9,7 @@ from .wanvideo.modules.t5 import T5EncoderModel from .wanvideo.modules.clip import CLIPModel from accelerate import init_empty_weights -from accelerate.utils import set_module_tensor_to_device +from .utils import set_module_tensor_to_device from .fp8_optimization import convert_linear_with_lora_and_scale @@ -54,7 +54,7 @@ class WanVideoModel(comfy.model_base.BaseModel): self.pipeline[k] = v try: - from comfy.latent_formats import Wan21 + from comfy.latent_formats import Wan21, Wan22 latent_format = Wan21 except: #for backwards compatibility log.warning("Wan21 latent format not found, update ComfyUI for better livepreview") @@ -62,11 +62,11 @@ except: #for backwards compatibility latent_format = HunyuanVideo class WanVideoModelConfig: - def __init__(self, dtype): + def __init__(self, dtype, latent_format=latent_format): self.unet_config = {} self.unet_extra_config = {} self.latent_format = latent_format - self.latent_format.latent_channels = 16 + #self.latent_format.latent_channels = 16 self.manual_cast_dtype = dtype self.sampling_settings = {"multiplier": 1.0} self.memory_usage_factor = 2.0 @@ -855,8 +855,18 @@ class WanVideoModelLoader: elif "control_adapter.conv.weight" in sd: model_type = "t2v" - num_heads = 40 if dim == 5120 else 12 - num_layers = 40 if dim == 5120 else 30 + out_dim = 16 + if dim == 5120: #14B + num_heads = 40 + num_layers = 40 + elif dim == 3072: #5B + num_heads = 24 + num_layers = 30 + out_dim = 48 + model_type = "t2v" #5B no img crossattn + else: #1.3B + num_heads = 12 + num_layers = 30 vace_layers, vace_in_dim = None, None if "vace_blocks.0.after_proj.weight" in sd: @@ -908,6 +918,8 @@ class WanVideoModelLoader: if dim == 1536: model_variant = "1_3B" + if dim == 3072: + log.info(f"5B model detected, no Teacache or MagCache coefficients available, consider using EasyCache for this model") log.info(f"Model variant detected: {model_variant}") TRANSFORMER_CONFIG= { @@ -920,7 +932,7 @@ class WanVideoModelLoader: "freq_dim": 256, "in_dim": in_channels, "model_type": model_type, - "out_dim": 16, + "out_dim": out_dim, "text_len": 512, "num_heads": num_heads, "num_layers": num_layers, @@ -1001,8 +1013,9 @@ class WanVideoModelLoader: transformer.add_proj = zero_module(torch.nn.Linear(inner_dim, inner_dim)) transformer.attn_conv_in = torch.nn.Conv3d(attn_cond_in_dim, inner_dim, kernel_size=transformer.patch_size, stride=transformer.patch_size) + latent_format=Wan22 if dim == 3072 else Wan21 comfy_model = WanVideoModel( - WanVideoModelConfig(base_dtype), + WanVideoModelConfig(base_dtype, latent_format=latent_format), model_type=comfy.model_base.ModelType.FLOW, device=device, ) @@ -1299,7 +1312,7 @@ class WanVideoVAELoader: DESCRIPTION = "Loads Wan VAE model from 'ComfyUI/models/vae'" def loadmodel(self, model_name, precision): - from .wanvideo.wan_video_vae import WanVideoVAE + from .wanvideo.wan_video_vae import WanVideoVAE, WanVideoVAE38 dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] #with open(os.path.join(script_directory, 'configs', 'hy_vae_config.json')) as f: @@ -1310,8 +1323,12 @@ class WanVideoVAELoader: has_model_prefix = any(k.startswith("model.") for k in vae_sd.keys()) if not has_model_prefix: vae_sd = {f"model.{k}": v for k, v in vae_sd.items()} - - vae = WanVideoVAE(dtype=dtype) + + if vae_sd["model.conv2.weight"].shape[0] == 16: + vae = WanVideoVAE(dtype=dtype) + elif vae_sd["model.conv2.weight"].shape[0] == 48: + vae = WanVideoVAE38(dtype=dtype) + vae.load_state_dict(vae_sd) vae.eval() vae.to(device = offload_device, dtype = dtype) diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 4f11a86..c790400 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -1537,9 +1537,7 @@ class WanModel(ModelMixin, ConfigMixin): else: expanded_timesteps = False - e = self.time_embedding( - sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(x.dtype) - ) # b, dim + e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(x.dtype)) # b, dim e0 = self.time_projection(e).unflatten(1, (6, self.dim)) # b, 6, dim if fps_embeds is not None: diff --git a/wanvideo/schedulers/basic_flowmatch.py b/wanvideo/schedulers/basic_flowmatch.py index a44d66e..4a3715e 100644 --- a/wanvideo/schedulers/basic_flowmatch.py +++ b/wanvideo/schedulers/basic_flowmatch.py @@ -52,11 +52,11 @@ class FlowMatchScheduler(): timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) - if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): - sigma_ = 1 if ( - self.inverse_timesteps or self.reverse_sigmas) else 0 - else: - sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) + # if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): + # sigma_ = 1 if ( + # self.inverse_timesteps or self.reverse_sigmas) else 0 + #else: + sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) prev_sample = sample + model_output * (sigma_ - sigma) return prev_sample diff --git a/wanvideo/wan_video_vae.py b/wanvideo/wan_video_vae.py index 0165a13..64c3e00 100644 --- a/wanvideo/wan_video_vae.py +++ b/wanvideo/wan_video_vae.py @@ -196,6 +196,75 @@ class Resample(nn.Module): nn.init.zeros_(conv.bias.data) + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size) + return x + + +class Resample38(Resample): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super(Resample, self).__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + else: + self.resample = nn.Identity() + class ResidualBlock(nn.Module): def __init__(self, in_dim, out_dim, dropout=0.0): @@ -274,6 +343,178 @@ class AttentionBlock(nn.Module): return x + identity +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample38(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False + ): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample38(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + class Encoder3d(nn.Module): def __init__(self, @@ -377,6 +618,122 @@ class Encoder3d(nn.Module): return x +class Encoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] if i < len(temperal_downsample) else False + ) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + class Decoder3d(nn.Module): def __init__(self, @@ -482,6 +839,108 @@ class Decoder3d(nn.Module): return x + +class Decoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock(in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1)) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1)) + + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + def count_conv3d(model): count = 0 for m in model.modules(): @@ -690,6 +1149,7 @@ class WanVideoVAE(nn.Module): self.mean = torch.tensor(mean) self.std = torch.tensor(std) self.scale = [self.mean, 1.0 / self.std] + self.z_dim = z_dim # init model self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) @@ -798,7 +1258,7 @@ class WanVideoVAE(nn.Module): if end_: out_T += 1 weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) - values = torch.zeros((1, 16, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) pbar = ProgressBar(len(tasks)) for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): @@ -910,4 +1370,122 @@ class WanVideoVAEStateDictConverter: state_dict = state_dict['model_state'] for name in state_dict: state_dict_['model.' + name] = state_dict[name] - return state_dict_ \ No newline at end of file + return state_dict_ + + +class VideoVAE38_(VideoVAE_): + + def __init__(self, + dim=160, + z_dim=48, + dec_dim=256, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + dtype=torch.bfloat16): + super(VideoVAE_, self).__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + self.dtype = dtype + + # modules + self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True) + else: + out_ = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + +class WanVideoVAE38(WanVideoVAE): + + def __init__(self, z_dim=48, dim=160, dtype=torch.bfloat16): + super(WanVideoVAE, self).__init__() + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + self.dtype = dtype + + # init model + self.model = VideoVAE38_(z_dim=z_dim, dim=dim, dtype=dtype).eval().requires_grad_(False) + self.upsampling_factor = 16 \ No newline at end of file