





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

			'''simple docstring'''







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





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return 3_2





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return 3_2





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return self.time_input_dim





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return self.time_input_dim * 4





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return 1_0_0





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				torch.manual_seed(0							)

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

				lowerCamelCase__								=		UNetaDConditionModel(**_a							)
				return model





			@property
			def 		__lowerCamelCase				(    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 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.dummy_unet
				lowerCamelCase__								=		self.dummy_movq

				lowerCamelCase__								=		{
				    '''num_train_timesteps''': 1_0_0_0,
				    '''beta_schedule''': '''linear''',
				    '''beta_start''': 0.0_0085,
				    '''beta_end''': 0.012,
				    '''clip_sample''': False,
				    '''set_alpha_to_one''': False,
				    '''steps_offset''': 0,
				    '''prediction_type''': '''epsilon''',
				    '''thresholding''': False,
				}

				lowerCamelCase__								=		DDIMScheduler(**_a							)

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

				return components





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase=0							):



				'''simple docstring'''







				lowerCamelCase__								=		floats_tensor((1, self.text_embedder_hidden_size)       ,  rng=random.Random(_a							)							).to(_a							)
				lowerCamelCase__								=		floats_tensor((1, self.text_embedder_hidden_size)       ,  rng=random.Random(seed + 1							)							).to(
				    _a							)
				# create init_image
				lowerCamelCase__								=		floats_tensor((1, 3, 6_4, 6_4)       ,  rng=random.Random(_a							)							).to(_a							)
				lowerCamelCase__								=		image.cpu().permute(0       ,  2       ,  3       ,  1							)[0]
				lowerCamelCase__								=		Image.fromarray(np.uinta(_a							)							).convert('''RGB'''							).resize((2_5_6, 2_5_6)							)
				# create hint
				lowerCamelCase__								=		floats_tensor((1, 3, 6_4, 6_4)       ,  rng=random.Random(_a							)							).to(_a							)

				if str(_a							).startswith('''mps'''							):
					lowerCamelCase__								=		torch.manual_seed(_a							)
				else:
					lowerCamelCase__								=		torch.Generator(device=_a							).manual_seed(_a							)
				lowerCamelCase__								=		{
				    '''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 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		'''cpu'''

				lowerCamelCase__								=		self.get_dummy_components()

				lowerCamelCase__								=		self.pipeline_class(**_a							)
				lowerCamelCase__								=		pipe.to(_a							)

				pipe.set_progress_bar_config(disable=_a							)

				lowerCamelCase__								=		pipe(**self.get_dummy_inputs(_a							)							)
				lowerCamelCase__								=		output.images

				lowerCamelCase__								=		pipe(
				    **self.get_dummy_inputs(_a							)       ,  return_dict=_a       ,  )[0]

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

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

				lowerCamelCase__								=		np.array(
				    [0.5498_5034, 0.5550_9365, 0.5256_1504, 0.557_0494, 0.559_3818, 0.526_3979, 0.5028_5643, 0.506_9846, 0.5119_6736]							)
				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 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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

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

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

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

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

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

				pipeline.set_progress_bar_config(disable=_a							)

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

				lowerCamelCase__				,	lowerCamelCase__								=		pipe_prior(
				    _a       ,  image=_a       ,  strength=0.85       ,  generator=_a       ,  negative_prompt=''''''       ,  ).to_tuple()

				lowerCamelCase__								=		pipeline(
				    image=_a       ,  image_embeds=_a       ,  negative_image_embeds=_a       ,  hint=_a       ,  generator=_a       ,  num_inference_steps=1_0_0       ,  height=5_1_2       ,  width=5_1_2       ,  strength=0.5       ,  output_type='''np'''       ,  )

				lowerCamelCase__								=		output.images[0]

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

				assert_mean_pixel_difference(_a       ,  _a							)









"""simple docstring"""




from typing import Dict, List, Optional, Tuple, 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_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
					import torch


lowerCAmelCase							:			int     							=							logging.get_logger(__name__)






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''
				__UpperCamelCase			  =						["pixel_values"]





				def __init__(       self		, _a = True		, _a = None		, _a = PILImageResampling.BILINEAR		, _a = True		, _a = None		, _a = True		, _a = 1 / 255		, _a = True		, _a = None		, _a = None		, **_a		, ):





							"""simple docstring"""
							super().__init__(**_a )
							lowerCamelCase			      =					size if size is not None else {"""shortest_edge""": 256}
							lowerCamelCase			      =					get_size_dict(_a		, default_to_square=_a )
							lowerCamelCase			      =					crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
							lowerCamelCase			      =					get_size_dict(_a		, param_name="""crop_size""" )
							lowerCamelCase			      =					do_resize
							lowerCamelCase			      =					size
							lowerCamelCase			      =					resample
							lowerCamelCase			      =					do_center_crop
							lowerCamelCase			      =					crop_size
							lowerCamelCase			      =					do_rescale
							lowerCamelCase			      =					rescale_factor
							lowerCamelCase			      =					do_normalize
							lowerCamelCase			      =					image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							lowerCamelCase			      =					image_std if image_std is not None else IMAGENET_STANDARD_STD





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a = PILImageResampling.BICUBIC		, _a = None		, **_a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					get_size_dict(_a		, default_to_square=_a )
							if "shortest_edge" not in size:
										raise ValueError(f'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}' )
							lowerCamelCase			      =					get_resize_output_image_size(_a		, size=size["""shortest_edge"""]		, default_to_square=_a )
							return resize(_a		, size=_a		, resample=_a		, data_format=_a		, **_a )





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a = None		, **_a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					get_size_dict(_a )
							if "height" not in size or "width" not in size:
										raise ValueError(f'The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}' )
							return center_crop(_a		, size=(size["""height"""], size["""width"""])		, data_format=_a		, **_a )





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a = None		, **_a ):





							"""simple docstring"""
							return rescale(_a		, scale=_a		, data_format=_a		, **_a )





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a		, _a = None		, **_a		, ):





							"""simple docstring"""
							return normalize(_a		, mean=_a		, std=_a		, data_format=_a		, **_a )





				def    _lowerCAmelCase				(       self		, _a		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = None		, _a = ChannelDimension.FIRST		, **_a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					do_resize if do_resize is not None else self.do_resize
							lowerCamelCase			      =					size if size is not None else self.size
							lowerCamelCase			      =					get_size_dict(_a		, default_to_square=_a )
							lowerCamelCase			      =					resample if resample is not None else self.resample
							lowerCamelCase			      =					do_center_crop if do_center_crop is not None else self.do_center_crop
							lowerCamelCase			      =					crop_size if crop_size is not None else self.crop_size
							lowerCamelCase			      =					get_size_dict(_a		, param_name="""crop_size""" )
							lowerCamelCase			      =					do_rescale if do_rescale is not None else self.do_rescale
							lowerCamelCase			      =					rescale_factor if rescale_factor is not None else self.rescale_factor
							lowerCamelCase			      =					do_normalize if do_normalize is not None else self.do_normalize
							lowerCamelCase			      =					image_mean if image_mean is not None else self.image_mean
							lowerCamelCase			      =					image_std if image_std is not None else self.image_std

							lowerCamelCase			      =					make_list_of_images(_a )

							if not valid_images(_a ):
										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:
										raise ValueError("""Size must be specified if do_resize is True.""" )

							if do_center_crop and crop_size is None:
										raise ValueError("""Crop size 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.
							lowerCamelCase			      =					[to_numpy_array(_a ) for image in images]

							if do_resize:
										lowerCamelCase			      =					[self.resize(image=_a		, size=_a		, resample=_a ) for image in images]

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

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

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

							lowerCamelCase			      =					[to_channel_dimension_format(_a		, _a ) for image in images]

							lowerCamelCase			      =					{"""pixel_values""": images}
							return BatchFeature(data=_a		, tensor_type=_a )





				def    _lowerCAmelCase				(       self		, _a		, _a = None ):





							"""simple docstring"""
							lowerCamelCase			      =					outputs.logits

							# Resize logits and compute semantic segmentation maps
							if target_sizes is not None:
										if len(_a ) != len(_a ):
													raise ValueError(
													    """Make sure that you pass in as many target sizes as the batch dimension of the logits""" )

										if is_torch_tensor(_a ):
													lowerCamelCase			      =					target_sizes.numpy()

										lowerCamelCase			      =					[]

										for idx in range(len(_a ) ):
													lowerCamelCase			      =					torch.nn.functional.interpolate(
													    logits[idx].unsqueeze(dim=0 )		, size=target_sizes[idx]		, mode="""bilinear"""		, align_corners=_a )
													lowerCamelCase			      =					resized_logits[0].argmax(dim=0 )
													semantic_segmentation.append(_a )
							else:
										lowerCamelCase			      =					logits.argmax(dim=1 )
										lowerCamelCase			      =					[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0] )]

							return semantic_segmentation



