





"""simple docstring"""





import json
import os
from typing import Optional

import numpy as np

from ...feature_extraction_utils import BatchFeature
from ...processing_utils import ProcessorMixin
from ...utils import logging
from ...utils.hub import get_file_from_repo
from ..auto import AutoTokenizer


lowerCAmelCase__   	=	logging.get_logger(__name__)





class 				_lowerCamelCase					(				UpperCamelCase__       ):
						UpperCAmelCase_       =	'''AutoTokenizer'''
						UpperCAmelCase_       =	['''tokenizer''']

						UpperCAmelCase_       =	{
						    '''semantic_prompt''': 1,
						    '''coarse_prompt''': 2,
						    '''fine_prompt''': 2,
						}


						def __init__(self      ,							__a      ,							__a=None						)       ->      str:
								super().__init__(UpperCamelCase_						)

								UpperCamelCase      =		speaker_embeddings


						@classmethod
						def 			snake_case_      (cls      ,							__a      ,							__a="speaker_embeddings_path.json"      ,							**__a						)       ->      Tuple:







								if speaker_embeddings_dict_path is not None:
										UpperCamelCase      =		get_file_from_repo(
										    UpperCamelCase_      ,							UpperCamelCase_      ,							subfolder=kwargs.pop("subfolder"      ,							UpperCamelCase_						)      ,							cache_dir=kwargs.pop("cache_dir"      ,							UpperCamelCase_						)      ,							force_download=kwargs.pop("force_download"      ,							UpperCamelCase_						)      ,							proxies=kwargs.pop("proxies"      ,							UpperCamelCase_						)      ,							resume_download=kwargs.pop("resume_download"      ,							UpperCamelCase_						)      ,							local_files_only=kwargs.pop("local_files_only"      ,							UpperCamelCase_						)      ,							use_auth_token=kwargs.pop("use_auth_token"      ,							UpperCamelCase_						)      ,							revision=kwargs.pop("revision"      ,							UpperCamelCase_						)      ,							)
										if speaker_embeddings_path is None:
												logger.warning(
												    F"`{os.path.join(UpperCamelCase_      ,							UpperCamelCase_						)}` does not exists\n                    , no preloaded speaker embeddings will be used - Make sure to provide a correct path to the json\n                    dictionnary if wanted, otherwise set `speaker_embeddings_dict_path=None`."						)
												UpperCamelCase      =		None
										else:
												with open(UpperCamelCase_						) as speaker_embeddings_json:
														UpperCamelCase      =		json.load(UpperCamelCase_						)
								else:
										UpperCamelCase      =		None

								UpperCamelCase      =		AutoTokenizer.from_pretrained(UpperCamelCase_      ,							**UpperCamelCase_						)

								return cls(tokenizer=UpperCamelCase_      ,							speaker_embeddings=UpperCamelCase_						)


						def 			snake_case_      (self      ,							__a      ,							__a="speaker_embeddings_path.json"      ,							__a="speaker_embeddings"      ,							__a = False      ,							**__a      ,							)       ->      Optional[Any]:







								if self.speaker_embeddings is not None:
										os.makedirs(os.path.join(UpperCamelCase_      ,							UpperCamelCase_      ,							"v2"						)      ,							exist_ok=UpperCamelCase_						)

										UpperCamelCase      =		{}

										UpperCamelCase      =		save_directory

										for prompt_key in self.speaker_embeddings:
												if prompt_key != "repo_or_path":
														UpperCamelCase      =		self._load_voice_preset(UpperCamelCase_						)

														UpperCamelCase      =		{}
														for key in self.speaker_embeddings[prompt_key]:
																np.save(
																    os.path.join(
																        embeddings_dict["repo_or_path"]      ,							UpperCamelCase_      ,							F"{prompt_key}_{key}"						)      ,							voice_preset[key]      ,							allow_pickle=UpperCamelCase_      ,							)
																UpperCamelCase      =		os.path.join(UpperCamelCase_      ,							F"{prompt_key}_{key}.npy"						)

														UpperCamelCase      =		tmp_dict

										with open(os.path.join(UpperCamelCase_      ,							UpperCamelCase_						)      ,							"w"						) as fp:
												json.dump(UpperCamelCase_      ,							UpperCamelCase_						)

								super().save_pretrained(UpperCamelCase_      ,							UpperCamelCase_      ,							**UpperCamelCase_						)


						def 			snake_case_      (self      ,							__a = None      ,							**__a						)       ->      Tuple:
								UpperCamelCase      =		self.speaker_embeddings[voice_preset]

								UpperCamelCase      =		{}
								for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
										if key not in voice_preset_paths:
												raise ValueError(
												    F"Voice preset unrecognized, missing {key} as a key in self.speaker_embeddings[{voice_preset}]."						)

										UpperCamelCase      =		get_file_from_repo(
										    self.speaker_embeddings.get("repo_or_path"      ,							"/"						)      ,							voice_preset_paths[key]      ,							subfolder=kwargs.pop("subfolder"      ,							UpperCamelCase_						)      ,							cache_dir=kwargs.pop("cache_dir"      ,							UpperCamelCase_						)      ,							force_download=kwargs.pop("force_download"      ,							UpperCamelCase_						)      ,							proxies=kwargs.pop("proxies"      ,							UpperCamelCase_						)      ,							resume_download=kwargs.pop("resume_download"      ,							UpperCamelCase_						)      ,							local_files_only=kwargs.pop("local_files_only"      ,							UpperCamelCase_						)      ,							use_auth_token=kwargs.pop("use_auth_token"      ,							UpperCamelCase_						)      ,							revision=kwargs.pop("revision"      ,							UpperCamelCase_						)      ,							)
										if path is None:
												raise ValueError(
												    F"`{os.path.join(self.speaker_embeddings.get('repo_or_path'      ,							'/'						)      ,							voice_preset_paths[key]						)}` does not exists\n                    , no preloaded voice preset will be used - Make sure to provide correct paths to the {voice_preset}\n                    embeddings."						)

										UpperCamelCase      =		np.load(UpperCamelCase_						)

								return voice_preset_dict


						def 			snake_case_      (self      ,							__a = None						)       ->      List[str]:







								for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
										if key not in voice_preset:
												raise ValueError(F"Voice preset unrecognized, missing {key} as a key."						)

										if not isinstance(voice_preset[key]      ,							np.ndarray						):
												raise ValueError(F"{key} voice preset must be a {str(self.preset_shape[key]						)}D ndarray."						)

										if len(voice_preset[key].shape						) != self.preset_shape[key]:
												raise ValueError(F"{key} voice preset must be a {str(self.preset_shape[key]						)}D ndarray."						)




						def __call__(self      ,							__a=None      ,							__a=None      ,							__a="pt"      ,							__a=2_56      ,							__a=False      ,							__a=True      ,							__a=False      ,							**__a      ,							)       ->      List[str]:







								if voice_preset is not None and not isinstance(UpperCamelCase_      ,							UpperCamelCase_						):
										if (
										    isinstance(UpperCamelCase_      ,							UpperCamelCase_						)
										    and self.speaker_embeddings is not None
										    and voice_preset in self.speaker_embeddings
										):
												UpperCamelCase      =		self._load_voice_preset(UpperCamelCase_						)

										else:
												if isinstance(UpperCamelCase_      ,							UpperCamelCase_						) and not voice_preset.endswith(".npz"						):
														UpperCamelCase      =		voice_preset + ".npz"

												UpperCamelCase      =		np.load(UpperCamelCase_						)

								if voice_preset is not None:
										self._validate_voice_preset_dict(UpperCamelCase_      ,							**UpperCamelCase_						)
										UpperCamelCase      =		BatchFeature(data=UpperCamelCase_      ,							tensor_type=UpperCamelCase_						)

								UpperCamelCase      =		self.tokenizer(
								    UpperCamelCase_      ,							return_tensors=UpperCamelCase_      ,							padding="max_length"      ,							max_length=UpperCamelCase_      ,							return_attention_mask=UpperCamelCase_      ,							return_token_type_ids=UpperCamelCase_      ,							add_special_tokens=UpperCamelCase_      ,							**UpperCamelCase_      ,							)

								if voice_preset is not None:
										UpperCamelCase      =		voice_preset

								return encoded_text










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        _a							(			UpperCamelCase__     ,	unittest.TestCase			):
							_lowercase	:			Union[str, Any]        = KandinskyVaaImgaImgPipeline
							_lowercase	:			Tuple        = ['''image_embeds''', '''negative_image_embeds''', '''image''']
							_lowercase	:			Any        = [
							    '''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	:			Optional[Any]        = False

							@property
							def       lowerCamelCase_ (							self:	Union[str, Any]							)			->   Dict:





									"""simple docstring"""







									return 32

							@property
							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Optional[Any]:





									"""simple docstring"""







									return 32

							@property
							def       lowerCamelCase_ (							self:	Any							)			->   Any:





									"""simple docstring"""







									return self.time_input_dim

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





									"""simple docstring"""







									return self.time_input_dim * 4

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





									"""simple docstring"""







									return 100

							@property
							def       lowerCamelCase_ (							self:	int							)			->   int:





									"""simple docstring"""







									torch.manual_seed(0							)

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

									lowercase__            =      UNetaDConditionModel(**UpperCamelCase_							)
									return model

							@property
							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Union[str, Any]:





									"""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:	Optional[Any]							)			->   int:





									"""simple docstring"""







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

							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      self.dummy_unet
									lowercase__            =      self.dummy_movq

									lowercase__            =      {
									    '''num_train_timesteps''': 1_000,
									    '''beta_schedule''': '''linear''',
									    '''beta_start''': 0.00085,
									    '''beta_end''': 0.012,
									    '''clip_sample''': False,
									    '''set_alpha_to_one''': False,
									    '''steps_offset''': 0,
									    '''prediction_type''': '''epsilon''',
									    '''thresholding''': False,
									}

									lowercase__            =      DDIMScheduler(**UpperCamelCase_							)

									lowercase__            =      {
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''movq''': movq,
									}

									return components

							def       lowerCamelCase_ (							self:	Optional[int]    ,						UpperCamelCase_:	Optional[Any]    ,						UpperCamelCase_:	Optional[int]=0							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      floats_tensor((1, self.text_embedder_hidden_size)    ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)
									lowercase__            =      floats_tensor((1, self.text_embedder_hidden_size)    ,						rng=random.Random(seed + 1							)							).to(
									    UpperCamelCase_							)
									# create init_image
									lowercase__            =      floats_tensor((1, 3, 64, 64)    ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)
									lowercase__            =      image.cpu().permute(0    ,						2    ,						3    ,						1							)[0]
									lowercase__            =      Image.fromarray(np.uinta(UpperCamelCase_							)							).convert('''RGB'''							).resize((256, 256)							)

									if str(UpperCamelCase_							).startswith('''mps'''							):
											lowercase__            =      torch.manual_seed(UpperCamelCase_							)
									else:
											lowercase__            =      torch.Generator(device=UpperCamelCase_							).manual_seed(UpperCamelCase_							)
									lowercase__            =      {
									    '''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       lowerCamelCase_ (							self:	Optional[int]							)			->   Dict:





									"""simple docstring"""







									lowercase__            =      '''cpu'''

									lowercase__            =      self.get_dummy_components()

									lowercase__            =      self.pipeline_class(**UpperCamelCase_							)
									lowercase__            =      pipe.to(UpperCamelCase_							)

									pipe.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__            =      pipe(**self.get_dummy_inputs(UpperCamelCase_							)							)
									lowercase__            =      output.images

									lowercase__            =      pipe(
									    **self.get_dummy_inputs(UpperCamelCase_							)    ,						return_dict=UpperCamelCase_    ,						)[0]

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

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

									lowercase__            =      np.array(
									    [0.6199778, 0.63984406, 0.46145785, 0.62944984, 0.5622215, 0.47306132, 0.47441456, 0.4607606, 0.48719263]							)
									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       lowerCamelCase_ (							self:	str							)			->   List[Any]:





									"""simple docstring"""







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


							def       lowerCamelCase_ (							self:	List[str]							)			->   Union[str, Any]:





									"""simple docstring"""







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

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

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

									lowercase__            =      KandinskyVaaImgaImgPipeline.from_pretrained(
									    '''kandinsky-community/kandinsky-2-2-decoder'''    ,						torch_dtype=torch.floataa							)
									lowercase__            =      pipeline.to(UpperCamelCase_							)

									pipeline.set_progress_bar_config(disable=UpperCamelCase_							)

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

									lowercase__            =      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'''    ,						)

									lowercase__            =      output.images[0]

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

									assert_mean_pixel_difference(UpperCamelCase_    ,						UpperCamelCase_							)


