



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      __lowerCAmelCase     (		UpperCAmelCase_							,     unittest.TestCase      ):
					"""simple docstring"""






					A__				:		Optional[Any]			       =    KandinskyVaaPipeline
					A__				:		int			       =    [
					    "image_embeds",
					    "negative_image_embeds",
					]
					A__				:		Union[str, Any]			       =    ["image_embeds", "negative_image_embeds"]
					A__				:		Union[str, Any]			       =    [
					    "generator",
					    "height",
					    "width",
					    "latents",
					    "guidance_scale",
					    "num_inference_steps",
					    "return_dict",
					    "guidance_scale",
					    "num_images_per_prompt",
					    "output_type",
					    "return_dict",
					]
					A__				:		Tuple			       =    False



					@property
					def     _a							(     self					:       int					):






								"""simple docstring"""


								return 32



					@property
					def     _a							(     self					:       int					):






								"""simple docstring"""


								return 32



					@property
					def     _a							(     self					:       Union[str, Any]					):






								"""simple docstring"""


								return self.time_input_dim



					@property
					def     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								return self.time_input_dim * 4



					@property
					def     _a							(     self					:       int					):






								"""simple docstring"""


								return 1_00



					@property
					def     _a							(     self					:       int					):






								"""simple docstring"""


								torch.manual_seed(0					)

								A__       =    {
								    '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,
								}

								A__       =    UNetaDConditionModel(**_snake_case					)
								return model



					@property
					def     _a							(     self					:       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     _a							(     self					:       Dict					):






								"""simple docstring"""


								torch.manual_seed(0					)
								A__       =    VQModel(**self.dummy_movq_kwargs					)
								return model



					def     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								A__       =    self.dummy_unet
								A__       =    self.dummy_movq

								A__       =    DDIMScheduler(
								    num_train_timesteps=10_00				,					beta_schedule='linear'				,					beta_start=0.0_0085				,					beta_end=0.012				,					clip_sample=_snake_case				,					set_alpha_to_one=_snake_case				,					steps_offset=1				,					prediction_type='epsilon'				,					thresholding=_snake_case				,					)

								A__       =    {
								    'unet': unet,
								    'scheduler': scheduler,
								    'movq': movq,
								}
								return components



					def     _a							(     self					:       Optional[int]				,					_snake_case					:       Optional[Any]				,					_snake_case					:       Optional[int]=0					):






								"""simple docstring"""


								A__       =    floats_tensor((1, self.text_embedder_hidden_size)				,					rng=random.Random(_snake_case					)					).to(_snake_case					)
								A__       =    floats_tensor((1, self.text_embedder_hidden_size)				,					rng=random.Random(seed + 1					)					).to(
								    _snake_case					)
								if str(_snake_case					).startswith('mps'					):
											A__       =    torch.manual_seed(_snake_case					)
								else:
											A__       =    torch.Generator(device=_snake_case					).manual_seed(_snake_case					)
								A__       =    {
								    '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     _a							(     self					:       Optional[int]					):






								"""simple docstring"""


								A__       =    'cpu'

								A__       =    self.get_dummy_components()

								A__       =    self.pipeline_class(**_snake_case					)
								A__       =    pipe.to(_snake_case					)

								pipe.set_progress_bar_config(disable=_snake_case					)

								A__       =    pipe(**self.get_dummy_inputs(_snake_case					)					)
								A__       =    output.images

								A__       =    pipe(
								    **self.get_dummy_inputs(_snake_case					)				,					return_dict=_snake_case				,					)[0]

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

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

								A__       =    np.array(
								    [0.623_7976, 1.0, 0.3644_1332, 1.0, 0.7063_9634, 0.2987_7186, 0.8565_2125, 0.521_6843, 0.5445_4046]					)

								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      __lowerCAmelCase     (		unittest.TestCase      ):
					"""simple docstring"""






					def     _a							(     self					:       List[Any]					):






								"""simple docstring"""


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



					def     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								A__       =    load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'					)

								A__       =    KandinskyVaaPriorPipeline.from_pretrained(
								    'kandinsky-community/kandinsky-2-2-prior'				,					torch_dtype=torch.floataa					)
								pipe_prior.to(_snake_case					)

								A__       =    KandinskyVaaPipeline.from_pretrained(
								    'kandinsky-community/kandinsky-2-2-decoder'				,					torch_dtype=torch.floataa					)
								A__       =    pipeline.to(_snake_case					)
								pipeline.set_progress_bar_config(disable=_snake_case					)

								A__       =    'red cat, 4k photo'

								A__       =    torch.Generator(device='cuda'					).manual_seed(0					)
								A__     , A__       =    pipe_prior(
								    _snake_case				,					generator=_snake_case				,					num_inference_steps=5				,					negative_prompt=''				,					).to_tuple()

								A__       =    torch.Generator(device='cuda'					).manual_seed(0					)
								A__       =    pipeline(
								    image_embeds=_snake_case				,					negative_image_embeds=_snake_case				,					generator=_snake_case				,					num_inference_steps=1_00				,					output_type='np'				,					)

								A__       =    output.images[0]

								assert image.shape == (5_12, 5_12, 3)

								assert_mean_pixel_difference(_snake_case				,					_snake_case					)









import pytest

import datasets


# Import fixture modules as plugins
SCREAMING_SNAKE_CASE__       =      ['''tests.fixtures.files''', '''tests.fixtures.hub''', '''tests.fixtures.fsspec''']


def A						(       __UpperCamelCase ,	__UpperCamelCase					)						->      Optional[int]:
			# Mark tests as "unit" by default if not marked as "integration" (or already marked as "unit")
			for item in items:
						if any(marker in item.keywords for marker in ['integration', 'unit']					):
									continue
						item.add_marker(pytest.mark.unit					)


def A						(       __UpperCamelCase					)						->      str:
			config.addinivalue_line('markers' ,	'torchaudio_latest: mark test to run with torchaudio>=0.12'					)


@pytest.fixture(autouse=__UpperCamelCase					)
def A						(       __UpperCamelCase ,	__UpperCamelCase					)						->      Union[str, Any]:
			# test_hf_cache_home = tmp_path_factory.mktemp("cache")  # TODO: why a cache dir per test function does not work?
			A__       =    tmp_path_factory.getbasetemp() / 'cache'
			A__       =    test_hf_cache_home / 'datasets'
			A__       =    test_hf_cache_home / 'metrics'
			A__       =    test_hf_cache_home / 'modules'
			monkeypatch.setattr('datasets.config.HF_DATASETS_CACHE' ,	str(__UpperCamelCase					)					)
			monkeypatch.setattr('datasets.config.HF_METRICS_CACHE' ,	str(__UpperCamelCase					)					)
			monkeypatch.setattr('datasets.config.HF_MODULES_CACHE' ,	str(__UpperCamelCase					)					)
			A__       =    test_hf_datasets_cache / 'downloads'
			monkeypatch.setattr('datasets.config.DOWNLOADED_DATASETS_PATH' ,	str(__UpperCamelCase					)					)
			A__       =    test_hf_datasets_cache / 'downloads' / 'extracted'
			monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_PATH' ,	str(__UpperCamelCase					)					)


@pytest.fixture(autouse=__UpperCamelCase ,	scope='session'					)
def A						(       )						->      Union[str, Any]:
			datasets.disable_progress_bar()


@pytest.fixture(autouse=__UpperCamelCase					)
def A						(       __UpperCamelCase					)						->      int:
			# don't take tests into account when counting downloads
			monkeypatch.setattr('datasets.config.HF_UPDATE_DOWNLOAD_COUNTS' ,	__UpperCamelCase					)






@pytest.fixture
def A						(       __UpperCamelCase					)						->      Any:
			# Required to suppress RemovedIn20Warning when feature(s) are not compatible with SQLAlchemy 2.0
			# To be removed once SQLAlchemy 2.0 supported
			monkeypatch.setattr('sqlalchemy.util.deprecations.SILENCE_UBER_WARNING' ,	__UpperCamelCase					)




