make style

Author: a-r-r-o-w

scripts/convert_cogvideox_to_diffusers.py

@@ -140,7 +140,7 @@ def convert_transformer(
     use_rotary_positional_embeddings: bool,
     i2v: bool,
     dtype: torch.dtype,
-    init_kwargs: Dict[str, Any]
+    init_kwargs: Dict[str, Any],
 ):
     PREFIX_KEY = "model.diffusion_model."
 
@@ -165,7 +165,7 @@ def convert_transformer(
             if special_key not in key:
                 continue
             handler_fn_inplace(key, original_state_dict)
-    
+
     transformer.load_state_dict(original_state_dict, strict=True)
     return transformer
 
@@ -201,7 +201,7 @@ def get_init_kwargs(version: str):
             "sample_width": 720 // vae_scale_factor_spatial,
             "sample_frames": 49,
         }
-    
+
     elif version == "1.5":
         vae_scale_factor_spatial = 8
         init_kwargs = {
@@ -214,7 +214,7 @@ def get_init_kwargs(version: str):
         }
     else:
         raise ValueError("Unsupported version of CogVideoX.")
-    
+
     return init_kwargs
 
 
@@ -245,8 +245,18 @@ def get_args():
     parser.add_argument("--scaling_factor", type=float, default=1.15258426, help="Scaling factor in the VAE")
     # For CogVideoX-2B, snr_shift_scale is 3.0. For 5B, it is 1.0
     parser.add_argument("--snr_shift_scale", type=float, default=3.0, help="Scaling factor in the VAE")
-    parser.add_argument("--i2v", action="store_true", default=False, help="Whether the model to be converted is the Image-to-Video version of CogVideoX.")
-    parser.add_argument("--version", choices=["1.0", "1.5"], default="1.0", help="Which version of CogVideoX to use for initializing default modeling parameters.")
+    parser.add_argument(
+        "--i2v",
+        action="store_true",
+        default=False,
+        help="Whether the model to be converted is the Image-to-Video version of CogVideoX.",
+    )
+    parser.add_argument(
+        "--version",
+        choices=["1.0", "1.5"],
+        default="1.0",
+        help="Which version of CogVideoX to use for initializing default modeling parameters.",
+    )
     return parser.parse_args()
 
 

src/diffusers/models/embeddings.py

@@ -17,7 +17,6 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-from einops import rearrange
 from torch import nn
 
 from ..utils import deprecate
@@ -377,7 +376,7 @@ def __init__(
         else:
             # CogVideoX 1.5 checkpoints
             self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim)
-        
+
         self.text_proj = nn.Linear(text_embed_dim, embed_dim)
 
         if use_positional_embeddings or use_learned_positional_embeddings:
@@ -429,7 +428,9 @@ def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):
             p_t = self.patch_size_t
 
             image_embeds = image_embeds.permute(0, 1, 3, 4, 2)
-            image_embeds = image_embeds.reshape(batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels)
+            image_embeds = image_embeds.reshape(
+                batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels
+            )
             image_embeds = image_embeds.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(4, 7).flatten(1, 3)
             image_embeds = self.proj(image_embeds)
 

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -308,7 +308,7 @@ def __init__(
         else:
             # For CogVideoX 1.5
             output_dim = patch_size * patch_size * patch_size_t * out_channels
-        
+
         self.proj_out = nn.Linear(inner_dim, output_dim)
 
         self.gradient_checkpointing = False
@@ -516,7 +516,9 @@ def custom_forward(*inputs):
             output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
             output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
         else:
-            output = hidden_states.reshape(batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p)
+            output = hidden_states.reshape(
+                batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p
+            )
             output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
             output = output[:, remaining_frames:]
 

src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py

@@ -449,7 +449,7 @@ def _prepare_rotary_positional_embeddings(
         base_size_width = self.transformer.config.sample_width // p
         base_size_height = self.transformer.config.sample_height // p
         base_num_frames = (num_frames + p_t - 1) // p_t
-        
+
         grid_crops_coords = get_resize_crop_region_for_grid(
             (grid_height, grid_width), base_size_width, base_size_height
         )