



'''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 			UpperCamelCase_ (							__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  _lowercase(  self    )			->			Dict:
								return 32





	@property
	def  _lowercase(  self    )			->			List[str]:
								return 32





	@property
	def  _lowercase(  self    )			->			str:
								return self.time_input_dim





	@property
	def  _lowercase(  self    )			->			List[str]:
								return self.time_input_dim * 4





	@property
	def  _lowercase(  self    )			->			Tuple:
								return 100





	@property
	def  _lowercase(  self    )			->			Union[str, Any]:
								torch.manual_seed(0    )

								UpperCAmelCase     :		Optional[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     :		Any							     =		UNetaDConditionModel(**A    )
								return model





	@property
	def  _lowercase(  self    )			->			Tuple:
								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  _lowercase(  self    )			->			Dict:
								torch.manual_seed(0    )
								UpperCAmelCase     :		str							     =		VQModel(**self.dummy_movq_kwargs    )
								return model





	def  _lowercase(  self    )			->			Tuple:
								UpperCAmelCase     :		Optional[Any]							     =		self.dummy_unet
								UpperCAmelCase     :		int							     =		self.dummy_movq

								UpperCAmelCase     :		List[str]							     =		{
								    """num_train_timesteps""": 1000,
								    """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,
								}

								UpperCAmelCase     :		int							     =		DDIMScheduler(**A    )

								UpperCAmelCase     :		Tuple							     =		{
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """movq""": movq,
								}

								return components





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

								if str(A    ).startswith("""mps"""    ):
															UpperCAmelCase     :		Any							     =		torch.manual_seed(A    )
								else:
															UpperCAmelCase     :		List[str]							     =		torch.Generator(device=A    ).manual_seed(A    )
								UpperCAmelCase     :		str							     =		{
								    """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  _lowercase(  self    )			->			List[str]:
								UpperCAmelCase     :		Dict							     =		"""cpu"""

								UpperCAmelCase     :		Optional[Any]							     =		self.get_dummy_components()

								UpperCAmelCase     :		Any							     =		self.pipeline_class(**A    )
								UpperCAmelCase     :		List[str]							     =		pipe.to(A    )

								pipe.set_progress_bar_config(disable=A    )

								UpperCAmelCase     :		Union[str, Any]							     =		pipe(**self.get_dummy_inputs(A    )    )
								UpperCAmelCase     :		Any							     =		output.images

								UpperCAmelCase     :		Any							     =		pipe(
								    **self.get_dummy_inputs(A    )		,		return_dict=A		,		)[0]

								UpperCAmelCase     :		str							     =		image[0, -3:, -3:, -1]
								UpperCAmelCase     :		int							     =		image_from_tuple[0, -3:, -3:, -1]

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

								UpperCAmelCase     :		Any							     =		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 			UpperCamelCase_ (							unittest.TestCase			):





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






	def  _lowercase(  self    )			->			Tuple:
								UpperCAmelCase     :		Tuple							     =		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     :		int							     =		init_image.resize((512, 512)    )

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

								UpperCAmelCase     :		Dict							     =		"""A robot, 4k photo"""

								UpperCAmelCase     :		int							     =		KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
								    """kandinsky-community/kandinsky-2-2-prior"""		,		torch_dtype=torch.floataa    )
								pipe_prior.to(A    )

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

								pipeline.set_progress_bar_config(disable=A    )

								UpperCAmelCase     :		Optional[int]							     =		torch.Generator(device="""cpu"""    ).manual_seed(0    )

								UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		pipe_prior(
								    A		,		image=A		,		strength=0.8_5		,		generator=A		,		negative_prompt=""""""		,		).to_tuple()

								UpperCAmelCase     :		str							     =		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"""		,		)

								UpperCAmelCase     :		List[Any]							     =		output.images[0]

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

								assert_mean_pixel_difference(A		,		A    )






'''simple docstring'''




import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoImageProcessor, ViTImageProcessor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / """utils"""))

from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


a   :			List[str]					  =    get_tests_dir("""fixtures""")

class 			UpperCamelCase_ (							unittest.TestCase			):





	def  _lowercase(  self    )			->			int:
								# A mock response for an HTTP head request to emulate server down
								UpperCAmelCase     :		Tuple							     =		mock.Mock()
								UpperCAmelCase     :		List[str]							     =		500
								UpperCAmelCase     :		Any							     =		{}
								UpperCAmelCase     :		List[str]							     =		HTTPError
								UpperCAmelCase     :		str							     =		{}

								# Download this model to make sure it's in the cache.
								UpperCAmelCase     :		Optional[int]							     =		ViTImageProcessor.from_pretrained("""hf-internal-testing/tiny-random-vit"""    )
								# Under the mock environment we get a 500 error when trying to reach the model.
								with mock.patch("""requests.Session.request"""		,		return_value=A    ) as mock_head:
															UpperCAmelCase     :		Optional[int]							     =		ViTImageProcessor.from_pretrained("""hf-internal-testing/tiny-random-vit"""    )
															# This check we did call the fake head request
															mock_head.assert_called()





	def  _lowercase(  self    )			->			Any:
								# This test is for deprecated behavior and can be removed in v5
								UpperCAmelCase     :		Tuple							     =		ViTImageProcessor.from_pretrained(
								    """https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json"""    )






	def  _lowercase(  self    )			->			Union[str, Any]:
								with self.assertRaises(A    ):
															# config is in subfolder, the following should not work without specifying the subfolder
															UpperCAmelCase     :		Any							     =		AutoImageProcessor.from_pretrained("""hf-internal-testing/stable-diffusion-all-variants"""    )

								UpperCAmelCase     :		Union[str, Any]							     =		AutoImageProcessor.from_pretrained(
								    """hf-internal-testing/stable-diffusion-all-variants"""		,		subfolder="""feature_extractor"""    )

								self.assertIsNotNone(A    )




@is_staging_test
class 			UpperCamelCase_ (							unittest.TestCase			):





	@classmethod
	def  _lowercase(  cls    )			->			Dict:
								UpperCAmelCase     :		Tuple							     =		TOKEN
								HfFolder.save_token(A    )





	@classmethod
	def  _lowercase(  cls    )			->			List[str]:
								try:
															delete_repo(token=cls._token		,		repo_id="""test-image-processor"""    )
								except HTTPError:
															pass

								try:
															delete_repo(token=cls._token		,		repo_id="""valid_org/test-image-processor-org"""    )
								except HTTPError:
															pass

								try:
															delete_repo(token=cls._token		,		repo_id="""test-dynamic-image-processor"""    )
								except HTTPError:
															pass





	def  _lowercase(  self    )			->			Optional[int]:
								UpperCAmelCase     :		Union[str, Any]							     =		ViTImageProcessor.from_pretrained(A    )
								image_processor.push_to_hub("""test-image-processor"""		,		use_auth_token=self._token    )

								UpperCAmelCase     :		Optional[int]							     =		ViTImageProcessor.from_pretrained(f'''{USER}/test-image-processor'''    )
								for k, v in image_processor.__dict__.items():
															self.assertEqual(A		,		getattr(A		,		A    )    )

								# Reset repo
								delete_repo(token=self._token		,		repo_id="""test-image-processor"""    )

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															image_processor.save_pretrained(
															    A		,		repo_id="""test-image-processor"""		,		push_to_hub=A		,		use_auth_token=self._token    )

								UpperCAmelCase     :		Tuple							     =		ViTImageProcessor.from_pretrained(f'''{USER}/test-image-processor'''    )
								for k, v in image_processor.__dict__.items():
															self.assertEqual(A		,		getattr(A		,		A    )    )





	def  _lowercase(  self    )			->			List[str]:
								UpperCAmelCase     :		List[str]							     =		ViTImageProcessor.from_pretrained(A    )
								image_processor.push_to_hub("""valid_org/test-image-processor"""		,		use_auth_token=self._token    )

								UpperCAmelCase     :		Tuple							     =		ViTImageProcessor.from_pretrained("""valid_org/test-image-processor"""    )
								for k, v in image_processor.__dict__.items():
															self.assertEqual(A		,		getattr(A		,		A    )    )

								# Reset repo
								delete_repo(token=self._token		,		repo_id="""valid_org/test-image-processor"""    )

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															image_processor.save_pretrained(
															    A		,		repo_id="""valid_org/test-image-processor-org"""		,		push_to_hub=A		,		use_auth_token=self._token    )

								UpperCAmelCase     :		int							     =		ViTImageProcessor.from_pretrained("""valid_org/test-image-processor-org"""    )
								for k, v in image_processor.__dict__.items():
															self.assertEqual(A		,		getattr(A		,		A    )    )






	def  _lowercase(  self    )			->			Optional[int]:
								CustomImageProcessor.register_for_auto_class()
								UpperCAmelCase     :		Optional[Any]							     =		CustomImageProcessor.from_pretrained(A    )

								image_processor.push_to_hub("""test-dynamic-image-processor"""		,		use_auth_token=self._token    )

								# This has added the proper auto_map field to the config
								self.assertDictEqual(
								    image_processor.auto_map		,		{"""AutoImageProcessor""": """custom_image_processing.CustomImageProcessor"""}		,		)

								UpperCAmelCase     :		Union[str, Any]							     =		AutoImageProcessor.from_pretrained(
								    f'''{USER}/test-dynamic-image-processor'''		,		trust_remote_code=A    )
								# Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
								self.assertEqual(new_image_processor.__class__.__name__		,		"""CustomImageProcessor"""    )

