




# 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        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ,							__SCREAMING_SNAKE_CASE   ,							__SCREAMING_SNAKE_CASE   ,							unittest.TestCase   ):







 '''simple docstring'''
 lowercase_		=      StableDiffusionControlNetImgaImgPipeline
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowercase_		=      IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}   )
 lowercase_		=      IMAGE_TO_IMAGE_IMAGE_PARAMS

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->int:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:				str        =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)
   lowerCamelCase__:				List[str]        =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)
   lowerCamelCase__:				List[Any]        =DDIMScheduler(
       beta_start=0.0_0085	,			beta_end=0.012	,			beta_schedule="scaled_linear"	,			clip_sample=_a	,			set_alpha_to_one=_a	,			)
   torch.manual_seed(0)
   lowerCamelCase__:				Tuple        =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)
   lowerCamelCase__:				List[Any]        =CLIPTextConfig(
       bos_token_id=0	,			eos_token_id=2	,			hidden_size=32	,			intermediate_size=37	,			layer_norm_eps=1E-0_5	,			num_attention_heads=4	,			num_hidden_layers=5	,			pad_token_id=1	,			vocab_size=1_000	,			)
   lowerCamelCase__:				Any        =CLIPTextModel(_a)
   lowerCamelCase__:				int        =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:				Optional[Any]        ={
       "unet": unet,
       "controlnet": controlnet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "safety_checker": None,
       "feature_extractor": None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       Optional[int]	,			UpperCAmelCase_    :       Any=0)	->Tuple:



   '''simple docstring'''



   if str(_a).startswith("mps"):
     lowerCamelCase__:				Any        =torch.manual_seed(_a)
   else:
     lowerCamelCase__:				Optional[int]        =torch.Generator(device=_a).manual_seed(_a)

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



   '''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 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->Dict:



   '''simple docstring'''



   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3)







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



   '''simple docstring'''



   self._test_inference_batch_single_identical(expected_max_diff=2E-3)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ,							__SCREAMING_SNAKE_CASE   ,							unittest.TestCase   ):







 '''simple docstring'''
 lowercase_		=      StableDiffusionControlNetImgaImgPipeline
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowercase_		=      frozenset([]   )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int])	->List[str]:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:				Optional[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(UpperCAmelCase_    :       Union[str, Any]):
     if isinstance(_a	,			torch.nn.Convad):
       torch.nn.init.normal(m.weight)
       m.bias.data.fill_(1.0)

   lowerCamelCase__:				List[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(_a)

   torch.manual_seed(0)
   lowerCamelCase__:				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(_a)

   torch.manual_seed(0)
   lowerCamelCase__:				str        =DDIMScheduler(
       beta_start=0.0_0085	,			beta_end=0.012	,			beta_schedule="scaled_linear"	,			clip_sample=_a	,			set_alpha_to_one=_a	,			)
   torch.manual_seed(0)
   lowerCamelCase__:				Union[str, Any]        =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)
   lowerCamelCase__:				Union[str, Any]        =CLIPTextConfig(
       bos_token_id=0	,			eos_token_id=2	,			hidden_size=32	,			intermediate_size=37	,			layer_norm_eps=1E-0_5	,			num_attention_heads=4	,			num_hidden_layers=5	,			pad_token_id=1	,			vocab_size=1_000	,			)
   lowerCamelCase__:				Union[str, Any]        =CLIPTextModel(_a)
   lowerCamelCase__:				Any        =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:				Union[str, Any]        =MultiControlNetModel([controlneta, controlneta])

   lowerCamelCase__:				str        ={
       "unet": unet,
       "controlnet": controlnet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "safety_checker": None,
       "feature_extractor": None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       Union[str, Any]=0)	->List[Any]:



   '''simple docstring'''



   if str(_a).startswith("mps"):
     lowerCamelCase__:				List[str]        =torch.manual_seed(_a)
   else:
     lowerCamelCase__:				str        =torch.Generator(device=_a).manual_seed(_a)

   lowerCamelCase__:				Dict        =2

   lowerCamelCase__:				List[Any]        =[
       randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)	,			generator=_a	,			device=torch.device(_a)	,			),
       randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)	,			generator=_a	,			device=torch.device(_a)	,			),
   ]

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



   '''simple docstring'''



   lowerCamelCase__:				Dict        =self.get_dummy_components()
   lowerCamelCase__:				int        =self.pipeline_class(**_a)
   pipe.to(_a)

   lowerCamelCase__:				Dict        =10.0
   lowerCamelCase__:				str        =4

   lowerCamelCase__:				List[Any]        =self.get_dummy_inputs(_a)
   lowerCamelCase__:				Any        =steps
   lowerCamelCase__:				Union[str, Any]        =scale
   lowerCamelCase__:				Dict        =pipe(**_a)[0]

   lowerCamelCase__:				Union[str, Any]        =self.get_dummy_inputs(_a)
   lowerCamelCase__:				Any        =steps
   lowerCamelCase__:				str        =scale
   lowerCamelCase__:				str        =pipe(**_a	,			control_guidance_start=0.1	,			control_guidance_end=0.2)[0]

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

   lowerCamelCase__:				Optional[Any]        =self.get_dummy_inputs(_a)
   lowerCamelCase__:				Optional[Any]        =steps
   lowerCamelCase__:				Optional[Any]        =scale
   lowerCamelCase__:				Tuple        =pipe(**_a	,			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 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->int:



   '''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 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->Any:



   '''simple docstring'''



   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->Optional[int]:



   '''simple docstring'''



   self._test_inference_batch_single_identical(expected_max_diff=2E-3)







 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				List[Any]        =self.get_dummy_components()
   lowerCamelCase__:				List[str]        =self.pipeline_class(**_a)
   pipe.to(_a)
   pipe.set_progress_bar_config(disable=_a)
   with tempfile.TemporaryDirectory() as tmpdir:
     try:
       # save_pretrained is not implemented for Multi-ControlNet
       pipe.save_pretrained(_a)
     except NotImplementedError:
       pass



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







 '''simple docstring'''

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



   '''simple docstring'''



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







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



   '''simple docstring'''



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

   lowerCamelCase__:				List[str]        =StableDiffusionControlNetImgaImgPipeline.from_pretrained(
       "runwayml/stable-diffusion-v1-5"	,			safety_checker=_a	,			controlnet=_a)
   pipe.enable_model_cpu_offload()
   pipe.set_progress_bar_config(disable=_a)

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

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

   lowerCamelCase__:				Optional[Any]        =output.images[0]

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

   lowerCamelCase__:				Optional[int]        =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








import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment



@require_huggingface_suite
class        _SCREAMING_SNAKE_CASE						(     unittest.TestCase   ):







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int])	->Tuple:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =inspect.getfile(accelerate.test_utils)
   lowerCamelCase__:				str        =os.path.sep.join(
       mod_file.split(os.path.sep)[:-1] + ["scripts", "external_deps", "test_metrics.py"])

   from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

   lowerCamelCase__:				str        =test_metrics

 @require_cpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Any:



   '''simple docstring'''



   debug_launcher(self.test_metrics.main	,			num_processes=1)

 @require_cpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->int:



   '''simple docstring'''



   debug_launcher(self.test_metrics.main)

 @require_single_gpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       Any)	->int:



   '''simple docstring'''



   self.test_metrics.main()







 @require_multi_gpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->Dict:



   '''simple docstring'''



   print(F"""Found {torch.cuda.device_count()} devices.""")
   lowerCamelCase__:				Optional[Any]        =["torchrun", F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
   with patch_environment(omp_num_threads=1):
     execute_subprocess_async(UpperCAmelCase_	,			env=os.environ.copy())

