





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      (a_       ,					unittest.TestCase   ):
   '''simple docstring'''




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






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







     """simple docstring"""






     return 32






   @property
   def        a_	(      self       :   str       )     ->     str:







     """simple docstring"""






     return 32






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







     """simple docstring"""






     return self.time_input_dim






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







     """simple docstring"""






     return self.time_input_dim * 4






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







     """simple docstring"""






     return 1_00






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







     """simple docstring"""






     torch.manual_seed(0       )

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

     A__  				=       UNetaDConditionModel(**lowercase_       )
     return model






   @property
   def        a_	(      self       :   Dict       )     ->     List[str]:







     """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        a_	(      self       :   Any       )     ->     List[Any]:







     """simple docstring"""






     torch.manual_seed(0       )
     A__  				=       VQModel(**self.dummy_movq_kwargs       )
     return model






   def        a_	(      self       :   List[Any]       )     ->     Optional[Any]:







     """simple docstring"""






     A__  				=       self.dummy_unet
     A__  				=       self.dummy_movq

     A__  				=       {
         """num_train_timesteps""": 10_00,
         """beta_schedule""": """linear""",
         """beta_start""": 0.0_0_0_8_5,
         """beta_end""": 0.0_1_2,
         """clip_sample""": False,
         """set_alpha_to_one""": False,
         """steps_offset""": 0,
         """prediction_type""": """epsilon""",
         """thresholding""": False,
     }

     A__  				=       DDIMScheduler(**lowercase_       )

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

     return components






   def        a_	(      self       :   Dict							,     __lowerCAmelCase       :   Optional[int]							,     __lowerCAmelCase       :   Optional[Any]=0       )     ->     str:







     """simple docstring"""






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

     if str(lowercase_       ).startswith("""mps"""       ):
       A__  				=       torch.manual_seed(lowercase_       )
     else:
       A__  				=       torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )
     A__  				=       {
         """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        a_	(      self       :   Union[str, Any]       )     ->     Optional[int]:







     """simple docstring"""






     A__  				=       """cpu"""

     A__  				=       self.get_dummy_components()

     A__  				=       self.pipeline_class(**lowercase_       )
     A__  				=       pipe.to(lowercase_       )

     pipe.set_progress_bar_config(disable=lowercase_       )

     A__  				=       pipe(**self.get_dummy_inputs(lowercase_       )       )
     A__  				=       output.images

     A__  				=       pipe(
         **self.get_dummy_inputs(lowercase_       )							,     return_dict=lowercase_							,     )[0]

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

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

     A__  				=       np.array(
         [0.5_4_9_8_5_0_3_4, 0.5_5_5_0_9_3_6_5, 0.5_2_5_6_1_5_0_4, 0.5_5_7_0_4_9_4, 0.5_5_9_3_8_1_8, 0.5_2_6_3_9_7_9, 0.5_0_2_8_5_6_4_3, 0.5_0_6_9_8_4_6, 0.5_1_1_9_6_7_3_6]       )
     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        a_	(      self       :   str       )     ->     Optional[int]:







     """simple docstring"""






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







   def        a_	(      self       :   str       )     ->     Any:







     """simple docstring"""






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

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

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

     A__  				=       """A robot, 4k photo"""

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

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

     pipeline.set_progress_bar_config(disable=lowercase_       )

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

     A__				,      A__  				=       pipe_prior(
         lowercase_							,     image=lowercase_							,     strength=0.8_5							,     generator=lowercase_							,     negative_prompt=""""""							,     ).to_tuple()

     A__  				=       pipeline(
         image=lowercase_							,     image_embeds=lowercase_							,     negative_image_embeds=lowercase_							,     hint=lowercase_							,     generator=lowercase_							,     num_inference_steps=1_00							,     height=5_12							,     width=5_12							,     strength=0.5							,     output_type="""np"""							,     )

     A__  				=       output.images[0]

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

     assert_mean_pixel_difference(lowercase_							,     lowercase_       )










from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
					import PIL


A   : int			=				logging.get_logger(__name__)


class   A      (SCREAMING_SNAKE_CASE   ):
			'''simple docstring'''




			__lowerCamelCase     :	Optional[Any]          =   ['''pixel_values''']






			def __init__(      self       :   Dict							,     __lowerCAmelCase       :   bool = True							,     __lowerCAmelCase       :   Dict[str, int] = None							,     __lowerCAmelCase       :   int = 0.9							,     __lowerCAmelCase       :   PILImageResampling = PILImageResampling.BICUBIC							,     __lowerCAmelCase       :   bool = True							,     __lowerCAmelCase       :   Dict[str, int] = None							,     __lowerCAmelCase       :   Union[int, float] = 1 / 2_55							,     __lowerCAmelCase       :   bool = True							,     __lowerCAmelCase       :   bool = True							,     __lowerCAmelCase       :   Optional[Union[float, List[float]]] = None							,     __lowerCAmelCase       :   Optional[Union[float, List[float]]] = None							,     **__lowerCAmelCase       :   Optional[int]							,     )     ->     None:







					"""simple docstring"""






					super().__init__(**__lowerCAmelCase       )
					A__  				=       size if size is not None else {"""shortest_edge""": 2_24}
					A__  				=       get_size_dict(__lowerCAmelCase							,     default_to_square=__lowerCAmelCase       )
					A__  				=       crop_size if crop_size is not None else {"""height""": 2_24, """width""": 2_24}
					A__  				=       get_size_dict(__lowerCAmelCase							,     param_name="""crop_size"""       )

					A__  				=       do_resize
					A__  				=       size
					A__  				=       crop_pct
					A__  				=       resample
					A__  				=       do_center_crop
					A__  				=       crop_size
					A__  				=       do_rescale
					A__  				=       rescale_factor
					A__  				=       do_normalize
					A__  				=       image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
					A__  				=       image_std if image_std is not None else IMAGENET_DEFAULT_STD






			def        a_	(      self       :   int							,     __lowerCAmelCase       :   np.ndarray							,     __lowerCAmelCase       :   Dict[str, int]							,     __lowerCAmelCase       :   Optional[float] = None							,     __lowerCAmelCase       :   PILImageResampling = PILImageResampling.BICUBIC							,     __lowerCAmelCase       :   Optional[Union[str, ChannelDimension]] = None							,     **__lowerCAmelCase       :   int							,     )     ->     np.ndarray:







					"""simple docstring"""






					A__  				=       get_size_dict(__lowerCAmelCase							,     default_to_square=__lowerCAmelCase       )
					if "shortest_edge" not in size and ("height" not in size or "width" not in size):
							raise ValueError(f'size must contain \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'       )
					if crop_pct is not None:
							if "shortest_edge" in size:
									A__  				=       int(size["""shortest_edge"""] / crop_pct       )
							elif "height" in size and "width" in size:
									if size["height"] == size["width"]:
											A__  				=       int(size["""height"""] / crop_pct       )
									else:
											A__  				=       (int(size["""height"""] / crop_pct       ), int(size["""width"""] / crop_pct       ))
							else:
									raise ValueError("""Invalid size for resize: {}""".format(__lowerCAmelCase       )       )

							A__  				=       get_resize_output_image_size(__lowerCAmelCase							,     size=__lowerCAmelCase							,     default_to_square=__lowerCAmelCase       )
					else:
							if "shortest_edge" in size:
									A__  				=       get_resize_output_image_size(__lowerCAmelCase							,     size=size["""shortest_edge"""]							,     default_to_square=__lowerCAmelCase       )
							elif "height" in size and "width" in size:
									A__  				=       (size["""height"""], size["""width"""])
							else:
									raise ValueError("""Invalid size for resize: {}""".format(__lowerCAmelCase       )       )

					return resize(__lowerCAmelCase							,     size=__lowerCAmelCase							,     resample=__lowerCAmelCase							,     data_format=__lowerCAmelCase							,     **__lowerCAmelCase       )






			def        a_	(      self       :   Optional[Any]							,     __lowerCAmelCase       :   np.ndarray							,     __lowerCAmelCase       :   Dict[str, int]							,     __lowerCAmelCase       :   Optional[Union[str, ChannelDimension]] = None							,     **__lowerCAmelCase       :   Optional[Any]							,     )     ->     np.ndarray:







					"""simple docstring"""






					A__  				=       get_size_dict(__lowerCAmelCase       )
					if "height" not in size or "width" not in size:
							raise ValueError(f'size must contain \'height\' and \'width\' as keys. Got {size.keys()}'       )
					return center_crop(__lowerCAmelCase							,     size=(size["""height"""], size["""width"""])							,     data_format=__lowerCAmelCase							,     **__lowerCAmelCase       )






			def        a_	(      self       :   str							,     __lowerCAmelCase       :   np.ndarray							,     __lowerCAmelCase       :   Union[int, float]							,     __lowerCAmelCase       :   Optional[Union[str, ChannelDimension]] = None							,     **__lowerCAmelCase       :   Dict							,     )     ->     List[str]:







					"""simple docstring"""






					return rescale(__lowerCAmelCase							,     scale=__lowerCAmelCase							,     data_format=__lowerCAmelCase							,     **__lowerCAmelCase       )






			def        a_	(      self       :   int							,     __lowerCAmelCase       :   np.ndarray							,     __lowerCAmelCase       :   Union[float, List[float]]							,     __lowerCAmelCase       :   Union[float, List[float]]							,     __lowerCAmelCase       :   Optional[Union[str, ChannelDimension]] = None							,     **__lowerCAmelCase       :   int							,     )     ->     np.ndarray:







					"""simple docstring"""






					return normalize(__lowerCAmelCase							,     mean=__lowerCAmelCase							,     std=__lowerCAmelCase							,     data_format=__lowerCAmelCase							,     **__lowerCAmelCase       )







			def        a_	(      self       :   Optional[Any]							,     __lowerCAmelCase       :   ImageInput							,     __lowerCAmelCase       :   bool = None							,     __lowerCAmelCase       :   Dict[str, int] = None							,     __lowerCAmelCase       :   int = None							,     __lowerCAmelCase       :   PILImageResampling = None							,     __lowerCAmelCase       :   bool = None							,     __lowerCAmelCase       :   Dict[str, int] = None							,     __lowerCAmelCase       :   bool = None							,     __lowerCAmelCase       :   float = None							,     __lowerCAmelCase       :   bool = None							,     __lowerCAmelCase       :   Optional[Union[float, List[float]]] = None							,     __lowerCAmelCase       :   Optional[Union[float, List[float]]] = None							,     __lowerCAmelCase       :   Optional[Union[str, TensorType]] = None							,     __lowerCAmelCase       :   ChannelDimension = ChannelDimension.FIRST							,     **__lowerCAmelCase       :   Tuple							,     )     ->     PIL.Image.Image:







					"""simple docstring"""






					A__  				=       do_resize if do_resize is not None else self.do_resize
					A__  				=       crop_pct if crop_pct is not None else self.crop_pct
					A__  				=       resample if resample is not None else self.resample
					A__  				=       do_center_crop if do_center_crop is not None else self.do_center_crop
					A__  				=       do_rescale if do_rescale is not None else self.do_rescale
					A__  				=       rescale_factor if rescale_factor is not None else self.rescale_factor
					A__  				=       do_normalize if do_normalize is not None else self.do_normalize
					A__  				=       image_mean if image_mean is not None else self.image_mean
					A__  				=       image_std if image_std is not None else self.image_std

					A__  				=       size if size is not None else self.size
					A__  				=       get_size_dict(__lowerCAmelCase							,     default_to_square=__lowerCAmelCase       )
					A__  				=       crop_size if crop_size is not None else self.crop_size
					A__  				=       get_size_dict(__lowerCAmelCase							,     param_name="""crop_size"""       )

					A__  				=       make_list_of_images(__lowerCAmelCase       )

					if not valid_images(__lowerCAmelCase       ):
							raise ValueError(
							    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
							    """torch.Tensor, tf.Tensor or jax.ndarray."""       )

					if do_resize and size is None or resample is None:
							raise ValueError("""Size and resample must be specified if do_resize is True."""       )

					if do_center_crop and crop_pct is None:
							raise ValueError("""Crop_pct must be specified if do_center_crop is True."""       )

					if do_rescale and rescale_factor is None:
							raise ValueError("""Rescale factor must be specified if do_rescale is True."""       )

					if do_normalize and (image_mean is None or image_std is None):
							raise ValueError("""Image mean and std must be specified if do_normalize is True."""       )

					# All transformations expect numpy arrays.
					A__  				=       [to_numpy_array(__lowerCAmelCase       ) for image in images]

					if do_resize:
							A__  				=       [self.resize(image=__lowerCAmelCase							,     size=__lowerCAmelCase							,     crop_pct=__lowerCAmelCase							,     resample=__lowerCAmelCase       ) for image in images]

					if do_center_crop:
							A__  				=       [self.center_crop(image=__lowerCAmelCase							,     size=__lowerCAmelCase       ) for image in images]

					if do_rescale:
							A__  				=       [self.rescale(image=__lowerCAmelCase							,     scale=__lowerCAmelCase       ) for image in images]

					if do_normalize:
							A__  				=       [self.normalize(image=__lowerCAmelCase							,     mean=__lowerCAmelCase							,     std=__lowerCAmelCase       ) for image in images]

					A__  				=       [to_channel_dimension_format(__lowerCAmelCase							,     __lowerCAmelCase       ) for image in images]

					A__  				=       {"""pixel_values""": images}
					return BatchFeature(data=__lowerCAmelCase							,     tensor_type=__lowerCAmelCase       )



