





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  A__    (							_A ,  unittest.TestCase	):




   """simple docstring"""



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

   @property
   def 							__lowercase		(					self)			->						Any:







          '''simple docstring'''

          return 32

   @property
   def 							__lowercase		(					self)			->						Tuple:







          '''simple docstring'''

          return 32

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







          '''simple docstring'''

          return self.time_input_dim

   @property
   def 							__lowercase		(					self)			->						Dict:







          '''simple docstring'''

          return self.time_input_dim * 4

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







          '''simple docstring'''

          return 100

   @property
   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          torch.manual_seed(0)

          a__						:	Optional[Any] =				{
              '''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__						:	Optional[int] =				UNetaDConditionModel(**lowercase)
          return model

   @property
   def 							__lowercase		(					self)			->						Optional[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 							__lowercase		(					self)			->						Any:







          '''simple docstring'''

          torch.manual_seed(0)
          a__						:	Optional[Any] =				VQModel(**self.dummy_movq_kwargs)
          return model

   def 							__lowercase		(					self)			->						Dict:







          '''simple docstring'''

          a__						:	Union[str, Any] =				self.dummy_unet
          a__						:	Optional[Any] =				self.dummy_movq

          a__						:	Dict =				DDIMScheduler(
              num_train_timesteps=1000					,  beta_schedule='linear'					,  beta_start=0.0_00_85					,  beta_end=0.0_12					,  clip_sample=lowercase					,  set_alpha_to_one=lowercase					,  steps_offset=1					,  prediction_type='epsilon'					,  thresholding=lowercase					,  )

          a__						:	List[str] =				{
              '''unet''': unet,
              '''scheduler''': scheduler,
              '''movq''': movq,
          }
          return components

   def 							__lowercase		(					self					,  lowercase					,  lowercase=0)			->						List[Any]:







          '''simple docstring'''

          a__						:	Any =				floats_tensor((1, self.text_embedder_hidden_size)					,  rng=random.Random(lowercase)).to(lowercase)
          a__						:	Dict =				floats_tensor((1, self.text_embedder_hidden_size)					,  rng=random.Random(seed + 1)).to(
              lowercase)
          if str(lowercase).startswith('mps'):
                 a__						:	Dict =				torch.manual_seed(lowercase)
          else:
                 a__						:	List[Any] =				torch.Generator(device=lowercase).manual_seed(lowercase)
          a__						:	List[str] =				{
              '''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 							__lowercase		(					self)			->						List[Any]:







          '''simple docstring'''

          a__						:	Union[str, Any] =				'''cpu'''

          a__						:	Optional[int] =				self.get_dummy_components()

          a__						:	Optional[int] =				self.pipeline_class(**lowercase)
          a__						:	List[Any] =				pipe.to(lowercase)

          pipe.set_progress_bar_config(disable=lowercase)

          a__						:	str =				pipe(**self.get_dummy_inputs(lowercase))
          a__						:	Union[str, Any] =				output.images

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

          a__						:	Tuple =				image[0, -3:, -3:, -1]
          a__						:	str =				image_from_tuple[0, -3:, -3:, -1]

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

          a__						:	Union[str, Any] =				np.array(
              [0.6_23_79_76, 1.0, 0.36_44_13_32, 1.0, 0.70_63_96_34, 0.29_87_71_86, 0.85_65_21_25, 0.5_21_68_43, 0.54_45_40_46])

          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	):




   """simple docstring"""



   def 							__lowercase		(					self)			->						Tuple:







          '''simple docstring'''

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

   def 							__lowercase		(					self)			->						List[str]:







          '''simple docstring'''

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

          a__						:	List[str] =				KandinskyVaaPriorPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-2-prior'					,  torch_dtype=torch.floataa)
          pipe_prior.to(lowercase)

          a__						:	int =				KandinskyVaaPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-2-decoder'					,  torch_dtype=torch.floataa)
          a__						:	Any =				pipeline.to(lowercase)
          pipeline.set_progress_bar_config(disable=lowercase)

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

          a__						:	Tuple =				torch.Generator(device='cuda').manual_seed(0)
          a__						:	Dict =				pipe_prior(
              lowercase					,  generator=lowercase					,  num_inference_steps=5					,  negative_prompt=''					,  ).to_tuple()

          a__						:	Any =				torch.Generator(device='cuda').manual_seed(0)
          a__						:	List[Any] =				pipeline(
              image_embeds=lowercase					,  negative_image_embeds=lowercase					,  generator=lowercase					,  num_inference_steps=100					,  output_type='np'					,  )

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

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

          assert_mean_pixel_difference(lowercase					,  lowercase)










"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__A		:       Union[str, Any]  				=		{}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:       Dict  				=		['''NllbTokenizer''']

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:       Dict  				=		['''NllbTokenizerFast''']


if TYPE_CHECKING:
						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_nllb import NllbTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_nllb_fast import NllbTokenizerFast

else:
						import sys

						__A		:       List[Any]  				=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



