




"""simple docstring"""


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 					lowercase(     __lowerCamelCase       ,			unittest.TestCase ):







 '''simple docstring'''





 lowercase__ 					=					KandinskyVaaPipeline
 lowercase__ 					=					[
     "image_embeds",
     "negative_image_embeds",
 ]
 lowercase__ 					=					["image_embeds", "negative_image_embeds"]
 lowercase__ 					=					[
     "generator",
     "height",
     "width",
     "latents",
     "guidance_scale",
     "num_inference_steps",
     "return_dict",
     "guidance_scale",
     "num_images_per_prompt",
     "output_type",
     "return_dict",
 ]
 lowercase__ 					=					False


 @property
 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



   return 32


 @property
 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



   return 32


 @property
 def  UpperCamelCase_   (				self:		List[Any]		):




   '''simple docstring'''



   return self.time_input_dim


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




   '''simple docstring'''



   return self.time_input_dim * 4


 @property
 def  UpperCamelCase_   (				self:		int		):




   '''simple docstring'''



   return 100


 @property
 def  UpperCamelCase_   (				self:		int		):




   '''simple docstring'''



   torch.manual_seed(0		)

   _snake_case							:		str								=  {
       """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,
   }

   _snake_case							:		Optional[int]								=  UNetaDConditionModel(**a_		)
   return model


 @property
 def  UpperCamelCase_   (				self:		int		):




   '''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  UpperCamelCase_   (				self:		Union[str, Any]		):




   '''simple docstring'''



   torch.manual_seed(0		)
   _snake_case							:		Union[str, Any]								=  VQModel(**self.dummy_movq_kwargs		)
   return model


 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   _snake_case							:		Union[str, Any]								=  self.dummy_unet
   _snake_case							:		int								=  self.dummy_movq

   _snake_case							:		List[Any]								=  DDIMScheduler(
       num_train_timesteps=1_000,      beta_schedule="""linear""",      beta_start=0.00_085,      beta_end=0.012,      clip_sample=a_,      set_alpha_to_one=a_,      steps_offset=1,      prediction_type="""epsilon""",      thresholding=a_,      )

   _snake_case							:		Tuple								=  {
       """unet""": unet,
       """scheduler""": scheduler,
       """movq""": movq,
   }
   return components


 def  UpperCamelCase_   (				self:		List[Any],      a_:		Optional[int],      a_:		Tuple=0		):




   '''simple docstring'''



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




   '''simple docstring'''



   _snake_case							:		str								=  """cpu"""

   _snake_case							:		int								=  self.get_dummy_components()

   _snake_case							:		Tuple								=  self.pipeline_class(**a_		)
   _snake_case							:		Optional[Any]								=  pipe.to(a_		)

   pipe.set_progress_bar_config(disable=a_		)

   _snake_case							:		List[str]								=  pipe(**self.get_dummy_inputs(a_		)		)
   _snake_case							:		Optional[Any]								=  output.images

   _snake_case							:		Optional[int]								=  pipe(
       **self.get_dummy_inputs(a_		),      return_dict=a_,      )[0]

   _snake_case							:		int								=  image[0, -3:, -3:, -1]
   _snake_case							:		Any								=  image_from_tuple[0, -3:, -3:, -1]

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

   _snake_case							:		int								=  np.array(
       [0.6_237_976, 1.0, 0.36_441_332, 1.0, 0.70_639_634, 0.29_877_186, 0.85_652_125, 0.5_216_843, 0.54_454_046]		)

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







 '''simple docstring'''





 def  UpperCamelCase_   (				self:		Optional[Any]		):




   '''simple docstring'''



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


 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



   _snake_case							:		Optional[int]								=  load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
       """/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy"""		)

   _snake_case							:		Dict								=  KandinskyVaaPriorPipeline.from_pretrained(
       """kandinsky-community/kandinsky-2-2-prior""",      torch_dtype=torch.floataa		)
   pipe_prior.to(a_		)

   _snake_case							:		Tuple								=  KandinskyVaaPipeline.from_pretrained(
       """kandinsky-community/kandinsky-2-2-decoder""",      torch_dtype=torch.floataa		)
   _snake_case							:		str								=  pipeline.to(a_		)
   pipeline.set_progress_bar_config(disable=a_		)

   _snake_case							:		Tuple								=  """red cat, 4k photo"""

   _snake_case							:		Tuple								=  torch.Generator(device="""cuda"""		).manual_seed(0		)
   _snake_case							:		Union[str, Any]								=  pipe_prior(
       a_,      generator=a_,      num_inference_steps=5,      negative_prompt="""""",      ).to_tuple()

   _snake_case							:		Dict								=  torch.Generator(device="""cuda"""		).manual_seed(0		)
   _snake_case							:		Optional[int]								=  pipeline(
       image_embeds=a_,      negative_image_embeds=a_,      generator=a_,      num_inference_steps=100,      output_type="""np""",      )

   _snake_case							:		List[Any]								=  output.images[0]

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

   assert_mean_pixel_difference(a_,      a_		)









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available


__UpperCAmelCase       :   Union[str, Any]				       =							{}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase       :   str				       =							["BartphoTokenizer"]

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

else:
      import sys

      __UpperCAmelCase       :   Optional[Any]				       =							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






