diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index fc60ff845ccf..75ddb21fb15d 100644 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import TYPE_CHECKING + from ..utils import _LazyModule, is_flax_available, is_torch_available @@ -40,7 +42,32 @@ _import_structure["unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"] _import_structure["vae_flax"] = ["FlaxAutoencoderKL"] -import sys +if TYPE_CHECKING: + if is_torch_available(): + from .adapter import MultiAdapter, T2IAdapter + from .autoencoder_asym_kl import AsymmetricAutoencoderKL + from .autoencoder_kl import AutoencoderKL + from .autoencoder_tiny import AutoencoderTiny + from .controlnet import ControlNetModel + from .dual_transformer_2d import DualTransformer2DModel + from .modeling_utils import ModelMixin + from .prior_transformer import PriorTransformer + from .t5_film_transformer import T5FilmDecoder + from .transformer_2d import Transformer2DModel + from .transformer_temporal import TransformerTemporalModel + from .unet_1d import UNet1DModel + from .unet_2d import UNet2DModel + from .unet_2d_condition import UNet2DConditionModel + from .unet_3d_condition import UNet3DConditionModel + from .vq_model import VQModel + + if is_flax_available(): + from .controlnet_flax import FlaxControlNetModel + from .unet_2d_condition_flax import FlaxUNet2DConditionModel + from .vae_flax import FlaxAutoencoderKL + +else: + import sys -sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index b237adae7d54..8bf0a98de893 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ..utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -13,8 +15,8 @@ # These modules contain pipelines from multiple libraries/frameworks -_import_structure = {"stable_diffusion": [], "latent_diffusion": [], "controlnet": []} _dummy_objects = {} +_import_structure = {"stable_diffusion": [], "latent_diffusion": [], "controlnet": []} try: if not is_torch_available(): @@ -23,7 +25,6 @@ from ..utils import dummy_pt_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_pt_objects)) - else: _import_structure["auto_pipeline"] = [ "AutoPipelineForImage2Image", @@ -42,7 +43,6 @@ _import_structure["repaint"] = ["RePaintPipeline"] _import_structure["score_sde_ve"] = ["ScoreSdeVePipeline"] _import_structure["stochastic_karras_ve"] = ["KarrasVePipeline"] - try: if not (is_torch_available() and is_librosa_available()): raise OptionalDependencyNotAvailable() @@ -50,10 +50,8 @@ from ..utils import dummy_torch_and_librosa_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_torch_and_librosa_objects)) - else: _import_structure["audio_diffusion"] = ["AudioDiffusionPipeline", "Mel"] - try: if not (is_torch_available() and is_transformers_available()): raise OptionalDependencyNotAvailable() @@ -61,11 +59,14 @@ 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["alt_diffusion"] = ["AltDiffusionImg2ImgPipeline", "AltDiffusionPipeline"] _import_structure["audioldm"] = ["AudioLDMPipeline"] - _import_structure["audioldm2"] = ["AudioLDM2Pipeline", "AudioLDM2ProjectionModel", "AudioLDM2UNet2DConditionModel"] + _import_structure["audioldm2"] = [ + "AudioLDM2Pipeline", + "AudioLDM2ProjectionModel", + "AudioLDM2UNet2DConditionModel", + ] _import_structure["controlnet"].extend( [ "StableDiffusionControlNetImg2ImgPipeline", @@ -117,6 +118,8 @@ "StableDiffusionDepth2ImgPipeline", "StableDiffusionDiffEditPipeline", "StableDiffusionGLIGENPipeline", + "StableDiffusionGLIGENPipeline", + "StableDiffusionGLIGENTextImagePipeline", "StableDiffusionImageVariationPipeline", "StableDiffusionImg2ImgPipeline", "StableDiffusionInpaintPipeline", @@ -133,8 +136,6 @@ "StableDiffusionUpscalePipeline", "StableUnCLIPImg2ImgPipeline", "StableUnCLIPPipeline", - "StableDiffusionGLIGENTextImagePipeline", - "StableDiffusionGLIGENPipeline", ] ) _import_structure["stable_diffusion_safe"] = ["StableDiffusionPipelineSafe"] @@ -169,8 +170,6 @@ "WuerstchenDecoderPipeline", "WuerstchenPriorPipeline", ] - - try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() @@ -178,10 +177,8 @@ from ..utils import dummy_onnx_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_onnx_objects)) - else: _import_structure["onnx_utils"] = ["OnnxRuntimeModel"] - try: if not (is_torch_available() and is_transformers_available() and is_onnx_available()): raise OptionalDependencyNotAvailable() @@ -189,7 +186,6 @@ from ..utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_onnx_objects)) - else: _import_structure["stable_diffusion"].extend( [ @@ -201,7 +197,6 @@ "StableDiffusionOnnxPipeline", ] ) - try: if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()): raise OptionalDependencyNotAvailable() @@ -209,10 +204,8 @@ 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["stable_diffusion"].extend(["StableDiffusionKDiffusionPipeline"]) - try: if not is_flax_available(): raise OptionalDependencyNotAvailable() @@ -220,11 +213,8 @@ from ..utils import dummy_flax_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_flax_objects)) - else: _import_structure["pipeline_flax_utils"] = ["FlaxDiffusionPipeline"] - - try: if not (is_flax_available() and is_transformers_available()): raise OptionalDependencyNotAvailable() @@ -232,7 +222,6 @@ from ..utils import dummy_flax_and_transformers_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) - else: _import_structure["controlnet"].extend(["FlaxStableDiffusionControlNetPipeline"]) _import_structure["stable_diffusion"].extend( @@ -249,19 +238,217 @@ from ..utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_transformers_and_torch_and_note_seq_objects)) - else: _import_structure["spectrogram_diffusion"] = ["MidiProcessor", "SpectrogramDiffusionPipeline"] +if TYPE_CHECKING: + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_pt_objects import * # noqa F403 + + else: + from .auto_pipeline import AutoPipelineForImage2Image, AutoPipelineForInpainting, AutoPipelineForText2Image + from .consistency_models import ConsistencyModelPipeline + from .dance_diffusion import DanceDiffusionPipeline + from .ddim import DDIMPipeline + from .ddpm import DDPMPipeline + from .dit import DiTPipeline + from .latent_diffusion import LDMSuperResolutionPipeline + from .latent_diffusion_uncond import LDMPipeline + from .pipeline_utils import AudioPipelineOutput, DiffusionPipeline, ImagePipelineOutput + from .pndm import PNDMPipeline + from .repaint import RePaintPipeline + from .score_sde_ve import ScoreSdeVePipeline + from .stochastic_karras_ve import KarrasVePipeline + + try: + if not (is_torch_available() and is_librosa_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_torch_and_librosa_objects import * + else: + from .audio_diffusion import AudioDiffusionPipeline, Mel + + try: + if not (is_torch_available() and is_transformers_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_torch_and_transformers_objects import * + else: + from .alt_diffusion import AltDiffusionImg2ImgPipeline, AltDiffusionPipeline + from .audioldm import AudioLDMPipeline + from .audioldm2 import AudioLDM2Pipeline, AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel + from .controlnet import ( + StableDiffusionControlNetImg2ImgPipeline, + StableDiffusionControlNetInpaintPipeline, + StableDiffusionControlNetPipeline, + StableDiffusionXLControlNetImg2ImgPipeline, + StableDiffusionXLControlNetInpaintPipeline, + StableDiffusionXLControlNetPipeline, + ) + from .deepfloyd_if import ( + IFImg2ImgPipeline, + IFImg2ImgSuperResolutionPipeline, + IFInpaintingPipeline, + IFInpaintingSuperResolutionPipeline, + IFPipeline, + IFSuperResolutionPipeline, + ) + from .kandinsky import ( + KandinskyCombinedPipeline, + KandinskyImg2ImgCombinedPipeline, + KandinskyImg2ImgPipeline, + KandinskyInpaintCombinedPipeline, + KandinskyInpaintPipeline, + KandinskyPipeline, + KandinskyPriorPipeline, + ) + from .kandinsky2_2 import ( + KandinskyV22CombinedPipeline, + KandinskyV22ControlnetImg2ImgPipeline, + KandinskyV22ControlnetPipeline, + KandinskyV22Img2ImgCombinedPipeline, + KandinskyV22Img2ImgPipeline, + KandinskyV22InpaintCombinedPipeline, + KandinskyV22InpaintPipeline, + KandinskyV22Pipeline, + KandinskyV22PriorEmb2EmbPipeline, + KandinskyV22PriorPipeline, + ) + from .latent_diffusion import LDMTextToImagePipeline + from .musicldm import MusicLDMPipeline + from .paint_by_example import PaintByExamplePipeline + from .semantic_stable_diffusion import SemanticStableDiffusionPipeline + from .shap_e import ShapEImg2ImgPipeline, ShapEPipeline + from .stable_diffusion import ( + CycleDiffusionPipeline, + StableDiffusionAttendAndExcitePipeline, + StableDiffusionDepth2ImgPipeline, + StableDiffusionDiffEditPipeline, + StableDiffusionGLIGENPipeline, + StableDiffusionGLIGENTextImagePipeline, + StableDiffusionImageVariationPipeline, + StableDiffusionImg2ImgPipeline, + StableDiffusionInpaintPipeline, + StableDiffusionInpaintPipelineLegacy, + StableDiffusionInstructPix2PixPipeline, + StableDiffusionLatentUpscalePipeline, + StableDiffusionLDM3DPipeline, + StableDiffusionModelEditingPipeline, + StableDiffusionPanoramaPipeline, + StableDiffusionParadigmsPipeline, + StableDiffusionPipeline, + StableDiffusionPix2PixZeroPipeline, + StableDiffusionSAGPipeline, + StableDiffusionUpscalePipeline, + StableUnCLIPImg2ImgPipeline, + StableUnCLIPPipeline, + ) + from .stable_diffusion_safe import StableDiffusionPipelineSafe + from .stable_diffusion_xl import ( + StableDiffusionXLImg2ImgPipeline, + StableDiffusionXLInpaintPipeline, + StableDiffusionXLInstructPix2PixPipeline, + StableDiffusionXLPipeline, + ) + from .t2i_adapter import StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline + from .text_to_video_synthesis import ( + TextToVideoSDPipeline, + TextToVideoZeroPipeline, + VideoToVideoSDPipeline, + ) + from .unclip import UnCLIPImageVariationPipeline, UnCLIPPipeline + from .unidiffuser import ( + ImageTextPipelineOutput, + UniDiffuserModel, + UniDiffuserPipeline, + UniDiffuserTextDecoder, + ) + from .versatile_diffusion import ( + VersatileDiffusionDualGuidedPipeline, + VersatileDiffusionImageVariationPipeline, + VersatileDiffusionPipeline, + VersatileDiffusionTextToImagePipeline, + ) + from .vq_diffusion import VQDiffusionPipeline + from .wuerstchen import ( + WuerstchenCombinedPipeline, + WuerstchenDecoderPipeline, + WuerstchenPriorPipeline, + ) + + try: + if not is_onnx_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_onnx_objects import * # noqa F403 + + else: + from .onnx_utils import OnnxRuntimeModel + + try: + if not (is_torch_available() and is_transformers_available() and is_onnx_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_torch_and_transformers_and_onnx_objects import * + else: + from .stable_diffusion import ( + OnnxStableDiffusionImg2ImgPipeline, + OnnxStableDiffusionInpaintPipeline, + OnnxStableDiffusionInpaintPipelineLegacy, + OnnxStableDiffusionPipeline, + OnnxStableDiffusionUpscalePipeline, + StableDiffusionOnnxPipeline, + ) + + try: + if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_torch_and_transformers_and_k_diffusion_objects import * + else: + from .stable_diffusion import StableDiffusionKDiffusionPipeline + + try: + if not is_flax_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_flax_objects import * # noqa F403 + else: + from .pipeline_flax_utils import FlaxDiffusionPipeline + + try: + if not (is_flax_available() and is_transformers_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ..utils.dummy_flax_and_transformers_objects import * + else: + from .controlnet import FlaxStableDiffusionControlNetPipeline + from .stable_diffusion import ( + FlaxStableDiffusionImg2ImgPipeline, + FlaxStableDiffusionInpaintPipeline, + FlaxStableDiffusionPipeline, + ) + + 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 + + else: + from .spectrogram_diffusion import MidiProcessor, SpectrogramDiffusionPipeline -import sys - +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) + 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) diff --git a/src/diffusers/pipelines/alt_diffusion/__init__.py b/src/diffusers/pipelines/alt_diffusion/__init__.py index c2e4db7eab1c..c0ad3b4a3486 100644 --- a/src/diffusers/pipelines/alt_diffusion/__init__.py +++ b/src/diffusers/pipelines/alt_diffusion/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,8 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -17,21 +19,34 @@ from ...utils import dummy_torch_and_transformers_objects _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) - else: - _import_structure["pipeline_output"] = ["AltDiffusionPipelineOutput"] _import_structure["modeling_roberta_series"] = ["RobertaSeriesModelWithTransformation"] _import_structure["pipeline_alt_diffusion"] = ["AltDiffusionPipeline"] _import_structure["pipeline_alt_diffusion_img2img"] = ["AltDiffusionImg2ImgPipeline"] -import sys + _import_structure["pipeline_output"] = ["AltDiffusionPipelineOutput"] +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 * -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) + else: + from .modeling_roberta_series import RobertaSeriesModelWithTransformation + from .pipeline_alt_diffusion import AltDiffusionPipeline + from .pipeline_alt_diffusion_img2img import AltDiffusionImg2ImgPipeline + from .pipeline_output import AltDiffusionPipelineOutput + +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) diff --git a/src/diffusers/pipelines/audio_diffusion/__init__.py b/src/diffusers/pipelines/audio_diffusion/__init__.py index 578a94693382..7880d02a79a3 100644 --- a/src/diffusers/pipelines/audio_diffusion/__init__.py +++ b/src/diffusers/pipelines/audio_diffusion/__init__.py @@ -1,18 +1,23 @@ -from ...utils import _LazyModule +from typing import TYPE_CHECKING +from ...utils import _LazyModule -_import_structure = {} -_dummy_objects = {} -_import_structure["mel"] = ["Mel"] -_import_structure["pipeline_audio_diffusion"] = ["AudioDiffusionPipeline"] +_import_structure = { + "mel": ["Mel"], + "pipeline_audio_diffusion": ["AudioDiffusionPipeline"], +} -import sys +if TYPE_CHECKING: + from .mel import Mel + from .pipeline_audio_diffusion import AudioDiffusionPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/audioldm/__init__.py b/src/diffusers/pipelines/audioldm/__init__.py index 2acd5c25ed75..57c0fe46c324 100644 --- a/src/diffusers/pipelines/audioldm/__init__.py +++ b/src/diffusers/pipelines/audioldm/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): @@ -20,19 +21,30 @@ ) _dummy_objects.update({"AudioLDMPipeline": AudioLDMPipeline}) - else: _import_structure["pipeline_audioldm"] = ["AudioLDMPipeline"] -import sys +if TYPE_CHECKING: + try: + if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import ( + AudioLDMPipeline, + ) -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + else: + from .pipeline_audioldm import AudioLDMPipeline +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/audioldm2/__init__.py b/src/diffusers/pipelines/audioldm2/__init__.py index 67001f8e44ca..50330c677452 100644 --- a/src/diffusers/pipelines/audioldm2/__init__.py +++ b/src/diffusers/pipelines/audioldm2/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -8,9 +10,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): @@ -19,11 +20,23 @@ from ...utils import dummy_torch_and_transformers_objects _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) - else: _import_structure["modeling_audioldm2"] = ["AudioLDM2ProjectionModel", "AudioLDM2UNet2DConditionModel"] _import_structure["pipeline_audioldm2"] = ["AudioLDM2Pipeline"] + +if TYPE_CHECKING: + try: + if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + + else: + from .modeling_audioldm2 import AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel + from .pipeline_audioldm2 import AudioLDM2Pipeline + +else: import sys sys.modules[__name__] = _LazyModule( @@ -32,3 +45,5 @@ _import_structure, module_spec=__spec__, ) + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/consistency_models/__init__.py b/src/diffusers/pipelines/consistency_models/__init__.py index d1d2ab59500b..83fd1341d82a 100644 --- a/src/diffusers/pipelines/consistency_models/__init__.py +++ b/src/diffusers/pipelines/consistency_models/__init__.py @@ -1,17 +1,21 @@ +from typing import TYPE_CHECKING + from ...utils import ( _LazyModule, ) -_import_structure = {} -_import_structure["pipeline_consistency_models"] = ["ConsistencyModelPipeline"] +_import_structure = {"pipeline_consistency_models": ["ConsistencyModelPipeline"]} -import sys +if TYPE_CHECKING: + from .pipeline_consistency_models import ConsistencyModelPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/controlnet/__init__.py b/src/diffusers/pipelines/controlnet/__init__.py index 60b3fa0b7539..5c551533f3a8 100644 --- a/src/diffusers/pipelines/controlnet/__init__.py +++ b/src/diffusers/pipelines/controlnet/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -8,50 +10,68 @@ ) -_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 # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) - - else: _import_structure["multicontrolnet"] = ["MultiControlNetModel"] _import_structure["pipeline_controlnet"] = ["StableDiffusionControlNetPipeline"] _import_structure["pipeline_controlnet_img2img"] = ["StableDiffusionControlNetImg2ImgPipeline"] _import_structure["pipeline_controlnet_inpaint"] = ["StableDiffusionControlNetInpaintPipeline"] + _import_structure["pipeline_controlnet_inpaint_sd_xl"] = ["StableDiffusionXLControlNetInpaintPipeline"] _import_structure["pipeline_controlnet_sd_xl"] = ["StableDiffusionXLControlNetPipeline"] _import_structure["pipeline_controlnet_sd_xl_img2img"] = ["StableDiffusionXLControlNetImg2ImgPipeline"] - _import_structure["pipeline_controlnet_inpaint_sd_xl"] = ["StableDiffusionXLControlNetInpaintPipeline"] - try: if not (is_transformers_available() and is_flax_available()): raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: from ...utils import dummy_flax_and_transformers_objects # noqa F403 _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) - else: _import_structure["pipeline_flax_controlnet"] = ["FlaxStableDiffusionControlNetPipeline"] -import sys +if TYPE_CHECKING: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .multicontrolnet import MultiControlNetModel + from .pipeline_controlnet import StableDiffusionControlNetPipeline + from .pipeline_controlnet_img2img import StableDiffusionControlNetImg2ImgPipeline + from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline + from .pipeline_controlnet_inpaint_sd_xl import StableDiffusionXLControlNetInpaintPipeline + from .pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline + from .pipeline_controlnet_sd_xl_img2img import StableDiffusionXLControlNetImg2ImgPipeline + + try: + if not (is_transformers_available() and is_flax_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_flax_and_transformers_objects import * # noqa F403 + else: + from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline -for name, value in _dummy_objects.items(): - setattr(sys.modules[__name__], name, value) + +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) diff --git a/src/diffusers/pipelines/dance_diffusion/__init__.py b/src/diffusers/pipelines/dance_diffusion/__init__.py index 39f213b35a04..c777d437060c 100644 --- a/src/diffusers/pipelines/dance_diffusion/__init__.py +++ b/src/diffusers/pipelines/dance_diffusion/__init__.py @@ -1,16 +1,18 @@ -from ...utils import _LazyModule - +from typing import TYPE_CHECKING -_import_structure = {} -_import_structure["pipeline_dance_diffusion"] = ["DanceDiffusionPipeline"] +from ...utils import _LazyModule -import sys +_import_structure = {"pipeline_dance_diffusion": ["DanceDiffusionPipeline"]} +if TYPE_CHECKING: + from .pipeline_dance_diffusion import DanceDiffusionPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/ddim/__init__.py b/src/diffusers/pipelines/ddim/__init__.py index 1715a2b6acbb..0121cd8f6dac 100644 --- a/src/diffusers/pipelines/ddim/__init__.py +++ b/src/diffusers/pipelines/ddim/__init__.py @@ -1,15 +1,18 @@ -from ...utils import _LazyModule +from typing import TYPE_CHECKING +from ...utils import _LazyModule -_import_structure = {} -_import_structure["pipeline_ddim"] = ["DDIMPipeline"] -import sys +_import_structure = {"pipeline_ddim": ["DDIMPipeline"]} +if TYPE_CHECKING: + from .pipeline_ddim import DDIMPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/ddpm/__init__.py b/src/diffusers/pipelines/ddpm/__init__.py index a3936af03a6a..f9320e0bc567 100644 --- a/src/diffusers/pipelines/ddpm/__init__.py +++ b/src/diffusers/pipelines/ddpm/__init__.py @@ -1,17 +1,21 @@ +from typing import TYPE_CHECKING + from ...utils import ( _LazyModule, ) -_import_structure = {} -_import_structure["pipeline_ddpm"] = ["DDPMPipeline"] +_import_structure = {"pipeline_ddpm": ["DDPMPipeline"]} -import sys +if TYPE_CHECKING: + from .pipeline_ddpm import DDPMPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/deepfloyd_if/__init__.py b/src/diffusers/pipelines/deepfloyd_if/__init__.py index a6d58cab9c81..bb0acffc6fa7 100644 --- a/src/diffusers/pipelines/deepfloyd_if/__init__.py +++ b/src/diffusers/pipelines/deepfloyd_if/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,19 +9,19 @@ ) -_import_structure = {} _dummy_objects = {} - -_import_structure["timesteps"] = [ - "fast27_timesteps", - "smart27_timesteps", - "smart50_timesteps", - "smart100_timesteps", - "smart185_timesteps", - "super27_timesteps", - "super40_timesteps", - "super100_timesteps", -] +_import_structure = { + "timesteps": [ + "fast27_timesteps", + "smart100_timesteps", + "smart185_timesteps", + "smart27_timesteps", + "smart50_timesteps", + "super100_timesteps", + "super27_timesteps", + "super40_timesteps", + ] +} try: if not (is_transformers_available() and is_torch_available()): @@ -28,28 +30,55 @@ 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"] = ["IFPipelineOutput"] _import_structure["pipeline_if"] = ["IFPipeline"] _import_structure["pipeline_if_img2img"] = ["IFImg2ImgPipeline"] _import_structure["pipeline_if_img2img_superresolution"] = ["IFImg2ImgSuperResolutionPipeline"] _import_structure["pipeline_if_inpainting"] = ["IFInpaintingPipeline"] _import_structure["pipeline_if_inpainting_superresolution"] = ["IFInpaintingSuperResolutionPipeline"] _import_structure["pipeline_if_superresolution"] = ["IFSuperResolutionPipeline"] + _import_structure["pipeline_output"] = ["IFPipelineOutput"] _import_structure["safety_checker"] = ["IFSafetyChecker"] _import_structure["watermark"] = ["IFWatermarker"] -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_if import IFPipeline + from .pipeline_if_img2img import IFImg2ImgPipeline + from .pipeline_if_img2img_superresolution import IFImg2ImgSuperResolutionPipeline + from .pipeline_if_inpainting import IFInpaintingPipeline + from .pipeline_if_inpainting_superresolution import IFInpaintingSuperResolutionPipeline + from .pipeline_if_superresolution import IFSuperResolutionPipeline + from .pipeline_output import IFPipelineOutput + from .safety_checker import IFSafetyChecker + from .timesteps import ( + fast27_timesteps, + smart27_timesteps, + smart50_timesteps, + smart100_timesteps, + smart185_timesteps, + super27_timesteps, + super40_timesteps, + super100_timesteps, + ) + from .watermark import IFWatermarker -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/dit/__init__.py b/src/diffusers/pipelines/dit/__init__.py index be3c74454393..a260779cafae 100644 --- a/src/diffusers/pipelines/dit/__init__.py +++ b/src/diffusers/pipelines/dit/__init__.py @@ -1,15 +1,19 @@ +from typing import TYPE_CHECKING + from ...utils import _LazyModule -_import_structure = {} -_import_structure["pipeline_dit"] = ["DiTPipeline"] +_import_structure = {"pipeline_dit": ["DiTPipeline"]} -import sys +if TYPE_CHECKING: + from .pipeline_dit import DiTPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/kandinsky/__init__.py b/src/diffusers/pipelines/kandinsky/__init__.py index cc4580721eff..63b34e16c95a 100644 --- a/src/diffusers/pipelines/kandinsky/__init__.py +++ b/src/diffusers/pipelines/kandinsky/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -18,7 +19,6 @@ 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_kandinsky"] = ["KandinskyPipeline"] _import_structure["pipeline_kandinsky_combined"] = [ @@ -32,15 +32,34 @@ _import_structure["text_encoder"] = ["MultilingualCLIP"] -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_kandinsky import KandinskyPipeline + from .pipeline_kandinsky_combined import ( + KandinskyCombinedPipeline, + KandinskyImg2ImgCombinedPipeline, + KandinskyInpaintCombinedPipeline, + ) + from .pipeline_kandinsky_img2img import KandinskyImg2ImgPipeline + from .pipeline_kandinsky_inpaint import KandinskyInpaintPipeline + from .pipeline_kandinsky_prior import KandinskyPriorPipeline, KandinskyPriorPipelineOutput + from .text_encoder import MultilingualCLIP -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/kandinsky2_2/__init__.py b/src/diffusers/pipelines/kandinsky2_2/__init__.py index 639d6ad977c2..461e3d25ca73 100644 --- a/src/diffusers/pipelines/kandinsky2_2/__init__.py +++ b/src/diffusers/pipelines/kandinsky2_2/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -18,7 +19,6 @@ 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_kandinsky2_2"] = ["KandinskyV22Pipeline"] _import_structure["pipeline_kandinsky2_2_combined"] = [ @@ -34,15 +34,36 @@ _import_structure["pipeline_kandinsky2_2_prior_emb2emb"] = ["KandinskyV22PriorEmb2EmbPipeline"] -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_kandinsky2_2 import KandinskyV22Pipeline + from .pipeline_kandinsky2_2_combined import ( + KandinskyV22CombinedPipeline, + KandinskyV22Img2ImgCombinedPipeline, + KandinskyV22InpaintCombinedPipeline, + ) + from .pipeline_kandinsky2_2_controlnet import KandinskyV22ControlnetPipeline + from .pipeline_kandinsky2_2_controlnet_img2img import KandinskyV22ControlnetImg2ImgPipeline + from .pipeline_kandinsky2_2_img2img import KandinskyV22Img2ImgPipeline + from .pipeline_kandinsky2_2_inpainting import KandinskyV22InpaintPipeline + from .pipeline_kandinsky2_2_prior import KandinskyV22PriorPipeline + from .pipeline_kandinsky2_2_prior_emb2emb import KandinskyV22PriorEmb2EmbPipeline -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/latent_diffusion/__init__.py b/src/diffusers/pipelines/latent_diffusion/__init__.py index a78e6622bcfe..bc6ac82217a3 100644 --- a/src/diffusers/pipelines/latent_diffusion/__init__.py +++ b/src/diffusers/pipelines/latent_diffusion/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -23,15 +24,26 @@ _import_structure["pipeline_latent_diffusion_superresolution"] = ["LDMSuperResolutionPipeline"] -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_latent_diffusion import LDMBertModel, LDMTextToImagePipeline + from .pipeline_latent_diffusion_superresolution import LDMSuperResolutionPipeline -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py b/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py index 73e5c703f61a..2dd64d353513 100644 --- a/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py +++ b/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py @@ -1,15 +1,18 @@ -from ...utils import _LazyModule +from typing import TYPE_CHECKING +from ...utils import _LazyModule -_import_structure = {} -_import_structure["pipeline_latent_diffusion_uncond"] = ["LDMPipeline"] -import sys +_import_structure = {"pipeline_latent_diffusion_uncond": ["LDMPipeline"]} +if TYPE_CHECKING: + from .pipeline_latent_diffusion_uncond import LDMPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/musicldm/__init__.py b/src/diffusers/pipelines/musicldm/__init__.py index 6228f763a53b..e49eb1f16d7d 100644 --- a/src/diffusers/pipelines/musicldm/__init__.py +++ b/src/diffusers/pipelines/musicldm/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -8,8 +10,8 @@ ) -_import_structure = {} _dummy_objects = {} +_import_structure = {} try: if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): @@ -18,19 +20,29 @@ 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_musicldm"] = ["MusicLDMPipeline"] -import sys +if TYPE_CHECKING: + try: + if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): + raise OptionalDependencyNotAvailable() -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_musicldm import MusicLDMPipeline + +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/paint_by_example/__init__.py b/src/diffusers/pipelines/paint_by_example/__init__.py index c19ce1036e3f..bfe4810e5ab5 100644 --- a/src/diffusers/pipelines/paint_by_example/__init__.py +++ b/src/diffusers/pipelines/paint_by_example/__init__.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import List, Optional, Union +from typing import TYPE_CHECKING, List, Optional, Union import numpy as np import PIL @@ -14,8 +14,8 @@ ) -_import_structure = {} _dummy_objects = {} +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -28,15 +28,27 @@ _import_structure["image_encoder"] = ["PaintByExampleImageEncoder"] _import_structure["pipeline_paint_by_example"] = ["PaintByExamplePipeline"] -import sys +if TYPE_CHECKING: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .image_encoder import PaintByExampleImageEncoder + from .pipeline_paint_by_example import PaintByExamplePipeline + +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/pndm/__init__.py b/src/diffusers/pipelines/pndm/__init__.py index 7374016c32d9..4764b1b7594a 100644 --- a/src/diffusers/pipelines/pndm/__init__.py +++ b/src/diffusers/pipelines/pndm/__init__.py @@ -1,16 +1,18 @@ -from ...utils import _LazyModule - +from typing import TYPE_CHECKING -_import_structure = {} -_import_structure["pipeline_pndm"] = ["PNDMPipeline"] +from ...utils import _LazyModule -import sys +_import_structure = {"pipeline_pndm": ["PNDMPipeline"]} +if TYPE_CHECKING: + from .pipeline_pndm import PNDMPipeline +else: + import sys -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/repaint/__init__.py b/src/diffusers/pipelines/repaint/__init__.py index 2a0eedf30bbf..ffed0c2ab05c 100644 --- a/src/diffusers/pipelines/repaint/__init__.py +++ b/src/diffusers/pipelines/repaint/__init__.py @@ -1,15 +1,19 @@ +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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/score_sde_ve/__init__.py b/src/diffusers/pipelines/score_sde_ve/__init__.py index 2cd7ac2bf440..192467afd11a 100644 --- a/src/diffusers/pipelines/score_sde_ve/__init__.py +++ b/src/diffusers/pipelines/score_sde_ve/__init__.py @@ -1,15 +1,19 @@ +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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/semantic_stable_diffusion/__init__.py b/src/diffusers/pipelines/semantic_stable_diffusion/__init__.py index 1b743ac3d58d..96842bc84225 100644 --- a/src/diffusers/pipelines/semantic_stable_diffusion/__init__.py +++ b/src/diffusers/pipelines/semantic_stable_diffusion/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -18,21 +19,30 @@ 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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/shap_e/__init__.py b/src/diffusers/pipelines/shap_e/__init__.py index 2a56148fee91..13a9fc1aa1cb 100644 --- a/src/diffusers/pipelines/shap_e/__init__.py +++ b/src/diffusers/pipelines/shap_e/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -33,15 +34,37 @@ "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, + ) -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/spectrogram_diffusion/__init__.py b/src/diffusers/pipelines/spectrogram_diffusion/__init__.py index e8bcf63c2986..e6e8393aa8af 100644 --- a/src/diffusers/pipelines/spectrogram_diffusion/__init__.py +++ b/src/diffusers/pipelines/spectrogram_diffusion/__init__.py @@ -1,4 +1,5 @@ # flake8: noqa +from typing import TYPE_CHECKING from ...utils import ( _LazyModule, is_note_seq_available, @@ -8,9 +9,8 @@ get_objects_from_module, ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -19,32 +19,56 @@ 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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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(): - setattr(sys.modules[__name__], name, value) + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/stable_diffusion/__init__.py b/src/diffusers/pipelines/stable_diffusion/__init__.py index f6f3327c5fb6..57dff8568203 100644 --- a/src/diffusers/pipelines/stable_diffusion/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -12,15 +14,12 @@ ) -_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 @@ 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 @@ _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 @@ _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 @@ _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 @@ 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 @@ 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 @@ 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 + + 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__, -) + 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) + 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) diff --git a/src/diffusers/pipelines/stable_diffusion_safe/__init__.py b/src/diffusers/pipelines/stable_diffusion_safe/__init__.py index 8ce71c01328a..67c6ab1f6686 100644 --- a/src/diffusers/pipelines/stable_diffusion_safe/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion_safe/__init__.py @@ -1,12 +1,19 @@ 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 - """ +_dummy_objects = {} +_additional_imports = {} +_import_structure = { + "pipeline_output": ["StableDiffusionSafePipelineOutput"], + "pipeline_stable_diffusion_safe": ["StableDiffusionPipelineSafe"], + "safety_checker": ["StableDiffusionSafetyChecker"], +} +_additional_imports.update({"SafetyConfig": SafetyConfig}) - 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] +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_output import StableDiffusionSafePipelineOutput + from .pipeline_stable_diffusion_safe import StableDiffusionPipelineSafe + from .safety_checker import SafeStableDiffusionSafetyChecker -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_stable_diffusion_safe import StableDiffusionPipelineSafe - from .safety_checker import SafeStableDiffusionSafetyChecker + 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) diff --git a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py new file mode 100644 index 000000000000..8567a304c696 --- /dev/null +++ b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py @@ -0,0 +1,34 @@ +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] diff --git a/src/diffusers/pipelines/stable_diffusion_xl/__init__.py b/src/diffusers/pipelines/stable_diffusion_xl/__init__.py index ebe12db15fd9..2c4bf44f8dec 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,10 +9,8 @@ ) -_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 @@ _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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/stochastic_karras_ve/__init__.py b/src/diffusers/pipelines/stochastic_karras_ve/__init__.py index 2f82b438c5e3..ebc6506f58b3 100644 --- a/src/diffusers/pipelines/stochastic_karras_ve/__init__.py +++ b/src/diffusers/pipelines/stochastic_karras_ve/__init__.py @@ -1,15 +1,19 @@ +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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/t2i_adapter/__init__.py b/src/diffusers/pipelines/t2i_adapter/__init__.py index b6e6ee724a67..db5dd4ff21b6 100644 --- a/src/diffusers/pipelines/t2i_adapter/__init__.py +++ b/src/diffusers/pipelines/t2i_adapter/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,8 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -21,14 +23,24 @@ _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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) -for name, value in _dummy_objects.items(): - setattr(sys.modules[__name__], name, value) + 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) diff --git a/src/diffusers/pipelines/text_to_video_synthesis/__init__.py b/src/diffusers/pipelines/text_to_video_synthesis/__init__.py index af3b9bfde1ce..a09a63476b7c 100644 --- a/src/diffusers/pipelines/text_to_video_synthesis/__init__.py +++ b/src/diffusers/pipelines/text_to_video_synthesis/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,8 +9,8 @@ ) -_import_structure = {} _dummy_objects = {} +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -17,19 +19,31 @@ 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 -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/unclip/__init__.py b/src/diffusers/pipelines/unclip/__init__.py index f546dbb5041d..6d6a6398bcec 100644 --- a/src/diffusers/pipelines/unclip/__init__.py +++ b/src/diffusers/pipelines/unclip/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_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 @@ _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 -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) +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) diff --git a/src/diffusers/pipelines/unidiffuser/__init__.py b/src/diffusers/pipelines/unidiffuser/__init__.py index ac0207b6045d..52bdb0c40552 100644 --- a/src/diffusers/pipelines/unidiffuser/__init__.py +++ b/src/diffusers/pipelines/unidiffuser/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -6,9 +8,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -22,21 +23,35 @@ _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( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/versatile_diffusion/__init__.py b/src/diffusers/pipelines/versatile_diffusion/__init__.py index 8fbe932b18a6..ba7019c24d94 100644 --- a/src/diffusers/pipelines/versatile_diffusion/__init__.py +++ b/src/diffusers/pipelines/versatile_diffusion/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,9 +9,8 @@ ) -_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 @@ _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 -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +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 _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/vq_diffusion/__init__.py b/src/diffusers/pipelines/vq_diffusion/__init__.py index 8917802c2694..b8fb7f55e8ce 100644 --- a/src/diffusers/pipelines/vq_diffusion/__init__.py +++ b/src/diffusers/pipelines/vq_diffusion/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -6,9 +8,8 @@ ) -_import_structure = {} _dummy_objects = {} - +_import_structure = {} try: if not (is_transformers_available() and is_torch_available()): @@ -28,12 +29,25 @@ 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 -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/pipelines/wuerstchen/__init__.py b/src/diffusers/pipelines/wuerstchen/__init__.py index 17da4c1ad8c5..13407f2cd10c 100644 --- a/src/diffusers/pipelines/wuerstchen/__init__.py +++ b/src/diffusers/pipelines/wuerstchen/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from ...utils import ( OptionalDependencyNotAvailable, _LazyModule, @@ -7,17 +9,16 @@ ) -_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 @@ _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 -sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, -) +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 270e10cdbe18..bbd943e56cec 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/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 @@ ) -_import_structure = {} _dummy_modules = {} +_import_structure = {} try: if not is_torch_available(): @@ -32,12 +34,7 @@ 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 @@ _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 @@ _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 @@ 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 @@ 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(): - setattr(sys.modules[__name__], name, value) + 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) diff --git a/utils/custom_init_isort.py b/utils/custom_init_isort.py index f8ef799c5e6c..e1e85974aeed 100644 --- a/utils/custom_init_isort.py +++ b/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_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 _inner(x): 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 noop(x): # 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 _replace(match): 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: