


"""simple docstring"""







import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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 				UpperCamelCase_			(				UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						int					      =    KandinskyInpaintPipeline
	snake_case__       :						List[Any]					      =    ["prompt", "image_embeds", "negative_image_embeds", "image", "mask_image"]
	snake_case__       :						Union[str, Any]					      =    [
	    "prompt",
	    "negative_prompt",
	    "image_embeds",
	    "negative_image_embeds",
	    "image",
	    "mask_image",
	]
	snake_case__       :						List[Any]					      =    [
	    "generator",
	    "height",
	    "width",
	    "latents",
	    "guidance_scale",
	    "negative_prompt",
	    "num_inference_steps",
	    "return_dict",
	    "guidance_scale",
	    "num_images_per_prompt",
	    "output_type",
	    "return_dict",
	]
	snake_case__       :						Optional[int]					      =    False




	@property
	def 						UpperCAmelCase_   (				self    :       Any     ) ->			str:
						return 3_2




	@property
	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Any:
						return 3_2




	@property
	def 						UpperCAmelCase_   (				self    :       Union[str, Any]     ) ->			List[Any]:
						return self.time_input_dim




	@property
	def 						UpperCAmelCase_   (				self    :       int     ) ->			int:
						return self.time_input_dim * 4




	@property
	def 						UpperCAmelCase_   (				self    :       str     ) ->			Any:
						return 1_0_0




	@property
	def 						UpperCAmelCase_   (				self    :       Optional[Any]     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       XLMRobertaTokenizerFast.from_pretrained("YiYiXu/tiny-random-mclip-base"     )
						return tokenizer




	@property
	def 						UpperCAmelCase_   (				self    :       Any     ) ->			List[Any]:
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       MCLIPConfig(
						    numDims=self.cross_attention_dim    ,    transformerDimensions=self.text_embedder_hidden_size    ,    hidden_size=self.text_embedder_hidden_size    ,    intermediate_size=3_7    ,    num_attention_heads=4    ,    num_hidden_layers=5    ,    vocab_size=1_0_0_5    ,    )

						__SCREAMING_SNAKE_CASE		     =       MultilingualCLIP(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       text_encoder.eval()

						return text_encoder




	@property
	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Optional[Any]:
						torch.manual_seed(0     )

						__SCREAMING_SNAKE_CASE		     =       {
						    "in_channels": 9,
						    # Out channels is double in channels because predicts mean and variance
						    "out_channels": 8,
						    "addition_embed_type": "text_image",
						    "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": "text_image_proj",
						    "cross_attention_dim": self.cross_attention_dim,
						    "attention_head_dim": 4,
						    "resnet_time_scale_shift": "scale_shift",
						    "class_embed_type": None,
						}

						__SCREAMING_SNAKE_CASE		     =       UNetaDConditionModel(**UpperCAmelCase__     )
						return model




	@property
	def 						UpperCAmelCase_   (				self    :       List[Any]     ) ->			Any:
						return {
						    "block_out_channels": [3_2, 6_4],
						    "down_block_types": ["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",
						    ],
						    "vq_embed_dim": 4,
						}




	@property
	def 						UpperCAmelCase_   (				self    :       Union[str, Any]     ) ->			Optional[Any]:
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       VQModel(**self.dummy_movq_kwargs     )
						return model




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Optional[Any]:
						__SCREAMING_SNAKE_CASE		     =       self.dummy_text_encoder
						__SCREAMING_SNAKE_CASE		     =       self.dummy_tokenizer
						__SCREAMING_SNAKE_CASE		     =       self.dummy_unet
						__SCREAMING_SNAKE_CASE		     =       self.dummy_movq

						__SCREAMING_SNAKE_CASE		     =       DDIMScheduler(
						    num_train_timesteps=1_0_0_0    ,    beta_schedule="linear"    ,    beta_start=0.00_085    ,    beta_end=0.012    ,    clip_sample=UpperCAmelCase__    ,    set_alpha_to_one=UpperCAmelCase__    ,    steps_offset=1    ,    prediction_type="epsilon"    ,    thresholding=UpperCAmelCase__    ,    )

						__SCREAMING_SNAKE_CASE		     =       {
						    "text_encoder": text_encoder,
						    "tokenizer": tokenizer,
						    "unet": unet,
						    "scheduler": scheduler,
						    "movq": movq,
						}

						return components




	def 						UpperCAmelCase_   (				self    :       Optional[Any]    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       Union[str, Any]=0     ) ->			Dict:
						__SCREAMING_SNAKE_CASE		     =       floats_tensor((1, self.cross_attention_dim)    ,    rng=random.Random(UpperCAmelCase__     )     ).to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       floats_tensor((1, self.cross_attention_dim)    ,    rng=random.Random(seed + 1     )     ).to(UpperCAmelCase__     )
						# create init_image
						__SCREAMING_SNAKE_CASE		     =       floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(UpperCAmelCase__     )     ).to(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       image.cpu().permute(0    ,    2    ,    3    ,    1     )[0]
						__SCREAMING_SNAKE_CASE		     =       Image.fromarray(np.uinta(UpperCAmelCase__     )     ).convert("RGB"     ).resize((2_5_6, 2_5_6)     )
						# create mask
						__SCREAMING_SNAKE_CASE		     =       np.ones((6_4, 6_4)    ,    dtype=np.floataa     )
						__SCREAMING_SNAKE_CASE		     =       0

						if str(UpperCAmelCase__     ).startswith("mps"     ):
											__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(UpperCAmelCase__     )
						else:
											__SCREAMING_SNAKE_CASE		     =       torch.Generator(device=UpperCAmelCase__     ).manual_seed(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       {
						    "prompt": "horse",
						    "image": init_image,
						    "mask_image": mask,
						    "image_embeds": image_embeds,
						    "negative_image_embeds": negative_image_embeds,
						    "generator": generator,
						    "height": 6_4,
						    "width": 6_4,
						    "num_inference_steps": 2,
						    "guidance_scale": 4.0,
						    "output_type": "np",
						}
						return inputs




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       "cpu"

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_components()

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

						pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

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

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

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

						print(F"""image.shape {image.shape}"""     )

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

						__SCREAMING_SNAKE_CASE		     =       np.array(
						    [0.8_326_919, 0.73_790_467, 0.20_918_581, 0.9_309_612, 0.5_511_791, 0.43_713_328, 0.5_513_321, 0.49_922_934, 0.59_497_786]     )

						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()}"""




	def 						UpperCAmelCase_   (				self    :       int     ) ->			Optional[int]:
						super().test_inference_batch_single_identical(expected_max_diff=3E-3     )




@slow
@require_torch_gpu
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       List[Any]     ) ->			str:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE		     =       load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy"     )

						__SCREAMING_SNAKE_CASE		     =       load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"     )
						__SCREAMING_SNAKE_CASE		     =       np.ones((7_6_8, 7_6_8)    ,    dtype=np.floataa     )
						__SCREAMING_SNAKE_CASE		     =       0

						__SCREAMING_SNAKE_CASE		     =       "a hat"

						__SCREAMING_SNAKE_CASE		     =       KandinskyPriorPipeline.from_pretrained(
						    "kandinsky-community/kandinsky-2-1-prior"    ,    torch_dtype=torch.floataa     )
						pipe_prior.to(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       KandinskyInpaintPipeline.from_pretrained(
						    "kandinsky-community/kandinsky-2-1-inpaint"    ,    torch_dtype=torch.floataa     )
						__SCREAMING_SNAKE_CASE		     =       pipeline.to(UpperCAmelCase__     )
						pipeline.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       torch.Generator(device="cpu"     ).manual_seed(0     )
						__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       pipe_prior(
						    UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=5    ,    negative_prompt=""    ,    ).to_tuple()

						__SCREAMING_SNAKE_CASE		     =       pipeline(
						    UpperCAmelCase__    ,    image=UpperCAmelCase__    ,    mask_image=UpperCAmelCase__    ,    image_embeds=UpperCAmelCase__    ,    negative_image_embeds=UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=1_0_0    ,    height=7_6_8    ,    width=7_6_8    ,    output_type="np"    ,    )

						__SCREAMING_SNAKE_CASE		     =       output.images[0]

						assert image.shape == (7_6_8, 7_6_8, 3)

						assert_mean_pixel_difference(UpperCAmelCase__    ,    UpperCAmelCase__     )




"""simple docstring"""







import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 				UpperCamelCase_			(				UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						Union[str, Any]					      =    DiTPipeline
	snake_case__       :						Union[str, Any]					      =    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
	snake_case__       :						Any					      =    PipelineTesterMixin.required_optional_params - {
	    "latents",
	    "num_images_per_prompt",
	    "callback",
	    "callback_steps",
	}
	snake_case__       :						Optional[Any]					      =    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
	snake_case__       :						Optional[Any]					      =    False




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Dict:
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       TransformeraDModel(
						    sample_size=1_6    ,    num_layers=2    ,    patch_size=4    ,    attention_head_dim=8    ,    num_attention_heads=2    ,    in_channels=4    ,    out_channels=8    ,    attention_bias=UpperCAmelCase__    ,    activation_fn="gelu-approximate"    ,    num_embeds_ada_norm=1_0_0_0    ,    norm_type="ada_norm_zero"    ,    norm_elementwise_affine=UpperCAmelCase__    ,    )
						__SCREAMING_SNAKE_CASE		     =       AutoencoderKL()
						__SCREAMING_SNAKE_CASE		     =       DDIMScheduler()
						__SCREAMING_SNAKE_CASE		     =       {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
						return components




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       str=0     ) ->			Optional[int]:
						if str(UpperCAmelCase__     ).startswith("mps"     ):
											__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(UpperCAmelCase__     )
						else:
											__SCREAMING_SNAKE_CASE		     =       torch.Generator(device=UpperCAmelCase__     ).manual_seed(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       {
						    "class_labels": [1],
						    "generator": generator,
						    "num_inference_steps": 2,
						    "output_type": "numpy",
						}
						return inputs




	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			List[str]:
						__SCREAMING_SNAKE_CASE		     =       "cpu"

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_components()
						__SCREAMING_SNAKE_CASE		     =       self.pipeline_class(**UpperCAmelCase__     )
						pipe.to(UpperCAmelCase__     )
						pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_inputs(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       pipe(**UpperCAmelCase__     ).images
						__SCREAMING_SNAKE_CASE		     =       image[0, -3:, -3:, -1]

						self.assertEqual(image.shape    ,    (1, 1_6, 1_6, 3)     )
						__SCREAMING_SNAKE_CASE		     =       np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]     )
						__SCREAMING_SNAKE_CASE		     =       np.abs(image_slice.flatten() - expected_slice     ).max()
						self.assertLessEqual(UpperCAmelCase__    ,    1E-3     )




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Union[str, Any]:
						self._test_inference_batch_single_identical(relax_max_difference=UpperCAmelCase__    ,    expected_max_diff=1E-3     )




	@unittest.skipIf(
	    torch_device != "cuda" or not is_xformers_available()    ,    reason="XFormers attention is only available with CUDA and `xformers` installed"    ,    )
	def 						UpperCAmelCase_   (				self    :       int     ) ->			Dict:
						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3     )




@require_torch_gpu
@slow
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Any:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def 						UpperCAmelCase_   (				self    :       Union[str, Any]     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )

						__SCREAMING_SNAKE_CASE		     =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"     )
						pipe.to("cuda"     )

						__SCREAMING_SNAKE_CASE		     =       ["vase", "umbrella", "white shark", "white wolf"]
						__SCREAMING_SNAKE_CASE		     =       pipe.get_label_ids(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       pipe(UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=4_0    ,    output_type="np"     ).images

						for word, image in zip(UpperCAmelCase__    ,    UpperCAmelCase__     ):
											__SCREAMING_SNAKE_CASE		     =       load_numpy(
											    F"""https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy"""     )
											assert np.abs((expected_image - image).max()     ) < 1E-2




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"     )
						__SCREAMING_SNAKE_CASE		     =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
						pipe.to("cuda"     )

						__SCREAMING_SNAKE_CASE		     =       ["vase", "umbrella"]
						__SCREAMING_SNAKE_CASE		     =       pipe.get_label_ids(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       pipe(UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=2_5    ,    output_type="np"     ).images

						for word, image in zip(UpperCAmelCase__    ,    UpperCAmelCase__     ):
											__SCREAMING_SNAKE_CASE		     =       load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    F"""/dit/{word}_512.npy"""     )

											assert np.abs((expected_image - image).max()     ) < 1E-1
