






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

					'''simple docstring'''

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



					def       lowerCAmelCase				(  self				:	Dict							)			->  Any:


												"""simple docstring"""

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

												__lowercase      :							Dict				=     {
												    """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				:	List[Any]		,     __a				:	str		,     __a				:	List[Any]=0							)			->  Optional[Any]:


												"""simple docstring"""

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

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


												"""simple docstring"""

												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3							)



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


												"""simple docstring"""

												self._test_inference_batch_single_identical(expected_max_diff=2E-3							)







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

					'''simple docstring'''

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



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


												"""simple docstring"""

												torch.manual_seed(0							)
												__lowercase      :							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				:	List[Any]							):
																			if isinstance(__a		,     torch.nn.Convad							):
																										torch.nn.init.normal(m.weight							)
																										m.bias.data.fill_(1.0							)

												__lowercase      :							Dict				=     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							)
												__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)		,     )
												controlneta.controlnet_down_blocks.apply(__a							)

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

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

												__lowercase      :							List[Any]				=     {
												    """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		,     __a				:	Union[str, Any]		,     __a				:	Optional[int]=0							)			->  Dict:


												"""simple docstring"""

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

												__lowercase      :							List[str]				=     2

												__lowercase      :							Optional[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							)		,     ),
												]

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


												"""simple docstring"""

												__lowercase      :							Dict				=     self.get_dummy_components()
												__lowercase      :							str				=     self.pipeline_class(**__a							)
												pipe.to(__a							)

												__lowercase      :							int				=     10.0
												__lowercase      :							Tuple				=     4

												__lowercase      :							Tuple				=     self.get_dummy_inputs(__a							)
												__lowercase      :							Tuple				=     steps
												__lowercase      :							Optional[int]				=     scale
												__lowercase      :							int				=     pipe(**__a							)[0]

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

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

												__lowercase      :							Any				=     self.get_dummy_inputs(__a							)
												__lowercase      :							List[str]				=     steps
												__lowercase      :							int				=     scale
												__lowercase      :							List[str]				=     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       lowerCAmelCase				(  self				:	Union[str, Any]							)			->  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				:	List[Any]							)			->  Tuple:


												"""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				:	Optional[Any]							)			->  str:


												"""simple docstring"""

												__lowercase      :							str				=     self.get_dummy_components()
												__lowercase      :							Dict				=     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 		lowerCAmelCase     ( unittest.TestCase			):

					'''simple docstring'''

					def       lowerCAmelCase				(  self				:	Dict							)			->  Any:


												"""simple docstring"""

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



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


												"""simple docstring"""

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

												__lowercase      :							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							)

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

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

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

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

												__lowercase      :							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







# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from typing import Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import SchedulerMixin

class 		lowerCAmelCase     ( __a      ,       __a			):

					'''simple docstring'''

					_A						:			str									=      1



					@register_to_config
					def __init__(  self				:	Optional[int]		,     __a				:	Tuple=2000		,     __a				:	List[str]=0.1		,     __a				:	str=20		,     __a				:	Optional[int]=1E-3							)			->  int:


												"""simple docstring"""

												__lowercase      :							Tuple				=     None
												__lowercase      :							Union[str, Any]				=     None
												__lowercase      :							int				=     None



					def       lowerCAmelCase				(  self				:	List[Any]		,     __a				:	Any		,     __a				:	Union[str, torch.device] = None							)			->  str:


												"""simple docstring"""

												__lowercase      :							List[str]				=     torch.linspace(1		,     self.config.sampling_eps		,     __a		,     device=__a							)



					def       lowerCAmelCase				(  self				:	Tuple		,     __a				:	List[Any]		,     __a				:	Tuple		,     __a				:	int		,     __a				:	Optional[int]=None							)			->  str:


												"""simple docstring"""

												if self.timesteps is None:
																			raise ValueError(
																			    """`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"""							)

												# TODO(Patrick) better comments + non-PyTorch
												# postprocess model score
												__lowercase      :							Dict				=     (
												    -0.25 * t**2 * (self.config.beta_max - self.config.beta_min) - 0.5 * t * self.config.beta_min
												)
												__lowercase      :							int				=     torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff							)							)
												__lowercase      :							Union[str, Any]				=     std.flatten()
												while len(std.shape							) < len(score.shape							):
																			__lowercase      :							Optional[Any]				=     std.unsqueeze(-1							)
												__lowercase      :							List[Any]				=     -score / std

												# compute
												__lowercase      :							Dict				=     -1.0 / len(self.timesteps							)

												__lowercase      :							int				=     self.config.beta_min + t * (self.config.beta_max - self.config.beta_min)
												__lowercase      :							List[Any]				=     beta_t.flatten()
												while len(beta_t.shape							) < len(x.shape							):
																			__lowercase      :							Union[str, Any]				=     beta_t.unsqueeze(-1							)
												__lowercase      :							List[str]				=     -0.5 * beta_t * x

												__lowercase      :							int				=     torch.sqrt(__a							)
												__lowercase      :							Union[str, Any]				=     drift - diffusion**2 * score
												__lowercase      :							Optional[Any]				=     x + drift * dt

												# add noise
												__lowercase      :							List[str]				=     randn_tensor(x.shape		,     layout=x.layout		,     generator=__a		,     device=x.device		,     dtype=x.dtype							)
												__lowercase      :							str				=     x_mean + diffusion * math.sqrt(-dt							) * noise

												return x, x_mean



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


												"""simple docstring"""

												return self.config.num_train_timesteps