





"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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  _a						(      _lowerCAmelCase	,				unittest.TestCase				):
				UpperCamelCase								=    KandinskyPipeline
				UpperCamelCase								=    [
				    """prompt""",
				    """image_embeds""",
				    """negative_image_embeds""",
				]
				UpperCamelCase								=    ["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
				UpperCamelCase								=    [
				    """generator""",
				    """height""",
				    """width""",
				    """latents""",
				    """guidance_scale""",
				    """negative_prompt""",
				    """num_inference_steps""",
				    """return_dict""",
				    """guidance_scale""",
				    """num_images_per_prompt""",
				    """output_type""",
				    """return_dict""",
				]
				UpperCamelCase								=    False






				@property
				def 				snake_case	( self		:    Optional[int]							)	->      str:




										'''simple docstring'''



										return 3_2






				@property
				def 				snake_case	( self		:    Dict							)	->      Any:




										'''simple docstring'''



										return 3_2






				@property
				def 				snake_case	( self		:    Optional[int]							)	->      Tuple:




										'''simple docstring'''



										return self.time_input_dim






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




										'''simple docstring'''



										return self.time_input_dim * 4






				@property
				def 				snake_case	( self		:    str							)	->      str:




										'''simple docstring'''



										return 1_0_0






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




										'''simple docstring'''



										_UpperCamelCase			:					List[str]     =    XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''							)
										return tokenizer






				@property
				def 				snake_case	( self		:    Tuple							)	->      Tuple:




										'''simple docstring'''



										torch.manual_seed(0							)
										_UpperCamelCase			:					Optional[Any]     =    MCLIPConfig(
										    numDims=self.cross_attention_dim, transformerDimensions=self.text_embedder_hidden_size, hidden_size=self.text_embedder_hidden_size, intermediate_size=3_7, num_attention_heads=4, num_hidden_layers=5, vocab_size=1_0_0_5, )

										_UpperCamelCase			:					str     =    MultilingualCLIP(_lowercase							)
										_UpperCamelCase			:					Dict     =    text_encoder.eval()

										return text_encoder






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




										'''simple docstring'''



										torch.manual_seed(0							)

										_UpperCamelCase			:					Dict     =    {
										    '''in_channels''': 4,
										    # 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,
										}

										_UpperCamelCase			:					Dict     =    UNetaDConditionModel(**_lowercase							)
										return model






				@property
				def 				snake_case	( self		:    Optional[Any]							)	->      Tuple:




										'''simple docstring'''



										return {
										    "block_out_channels": [3_2, 6_4],
										    "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": 1_2,
										    "out_channels": 3,
										    "up_block_types": [
										        "AttnUpDecoderBlock2D",
										        "UpDecoderBlock2D",
										    ],
										    "vq_embed_dim": 4,
										}






				@property
				def 				snake_case	( self		:    Optional[int]							)	->      str:




										'''simple docstring'''



										torch.manual_seed(0							)
										_UpperCamelCase			:					Tuple     =    VQModel(**self.dummy_movq_kwargs							)
										return model






				def 				snake_case	( self		:    Optional[Any]							)	->      Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase			:					Optional[Any]     =    self.dummy_text_encoder
										_UpperCamelCase			:					List[str]     =    self.dummy_tokenizer
										_UpperCamelCase			:					Dict     =    self.dummy_unet
										_UpperCamelCase			:					str     =    self.dummy_movq

										_UpperCamelCase			:					str     =    DDIMScheduler(
										    num_train_timesteps=1_0_0_0, beta_schedule='''linear''', beta_start=0.00_085, beta_end=0.012, clip_sample=_lowercase, set_alpha_to_one=_lowercase, steps_offset=1, prediction_type='''epsilon''', thresholding=_lowercase, )

										_UpperCamelCase			:					Any     =    {
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										    '''movq''': movq,
										}
										return components






				def 				snake_case	( self		:    Dict, lowerCAmelCase__		:    int, lowerCAmelCase__		:    Union[str, Any]=0							)	->      Dict:




										'''simple docstring'''



										_UpperCamelCase			:					int     =    floats_tensor((1, self.cross_attention_dim), rng=random.Random(_lowercase							)							).to(_lowercase							)
										_UpperCamelCase			:					Dict     =    floats_tensor((1, self.cross_attention_dim), rng=random.Random(seed + 1							)							).to(_lowercase							)
										if str(_lowercase							).startswith('''mps'''							):
																_UpperCamelCase			:					Tuple     =    torch.manual_seed(_lowercase							)
										else:
																_UpperCamelCase			:					Union[str, Any]     =    torch.Generator(device=_lowercase							).manual_seed(_lowercase							)
										_UpperCamelCase			:					Optional[Any]     =    {
										    '''prompt''': '''horse''',
										    '''image_embeds''': image_embeds,
										    '''negative_image_embeds''': negative_image_embeds,
										    '''generator''': generator,
										    '''height''': 6_4,
										    '''width''': 6_4,
										    '''guidance_scale''': 4.0,
										    '''num_inference_steps''': 2,
										    '''output_type''': '''np''',
										}
										return inputs






				def 				snake_case	( self		:    Optional[Any]							)	->      Any:




										'''simple docstring'''



										_UpperCamelCase			:					Optional[int]     =    '''cpu'''

										_UpperCamelCase			:					Any     =    self.get_dummy_components()

										_UpperCamelCase			:					int     =    self.pipeline_class(**_lowercase							)
										_UpperCamelCase			:					Dict     =    pipe.to(_lowercase							)

										pipe.set_progress_bar_config(disable=_lowercase							)

										_UpperCamelCase			:					Tuple     =    pipe(**self.get_dummy_inputs(_lowercase							)							)
										_UpperCamelCase			:					Union[str, Any]     =    output.images

										_UpperCamelCase			:					Union[str, Any]     =    pipe(
										    **self.get_dummy_inputs(_lowercase							), return_dict=_lowercase, )[0]

										_UpperCamelCase			:					Union[str, Any]     =    image[0, -3:, -3:, -1]
										_UpperCamelCase			:					List[str]     =    image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 6_4, 6_4, 3)

										_UpperCamelCase			:					List[Any]     =    np.array(
										    [0.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]							)

										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  _a						(      unittest.TestCase				):






				def 				snake_case	( self		:    Tuple							)	->      int:




										'''simple docstring'''



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






				def 				snake_case	( self		:    Dict							)	->      Optional[int]:




										'''simple docstring'''



										_UpperCamelCase			:					Dict     =    load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/kandinsky/kandinsky_text2img_cat_fp16.npy'''							)

										_UpperCamelCase			:					int     =    KandinskyPriorPipeline.from_pretrained(
										    '''kandinsky-community/kandinsky-2-1-prior''', torch_dtype=torch.floataa							)
										pipe_prior.to(_lowercase							)

										_UpperCamelCase			:					Tuple     =    KandinskyPipeline.from_pretrained('''kandinsky-community/kandinsky-2-1''', torch_dtype=torch.floataa							)
										_UpperCamelCase			:					Dict     =    pipeline.to(_lowercase							)
										pipeline.set_progress_bar_config(disable=_lowercase							)

										_UpperCamelCase			:					Dict     =    '''red cat, 4k photo'''

										_UpperCamelCase			:					int     =    torch.Generator(device='''cuda'''							).manual_seed(0							)
										_UpperCamelCase   ,							_UpperCamelCase			:					Optional[int]     =    pipe_prior(
										    _lowercase, generator=_lowercase, num_inference_steps=5, negative_prompt='''''', ).to_tuple()

										_UpperCamelCase			:					Any     =    torch.Generator(device='''cuda'''							).manual_seed(0							)
										_UpperCamelCase			:					Optional[Any]     =    pipeline(
										    _lowercase, image_embeds=_lowercase, negative_image_embeds=_lowercase, generator=_lowercase, num_inference_steps=1_0_0, output_type='''np''', )

										_UpperCamelCase			:					List[str]     =    output.images[0]

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

										assert_mean_pixel_difference(_lowercase, _lowercase							)






from typing import Callable




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       )->		List[Any]:



    """simple docstring"""






    def wrapper(SCREAMING_SNAKE_CASE_       ):
        UpperCamelCase_          =		function_with_docstring.__doc__
        return target_function

    return wrapper
