cleanup

Author: Ednaordinary

src/diffusers/__init__.py

@@ -923,8 +923,8 @@
             BriaFiboTransformer2DModel,
             BriaTransformer2DModel,
             CacheMixin,
-            ChromaTransformer2DModel,
             ChromaRadianceTransformer2DModel,
+            ChromaTransformer2DModel,
             ChronoEditTransformer3DModel,
             CogVideoXTransformer3DModel,
             CogView3PlusTransformer2DModel,
@@ -1133,8 +1133,8 @@
             BriaFiboPipeline,
             BriaPipeline,
             ChromaImg2ImgPipeline,
-            ChromaRadiancePipeline,
             ChromaPipeline,
+            ChromaRadiancePipeline,
             ChronoEditPipeline,
             CLIPImageProjection,
             CogVideoXFunControlPipeline,

src/diffusers/loaders/single_file_model.py

@@ -31,8 +31,8 @@
     convert_animatediff_checkpoint_to_diffusers,
     convert_auraflow_transformer_checkpoint_to_diffusers,
     convert_autoencoder_dc_checkpoint_to_diffusers,
-    convert_chroma_transformer_checkpoint_to_diffusers,
     convert_chroma_radiance_transformer_checkpoint_to_diffusers,
+    convert_chroma_transformer_checkpoint_to_diffusers,
     convert_controlnet_checkpoint,
     convert_cosmos_transformer_checkpoint_to_diffusers,
     convert_flux2_transformer_checkpoint_to_diffusers,

src/diffusers/loaders/single_file_utils.py

@@ -663,9 +663,7 @@ def infer_diffusers_model_type(checkpoint):
         model_type = "flux-2-dev"
 
     elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux"]):
-        if any(
-            c in checkpoint for c in ["distilled_guidance_layer.in_proj.bias"]
-        ):
+        if any(c in checkpoint for c in ["distilled_guidance_layer.in_proj.bias"]):
             # Should be updated once a repo exists
             # if any(h in checkpoint for h in ["nerf_blocks.0.param_generator.bias"]):
             #     model_type = "chroma-radiance"
@@ -3556,6 +3554,7 @@ def swap_scale_shift(weight):
 
     return converted_state_dict
 
+
 def convert_chroma_radiance_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
     converted_state_dict = {}
     keys = list(checkpoint.keys())
@@ -3715,30 +3714,20 @@ def swap_scale_shift(weight):
         # output projections.
         converted_state_dict[f"{block_prefix}proj_out.weight"] = checkpoint.pop(f"single_blocks.{i}.linear2.weight")
         converted_state_dict[f"{block_prefix}proj_out.bias"] = checkpoint.pop(f"single_blocks.{i}.linear2.bias")
-    
+
     # nerf
-    
-    converted_state_dict["nerf.nerf_embedder.embedder.0.bias"] = checkpoint.pop(
-        "nerf_image_embedder.embedder.0.bias"
-    )
+
+    converted_state_dict["nerf.nerf_embedder.embedder.0.bias"] = checkpoint.pop("nerf_image_embedder.embedder.0.bias")
     converted_state_dict["nerf.nerf_embedder.embedder.0.weight"] = checkpoint.pop(
         "nerf_image_embedder.embedder.0.weight"
     )
-    converted_state_dict["nerf.final_layer.conv.bias"] = checkpoint.pop(
-        "nerf_final_layer_conv.conv.bias"
-    )
-    converted_state_dict["nerf.final_layer.conv.weight"] = checkpoint.pop(
-        "nerf_final_layer_conv.conv.weight"
-    )
-    converted_state_dict["nerf.final_layer.norm.weight"] = checkpoint.pop(
-        "nerf_final_layer_conv.norm.scale"
-    )
+    converted_state_dict["nerf.final_layer.conv.bias"] = checkpoint.pop("nerf_final_layer_conv.conv.bias")
+    converted_state_dict["nerf.final_layer.conv.weight"] = checkpoint.pop("nerf_final_layer_conv.conv.weight")
+    converted_state_dict["nerf.final_layer.norm.weight"] = checkpoint.pop("nerf_final_layer_conv.norm.scale")
 
     for i in range(num_nerf_layers):
         block_prefix = f"nerf.blocks.{i}."
-        converted_state_dict[f"{block_prefix}norm.weight"] = checkpoint.pop(
-            f"nerf_blocks.{i}.norm.scale"
-        )
+        converted_state_dict[f"{block_prefix}norm.weight"] = checkpoint.pop(f"nerf_blocks.{i}.norm.scale")
         converted_state_dict[f"{block_prefix}param_generator.bias"] = checkpoint.pop(
             f"nerf_blocks.{i}.param_generator.bias"
         )
@@ -3747,16 +3736,13 @@ def swap_scale_shift(weight):
         )
 
     # patch
-    
-    converted_state_dict["x_embedder_patch.bias"] = checkpoint.pop(
-        "img_in_patch.bias"
-    )
-    converted_state_dict["x_embedder_patch.weight"] = checkpoint.pop(
-        "img_in_patch.weight"
-    )
+
+    converted_state_dict["x_embedder_patch.bias"] = checkpoint.pop("img_in_patch.bias")
+    converted_state_dict["x_embedder_patch.weight"] = checkpoint.pop("img_in_patch.weight")
 
     return converted_state_dict
 
+
 def convert_cosmos_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
     converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())}
 

src/diffusers/models/__init__.py

@@ -86,7 +86,10 @@
     _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"]
     _import_structure["transformers.transformer_bria"] = ["BriaTransformer2DModel"]
     _import_structure["transformers.transformer_bria_fibo"] = ["BriaFiboTransformer2DModel"]
-    _import_structure["transformers.transformer_chroma"] = ["ChromaTransformer2DModel", "ChromaRadianceTransformer2DModel"]
+    _import_structure["transformers.transformer_chroma"] = [
+        "ChromaTransformer2DModel",
+        "ChromaRadianceTransformer2DModel",
+    ]
     _import_structure["transformers.transformer_chronoedit"] = ["ChronoEditTransformer3DModel"]
     _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"]
     _import_structure["transformers.transformer_cogview4"] = ["CogView4Transformer2DModel"]
@@ -184,8 +187,8 @@
             AuraFlowTransformer2DModel,
             BriaFiboTransformer2DModel,
             BriaTransformer2DModel,
-            ChromaTransformer2DModel,
             ChromaRadianceTransformer2DModel,
+            ChromaTransformer2DModel,
             ChronoEditTransformer3DModel,
             CogVideoXTransformer3DModel,
             CogView3PlusTransformer2DModel,

src/diffusers/models/transformers/transformer_chroma.py

@@ -35,6 +35,7 @@
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
+
 class Nerf(nn.Module):
     def __init__(
         self,
@@ -44,7 +45,7 @@ def __init__(
         transformer_hidden_size: int,
         max_freqs: int,
         mlp_ratio: int,
-        eps = 1e-6,
+        eps=1e-6,
     ):
         super().__init__()
         self.nerf_embedder = NerfEmbedder(
@@ -69,6 +70,7 @@ def __init__(
             eps=eps,
         )
         self.transformer_hidden_size = transformer_hidden_size
+
     def __call__(
         self,
         pixels,
@@ -77,20 +79,20 @@ def __call__(
         num_patches,
     ):
         batch_size, channels, height, width = pixels.shape
-        
+
         pixels = nn.functional.unfold(pixels, kernel_size=patch_size, stride=patch_size)
         pixels = pixels.transpose(1, 2)
-        
+
         hidden = latents.reshape(batch_size * num_patches, self.transformer_hidden_size)
         pixels = pixels.reshape(batch_size * num_patches, channels, patch_size**2).transpose(1, 2)
-        
+
         # Get pixel embeddings
         latents_dct = self.nerf_embedder(pixels)
-        
+
         # Pass through blocks
         for block in self.blocks:
             latents_dct = block(latents_dct, hidden)
-        
+
         latents_dct = latents_dct.transpose(1, 2).reshape(batch_size, num_patches, -1).transpose(1, 2)
         latents_dct = nn.functional.fold(
             latents_dct,
@@ -100,6 +102,7 @@ def __call__(
         )
         return self.final_layer(latents_dct)
 
+
 class NerfEmbedder(nn.Module):
     def __init__(
         self,
@@ -111,6 +114,7 @@ def __init__(
         self.max_freqs = max_freqs
         self.hidden_size = hidden_size
         self.embedder = nn.Sequential(nn.Linear(in_channels + max_freqs**2, hidden_size))
+
     def fetch_pos(self, patch_size) -> torch.Tensor:
         pos_x = torch.linspace(0, 1, patch_size)
         pos_y = torch.linspace(0, 1, patch_size)
@@ -123,8 +127,9 @@ def fetch_pos(self, patch_size) -> torch.Tensor:
         coeffs = (1 + freqs_x * freqs_y) ** -1
         dct_x = torch.cos(pos_x * freqs_x * torch.pi)
         dct_y = torch.cos(pos_y * freqs_y * torch.pi)
-        dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+        dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2)
         return dct
+
     def __call__(self, inputs: torch.Tensor) -> torch.Tensor:
         batch, pixels, channels = inputs.shape
         patch_size = int(pixels**0.5)
@@ -134,13 +139,15 @@ def __call__(self, inputs: torch.Tensor) -> torch.Tensor:
         inputs = torch.cat((inputs, dct), dim=-1)
         return self.embedder(inputs)
 
+
 class NerfGLUBlock(nn.Module):
     def __init__(self, transformer_hidden_size: int, nerf_hidden_size: int, mlp_ratio, eps):
         super().__init__()
         total_params = 3 * nerf_hidden_size**2 * mlp_ratio
         self.param_generator = nn.Linear(transformer_hidden_size, total_params)
         self.norm = RMSNorm(nerf_hidden_size, eps=eps)
         self.mlp_ratio = mlp_ratio
+
     def forward(self, x: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
         batch_size, num_x, hidden_size_x = x.shape
         mlp_params = self.param_generator(s)
@@ -156,6 +163,7 @@ def forward(self, x: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
         x = torch.bmm(torch.nn.functional.silu(torch.bmm(x, fc1_gate)) * torch.bmm(x, fc1_value), fc2)
         return x + res_x
 
+
 class NerfFinalLayer(nn.Module):
     def __init__(self, hidden_size: int, out_channels: int, eps):
         super().__init__()
@@ -166,9 +174,11 @@ def __init__(self, hidden_size: int, out_channels: int, eps):
             kernel_size=3,
             padding=1,
         )
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.conv(self.norm(x.movedim(1, -1)).movedim(-1, 1))
 
+
 class ChromaAdaLayerNormZeroPruned(nn.Module):
     r"""
     Norm layer adaptive layer norm zero (adaLN-Zero).
@@ -658,7 +668,7 @@ def forward(
                 logger.warning(
                     "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
                 )
-        
+
         hidden_states = self.x_embedder(hidden_states)
 
         timestep = timestep.to(hidden_states.dtype) * 1000
@@ -773,6 +783,7 @@ def forward(
 
         return Transformer2DModelOutput(sample=output)
 
+
 class ChromaRadianceTransformer2DModel(
     ModelMixin,
     ConfigMixin,
@@ -850,7 +861,7 @@ def __init__(
             hidden_dim=approximator_hidden_dim,
             n_layers=approximator_layers,
         )
-        
+
         self.nerf = Nerf(
             in_channels,
             nerf_layers,
@@ -859,7 +870,7 @@ def __init__(
             nerf_max_freqs,
             nerf_mlp_ratio,
         )
-        
+
         self.x_embedder_patch = nn.Conv2d(
             in_channels,
             self.inner_dim,
@@ -932,7 +943,6 @@ def forward(
             If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
             `tuple` where the first element is the sample tensor.
         """
-        print(self.device)
         pixels = hidden_states.to(self.device)
         if joint_attention_kwargs is not None:
             joint_attention_kwargs = joint_attention_kwargs.copy()

src/diffusers/models/transformers/transformer_flux.py

@@ -115,8 +115,6 @@ def __call__(
             query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
             key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
 
-        print("query", query.shape, "key", key.shape, "value", value.shape)
-
         hidden_states = dispatch_attention_fn(
             query,
             key,

src/diffusers/pipelines/chroma/__init__.py

@@ -33,8 +33,8 @@
         from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
     else:
         from .pipeline_chroma import ChromaPipeline
-        from .pipeline_chroma_radiance import ChromaRadiancePipeline
         from .pipeline_chroma_img2img import ChromaImg2ImgPipeline
+        from .pipeline_chroma_radiance import ChromaRadiancePipeline
 else:
     import sys
 

src/diffusers/pipelines/chroma/pipeline_chroma_radiance.py

@@ -21,11 +21,10 @@
 
 from ...image_processor import PipelineImageInput, VaeImageProcessor
 from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
-from ...models import AutoencoderKL, ChromaTransformer2DModel, ChromaRadianceTransformer2DModel
+from ...models import ChromaRadianceTransformer2DModel
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import (
     USE_PEFT_BACKEND,
-    deprecate,
     is_torch_xla_available,
     logging,
     replace_example_docstring,
@@ -147,6 +146,7 @@ def retrieve_timesteps(
         timesteps = scheduler.timesteps
     return timesteps, num_inference_steps
 
+
 class ChromaRadiancePipeline(
     DiffusionPipeline,
     FluxLoraLoaderMixin,
@@ -420,7 +420,6 @@ def check_inputs(
         callback_on_step_end_tensor_inputs=None,
         max_sequence_length=None,
     ):
-
         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
         ):
@@ -503,9 +502,7 @@ def prepare_latents(
         latents=None,
         patch_size=2,
     ):
-
         shape = (batch_size, num_channels_latents, height, width)
-        print(shape)
 
         if latents is not None:
             latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
@@ -518,9 +515,11 @@ def prepare_latents(
             )
 
         latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        #latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        # latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
 
-        latent_image_ids = self._prepare_latent_image_ids(batch_size, height // patch_size, width // patch_size, device, dtype)
+        latent_image_ids = self._prepare_latent_image_ids(
+            batch_size, height // patch_size, width // patch_size, device, dtype
+        )
 
         return latents, latent_image_ids
 
@@ -822,7 +821,6 @@ def __call__(
                 batch_size * num_images_per_prompt,
             )
 
-
         # 6. Denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -887,12 +885,12 @@ def __call__(
 
         self._current_timestep = None
 
-        # 7. 
+        # 7.
 
         if output_type == "latent":
             image = latents
         else:
-            #image = self._unpack_latents(image, height, width)
+            # image = self._unpack_latents(image, height, width)
             image = self.image_processor.postprocess(latents, output_type=output_type)
 
         # Offload all models

src/diffusers/utils/dummy_pt_objects.py

@@ -662,6 +662,7 @@ def from_config(cls, *args, **kwargs):
     def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch"])
 
+
 class ChromaRadianceTransformer2DModel(metaclass=DummyObject):
     _backends = ["torch"]
 
@@ -676,6 +677,7 @@ def from_config(cls, *args, **kwargs):
     def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch"])
 
+
 class ChronoEditTransformer3DModel(metaclass=DummyObject):
     _backends = ["torch"]