

"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    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   __lowercase      (     __lowerCamelCase						,      unittest.TestCase       ):
  snake_case_												=							KandinskyVaaControlnetPipeline
  snake_case_												=							["""image_embeds""", """negative_image_embeds""", """hint"""]
  snake_case_												=							["""image_embeds""", """negative_image_embeds""", """hint"""]
  snake_case_												=							[
      """generator""",
      """height""",
      """width""",
      """latents""",
      """guidance_scale""",
      """num_inference_steps""",
      """return_dict""",
      """guidance_scale""",
      """num_images_per_prompt""",
      """output_type""",
      """return_dict""",
  ]
  snake_case_												=							False

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






   '''simple docstring'''
   return 32

  @property
  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   return 32

  @property
  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   return self.time_input_dim

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






   '''simple docstring'''
   return self.time_input_dim * 4

  @property
  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   return 100

  @property
  def  __lowercase							(	self		:      Any     ):






   '''simple docstring'''
   torch.manual_seed(0     )

   UpperCAmelCase__       :	Tuple        =       {
       """in_channels""": 8,
       # Out channels is double in channels because predicts mean and variance
       """out_channels""": 8,
       """addition_embed_type""": """image_hint""",
       """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,
   }

   UpperCAmelCase__       :	int        =       UNetaDConditionModel(**A     )
   return model

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






   '''simple docstring'''
   return {
       "block_out_channels": [32, 32, 64, 64],
       "down_block_types": [
           "DownEncoderBlock2D",
           "DownEncoderBlock2D",
           "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", "UpDecoderBlock2D", "UpDecoderBlock2D"],
       "vq_embed_dim": 4,
   }

  @property
  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   torch.manual_seed(0     )
   UpperCAmelCase__       :	str        =       VQModel(**self.dummy_movq_kwargs     )
   return model

  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       self.dummy_unet
   UpperCAmelCase__       :	List[Any]        =       self.dummy_movq

   UpperCAmelCase__       :	List[Any]        =       DDIMScheduler(
       num_train_timesteps=1_000       ,beta_schedule="""linear"""       ,beta_start=0.0_0_0_8_5       ,beta_end=0.0_1_2       ,clip_sample=A       ,set_alpha_to_one=A       ,steps_offset=1       ,prediction_type="""epsilon"""       ,thresholding=A       ,)

   UpperCAmelCase__       :	Optional[Any]        =       {
       """unet""": unet,
       """scheduler""": scheduler,
       """movq""": movq,
   }
   return components

  def  __lowercase							(	self		:      str       ,A		:      Optional[Any]       ,A		:      Any=0     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       floats_tensor((1, self.text_embedder_hidden_size)       ,rng=random.Random(A     )     ).to(A     )
   UpperCAmelCase__       :	Union[str, Any]        =       floats_tensor((1, self.text_embedder_hidden_size)       ,rng=random.Random(seed + 1     )     ).to(
       A     )

   # create hint
   UpperCAmelCase__       :	int        =       floats_tensor((1, 3, 64, 64)       ,rng=random.Random(A     )     ).to(A     )

   if str(A     ).startswith("""mps"""     ):
    UpperCAmelCase__       :	Optional[int]        =       torch.manual_seed(A     )
   else:
    UpperCAmelCase__       :	Dict        =       torch.Generator(device=A     ).manual_seed(A     )
   UpperCAmelCase__       :	Dict        =       {
       """image_embeds""": image_embeds,
       """negative_image_embeds""": negative_image_embeds,
       """hint""": hint,
       """generator""": generator,
       """height""": 64,
       """width""": 64,
       """guidance_scale""": 4.0,
       """num_inference_steps""": 2,
       """output_type""": """np""",
   }
   return inputs




  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       """cpu"""

   UpperCAmelCase__       :	List[Any]        =       self.get_dummy_components()

   UpperCAmelCase__       :	Union[str, Any]        =       self.pipeline_class(**A     )
   UpperCAmelCase__       :	Optional[int]        =       pipe.to(A     )

   pipe.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	Optional[int]        =       pipe(**self.get_dummy_inputs(A     )     )
   UpperCAmelCase__       :	Tuple        =       output.images

   UpperCAmelCase__       :	Dict        =       pipe(
       **self.get_dummy_inputs(A     )       ,return_dict=A       ,)[0]

   UpperCAmelCase__       :	Tuple        =       image[0, -3:, -3:, -1]
   UpperCAmelCase__       :	Dict        =       image_from_tuple[0, -3:, -3:, -1]

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

   UpperCAmelCase__       :	Optional[int]        =       np.array(
       [0.6_9_5_9_8_2_6, 0.8_6_8_2_7_9, 0.7_5_5_8_0_9_2, 0.6_8_7_6_9_4_6_7, 0.8_5_8_0_5_8_0_4, 0.6_5_9_7_7_4_9_6, 0.4_4_8_8_5_3_0_2, 0.5_9_5_9_1_1_1, 0.4_2_5_1_5_9_5]     )

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

  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   # clean up the VRAM after each test
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()




  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	int        =       load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
       """/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy"""     )

   UpperCAmelCase__       :	int        =       load_image(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
       """/kandinskyv22/hint_image_cat.png"""     )
   UpperCAmelCase__       :	int        =       torch.from_numpy(np.array(A     )     ).float() / 2_5_5.0
   UpperCAmelCase__       :	Union[str, Any]        =       hint.permute(2       ,0       ,1     ).unsqueeze(0     )

   UpperCAmelCase__       :	List[str]        =       KandinskyVaaPriorPipeline.from_pretrained(
       """kandinsky-community/kandinsky-2-2-prior"""       ,torch_dtype=torch.floataa     )
   pipe_prior.to(A     )

   UpperCAmelCase__       :	List[Any]        =       KandinskyVaaControlnetPipeline.from_pretrained(
       """kandinsky-community/kandinsky-2-2-controlnet-depth"""       ,torch_dtype=torch.floataa     )
   UpperCAmelCase__       :	int        =       pipeline.to(A     )
   pipeline.set_progress_bar_config(disable=A     )

   UpperCAmelCase__       :	Optional[Any]        =       """A robot, 4k photo"""

   UpperCAmelCase__       :	List[Any]        =       torch.Generator(device="""cuda"""     ).manual_seed(0     )
   UpperCAmelCase__ ,   UpperCAmelCase__       :	Tuple        =       pipe_prior(
       A       ,generator=A       ,num_inference_steps=5       ,negative_prompt=""""""       ,).to_tuple()

   UpperCAmelCase__       :	List[str]        =       torch.Generator(device="""cuda"""     ).manual_seed(0     )
   UpperCAmelCase__       :	int        =       pipeline(
       image_embeds=A       ,negative_image_embeds=A       ,hint=A       ,generator=A       ,num_inference_steps=100       ,output_type="""np"""       ,)

   UpperCAmelCase__       :	Any        =       output.images[0]

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

   assert_mean_pixel_difference(A       ,A     )












"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase     :						Optional[Any]						       =  logging.get_logger(__name__)

__UpperCamelCase     :						List[Any]						       =  {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}







class       a							(   a__  ):
 snake_case__             =       '''megatron-bert'''



 def __init__(     self				,       _snake_case=2_90_56				,       _snake_case=10_24				,       _snake_case=24				,       _snake_case=16				,       _snake_case=40_96				,       _snake_case="gelu"				,       _snake_case=0.1				,       _snake_case=0.1				,       _snake_case=5_12				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=1E-12				,       _snake_case=0				,       _snake_case="absolute"				,       _snake_case=True				,       **_snake_case				,       ):







  """simple docstring"""






  super().__init__(pad_token_id=_snake_case				,       **_snake_case     )

  lowerCAmelCase    				= vocab_size
  lowerCAmelCase    				= hidden_size
  lowerCAmelCase    				= num_hidden_layers
  lowerCAmelCase    				= num_attention_heads
  lowerCAmelCase    				= hidden_act
  lowerCAmelCase    				= intermediate_size
  lowerCAmelCase    				= hidden_dropout_prob
  lowerCAmelCase    				= attention_probs_dropout_prob
  lowerCAmelCase    				= max_position_embeddings
  lowerCAmelCase    				= type_vocab_size
  lowerCAmelCase    				= initializer_range
  lowerCAmelCase    				= layer_norm_eps
  lowerCAmelCase    				= position_embedding_type
  lowerCAmelCase    				= use_cache





