


"""simple docstring"""




__UpperCamelCase						=      '''0.18.2'''

from .configuration_utils import ConfigMixin
from .utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_onnx_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    logging,
)


try:
	if not is_onnx_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_onnx_objects import *  # noqa F403
else:
	from .pipelines import OnnxRuntimeModel

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_pt_objects import *  # noqa F403
else:
	from .models import (
	    AutoencoderKL,
	    ControlNetModel,
	    ModelMixin,
	    PriorTransformer,
	    TaFilmDecoder,
	    TransformeraDModel,
	    UNetaDModel,
	    UNetaDConditionModel,
	    UNetaDModel,
	    UNetaDConditionModel,
	    VQModel,
	)
	from .optimization import (
	    get_constant_schedule,
	    get_constant_schedule_with_warmup,
	    get_cosine_schedule_with_warmup,
	    get_cosine_with_hard_restarts_schedule_with_warmup,
	    get_linear_schedule_with_warmup,
	    get_polynomial_decay_schedule_with_warmup,
	    get_scheduler,
	)
	from .pipelines import (
	    AudioPipelineOutput,
	    ConsistencyModelPipeline,
	    DanceDiffusionPipeline,
	    DDIMPipeline,
	    DDPMPipeline,
	    DiffusionPipeline,
	    DiTPipeline,
	    ImagePipelineOutput,
	    KarrasVePipeline,
	    LDMPipeline,
	    LDMSuperResolutionPipeline,
	    PNDMPipeline,
	    RePaintPipeline,
	    ScoreSdeVePipeline,
	)
	from .schedulers import (
	    CMStochasticIterativeScheduler,
	    DDIMInverseScheduler,
	    DDIMParallelScheduler,
	    DDIMScheduler,
	    DDPMParallelScheduler,
	    DDPMScheduler,
	    DEISMultistepScheduler,
	    DPMSolverMultistepInverseScheduler,
	    DPMSolverMultistepScheduler,
	    DPMSolverSinglestepScheduler,
	    EulerAncestralDiscreteScheduler,
	    EulerDiscreteScheduler,
	    HeunDiscreteScheduler,
	    IPNDMScheduler,
	    KarrasVeScheduler,
	    KDPMaAncestralDiscreteScheduler,
	    KDPMaDiscreteScheduler,
	    PNDMScheduler,
	    RePaintScheduler,
	    SchedulerMixin,
	    ScoreSdeVeScheduler,
	    UnCLIPScheduler,
	    UniPCMultistepScheduler,
	    VQDiffusionScheduler,
	)
	from .training_utils import EMAModel

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 .schedulers 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 .schedulers import DPMSolverSDEScheduler

try:
	if not (is_torch_available() and is_transformers_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
	from .pipelines import (
	    AltDiffusionImgaImgPipeline,
	    AltDiffusionPipeline,
	    AudioLDMPipeline,
	    CycleDiffusionPipeline,
	    IFImgaImgPipeline,
	    IFImgaImgSuperResolutionPipeline,
	    IFInpaintingPipeline,
	    IFInpaintingSuperResolutionPipeline,
	    IFPipeline,
	    IFSuperResolutionPipeline,
	    ImageTextPipelineOutput,
	    KandinskyImgaImgPipeline,
	    KandinskyInpaintPipeline,
	    KandinskyPipeline,
	    KandinskyPriorPipeline,
	    KandinskyVaaControlnetImgaImgPipeline,
	    KandinskyVaaControlnetPipeline,
	    KandinskyVaaImgaImgPipeline,
	    KandinskyVaaInpaintPipeline,
	    KandinskyVaaPipeline,
	    KandinskyVaaPriorEmbaEmbPipeline,
	    KandinskyVaaPriorPipeline,
	    LDMTextToImagePipeline,
	    PaintByExamplePipeline,
	    SemanticStableDiffusionPipeline,
	    ShapEImgaImgPipeline,
	    ShapEPipeline,
	    StableDiffusionAttendAndExcitePipeline,
	    StableDiffusionControlNetImgaImgPipeline,
	    StableDiffusionControlNetInpaintPipeline,
	    StableDiffusionControlNetPipeline,
	    StableDiffusionDepthaImgPipeline,
	    StableDiffusionDiffEditPipeline,
	    StableDiffusionImageVariationPipeline,
	    StableDiffusionImgaImgPipeline,
	    StableDiffusionInpaintPipeline,
	    StableDiffusionInpaintPipelineLegacy,
	    StableDiffusionInstructPixaPixPipeline,
	    StableDiffusionLatentUpscalePipeline,
	    StableDiffusionLDMaDPipeline,
	    StableDiffusionModelEditingPipeline,
	    StableDiffusionPanoramaPipeline,
	    StableDiffusionParadigmsPipeline,
	    StableDiffusionPipeline,
	    StableDiffusionPipelineSafe,
	    StableDiffusionPixaPixZeroPipeline,
	    StableDiffusionSAGPipeline,
	    StableDiffusionUpscalePipeline,
	    StableUnCLIPImgaImgPipeline,
	    StableUnCLIPPipeline,
	    TextToVideoSDPipeline,
	    TextToVideoZeroPipeline,
	    UnCLIPImageVariationPipeline,
	    UnCLIPPipeline,
	    UniDiffuserModel,
	    UniDiffuserPipeline,
	    UniDiffuserTextDecoder,
	    VersatileDiffusionDualGuidedPipeline,
	    VersatileDiffusionImageVariationPipeline,
	    VersatileDiffusionPipeline,
	    VersatileDiffusionTextToImagePipeline,
	    VideoToVideoSDPipeline,
	    VQDiffusionPipeline,
	)

try:
	if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import *  # noqa F403
else:
	from .pipelines import StableDiffusionXLImgaImgPipeline, StableDiffusionXLPipeline

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 *  # noqa F403
else:
	from .pipelines import StableDiffusionKDiffusionPipeline

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 *  # noqa F403
else:
	from .pipelines import (
	    OnnxStableDiffusionImgaImgPipeline,
	    OnnxStableDiffusionInpaintPipeline,
	    OnnxStableDiffusionInpaintPipelineLegacy,
	    OnnxStableDiffusionPipeline,
	    OnnxStableDiffusionUpscalePipeline,
	    StableDiffusionOnnxPipeline,
	)

try:
	if not (is_torch_available() and is_librosa_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
else:
	from .pipelines import AudioDiffusionPipeline, Mel

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 .pipelines import SpectrogramDiffusionPipeline

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_flax_objects import *  # noqa F403
else:
	from .models.controlnet_flax import FlaxControlNetModel
	from .models.modeling_flax_utils import FlaxModelMixin
	from .models.unet_ad_condition_flax import FlaxUNetaDConditionModel
	from .models.vae_flax import FlaxAutoencoderKL
	from .pipelines import FlaxDiffusionPipeline
	from .schedulers import (
	    FlaxDDIMScheduler,
	    FlaxDDPMScheduler,
	    FlaxDPMSolverMultistepScheduler,
	    FlaxKarrasVeScheduler,
	    FlaxLMSDiscreteScheduler,
	    FlaxPNDMScheduler,
	    FlaxSchedulerMixin,
	    FlaxScoreSdeVeScheduler,
	)


try:
	if not (is_flax_available() and is_transformers_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
else:
	from .pipelines import (
	    FlaxStableDiffusionControlNetPipeline,
	    FlaxStableDiffusionImgaImgPipeline,
	    FlaxStableDiffusionInpaintPipeline,
	    FlaxStableDiffusionPipeline,
	)

try:
	if not (is_note_seq_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from .utils.dummy_note_seq_objects import *  # noqa F403
else:
	from .pipelines import MidiProcessor









"""simple docstring"""





import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class 		lowercase					:





      def     _snake_case			(       self )		->     List[Any]:
            torch.manual_seed(0 )
            lowerCAmelCase  						=       TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       UNetaDConditionModel(
                sample_size=32      ,						layers_per_block=1      ,						block_out_channels=[32, 64]      ,						down_block_types=[
                    """ResnetDownsampleBlock2D""",
                    """SimpleCrossAttnDownBlock2D""",
                ]      ,						mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""      ,						up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]      ,						in_channels=3      ,						out_channels=6      ,						cross_attention_dim=32      ,						encoder_hid_dim=32      ,						attention_head_dim=8      ,						addition_embed_type="""text"""      ,						addition_embed_type_num_heads=2      ,						cross_attention_norm="""group_norm"""      ,						resnet_time_scale_shift="""scale_shift"""      ,						act_fn="""gelu"""      ,						)
            unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            torch.manual_seed(0 )
            lowerCAmelCase  						=       DDPMScheduler(
                num_train_timesteps=1_000      ,						beta_schedule="""squaredcos_cap_v2"""      ,						beta_start=0.0_001      ,						beta_end=0.02      ,						thresholding=lowercase      ,						dynamic_thresholding_ratio=0.95      ,						sample_max_value=1.0      ,						prediction_type="""epsilon"""      ,						variance_type="""learned_range"""      ,						)

            torch.manual_seed(0 )
            lowerCAmelCase  						=       IFWatermarker()

            return {
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "unet": unet,
                "scheduler": scheduler,
                "watermarker": watermarker,
                "safety_checker": None,
                "feature_extractor": None,
            }





      def     _snake_case			(       self )		->     int:
            torch.manual_seed(0 )
            lowerCAmelCase  						=       TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       UNetaDConditionModel(
                sample_size=32      ,						layers_per_block=[1, 2]      ,						block_out_channels=[32, 64]      ,						down_block_types=[
                    """ResnetDownsampleBlock2D""",
                    """SimpleCrossAttnDownBlock2D""",
                ]      ,						mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""      ,						up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]      ,						in_channels=6      ,						out_channels=6      ,						cross_attention_dim=32      ,						encoder_hid_dim=32      ,						attention_head_dim=8      ,						addition_embed_type="""text"""      ,						addition_embed_type_num_heads=2      ,						cross_attention_norm="""group_norm"""      ,						resnet_time_scale_shift="""scale_shift"""      ,						act_fn="""gelu"""      ,						class_embed_type="""timestep"""      ,						mid_block_scale_factor=1.414      ,						time_embedding_act_fn="""gelu"""      ,						time_embedding_dim=32      ,						)
            unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            torch.manual_seed(0 )
            lowerCAmelCase  						=       DDPMScheduler(
                num_train_timesteps=1_000      ,						beta_schedule="""squaredcos_cap_v2"""      ,						beta_start=0.0_001      ,						beta_end=0.02      ,						thresholding=lowercase      ,						dynamic_thresholding_ratio=0.95      ,						sample_max_value=1.0      ,						prediction_type="""epsilon"""      ,						variance_type="""learned_range"""      ,						)

            torch.manual_seed(0 )
            lowerCAmelCase  						=       DDPMScheduler(
                num_train_timesteps=1_000      ,						beta_schedule="""squaredcos_cap_v2"""      ,						beta_start=0.0_001      ,						beta_end=0.02      ,						)

            torch.manual_seed(0 )
            lowerCAmelCase  						=       IFWatermarker()

            return {
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "unet": unet,
                "scheduler": scheduler,
                "image_noising_scheduler": image_noising_scheduler,
                "watermarker": watermarker,
                "safety_checker": None,
                "feature_extractor": None,
            }





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.get_dummy_components()
            lowerCAmelCase  						=       self.pipeline_class(**lowercase )
            pipe.to(lowercase )
            pipe.set_progress_bar_config(disable=lowercase )

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )

            lowerCAmelCase  						=       inputs["""prompt"""]
            lowerCAmelCase  						=       inputs["""generator"""]
            lowerCAmelCase  						=       inputs["""num_inference_steps"""]
            lowerCAmelCase  						=       inputs["""output_type"""]

            if "image" in inputs:
                  lowerCAmelCase  						=       inputs["""image"""]
            else:
                  lowerCAmelCase  						=       None

            if "mask_image" in inputs:
                  lowerCAmelCase  						=       inputs["""mask_image"""]
            else:
                  lowerCAmelCase  						=       None

            if "original_image" in inputs:
                  lowerCAmelCase  						=       inputs["""original_image"""]
            else:
                  lowerCAmelCase  						=       None

            lowerCAmelCase						,		lowerCAmelCase  						=       pipe.encode_prompt(lowercase )

            # inputs with prompt converted to embeddings
            lowerCAmelCase  						=       {
                """prompt_embeds""": prompt_embeds,
                """negative_prompt_embeds""": negative_prompt_embeds,
                """generator""": generator,
                """num_inference_steps""": num_inference_steps,
                """output_type""": output_type,
            }

            if image is not None:
                  lowerCAmelCase  						=       image

            if mask_image is not None:
                  lowerCAmelCase  						=       mask_image

            if original_image is not None:
                  lowerCAmelCase  						=       original_image

            # set all optional components to None
            for optional_component in pipe._optional_components:
                  setattr(lowercase      ,						lowercase      ,						lowercase )

            lowerCAmelCase  						=       pipe(**lowercase )[0]

            with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(lowercase )
                  lowerCAmelCase  						=       self.pipeline_class.from_pretrained(lowercase )
                  pipe_loaded.to(lowercase )
                  pipe_loaded.set_progress_bar_config(disable=lowercase )

            pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            for optional_component in pipe._optional_components:
                  self.assertTrue(
                      getattr(lowercase      ,						lowercase ) is None      ,						f'`{optional_component}` did not stay set to None after loading.'      ,						)

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )

            lowerCAmelCase  						=       inputs["""generator"""]
            lowerCAmelCase  						=       inputs["""num_inference_steps"""]
            lowerCAmelCase  						=       inputs["""output_type"""]

            # inputs with prompt converted to embeddings
            lowerCAmelCase  						=       {
                """prompt_embeds""": prompt_embeds,
                """negative_prompt_embeds""": negative_prompt_embeds,
                """generator""": generator,
                """num_inference_steps""": num_inference_steps,
                """output_type""": output_type,
            }

            if image is not None:
                  lowerCAmelCase  						=       image

            if mask_image is not None:
                  lowerCAmelCase  						=       mask_image

            if original_image is not None:
                  lowerCAmelCase  						=       original_image

            lowerCAmelCase  						=       pipe_loaded(**lowercase )[0]

            lowerCAmelCase  						=       np.abs(to_np(lowercase ) - to_np(lowercase ) ).max()
            self.assertLess(lowercase      ,						1e-4 )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.get_dummy_components()
            lowerCAmelCase  						=       self.pipeline_class(**lowercase )
            pipe.to(lowercase )
            pipe.set_progress_bar_config(disable=lowercase )

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )
            lowerCAmelCase  						=       pipe(**lowercase )[0]

            with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(lowercase )
                  lowerCAmelCase  						=       self.pipeline_class.from_pretrained(lowercase )
                  pipe_loaded.to(lowercase )
                  pipe_loaded.set_progress_bar_config(disable=lowercase )

            pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )
            lowerCAmelCase  						=       pipe_loaded(**lowercase )[0]

            lowerCAmelCase  						=       np.abs(to_np(lowercase ) - to_np(lowercase ) ).max()
            self.assertLess(lowercase      ,						1e-4 )






