add 5b t2v

Author: yiyixuxu

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -236,7 +236,7 @@ def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None:
         super().__init__()
         self.dim = dim
         self.mode = mode
-        
+
         # default to dim //2
         if upsample_out_dim is None:
             upsample_out_dim = dim // 2
@@ -524,7 +524,7 @@ class WanEncoder3d(nn.Module):
 
     def __init__(
         self,
-        in_channels: int = 3,   
+        in_channels: int = 3,
         dim=128,
         z_dim=4,
         dim_mult=[1, 2, 4, 4],
@@ -558,10 +558,10 @@ def __init__(
             if is_residual:
                 self.down_blocks.append(
                     WanResidualDownBlock(
-                        in_dim, 
-                        out_dim, 
-                        dropout, 
-                        num_res_blocks, 
+                        in_dim,
+                        out_dim,
+                        dropout,
+                        num_res_blocks,
                         temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False,
                         down_flag=i != len(dim_mult) - 1,
                         )
@@ -708,10 +708,10 @@ def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
                 x = self.upsampler(x, feat_cache, feat_idx)
             else:
                 x = self.upsampler(x)
-        
+
         if self.avg_shortcut is not None:
             x = x + self.avg_shortcut(x_copy, first_chunk=first_chunk)
-        
+
         return x
 
 class WanUpBlock(nn.Module):
@@ -912,10 +912,9 @@ def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
         return x
 
 
-# YiYi TODO: refactor this
-from einops import rearrange
-
 def patchify(x, patch_size):
+    # YiYi TODO: refactor this
+    from einops import rearrange
     if patch_size == 1:
         return x
     if x.dim() == 4:
@@ -935,6 +934,8 @@ def patchify(x, patch_size):
 
 
 def unpatchify(x, patch_size):
+    # YiYi TODO: refactor this
+    from einops import rearrange
     if patch_size == 1:
         return x
 

src/diffusers/models/transformers/transformer_wan.py

@@ -171,7 +171,7 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: Optional[torch.Tensor] = None,
         timestep_seq_len: Optional[int] = None,
-    ):  
+    ):
         timestep = self.timesteps_proj(timestep)
         if timestep_seq_len is not None:
             timestep = timestep.unflatten(0, (1, timestep_seq_len))
@@ -518,7 +518,7 @@ def forward(
                 hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
 
         # 5. Output norm, projection & unpatchify
-        if temb.ndim ==3: 
+        if temb.ndim ==3:
             # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
             shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
             shift = shift.squeeze(2)

src/diffusers/pipelines/wan/pipeline_wan.py

@@ -318,7 +318,7 @@ def check_inputs(
             not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
         ):
             raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
-        
+
         if self.config.boundary_ratio is None and guidance_scale_2 is not None:
             raise ValueError("`guidance_scale_2` is only supported when the pipeline's `boundary_ratio` is not None.")