






# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

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 						lowerCAmelCase   (       UpperCamelCase_					,				UpperCamelCase_					,				UpperCamelCase_					,				unittest.TestCase	):

					'''simple docstring'''

					_A  :	Optional[Any]     =					StableDiffusionControlNetImgaImgPipeline
					_A  :	List[str]     =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
					_A  :	Union[str, Any]     =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					_A  :	Dict     =					IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}	)
					_A  :	Union[str, Any]     =					IMAGE_TO_IMAGE_IMAGE_PARAMS



					def 					lowerCAmelCase (     self :		List[str]							)	->     Dict:


												"""simple docstring"""

												torch.manual_seed(0							)
												__lowercase					:							Tuple			=   UNetaDConditionModel(
												    block_out_channels=(32, 64)						,   layers_per_block=2						,   sample_size=32						,   in_channels=4						,   out_channels=4						,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,   up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")						,   cross_attention_dim=32						,   )
												torch.manual_seed(0							)
												__lowercase					:							int			=   ControlNetModel(
												    block_out_channels=(32, 64)						,   layers_per_block=2						,   in_channels=4						,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,   cross_attention_dim=32						,   conditioning_embedding_out_channels=(16, 32)						,   )
												torch.manual_seed(0							)
												__lowercase					:							List[str]			=   DDIMScheduler(
												    beta_start=0.00085						,   beta_end=0.012						,   beta_schedule="""scaled_linear"""						,   clip_sample=UpperCamelCase__						,   set_alpha_to_one=UpperCamelCase__						,   )
												torch.manual_seed(0							)
												__lowercase					:							Dict			=   AutoencoderKL(
												    block_out_channels=[32, 64]						,   in_channels=3						,   out_channels=3						,   down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]						,   up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]						,   latent_channels=4						,   )
												torch.manual_seed(0							)
												__lowercase					:							Union[str, Any]			=   CLIPTextConfig(
												    bos_token_id=0						,   eos_token_id=2						,   hidden_size=32						,   intermediate_size=37						,   layer_norm_eps=1E-05						,   num_attention_heads=4						,   num_hidden_layers=5						,   pad_token_id=1						,   vocab_size=1000						,   )
												__lowercase					:							Optional[int]			=   CLIPTextModel(UpperCamelCase__							)
												__lowercase					:							str			=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

												__lowercase					:							Optional[int]			=   {
												    '''unet''': unet,
												    '''controlnet''': controlnet,
												    '''scheduler''': scheduler,
												    '''vae''': vae,
												    '''text_encoder''': text_encoder,
												    '''tokenizer''': tokenizer,
												    '''safety_checker''': None,
												    '''feature_extractor''': None,
												}
												return components



					def 					lowerCAmelCase (     self :		Any						,   __a :		Any						,   __a :		int=0							)	->     int:


												"""simple docstring"""

												if str(UpperCamelCase__							).startswith("""mps"""							):
																			__lowercase					:							List[Any]			=   torch.manual_seed(UpperCamelCase__							)
												else:
																			__lowercase					:							int			=   torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)

												__lowercase					:							Optional[int]			=   2
												__lowercase					:							Any			=   randn_tensor(
												    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,   generator=UpperCamelCase__						,   device=torch.device(UpperCamelCase__							)						,   )
												__lowercase					:							List[str]			=   floats_tensor(control_image.shape						,   rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
												__lowercase					:							List[str]			=   image.cpu().permute(0						,   2						,   3						,   1							)[0]
												__lowercase					:							Any			=   Image.fromarray(np.uinta(UpperCamelCase__							)							).convert("""RGB"""							).resize((64, 64)							)
												__lowercase					:							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 					lowerCAmelCase (     self :		Dict							)	->     List[str]:


												"""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 					lowerCAmelCase (     self :		Tuple							)	->     Dict:


												"""simple docstring"""

												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3							)



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


												"""simple docstring"""

												self._test_inference_batch_single_identical(expected_max_diff=2E-3							)







class 						lowerCAmelCase   (       UpperCamelCase_					,				UpperCamelCase_					,				unittest.TestCase	):

					'''simple docstring'''

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



					def 					lowerCAmelCase (     self :		Any							)	->     Optional[Any]:


												"""simple docstring"""

												torch.manual_seed(0							)
												__lowercase					:							Union[str, Any]			=   UNetaDConditionModel(
												    block_out_channels=(32, 64)						,   layers_per_block=2						,   sample_size=32						,   in_channels=4						,   out_channels=4						,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,   up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")						,   cross_attention_dim=32						,   )
												torch.manual_seed(0							)

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

												__lowercase					:							Tuple			=   ControlNetModel(
												    block_out_channels=(32, 64)						,   layers_per_block=2						,   in_channels=4						,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,   cross_attention_dim=32						,   conditioning_embedding_out_channels=(16, 32)						,   )
												controlneta.controlnet_down_blocks.apply(UpperCamelCase__							)

												torch.manual_seed(0							)
												__lowercase					:							Union[str, Any]			=   ControlNetModel(
												    block_out_channels=(32, 64)						,   layers_per_block=2						,   in_channels=4						,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,   cross_attention_dim=32						,   conditioning_embedding_out_channels=(16, 32)						,   )
												controlneta.controlnet_down_blocks.apply(UpperCamelCase__							)

												torch.manual_seed(0							)
												__lowercase					:							Optional[int]			=   DDIMScheduler(
												    beta_start=0.00085						,   beta_end=0.012						,   beta_schedule="""scaled_linear"""						,   clip_sample=UpperCamelCase__						,   set_alpha_to_one=UpperCamelCase__						,   )
												torch.manual_seed(0							)
												__lowercase					:							str			=   AutoencoderKL(
												    block_out_channels=[32, 64]						,   in_channels=3						,   out_channels=3						,   down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]						,   up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]						,   latent_channels=4						,   )
												torch.manual_seed(0							)
												__lowercase					:							Any			=   CLIPTextConfig(
												    bos_token_id=0						,   eos_token_id=2						,   hidden_size=32						,   intermediate_size=37						,   layer_norm_eps=1E-05						,   num_attention_heads=4						,   num_hidden_layers=5						,   pad_token_id=1						,   vocab_size=1000						,   )
												__lowercase					:							Any			=   CLIPTextModel(UpperCamelCase__							)
												__lowercase					:							Dict			=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

												__lowercase					:							Optional[int]			=   MultiControlNetModel([controlneta, controlneta]							)

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



					def 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		Optional[Any]						,   __a :		Union[str, Any]=0							)	->     List[Any]:


												"""simple docstring"""

												if str(UpperCamelCase__							).startswith("""mps"""							):
																			__lowercase					:							Tuple			=   torch.manual_seed(UpperCamelCase__							)
												else:
																			__lowercase					:							Dict			=   torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)

												__lowercase					:							Optional[Any]			=   2

												__lowercase					:							int			=   [
												    randn_tensor(
												        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,   generator=UpperCamelCase__						,   device=torch.device(UpperCamelCase__							)						,   ),
												    randn_tensor(
												        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,   generator=UpperCamelCase__						,   device=torch.device(UpperCamelCase__							)						,   ),
												]

												__lowercase					:							Any			=   floats_tensor(control_image[0].shape						,   rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
												__lowercase					:							List[str]			=   image.cpu().permute(0						,   2						,   3						,   1							)[0]
												__lowercase					:							List[str]			=   Image.fromarray(np.uinta(UpperCamelCase__							)							).convert("""RGB"""							).resize((64, 64)							)
												__lowercase					:							Dict			=   {
												    '''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 					lowerCAmelCase (     self :		Optional[Any]							)	->     Dict:


												"""simple docstring"""

												__lowercase					:							Dict			=   self.get_dummy_components()
												__lowercase					:							Union[str, Any]			=   self.pipeline_class(**UpperCamelCase__							)
												pipe.to(UpperCamelCase__							)

												__lowercase					:							Union[str, Any]			=   10.0
												__lowercase					:							Dict			=   4

												__lowercase					:							Dict			=   self.get_dummy_inputs(UpperCamelCase__							)
												__lowercase					:							Any			=   steps
												__lowercase					:							Optional[Any]			=   scale
												__lowercase					:							Dict			=   pipe(**UpperCamelCase__							)[0]

												__lowercase					:							Union[str, Any]			=   self.get_dummy_inputs(UpperCamelCase__							)
												__lowercase					:							List[str]			=   steps
												__lowercase					:							Dict			=   scale
												__lowercase					:							Dict			=   pipe(**UpperCamelCase__						,   control_guidance_start=0.1						,   control_guidance_end=0.2							)[0]

												__lowercase					:							Optional[int]			=   self.get_dummy_inputs(UpperCamelCase__							)
												__lowercase					:							int			=   steps
												__lowercase					:							Tuple			=   scale
												__lowercase					:							Optional[Any]			=   pipe(**UpperCamelCase__						,   control_guidance_start=[0.1, 0.3]						,   control_guidance_end=[0.2, 0.7]							)[0]

												__lowercase					:							Dict			=   self.get_dummy_inputs(UpperCamelCase__							)
												__lowercase					:							Union[str, Any]			=   steps
												__lowercase					:							Union[str, Any]			=   scale
												__lowercase					:							str			=   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 					lowerCAmelCase (     self :		str							)	->     Optional[Any]:


												"""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 					lowerCAmelCase (     self :		Tuple							)	->     List[str]:


												"""simple docstring"""

												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3							)



					def 					lowerCAmelCase (     self :		Optional[Any]							)	->     Union[str, Any]:


												"""simple docstring"""

												self._test_inference_batch_single_identical(expected_max_diff=2E-3							)



					def 					lowerCAmelCase (     self :		Any							)	->     Union[str, Any]:


												"""simple docstring"""

												__lowercase					:							Optional[int]			=   self.get_dummy_components()
												__lowercase					:							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 						lowerCAmelCase   (       unittest.TestCase	):

					'''simple docstring'''

					def 					lowerCAmelCase (     self :		Optional[int]							)	->     Dict:


												"""simple docstring"""

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



					def 					lowerCAmelCase (     self :		Optional[int]							)	->     int:


												"""simple docstring"""

												__lowercase					:							int			=   ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""							)

												__lowercase					:							Any			=   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__							)

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

												__lowercase					:							Any			=   pipe(
												    UpperCamelCase__						,   UpperCamelCase__						,   control_image=UpperCamelCase__						,   generator=UpperCamelCase__						,   output_type="""np"""						,   num_inference_steps=50						,   strength=0.6						,   )

												__lowercase					:							Tuple			=   output.images[0]

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

												__lowercase					:							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
from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
					import accelerate
def       A	(    _lowercase						):
						if not is_accelerate_available():
												return method
						SCREAMING_SNAKE_CASE							:    int							=							version.parse(accelerate.__version__						).base_version
						if version.parse(_lowercase						) < version.parse('''0.17.0'''						):
												return method

						def wrapper(self		,  *_lowercase		,  **_lowercase						):
												if hasattr(self		,  '''_hf_hook'''						) and hasattr(self._hf_hook		,  '''pre_forward'''						):
																		self._hf_hook.pre_forward(self						)
												return method(self		,  *_lowercase		,  **_lowercase						)

						return wrapper





