fix code & mask style & fix-copies

Author: junqiangwu

docs/source/en/_toctree.yml

@@ -373,6 +373,8 @@
         title: LuminaNextDiT2DModel
       - local: api/models/mochi_transformer3d
         title: MochiTransformer3DModel
+      - local: api/models/longcat_image_transformer2d
+        title: LongCatImageTransformer2DModel
       - local: api/models/omnigen_transformer
         title: OmniGenTransformer2DModel
       - local: api/models/ovisimage_transformer2d
@@ -567,6 +569,8 @@
         title: Lumina 2.0
       - local: api/pipelines/lumina
         title: Lumina-T2X
+      - local: api/pipelines/longcat_image
+        title: LongCat-Image
       - local: api/pipelines/marigold
         title: Marigold
       - local: api/pipelines/panorama

src/diffusers/__init__.py

@@ -973,10 +973,10 @@
             Kandinsky3UNet,
             Kandinsky5Transformer3DModel,
             LatteTransformer3DModel,
+            LongCatImageTransformer2DModel,
             LTXVideoTransformer3DModel,
             Lumina2Transformer2DModel,
             LuminaNextDiT2DModel,
-            LongCatImageTransformer2DModel,
             MochiTransformer3DModel,
             ModelMixin,
             MotionAdapter,
@@ -1241,12 +1241,12 @@
             LDMTextToImagePipeline,
             LEditsPPPipelineStableDiffusion,
             LEditsPPPipelineStableDiffusionXL,
+            LongCatImageEditPipeline,
+            LongCatImagePipeline,
             LTXConditionPipeline,
             LTXImageToVideoPipeline,
             LTXLatentUpsamplePipeline,
             LTXPipeline,
-            LongCatImagePipeline,
-            LongCatImageEditPipeline,
             LucyEditPipeline,
             Lumina2Pipeline,
             Lumina2Text2ImgPipeline,

src/diffusers/models/__init__.py

@@ -209,10 +209,10 @@
             HunyuanVideoTransformer3DModel,
             Kandinsky5Transformer3DModel,
             LatteTransformer3DModel,
+            LongCatImageTransformer2DModel,
             LTXVideoTransformer3DModel,
             Lumina2Transformer2DModel,
             LuminaNextDiT2DModel,
-            LongCatImageTransformer2DModel,
             MochiTransformer3DModel,
             OmniGenTransformer2DModel,
             OvisImageTransformer2DModel,

src/diffusers/models/transformers/__init__.py

@@ -33,9 +33,9 @@
     from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel
     from .transformer_hunyuanimage import HunyuanImageTransformer2DModel
     from .transformer_kandinsky import Kandinsky5Transformer3DModel
+    from .transformer_longcat_image import LongCatImageTransformer2DModel
     from .transformer_ltx import LTXVideoTransformer3DModel
     from .transformer_lumina2 import Lumina2Transformer2DModel
-    from .transformer_longcat_image import LongCatImageTransformer2DModel
     from .transformer_mochi import MochiTransformer3DModel
     from .transformer_omnigen import OmniGenTransformer2DModel
     from .transformer_ovis_image import OvisImageTransformer2DModel
@@ -48,4 +48,4 @@
     from .transformer_wan import WanTransformer3DModel
     from .transformer_wan_animate import WanAnimateTransformer3DModel
     from .transformer_wan_vace import WanVACETransformer3DModel
-    from .transformer_z_image import ZImageTransformer2DModel
+    from .transformer_z_image import ZImageTransformer2DModel

src/diffusers/models/transformers/transformer_longcat_image.py

@@ -15,7 +15,6 @@
 import inspect
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -24,14 +23,14 @@
 from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
 from ...utils import is_torch_npu_available, logging
 from ...utils.torch_utils import maybe_allow_in_graph
-from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
-from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
+from ..attention import AttentionModuleMixin, FeedForward
 from ..attention_dispatch import dispatch_attention_fn
 from ..cache_utils import CacheMixin
+from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
-from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
+
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -381,7 +380,6 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
         return freqs_cos, freqs_sin
 
 
-
 class LongCatImageTimestepEmbeddings(nn.Module):
     def __init__(self, embedding_dim):
         super().__init__()
@@ -394,14 +392,15 @@ def forward(self, timestep, hidden_dtype):
         timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))  # (N, D)
 
         return timesteps_emb
-    
+
 
 class LongCatImageTransformer2DModel(
     ModelMixin,
     ConfigMixin,
     PeftAdapterMixin,
     FromOriginalModelMixin,
-    CacheMixin, ):
+    CacheMixin,
+):
     """
     The Transformer model introduced in Longcat-Image.
     """
@@ -455,10 +454,8 @@ def __init__(
             ]
         )
 
-        self.norm_out = AdaLayerNormContinuous(
-            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
-        self.proj_out = nn.Linear(
-            self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
 
         self.gradient_checkpointing = False
         self.use_checkpoint = [True] * num_layers
@@ -498,10 +495,9 @@ def forward(
 
         timestep = timestep.to(hidden_states.dtype) * 1000
 
-        temb = self.time_embed( timestep, hidden_states.dtype )
+        temb = self.time_embed(timestep, hidden_states.dtype)
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
-
         ids = torch.cat((txt_ids, img_ids), dim=0)
         if is_torch_npu_available():
             freqs_cos, freqs_sin = self.pos_embed(ids.cpu())
@@ -528,7 +524,7 @@ def forward(
 
         for index_block, block in enumerate(self.single_transformer_blocks):
             if torch.is_grad_enabled() and self.gradient_checkpointing and self.use_single_checkpoint[index_block]:
-                encoder_hidden_states,hidden_states = self._gradient_checkpointing_func(
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
                     block,
                     hidden_states,
                     encoder_hidden_states,

src/diffusers/pipelines/__init__.py

@@ -291,7 +291,7 @@
     _import_structure["lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"]
     _import_structure["lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"]
     _import_structure["lucy"] = ["LucyEditPipeline"]
-    _import_structure["longcat_image"] = ["LongCatImagePipeline","LongCatImageEditPipeline"]
+    _import_structure["longcat_image"] = ["LongCatImagePipeline", "LongCatImageEditPipeline"]
     _import_structure["marigold"].extend(
         [
             "MarigoldDepthPipeline",
@@ -719,11 +719,11 @@
             LEditsPPPipelineStableDiffusion,
             LEditsPPPipelineStableDiffusionXL,
         )
+        from .longcat_image import LongCatImageEditPipeline, LongCatImagePipeline
         from .ltx import LTXConditionPipeline, LTXImageToVideoPipeline, LTXLatentUpsamplePipeline, LTXPipeline
         from .lucy import LucyEditPipeline
         from .lumina import LuminaPipeline, LuminaText2ImgPipeline
         from .lumina2 import Lumina2Pipeline, Lumina2Text2ImgPipeline
-        from .longcat_image import LongCatImagePipeline,LongCatImageEditPipeline
         from .marigold import (
             MarigoldDepthPipeline,
             MarigoldIntrinsicsPipeline,

src/diffusers/pipelines/longcat_image/__init__.py

@@ -33,9 +33,9 @@
     except OptionalDependencyNotAvailable:
         from ...utils.dummy_torch_and_transformers_objects import *
     else:
-        from .pipeline_output import LongCatImagePipelineOutput
         from .pipeline_longcat_image import LongCatImagePipeline
         from .pipeline_longcat_image_edit import LongCatImageEditPipeline
+        from .pipeline_output import LongCatImagePipelineOutput
 
 else:
     import sys