


# 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 			_a				(     UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__	,      unittest.TestCase     ):
							_lowercase      :     List[str]     =					StableDiffusionControlNetImgaImgPipeline
							_lowercase      :     int     =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
							_lowercase      :     Union[str, Any]     =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							_lowercase      :     Tuple     =					IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}     )
							_lowercase      :     Tuple     =					IMAGE_TO_IMAGE_IMAGE_PARAMS

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





									"""simple docstring"""







									torch.manual_seed(0				)
									lowercase__				= 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__				= 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__				= 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__				= 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__				= 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       ,		)
									lowercase__				= CLIPTextModel(UpperCamelCase_				)
									lowercase__				= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

									lowercase__				= {
									    '''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       ,		UpperCamelCase_:      List[Any]       ,		UpperCamelCase_:      Any=0				)     ->  Union[str, Any]:





									"""simple docstring"""







									if str(UpperCamelCase_				).startswith('''mps'''				):
											lowercase__				= torch.manual_seed(UpperCamelCase_				)
									else:
											lowercase__				= torch.Generator(device=UpperCamelCase_				).manual_seed(UpperCamelCase_				)

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





									"""simple docstring"""







									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3				)


							def  lowerCamelCase_					(				self:      str				)     ->  Any:





									"""simple docstring"""







									self._test_inference_batch_single_identical(expected_max_diff=2E-3				)


class 			_a				(     UpperCamelCase__	,      UpperCamelCase__	,      unittest.TestCase     ):
							_lowercase      :     Optional[int]     =					StableDiffusionControlNetImgaImgPipeline
							_lowercase      :     List[Any]     =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
							_lowercase      :     List[Any]     =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							_lowercase      :     Dict     =					frozenset([]     )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

							def  lowerCamelCase_					(				self:      Dict				)     ->  List[Any]:





									"""simple docstring"""







									torch.manual_seed(0				)
									lowercase__				= 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_:      Optional[int]				):
											if isinstance(UpperCamelCase_       ,		torch.nn.Convad				):
													torch.nn.init.normal(m.weight				)
													m.bias.data.fill_(1.0				)

									lowercase__				= 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__				= 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__				= 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__				= 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__				= 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       ,		)
									lowercase__				= CLIPTextModel(UpperCamelCase_				)
									lowercase__				= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

									lowercase__				= MultiControlNetModel([controlneta, controlneta]				)

									lowercase__				= {
									    '''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:      Tuple       ,		UpperCamelCase_:      Optional[Any]       ,		UpperCamelCase_:      List[Any]=0				)     ->  Union[str, Any]:





									"""simple docstring"""







									if str(UpperCamelCase_				).startswith('''mps'''				):
											lowercase__				= torch.manual_seed(UpperCamelCase_				)
									else:
											lowercase__				= torch.Generator(device=UpperCamelCase_				).manual_seed(UpperCamelCase_				)

									lowercase__				= 2

									lowercase__				= [
									    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__				= floats_tensor(control_image[0].shape       ,		rng=random.Random(UpperCamelCase_				)				).to(UpperCamelCase_				)
									lowercase__				= image.cpu().permute(0       ,		2       ,		3       ,		1				)[0]
									lowercase__				= Image.fromarray(np.uinta(UpperCamelCase_				)				).convert('''RGB'''				).resize((64, 64)				)
									lowercase__				= {
									    '''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]				)     ->  int:





									"""simple docstring"""







									lowercase__				= self.get_dummy_components()
									lowercase__				= self.pipeline_class(**UpperCamelCase_				)
									pipe.to(UpperCamelCase_				)

									lowercase__				= 10.0
									lowercase__				= 4

									lowercase__				= self.get_dummy_inputs(UpperCamelCase_				)
									lowercase__				= steps
									lowercase__				= scale
									lowercase__				= pipe(**UpperCamelCase_				)[0]

									lowercase__				= self.get_dummy_inputs(UpperCamelCase_				)
									lowercase__				= steps
									lowercase__				= scale
									lowercase__				= pipe(**UpperCamelCase_       ,		control_guidance_start=0.1       ,		control_guidance_end=0.2				)[0]

									lowercase__				= self.get_dummy_inputs(UpperCamelCase_				)
									lowercase__				= steps
									lowercase__				= scale
									lowercase__				= pipe(**UpperCamelCase_       ,		control_guidance_start=[0.1, 0.3]       ,		control_guidance_end=[0.2, 0.7]				)[0]

									lowercase__				= self.get_dummy_inputs(UpperCamelCase_				)
									lowercase__				= steps
									lowercase__				= scale
									lowercase__				= 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:      Optional[Any]				)     ->  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:      Dict				)     ->  Any:





									"""simple docstring"""







									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3				)

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





									"""simple docstring"""







									self._test_inference_batch_single_identical(expected_max_diff=2E-3				)


							def  lowerCamelCase_					(				self:      List[str]				)     ->  Tuple:





									"""simple docstring"""







									lowercase__				= self.get_dummy_components()
									lowercase__				= 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 			_a				(     unittest.TestCase     ):

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





									"""simple docstring"""







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


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





									"""simple docstring"""







									lowercase__				= ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''				)

									lowercase__				= 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__				= torch.Generator(device='''cpu'''				).manual_seed(0				)
									lowercase__				= '''evil space-punk bird'''
									lowercase__				= load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''				).resize((512, 512)				)
									lowercase__				= load_image(
									    '''https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'''				).resize((512, 512)				)

									lowercase__				= pipe(
									    UpperCamelCase_       ,		UpperCamelCase_       ,		control_image=UpperCamelCase_       ,		generator=UpperCamelCase_       ,		output_type='''np'''       ,		num_inference_steps=50       ,		strength=0.6       ,		)

									lowercase__				= output.images[0]

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

									lowercase__				= 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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


_lowercase  =			{
    """configuration_groupvit""": [
        """GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """GroupViTConfig""",
        """GroupViTOnnxConfig""",
        """GroupViTTextConfig""",
        """GroupViTVisionConfig""",
    ],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowercase  =			[
	    """GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """GroupViTModel""",
	    """GroupViTPreTrainedModel""",
	    """GroupViTTextModel""",
	    """GroupViTVisionModel""",
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowercase  =			[
	    """TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """TFGroupViTModel""",
	    """TFGroupViTPreTrainedModel""",
	    """TFGroupViTTextModel""",
	    """TFGroupViTVisionModel""",
	]

if TYPE_CHECKING:
	from .configuration_groupvit import (
	    GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    GroupViTConfig,
	    GroupViTOnnxConfig,
	    GroupViTTextConfig,
	    GroupViTVisionConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_groupvit import (
		    GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    GroupViTModel,
		    GroupViTPreTrainedModel,
		    GroupViTTextModel,
		    GroupViTVisionModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_groupvit import (
		    TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFGroupViTModel,
		    TFGroupViTPreTrainedModel,
		    TFGroupViTTextModel,
		    TFGroupViTVisionModel,
		)

else:
	import sys

	_lowercase  =			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

