


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 lowerCAmelCase_		(				lowercase__      ,		unittest.TestCase ):

	'''simple docstring'''

	__snake_case				=						KandinskyVaaControlnetImgaImgPipeline
	__snake_case				=						["image_embeds", "negative_image_embeds", "image", "hint"]
	__snake_case				=						["image_embeds", "negative_image_embeds", "image", "hint"]
	__snake_case				=						[
	    "generator",
	    "height",
	    "width",
	    "strength",
	    "guidance_scale",
	    "num_inference_steps",
	    "return_dict",
	    "guidance_scale",
	    "num_images_per_prompt",
	    "output_type",
	    "return_dict",
	]
	__snake_case				=						False



	@property
	def 		UpperCamelCase__  (			self      ):
							return 32



	@property
	def 		UpperCamelCase__  (			self      ):
							return 32



	@property
	def 		UpperCamelCase__  (			self      ):
							return self.time_input_dim



	@property
	def 		UpperCamelCase__  (			self      ):
							return self.time_input_dim * 4



	@property
	def 		UpperCamelCase__  (			self      ):
							return 1_00



	@property
	def 		UpperCamelCase__  (			self      ):
							torch.manual_seed(0      )

							snake_case_ =   {
							    '''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,
							}

							snake_case_ =   UNetaDConditionModel(**UpperCAmelCase__      )
							return model



	@property
	def 		UpperCamelCase__  (			self      ):
							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 		UpperCamelCase__  (			self      ):
							torch.manual_seed(0      )
							snake_case_ =   VQModel(**self.dummy_movq_kwargs      )
							return model



	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   self.dummy_unet
							snake_case_ =   self.dummy_movq

							snake_case_ =   {
							    '''num_train_timesteps''': 10_00,
							    '''beta_schedule''': '''linear''',
							    '''beta_start''': 0.00_085,
							    '''beta_end''': 0.012,
							    '''clip_sample''': False,
							    '''set_alpha_to_one''': False,
							    '''steps_offset''': 0,
							    '''prediction_type''': '''epsilon''',
							    '''thresholding''': False,
							}

							snake_case_ =   DDIMScheduler(**UpperCAmelCase__      )

							snake_case_ =   {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''movq''': movq,
							}

							return components



	def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase=0      ):
							snake_case_ =   floats_tensor((1, self.text_embedder_hidden_size)	,   rng=random.Random(UpperCAmelCase__      )      ).to(UpperCAmelCase__      )
							snake_case_ =   floats_tensor((1, self.text_embedder_hidden_size)	,   rng=random.Random(seed + 1      )      ).to(
							    UpperCAmelCase__      )
							# create init_image
							snake_case_ =   floats_tensor((1, 3, 64, 64)	,   rng=random.Random(UpperCAmelCase__      )      ).to(UpperCAmelCase__      )
							snake_case_ =   image.cpu().permute(0	,   2	,   3	,   1      )[0]
							snake_case_ =   Image.fromarray(np.uinta(UpperCAmelCase__      )      ).convert('''RGB'''      ).resize((2_56, 2_56)      )
							# create hint
							snake_case_ =   floats_tensor((1, 3, 64, 64)	,   rng=random.Random(UpperCAmelCase__      )      ).to(UpperCAmelCase__      )

							if str(UpperCAmelCase__      ).startswith('''mps'''      ):
													snake_case_ =   torch.manual_seed(UpperCAmelCase__      )
							else:
													snake_case_ =   torch.Generator(device=UpperCAmelCase__      ).manual_seed(UpperCAmelCase__      )
							snake_case_ =   {
							    '''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 		UpperCamelCase__  (			self      ):
							snake_case_ =   '''cpu'''

							snake_case_ =   self.get_dummy_components()

							snake_case_ =   self.pipeline_class(**UpperCAmelCase__      )
							snake_case_ =   pipe.to(UpperCAmelCase__      )

							pipe.set_progress_bar_config(disable=UpperCAmelCase__      )

							snake_case_ =   pipe(**self.get_dummy_inputs(UpperCAmelCase__      )      )
							snake_case_ =   output.images

							snake_case_ =   pipe(
							    **self.get_dummy_inputs(UpperCAmelCase__      )	,   return_dict=UpperCAmelCase__	,   )[0]

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

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

							snake_case_ =   np.array(
							    [0.54_985_034, 0.55_509_365, 0.52_561_504, 0.5_570_494, 0.5_593_818, 0.5_263_979, 0.50_285_643, 0.5_069_846, 0.51_196_736]      )
							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 lowerCAmelCase_		(				unittest.TestCase ):

	'''simple docstring'''

	def 		UpperCamelCase__  (			self      ):
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   load_numpy(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							    '''/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy'''      )

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

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

							snake_case_ =   '''A robot, 4k photo'''

							snake_case_ =   KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
							    '''kandinsky-community/kandinsky-2-2-prior'''	,   torch_dtype=torch.floataa      )
							pipe_prior.to(UpperCAmelCase__      )

							snake_case_ =   KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
							    '''kandinsky-community/kandinsky-2-2-controlnet-depth'''	,   torch_dtype=torch.floataa      )
							snake_case_ =   pipeline.to(UpperCAmelCase__      )

							pipeline.set_progress_bar_config(disable=UpperCAmelCase__      )

							snake_case_ =   torch.Generator(device='''cpu'''      ).manual_seed(0      )

							snake_case_ =   pipe_prior(
							    UpperCAmelCase__	,   image=UpperCAmelCase__	,   strength=0.85	,   generator=UpperCAmelCase__	,   negative_prompt=''''''	,   ).to_tuple()

							snake_case_ =   pipeline(
							    image=UpperCAmelCase__	,   image_embeds=UpperCAmelCase__	,   negative_image_embeds=UpperCAmelCase__	,   hint=UpperCAmelCase__	,   generator=UpperCAmelCase__	,   num_inference_steps=1_00	,   height=5_12	,   width=5_12	,   strength=0.5	,   output_type='''np'''	,   )

							snake_case_ =   output.images[0]

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

							assert_mean_pixel_difference(UpperCAmelCase__	,   UpperCAmelCase__      )



# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
   import jax
   import jaxlib

UpperCAmelCase    =					get_logger()

UpperCAmelCase    =					None







class lowerCAmelCase_		(				TensorFormatter[Mapping, "jax.Array", Mapping] ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   **_UpperCAmelCase      ):
       super().__init__(features=_UpperCAmelCase      )
       import jax
       from jaxlib.xla_client import Device

       if isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ):
             raise ValueError(
                 F'''Expected {device} to be a `str` not {type(_UpperCAmelCase      )}, as `jaxlib.xla_extension.Device` '''
                 '''is not serializable neither with `pickle` nor with `dill`. Instead you can surround '''
                 '''the device with `str()` to get its string identifier that will be internally mapped '''
                 '''to the actual `jaxlib.xla_extension.Device`.'''      )
       snake_case_ =   device if isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ) else str(jax.devices()[0]      )
       # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
       # with `pickle` nor with `dill`, so we need to use a global variable instead
       global DEVICE_MAPPING
       if DEVICE_MAPPING is None:
             snake_case_ =   self._map_devices_to_str()
       if self.device not in list(DEVICE_MAPPING.keys()      ):
             logger.warning(
                 F'''Device with string identifier {self.device} not listed among the available '''
                 F'''devices: {list(DEVICE_MAPPING.keys()      )}, so falling back to the default '''
                 F'''device: {str(jax.devices()[0]      )}.'''      )
             snake_case_ =   str(jax.devices()[0]      )
       snake_case_ =   jnp_array_kwargs



 @staticmethod
 def 		UpperCamelCase__  (			):
       import jax

       return {str(_UpperCAmelCase      ): device for device in jax.devices()}



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       import jax
       import jax.numpy as jnp

       if isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ) and column:
             if all(
                 isinstance(_UpperCAmelCase	,   jax.Array      ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column      ):
                   return jnp.stack(_UpperCAmelCase	,   axis=0      )
       return column



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       import jax
       import jax.numpy as jnp

       if isinstance(_UpperCAmelCase	,   (str, bytes, type(_UpperCAmelCase      ))      ):
             return value
       elif isinstance(_UpperCAmelCase	,   (np.character, np.ndarray)      ) and np.issubdtype(value.dtype	,   np.character      ):
             return value.tolist()

       snake_case_ =   {}

       if isinstance(_UpperCAmelCase	,   (np.number, np.ndarray)      ) and np.issubdtype(value.dtype	,   np.integer      ):
             # the default int precision depends on the jax config
             # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
             if jax.config.jax_enable_xaa:
                   snake_case_ =   {'''dtype''': jnp.intaa}
             else:
                   snake_case_ =   {'''dtype''': jnp.intaa}
       elif isinstance(_UpperCAmelCase	,   (np.number, np.ndarray)      ) and np.issubdtype(value.dtype	,   np.floating      ):
             snake_case_ =   {'''dtype''': jnp.floataa}
       elif config.PIL_AVAILABLE and "PIL" in sys.modules:
             import PIL.Image

             if isinstance(_UpperCAmelCase	,   PIL.Image.Image      ):
                   snake_case_ =   np.asarray(_UpperCAmelCase      )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
       global DEVICE_MAPPING
       if DEVICE_MAPPING is None:
             snake_case_ =   self._map_devices_to_str()

       with jax.default_device(DEVICE_MAPPING[self.device]      ):
             # calling jnp.array on a np.ndarray does copy the data
             # see https://github.com/google/jax/issues/4486
             return jnp.array(_UpperCAmelCase	,   **{**default_dtype, **self.jnp_array_kwargs}      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       import jax

       # support for torch, tf, jax etc.
       if config.TORCH_AVAILABLE and "torch" in sys.modules:
             import torch

             if isinstance(_UpperCAmelCase	,   torch.Tensor      ):
                   return self._tensorize(data_struct.detach().cpu().numpy()[()]      )
       if hasattr(_UpperCAmelCase	,   '''__array__'''      ) and not isinstance(_UpperCAmelCase	,   jax.Array      ):
             snake_case_ =   data_struct.__array__()
       # support for nested types like struct of list of struct
       if isinstance(_UpperCAmelCase	,   np.ndarray      ):
             if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
                   return self._consolidate([self.recursive_tensorize(_UpperCAmelCase      ) for substruct in data_struct]      )
       elif isinstance(_UpperCAmelCase	,   (list, tuple)      ):
             return self._consolidate([self.recursive_tensorize(_UpperCAmelCase      ) for substruct in data_struct]      )
       return self._tensorize(_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       return map_nested(self._recursive_tensorize	,   _UpperCAmelCase	,   map_list=_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.numpy_arrow_extractor().extract_row(_UpperCAmelCase      )
       snake_case_ =   self.python_features_decoder.decode_row(_UpperCAmelCase      )
       return self.recursive_tensorize(_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.numpy_arrow_extractor().extract_column(_UpperCAmelCase      )
       snake_case_ =   self.python_features_decoder.decode_column(_UpperCAmelCase	,   pa_table.column_names[0]      )
       snake_case_ =   self.recursive_tensorize(_UpperCAmelCase      )
       snake_case_ =   self._consolidate(_UpperCAmelCase      )
       return column





 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.numpy_arrow_extractor().extract_batch(_UpperCAmelCase      )
       snake_case_ =   self.python_features_decoder.decode_batch(_UpperCAmelCase      )
       snake_case_ =   self.recursive_tensorize(_UpperCAmelCase      )
       for column_name in batch:
             snake_case_ =   self._consolidate(batch[column_name]      )
       return batch