style and quality pass

Author: galbria

src/diffusers/models/transformers/transformer_bria.py

@@ -3,7 +3,9 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from packaging import version
 
+from diffusers import __version__ as diffusers_version
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
 from diffusers.models.embeddings import TimestepEmbedding, get_timestep_embedding
@@ -13,9 +15,6 @@
 from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
 from diffusers.pipelines.bria.bria_utils import FluxPosEmbed as EmbedND
 from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
-from diffusers import __version__ as diffusers_version
-from packaging import version
-
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -246,7 +245,7 @@ def __init__(
     # def _set_gradient_checkpointing(self, module, enable=False):
     #     if hasattr(module, "gradient_checkpointing"):
     #         module.gradient_checkpointing = enable
-    
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -324,14 +323,13 @@ def forward(
             temb += self.guidance_embed(guidance, dtype=hidden_states.dtype)
 
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
-        
+
         if len(txt_ids.shape) == 3:
             txt_ids = txt_ids[0]
 
         if len(img_ids.shape) == 3:
             img_ids = img_ids[0]
 
-
         ids = torch.cat((txt_ids, img_ids), dim=0)
         image_rotary_emb = self.pos_embed(ids)
 
@@ -408,9 +406,9 @@ def custom_forward(*inputs):
             else:
                 if version.parse(diffusers_version) < version.parse("0.35.0.dev0"):
                     hidden_states = block(
-                    hidden_states=hidden_states,
-                    temb=temb,
-                    image_rotary_emb=image_rotary_emb,
+                        hidden_states=hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
                     )
                 else:
                     encoder_hidden_states, hidden_states = block(

src/diffusers/pipelines/bria/pipeline_bria.py

@@ -9,7 +9,6 @@
     T5TokenizerFast,
 )
 
-import diffusers
 from diffusers import AutoencoderKL, DDIMScheduler, EulerAncestralDiscreteScheduler
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.loaders import FluxLoraLoaderMixin
@@ -83,6 +82,7 @@ class BriaPipeline(FluxPipeline):
             Tokenizer of class
             [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
     """
+
     model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae"
     _optional_components = ["image_encoder", "feature_extractor"]
     _callback_tensor_inputs = ["latents", "prompt_embeds"]
@@ -455,8 +455,6 @@ def __call__(
             latent_image_ids = latent_image_ids[0]
         if len(text_ids.shape) == 3:
             text_ids = text_ids[0]
-            
-
 
         # 6. Denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:

tests/models/transformers/test_models_transformer_bria.py

@@ -120,7 +120,6 @@ def output_shape(self):
 
     def prepare_init_args_and_inputs_for_common(self):
         init_dict = {
-            
             "patch_size": 1,
             "in_channels": 4,
             "num_layers": 1,
@@ -130,7 +129,6 @@ def prepare_init_args_and_inputs_for_common(self):
             "joint_attention_dim": 32,
             "pooled_projection_dim": None,
             "axes_dims_rope": [0, 4, 4],
-        
         }
 
         inputs_dict = self.dummy_input
@@ -163,8 +161,7 @@ def test_deprecated_inputs_img_txt_ids_3d(self):
             torch.allclose(output_1, output_2, atol=1e-5),
             msg="output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) are not equal as them as 2d inputs",
         )
-        
-        
+
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"BriaTransformer2DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

tests/pipelines/bria/test_pipeline_bria.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import gc
-import unittest
 import tempfile
+import unittest
 
 import numpy as np
 import torch
@@ -32,18 +32,17 @@
     enable_full_determinism,
     nightly,
     numpy_cosine_similarity_distance,
-    require_torch_gpu,
     require_accelerator,
+    require_torch_gpu,
     slow,
     torch_device,
 )
 
-
 # from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist
-from tests.pipelines.test_pipelines_common import PipelineTesterMixin,to_np
+from tests.pipelines.test_pipelines_common import PipelineTesterMixin, to_np
 
-enable_full_determinism()
 
+enable_full_determinism()
 
 
 class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
@@ -56,7 +55,7 @@ class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
     test_xformers_attention = False
     test_layerwise_casting = True
     test_group_offloading = True
-    
+
     def get_dummy_components(self):
         torch.manual_seed(0)
         transformer = BriaTransformer2DModel(
@@ -123,8 +122,10 @@ def get_dummy_inputs(self, device, seed=0):
             "output_type": "np",
         }
         return inputs
+
     def test_encode_prompt_works_in_isolation(self):
         pass
+
     def test_bria_different_prompts(self):
         pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
         inputs = self.get_dummy_inputs(torch_device)
@@ -135,9 +136,6 @@ def test_bria_different_prompts(self):
         max_diff = np.abs(output_same_prompt - output_different_prompts).max()
         assert max_diff > 1e-6
 
-    
-
-    
     def test_image_output_shape(self):
         pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
         inputs = self.get_dummy_inputs(torch_device)
@@ -194,11 +192,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2):
         self.assertLess(
             max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading."
         )
-        
-        
-
 
-    
     def test_bria_image_output_shape(self):
         pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
         inputs = self.get_dummy_inputs(torch_device)
@@ -217,9 +211,9 @@ def test_to_dtype(self):
         components = self.get_dummy_components()
         pipe = self.pipeline_class(**components)
         pipe.set_progress_bar_config(disable=None)
-    
+
         model_dtypes = [component.dtype for component in components.values() if hasattr(component, "dtype")]
-        self.assertTrue([dtype == torch.float32 for dtype in model_dtypes] == [False,True,True])
+        self.assertTrue([dtype == torch.float32 for dtype in model_dtypes] == [False, True, True])
 
     def test_torch_dtype_dict(self):
         components = self.get_dummy_components()
@@ -243,8 +237,6 @@ def test_torch_dtype_dict(self):
             self.assertEqual(loaded_pipe.text_encoder.dtype, torch.float16)
             self.assertEqual(loaded_pipe.vae.dtype, torch.float16)
 
-    
-
 
 @slow
 @require_torch_gpu
@@ -265,9 +257,7 @@ def tearDown(self):
     def get_inputs(self, device, seed=0):
         generator = torch.Generator(device="cpu").manual_seed(seed)
         prompt_embeds = torch.load(
-            hf_hub_download(
-                repo_id="diffusers/test-slices", repo_type="dataset", filename="bria_prompt_embeds.pt"
-            )
+            hf_hub_download(repo_id="diffusers/test-slices", repo_type="dataset", filename="bria_prompt_embeds.pt")
         ).to(device)
         return {
             "prompt_embeds": prompt_embeds,
@@ -324,7 +314,7 @@ def test_bria_inference_bf16(self):
         )
         max_diff = numpy_cosine_similarity_distance(expected_slice, image_slice)
         self.assertLess(max_diff, 1e-4, f"Image slice is different from expected slice: {max_diff:.4f}")
-    
+
     def test_to_dtype(self):
         components = self.get_dummy_components()
         pipe = self.pipeline_class(**components)
@@ -396,4 +386,3 @@ def test_bria_inference(self):
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, image_slice)
         self.assertLess(max_diff, 1e-4, f"Image slice is different from expected slice: {max_diff:.4f}")
-