





import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, 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 SCREAMING_SNAKE_CASE_  (   snake_case_    ,			unittest.TestCase   ):
						__magic_name__:     str	 =    KandinskyVaaPipeline
						__magic_name__:     str	 =    [
						    "image_embeds",
						    "negative_image_embeds",
						]
						__magic_name__:     str	 =    ["image_embeds", "negative_image_embeds"]
						__magic_name__:     int	 =    [
						    "generator",
						    "height",
						    "width",
						    "latents",
						    "guidance_scale",
						    "num_inference_steps",
						    "return_dict",
						    "guidance_scale",
						    "num_images_per_prompt",
						    "output_type",
						    "return_dict",
						]
						__magic_name__:     Dict	 =    False






						@property
						def UpperCAmelCase_			( self : Tuple		)							->      int:







							"""simple docstring"""


							return 32






						@property
						def UpperCAmelCase_			( self : Tuple		)							->      Any:







							"""simple docstring"""


							return 32






						@property
						def UpperCAmelCase_			( self : List[str]		)							->      Tuple:







							"""simple docstring"""


							return self.time_input_dim






						@property
						def UpperCAmelCase_			( self : Optional[Any]		)							->      str:







							"""simple docstring"""


							return self.time_input_dim * 4






						@property
						def UpperCAmelCase_			( self : List[Any]		)							->      Any:







							"""simple docstring"""


							return 100






						@property
						def UpperCAmelCase_			( self : List[str]		)							->      Union[str, Any]:







							"""simple docstring"""


							torch.manual_seed(0		)

							snake_case_						:     Dict     =		{
							    'in_channels': 4,
							    # Out channels is double in channels because predicts mean and variance
							    'out_channels': 8,
							    'addition_embed_type': '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': 'image_proj',
							    'cross_attention_dim': self.cross_attention_dim,
							    'attention_head_dim': 4,
							    'resnet_time_scale_shift': 'scale_shift',
							    'class_embed_type': None,
							}

							snake_case_						:     Dict     =		UNetaDConditionModel(**_A		)
							return model






						@property
						def UpperCAmelCase_			( self : List[Any]		)							->      List[str]:







							"""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 UpperCAmelCase_			( self : str		)							->      int:







							"""simple docstring"""


							torch.manual_seed(0		)
							snake_case_						:     List[str]     =		VQModel(**self.dummy_movq_kwargs		)
							return model






						def UpperCAmelCase_			( self : Dict		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     Any     =		self.dummy_unet
							snake_case_						:     List[Any]     =		self.dummy_movq

							snake_case_						:     Any     =		DDIMScheduler(
							    num_train_timesteps=1000						,				beta_schedule='linear'						,				beta_start=0.0_0_0_8_5						,				beta_end=0.0_1_2						,				clip_sample=_A						,				set_alpha_to_one=_A						,				steps_offset=1						,				prediction_type='epsilon'						,				thresholding=_A						,				)

							snake_case_						:     Dict     =		{
							    'unet': unet,
							    'scheduler': scheduler,
							    'movq': movq,
							}
							return components






						def UpperCAmelCase_			( self : List[Any]						,				_A : Tuple						,				_A : List[str]=0		)							->      int:







							"""simple docstring"""


							snake_case_						:     Union[str, Any]     =		floats_tensor((1, self.text_embedder_hidden_size)						,				rng=random.Random(_A		)		).to(_A		)
							snake_case_						:     Tuple     =		floats_tensor((1, self.text_embedder_hidden_size)						,				rng=random.Random(seed + 1		)		).to(
							    _A		)
							if str(_A		).startswith('mps'		):
								snake_case_						:     Union[str, Any]     =		torch.manual_seed(_A		)
							else:
								snake_case_						:     Optional[Any]     =		torch.Generator(device=_A		).manual_seed(_A		)
							snake_case_						:     List[Any]     =		{
							    'image_embeds': image_embeds,
							    'negative_image_embeds': negative_image_embeds,
							    'generator': generator,
							    'height': 64,
							    'width': 64,
							    'guidance_scale': 4.0,
							    'num_inference_steps': 2,
							    'output_type': 'np',
							}
							return inputs






						def UpperCAmelCase_			( self : Dict		)							->      str:







							"""simple docstring"""


							snake_case_						:     List[Any]     =		'cpu'

							snake_case_						:     List[str]     =		self.get_dummy_components()

							snake_case_						:     Union[str, Any]     =		self.pipeline_class(**_A		)
							snake_case_						:     Dict     =		pipe.to(_A		)

							pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     List[str]     =		pipe(**self.get_dummy_inputs(_A		)		)
							snake_case_						:     Any     =		output.images

							snake_case_						:     int     =		pipe(
							    **self.get_dummy_inputs(_A		)						,				return_dict=_A						,				)[0]

							snake_case_						:     List[Any]     =		image[0, -3:, -3:, -1]
							snake_case_						:     Dict     =		image_from_tuple[0, -3:, -3:, -1]

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

							snake_case_						:     List[str]     =		np.array(
							    [0.6_2_3_7_9_7_6, 1.0, 0.3_6_4_4_1_3_3_2, 1.0, 0.7_0_6_3_9_6_3_4, 0.2_9_8_7_7_1_8_6, 0.8_5_6_5_2_1_2_5, 0.5_2_1_6_8_4_3, 0.5_4_4_5_4_0_4_6]		)

							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()}"""
@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):






						def UpperCAmelCase_			( self : Tuple		)							->      Dict:







							"""simple docstring"""


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






						def UpperCAmelCase_			( self : Any		)							->      Optional[int]:







							"""simple docstring"""


							snake_case_						:     List[str]     =		load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'		)

							snake_case_						:     int     =		KandinskyVaaPriorPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-prior'						,				torch_dtype=torch.floataa		)
							pipe_prior.to(_A		)

							snake_case_						:     int     =		KandinskyVaaPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-decoder'						,				torch_dtype=torch.floataa		)
							snake_case_						:     Tuple     =		pipeline.to(_A		)
							pipeline.set_progress_bar_config(disable=_A		)

							snake_case_						:     int     =		'red cat, 4k photo'

							snake_case_						:     List[Any]     =		torch.Generator(device='cuda'		).manual_seed(0		)
							snake_case_  ,snake_case_						:     Tuple     =		pipe_prior(
							    _A						,				generator=_A						,				num_inference_steps=5						,				negative_prompt=''						,				).to_tuple()

							snake_case_						:     str     =		torch.Generator(device='cuda'		).manual_seed(0		)
							snake_case_						:     Dict     =		pipeline(
							    image_embeds=_A						,				negative_image_embeds=_A						,				generator=_A						,				num_inference_steps=100						,				output_type='np'						,				)

							snake_case_						:     Tuple     =		output.images[0]

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

							assert_mean_pixel_difference(_A						,				_A		)








_SCREAMING_SNAKE_CASE 		=						{
    """Pillow""": """Pillow""",
    """accelerate""": """accelerate>=0.11.0""",
    """compel""": """compel==0.1.8""",
    """black""": """black~=23.1""",
    """datasets""": """datasets""",
    """filelock""": """filelock""",
    """flax""": """flax>=0.4.1""",
    """hf-doc-builder""": """hf-doc-builder>=0.3.0""",
    """huggingface-hub""": """huggingface-hub>=0.13.2""",
    """requests-mock""": """requests-mock==1.10.0""",
    """importlib_metadata""": """importlib_metadata""",
    """invisible-watermark""": """invisible-watermark""",
    """isort""": """isort>=5.5.4""",
    """jax""": """jax>=0.2.8,!=0.3.2""",
    """jaxlib""": """jaxlib>=0.1.65""",
    """Jinja2""": """Jinja2""",
    """k-diffusion""": """k-diffusion>=0.0.12""",
    """torchsde""": """torchsde""",
    """note_seq""": """note_seq""",
    """librosa""": """librosa""",
    """numpy""": """numpy""",
    """omegaconf""": """omegaconf""",
    """parameterized""": """parameterized""",
    """protobuf""": """protobuf>=3.20.3,<4""",
    """pytest""": """pytest""",
    """pytest-timeout""": """pytest-timeout""",
    """pytest-xdist""": """pytest-xdist""",
    """ruff""": """ruff>=0.0.241""",
    """safetensors""": """safetensors""",
    """sentencepiece""": """sentencepiece>=0.1.91,!=0.1.92""",
    """scipy""": """scipy""",
    """onnx""": """onnx""",
    """regex""": """regex!=2019.12.17""",
    """requests""": """requests""",
    """tensorboard""": """tensorboard""",
    """torch""": """torch>=1.4""",
    """torchvision""": """torchvision""",
    """transformers""": """transformers>=4.25.1""",
    """urllib3""": """urllib3<=2.0.0""",
}

