

'''simple docstring'''

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()







class   UpperCamelCase__    (   _UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							unittest.TestCase   ):


		"""simple docstring"""

		SCREAMING_SNAKE_CASE__  :	int   =							StableDiffusionControlNetImgaImgPipeline
		SCREAMING_SNAKE_CASE__  :	Tuple   =							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
		SCREAMING_SNAKE_CASE__  :	Tuple   =							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
		SCREAMING_SNAKE_CASE__  :	Union[str, Any]   =							IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}   )
		SCREAMING_SNAKE_CASE__  :	Union[str, Any]   =							IMAGE_TO_IMAGE_IMAGE_PARAMS







		def       A_       ( self			):


							'''simple docstring'''

							torch.manual_seed(0			)
							UpperCAmelCase		:							List[str]   					=   UNetaDConditionModel(
							    block_out_channels=(3_2, 6_4)    ,			layers_per_block=2    ,			sample_size=3_2    ,			in_channels=4    ,			out_channels=4    ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")    ,			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")    ,			cross_attention_dim=3_2    ,			)
							torch.manual_seed(0			)
							UpperCAmelCase		:							Optional[Any]   					=   ControlNetModel(
							    block_out_channels=(3_2, 6_4)    ,			layers_per_block=2    ,			in_channels=4    ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")    ,			cross_attention_dim=3_2    ,			conditioning_embedding_out_channels=(1_6, 3_2)    ,			)
							torch.manual_seed(0			)
							UpperCAmelCase		:							List[str]   					=   DDIMScheduler(
							    beta_start=0.0_0085    ,			beta_end=0.012    ,			beta_schedule="scaled_linear"    ,			clip_sample=_UpperCAmelCase    ,			set_alpha_to_one=_UpperCAmelCase    ,			)
							torch.manual_seed(0			)
							UpperCAmelCase		:							List[str]   					=   AutoencoderKL(
							    block_out_channels=[3_2, 6_4]    ,			in_channels=3    ,			out_channels=3    ,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]    ,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]    ,			latent_channels=4    ,			)
							torch.manual_seed(0			)
							UpperCAmelCase		:							List[Any]   					=   CLIPTextConfig(
							    bos_token_id=0    ,			eos_token_id=2    ,			hidden_size=3_2    ,			intermediate_size=3_7    ,			layer_norm_eps=1e-05    ,			num_attention_heads=4    ,			num_hidden_layers=5    ,			pad_token_id=1    ,			vocab_size=1_0_0_0    ,			)
							UpperCAmelCase		:							str   					=   CLIPTextModel(_UpperCAmelCase			)
							UpperCAmelCase		:							List[str]   					=   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

							UpperCAmelCase		:							Tuple   					=   {
							    'unet': unet,
							    'controlnet': controlnet,
							    'scheduler': scheduler,
							    'vae': vae,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'safety_checker': None,
							    'feature_extractor': None,
							}
							return components







		def       A_       ( self    ,			snake_case    ,			snake_case=0			):


							'''simple docstring'''

							if str(_UpperCAmelCase			).startswith("mps"			):
												UpperCAmelCase		:							Optional[Any]   					=   torch.manual_seed(_UpperCAmelCase			)
							else:
												UpperCAmelCase		:							Dict   					=   torch.Generator(device=_UpperCAmelCase			).manual_seed(_UpperCAmelCase			)

							UpperCAmelCase		:							int   					=   2
							UpperCAmelCase		:							str   					=   randn_tensor(
							    (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)    ,			generator=_UpperCAmelCase    ,			device=torch.device(_UpperCAmelCase			)    ,			)
							UpperCAmelCase		:							int   					=   floats_tensor(control_image.shape    ,			rng=random.Random(_UpperCAmelCase			)			).to(_UpperCAmelCase			)
							UpperCAmelCase		:							int   					=   image.cpu().permute(0    ,			2    ,			3    ,			1			)[0]
							UpperCAmelCase		:							List[Any]   					=   Image.fromarray(np.uinta(_UpperCAmelCase			)			).convert("RGB"			).resize((6_4, 6_4)			)
							UpperCAmelCase		:							List[str]   					=   {
							    'prompt': 'A painting of a squirrel eating a burger',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 6.0,
							    'output_type': 'numpy',
							    'image': image,
							    'control_image': control_image,
							}

							return inputs







		def       A_       ( self			):


							'''simple docstring'''

							return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3			)







		@unittest.skipIf(
		    torch_device != "cuda" or not is_xformers_available()    ,			reason="XFormers attention is only available with CUDA and `xformers` installed"    ,			)
		def       A_       ( self			):


							'''simple docstring'''

							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3			)







		def       A_       ( self			):


							'''simple docstring'''

							self._test_inference_batch_single_identical(expected_max_diff=2e-3			)







class   UpperCamelCase__    (   _UpperCamelCase							,							_UpperCamelCase							,							unittest.TestCase   ):


		"""simple docstring"""

		SCREAMING_SNAKE_CASE__  :	Optional[int]   =							StableDiffusionControlNetImgaImgPipeline
		SCREAMING_SNAKE_CASE__  :	Optional[Any]   =							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
		SCREAMING_SNAKE_CASE__  :	Optional[int]   =							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
		SCREAMING_SNAKE_CASE__  :	List[Any]   =							frozenset([]   )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess







		def       A_       ( self			):


							'''simple docstring'''

							torch.manual_seed(0			)
							UpperCAmelCase		:							Dict   					=   UNetaDConditionModel(
							    block_out_channels=(3_2, 6_4)    ,			layers_per_block=2    ,			sample_size=3_2    ,			in_channels=4    ,			out_channels=4    ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")    ,			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")    ,			cross_attention_dim=3_2    ,			)
							torch.manual_seed(0			)

							def init_weights(snake_case			):
												if isinstance(_UpperCAmelCase    ,			torch.nn.Convad			):
																	torch.nn.init.normal(m.weight			)
																	m.bias.data.fill_(1.0			)

							UpperCAmelCase		:							Union[str, Any]   					=   ControlNetModel(
							    block_out_channels=(3_2, 6_4)    ,			layers_per_block=2    ,			in_channels=4    ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")    ,			cross_attention_dim=3_2    ,			conditioning_embedding_out_channels=(1_6, 3_2)    ,			)
							controlneta.controlnet_down_blocks.apply(_UpperCAmelCase			)

							torch.manual_seed(0			)
							UpperCAmelCase		:							Union[str, Any]   					=   ControlNetModel(
							    block_out_channels=(3_2, 6_4)    ,			layers_per_block=2    ,			in_channels=4    ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")    ,			cross_attention_dim=3_2    ,			conditioning_embedding_out_channels=(1_6, 3_2)    ,			)
							controlneta.controlnet_down_blocks.apply(_UpperCAmelCase			)

							torch.manual_seed(0			)
							UpperCAmelCase		:							Optional[int]   					=   DDIMScheduler(
							    beta_start=0.0_0085    ,			beta_end=0.012    ,			beta_schedule="scaled_linear"    ,			clip_sample=_UpperCAmelCase    ,			set_alpha_to_one=_UpperCAmelCase    ,			)
							torch.manual_seed(0			)
							UpperCAmelCase		:							List[str]   					=   AutoencoderKL(
							    block_out_channels=[3_2, 6_4]    ,			in_channels=3    ,			out_channels=3    ,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]    ,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]    ,			latent_channels=4    ,			)
							torch.manual_seed(0			)
							UpperCAmelCase		:							Dict   					=   CLIPTextConfig(
							    bos_token_id=0    ,			eos_token_id=2    ,			hidden_size=3_2    ,			intermediate_size=3_7    ,			layer_norm_eps=1e-05    ,			num_attention_heads=4    ,			num_hidden_layers=5    ,			pad_token_id=1    ,			vocab_size=1_0_0_0    ,			)
							UpperCAmelCase		:							Optional[int]   					=   CLIPTextModel(_UpperCAmelCase			)
							UpperCAmelCase		:							Dict   					=   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

							UpperCAmelCase		:							str   					=   MultiControlNetModel([controlneta, controlneta]			)

							UpperCAmelCase		:							Tuple   					=   {
							    'unet': unet,
							    'controlnet': controlnet,
							    'scheduler': scheduler,
							    'vae': vae,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'safety_checker': None,
							    'feature_extractor': None,
							}
							return components







		def       A_       ( self    ,			snake_case    ,			snake_case=0			):


							'''simple docstring'''

							if str(_UpperCAmelCase			).startswith("mps"			):
												UpperCAmelCase		:							List[str]   					=   torch.manual_seed(_UpperCAmelCase			)
							else:
												UpperCAmelCase		:							Dict   					=   torch.Generator(device=_UpperCAmelCase			).manual_seed(_UpperCAmelCase			)

							UpperCAmelCase		:							str   					=   2

							UpperCAmelCase		:							Union[str, Any]   					=   [
							    randn_tensor(
							        (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)    ,			generator=_UpperCAmelCase    ,			device=torch.device(_UpperCAmelCase			)    ,			),
							    randn_tensor(
							        (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)    ,			generator=_UpperCAmelCase    ,			device=torch.device(_UpperCAmelCase			)    ,			),
							]

							UpperCAmelCase		:							str   					=   floats_tensor(control_image[0].shape    ,			rng=random.Random(_UpperCAmelCase			)			).to(_UpperCAmelCase			)
							UpperCAmelCase		:							int   					=   image.cpu().permute(0    ,			2    ,			3    ,			1			)[0]
							UpperCAmelCase		:							Optional[Any]   					=   Image.fromarray(np.uinta(_UpperCAmelCase			)			).convert("RGB"			).resize((6_4, 6_4)			)
							UpperCAmelCase		:							Tuple   					=   {
							    'prompt': 'A painting of a squirrel eating a burger',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 6.0,
							    'output_type': 'numpy',
							    'image': image,
							    'control_image': control_image,
							}

							return inputs







		def       A_       ( self			):


							'''simple docstring'''

							UpperCAmelCase		:							Dict   					=   self.get_dummy_components()
							UpperCAmelCase		:							List[str]   					=   self.pipeline_class(**_UpperCAmelCase			)
							pipe.to(_UpperCAmelCase			)

							UpperCAmelCase		:							Any   					=   10.0
							UpperCAmelCase		:							Optional[Any]   					=   4

							UpperCAmelCase		:							Tuple   					=   self.get_dummy_inputs(_UpperCAmelCase			)
							UpperCAmelCase		:							Optional[Any]   					=   steps
							UpperCAmelCase		:							Optional[int]   					=   scale
							UpperCAmelCase		:							Any   					=   pipe(**_UpperCAmelCase			)[0]

							UpperCAmelCase		:							Any   					=   self.get_dummy_inputs(_UpperCAmelCase			)
							UpperCAmelCase		:							Tuple   					=   steps
							UpperCAmelCase		:							List[str]   					=   scale
							UpperCAmelCase		:							List[Any]   					=   pipe(**_UpperCAmelCase    ,			control_guidance_start=0.1    ,			control_guidance_end=0.2			)[0]

							UpperCAmelCase		:							str   					=   self.get_dummy_inputs(_UpperCAmelCase			)
							UpperCAmelCase		:							Optional[Any]   					=   steps
							UpperCAmelCase		:							int   					=   scale
							UpperCAmelCase		:							List[str]   					=   pipe(**_UpperCAmelCase    ,			control_guidance_start=[0.1, 0.3]    ,			control_guidance_end=[0.2, 0.7]			)[0]

							UpperCAmelCase		:							str   					=   self.get_dummy_inputs(_UpperCAmelCase			)
							UpperCAmelCase		:							int   					=   steps
							UpperCAmelCase		:							str   					=   scale
							UpperCAmelCase		:							int   					=   pipe(**_UpperCAmelCase    ,			control_guidance_start=0.4    ,			control_guidance_end=[0.5, 0.8]			)[0]

							# make sure that all outputs are different
							assert np.sum(np.abs(output_a - output_a			)			) > 1e-3
							assert np.sum(np.abs(output_a - output_a			)			) > 1e-3
							assert np.sum(np.abs(output_a - output_a			)			) > 1e-3







		def       A_       ( self			):


							'''simple docstring'''

							return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3			)







		@unittest.skipIf(
		    torch_device != "cuda" or not is_xformers_available()    ,			reason="XFormers attention is only available with CUDA and `xformers` installed"    ,			)
		def       A_       ( self			):


							'''simple docstring'''

							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3			)







		def       A_       ( self			):


							'''simple docstring'''

							self._test_inference_batch_single_identical(expected_max_diff=2e-3			)







		def       A_       ( self			):


							'''simple docstring'''

							UpperCAmelCase		:							Optional[int]   					=   self.get_dummy_components()
							UpperCAmelCase		:							Dict   					=   self.pipeline_class(**_UpperCAmelCase			)
							pipe.to(_UpperCAmelCase			)
							pipe.set_progress_bar_config(disable=_UpperCAmelCase			)
							with tempfile.TemporaryDirectory() as tmpdir:
												try:
																	# save_pretrained is not implemented for Multi-ControlNet
																	pipe.save_pretrained(_UpperCAmelCase			)
												except NotImplementedError:
																	pass







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


		"""simple docstring"""

		def       A_       ( self			):


							'''simple docstring'''

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







		def       A_       ( self			):


							'''simple docstring'''

							UpperCAmelCase		:							Tuple   					=   ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"			)

							UpperCAmelCase		:							Tuple   					=   StableDiffusionControlNetImgaImgPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"    ,			safety_checker=_UpperCAmelCase    ,			controlnet=_UpperCAmelCase			)
							pipe.enable_model_cpu_offload()
							pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

							UpperCAmelCase		:							Tuple   					=   torch.Generator(device="cpu"			).manual_seed(0			)
							UpperCAmelCase		:							Union[str, Any]   					=   'evil space-punk bird'
							UpperCAmelCase		:							str   					=   load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"			).resize((5_1_2, 5_1_2)			)
							UpperCAmelCase		:							Optional[int]   					=   load_image(
							    "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"			).resize((5_1_2, 5_1_2)			)

							UpperCAmelCase		:							Optional[Any]   					=   pipe(
							    _UpperCAmelCase    ,			_UpperCAmelCase    ,			control_image=_UpperCAmelCase    ,			generator=_UpperCAmelCase    ,			output_type="np"    ,			num_inference_steps=5_0    ,			strength=0.6    ,			)

							UpperCAmelCase		:							Tuple   					=   output.images[0]

							assert image.shape == (5_1_2, 5_1_2, 3)

							UpperCAmelCase		:							Union[str, Any]   					=   load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy"			)

							assert np.abs(expected_image - image			).max() < 9e-2





'''simple docstring'''

from pathlib import Path

import cva
import numpy as np
from matplotlib import pyplot as plt



def      __lowerCamelCase     (      lowerCAmelCase_						,			lowerCAmelCase_						,			lowerCAmelCase_						,			lowerCAmelCase_						,			lowerCAmelCase_  )			->      np.ndarray:
		_a     :      Union[str, Any]        =							cva.getAffineTransform(lowerCAmelCase_						,			lowerCAmelCase_  )
		return cva.warpAffine(lowerCAmelCase_						,			lowerCAmelCase_						,			(rows, cols)  )


if __name__ == "__main__":
							# read original image
							__lowerCAmelCase									=   cva.imread(
							    str(Path(__file__).resolve().parent.parent / '''image_data''' / '''lena.jpg''')
							)
							# turn image in gray scale value
							__lowerCAmelCase									=   cva.cvtColor(image, cva.COLOR_BGR2GRAY)
							# get image shape
							__lowerCAmelCase			,					__lowerCAmelCase									=   gray_img.shape

							# set different points to rotate image
							__lowerCAmelCase									=   np.array([[50, 50], [200, 50], [50, 200]], np.floataa)
							__lowerCAmelCase									=   np.array([[10, 100], [200, 50], [100, 250]], np.floataa)
							__lowerCAmelCase									=   np.array([[50, 50], [150, 50], [120, 200]], np.floataa)
							__lowerCAmelCase									=   np.array([[10, 100], [80, 50], [180, 250]], np.floataa)

							# add all rotated images in a list
							__lowerCAmelCase									=   [
							    gray_img,
							    get_rotation(gray_img, ptsa, ptsa, img_rows, img_cols),
							    get_rotation(gray_img, ptsa, ptsa, img_rows, img_cols),
							    get_rotation(gray_img, ptsa, ptsa, img_rows, img_cols),
							]

							# plot different image rotations
							__lowerCAmelCase									=   plt.figure(1)
							__lowerCAmelCase									=   ['''Original''', '''Rotation 1''', '''Rotation 2''', '''Rotation 3''']
							for i, image in enumerate(images):
														plt.subplot(2, 2, i + 1), plt.imshow(image, '''gray''')
														plt.title(titles[i])
														plt.axis('''off''')
														plt.subplots_adjust(left=0.0, bottom=0.05, right=1.0, top=0.95)
							plt.show()




