'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    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 					__SCREAMING_SNAKE_CASE						(	_UpperCamelCase						,     unittest.TestCase	):







  '''simple docstring'''






  lowerCamelCase_		:Optional[Any]						   =		KandinskyVaaControlnetImgaImgPipeline
  lowerCamelCase_		:Optional[int]						   =		['image_embeds', 'negative_image_embeds', 'image', 'hint']
  lowerCamelCase_		:Any						   =		['image_embeds', 'negative_image_embeds', 'image', 'hint']
  lowerCamelCase_		:Union[str, Any]						   =		[
      'generator',
      'height',
      'width',
      'strength',
      'guidance_scale',
      'num_inference_steps',
      'return_dict',
      'guidance_scale',
      'num_images_per_prompt',
      'output_type',
      'return_dict',
  ]
  lowerCamelCase_		:Optional[int]						   =		False


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return 3_2


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return 3_2


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return self.time_input_dim


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return self.time_input_dim * 4


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return 1_0_0


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       torch.manual_seed(0						)

       UpperCAmelCase_							:       List[str]       =				{
           '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_							:       Dict       =				UNetaDConditionModel(**_SCREAMING_SNAKE_CASE						)
       return model


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       torch.manual_seed(0						)
       UpperCAmelCase_							:       str       =				VQModel(**self.dummy_movq_kwargs						)
       return model


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Any       =				self.dummy_unet
       UpperCAmelCase_							:       Any       =				self.dummy_movq

       UpperCAmelCase_							:       int       =				{
           'num_train_timesteps': 1_0_0_0,
           'beta_schedule': 'linear',
           'beta_start': 0.0_00_85,
           'beta_end': 0.0_12,
           'clip_sample': False,
           'set_alpha_to_one': False,
           'steps_offset': 0,
           'prediction_type': 'epsilon',
           'thresholding': False,
       }

       UpperCAmelCase_							:       List[Any]       =				DDIMScheduler(**_SCREAMING_SNAKE_CASE						)

       UpperCAmelCase_							:       Optional[int]       =				{
           'unet': unet,
           'scheduler': scheduler,
           'movq': movq,
       }

       return components


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_=0						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[int]       =				floats_tensor((1, self.text_embedder_hidden_size)				,		rng=random.Random(_SCREAMING_SNAKE_CASE						)						).to(_SCREAMING_SNAKE_CASE						)
       UpperCAmelCase_							:       Optional[Any]       =				floats_tensor((1, self.text_embedder_hidden_size)				,		rng=random.Random(seed + 1						)						).to(
           _SCREAMING_SNAKE_CASE						)
       # create init_image
       UpperCAmelCase_							:       Union[str, Any]       =				floats_tensor((1, 3, 6_4, 6_4)				,		rng=random.Random(_SCREAMING_SNAKE_CASE						)						).to(_SCREAMING_SNAKE_CASE						)
       UpperCAmelCase_							:       Tuple       =				image.cpu().permute(0				,		2				,		3				,		1						)[0]
       UpperCAmelCase_							:       Optional[Any]       =				Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE						)						).convert('RGB'						).resize((2_5_6, 2_5_6)						)
       # create hint
       UpperCAmelCase_							:       List[str]       =				floats_tensor((1, 3, 6_4, 6_4)				,		rng=random.Random(_SCREAMING_SNAKE_CASE						)						).to(_SCREAMING_SNAKE_CASE						)

       if str(_SCREAMING_SNAKE_CASE						).startswith('mps'						):
            UpperCAmelCase_							:       Union[str, Any]       =				torch.manual_seed(_SCREAMING_SNAKE_CASE						)
       else:
            UpperCAmelCase_							:       Tuple       =				torch.Generator(device=_SCREAMING_SNAKE_CASE						).manual_seed(_SCREAMING_SNAKE_CASE						)
       UpperCAmelCase_							:       Dict       =				{
           'image': init_image,
           'image_embeds': image_embeds,
           'negative_image_embeds': negative_image_embeds,
           'hint': hint,
           'generator': generator,
           'height': 6_4,
           'width': 6_4,
           'num_inference_steps': 1_0,
           'guidance_scale': 7.0,
           'strength': 0.2,
           'output_type': 'np',
       }
       return inputs






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       int       =				'cpu'

       UpperCAmelCase_							:       str       =				self.get_dummy_components()

       UpperCAmelCase_							:       Union[str, Any]       =				self.pipeline_class(**_SCREAMING_SNAKE_CASE						)
       UpperCAmelCase_							:       Union[str, Any]       =				pipe.to(_SCREAMING_SNAKE_CASE						)

       pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       UpperCAmelCase_							:       Any       =				pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)						)
       UpperCAmelCase_							:       List[str]       =				output.images

       UpperCAmelCase_							:       Dict       =				pipe(
           **self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)				,		return_dict=_SCREAMING_SNAKE_CASE				,		)[0]

       UpperCAmelCase_							:       Dict       =				image[0, -3:, -3:, -1]
       UpperCAmelCase_							:       Optional[Any]       =				image_from_tuple[0, -3:, -3:, -1]

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

       UpperCAmelCase_							:       Optional[Any]       =				np.array(
           [0.54_98_50_34, 0.55_50_93_65, 0.52_56_15_04, 0.5_57_04_94, 0.5_59_38_18, 0.5_26_39_79, 0.50_28_56_43, 0.5_06_98_46, 0.51_19_67_36]						)
       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 					__SCREAMING_SNAKE_CASE						(	unittest.TestCase	):







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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

       UpperCAmelCase_							:       Optional[int]       =				load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'						)
       UpperCAmelCase_							:       Optional[Any]       =				init_image.resize((5_1_2, 5_1_2)						)

       UpperCAmelCase_							:       List[str]       =				load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/kandinskyv22/hint_image_cat.png'						)
       UpperCAmelCase_							:       Optional[Any]       =				torch.from_numpy(np.array(_SCREAMING_SNAKE_CASE						)						).float() / 2_55.0
       UpperCAmelCase_							:       int       =				hint.permute(2				,		0				,		1						).unsqueeze(0						)

       UpperCAmelCase_							:       Dict       =				'A robot, 4k photo'

       UpperCAmelCase_							:       Union[str, Any]       =				KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
           'kandinsky-community/kandinsky-2-2-prior'				,		torch_dtype=torch.floataa						)
       pipe_prior.to(_SCREAMING_SNAKE_CASE						)

       UpperCAmelCase_							:       List[Any]       =				KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
           'kandinsky-community/kandinsky-2-2-controlnet-depth'				,		torch_dtype=torch.floataa						)
       UpperCAmelCase_							:       int       =				pipeline.to(_SCREAMING_SNAKE_CASE						)

       pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       UpperCAmelCase_							:       List[str]       =				torch.Generator(device='cpu'						).manual_seed(0						)

       UpperCAmelCase_       ,      UpperCAmelCase_							:       Any       =				pipe_prior(
           _SCREAMING_SNAKE_CASE				,		image=_SCREAMING_SNAKE_CASE				,		strength=0.85				,		generator=_SCREAMING_SNAKE_CASE				,		negative_prompt=''				,		).to_tuple()

       UpperCAmelCase_							:       Dict       =				pipeline(
           image=_SCREAMING_SNAKE_CASE				,		image_embeds=_SCREAMING_SNAKE_CASE				,		negative_image_embeds=_SCREAMING_SNAKE_CASE				,		hint=_SCREAMING_SNAKE_CASE				,		generator=_SCREAMING_SNAKE_CASE				,		num_inference_steps=1_0_0				,		height=5_1_2				,		width=5_1_2				,		strength=0.5				,		output_type='np'				,		)

       UpperCAmelCase_							:       str       =				output.images[0]

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

       assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE						)







'''simple docstring'''


from typing import List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case__					:      List[Any]				      =  logging.get_logger(__name__)

snake_case__					:      Optional[Any]				      =  {
    '''huggingface/autoformer-tourism-monthly''': '''https://huggingface.co/huggingface/autoformer-tourism-monthly/resolve/main/config.json''',
}






class 					__SCREAMING_SNAKE_CASE						(	lowerCamelCase_	):







  '''simple docstring'''






  lowerCamelCase_		:List[str]						   =		'''autoformer'''
  lowerCamelCase_		:Optional[Any]						   =		{
      '''hidden_size''': '''d_model''',
      '''num_attention_heads''': '''encoder_attention_heads''',
      '''num_hidden_layers''': '''encoder_layers''',
  }


  def __init__(			self				,		snake_case_ = None				,		snake_case_ = None				,		snake_case_ = "student_t"				,		snake_case_ = "nll"				,		snake_case_ = 1				,		snake_case_ = [1, 2, 3, 4, 5, 6, 7]				,		snake_case_ = True				,		snake_case_ = 0				,		snake_case_ = 0				,		snake_case_ = 0				,		snake_case_ = 0				,		snake_case_ = None				,		snake_case_ = None				,		snake_case_ = 6_4				,		snake_case_ = 2				,		snake_case_ = 2				,		snake_case_ = 2				,		snake_case_ = 2				,		snake_case_ = 3_2				,		snake_case_ = 3_2				,		snake_case_ = "gelu"				,		snake_case_ = 0.1				,		snake_case_ = 0.1				,		snake_case_ = 0.1				,		snake_case_ = 0.1				,		snake_case_ = 0.1				,		snake_case_ = 1_0_0				,		snake_case_ = 0.02				,		snake_case_ = True				,		snake_case_=True				,		snake_case_ = 1_0				,		snake_case_ = 2_5				,		snake_case_ = 3				,		**snake_case_				,		):




       '''simple docstring'''

       UpperCAmelCase_							:       List[Any]       =				prediction_length
       UpperCAmelCase_							:       List[str]       =				context_length if context_length is not None else prediction_length
       UpperCAmelCase_							:       Optional[int]       =				distribution_output
       UpperCAmelCase_							:       Optional[int]       =				loss
       UpperCAmelCase_							:       Union[str, Any]       =				input_size
       UpperCAmelCase_							:       int       =				num_time_features
       UpperCAmelCase_							:       List[str]       =				lags_sequence
       UpperCAmelCase_							:       Any       =				scaling
       UpperCAmelCase_							:       Any       =				num_dynamic_real_features
       UpperCAmelCase_							:       int       =				num_static_real_features
       UpperCAmelCase_							:       Optional[Any]       =				num_static_categorical_features
       if cardinality is not None and num_static_categorical_features > 0:
            if len(snake_case_						) != num_static_categorical_features:
                 raise ValueError(
                     'The cardinality should be a list of the same length as `num_static_categorical_features`'						)
            UpperCAmelCase_							:       List[Any]       =				cardinality
       else:
            UpperCAmelCase_							:       Optional[int]       =				[0]
       if embedding_dimension is not None and num_static_categorical_features > 0:
            if len(snake_case_						) != num_static_categorical_features:
                 raise ValueError(
                     'The embedding dimension should be a list of the same length as `num_static_categorical_features`'						)
            UpperCAmelCase_							:       List[str]       =				embedding_dimension
       else:
            UpperCAmelCase_							:       List[Any]       =				[min(5_0				,		(cat + 1) // 2						) for cat in self.cardinality]
       UpperCAmelCase_							:       List[str]       =				num_parallel_samples

       # Transformer architecture configuration
       UpperCAmelCase_							:       Union[str, Any]       =				input_size * len(self.lags_sequence						) + self._number_of_features
       UpperCAmelCase_							:       str       =				d_model
       UpperCAmelCase_							:       str       =				encoder_attention_heads
       UpperCAmelCase_							:       str       =				decoder_attention_heads
       UpperCAmelCase_							:       str       =				encoder_ffn_dim
       UpperCAmelCase_							:       str       =				decoder_ffn_dim
       UpperCAmelCase_							:       str       =				encoder_layers
       UpperCAmelCase_							:       str       =				decoder_layers

       UpperCAmelCase_							:       str       =				dropout
       UpperCAmelCase_							:       Optional[int]       =				attention_dropout
       UpperCAmelCase_							:       Tuple       =				activation_dropout
       UpperCAmelCase_							:       Any       =				encoder_layerdrop
       UpperCAmelCase_							:       Tuple       =				decoder_layerdrop

       UpperCAmelCase_							:       List[str]       =				activation_function
       UpperCAmelCase_							:       Tuple       =				init_std

       UpperCAmelCase_							:       Union[str, Any]       =				use_cache

       # Autoformer
       UpperCAmelCase_							:       Any       =				label_length
       UpperCAmelCase_							:       Union[str, Any]       =				moving_average
       UpperCAmelCase_							:       Tuple       =				autocorrelation_factor

       super().__init__(is_encoder_decoder=snake_case_				,		**snake_case_						)






  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return (
           sum(self.embedding_dimension						)
           + self.num_dynamic_real_features
           + self.num_time_features
           + self.num_static_real_features
           + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
       )






