From 7c68298b5f54a2ea5f9b4a1a9119eeea57d349bd Mon Sep 17 00:00:00 2001 From: linjunfan Date: Fri, 14 Aug 2026 16:07:38 +0800 Subject: [PATCH 1/4] modify wan2.1 to wan2.2 --- .../models/wan_i2v/custom_transformer.py | 251 ++++++++++------- core/finetune/models/wan_i2v/sft_trainer.py | 264 ++++++++---------- 2 files changed, 274 insertions(+), 241 deletions(-) diff --git a/core/finetune/models/wan_i2v/custom_transformer.py b/core/finetune/models/wan_i2v/custom_transformer.py index e61f5b6..d19be8a 100644 --- a/core/finetune/models/wan_i2v/custom_transformer.py +++ b/core/finetune/models/wan_i2v/custom_transformer.py @@ -12,6 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +#--------------------------------------------------------------------------- +# Changes from Wan2.1 (custom_transformer.py) to Wan2.2: +# +# (1) Per-timestep modulation: e [B,L,6,C] replaces global temb [B,C] +# in WanTransformerBlock and output head. +# (2) Per-position time embedding: condition_embedder sub-components +# called with [B*L] instead of [B], producing e_mod [B,L,6,dim]. +# +# All other components are IDENTICAL to the Wan2.1 version. +# --------------------------------------------------------------------------- + import math from typing import Any, Dict, Optional, Tuple, Union import os @@ -33,6 +44,10 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# ============================================================================== +# WanAttnProcessor2_0 — UNCHANGED from Wan2.1 +# ============================================================================== + class WanAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): @@ -101,9 +116,8 @@ def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor): hidden_states_img = hidden_states_img.type_as(query) if cos_sim is not None and not do_kv_cache: - cos_sim = cos_sim + 1.0 - attn_mask_from_cos_sim = torch.log(cos_sim + 1e-6) - attention_mask_GGA = attention_mask_GGA[0,:,0] + attn_mask_from_cos_sim = cos_sim # 已在 forward 中完成 +1.0 / log 变换 + attention_mask_GGA = attention_mask_GGA[0,:,0] #! Exo attention hidden_states_exo = F.scaled_dot_product_attention( query[:, :, attention_mask_GGA==1], @@ -121,7 +135,7 @@ def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor): attn_mask=attn_mask_ego.to(query.dtype), dropout_p=0.0, is_causal=False - ) + ) #! Combine exo and ego hidden_states = torch.zeros_like(query) @@ -145,6 +159,10 @@ def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor): return hidden_states, kv_cache +# ============================================================================== +# WanImageEmbedding — UNCHANGED from Wan2.1 +# ============================================================================== + class WanImageEmbedding(torch.nn.Module): def __init__(self, in_features: int, out_features: int, pos_embed_seq_len=None): super().__init__() @@ -169,6 +187,10 @@ def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: return hidden_states +# ============================================================================== +# WanTimeTextImageEmbedding — UNCHANGED from Wan2.1 +# ============================================================================== + class WanTimeTextImageEmbedding(nn.Module): def __init__( self, @@ -196,13 +218,22 @@ def forward( timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: Optional[torch.Tensor] = None, + timestep_seq_len: Optional[int] = None, ): + if timestep_seq_len is not None: + batch_size = timestep.shape[0] + timestep = timestep.unsqueeze(1).expand(batch_size, timestep_seq_len).flatten() + timestep = self.timesteps_proj(timestep) time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: timestep = timestep.to(time_embedder_dtype) temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + + if timestep_seq_len is not None: + temb = temb.unflatten(0, (batch_size, timestep_seq_len)) + timestep_proj = self.time_proj(self.act_fn(temb)) encoder_hidden_states = self.text_embedder(encoder_hidden_states) @@ -212,6 +243,10 @@ def forward( return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image +# ============================================================================== +# WanRotaryPosEmbed — UNCHANGED from Wan2.1 +# ============================================================================== + class WanRotaryPosEmbed(nn.Module): def __init__( self, attention_head_dim: int, patch_size: Tuple[int, int, int], max_seq_len: int, theta: float = 10000.0 @@ -256,6 +291,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return freqs +# ============================================================================== +# SelfAttention — UNCHANGED from Wan2.1 +# ============================================================================== + class SelfAttention(Attention): def __init__( self, @@ -310,6 +349,10 @@ def forward( return hidden_states +# ============================================================================== +# CrossAttention — UNCHANGED from Wan2.1 +# ============================================================================== + class CrossAttention(Attention): def __init__( self, @@ -353,7 +396,13 @@ def forward( attention_mask=None, ) return hidden_states - + + +# ============================================================================== +# WanTransformerBlock — [CHANGED] per-timestep modulation (Wan2.2) +# ============================================================================== +# Wan2.1: global temb [B,6,C] broadcast to all positions +# Wan2.2: per-position e [B,L,6,C] → chunk(6,dim=2) → 6×[B,L,1,C] → squeeze(2) class WanTransformerBlock(nn.Module): def __init__( @@ -404,6 +453,7 @@ def __init__( self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + # [Wan2.2] modulation [1,6,dim] — same shape & init as Wan2.1 scale_shift_table self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward( @@ -417,66 +467,51 @@ def forward( cos_sim: Optional[torch.Tensor] = None, do_kv_cache: bool = False, ) -> torch.Tensor: + # [Wan2.2] Per-timestep modulation: [1,6,dim] + [B,L,6,dim] → chunk(6,dim=2) → 6×[B,L,1,dim] + # (vs Wan2.1: [1,6,dim] + [B,1,dim] → chunk(6,dim=1) → 6×[B,1,dim] broadcasting to all positions) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table + temb.float() - ).chunk(6, dim=1) + self.scale_shift_table.unsqueeze(0) + temb.float() + ).chunk(6, dim=2) + # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) - attn_output = self.attn1(hidden_states=norm_hidden_states, rotary_emb=rotary_emb, - attention_GGA=attention_GGA, attention_mask_GGA=attention_mask_GGA, cos_sim=cos_sim, + norm_hidden_states = ( + self.norm1(hidden_states.float()) * (1 + scale_msa.squeeze(2)) + shift_msa.squeeze(2) + ).type_as(hidden_states) + attn_output = self.attn1(hidden_states=norm_hidden_states, rotary_emb=rotary_emb, + attention_GGA=attention_GGA, attention_mask_GGA=attention_mask_GGA, cos_sim=cos_sim, do_kv_cache=do_kv_cache) - hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + hidden_states = (hidden_states.float() + attn_output.float() * gate_msa.squeeze(2)).type_as(hidden_states) # 2. Cross-attention norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) - attn_output = self.attn2(hidden_states=norm_hidden_states, encoder_hidden_states=encoder_hidden_states ) #, attention_GGA=attention_GGA, attention_mask_GGA=attention_mask_GGA) + attn_output = self.attn2(hidden_states=norm_hidden_states, encoder_hidden_states=encoder_hidden_states ) hidden_states = hidden_states + attn_output # 3. Feed-forward - norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( - hidden_states - ) + norm_hidden_states = ( + self.norm3(hidden_states.float()) * (1 + c_scale_msa.squeeze(2)) + c_shift_msa.squeeze(2) + ).type_as(hidden_states) ff_output = self.ffn(norm_hidden_states) - hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa.squeeze(2)).type_as(hidden_states) return hidden_states +# ============================================================================== +# WanTransformer3DModel_GGA — [CHANGED] per-position time embedding + output head (Wan2.2) +# ============================================================================== + class WanTransformer3DModel_GGA(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin): r""" - A Transformer model for video-like data used in the Wan model. - - Args: - patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): - 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). - num_attention_heads (`int`, defaults to `40`): - Fixed length for text embeddings. - attention_head_dim (`int`, defaults to `128`): - The number of channels in each head. - in_channels (`int`, defaults to `16`): - The number of channels in the input. - out_channels (`int`, defaults to `16`): - The number of channels in the output. - text_dim (`int`, defaults to `512`): - Input dimension for text embeddings. - freq_dim (`int`, defaults to `256`): - Dimension for sinusoidal time embeddings. - ffn_dim (`int`, defaults to `13824`): - Intermediate dimension in feed-forward network. - num_layers (`int`, defaults to `40`): - The number of layers of transformer blocks to use. - window_size (`Tuple[int]`, defaults to `(-1, -1)`): - Window size for local attention (-1 indicates global attention). - cross_attn_norm (`bool`, defaults to `True`): - Enable cross-attention normalization. - qk_norm (`bool`, defaults to `True`): - Enable query/key normalization. - eps (`float`, defaults to `1e-6`): - Epsilon value for normalization layers. - add_img_emb (`bool`, defaults to `False`): - Whether to use img_emb. - added_kv_proj_dim (`int`, *optional*, defaults to `None`): - The number of channels to use for the added key and value projections. If `None`, no projection is used. + Wan2.2 I2V-A14B GGA Transformer. + + Derived from Wan2.1 ``WanTransformer3DModel_GGA`` (custom_transformer.py). + All __init__ and forward parameters are identical to Wan2.1. + + Wan2.2 changes from Wan2.1: + - Per-timestep modulation: e [B,L,6,C] → blocks + output head + - Per-position time embedding: condition_embedder sub-components called + with [B*L] instead of [B] in forward() """ _supports_gradient_checkpointing = True @@ -510,12 +545,14 @@ def __init__( inner_dim = num_attention_heads * attention_head_dim out_channels = out_channels or in_channels - # 1. Patch & position embedding + # 1. Patch & position embedding (unchanged) self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) - # 2. Condition embeddings - # image_embedding_dim=1280 for I2V model + # 2. Condition embeddings — diffusers standard components (Timesteps, + # TimestepEmbedding, PixArtAlphaTextProjection). Same as Wan2.1. + # For Wan2.2 per-position time embedding, these sub-components are + # called with [B*L] input in forward() instead of here. self.condition_embedder = WanTimeTextImageEmbedding( dim=inner_dim, time_freq_dim=freq_dim, @@ -538,6 +575,7 @@ def __init__( # 4. Output norm & projection self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + # [Wan2.2] head_modulation [1,2,dim] (same shape & init as Wan2.1 scale_shift_table) self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False @@ -557,12 +595,12 @@ def forward( cos_sim_scaling_factor: float = 1.0, do_kv_cache: bool = False, do_custom_test: bool = False, - save_attn_map: bool = True, - save_attn_dir_name: Optional[str] = None, - save_attn_map_mask: Optional[torch.Tensor] = None, - save_attn_map_avg: Optional[torch.Tensor] = None, - save_attn_dir_step: str = "", - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + # save_attn_map: bool = True, + # save_attn_dir_name: Optional[str] = None, + # save_attn_map_mask: Optional[torch.Tensor] = None, + # save_attn_map_avg: Optional[torch.Tensor] = None, + # save_attn_dir_step: str = "", + ) -> Union[torch.Tensor, Transformer2DModelOutput]: if attention_kwargs is not None: attention_kwargs = attention_kwargs.copy() lora_scale = attention_kwargs.pop("scale", 1.0) @@ -584,19 +622,20 @@ def forward( post_patch_height = height // p_h post_patch_width = width // p_w - + # ── RoPE + Patch embedding (unchanged) ────────────────────────── rotary_emb = self.rope(hidden_states) hidden_states = self.patch_embedding(hidden_states) - if save_attn_map: # B, C, F, H, W - save_attn_map_mask = torch.zeros_like(hidden_states[..., -(post_patch_height - post_patch_width):]) - save_attn_map_mask[...,:1,:,:] = 1 - save_attn_map_mask = save_attn_map_mask.flatten(2).transpose(1, 2) + # if save_attn_map: # B, C, F, H, W + # save_attn_map_mask = torch.zeros_like(hidden_states[..., -(post_patch_height - post_patch_width):]) + # save_attn_map_mask[...,:1,:,:] = 1 + # save_attn_map_mask = save_attn_map_mask.flatten(2).transpose(1, 2) hidden_states = hidden_states.flatten(2).transpose(1, 2) + seq_len = hidden_states.size(1) - - if attention_GGA is not None: + # ── GGA preparation (unchanged) ───────────────────────────────── + if attention_GGA is not None: if len(point_vecs_per_frame.shape) != 5: point_vecs_per_frame = point_vecs_per_frame.squeeze() @@ -613,7 +652,7 @@ def forward( count_include_pad=False, ).to(dtype=hidden_states.dtype) - + attention_GGA = attention_GGA.permute(0,4,1,2,3) # B, C, F, H, W attention_GGA = F.avg_pool3d( attention_GGA, @@ -636,59 +675,75 @@ def forward( frame_cut = H * W cos_sim = None - cos_sim = torch.zeros((1, attention_GGA.shape[1], attention_GGA.shape[1]), device=hidden_states.device) #! 1, 28028, 28028 - for i in range(FF): - cos_sim[:, i*frame_cut:(i+1)*frame_cut,:] = torch.matmul(attention_GGA[:,i*frame_cut:(i+1)*frame_cut,:], point_vecs_per_frame[i:i+1].transpose(-1, -2)).to(hidden_states.device) - - cos_sim = torch.clamp(cos_sim, min=-1.0, max=1.0) - - #! Scaling with hyperparameter - if cos_sim_scaling_factor > 1.0: - mask = cos_sim > 0 - cos_sim[mask] = cos_sim[mask] * cos_sim_scaling_factor - else: - cos_sim = cos_sim * cos_sim_scaling_factor - - attention_mask_GGA_ = attention_mask_GGA[:,:,0] #! 1, 28028 - mask = (attention_mask_GGA_[:, :, None] == 1).to(hidden_states.device) - cos_sim.masked_fill_(mask, 0.0) + # cos_sim 是 attention mask、无需梯度：no_grad 使原地操作合法，并规避 checkpoint + # 重算时的 autograd version 检查；dtype=bf16 将 [1,N,N] 张量减半且与 SDPA mask 一致。 + with torch.no_grad(): + cos_sim = torch.zeros((1, attention_GGA.shape[1], attention_GGA.shape[1]), device=hidden_states.device, dtype=hidden_states.dtype) #! 1, 28028, 28028 + for i in range(FF): + cos_sim[:, i*frame_cut:(i+1)*frame_cut,:] = torch.matmul(attention_GGA[:,i*frame_cut:(i+1)*frame_cut,:], point_vecs_per_frame[i:i+1].transpose(-1, -2)).to(hidden_states.device) + + cos_sim = torch.clamp(cos_sim, min=-1.0, max=1.0) + + #! Scaling with hyperparameter + if cos_sim_scaling_factor > 1.0: + mask = cos_sim > 0 + cos_sim[mask] = cos_sim[mask] * cos_sim_scaling_factor + else: + cos_sim = cos_sim * cos_sim_scaling_factor + + attention_mask_GGA_ = attention_mask_GGA[:,:,0] #! 1, 28028 + mask = (attention_mask_GGA_[:, :, None] == 1).to(hidden_states.device) + cos_sim.masked_fill_(mask, 0.0) + # log 变换移到 checkpoint 外：原地 add_/log_ 不触碰 checkpoint 保存的输入 + cos_sim = cos_sim.add_(1.0).add_(1e-6).log_() else: cos_sim = None + + # ── [Wan2.2] Per-position time embedding ────────────────────── + # Reuse condition_embedder's Timesteps + TimestepEmbedding, + # but call with [B*L] instead of [B] for per-position encoding. + # (vs Wan2.1: self.condition_embedder(timestep, ...) → global temb + timestep_proj) + temb, e_mod, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image, + timestep_seq_len=seq_len) - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image - ) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) + e_mod = e_mod.unflatten(-1, (6, self.config.num_attention_heads * self.config.attention_head_dim)) # [B, L, 6, dim] + + # ── Text + optional Image embedding ───────────────────────── if encoder_hidden_states_image is not None: + encoder_hidden_states_image = self.condition_embedder.image_embedder(encoder_hidden_states_image) encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - # 4. Transformer blocks + # ── Transformer blocks ────────────────────────────────────────── + # [Wan2.2] Pass e_mod [B,L,6,C] instead of timestep_proj [B,6,C] cur_block = 0 if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, attention_GGA, attention_mask_GGA, cos_sim) + block, hidden_states, encoder_hidden_states, e_mod, + rotary_emb, attention_GGA, attention_mask_GGA, cos_sim) cur_block += 1 else: for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, attention_GGA, attention_mask_GGA, cos_sim=cos_sim, do_kv_cache=do_kv_cache) + hidden_states = block( + hidden_states, encoder_hidden_states, e_mod, rotary_emb, + attention_GGA, attention_mask_GGA, cos_sim=cos_sim, do_kv_cache=do_kv_cache) cur_block += 1 - - # 5. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - # Move the shift and scale tensors to the same device as hidden_states. - # When using multi-GPU inference via accelerate these will be on the - # first device rather than the last device, which hidden_states ends up - # on. - shift = shift.to(hidden_states.device) - scale = scale.to(hidden_states.device) + # ── [Wan2.2] Output head with per-timestep modulation ─────────── + # (vs Wan2.1: global temb + scale_shift_table [1,2,dim]) + head_e = temb.unsqueeze(2) # [B, L, 1, dim] + shift, scale = (self.scale_shift_table.unsqueeze(0) + head_e.float()).chunk(2, dim=2) + + shift = shift.squeeze(2).to(hidden_states.device) + scale = scale.squeeze(2).to(hidden_states.device) hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) hidden_states = self.proj_out(hidden_states) + # ── Unpatchify (unchanged) ────────────────────────────────────── hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 ) diff --git a/core/finetune/models/wan_i2v/sft_trainer.py b/core/finetune/models/wan_i2v/sft_trainer.py index 14e0d6e..bb4769f 100644 --- a/core/finetune/models/wan_i2v/sft_trainer.py +++ b/core/finetune/models/wan_i2v/sft_trainer.py @@ -6,7 +6,7 @@ FlowMatchEulerDiscreteScheduler, WanImageToVideoPipeline, ) -from core.finetune.models.wan_i2v.custom_transformer import WanTransformer3DModel_GGA as WanTransformer3DModel +from core.finetune.models.wan_i2v.custom_transformer import WanTransformer3DModel_GGA as WanTransformer3DModel from diffusers.utils import ( replace_example_docstring, logging, @@ -54,6 +54,7 @@ from core.finetune.trainer import Trainer from core.finetune.utils import unwrap_model + from ..utils import register from diffusers import ( @@ -65,14 +66,14 @@ def generate_uniform_pointmap(height, width): x = np.linspace(-1, 1, width) y = np.linspace(-1, 1, height) xv, yv = np.meshgrid(x, y) - + # Create a spatially varying Z value: increases from bottom (row 0) to top (row height-1) zv = 1 - np.linspace(0, 1, height)[:, None] # shape (H, 1) zv = np.repeat(zv, width, axis=1) # shape (H, W) - + # Stack to get (H, W, 3) array pointmap = np.stack([xv, yv, zv], axis=-1) - + # Normalize XYZ to [0, 1] for image saving # X and Y are already in [-1, 1], so map to [0, 1] pointmap[..., 0] = (pointmap[..., 0] + 1) / 2 @@ -98,14 +99,27 @@ def __init__( self, tokenizer: AutoTokenizer, text_encoder: UMT5EncoderModel, - image_encoder: CLIPVisionModel, - image_processor: CLIPImageProcessor, + # image_encoder: CLIPVisionModel, + # image_processor: CLIPImageProcessor, transformer: WanTransformer3DModel, vae: AutoencoderKLWan, scheduler: FlowMatchEulerDiscreteScheduler, + transformer_2: WanTransformer3DModel, + boundary_ratio: float = 0.900, ): - super().__init__(tokenizer, text_encoder, image_encoder, image_processor, transformer, vae, scheduler) + # [Wan2.2] image_encoder/image_processor accepted for API compatibility but passed as None (no CLIP) + super().__init__( + tokenizer=tokenizer, + text_encoder=text_encoder, + image_encoder=None, + image_processor=None, + transformer=transformer, + vae=vae, + scheduler=scheduler, + ) self._shared_latent_exo_randn = None + self.transformer_2 = transformer_2 + self.boundary_ratio = boundary_ratio @override def prepare_latents( @@ -115,18 +129,18 @@ def prepare_latents( num_channels_latents: int = 16, height: int = 480, exo_width: int = 784, - ego_width: int = 448, + ego_width: int = 448, num_frames: int = 49, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.Tensor] = None, last_images: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> Tuple[torch.Tensor, torch.Tensor, Tuple[int, int]]: # Calculate latent dimensions for each view type num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 latent_height = height // self.vae_scale_factor_spatial - exo_latent_width = exo_width // self.vae_scale_factor_spatial + exo_latent_width = exo_width // self.vae_scale_factor_spatial ego_latent_width = ego_width // self.vae_scale_factor_spatial total_latent_width = exo_latent_width + ego_latent_width @@ -150,17 +164,17 @@ def prepare_latents( # Add batch dimension if missing (video_processor outputs [C, F, H, W], VAE expects [B, C, F, H, W]) if ego_prior_video.dim() == 4: ego_prior_video = ego_prior_video.unsqueeze(0) # Add batch dimension - + ego_prior_latents = self.encode_video(ego_prior_video.to(device=device, dtype=self.vae.dtype)) - + # Calculate ego region boundaries - latent_total_width = self._shared_latent_exo_randn.shape[-1] + latent_total_width = self._shared_latent_exo_randn.shape[-1] ego_latent_width = ego_prior_latents.shape[-1] exo_latent_width = latent_total_width - ego_latent_width - + # Create a copy to avoid modifying the shared variable video_condition = self._shared_latent_exo_randn.clone() - + # Replace ego region (right part) with ego_prior_latents video_condition[:, :, :, :, exo_latent_width:] = ego_prior_latents @@ -173,10 +187,10 @@ def prepare_latents( # Create mask for exo and ego views # Use num_frames (target frame count) like the parent class, not frame_num (input frame count) mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, total_latent_width) - + # Set exo_view mask (left part) to 1.0 mask_lat_size[:, :, :, :, :exo_latent_width] = 1.0 - # Set ego_view mask (right part) to 0.0 + # Set ego_view mask (right part) to 0.0 mask_lat_size[:, :, :, :, exo_latent_width:] = 0.0 if last_images is None: @@ -191,18 +205,18 @@ def prepare_latents( mask_lat_size = mask_lat_size.to(latent_condition.device) return latents, torch.concat([mask_lat_size, latent_condition], dim=1), (exo_latent_width, ego_latent_width) - + def encode_video(self, video: torch.Tensor) -> torch.Tensor: """Encode video to latent space (copied from Trainer)""" # shape of input video: [B, C, F, H, W] vae = self.vae video = video.to(vae.device, dtype=vae.dtype) - + # Ensure video is in [B, C, F, H, W] format if video.dim() == 5 and video.shape[1] != 3: # If video is [B, F, C, H, W], transpose to [B, C, F, H, W] video = video.permute(0, 2, 1, 3, 4) - + latent_dist = vae.encode(video).latent_dist latent = latent_dist.sample() latents_mean = ( @@ -215,87 +229,53 @@ def encode_video(self, video: torch.Tensor) -> torch.Tensor: ) latent = (latent - latents_mean) * latents_std return latent - - def _setup_ic_lora_latent(self, exo_video, generator, device, height, width, + + + def _setup_ic_lora_latent(self, exo_video, generator, device, height, width, num_frames, batch_size=1, num_videos_per_prompt=1): """Setup shared latent for IC-LoRA (exo + random ego)""" if exo_video is None: return - + # Calculate latent dimensions latent_height = height // self.vae_scale_factor_spatial num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 num_channels_latents = self.vae.config.z_dim - - # Actual ego = 448 pixels -> 56 latents + + # Actual ego = 448 pixels -> 56 latents ego_pixel_width = 448 ego_latent_width = ego_pixel_width // self.vae_scale_factor_spatial - + batch_size = batch_size * num_videos_per_prompt # prepare_latent에서 bsz를 이렇게 넘기고 있음 - + # Encode exo_video using encode_video method # Add batch dimension if needed: [C, F, H, W] -> [1, C, F, H, W] if exo_video.dim() == 4: exo_video = exo_video.unsqueeze(0) exo_latents = self.encode_video(exo_video) # [B, C, F, H, exo_latent_width] exo_latents = exo_latents.repeat(batch_size, 1, 1, 1, 1) - + # Generate random ego latent ego_shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, ego_latent_width) # Generate on CPU with generator, then move to device ego_latents = torch.randn(ego_shape, generator=generator, dtype=exo_latents.dtype).to(device) - + # Concatenate exo + ego latents self._shared_latent_exo_randn = torch.cat([exo_latents, ego_latents], dim=-1) - + print(f"Setup IC-LoRA latent: exo_shape={exo_latents.shape}, ego_shape={ego_latents.shape}, total_shape={self._shared_latent_exo_randn.shape}") @override def encode_image( self, - exo_video: PipelineImageInput, + image: PipelineImageInput, device: Optional[torch.device] = None, ): """ - Override encode_image to use shared latent for image embedding - Input: exo video only - Return: CLIPEncode(decode(encode(exo video first frame)+randn_tensor)) + [Wan2.2] Override: Wan2.2 does NOT use CLIP image encoder. + Returns None — all conditioning is via concat in latent space. """ - device = device or self._execution_device - - # Use full latent for temporal consistency (not just first frame) - full_latent = self._shared_latent_exo_randn # [B, C, F, H, W] - - # Decode latent back to pixel space - # Remove VAE normalization first - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(full_latent.device, full_latent.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - full_latent.device, full_latent.dtype - ) - - # Denormalize - denorm_latent = full_latent / latents_std + latents_mean - - # Decode full video for temporal consistency, then extract first frame - decoded_video = self.vae.decode(denorm_latent.to(self.vae.dtype), return_dict=False)[0] - decoded_image = decoded_video[:, :, 0, :, :] # Extract first frame [B, C, H, W] - - print(f"Decoded image shape: {decoded_image.shape} for CLIP encoding") - - # Convert to float32 for CLIP compatibility (BFloat16 not supported by transformers image processor) - decoded_image = decoded_image.to(dtype=torch.float32) - - # Convert from [-1, 1] to [0, 1] range for CLIP processor - decoded_image = (decoded_image + 1.0) / 2.0 - - # Use decoded image for CLIP encoding - processor handles all normalization - image = self.image_processor(images=decoded_image, return_tensors="pt", do_rescale=False).to(device) - image_embeds = self.image_encoder(**image, output_hidden_states=True) - return image_embeds.hidden_states[-2] + return None @override def check_inputs( @@ -361,9 +341,10 @@ def __call__( negative_prompt: Union[str, List[str]] = None, height: int = 448, width: int = 784+448, #1232 - num_frames: int = 49, + num_frames: int = 81, num_inference_steps: int = 50, guidance_scale: float = 5.0, + guidance_scale_2: Optional[float] = None, num_videos_per_prompt: Optional[int] = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.Tensor] = None, @@ -468,7 +449,7 @@ def __call__( callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs device = self._execution_device - + # Preprocess exo_video for IC-LoRA setup if not isinstance(exo_video, torch.Tensor): # Truncate frames and convert PipelineImageInput to tensor and resize to exo dimensions @@ -476,7 +457,7 @@ def __call__( exo_video_processed = self.video_processor.preprocess(exo_video_truncated, height=height, width=width-height).to(device, dtype=torch.float32) else: exo_video_processed = exo_video.to(device, dtype=torch.float32) - + # Setup IC-LoRA latent at the beginning self._setup_ic_lora_latent( exo_video=exo_video_processed, @@ -510,7 +491,11 @@ def __call__( num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 num_frames = max(num_frames, 1) + if self.config.boundary_ratio is not None and guidance_scale_2 is None: + guidance_scale_2 = guidance_scale + self._guidance_scale = guidance_scale + self._guidance_scale_2 = guidance_scale_2 self._attention_kwargs = attention_kwargs self._current_timestep = None self._interrupt = False @@ -535,20 +520,12 @@ def __call__( device=device, ) - # Encode image embedding + # [Wan2.2] No CLIP: skip encode_image() and image_embeds handling transformer_dtype = self.transformer.dtype prompt_embeds = prompt_embeds.to(transformer_dtype) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) - if image_embeds is None: - if last_image is None: - image_embeds = self.encode_image(exo_video, device) - else: - image_embeds = self.encode_image([exo_video, last_image], device) - image_embeds = image_embeds.repeat(batch_size, 1, 1) - image_embeds = image_embeds.to(transformer_dtype) - # 4. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps @@ -571,7 +548,7 @@ def __call__( # Calculate individual widths from total width exo_width = width-height ego_width = height - + latents, condition, (exo_latent_width, ego_latent_width) = self.prepare_latents( ego_prior_video, batch_size * num_videos_per_prompt, @@ -591,6 +568,11 @@ def __call__( num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order self._num_timesteps = len(timesteps) + if self.config.boundary_ratio is not None: + boundary_timestep = self.config.boundary_ratio * self.scheduler.config.num_train_timesteps + else: + boundary_timestep = None + with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: @@ -600,11 +582,20 @@ def __call__( latent_model_input = torch.cat([latents, condition], dim=1).to(transformer_dtype) timestep = t.expand(latents.shape[0]) - noise_pred = self.transformer( + # [Wan2.2] MoE: select model for this timestep + if boundary_timestep is None or t >= boundary_timestep: + # wan2.1 or high-noise stage in wan2.2 + current_model = self.transformer + current_guidance_scale = guidance_scale + else: + # low-noise stage in wan2.2 + current_model = self.transformer_2 + current_guidance_scale = guidance_scale_2 + + noise_pred = current_model( hidden_states=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds, - encoder_hidden_states_image=image_embeds, attention_kwargs=attention_kwargs, return_dict=False, attention_GGA=attention_GGA, @@ -619,11 +610,10 @@ def __call__( )[0] if self.do_classifier_free_guidance: - noise_uncond = self.transformer( + noise_uncond = current_model( hidden_states=latent_model_input, timestep=timestep, encoder_hidden_states=negative_prompt_embeds, - encoder_hidden_states_image=image_embeds, attention_kwargs=attention_kwargs, return_dict=False, attention_GGA=attention_GGA, @@ -636,7 +626,7 @@ def __call__( save_attn_dir_name=save_attn_dir_name, save_attn_dir_step=f'step_{i:02d}', )[0] - noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond) # Set exo region (left part) of noise_pred to 0 to prevent exo latents from being updated # ego_latent_width = ego_width // self.vae_scale_factor_spatial # ego_width = height = 448 @@ -678,11 +668,11 @@ def __call__( # video = self.vae.decode(latents, return_dict=False)[0] latent_exo = latents[:, :, :, :, :exo_latent_width] # [B, C, F, H, exo_latent_width] latent_ego = latents[:, :, :, :, exo_latent_width:] # [B, C, F, H, ego_latent_width] - + # Decode each part separately video_exo = self.vae.decode(latent_exo, return_dict=False)[0] # [B, 3, F, H_exo, W_exo] video_ego = self.vae.decode(latent_ego, return_dict=False)[0] # [B, 3, F, H_ego, W_ego] - + # Concatenate along width dimension video = torch.cat([video_exo, video_ego], dim=-1) # [B, 3, F, H, W_exo + W_ego] video = self.video_processor.postprocess_video(video, output_type=output_type) @@ -698,34 +688,39 @@ def __call__( return WanPipelineOutput(frames=video) class WanI2VSftTrainer(Trainer): - UNLOAD_LIST = ["text_encoder", "image_encoder", "image_processor"] + UNLOAD_LIST = ["text_encoder"] # [Wan2.2] no CLIP @override def load_components(self) -> Dict[str, Any]: components = Components() model_path = str(self.args.model_path) - components.pipeline_cls = WanImageToVideoPipeline + components.pipeline_cls = WanWidthConcatImageToVideoPipeline components.tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") components.text_encoder = UMT5EncoderModel.from_pretrained(model_path, subfolder="text_encoder") - components.transformer = WanTransformer3DModel.from_pretrained(model_path, subfolder="transformer") + # [Wan2.2] MoE: load TWO transformers (see wan_i2v_A14B config) + components.transformer = WanTransformer3DModel.from_pretrained( + model_path, subfolder="transformer", torch_dtype=self.state.weight_dtype) # high noise model + + components.transformer_2 = WanTransformer3DModel.from_pretrained( + model_path, subfolder="transformer_2", torch_dtype=self.state.weight_dtype) # low noise model + components.boundary_ratio = getattr(self.args, "boundary_ratio", 0.900) components.vae = AutoencoderKLWan.from_pretrained(model_path, subfolder="vae") components.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler") - components.image_encoder = CLIPVisionModel.from_pretrained(model_path, subfolder="image_encoder") - - components.image_processor = CLIPImageProcessor.from_pretrained(model_path, subfolder="image_processor") + # [Wan2.2] image_encoder/image_processor omitted (Wan2.2 has no CLIP) return components @override def prepare_models(self) -> None: self.state.transformer_config = self.components.transformer.config + self.state.transformer_2_config = self.components.transformer_2.config @override def initialize_pipeline(self) -> WanWidthConcatImageToVideoPipeline: @@ -735,8 +730,10 @@ def initialize_pipeline(self) -> WanWidthConcatImageToVideoPipeline: vae=self.components.vae, transformer=unwrap_model(self.accelerator, self.components.transformer), scheduler=self.components.scheduler, - image_encoder=self.components.image_encoder, - image_processor=self.components.image_processor, + image_encoder=None, + image_processor=None, + transformer_2=unwrap_model(self.accelerator, self.components.transformer_2), + boundary_ratio=self.components.boundary_ratio, ) return pipe @@ -775,41 +772,19 @@ def encode_text(self, prompt: str) -> torch.Tensor: @override def encode_first_frame(self, image: torch.Tensor) -> torch.Tensor: """ - Encode image using CLIP image encoder for dataset preprocessing - - Args: - image: [C, H, W] single image tensor - - Returns: - image_embedding: [1, sequence_length, embedding_dimension] + [Wan2.2] Wan2.2 has no CLIP image encoder. Returns None. """ - # Ensure image is 3D [C, H, W] - if image.dim() != 3: - raise ValueError(f"Expected 3D image tensor [C, H, W], got {image.dim()}D tensor with shape {image.shape}") - - # Add batch dimension [1, C, H, W] - image = image.unsqueeze(0) - - # Use CLIP image processor and encoder - # Ensure image encoder is on the same device as inputs to avoid type/device mismatch - self.components.image_encoder = self.components.image_encoder.to(self.accelerator.device) - self.components.image_encoder.eval() - - processed_image = self.components.image_processor(images=image, return_tensors="pt") - processed_image = {k: v.to(self.accelerator.device) for k, v in processed_image.items()} - - image_embeds = self.components.image_encoder(**processed_image, output_hidden_states=True) - return image_embeds.hidden_states[-2] # [1, sequence_length, embedding_dimension] + return None @override def collate_fn(self, samples: List[Dict[str, Any]]) -> Dict[str, Any]: - ret = {"encoded_exo_ego_gt_video": [], "prompt_embedding": [], "encoded_exo_ego_prior_video": [], "image_embedding": [], "attention_GGA": [], "attention_mask_GGA": [], "point_vecs_per_frame": [], "cam_rays": []} + ret = {"encoded_exo_ego_gt_video": [], "prompt_embedding": [], "encoded_exo_ego_prior_video": [], "attention_GGA": [], "attention_mask_GGA": [], "point_vecs_per_frame": [], "cam_rays": []} for sample in samples: encoded_exo_ego_gt_video = sample["encoded_exo_ego_gt_video"] # exo+ego_gt concatenated and encoded prompt_embedding = sample["prompt_embedding"] encoded_exo_ego_prior_video = sample["encoded_exo_ego_prior_video"] # exo+ego_prior concatenated and encoded - image_embedding = sample["image_embedding"] # from exo_video + # [Wan2.2] image_embedding accepted but unused (no CLIP) attention_GGA = sample["attention_GGA"] attention_mask_GGA = sample["attention_mask_GGA"] @@ -820,7 +795,6 @@ def collate_fn(self, samples: List[Dict[str, Any]]) -> Dict[str, Any]: ret["encoded_exo_ego_gt_video"].append(encoded_exo_ego_gt_video) ret["prompt_embedding"].append(prompt_embedding) ret["encoded_exo_ego_prior_video"].append(encoded_exo_ego_prior_video) - ret["image_embedding"].append(image_embedding) if attention_GGA is not None: ret["attention_GGA"].append(attention_GGA) ret["attention_mask_GGA"].append(attention_mask_GGA) @@ -830,7 +804,7 @@ def collate_fn(self, samples: List[Dict[str, Any]]) -> Dict[str, Any]: ret["encoded_exo_ego_gt_video"] = torch.stack(ret["encoded_exo_ego_gt_video"]) ret["prompt_embedding"] = torch.stack(ret["prompt_embedding"]) ret["encoded_exo_ego_prior_video"] = torch.stack(ret["encoded_exo_ego_prior_video"]) # [B, C, F, H, W_exo+W_ego] - ret["image_embedding"] = torch.stack(ret["image_embedding"]) + # [Wan2.2] image_embedding stays as empty list (no CLIP) if len(ret["attention_GGA"]) > 0: ret["attention_GGA"] = torch.stack(ret["attention_GGA"]) @@ -851,14 +825,18 @@ def get_sigmas(self, timesteps, n_dim=4, dtype=torch.float32): @override def compute_loss(self, batch) -> torch.Tensor: - # Unwrap in case the transformer is wrapped by DDP/Accelerate - transformer = unwrap_model(self.accelerator, self.components.transformer) - transformer_dtype = transformer.dtype + # [Wan2.2] MoE: uniformly sample timestep from [0, 999]; + # the timestep itself is the expert router: + # t < boundary_ratio * 1000 → low_noise_model + # t >= boundary_ratio * 1000 → high_noise_model + # (vs Wan2.1 sft_trainer: timestep sampled from full range unconditionally) + boundary_step = int(self.components.boundary_ratio * self.components.scheduler.config.num_train_timesteps) + + transformer_dtype = self.components.transformer.dtype prompt_embedding = batch["prompt_embedding"].to(transformer_dtype) latent = batch["encoded_exo_ego_gt_video"].to(transformer_dtype) latent_condition = batch["encoded_exo_ego_prior_video"].to(transformer_dtype) - image_embedding = batch["image_embedding"].to(transformer_dtype) ##### attention_GGA = batch["attention_GGA"] if "attention_GGA" in batch else None @@ -878,22 +856,22 @@ def compute_loss(self, batch) -> torch.Tensor: # Get actual number of frames from video_condition actual_num_frames = (num_frames - 1) * vae_scale_factor_temporal + 1 - + mask_lat_size = torch.ones(latent_condition.shape[0], 1, actual_num_frames, latent_condition.shape[3], latent_condition.shape[4]) - + exo_width, ego_width = 784, 448 #! HARD CODING vae_scale_factor_spatial = 2 ** len(self.components.vae.config.temperal_downsample) # 8 - exo_latent_width = exo_width // vae_scale_factor_spatial + exo_latent_width = exo_width // vae_scale_factor_spatial ego_latent_width = ego_width // vae_scale_factor_spatial # Set exo_view mask (left part) to 1.0 mask_lat_size[:, :, :, :, :exo_latent_width] = 1.0 - # Set ego_view mask (right part) to 0.0 + # Set ego_view mask (right part) to 0.0 mask_lat_size[:, :, :, :, exo_latent_width:] = 0.0 - + # Use the same masking logic as WanWidthConcatImageToVideoPipeline.prepare_latents # Apply temporal masking (only first frame is visible, rest are masked) mask_lat_size[:, :, list(range(1, actual_num_frames - 1))] = 0 - + first_frame_mask = mask_lat_size[:, :, 0:1] first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=vae_scale_factor_temporal) mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) @@ -903,26 +881,26 @@ def compute_loss(self, batch) -> torch.Tensor: condition = torch.concat([mask_lat_size, latent_condition], dim=1) # [B, 20, 13, latent_H, latent_W] - # Sample a random timestep for each sample - timesteps_idx = torch.randint(0, self.components.scheduler.config.num_train_timesteps, (batch_size,)) - timesteps_idx = timesteps_idx.long() + # [Wan2.2] Uniform timestep sampling; timestep itself routes to expert + timesteps_idx = torch.randint(0, self.components.scheduler.config.num_train_timesteps, (batch_size,)).long() + batch_boundary = timesteps_idx[0].item() < boundary_step + current_model = self.components.transformer_2 if batch_boundary else self.components.transformer timesteps = self.components.scheduler.timesteps[timesteps_idx].to(device=latent.device) sigmas = self.get_sigmas(timesteps, n_dim=latent.ndim, dtype=latent.dtype) # Add noise to latent noise = torch.randn_like(latent) noisy_latents = (1.0 - sigmas) * latent + sigmas * noise - + # Apply IC-LoRA style masking: keep exo region (left) as original, only noise ego region (right) noisy_latents[:, :, :, :, :-ego_latent_width] = latent[:, :, :, :, :-ego_latent_width] - + target = noise - latent latent_model_input = torch.cat([noisy_latents, condition], dim=1) # [B, 36, 13, latent_H, latent_W] - predicted_noise = self.components.transformer( + predicted_noise = current_model( hidden_states=latent_model_input, encoder_hidden_states=prompt_embedding, - encoder_hidden_states_image=image_embedding, timestep=timesteps, return_dict=False, attention_GGA=attention_GGA, @@ -935,7 +913,7 @@ def compute_loss(self, batch) -> torch.Tensor: ego_predicted_noise = predicted_noise[:, :, :, :, -ego_latent_width:] ego_target = target[:, :, :, :, -ego_latent_width:] - + loss = torch.mean(((ego_predicted_noise.float() - ego_target.float()) ** 2).reshape(batch_size, -1), dim=1) loss = loss.mean() From 6e157fdb105bf9cf577f60dd70c4f72704baa132 Mon Sep 17 00:00:00 2001 From: linjunfan Date: Fri, 14 Aug 2026 16:36:54 +0800 Subject: [PATCH 2/4] refactor trainer --- core/finetune/trainer.py | 212 ++++++++++++++++++++++++++------------- 1 file changed, 142 insertions(+), 70 deletions(-) diff --git a/core/finetune/trainer.py b/core/finetune/trainer.py index cd8e286..30543cf 100644 --- a/core/finetune/trainer.py +++ b/core/finetune/trainer.py @@ -5,7 +5,7 @@ from datetime import timedelta from pathlib import Path from typing import Any, Dict, List, Tuple - +import os import diffusers import torch import transformers @@ -57,7 +57,11 @@ "fp16": torch.float16, # FP16 is Only Support for CogVideoX-2B "bf16": torch.bfloat16, } - +from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy, transformer_auto_wrap_policy, _or_policy +from accelerate import FullyShardedDataParallelPlugin +from torch.distributed.fsdp import BackwardPrefetch +from functools import partial +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP class Trainer: # If set, should be a list of components to unload (refer to `Components``) @@ -155,6 +159,8 @@ def prepare_models(self) -> None: self.components.vae.enable_tiling() self.state.transformer_config = self.components.transformer.config + if self.components.transformer_2 is not None: + self.state.transformer_2_config = self.components.transformer_2.config def prepare_dataset(self) -> None: logger.info("Initializing dataset and dataloader") @@ -166,8 +172,7 @@ def prepare_dataset(self) -> None: max_num_frames=self.state.train_frames, height=self.state.train_height, width=self.state.train_width, - trainer=self, - + trainer=self, ) else: raise ValueError(f"Invalid model type: {self.args.model_type}") @@ -225,7 +230,7 @@ def prepare_trainable_parameters(self): # For SFT, we train all the parameters in transformer model for attr_name, component in vars(self.components).items(): if hasattr(component, "requires_grad_"): - if self.args.training_type == "sft" and attr_name == "transformer": + if self.args.training_type == "sft" and attr_name in ["transformer", "transformer_2"]: component.requires_grad_(True) else: component.requires_grad_(False) @@ -237,29 +242,53 @@ def prepare_trainable_parameters(self): init_lora_weights=True, target_modules=self.args.target_modules, ) + transformer_2_lora_config = LoraConfig( + r=self.args.rank, + lora_alpha=self.args.lora_alpha, + init_lora_weights=True, + target_modules=self.args.target_modules, + ) self.components.transformer.add_adapter(transformer_lora_config) - self.__prepare_saving_loading_hooks(transformer_lora_config) + self.components.transformer_2.add_adapter(transformer_2_lora_config) + # LoRA权重与基座格式对齐 + self.components.transformer = self.components.transformer.to(dtype=torch.bfloat16) + self.components.transformer_2 = self.components.transformer_2.to(dtype=torch.bfloat16) + self.__prepare_saving_loading_hooks(transformer_lora_config, transformer_2_lora_config) for name, param in self.components.transformer.named_parameters(): if 'learnable_domain_embeddings' in name: param.requires_grad_(True) - logger.info(f"Training {name} after adding LoRA") + logger.info(f"Training transfomer {name} after adding LoRA") + for name, param in self.components.transformer_2.named_parameters(): + if 'learnable_domain_embeddings' in name: + param.requires_grad_(True) + logger.info(f"Training transformer_2 {name} after adding LoRA") # Load components needed for training to GPU (except transformer), and cast them to the specified data type - ignore_list = ["transformer"] + self.UNLOAD_LIST + ignore_list = ["transformer", "transformer_2"] + self.UNLOAD_LIST + if self.args.unload_vae_after_precompute: + ignore_list.append("vae") + logger.info("Keeping VAE on CPU after latent precompute") self.__move_components_to_device(dtype=weight_dtype, ignore_list=ignore_list) if self.args.gradient_checkpointing: self.components.transformer.enable_gradient_checkpointing() + if self.components.transformer_2 is not None: + self.components.transformer_2.enable_gradient_checkpointing() def prepare_optimizer(self) -> None: logger.info("Initializing optimizer and lr scheduler") # Make sure the trainable params are in float32 - cast_training_params([self.components.transformer], dtype=torch.float32) + # models_to_cast = [self.components.transformer] + # if self.components.transformer_2 is not None: + # models_to_cast.append(self.components.transformer_2) + # cast_training_params(models_to_cast, dtype=torch.float32) # For LoRA, we only want to train the LoRA weights # For SFT, we want to train all the parameters trainable_parameters = list(filter(lambda p: p.requires_grad, self.components.transformer.parameters())) + if self.components.transformer_2 is not None: + trainable_parameters += list(filter(lambda p: p.requires_grad, self.components.transformer_2.parameters())) transformer_parameters_with_lr = { "params": trainable_parameters, "lr": self.args.learning_rate, @@ -318,9 +347,12 @@ def prepare_optimizer(self) -> None: self.lr_scheduler = lr_scheduler def prepare_for_training(self) -> None: - self.components.transformer, self.optimizer, self.data_loader, self.lr_scheduler = self.accelerator.prepare( - self.components.transformer, self.optimizer, self.data_loader, self.lr_scheduler - ) + if self.components.transformer_2 is not None: + self.components.transformer, self.components.transformer_2, self.optimizer, self.data_loader, self.lr_scheduler = self.accelerator.prepare( + self.components.transformer, self.components.transformer_2, self.optimizer, self.data_loader, self.lr_scheduler) + else: + self.components.transformer, self.optimizer, self.data_loader, self.lr_scheduler = self.accelerator.prepare( + self.components.transformer, self.optimizer, self.data_loader, self.lr_scheduler) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(self.data_loader) / self.args.gradient_accumulation_steps) @@ -400,8 +432,11 @@ def train(self) -> None: logger.debug(f"Starting epoch ({epoch + 1}/{self.args.train_epochs})") self.components.transformer.train() - models_to_accumulate = [self.components.transformer] - + if self.components.transformer_2 is not None: + self.components.transformer_2.train() + models_to_accumulate = [self.components.transformer, self.components.transformer_2] + else: + models_to_accumulate = [self.components.transformer] for step, batch in enumerate(self.data_loader): logger.debug(f"Starting step {step + 1}") logs = {} @@ -413,13 +448,18 @@ def train(self) -> None: if accelerator.sync_gradients: if accelerator.distributed_type == DistributedType.DEEPSPEED: - grad_norm = self.components.transformer.get_global_grad_norm() - # In some cases the grad norm may not return a float - if torch.is_tensor(grad_norm): - grad_norm = grad_norm.item() + grad_norms = [self.components.transformer.get_global_grad_norm()] + if self.components.transformer_2 is not None: + grad_norms.append(self.components.transformer_2.get_global_grad_norm()) + grad_norm = max( + g.item() if torch.is_tensor(g) else g + for g in grad_norms) else: + all_transformer_parameters = list(self.components.transformer.parameters()) + if self.components.transformer_2 is not None: + all_transformer_parameters = list(self.components.transformer.parameters()) + list(self.components.transformer_2.parameters()) grad_norm = accelerator.clip_grad_norm_( - self.components.transformer.parameters(), self.args.max_grad_norm + all_transformer_parameters, self.args.max_grad_norm ) if torch.is_tensor(grad_norm): grad_norm = grad_norm.item() @@ -449,7 +489,10 @@ def train(self) -> None: logger.info(f"Memory after epoch {epoch + 1}: {json.dumps(memory_statistics, indent=4)}") accelerator.wait_for_everyone() - self.__maybe_save_checkpoint(global_step, must_save=True) + if global_step % self.args.checkpointing_steps == 0: + logger.info(f"Skipping duplicate final checkpoint save for step {global_step}") + else: + self.__maybe_save_checkpoint(global_step, must_save=True) del self.components @@ -475,9 +518,6 @@ def collate_fn(self, examples: List[Dict[str, Any]]): def load_components(self) -> Components: raise NotImplementedError - def initialize_pipeline(self) -> DiffusionPipeline: - raise NotImplementedError - def encode_video(self, video: torch.Tensor) -> torch.Tensor: # shape of input video: [B, C, F, H, W], where B = 1 # shape of output video: [B, C', F', H', W'], where B = 1 @@ -487,11 +527,6 @@ def encode_text(self, text: str) -> torch.Tensor: # shape of output text: [batch size, sequence length, embedding dimension] raise NotImplementedError - def encode_first_frame(self, image: torch.Tensor) -> torch.Tensor: - # shape of input image: [C, H, W] single image - # shape of output: [batch_size, sequence_length, embedding_dimension] - raise NotImplementedError - def compute_loss(self, batch) -> torch.Tensor: raise NotImplementedError @@ -521,63 +556,100 @@ def __move_components_to_cpu(self, unload_list: List[str] = []): if name in unload_list: setattr(self.components, name, component.to("cpu")) - def __prepare_saving_loading_hooks(self, transformer_lora_config): + def __prepare_saving_loading_hooks(self, transformer_lora_config, transformer_2_lora_config = None): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + """ + register save/load hook。每次 checkpoint 保存时，分别写入 transformer_lora.bin 和 + transformer_2_lora.bin; 恢复时从这两个文件加载。 + + Args: + transformer_lora_config: transformer 的 LoRA 配置 + transformer_2_lora_config: transformer_2 的 LoRA 配置（可选，默认同 transformer） + """ + if transformer_2_lora_config is None: + transformer_2_lora_config = transformer_lora_config + + lora_configs = { + "transformer": transformer_lora_config, + "transformer_2": transformer_2_lora_config, + } + # save hook def save_model_hook(models, weights, output_dir): if self.accelerator.is_main_process: - transformer_lora_layers_to_save = None + # 使用 unwrapped 后的对象身份 (id) 做精确匹配，避免两个同构模型 + # 因 isinstance 无法区分而依赖 models 列表顺序的隐式风险 + transformer_ref = unwrap_model(self.accelerator, self.components.transformer) + ref_to_key = {id(transformer_ref): "transformer"} + + if self.components.transformer_2 is not None: + transformer_2_ref = unwrap_model(self.accelerator, self.components.transformer_2) + ref_to_key[id(transformer_2_ref)] = "transformer_2" for model in models: - if isinstance( - unwrap_model(self.accelerator, model), - type(unwrap_model(self.accelerator, self.components.transformer)), - ): - model = unwrap_model(self.accelerator, model) - transformer_lora_layers_to_save = get_peft_model_state_dict(model) - else: + unwrapped = unwrap_model(self.accelerator, model) + key = ref_to_key.get(id(unwrapped)) + if key is None: raise ValueError(f"Unexpected save model: {model.__class__}") + lora_state_dict = get_peft_model_state_dict(unwrapped) + save_path = os.path.join(output_dir, f"{key}_lora.bin") + torch.save(lora_state_dict, save_path) + logger.info(f"Saved {key} LoRA weights → {save_path}") + # make sure to pop weight so that corresponding model is not saved again if weights: weights.pop() - self.components.pipeline_cls.save_lora_weights( - output_dir, - transformer_lora_layers=transformer_lora_layers_to_save, - ) def load_model_hook(models, input_dir): - if not self.accelerator.distributed_type == DistributedType.DEEPSPEED: - while len(models) > 0: - model = models.pop() - if isinstance( - unwrap_model(self.accelerator, model), - type(unwrap_model(self.accelerator, self.components.transformer)), - ): - transformer_ = unwrap_model(self.accelerator, model) - else: - raise ValueError(f"Unexpected save model: {unwrap_model(self.accelerator, model).__class__}") - else: - transformer_ = unwrap_model(self.accelerator, self.components.transformer).__class__.from_pretrained( - self.args.model_path, subfolder="transformer" - ) - transformer_.add_adapter(transformer_lora_config) - - lora_state_dict = self.components.pipeline_cls.lora_state_dict(input_dir) - transformer_state_dict = { - f'{k.replace("transformer.", "")}': v - for k, v in lora_state_dict.items() - if k.startswith("transformer.") - } - incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") - if incompatible_keys is not None: - # check only for unexpected keys - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - if unexpected_keys: + # 使用 unwrapped 后的对象身份 (id) 做精确匹配，与 save 侧一致， + # 完全消除对 models 列表顺序的隐式依赖 + transformer_ref = unwrap_model(self.accelerator, self.components.transformer) + ref_to_key = {id(transformer_ref): "transformer"} + + if self.components.transformer_2 is not None: + transformer_2_ref = unwrap_model(self.accelerator, self.components.transformer_2) + ref_to_key[id(transformer_2_ref)] = "transformer_2" + + model_map = {} + while len(models) > 0: + model = models.pop() + unwrapped = unwrap_model(self.accelerator, model) + key = ref_to_key.get(id(unwrapped)) + if key is None: + raise ValueError(f"Unexpected load model: {unwrapped.__class__}") + if key in model_map: logger.warning( - f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " - f" {unexpected_keys}. " + f"Duplicate model identity match for '{key}' — overwriting." + " This may indicate the same model was passed twice in models list." ) + model_map[key] = unwrapped + + # 对每个 transformer，从独立 .bin 文件加载 LoRA 权重 + for key in ["transformer", "transformer_2"]: + if key not in model_map: + logger.warning(f"Model {key} not found in checkpoint models list, skipping") + continue + transformer_ = model_map[key] + save_path = os.path.join(input_dir, f"{key}_lora.bin") + if not os.path.exists(save_path): + logger.warning(f"LoRA weights file not found: {save_path}, skipping {key}") + continue + lora_state_dict = torch.load(save_path, map_location="cpu") + incompatible_keys = set_peft_model_state_dict( + transformer_, lora_state_dict, adapter_name="default" + ) + if incompatible_keys is not None: + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading {key} adapter weights → unexpected keys: {unexpected_keys}" + ) + missing_keys = getattr(incompatible_keys, "missing_keys", None) + if missing_keys: + logger.warning( + f"Loading {key} adapter weights → missing keys: {missing_keys}" + ) self.accelerator.register_save_state_pre_hook(save_model_hook) self.accelerator.register_load_state_pre_hook(load_model_hook) From 93e0f4bf64effa7d5e37f137af97f751eb768a7f Mon Sep 17 00:00:00 2001 From: linjunfan Date: Fri, 14 Aug 2026 17:30:28 +0800 Subject: [PATCH 3/4] refactor schemas --- core/finetune/schemas/args.py | 2 ++ core/finetune/schemas/components.py | 5 +++++ core/finetune/schemas/state.py | 4 ++-- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/core/finetune/schemas/args.py b/core/finetune/schemas/args.py index c9f4a3f..1d51e5b 100644 --- a/core/finetune/schemas/args.py +++ b/core/finetune/schemas/args.py @@ -68,6 +68,7 @@ class Args(BaseModel): gradient_checkpointing: bool = True enable_slicing: bool = True enable_tiling: bool = True + unload_vae_after_precompute: bool = False nccl_timeout: int = 1800 ########## Lora ########## @@ -165,6 +166,7 @@ def parse_args(cls): parser.add_argument("--gradient_checkpointing", type=bool, default=True) parser.add_argument("--enable_slicing", type=bool, default=True) parser.add_argument("--enable_tiling", type=bool, default=True) + parser.add_argument("--unload_vae_after_precompute", action="store_true") parser.add_argument("--nccl_timeout", type=int, default=1800) # LoRA parameters diff --git a/core/finetune/schemas/components.py b/core/finetune/schemas/components.py index 19fef79..7513ce7 100644 --- a/core/finetune/schemas/components.py +++ b/core/finetune/schemas/components.py @@ -22,6 +22,7 @@ class Components(BaseModel): # Denoiser transformer: Any = None + transformer_2: Any = None unet: Any = None # Scheduler @@ -46,6 +47,7 @@ class Wan_Components(BaseModel): # Denoiser transformer: Any = None + transformer_2: Any = None unet: Any = None # Scheduler @@ -56,4 +58,7 @@ class Wan_Components(BaseModel): # Image Processor image_processor: Any = None + + # Boundary ratio + boundary_ratio: float | None = None \ No newline at end of file diff --git a/core/finetune/schemas/state.py b/core/finetune/schemas/state.py index 5b054c5..3e0b383 100644 --- a/core/finetune/schemas/state.py +++ b/core/finetune/schemas/state.py @@ -13,8 +13,8 @@ class State(BaseModel): train_width: int transformer_config: Dict[str, Any] = None - - weight_dtype: torch.dtype = torch.float32 # dtype for mixed precision training + transformer_2_config: Dict[str, Any] = None + weight_dtype: torch.dtype = torch.bfloat16 # dtype for mixed precision training num_trainable_parameters: int = 0 overwrote_max_train_steps: bool = False num_update_steps_per_epoch: int = 0 From 7814ca912006ebab4bec7bcd3133b911d0cc88c0 Mon Sep 17 00:00:00 2001 From: linjunfan Date: Fri, 14 Aug 2026 17:37:15 +0800 Subject: [PATCH 4/4] modify --- core/finetune/schemas/state.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/finetune/schemas/state.py b/core/finetune/schemas/state.py index 3e0b383..f64600c 100644 --- a/core/finetune/schemas/state.py +++ b/core/finetune/schemas/state.py @@ -14,7 +14,7 @@ class State(BaseModel): transformer_config: Dict[str, Any] = None transformer_2_config: Dict[str, Any] = None - weight_dtype: torch.dtype = torch.bfloat16 # dtype for mixed precision training + weight_dtype: torch.dtype = torch.bfloat32 # dtype for mixed precision training num_trainable_parameters: int = 0 overwrote_max_train_steps: bool = False num_update_steps_per_epoch: int = 0