




'''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 			_snake_case		(     _UpperCAmelCase						,						unittest.TestCase							):
			lowerCAmelCase_					:				Dict							=   KandinskyVaaControlnetImgaImgPipeline
			lowerCAmelCase_					:				int							=   ['image_embeds', 'negative_image_embeds', 'image', 'hint']
			lowerCAmelCase_					:				Any							=   ['image_embeds', 'negative_image_embeds', 'image', 'hint']
			lowerCAmelCase_					:				int							=   [
			    'generator',
			    'height',
			    'width',
			    'strength',
			    'guidance_scale',
			    'num_inference_steps',
			    'return_dict',
			    'guidance_scale',
			    'num_images_per_prompt',
			    'output_type',
			    'return_dict',
			]
			lowerCAmelCase_					:				List[Any]							=   False




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

										'''simple docstring'''




										return 32




			@property
			def        lowerCAmelCase__    ( self							)    ->    Dict:

										'''simple docstring'''




										return 32




			@property
			def        lowerCAmelCase__    ( self							)    ->    Tuple:

										'''simple docstring'''




										return self.time_input_dim




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

										'''simple docstring'''




										return self.time_input_dim * 4




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

										'''simple docstring'''




										return 100




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

										'''simple docstring'''




										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(**SCREAMING_SNAKE_CASE_							)
										return model




			@property
			def        lowerCAmelCase__    ( self							)    ->    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        lowerCAmelCase__    ( self							)    ->    List[Any]:

										'''simple docstring'''




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




			def        lowerCAmelCase__    ( self							)    ->    Tuple:

										'''simple docstring'''




										snake_case_      =				self.dummy_unet
										snake_case_      =				self.dummy_movq

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

										snake_case_      =				DDIMScheduler(**SCREAMING_SNAKE_CASE_							)

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

										return components




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__=0							)    ->    int:

										'''simple docstring'''




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

										if str(SCREAMING_SNAKE_CASE_							).startswith("mps"							):
																	snake_case_      =				torch.manual_seed(SCREAMING_SNAKE_CASE_							)
										else:
																	snake_case_      =				torch.Generator(device=SCREAMING_SNAKE_CASE_							).manual_seed(SCREAMING_SNAKE_CASE_							)
										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        lowerCAmelCase__    ( self							)    ->    Union[str, Any]:

										'''simple docstring'''




										snake_case_      =				"""cpu"""

										snake_case_      =				self.get_dummy_components()

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

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_							)

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

										snake_case_      =				pipe(
										    **self.get_dummy_inputs(SCREAMING_SNAKE_CASE_							)    ,							return_dict=SCREAMING_SNAKE_CASE_    ,							)[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.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 			_snake_case		(     unittest.TestCase							):




			def        lowerCAmelCase__    ( self							)    ->    Tuple:

										'''simple docstring'''




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




			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




										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((512, 512)							)

										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(SCREAMING_SNAKE_CASE_							)							).float() / 2_5_5.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(SCREAMING_SNAKE_CASE_							)

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

										pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_							)

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

										snake_case_      =				pipe_prior(
										    SCREAMING_SNAKE_CASE_    ,							image=SCREAMING_SNAKE_CASE_    ,							strength=0.8_5    ,							generator=SCREAMING_SNAKE_CASE_    ,							negative_prompt=""    ,							).to_tuple()

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

										snake_case_      =				output.images[0]

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

										assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_    ,							SCREAMING_SNAKE_CASE_							)










'''simple docstring'''





from ..utils import DummyObject, requires_backends




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Dict							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    str:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Dict							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Optional[Any]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    str:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[int]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Optional[Any]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				List[str]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[int]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				int							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Any:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Optional[Any]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Any:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    int:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Dict							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[int]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				List[str]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				int							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    int:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Any							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Tuple:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    str:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				str							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Tuple							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    int:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




