make style

Author: staoxiao

src/diffusers/models/transformers/transformer_omnigen.py

@@ -16,9 +16,9 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
+import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
-import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import PeftAdapterMixin
@@ -91,7 +91,11 @@ def __init__(
         self.pos_embed_max_size = pos_embed_max_size
 
         pos_embed = get_2d_sincos_pos_embed(
-            embed_dim, self.pos_embed_max_size, base_size=base_size, interpolation_scale=self.interpolation_scale, output_type="pt"
+            embed_dim,
+            self.pos_embed_max_size,
+            base_size=base_size,
+            interpolation_scale=self.interpolation_scale,
+            output_type="pt",
         )
         self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=True)
 
@@ -227,7 +231,7 @@ def apply_rotary_emb(
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
     """
-    
+
     cos, sin = freqs_cis  # [S, D]
     if len(cos.shape) == 2:
         cos = cos[None, None]
@@ -241,10 +245,10 @@ def apply_rotary_emb(
     x1 = x[..., : x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2 :]
     x_rotated = torch.cat((-x2, x1), dim=-1)
-       
+
     out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
     return out
-    
+
 
 class OmniGenAttnProcessor2_0:
     r"""
@@ -264,7 +268,6 @@ def __call__(
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-
         batch_size, sequence_length, _ = hidden_states.shape
 
         # Get Query-Key-Value Pair
@@ -674,9 +677,13 @@ def forward(
         image_rotary_emb = self.rotary_emb(hidden_states, position_ids)
         for decoder_layer in self.layers:
             if torch.is_grad_enabled() and self.gradient_checkpointing:
-                hidden_states = self._gradient_checkpointing_func(decoder_layer, hidden_states, attention_mask, image_rotary_emb)
+                hidden_states = self._gradient_checkpointing_func(
+                    decoder_layer, hidden_states, attention_mask, image_rotary_emb
+                )
             else:
-                hidden_states = decoder_layer(hidden_states, attention_mask=attention_mask, image_rotary_emb=image_rotary_emb)
+                hidden_states = decoder_layer(
+                    hidden_states, attention_mask=attention_mask, image_rotary_emb=image_rotary_emb
+                )
 
         hidden_states = self.norm(hidden_states)
 

src/diffusers/pipelines/consisid/pipeline_consisid.py

@@ -48,9 +48,14 @@
         >>> from huggingface_hub import snapshot_download
 
         >>> snapshot_download(repo_id="BestWishYsh/ConsisID-preview", local_dir="BestWishYsh/ConsisID-preview")
-        >>> face_helper_1, face_helper_2, face_clip_model, face_main_model, eva_transform_mean, eva_transform_std = (
-        ...     prepare_face_models("BestWishYsh/ConsisID-preview", device="cuda", dtype=torch.bfloat16)
-        ... )
+        >>> (
+        ...     face_helper_1,
+        ...     face_helper_2,
+        ...     face_clip_model,
+        ...     face_main_model,
+        ...     eva_transform_mean,
+        ...     eva_transform_std,
+        ... ) = prepare_face_models("BestWishYsh/ConsisID-preview", device="cuda", dtype=torch.bfloat16)
         >>> pipe = ConsisIDPipeline.from_pretrained("BestWishYsh/ConsisID-preview", torch_dtype=torch.bfloat16)
         >>> pipe.to("cuda")
 

src/diffusers/pipelines/omnigen/pipeline_omnigen.py

@@ -16,7 +16,6 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import numpy as np
-import PIL
 import torch
 from transformers import LlamaTokenizer
 
@@ -223,9 +222,9 @@ def check_inputs(
         if use_input_image_size_as_output:
             if input_images is None or input_images[0] is None:
                 raise ValueError(
-                                f"`use_input_image_size_as_output` is set to True, but no input image was found. If you are performing a text-to-image task, please set it to False."
-                            )
-                            
+                    "`use_input_image_size_as_output` is set to True, but no input image was found. If you are performing a text-to-image task, please set it to False."
+                )
+
         if callback_on_step_end_tensor_inputs is not None and not all(
             k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
         ):

src/diffusers/pipelines/omnigen/processor_omnigen.py

@@ -62,7 +62,7 @@ def __init__(self, text_tokenizer, max_image_size: int = 1024):
         )
 
         self.collator = OmniGenCollator()
-    
+
     def reset_max_image_size(self, max_image_size):
         self.max_image_size = max_image_size
         self.image_transform = transforms.Compose(

tests/models/transformers/test_models_transformer_omnigen.py

@@ -25,6 +25,7 @@
 
 enable_full_determinism()
 
+
 class OmniGenTransformerTests(ModelTesterMixin, unittest.TestCase):
     model_class = OmniGenTransformer2DModel
     main_input_name = "hidden_states"
@@ -42,11 +43,11 @@ def dummy_input(self):
         timestep = torch.rand(size=(batch_size,), dtype=hidden_states.dtype).to(torch_device)
         input_ids = torch.randint(0, 10, (batch_size, sequence_length)).to(torch_device)
         input_img_latents = [torch.randn((1, num_channels, height, width)).to(torch_device)]
-        input_image_sizes = {0: [[0, 0+height*width//2//2]]}
+        input_image_sizes = {0: [[0, 0 + height * width // 2 // 2]]}
 
-        attn_seq_length = sequence_length + 1 + height*width//2//2
+        attn_seq_length = sequence_length + 1 + height * width // 2 // 2
         attention_mask = torch.ones((batch_size, attn_seq_length, attn_seq_length)).to(torch_device)
-        position_ids = torch.LongTensor([list(range(attn_seq_length))]*batch_size).to(torch_device)
+        position_ids = torch.LongTensor([list(range(attn_seq_length))] * batch_size).to(torch_device)
 
         return {
             "hidden_states": hidden_states,
@@ -77,12 +78,11 @@ def prepare_init_args_and_inputs_for_common(self):
             "vocab_size": 100,
             "in_channels": 4,
             "time_step_dim": 4,
-            "rope_scaling": {"long_factor": list(range(1, 3)), "short_factor": list(range(1, 3))}
+            "rope_scaling": {"long_factor": list(range(1, 3)), "short_factor": list(range(1, 3))},
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
 
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"OmniGenTransformer2DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
-