






"""simple docstring"""
# 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  (				__lowerCAmelCase       ,       __lowerCAmelCase       ,       __lowerCAmelCase       ,       unittest.TestCase						):





	"""simple docstring"""




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





	def    lowerCamelCase(self							):
		torch.manual_seed(0							)
		A_		:						Dict	    =   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							)
		A_		:						Optional[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							)
		A_		:						Optional[Any]	    =   DDIMScheduler(
		    beta_start=0.00085     ,     beta_end=0.012     ,     beta_schedule="""scaled_linear"""     ,     clip_sample=lowerCamelCase__     ,     set_alpha_to_one=lowerCamelCase__     ,     )
		torch.manual_seed(0							)
		A_		:						List[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							)
		A_		:						Tuple	    =   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=1000     ,     )
		A_		:						Any	    =   CLIPTextModel(lowerCamelCase__							)
		A_		:						int	    =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

		A_		:						List[str]	    =   {
		    '''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     ,     lowerCAmelCase_     ,     lowerCAmelCase_=0							):






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

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





	def    lowerCamelCase(self							):
		self._test_inference_batch_single_identical(expected_max_diff=2e-3							)






class      SCREAMING_SNAKE_CASE  (				__lowerCAmelCase       ,       __lowerCAmelCase       ,       unittest.TestCase						):





	"""simple docstring"""




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





	def    lowerCamelCase(self							):
		torch.manual_seed(0							)
		A_		:						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(lowerCAmelCase_							):
			if isinstance(lowerCamelCase__     ,     torch.nn.Convad							):
				torch.nn.init.normal(m.weight							)
				m.bias.data.fill_(1.0							)

		A_		:						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(lowerCamelCase__							)

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

		torch.manual_seed(0							)
		A_		:						int	    =   DDIMScheduler(
		    beta_start=0.00085     ,     beta_end=0.012     ,     beta_schedule="""scaled_linear"""     ,     clip_sample=lowerCamelCase__     ,     set_alpha_to_one=lowerCamelCase__     ,     )
		torch.manual_seed(0							)
		A_		:						List[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							)
		A_		:						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=1000     ,     )
		A_		:						Union[str, Any]	    =   CLIPTextModel(lowerCamelCase__							)
		A_		:						Optional[int]	    =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

		A_		:						List[str]	    =   MultiControlNetModel([controlneta, controlneta]							)

		A_		:						List[str]	    =   {
		    '''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     ,     lowerCAmelCase_     ,     lowerCAmelCase_=0							):






		if str(lowerCamelCase__							).startswith("""mps"""							):
			A_		:						Tuple	    =   torch.manual_seed(lowerCamelCase__							)
		else:
			A_		:						Union[str, Any]	    =   torch.Generator(device=lowerCamelCase__							).manual_seed(lowerCamelCase__							)

		A_		:						str	    =   2

		A_		:						Tuple	    =   [
		    randn_tensor(
		        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)     ,     generator=lowerCamelCase__     ,     device=torch.device(lowerCamelCase__							)     ,     ),
		    randn_tensor(
		        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)     ,     generator=lowerCamelCase__     ,     device=torch.device(lowerCamelCase__							)     ,     ),
		]

		A_		:						List[str]	    =   floats_tensor(control_image[0].shape     ,     rng=random.Random(lowerCamelCase__							)							).to(lowerCamelCase__							)
		A_		:						Tuple	    =   image.cpu().permute(0     ,     2     ,     3     ,     1							)[0]
		A_		:						Optional[int]	    =   Image.fromarray(np.uinta(lowerCamelCase__							)							).convert("""RGB"""							).resize((64, 64)							)
		A_		:						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    lowerCamelCase(self							):
		A_		:						List[Any]	    =   self.get_dummy_components()
		A_		:						List[Any]	    =   self.pipeline_class(**lowerCamelCase__							)
		pipe.to(lowerCamelCase__							)

		A_		:						Union[str, Any]	    =   10.0
		A_		:						Optional[int]	    =   4

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

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

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

		A_		:						Tuple	    =   self.get_dummy_inputs(lowerCamelCase__							)
		A_		:						Optional[int]	    =   steps
		A_		:						Optional[Any]	    =   scale
		A_		:						Optional[int]	    =   pipe(**lowerCamelCase__     ,     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							):
		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							):
		self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3							)





	def    lowerCamelCase(self							):
		self._test_inference_batch_single_identical(expected_max_diff=2e-3							)





	def    lowerCamelCase(self							):
		A_		:						Any	    =   self.get_dummy_components()
		A_		:						int	    =   self.pipeline_class(**lowerCamelCase__							)
		pipe.to(lowerCamelCase__							)
		pipe.set_progress_bar_config(disable=lowerCamelCase__							)
		with tempfile.TemporaryDirectory() as tmpdir:
			try:
				# save_pretrained is not implemented for Multi-ControlNet
				pipe.save_pretrained(lowerCamelCase__							)
			except NotImplementedError:
				pass






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





	"""simple docstring"""




	def    lowerCamelCase(self							):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def    lowerCamelCase(self							):
		A_		:						str	    =   ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""							)

		A_		:						int	    =   StableDiffusionControlNetImgaImgPipeline.from_pretrained(
		    """runwayml/stable-diffusion-v1-5"""     ,     safety_checker=lowerCamelCase__     ,     controlnet=lowerCamelCase__							)
		pipe.enable_model_cpu_offload()
		pipe.set_progress_bar_config(disable=lowerCamelCase__							)

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

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

		A_		:						Dict	    =   output.images[0]

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

		A_		:						Optional[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"""
def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__					):
	return not any(
	    neighbour == 1 and colored_vertices[i] == color
	    for i, neighbour in enumerate(snake_case__					)					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__					):

	# Base Case
	if index == len(snake_case__					):
		return True

	# Recursive Step
	for i in range(snake_case__					):
		if valid_coloring(graph[index]   ,   snake_case__   ,   snake_case__					):
			# Color current vertex
			A_		:						Dict	    =   i
			# Validate coloring
			if util_color(snake_case__   ,   snake_case__   ,   snake_case__   ,   index + 1					):
				return True
			# Backtrack
			A_		:						Union[str, Any]	    =   -1
	return False






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__					):
	A_		:						int	    =   [-1] * len(snake_case__					)

	if util_color(snake_case__   ,   snake_case__   ,   snake_case__   ,   0					):
		return colored_vertices

	return []

