



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class        __lowerCAmelCase   (		__lowercase     ,		unittest.TestCase      ):
 lowerCAmelCase__				    =     KandinskyInpaintPipeline
 lowerCAmelCase__				    =     ["prompt", "image_embeds", "negative_image_embeds", "image", "mask_image"]
 lowerCAmelCase__				    =     [
     "prompt",
     "negative_prompt",
     "image_embeds",
     "negative_image_embeds",
     "image",
     "mask_image",
 ]
 lowerCAmelCase__				    =     [
     "generator",
     "height",
     "width",
     "latents",
     "guidance_scale",
     "negative_prompt",
     "num_inference_steps",
     "return_dict",
     "guidance_scale",
     "num_images_per_prompt",
     "output_type",
     "return_dict",
 ]
 lowerCAmelCase__				    =     False




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return 32




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return 32




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return self.time_input_dim




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return self.time_input_dim * 4




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return 100




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base''' )
    return tokenizer




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    torch.manual_seed(0 )
    __lowerCamelCase  							=			MCLIPConfig(
        numDims=self.cross_attention_dim		, transformerDimensions=self.text_embedder_hidden_size		, hidden_size=self.text_embedder_hidden_size		, intermediate_size=37		, num_attention_heads=4		, num_hidden_layers=5		, vocab_size=1005		, )

    __lowerCamelCase  							=			MultilingualCLIP(_a )
    __lowerCamelCase  							=			text_encoder.eval()

    return text_encoder




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    torch.manual_seed(0 )

    __lowerCamelCase  							=			{
        '''in_channels''': 9,
        # Out channels is double in channels because predicts mean and variance
        '''out_channels''': 8,
        '''addition_embed_type''': '''text_image''',
        '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
        '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
        '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
        '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
        '''layers_per_block''': 1,
        '''encoder_hid_dim''': self.text_embedder_hidden_size,
        '''encoder_hid_dim_type''': '''text_image_proj''',
        '''cross_attention_dim''': self.cross_attention_dim,
        '''attention_head_dim''': 4,
        '''resnet_time_scale_shift''': '''scale_shift''',
        '''class_embed_type''': None,
    }

    __lowerCamelCase  							=			UNetaDConditionModel(**_a )
    return model




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return {
        "block_out_channels": [32, 64],
        "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
        "in_channels": 3,
        "latent_channels": 4,
        "layers_per_block": 1,
        "norm_num_groups": 8,
        "norm_type": "spatial",
        "num_vq_embeddings": 12,
        "out_channels": 3,
        "up_block_types": [
            "AttnUpDecoderBlock2D",
            "UpDecoderBlock2D",
        ],
        "vq_embed_dim": 4,
    }




 @property
 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    torch.manual_seed(0 )
    __lowerCamelCase  							=			VQModel(**self.dummy_movq_kwargs )
    return model




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.dummy_text_encoder
    __lowerCamelCase  							=			self.dummy_tokenizer
    __lowerCamelCase  							=			self.dummy_unet
    __lowerCamelCase  							=			self.dummy_movq

    __lowerCamelCase  							=			DDIMScheduler(
        num_train_timesteps=1000		, beta_schedule='''linear'''		, beta_start=0.00_085		, beta_end=0.012		, clip_sample=_a		, set_alpha_to_one=_a		, steps_offset=1		, prediction_type='''epsilon'''		, thresholding=_a		, )

    __lowerCamelCase  							=			{
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''movq''': movq,
    }

    return components




 def 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase=0 ):



    '''simple docstring'''
    __lowerCamelCase  							=			floats_tensor((1, self.cross_attention_dim)		, rng=random.Random(_a ) ).to(_a )
    __lowerCamelCase  							=			floats_tensor((1, self.cross_attention_dim)		, rng=random.Random(seed + 1 ) ).to(_a )
    # create init_image
    __lowerCamelCase  							=			floats_tensor((1, 3, 64, 64)		, rng=random.Random(_a ) ).to(_a )
    __lowerCamelCase  							=			image.cpu().permute(0		, 2		, 3		, 1 )[0]
    __lowerCamelCase  							=			Image.fromarray(np.uinta(_a ) ).convert('''RGB''' ).resize((256, 256) )
    # create mask
    __lowerCamelCase  							=			np.ones((64, 64)		, dtype=np.floataa )
    __lowerCamelCase  							=			0

    if str(_a ).startswith('''mps''' ):
       __lowerCamelCase  							=			torch.manual_seed(_a )
    else:
       __lowerCamelCase  							=			torch.Generator(device=_a ).manual_seed(_a )
    __lowerCamelCase  							=			{
        '''prompt''': '''horse''',
        '''image''': init_image,
        '''mask_image''': mask,
        '''image_embeds''': image_embeds,
        '''negative_image_embeds''': negative_image_embeds,
        '''generator''': generator,
        '''height''': 64,
        '''width''': 64,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 4.0,
        '''output_type''': '''np''',
    }
    return inputs




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			'''cpu'''

    __lowerCamelCase  							=			self.get_dummy_components()

    __lowerCamelCase  							=			self.pipeline_class(**_a )
    __lowerCamelCase  							=			pipe.to(_a )

    pipe.set_progress_bar_config(disable=_a )

    __lowerCamelCase  							=			pipe(**self.get_dummy_inputs(_a ) )
    __lowerCamelCase  							=			output.images

    __lowerCamelCase  							=			pipe(
        **self.get_dummy_inputs(_a )		, return_dict=_a		, )[0]

    __lowerCamelCase  							=			image[0, -3:, -3:, -1]
    __lowerCamelCase  							=			image_from_tuple[0, -3:, -3:, -1]

    print(F"""image.shape {image.shape}""" )

    assert image.shape == (1, 64, 64, 3)

    __lowerCamelCase  							=			np.array(
        [0.8_326_919, 0.73_790_467, 0.20_918_581, 0.9_309_612, 0.5_511_791, 0.43_713_328, 0.5_513_321, 0.49_922_934, 0.59_497_786] )

    assert (
        np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
    ), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
    assert (
        np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2
    ), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    super().test_inference_batch_single_identical(expected_max_diff=3E-3 )



@slow
@require_torch_gpu
class        __lowerCAmelCase   (		unittest.TestCase      ):




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy''' )

    __lowerCamelCase  							=			load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png''' )
    __lowerCamelCase  							=			np.ones((768, 768)		, dtype=np.floataa )
    __lowerCamelCase  							=			0

    __lowerCamelCase  							=			'''a hat'''

    __lowerCamelCase  							=			KandinskyPriorPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-1-prior'''		, torch_dtype=torch.floataa )
    pipe_prior.to(_a )

    __lowerCamelCase  							=			KandinskyInpaintPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-1-inpaint'''		, torch_dtype=torch.floataa )
    __lowerCamelCase  							=			pipeline.to(_a )
    pipeline.set_progress_bar_config(disable=_a )

    __lowerCamelCase  							=			torch.Generator(device='''cpu''' ).manual_seed(0 )
    __lowerCamelCase  							=			pipe_prior(
        _a		, generator=_a		, num_inference_steps=5		, negative_prompt=''''''		, ).to_tuple()

    __lowerCamelCase  							=			pipeline(
        _a		, image=_a		, mask_image=_a		, image_embeds=_a		, negative_image_embeds=_a		, generator=_a		, num_inference_steps=100		, height=768		, width=768		, output_type='''np'''		, )

    __lowerCamelCase  							=			output.images[0]

    assert image.shape == (768, 768, 3)

    assert_mean_pixel_difference(_a		, _a )









a__				=		'''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



