[Utils] Correct custom init sort (#4967)

* [Utils] Correct custom init sort

* [Utils] Correct custom init sort

* [Utils] Correct custom init sort

* add type checking

* fix custom init sort

* fix test

* fix tests

---------
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

src/diffusers/pipelines/repaint/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import _LazyModule
 
 
-_import_structure = {}
-_import_structure["pipeline_repaint"] = ["RePaintPipeline"]
+_import_structure = {"pipeline_repaint": ["RePaintPipeline"]}
 
-import sys
+if TYPE_CHECKING:
+    from .pipeline_repaint import RePaintPipeline
 
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )

src/diffusers/pipelines/score_sde_ve/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import _LazyModule
 
 
-_import_structure = {}
-_import_structure["pipeline_score_sde_ve"] = ["ScoreSdeVePipeline"]
+_import_structure = {"pipeline_score_sde_ve": ["ScoreSdeVePipeline"]}
 
-import sys
+if TYPE_CHECKING:
+    from .pipeline_score_sde_ve import ScoreSdeVePipeline
 
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )

src/diffusers/pipelines/semantic_stable_diffusion/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,9 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -18,21 +19,30 @@ except OptionalDependencyNotAvailable:
     from ...utils import dummy_torch_and_transformers_objects  # noqa F403
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
-
 else:
     _import_structure["pipeline_output"] = ["SemanticStableDiffusionPipelineOutput"]
     _import_structure["pipeline_semantic_stable_diffusion"] = ["SemanticStableDiffusionPipeline"]
 
 
-import sys
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
 
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *
+    else:
+        from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline
 
-sys.modules[__name__] = _LazyModule(
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/shap_e/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,9 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -33,15 +34,37 @@ else:
         "VoidNeRFModel",
     ]
 
-import sys
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *
+    else:
+        from .camera import create_pan_cameras
+        from .pipeline_shap_e import ShapEPipeline
+        from .pipeline_shap_e_img2img import ShapEImg2ImgPipeline
+        from .renderer import (
+            BoundingBoxVolume,
+            ImportanceRaySampler,
+            MLPNeRFModelOutput,
+            MLPNeRSTFModel,
+            ShapEParamsProjModel,
+            ShapERenderer,
+            StratifiedRaySampler,
+            VoidNeRFModel,
+        )
 
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/spectrogram_diffusion/__init__.py

 # flake8: noqa
+from typing import TYPE_CHECKING
 from ...utils import (
     _LazyModule,
     is_note_seq_available,
@@ -8,9 +9,8 @@ from ...utils import (
     get_objects_from_module,
 )
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -19,32 +19,56 @@ except OptionalDependencyNotAvailable:
     from ...utils import dummy_torch_and_transformers_objects  # noqa F403
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
-
 else:
-    _import_structure["notes_encoder"] = ["SpectrogramNotesEncoder"]
     _import_structure["continous_encoder"] = ["SpectrogramContEncoder"]
+    _import_structure["notes_encoder"] = ["SpectrogramNotesEncoder"]
     _import_structure["pipeline_spectrogram_diffusion"] = [
         "SpectrogramContEncoder",
         "SpectrogramDiffusionPipeline",
         "T5FilmDecoder",
     ]
-
 try:
     if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
         raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
-    from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
+    from ...utils import dummy_transformers_and_torch_and_note_seq_objects
+
+    _dummy_objects.update(get_objects_from_module(dummy_transformers_and_torch_and_note_seq_objects))
 else:
     _import_structure["midi_utils"] = ["MidiProcessor"]
 
-import sys
 
-sys.modules[__name__] = _LazyModule(
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *
+    else:
+        from .pipeline_spectrogram_diffusion import SpectrogramDiffusionPipeline
+        from .pipeline_spectrogram_diffusion import SpectrogramContEncoder
+        from .pipeline_spectrogram_diffusion import SpectrogramNotesEncoder
+        from .pipeline_spectrogram_diffusion import T5FilmDecoder
+
+    try:
+        if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *
+
+    else:
+        from .midi_utils import MidiProcessor
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/stable_diffusion/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -12,15 +14,12 @@ from ...utils import (
 )
 
 
-_import_structure = {}
-_additional_imports = {}
 _dummy_objects = {}
-
-_import_structure["pipeline_output"] = ["StableDiffusionPipelineOutput"]
+_additional_imports = {}
+_import_structure = {"pipeline_output": ["StableDiffusionPipelineOutput"]}
 
 if is_transformers_available() and is_flax_available():
     _import_structure["pipeline_output"].extend(["FlaxStableDiffusionPipelineOutput"])
-
 try:
     if not (is_transformers_available() and is_torch_available()):
         raise OptionalDependencyNotAvailable()
@@ -28,12 +27,14 @@ except OptionalDependencyNotAvailable:
     from ...utils import dummy_torch_and_transformers_objects  # noqa F403
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
-
 else:
+    _import_structure["clip_image_project_model"] = ["CLIPImageProjection"]
     _import_structure["pipeline_cycle_diffusion"] = ["CycleDiffusionPipeline"]
     _import_structure["pipeline_stable_diffusion"] = ["StableDiffusionPipeline"]
     _import_structure["pipeline_stable_diffusion_attend_and_excite"] = ["StableDiffusionAttendAndExcitePipeline"]
     _import_structure["pipeline_stable_diffusion_gligen"] = ["StableDiffusionGLIGENPipeline"]
+    _import_structure["pipeline_stable_diffusion_gligen"] = ["StableDiffusionGLIGENPipeline"]
+    _import_structure["pipeline_stable_diffusion_gligen_text_image"] = ["StableDiffusionGLIGENTextImagePipeline"]
     _import_structure["pipeline_stable_diffusion_img2img"] = ["StableDiffusionImg2ImgPipeline"]
     _import_structure["pipeline_stable_diffusion_inpaint"] = ["StableDiffusionInpaintPipeline"]
     _import_structure["pipeline_stable_diffusion_inpaint_legacy"] = ["StableDiffusionInpaintPipelineLegacy"]
@@ -49,10 +50,6 @@ else:
     _import_structure["pipeline_stable_unclip_img2img"] = ["StableUnCLIPImg2ImgPipeline"]
     _import_structure["safety_checker"] = ["StableDiffusionSafetyChecker"]
     _import_structure["stable_unclip_image_normalizer"] = ["StableUnCLIPImageNormalizer"]
-    _import_structure["pipeline_stable_diffusion_gligen_text_image"] = ["StableDiffusionGLIGENTextImagePipeline"]
-    _import_structure["pipeline_stable_diffusion_gligen"] = ["StableDiffusionGLIGENPipeline"]
-    _import_structure["clip_image_project_model"] = ["CLIPImageProjection"]
-
 try:
     if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
         raise OptionalDependencyNotAvailable()
@@ -62,8 +59,6 @@ except OptionalDependencyNotAvailable:
     _dummy_objects.update({"StableDiffusionImageVariationPipeline": StableDiffusionImageVariationPipeline})
 else:
     _import_structure["pipeline_stable_diffusion_image_variation"] = ["StableDiffusionImageVariationPipeline"]
-
-
 try:
     if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.26.0")):
         raise OptionalDependencyNotAvailable()
@@ -85,8 +80,6 @@ else:
     _import_structure["pipeline_stable_diffusion_depth2img"] = ["StableDiffusionDepth2ImgPipeline"]
     _import_structure["pipeline_stable_diffusion_diffedit"] = ["StableDiffusionDiffEditPipeline"]
     _import_structure["pipeline_stable_diffusion_pix2pix_zero"] = ["StableDiffusionPix2PixZeroPipeline"]
-
-
 try:
     if not (
         is_torch_available()
@@ -99,10 +92,8 @@ except OptionalDependencyNotAvailable:
     from ...utils import dummy_torch_and_transformers_and_k_diffusion_objects  # noqa F403
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_k_diffusion_objects))
-
 else:
     _import_structure["pipeline_stable_diffusion_k_diffusion"] = ["StableDiffusionKDiffusionPipeline"]
-
 try:
     if not (is_transformers_available() and is_onnx_available()):
         raise OptionalDependencyNotAvailable()
@@ -110,7 +101,6 @@ except OptionalDependencyNotAvailable:
     from ...utils import dummy_onnx_objects  # noqa F403
 
     _dummy_objects.update(get_objects_from_module(dummy_onnx_objects))
-
 else:
     _import_structure["pipeline_onnx_stable_diffusion"] = [
         "OnnxStableDiffusionPipeline",
@@ -125,23 +115,120 @@ if is_transformers_available() and is_flax_available():
     from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
 
     _additional_imports.update({"PNDMSchedulerState": PNDMSchedulerState})
-
     _import_structure["pipeline_flax_stable_diffusion"] = ["FlaxStableDiffusionPipeline"]
     _import_structure["pipeline_flax_stable_diffusion_img2img"] = ["FlaxStableDiffusionImg2ImgPipeline"]
     _import_structure["pipeline_flax_stable_diffusion_inpaint"] = ["FlaxStableDiffusionInpaintPipeline"]
     _import_structure["safety_checker_flax"] = ["FlaxStableDiffusionSafetyChecker"]
 
-import sys
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
 
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *
+
+    else:
+        from .clip_image_project_model import CLIPImageProjection
+        from .pipeline_stable_diffusion import (
+            StableDiffusionPipeline,
+            StableDiffusionPipelineOutput,
+            StableDiffusionSafetyChecker,
+        )
+        from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
+        from .pipeline_stable_diffusion_gligen import StableDiffusionGLIGENPipeline
+        from .pipeline_stable_diffusion_gligen_text_image import StableDiffusionGLIGENTextImagePipeline
+        from .pipeline_stable_diffusion_img2img import StableDiffusionImg2ImgPipeline
+        from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
+        from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
+        from .pipeline_stable_diffusion_instruct_pix2pix import StableDiffusionInstructPix2PixPipeline
+        from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
+        from .pipeline_stable_diffusion_ldm3d import StableDiffusionLDM3DPipeline
+        from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
+        from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
+        from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
+        from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
+        from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
+        from .pipeline_stable_unclip import StableUnCLIPPipeline
+        from .pipeline_stable_unclip_img2img import StableUnCLIPImg2ImgPipeline
+        from .safety_checker import StableDiffusionSafetyChecker
+        from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
+
+    try:
+        if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
+    else:
+        from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline
 
-sys.modules[__name__] = _LazyModule(
+    try:
+        if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.26.0")):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import (
+            StableDiffusionDepth2ImgPipeline,
+            StableDiffusionDiffEditPipeline,
+            StableDiffusionPix2PixZeroPipeline,
+        )
+    else:
+        from .pipeline_stable_diffusion_depth2img import StableDiffusionDepth2ImgPipeline
+        from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
+        from .pipeline_stable_diffusion_pix2pix_zero import StableDiffusionPix2PixZeroPipeline
+
+    try:
+        if not (
+            is_torch_available()
+            and is_transformers_available()
+            and is_k_diffusion_available()
+            and is_k_diffusion_version(">=", "0.0.12")
+        ):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *
+    else:
+        from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline
+
+    try:
+        if not (is_transformers_available() and is_onnx_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_onnx_objects import *
+    else:
+        from .pipeline_onnx_stable_diffusion import (
+            OnnxStableDiffusionImg2ImgPipeline,
+            OnnxStableDiffusionInpaintPipeline,
+            OnnxStableDiffusionInpaintPipelineLegacy,
+            OnnxStableDiffusionPipeline,
+            OnnxStableDiffusionUpscalePipeline,
+            StableDiffusionOnnxPipeline,
+        )
+
+    try:
+        if not (is_transformers_available() and is_flax_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_flax_objects import *
+    else:
+        from .pipeline_flax_stable_diffusion import (
+            FlaxStableDiffusionImg2ImgPipeline,
+            FlaxStableDiffusionInpaintPipeline,
+            FlaxStableDiffusionPipeline,
+            FlaxStableDiffusionSafetyChecker,
+        )
+        from .pipeline_output import FlaxStableDiffusionPipelineOutput
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)
-for name, value in _additional_imports.items():
+    for name, value in _additional_imports.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/stable_diffusion_safe/__init__.py

 from dataclasses import dataclass
 from enum import Enum
-from typing import List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
 import numpy as np
 import PIL
 from PIL import Image
 
-from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available
+from ...utils import (
+    BaseOutput,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    get_objects_from_module,
+    is_torch_available,
+    is_transformers_available,
+)
 
 
 @dataclass
@@ -41,36 +48,47 @@ class SafetyConfig(object):
     }
 
 
-@dataclass
-class StableDiffusionSafePipelineOutput(BaseOutput):
-    """
-    Output class for Safe Stable Diffusion pipelines.
-
-    Args:
-        images (`List[PIL.Image.Image]` or `np.ndarray`)
-            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
-            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
-        nsfw_content_detected (`List[bool]`)
-            List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
-            (nsfw) content, or `None` if safety checking could not be performed.
-        images (`List[PIL.Image.Image]` or `np.ndarray`)
-            List of denoised PIL images that were flagged by the safety checker any may contain "not-safe-for-work"
-            (nsfw) content, or `None` if no safety check was performed or no images were flagged.
-        applied_safety_concept (`str`)
-            The safety concept that was applied for safety guidance, or `None` if safety guidance was disabled
-    """
-
-    images: Union[List[PIL.Image.Image], np.ndarray]
-    nsfw_content_detected: Optional[List[bool]]
-    unsafe_images: Optional[Union[List[PIL.Image.Image], np.ndarray]]
-    applied_safety_concept: Optional[str]
+_dummy_objects = {}
+_additional_imports = {}
+_import_structure = {
+    "pipeline_output": ["StableDiffusionSafePipelineOutput"],
+    "pipeline_stable_diffusion_safe": ["StableDiffusionPipelineSafe"],
+    "safety_checker": ["StableDiffusionSafetyChecker"],
+}
+_additional_imports.update({"SafetyConfig": SafetyConfig})
 
 
-try:
+if TYPE_CHECKING:
+    try:
         if not (is_transformers_available() and is_torch_available()):
             raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
+    except OptionalDependencyNotAvailable:
         from ...utils.dummy_torch_and_transformers_objects import *
-else:
+    else:
+        from .pipeline_output import StableDiffusionSafePipelineOutput
         from .pipeline_stable_diffusion_safe import StableDiffusionPipelineSafe
         from .safety_checker import SafeStableDiffusionSafetyChecker
+
+else:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils import dummy_torch_and_transformers_objects
+
+        _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+
+    else:
+        import sys
+
+        sys.modules[__name__] = _LazyModule(
+            __name__,
+            globals()["__file__"],
+            _import_structure,
+            module_spec=__spec__,
+        )
+
+        for name, value in _dummy_objects.items():
+            setattr(sys.modules[__name__], name, value)
+        for name, value in _additional_imports.items():
+            setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py

+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import numpy as np
+import PIL
+
+from ...utils import (
+    BaseOutput,
+)
+
+
+@dataclass
+class StableDiffusionSafePipelineOutput(BaseOutput):
+    """
+    Output class for Safe Stable Diffusion pipelines.
+
+    Args:
+        images (`List[PIL.Image.Image]` or `np.ndarray`)
+            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+        nsfw_content_detected (`List[bool]`)
+            List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, or `None` if safety checking could not be performed.
+        images (`List[PIL.Image.Image]` or `np.ndarray`)
+            List of denoised PIL images that were flagged by the safety checker any may contain "not-safe-for-work"
+            (nsfw) content, or `None` if no safety check was performed or no images were flagged.
+        applied_safety_concept (`str`)
+            The safety concept that was applied for safety guidance, or `None` if safety guidance was disabled
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    nsfw_content_detected: Optional[List[bool]]
+    unsafe_images: Optional[Union[List[PIL.Image.Image], np.ndarray]]
+    applied_safety_concept: Optional[str]

src/diffusers/pipelines/stable_diffusion_xl/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,10 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
-_import_structure["pipeline_output"] = ["StableDiffusionXLPipelineOutput"]
+_import_structure = {"pipeline_output": ["StableDiffusionXLPipelineOutput"]}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -25,15 +25,28 @@ else:
     _import_structure["pipeline_stable_diffusion_xl_inpaint"] = ["StableDiffusionXLInpaintPipeline"]
     _import_structure["pipeline_stable_diffusion_xl_instruct_pix2pix"] = ["StableDiffusionXLInstructPix2PixPipeline"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
+        from .pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipeline
+        from .pipeline_stable_diffusion_xl_inpaint import StableDiffusionXLInpaintPipeline
+        from .pipeline_stable_diffusion_xl_instruct_pix2pix import StableDiffusionXLInstructPix2PixPipeline
 
-sys.modules[__name__] = _LazyModule(
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/stochastic_karras_ve/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import _LazyModule
 
 
-_import_structure = {}
-_import_structure["pipeline_stochastic_karras_ve"] = ["KarrasVePipeline"]
+_import_structure = {"pipeline_stochastic_karras_ve": ["KarrasVePipeline"]}
 
-import sys
+if TYPE_CHECKING:
+    from .pipeline_stochastic_karras_ve import KarrasVePipeline
 
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )

src/diffusers/pipelines/t2i_adapter/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,8 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -21,14 +23,24 @@ else:
     _import_structure["pipeline_stable_diffusion_adapter"] = ["StableDiffusionAdapterPipeline"]
     _import_structure["pipeline_stable_diffusion_xl_adapter"] = ["StableDiffusionXLAdapterPipeline"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .pipeline_stable_diffusion_adapter import StableDiffusionAdapterPipeline
+        from .pipeline_stable_diffusion_xl_adapter import StableDiffusionXLAdapterPipeline
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
-for name, value in _dummy_objects.items():
+    )
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/text_to_video_synthesis/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,8 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -17,19 +19,31 @@ except OptionalDependencyNotAvailable:
     from ...utils import dummy_torch_and_transformers_objects  # noqa F403
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
-
 else:
     _import_structure["pipeline_output"] = ["TextToVideoSDPipelineOutput"]
     _import_structure["pipeline_text_to_video_synth"] = ["TextToVideoSDPipeline"]
     _import_structure["pipeline_text_to_video_synth_img2img"] = ["VideoToVideoSDPipeline"]
     _import_structure["pipeline_text_to_video_zero"] = ["TextToVideoZeroPipeline"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .pipeline_output import TextToVideoSDPipelineOutput
+        from .pipeline_text_to_video_synth import TextToVideoSDPipeline
+        from .pipeline_text_to_video_synth_img2img import VideoToVideoSDPipeline
+        from .pipeline_text_to_video_zero import TextToVideoZeroPipeline
+
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )

src/diffusers/pipelines/unclip/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,9 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
@@ -25,14 +26,26 @@ else:
     _import_structure["pipeline_unclip_image_variation"] = ["UnCLIPImageVariationPipeline"]
     _import_structure["text_proj"] = ["UnCLIPTextProjModel"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .pipeline_unclip import UnCLIPPipeline
+        from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
+        from .text_proj import UnCLIPTextProjModel
+
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
-for name, value in _dummy_objects.items():
+    )
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/unidiffuser/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -6,9 +8,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -22,21 +23,35 @@ except OptionalDependencyNotAvailable:
     _dummy_objects.update(
         {"ImageTextPipelineOutput": ImageTextPipelineOutput, "UniDiffuserPipeline": UniDiffuserPipeline}
     )
-
 else:
     _import_structure["modeling_text_decoder"] = ["UniDiffuserTextDecoder"]
     _import_structure["modeling_uvit"] = ["UniDiffuserModel", "UTransformer2DModel"]
     _import_structure["pipeline_unidiffuser"] = ["ImageTextPipelineOutput", "UniDiffuserPipeline"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import (
+            ImageTextPipelineOutput,
+            UniDiffuserPipeline,
+        )
+    else:
+        from .modeling_text_decoder import UniDiffuserTextDecoder
+        from .modeling_uvit import UniDiffuserModel, UTransformer2DModel
+        from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline
 
-sys.modules[__name__] = _LazyModule(
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/versatile_diffusion/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,9 +9,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
@@ -37,15 +38,33 @@ else:
     _import_structure["pipeline_versatile_diffusion_image_variation"] = ["VersatileDiffusionImageVariationPipeline"]
     _import_structure["pipeline_versatile_diffusion_text_to_image"] = ["VersatileDiffusionTextToImagePipeline"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import (
+            VersatileDiffusionDualGuidedPipeline,
+            VersatileDiffusionImageVariationPipeline,
+            VersatileDiffusionPipeline,
+            VersatileDiffusionTextToImagePipeline,
+        )
+    else:
+        from .pipeline_versatile_diffusion import VersatileDiffusionPipeline
+        from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline
+        from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline
+        from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline
+
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )
 
-for name, value in _dummy_objects.items():
+    for name, value in _dummy_objects.items():
         setattr(sys.modules[__name__], name, value)

src/diffusers/pipelines/vq_diffusion/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -6,9 +8,8 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
-
+_import_structure = {}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -28,12 +29,25 @@ except OptionalDependencyNotAvailable:
 else:
     _import_structure["pipeline_vq_diffusion"] = ["LearnedClassifierFreeSamplingEmbeddings", "VQDiffusionPipeline"]
 
-import sys
 
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import (
+            LearnedClassifierFreeSamplingEmbeddings,
+            VQDiffusionPipeline,
+        )
+    else:
+        from .pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings, VQDiffusionPipeline
+
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )

src/diffusers/pipelines/wuerstchen/__init__.py

+from typing import TYPE_CHECKING
+
 from ...utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
@@ -7,17 +9,16 @@ from ...utils import (
 )
 
 
-_import_structure = {}
 _dummy_objects = {}
+_import_structure = {}
+
 try:
     if not (is_transformers_available() and is_torch_available()):
         raise OptionalDependencyNotAvailable()
-
 except OptionalDependencyNotAvailable:
     from ...utils import dummy_torch_and_transformers_objects
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
-
 else:
     _import_structure["modeling_paella_vq_model"] = ["PaellaVQModel"]
     _import_structure["modeling_wuerstchen_diffnext"] = ["WuerstchenDiffNeXt"]
@@ -27,12 +28,26 @@ else:
     _import_structure["pipeline_wuerstchen_prior"] = ["DEFAULT_STAGE_C_TIMESTEPS", "WuerstchenPriorPipeline"]
 
 
-import sys
+if TYPE_CHECKING:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .modeling_paella_vq_model import PaellaVQModel
+        from .modeling_wuerstchen_diffnext import WuerstchenDiffNeXt
+        from .modeling_wuerstchen_prior import WuerstchenPrior
+        from .pipeline_wuerstchen import WuerstchenDecoderPipeline
+        from .pipeline_wuerstchen_combined import WuerstchenCombinedPipeline
+        from .pipeline_wuerstchen_prior import WuerstchenPriorPipeline
 
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(
+    sys.modules[__name__] = _LazyModule(
         __name__,
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
-)
+    )

src/diffusers/schedulers/__init__.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
 
 from ..utils import (
     OptionalDependencyNotAvailable,
     _LazyModule,
+    get_objects_from_module,
     is_flax_available,
     is_scipy_available,
     is_torch_available,
@@ -23,8 +25,8 @@ from ..utils import (
 )
 
 
-_import_structure = {}
 _dummy_modules = {}
+_import_structure = {}
 
 try:
     if not is_torch_available():
@@ -32,12 +34,7 @@ try:
 except OptionalDependencyNotAvailable:
     from ..utils import dummy_pt_objects  # noqa F403
 
-    modules = {}
-    for name in dir(dummy_pt_objects):
-        if (not name.endswith("Scheduler")) or name.startswith("_"):
-            continue
-        modules[name] = getattr(dummy_pt_objects, name)
-    _dummy_modules.update(modules)
+    _dummy_modules.update(get_objects_from_module(dummy_pt_objects))
 
 else:
     _import_structure["scheduling_consistency_models"] = ["CMStochasticIterativeScheduler"]
@@ -46,6 +43,7 @@ else:
     _import_structure["scheduling_ddim_parallel"] = ["DDIMParallelScheduler"]
     _import_structure["scheduling_ddpm"] = ["DDPMScheduler"]
     _import_structure["scheduling_ddpm_parallel"] = ["DDPMParallelScheduler"]
+    _import_structure["scheduling_ddpm_wuerstchen"] = ["DDPMWuerstchenScheduler"]
     _import_structure["scheduling_deis_multistep"] = ["DEISMultistepScheduler"]
     _import_structure["scheduling_dpmsolver_multistep"] = ["DPMSolverMultistepScheduler"]
     _import_structure["scheduling_dpmsolver_multistep_inverse"] = ["DPMSolverMultistepInverseScheduler"]
@@ -65,13 +63,15 @@ else:
     _import_structure["scheduling_unipc_multistep"] = ["UniPCMultistepScheduler"]
     _import_structure["scheduling_utils"] = ["KarrasDiffusionSchedulers", "SchedulerMixin"]
     _import_structure["scheduling_vq_diffusion"] = ["VQDiffusionScheduler"]
-    _import_structure["scheduling_ddpm_wuerstchen"] = ["DDPMWuerstchenScheduler"]
 
 try:
     if not is_flax_available():
         raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
-    from ..utils.dummy_flax_objects import *  # noqa F403
+    from ..utils import dummy_flax_objects  # noqa F403
+
+    _dummy_modules.update(get_objects_from_module(dummy_flax_objects))
+
 else:
     _import_structure["scheduling_ddim_flax"] = ["FlaxDDIMScheduler"]
     _import_structure["scheduling_ddpm_flax"] = ["FlaxDDPMScheduler"]
@@ -94,13 +94,7 @@ try:
 except OptionalDependencyNotAvailable:
     from ..utils import dummy_torch_and_scipy_objects  # noqa F403
 
-    modules = {}
-    for name in dir(dummy_torch_and_scipy_objects):
-        if (not name.endswith("Scheduler")) or name.startswith("_"):
-            continue
-        modules[name] = getattr(dummy_torch_and_scipy_objects, name)
-
-    _dummy_modules.update(modules)
+    _dummy_modules.update(get_objects_from_module(dummy_torch_and_scipy_objects))
 
 else:
     _import_structure["scheduling_lms_discrete"] = ["LMSDiscreteScheduler"]
@@ -111,21 +105,92 @@ try:
 except OptionalDependencyNotAvailable:
     from ..utils import dummy_torch_and_torchsde_objects  # noqa F403
 
-    modules = {}
-    for name in dir(dummy_torch_and_torchsde_objects):
-        if (not name.endswith("Scheduler")) or name.startswith("_"):
-            continue
-        modules[name] = getattr(dummy_torch_and_torchsde_objects, name)
-
-    _dummy_modules.update(modules)
-
+    _dummy_modules.update(get_objects_from_module(dummy_torch_and_torchsde_objects))
 
 else:
     _import_structure["scheduling_dpmsolver_sde"] = ["DPMSolverSDEScheduler"]
 
-import sys
+if TYPE_CHECKING:
+    from ..utils import (
+        OptionalDependencyNotAvailable,
+        is_flax_available,
+        is_scipy_available,
+        is_torch_available,
+        is_torchsde_available,
+    )
 
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ..utils.dummy_pt_objects import *  # noqa F403
+    else:
+        from .scheduling_consistency_models import CMStochasticIterativeScheduler
+        from .scheduling_ddim import DDIMScheduler
+        from .scheduling_ddim_inverse import DDIMInverseScheduler
+        from .scheduling_ddim_parallel import DDIMParallelScheduler
+        from .scheduling_ddpm import DDPMScheduler
+        from .scheduling_ddpm_parallel import DDPMParallelScheduler
+        from .scheduling_ddpm_wuerstchen import DDPMWuerstchenScheduler
+        from .scheduling_deis_multistep import DEISMultistepScheduler
+        from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
+        from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
+        from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+        from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
+        from .scheduling_euler_discrete import EulerDiscreteScheduler
+        from .scheduling_heun_discrete import HeunDiscreteScheduler
+        from .scheduling_ipndm import IPNDMScheduler
+        from .scheduling_k_dpm_2_ancestral_discrete import KDPM2AncestralDiscreteScheduler
+        from .scheduling_k_dpm_2_discrete import KDPM2DiscreteScheduler
+        from .scheduling_karras_ve import KarrasVeScheduler
+        from .scheduling_pndm import PNDMScheduler
+        from .scheduling_repaint import RePaintScheduler
+        from .scheduling_sde_ve import ScoreSdeVeScheduler
+        from .scheduling_sde_vp import ScoreSdeVpScheduler
+        from .scheduling_unclip import UnCLIPScheduler
+        from .scheduling_unipc_multistep import UniPCMultistepScheduler
+        from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+        from .scheduling_vq_diffusion import VQDiffusionScheduler
+
+    try:
+        if not is_flax_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ..utils.dummy_flax_objects import *  # noqa F403
+    else:
+        from .scheduling_ddim_flax import FlaxDDIMScheduler
+        from .scheduling_ddpm_flax import FlaxDDPMScheduler
+        from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
+        from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
+        from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
+        from .scheduling_pndm_flax import FlaxPNDMScheduler
+        from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
+        from .scheduling_utils_flax import (
+            FlaxKarrasDiffusionSchedulers,
+            FlaxSchedulerMixin,
+            FlaxSchedulerOutput,
+            broadcast_to_shape_from_left,
+        )
+
+    try:
+        if not (is_torch_available() and is_scipy_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
+    else:
+        from .scheduling_lms_discrete import LMSDiscreteScheduler
+
+    try:
+        if not (is_torch_available() and is_torchsde_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
+    else:
+        from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler
+
+else:
+    import sys
 
-sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
-for name, value in _dummy_modules.items():
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
+    for name, value in _dummy_modules.items():
         setattr(sys.modules[__name__], name, value)

utils/custom_init_isort.py

@@ -12,12 +12,35 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+Utility that sorts the imports in the custom inits of Diffusers. Diffusers uses init files that delay the
+import of an object to when it's actually needed. This is to avoid the main init importing all models, which would
+make the line `import transformers` very slow when the user has all optional dependencies installed. The inits with
+delayed imports have two halves: one definining a dictionary `_import_structure` which maps modules to the name of the
+objects in each module, and one in `TYPE_CHECKING` which looks like a normal init for type-checkers. `isort` or `ruff`
+properly sort the second half which looks like traditionl imports, the goal of this script is to sort the first half.
 
+Use from the root of the repo with:
+
+```bash
+python utils/custom_init_isort.py
+```
+
+which will auto-sort the imports (used in `make style`).
+
+For a check only (as used in `make quality`) run:
+
+```bash
+python utils/custom_init_isort.py --check_only
+```
+"""
 import argparse
 import os
 import re
+from typing import Any, Callable, List, Optional
 
 
+# Path is defined with the intent you should run this script from the root of the repo.
 PATH_TO_TRANSFORMERS = "src/diffusers"
 
 # Pattern that looks at the indentation in a line.
@@ -32,17 +55,30 @@ _re_strip_line = re.compile(r'^\s*"([^"]+)",\s*$')
 _re_bracket_content = re.compile(r"\[([^\]]+)\]")
 
 
-def get_indent(line):
-    """Returns the indent in `line`."""
+def get_indent(line: str) -> str:
+    """Returns the indent in  given line (as string)."""
     search = _re_indent.search(line)
     return "" if search is None else search.groups()[0]
 
 
-def split_code_in_indented_blocks(code, indent_level="", start_prompt=None, end_prompt=None):
+def split_code_in_indented_blocks(
+    code: str, indent_level: str = "", start_prompt: Optional[str] = None, end_prompt: Optional[str] = None
+) -> List[str]:
     """
-    Split `code` into its indented blocks, starting at `indent_level`. If provided, begins splitting after
-    `start_prompt` and stops at `end_prompt` (but returns what's before `start_prompt` as a first block and what's
-    after `end_prompt` as a last block, so `code` is always the same as joining the result of this function).
+    Split some code into its indented blocks, starting at a given level.
+
+    Args:
+        code (`str`): The code to split.
+        indent_level (`str`): The indent level (as string) to use for identifying the blocks to split.
+        start_prompt (`str`, *optional*): If provided, only starts splitting at the line where this text is.
+        end_prompt (`str`, *optional*): If provided, stops splitting at a line where this text is.
+
+    Warning:
+        The text before `start_prompt` or after `end_prompt` (if provided) is not ignored, just not split. The input `code`
+        can thus be retrieved by joining the result.
+
+    Returns:
+        `List[str]`: The list of blocks.
     """
     # Let's split the code into lines and move to start_index.
     index = 0
@@ -54,12 +90,17 @@ def split_code_in_indented_blocks(code, indent_level="", start_prompt=None, end_
     else:
         blocks = []
 
-    # We split into blocks until we get to the `end_prompt` (or the end of the block).
+    # This variable contains the block treated at a given time.
     current_block = [lines[index]]
     index += 1
+    # We split into blocks until we get to the `end_prompt` (or the end of the file).
     while index < len(lines) and (end_prompt is None or not lines[index].startswith(end_prompt)):
+        # We have a non-empty line with the proper indent -> start of a new block
         if len(lines[index]) > 0 and get_indent(lines[index]) == indent_level:
+            # Store the current block in the result and rest. There are two cases: the line is part of the block (like
+            # a closing parenthesis) or not.
             if len(current_block) > 0 and get_indent(current_block[-1]).startswith(indent_level + " "):
+                # Line is part of the current block
                 current_block.append(lines[index])
                 blocks.append("\n".join(current_block))
                 if index < len(lines) - 1:
@@ -68,9 +109,11 @@ def split_code_in_indented_blocks(code, indent_level="", start_prompt=None, end_
                 else:
                     current_block = []
             else:
+                # Line is not part of the current block
                 blocks.append("\n".join(current_block))
                 current_block = [lines[index]]
         else:
+            # Just add the line to the current block
             current_block.append(lines[index])
         index += 1
 
@@ -85,8 +128,10 @@ def split_code_in_indented_blocks(code, indent_level="", start_prompt=None, end_
     return blocks
 
 
-def ignore_underscore(key):
-    "Wraps a `key` (that maps an object to string) to lower case and remove underscores."
+def ignore_underscore_and_lowercase(key: Callable[[Any], str]) -> Callable[[Any], str]:
+    """
+    Wraps a key function (as used in a sort) to lowercase and ignore underscores.
+    """
 
     def _inner(x):
         return key(x).lower().replace("_", "")
@@ -94,8 +139,21 @@ def ignore_underscore(key):
     return _inner
 
 
-def sort_objects(objects, key=None):
-    "Sort a list of `objects` following the rules of isort. `key` optionally maps an object to a str."
+def sort_objects(objects: List[Any], key: Optional[Callable[[Any], str]] = None) -> List[Any]:
+    """
+    Sort a list of objects following the rules of isort (all uppercased first, camel-cased second and lower-cased
+    last).
+
+    Args:
+        objects (`List[Any]`):
+            The list of objects to sort.
+        key (`Callable[[Any], str]`, *optional*):
+            A function taking an object as input and returning a string, used to sort them by alphabetical order.
+            If not provided, will default to noop (so a `key` must be provided if the `objects` are not of type string).
+
+    Returns:
+        `List[Any]`: The sorted list with the same elements as in the inputs
+    """
 
     # If no key is provided, we use a noop.
     def noop(x):
@@ -110,18 +168,26 @@ def sort_objects(objects, key=None):
     # Functions begin with a lowercase, they go last.
     functions = [obj for obj in objects if not key(obj)[0].isupper()]
 
-    key1 = ignore_underscore(key)
+    # Then we sort each group.
+    key1 = ignore_underscore_and_lowercase(key)
     return sorted(constants, key=key1) + sorted(classes, key=key1) + sorted(functions, key=key1)
 
 
-def sort_objects_in_import(import_statement):
+def sort_objects_in_import(import_statement: str) -> str:
     """
-    Return the same `import_statement` but with objects properly sorted.
+    Sorts the imports in a single import statement.
+
+    Args:
+        import_statement (`str`): The import statement in which to sort the imports.
+
+    Returns:
+        `str`: The same as the input, but with objects properly sorted.
     """
 
     # This inner function sort imports between [ ].
     def _replace(match):
         imports = match.groups()[0]
+        # If there is one import only, nothing to do.
         if "," not in imports:
             return f"[{imports}]"
         keys = [part.strip().replace('"', "") for part in imports.split(",")]
@@ -165,13 +231,18 @@ def sort_objects_in_import(import_statement):
         return import_statement
 
 
-def sort_imports(file, check_only=True):
+def sort_imports(file: str, check_only: bool = True):
     """
-    Sort `_import_structure` imports in `file`, `check_only` determines if we only check or overwrite.
+    Sort the imports defined in the `_import_structure` of a given init.
+
+    Args:
+        file (`str`): The path to the init to check/fix.
+        check_only (`bool`, *optional*, defaults to `True`): Whether or not to just check (and not auto-fix) the init.
     """
-    with open(file, "r") as f:
+    with open(file, encoding="utf-8") as f:
         code = f.read()
 
+    # If the file is not a custom init, there is nothing to do.
     if "_import_structure" not in code:
         return
 
@@ -180,7 +251,7 @@ def sort_imports(file, check_only=True):
         code, start_prompt="_import_structure = {", end_prompt="if TYPE_CHECKING:"
     )
 
-    # We ignore block 0 (everything until start_prompt) and the last block (everything after end_prompt).
+    # We ignore block 0 (everything untils start_prompt) and the last block (everything after end_prompt).
     for block_idx in range(1, len(main_blocks) - 1):
         # Check if the block contains some `_import_structure`s thingy to sort.
         block = main_blocks[block_idx]
@@ -203,7 +274,7 @@ def sort_imports(file, check_only=True):
         # Slit the internal block into blocks of indent level 1.
         internal_blocks = split_code_in_indented_blocks(internal_block_code, indent_level=indent)
         # We have two categories of import key: list or _import_structure[key].append/extend
-        pattern = _re_direct_key if "_import_structure" in block_lines[0] else _re_indirect_key
+        pattern = _re_direct_key if "_import_structure = {" in block_lines[0] else _re_indirect_key
         # Grab the keys, but there is a trap: some lines are empty or just comments.
         keys = [(pattern.search(b).groups()[0] if pattern.search(b) is not None else None) for b in internal_blocks]
         # We only sort the lines with a key.
@@ -229,11 +300,17 @@ def sort_imports(file, check_only=True):
             return True
         else:
             print(f"Overwriting {file}.")
-            with open(file, "w") as f:
+            with open(file, "w", encoding="utf-8") as f:
                 f.write("\n".join(main_blocks))
 
 
 def sort_imports_in_all_inits(check_only=True):
+    """
+    Sort the imports defined in the `_import_structure` of all inits in the repo.
+
+    Args:
+        check_only (`bool`, *optional*, defaults to `True`): Whether or not to just check (and not auto-fix) the init.
+    """
     failures = []
     for root, _, files in os.walk(PATH_TO_TRANSFORMERS):
         if "__init__.py" in files: