






"""simple docstring"""


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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 						__UpperCamelCase     (     UpperCamelCase   ,     unittest.TestCase						):
					lowercase_    :				Dict      =  KandinskyVaaImgaImgPipeline
					lowercase_    :				str      =  ["""image_embeds""", """negative_image_embeds""", """image"""]
					lowercase_    :				Optional[int]      =  [
					    """image_embeds""",
					    """negative_image_embeds""",
					    """image""",
					]
					lowercase_    :				Union[str, Any]      =  [
					    """generator""",
					    """height""",
					    """width""",
					    """strength""",
					    """guidance_scale""",
					    """num_inference_steps""",
					    """return_dict""",
					    """guidance_scale""",
					    """num_images_per_prompt""",
					    """output_type""",
					    """return_dict""",
					]
					lowercase_    :				str      =  False
					@property
					def        UpperCAmelCase__ (							self							:	int   )	->       int:
											return 32
					@property
					def        UpperCAmelCase__ (							self							:	List[str]   )	->       Any:
											return 32
					@property
					def        UpperCAmelCase__ (							self							:	int   )	->       List[str]:
											return self.time_input_dim
					@property
					def        UpperCAmelCase__ (							self							:	List[str]   )	->       List[str]:
											return self.time_input_dim * 4
					@property
					def        UpperCAmelCase__ (							self							:	List[Any]   )	->       Any:
											return 100
					@property
					def        UpperCAmelCase__ (							self							:	Any   )	->       List[Any]:
											torch.manual_seed(0   )

											lowerCAmelCase   :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,
											}

											lowerCAmelCase   :Optional[int]						=			UNetaDConditionModel(**UpperCAmelCase   )
											return model
					@property
					def        UpperCAmelCase__ (							self							:	str   )	->       List[Any]:
											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							:	List[str]   )	->       int:
											torch.manual_seed(0   )
											lowerCAmelCase   :int						=			VQModel(**self.dummy_movq_kwargs   )
											return model
					def        UpperCAmelCase__ (							self							:	Dict   )	->       Dict:
											lowerCAmelCase   :Any						=			self.dummy_unet
											lowerCAmelCase   :List[str]						=			self.dummy_movq

											lowerCAmelCase   :Dict						=			{
											    'num_train_timesteps': 1000,
											    'beta_schedule': 'linear',
											    'beta_start': 0.0_0_0_8_5,
											    'beta_end': 0.0_1_2,
											    'clip_sample': False,
											    'set_alpha_to_one': False,
											    'steps_offset': 0,
											    'prediction_type': 'epsilon',
											    'thresholding': False,
											}

											lowerCAmelCase   :Union[str, Any]						=			DDIMScheduler(**UpperCAmelCase   )

											lowerCAmelCase   :Dict						=			{
											    'unet': unet,
											    'scheduler': scheduler,
											    'movq': movq,
											}

											return components
					def        UpperCAmelCase__ (							self							:	str     , UpperCAmelCase							:	int     , UpperCAmelCase							:	Optional[Any]=0   )	->       Tuple:
											lowerCAmelCase   :int						=			floats_tensor((1, self.text_embedder_hidden_size)     , rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Dict						=			floats_tensor((1, self.text_embedder_hidden_size)     , rng=random.Random(seed + 1   )   ).to(
											    UpperCAmelCase   )
											# create init_image
											lowerCAmelCase   :Union[str, Any]						=			floats_tensor((1, 3, 64, 64)     , rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
											lowerCAmelCase   :Tuple						=			image.cpu().permute(0     , 2     , 3     , 1   )[0]
											lowerCAmelCase   :List[str]						=			Image.fromarray(np.uinta(UpperCAmelCase   )   ).convert('RGB'   ).resize((256, 256)   )

											if str(UpperCAmelCase   ).startswith('mps'   ):
																	lowerCAmelCase   :int						=			torch.manual_seed(UpperCAmelCase   )
											else:
																	lowerCAmelCase   :Optional[Any]						=			torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
											lowerCAmelCase   :Optional[int]						=			{
											    'image': init_image,
											    'image_embeds': image_embeds,
											    'negative_image_embeds': negative_image_embeds,
											    'generator': generator,
											    'height': 64,
											    'width': 64,
											    'num_inference_steps': 10,
											    'guidance_scale': 7.0,
											    'strength': 0.2,
											    'output_type': 'np',
											}
											return inputs
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       List[str]:
											lowerCAmelCase   :Dict						=			'cpu'

											lowerCAmelCase   :str						=			self.get_dummy_components()

											lowerCAmelCase   :List[str]						=			self.pipeline_class(**UpperCAmelCase   )
											lowerCAmelCase   :Tuple						=			pipe.to(UpperCAmelCase   )

											pipe.set_progress_bar_config(disable=UpperCAmelCase   )

											lowerCAmelCase   :Dict						=			pipe(**self.get_dummy_inputs(UpperCAmelCase   )   )
											lowerCAmelCase   :str						=			output.images

											lowerCAmelCase   :List[str]						=			pipe(
											    **self.get_dummy_inputs(UpperCAmelCase   )     , return_dict=UpperCAmelCase     , )[0]

											lowerCAmelCase   :Optional[Any]						=			image[0, -3:, -3:, -1]
											lowerCAmelCase   :Optional[Any]						=			image_from_tuple[0, -3:, -3:, -1]

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

											lowerCAmelCase   :Any						=			np.array(
											    [0.6_1_9_9_7_7_8, 0.6_3_9_8_4_4_0_6, 0.4_6_1_4_5_7_8_5, 0.6_2_9_4_4_9_8_4, 0.5_6_2_2_2_1_5, 0.4_7_3_0_6_1_3_2, 0.4_7_4_4_1_4_5_6, 0.4_6_0_7_6_0_6, 0.4_8_7_1_9_2_6_3]   )
											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 						__UpperCamelCase     (     unittest.TestCase						):
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       List[str]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
					def        UpperCAmelCase__ (							self							:	List[Any]   )	->       int:
											lowerCAmelCase   :str						=			load_numpy(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
											    '/kandinskyv22/kandinskyv22_img2img_frog.npy'   )

											lowerCAmelCase   :int						=			load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'   )
											lowerCAmelCase   :Optional[int]						=			'A red cartoon frog, 4k'

											lowerCAmelCase   :Union[str, Any]						=			KandinskyVaaPriorPipeline.from_pretrained(
											    'kandinsky-community/kandinsky-2-2-prior'     , torch_dtype=torch.floataa   )
											pipe_prior.to(UpperCAmelCase   )

											lowerCAmelCase   :str						=			KandinskyVaaImgaImgPipeline.from_pretrained(
											    'kandinsky-community/kandinsky-2-2-decoder'     , torch_dtype=torch.floataa   )
											lowerCAmelCase   :Optional[Any]						=			pipeline.to(UpperCAmelCase   )

											pipeline.set_progress_bar_config(disable=UpperCAmelCase   )

											lowerCAmelCase   :Optional[int]						=			torch.Generator(device='cpu'   ).manual_seed(0   )
											lowerCAmelCase	, lowerCAmelCase   :Tuple						=			pipe_prior(
											    UpperCAmelCase     , generator=UpperCAmelCase     , num_inference_steps=5     , negative_prompt=''     , ).to_tuple()

											lowerCAmelCase   :Tuple						=			pipeline(
											    image=UpperCAmelCase     , image_embeds=UpperCAmelCase     , negative_image_embeds=UpperCAmelCase     , generator=UpperCAmelCase     , num_inference_steps=100     , height=768     , width=768     , strength=0.2     , output_type='np'     , )

											lowerCAmelCase   :Tuple						=			output.images[0]

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

											assert_mean_pixel_difference(UpperCAmelCase     , UpperCAmelCase   )







"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


__SCREAMING_SNAKE_CASE   =   {
    'configuration_transfo_xl': ['TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TransfoXLConfig'],
    'tokenization_transfo_xl': ['TransfoXLCorpus', 'TransfoXLTokenizer'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE   =   [
		    'TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'AdaptiveEmbedding',
		    'TransfoXLForSequenceClassification',
		    'TransfoXLLMHeadModel',
		    'TransfoXLModel',
		    'TransfoXLPreTrainedModel',
		    'load_tf_weights_in_transfo_xl',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE   =   [
		    'TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'TFAdaptiveEmbedding',
		    'TFTransfoXLForSequenceClassification',
		    'TFTransfoXLLMHeadModel',
		    'TFTransfoXLMainLayer',
		    'TFTransfoXLModel',
		    'TFTransfoXLPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
		from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_transfo_xl import (
				    TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
				    AdaptiveEmbedding,
				    TransfoXLForSequenceClassification,
				    TransfoXLLMHeadModel,
				    TransfoXLModel,
				    TransfoXLPreTrainedModel,
				    load_tf_weights_in_transfo_xl,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_transfo_xl import (
				    TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFAdaptiveEmbedding,
				    TFTransfoXLForSequenceClassification,
				    TFTransfoXLLMHeadModel,
				    TFTransfoXLMainLayer,
				    TFTransfoXLModel,
				    TFTransfoXLPreTrainedModel,
				)

else:
		import sys

		__SCREAMING_SNAKE_CASE   =   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)