




"""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__				(	__magic_name__							,    unittest.TestCase   ):
			lowercase_    =	KandinskyVaaControlnetImgaImgPipeline
			lowercase_    =	["image_embeds", "negative_image_embeds", "image", "hint"]
			lowercase_    =	["image_embeds", "negative_image_embeds", "image", "hint"]
			lowercase_    =	[
			    "generator",
			    "height",
			    "width",
			    "strength",
			    "guidance_scale",
			    "num_inference_steps",
			    "return_dict",
			    "guidance_scale",
			    "num_images_per_prompt",
			    "output_type",
			    "return_dict",
			]
			lowercase_    =	False


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


						"""simple docstring"""




						return 32


			@property
			def       a_     (      self      :    Dict):


						"""simple docstring"""




						return 32


			@property
			def       a_     (      self      :    Union[str, Any]):


						"""simple docstring"""




						return self.time_input_dim


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


						"""simple docstring"""




						return self.time_input_dim * 4


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


						"""simple docstring"""




						return 100


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


						"""simple docstring"""




						torch.manual_seed(0)

						__UpperCAmelCase	: List[Any]		=					{
						    "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	: Optional[int]		=					UNetaDConditionModel(**UpperCamelCase_)
						return model


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


						"""simple docstring"""




						torch.manual_seed(0)
						__UpperCAmelCase	: int		=					VQModel(**self.dummy_movq_kwargs)
						return model


			def       a_     (      self      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					self.dummy_unet
						__UpperCAmelCase	: Union[str, Any]		=					self.dummy_movq

						__UpperCAmelCase	: int		=					{
						    "num_train_timesteps": 1000,
						    "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	: Tuple		=					DDIMScheduler(**UpperCamelCase_)

						__UpperCAmelCase	: Any		=					{
						    "unet": unet,
						    "scheduler": scheduler,
						    "movq": movq,
						}

						return components


			def       a_     (      self      :    str      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Optional[int]=0):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					floats_tensor((1, self.text_embedder_hidden_size)      ,  rng=random.Random(UpperCamelCase_)).to(UpperCamelCase_)
						__UpperCAmelCase	: List[str]		=					floats_tensor((1, self.text_embedder_hidden_size)      ,  rng=random.Random(seed + 1)).to(
						    UpperCamelCase_)
						# create init_image
						__UpperCAmelCase	: str		=					floats_tensor((1, 3, 64, 64)      ,  rng=random.Random(UpperCamelCase_)).to(UpperCamelCase_)
						__UpperCAmelCase	: str		=					image.cpu().permute(0      ,  2      ,  3      ,  1)[0]
						__UpperCAmelCase	: Dict		=					Image.fromarray(np.uinta(UpperCamelCase_)).convert("RGB").resize((256, 256))
						# create hint
						__UpperCAmelCase	: Optional[int]		=					floats_tensor((1, 3, 64, 64)      ,  rng=random.Random(UpperCamelCase_)).to(UpperCamelCase_)

						if str(UpperCamelCase_).startswith("mps"):
									__UpperCAmelCase	: List[str]		=					torch.manual_seed(UpperCamelCase_)
						else:
									__UpperCAmelCase	: Optional[int]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					{
						    "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      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					"cpu"

						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components()

						__UpperCAmelCase	: Optional[int]		=					self.pipeline_class(**UpperCamelCase_)
						__UpperCAmelCase	: int		=					pipe.to(UpperCamelCase_)

						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Optional[Any]		=					pipe(**self.get_dummy_inputs(UpperCamelCase_))
						__UpperCAmelCase	: List[Any]		=					output.images

						__UpperCAmelCase	: Optional[Any]		=					pipe(
						    **self.get_dummy_inputs(UpperCamelCase_)      ,  return_dict=UpperCamelCase_      ,  )[0]

						__UpperCAmelCase	: List[Any]		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[str]		=					image_from_tuple[0, -3:, -3:, -1]

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

						__UpperCAmelCase	: Optional[int]		=					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       a_     (      self      :    Optional[int]):


						"""simple docstring"""




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


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy")

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

						__UpperCAmelCase	: Dict		=					load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/kandinskyv22/hint_image_cat.png")
						__UpperCAmelCase	: Tuple		=					torch.from_numpy(np.array(UpperCamelCase_)).float() / 255.0
						__UpperCAmelCase	: Union[str, Any]		=					hint.permute(2      ,  0      ,  1).unsqueeze(0)

						__UpperCAmelCase	: Optional[int]		=					"A robot, 4k photo"

						__UpperCAmelCase	: Dict		=					KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
						    "kandinsky-community/kandinsky-2-2-prior"      ,  torch_dtype=torch.floataa)
						pipe_prior.to(UpperCamelCase_)

						__UpperCAmelCase	: List[Any]		=					KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
						    "kandinsky-community/kandinsky-2-2-controlnet-depth"      ,  torch_dtype=torch.floataa)
						__UpperCAmelCase	: List[str]		=					pipeline.to(UpperCamelCase_)

						pipeline.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Any		=					torch.Generator(device="cpu").manual_seed(0)

						__UpperCAmelCase	: Optional[int]		=					pipe_prior(
						    UpperCamelCase_      ,  image=UpperCamelCase_      ,  strength=0.85      ,  generator=UpperCamelCase_      ,  negative_prompt=""      ,  ).to_tuple()

						__UpperCAmelCase	: str		=					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	: Optional[int]		=					output.images[0]

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

						assert_mean_pixel_difference(UpperCamelCase_      ,  UpperCamelCase_)












"""simple docstring"""





from ...processing_utils import ProcessorMixin



class      a__				(	__magic_name__   ):
			lowercase_    =	["image_processor", "feature_extractor"]
			lowercase_    =	"TvltImageProcessor"
			lowercase_    =	"TvltFeatureExtractor"


			def __init__(      self      :    Tuple      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Dict):


						"""simple docstring"""




						super().__init__(image_processor=UpperCamelCase_      ,  feature_extractor=UpperCamelCase_)

						__UpperCAmelCase	: Optional[int]		=					image_processor
						__UpperCAmelCase	: Dict		=					feature_extractor


			def __call__(      self      :    Tuple      ,  UpperCamelCase_      :    str=None      ,  UpperCamelCase_      :    Optional[int]=None      ,  UpperCamelCase_      :    Optional[int]=None      ,  UpperCamelCase_      :    str=None      ,  UpperCamelCase_      :    List[str]=False      ,  UpperCamelCase_      :    Optional[int]=False      ,  *UpperCamelCase_      :    List[str]      ,  **UpperCamelCase_      :    str      ,  ):


						"""simple docstring"""




						if images is None and audio is None:
									raise ValueError("You need to specify either an `images` or `audio` input to process.")

						__UpperCAmelCase	: Optional[Any]		=					None
						if images is not None:
									__UpperCAmelCase	: int		=					self.image_processor(UpperCamelCase_      ,  mask_pixel=UpperCamelCase_      ,  *UpperCamelCase_      ,  **UpperCamelCase_)
						if images_mixed is not None:
									__UpperCAmelCase	: int		=					self.image_processor(UpperCamelCase_      ,  is_mixed=UpperCamelCase_      ,  *UpperCamelCase_      ,  **UpperCamelCase_)
						if audio is not None:
									__UpperCAmelCase	: List[Any]		=					self.feature_extractor(
									    UpperCamelCase_      ,  *UpperCamelCase_      ,  sampling_rate=UpperCamelCase_      ,  mask_audio=UpperCamelCase_      ,  **UpperCamelCase_)

						__UpperCAmelCase	: List[str]		=					{}
						if audio is not None:
									output_dict.update(UpperCamelCase_)
						if images is not None:
									output_dict.update(UpperCamelCase_)
						if images_mixed_dict is not None:
									output_dict.update(UpperCamelCase_)
						return output_dict


			@property
			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.image_processor.model_input_names
						__UpperCAmelCase	: List[Any]		=					self.feature_extractor.model_input_names
						return list(dict.fromkeys(image_processor_input_names + feature_extractor_input_names))


