



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


						'''simple docstring'''





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



						@property
						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							return 32



						@property
						def     _snake_case	(	self	:							List[Any]	):






							'''simple docstring'''



							return 32



						@property
						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							return self.time_input_dim



						@property
						def     _snake_case	(	self	:							List[Any]	):






							'''simple docstring'''



							return self.time_input_dim * 4



						@property
						def     _snake_case	(	self	:							Any	):






							'''simple docstring'''



							return 100



						@property
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							torch.manual_seed(0	)

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

							__lowercase		       =			UNetaDConditionModel(**_A	)
							return model



						@property
						def     _snake_case	(	self	:							int	):






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






							'''simple docstring'''



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



						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							__lowercase		       =			self.dummy_unet
							__lowercase		       =			self.dummy_movq

							__lowercase		       =			{
							    "num_train_timesteps": 1_000,
							    "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,
							}

							__lowercase		       =			DDIMScheduler(**_A	)

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

							return components



						def     _snake_case	(	self	:							Tuple							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Dict=0	):






							'''simple docstring'''



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

							if str(_A	).startswith("mps"	):
								__lowercase		       =			torch.manual_seed(_A	)
							else:
								__lowercase		       =			torch.Generator(device=_A	).manual_seed(_A	)
							__lowercase		       =			{
							    "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     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							__lowercase		       =			"cpu"

							__lowercase		       =			self.get_dummy_components()

							__lowercase		       =			self.pipeline_class(**_A	)
							__lowercase		       =			pipe.to(_A	)

							pipe.set_progress_bar_config(disable=_A	)

							__lowercase		       =			pipe(**self.get_dummy_inputs(_A	)	)
							__lowercase		       =			output.images

							__lowercase		       =			pipe(
							    **self.get_dummy_inputs(_A	)							,      return_dict=_A							,      )[0]

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

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

							__lowercase		       =			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     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



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



						def     _snake_case	(	self	:							Optional[int]	):






							'''simple docstring'''



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

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

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

							__lowercase		       =			"A robot, 4k photo"

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

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

							pipeline.set_progress_bar_config(disable=_A	)

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

							__lowercase       ,    __lowercase		       =			pipe_prior(
							    _A							,      image=_A							,      strength=0.85							,      generator=_A							,      negative_prompt=""							,      ).to_tuple()

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

							__lowercase		       =			output.images[0]

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

							assert_mean_pixel_difference(_A							,      _A	)




import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin

class   _A   ( _lowercase  ,						_lowercase	):


						'''simple docstring'''





						@register_to_config
						def __init__(	self	:							Optional[Any]							,      *,
						    lowerCamelCase	:							int = 4							,      lowerCamelCase	:							int = 768							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Optional[int]							,      ):






							'''simple docstring'''



							super().__init__()

							__lowercase		       =			nn.Parameter(torch.zeros(lowerCamelCase	)	)

							# parameters for additional clip time embeddings
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)

							# parameters for encoder hidden states
							__lowercase		       =			clip_extra_context_tokens
							__lowercase		       =			nn.Linear(
							    lowerCamelCase							,      self.clip_extra_context_tokens * cross_attention_dim	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.LayerNorm(lowerCamelCase	)



						def     _snake_case	(	self	:							Union[str, Any]							,      *, lowerCamelCase	:							Any							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Tuple	):






							'''simple docstring'''



							if do_classifier_free_guidance:
								# Add the classifier free guidance embeddings to the image embeddings
								__lowercase		       =			image_embeddings.shape[0]
								__lowercase		       =			self.learned_classifier_free_guidance_embeddings.unsqueeze(0	)
								__lowercase		       =			classifier_free_guidance_embeddings.expand(
								    lowerCamelCase							,      -1	)
								__lowercase		       =			torch.cat([classifier_free_guidance_embeddings, image_embeddings]							,      dim=0	)

							# The image embeddings batch size and the text embeddings batch size are equal
							assert image_embeddings.shape[0] == prompt_embeds.shape[0]

							__lowercase		       =			prompt_embeds.shape[0]

							# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
							# adding CLIP embeddings to the existing timestep embedding, ...
							__lowercase		       =			self.embedding_proj(lowerCamelCase	)
							__lowercase		       =			self.clip_image_embeddings_project_to_time_embeddings(lowerCamelCase	)
							__lowercase		       =			time_projected_image_embeddings + time_projected_prompt_embeds

							# ... and by projecting CLIP embeddings into four
							# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
							__lowercase		       =			self.clip_extra_context_tokens_proj(lowerCamelCase	)
							__lowercase		       =			clip_extra_context_tokens.reshape(lowerCamelCase							,      -1							,      self.clip_extra_context_tokens	)
							__lowercase		       =			clip_extra_context_tokens.permute(0							,      2							,      1	)

							__lowercase		       =			self.encoder_hidden_states_proj(lowerCamelCase	)
							__lowercase		       =			self.text_encoder_hidden_states_norm(lowerCamelCase	)
							__lowercase		       =			torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]							,      dim=1	)

							return text_encoder_hidden_states, additive_clip_time_embeddings



