
"""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   __a  (      _lowerCAmelCase    ,     unittest.TestCase	):
       UpperCamelCase_		:					List[Any]   =			KandinskyVaaControlnetImgaImgPipeline
       UpperCamelCase_		:					Optional[Any]   =			['''image_embeds''', '''negative_image_embeds''', '''image''', '''hint''']
       UpperCamelCase_		:					Union[str, Any]   =			['''image_embeds''', '''negative_image_embeds''', '''image''', '''hint''']
       UpperCamelCase_		:					List[str]   =			[
           '''generator''',
           '''height''',
           '''width''',
           '''strength''',
           '''guidance_scale''',
           '''num_inference_steps''',
           '''return_dict''',
           '''guidance_scale''',
           '''num_images_per_prompt''',
           '''output_type''',
           '''return_dict''',
       ]
       UpperCamelCase_		:					Optional[Any]   =			False






       @property
       def 	_SCREAMING_SNAKE_CASE       (  self :							Dict			)->  Any:







              """simple docstring"""
              return 32






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







              """simple docstring"""
              return 32






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







              """simple docstring"""
              return self.time_input_dim






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







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






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







              """simple docstring"""
              return 100






       @property
       def 	_SCREAMING_SNAKE_CASE       (  self :							Union[str, Any]			)->  Tuple:







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

              UpperCamelCase					 =		{
                  "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					 =		UNetaDConditionModel(**UpperCAmelCase_			)
              return model






       @property
       def 	_SCREAMING_SNAKE_CASE       (  self :							List[str]			)->  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 	_SCREAMING_SNAKE_CASE       (  self :							Any			)->  Any:







              """simple docstring"""
              torch.manual_seed(0			)
              UpperCamelCase					 =		VQModel(**self.dummy_movq_kwargs			)
              return model






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







              """simple docstring"""
              UpperCamelCase					 =		self.dummy_unet
              UpperCamelCase					 =		self.dummy_movq

              UpperCamelCase					 =		{
                  "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,
              }

              UpperCamelCase					 =		DDIMScheduler(**UpperCAmelCase_			)

              UpperCamelCase					 =		{
                  "unet": unet,
                  "scheduler": scheduler,
                  "movq": movq,
              }

              return components






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[Any]       ,      UpperCAmelCase_ :							Dict       ,      UpperCAmelCase_ :							Any=0			)->  Tuple:







              """simple docstring"""
              UpperCamelCase					 =		floats_tensor((1, self.text_embedder_hidden_size)       ,      rng=random.Random(UpperCAmelCase_			)			).to(UpperCAmelCase_			)
              UpperCamelCase					 =		floats_tensor((1, self.text_embedder_hidden_size)       ,      rng=random.Random(seed + 1			)			).to(
                  UpperCAmelCase_			)
              # create init_image
              UpperCamelCase					 =		floats_tensor((1, 3, 64, 64)       ,      rng=random.Random(UpperCAmelCase_			)			).to(UpperCAmelCase_			)
              UpperCamelCase					 =		image.cpu().permute(0       ,      2       ,      3       ,      1			)[0]
              UpperCamelCase					 =		Image.fromarray(np.uinta(UpperCAmelCase_			)			).convert("RGB"			).resize((256, 256)			)
              # create hint
              UpperCamelCase					 =		floats_tensor((1, 3, 64, 64)       ,      rng=random.Random(UpperCAmelCase_			)			).to(UpperCAmelCase_			)

              if str(UpperCAmelCase_			).startswith("mps"			):
                     UpperCamelCase					 =		torch.manual_seed(UpperCAmelCase_			)
              else:
                     UpperCamelCase					 =		torch.Generator(device=UpperCAmelCase_			).manual_seed(UpperCAmelCase_			)
              UpperCamelCase					 =		{
                  "image": init_image,
                  "image_embeds": image_embeds,
                  "negative_image_embeds": negative_image_embeds,
                  "hint": hint,
                  "generator": generator,
                  "height": 64,
                  "width": 64,
                  "num_inference_steps": 10,
                  "guidance_scale": 7.0,
                  "strength": 0.2,
                  "output_type": "np",
              }
              return inputs






       def 	_SCREAMING_SNAKE_CASE       (  self :							Dict			)->  str:







              """simple docstring"""
              UpperCamelCase					 =		"cpu"

              UpperCamelCase					 =		self.get_dummy_components()

              UpperCamelCase					 =		self.pipeline_class(**UpperCAmelCase_			)
              UpperCamelCase					 =		pipe.to(UpperCAmelCase_			)

              pipe.set_progress_bar_config(disable=UpperCAmelCase_			)

              UpperCamelCase					 =		pipe(**self.get_dummy_inputs(UpperCAmelCase_			)			)
              UpperCamelCase					 =		output.images

              UpperCamelCase					 =		pipe(
                  **self.get_dummy_inputs(UpperCAmelCase_			)       ,      return_dict=UpperCAmelCase_       ,      )[0]

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

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

              UpperCamelCase					 =		np.array(
                  [0.54985034, 0.55509365, 0.52561504, 0.5570494, 0.5593818, 0.5263979, 0.50285643, 0.5069846, 0.51196736]			)
              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 	_SCREAMING_SNAKE_CASE       (  self :							Union[str, Any]			)->  Union[str, Any]:







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






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







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

              UpperCamelCase					 =		load_image(
                  "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"			)
              UpperCamelCase					 =		init_image.resize((512, 512)			)

              UpperCamelCase					 =		load_image(
                  "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                  "/kandinskyv22/hint_image_cat.png"			)
              UpperCamelCase					 =		torch.from_numpy(np.array(UpperCAmelCase_			)			).float() / 255.0
              UpperCamelCase					 =		hint.permute(2       ,      0       ,      1			).unsqueeze(0			)

              UpperCamelCase					 =		"A robot, 4k photo"

              UpperCamelCase					 =		KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
                  "kandinsky-community/kandinsky-2-2-prior"       ,      torch_dtype=torch.floataa			)
              pipe_prior.to(UpperCAmelCase_			)

              UpperCamelCase					 =		KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
                  "kandinsky-community/kandinsky-2-2-controlnet-depth"       ,      torch_dtype=torch.floataa			)
              UpperCamelCase					 =		pipeline.to(UpperCAmelCase_			)

              pipeline.set_progress_bar_config(disable=UpperCAmelCase_			)

              UpperCamelCase					 =		torch.Generator(device="cpu"			).manual_seed(0			)

              UpperCamelCase ,       UpperCamelCase					 =		pipe_prior(
                  UpperCAmelCase_       ,      image=UpperCAmelCase_       ,      strength=0.85       ,      generator=UpperCAmelCase_       ,      negative_prompt=""       ,      ).to_tuple()

              UpperCamelCase					 =		pipeline(
                  image=UpperCAmelCase_       ,      image_embeds=UpperCAmelCase_       ,      negative_image_embeds=UpperCAmelCase_       ,      hint=UpperCAmelCase_       ,      generator=UpperCAmelCase_       ,      num_inference_steps=100       ,      height=512       ,      width=512       ,      strength=0.5       ,      output_type="np"       ,      )

              UpperCamelCase					 =		output.images[0]

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

              assert_mean_pixel_difference(UpperCAmelCase_       ,      UpperCAmelCase_			)







"""simple docstring"""


import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)



class   __a  (      unittest.TestCase	):






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[Any]			)->  int:







              """simple docstring"""
              UpperCamelCase					 =		["a", "b", "c"]

              # Defaults to last layer if both are None
              UpperCamelCase ,       UpperCamelCase					 =		get_aligned_output_features_output_indices(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)
              self.assertEqual(UpperCAmelCase_       ,      ["c"]			)
              self.assertEqual(UpperCAmelCase_       ,      [2]			)

              # Out indices set to match out features
              UpperCamelCase ,       UpperCamelCase					 =		get_aligned_output_features_output_indices(["a", "c"]       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)
              self.assertEqual(UpperCAmelCase_       ,      ["a", "c"]			)
              self.assertEqual(UpperCAmelCase_       ,      [0, 2]			)

              # Out features set to match out indices
              UpperCamelCase ,       UpperCamelCase					 =		get_aligned_output_features_output_indices(UpperCAmelCase_       ,      [0, 2]       ,      UpperCAmelCase_			)
              self.assertEqual(UpperCAmelCase_       ,      ["a", "c"]			)
              self.assertEqual(UpperCAmelCase_       ,      [0, 2]			)

              # Out features selected from negative indices
              UpperCamelCase ,       UpperCamelCase					 =		get_aligned_output_features_output_indices(UpperCAmelCase_       ,      [-3, -1]       ,      UpperCAmelCase_			)
              self.assertEqual(UpperCAmelCase_       ,      ["a", "c"]			)
              self.assertEqual(UpperCAmelCase_       ,      [-3, -1]			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							int			)->  Union[str, Any]:







              """simple docstring"""
              # Stage names must be set
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(["a", "b"]       ,      (0, 1)       ,      UpperCAmelCase_			)

              # Out features must be a list
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(("a", "b")       ,      (0, 1)       ,      ["a", "b"]			)

              # Out features must be a subset of stage names
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(["a", "b"]       ,      (0, 1)       ,      ["a"]			)

              # Out indices must be a list or tuple
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(UpperCAmelCase_       ,      0       ,      ["a", "b"]			)

              # Out indices must be a subset of stage names
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(UpperCAmelCase_       ,      (0, 1)       ,      ["a"]			)

              # Out features and out indices must be the same length
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(["a", "b"]       ,      (0,)       ,      ["a", "b", "c"]			)

              # Out features should match out indices
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(["a", "b"]       ,      (0, 2)       ,      ["a", "b", "c"]			)

              # Out features and out indices should be in order
              with self.assertRaises(UpperCAmelCase_			):
                     verify_out_features_out_indices(["b", "a"]       ,      (0, 1)       ,      ["a", "b"]			)

              # Check passes with valid inputs
              verify_out_features_out_indices(["a", "b", "d"]       ,      (0, 1, -1)       ,      ["a", "b", "c", "d"]			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							Dict			)->  List[Any]:







              """simple docstring"""
              UpperCamelCase					 =		BackboneMixin()

              UpperCamelCase					 =		["a", "b", "c"]
              UpperCamelCase					 =		["a", "c"]
              UpperCamelCase					 =		[0, 2]

              # Check that the output features and indices are set correctly
              self.assertEqual(backbone.out_features       ,      ["a", "c"]			)
              self.assertEqual(backbone.out_indices       ,      [0, 2]			)

              # Check out features and indices are updated correctly
              UpperCamelCase					 =		["a", "b"]
              self.assertEqual(backbone.out_features       ,      ["a", "b"]			)
              self.assertEqual(backbone.out_indices       ,      [0, 1]			)

              UpperCamelCase					 =		[-3, -1]
              self.assertEqual(backbone.out_features       ,      ["a", "c"]			)
              self.assertEqual(backbone.out_indices       ,      [-3, -1]			)





