'''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     (					SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , unittest.TestCase	):







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





				def  snake_case__				(						self					:						List[Any]					)							->Any:


										'''simple docstring'''




										torch.manual_seed(0					)
										_UpperCamelCase			:			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					)
										_UpperCamelCase			:			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)       ,		)
										torch.manual_seed(0					)
										_UpperCamelCase			:			Optional[int]   			=						DDIMScheduler(
										    beta_start=0.0_0_0_8_5       ,		beta_end=0.0_1_2       ,		beta_schedule="scaled_linear"       ,		clip_sample=lowercase__       ,		set_alpha_to_one=lowercase__       ,		)
										torch.manual_seed(0					)
										_UpperCamelCase			:			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					)
										_UpperCamelCase			:			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=1_000       ,		)
										_UpperCamelCase			:			Tuple   			=						CLIPTextModel(lowercase__					)
										_UpperCamelCase			:			Optional[Any]   			=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

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





				def  snake_case__				(						self					:						List[Any]       ,		lowercase__					:						Tuple       ,		lowercase__					:						Union[str, Any]=0					)							->Optional[Any]:


										'''simple docstring'''




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

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


										'''simple docstring'''




										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)







				def  snake_case__				(						self					:						Tuple					)							->Optional[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'''
				UpperCAmelCase__							=							StableDiffusionControlNetImgaImgPipeline
				UpperCAmelCase__							=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
				UpperCAmelCase__							=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				UpperCAmelCase__							=							frozenset([]	)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess





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


										'''simple docstring'''




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

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

										_UpperCamelCase			:			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)       ,		)
										controlneta.controlnet_down_blocks.apply(lowercase__					)

										torch.manual_seed(0					)
										_UpperCamelCase			:			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(lowercase__					)

										torch.manual_seed(0					)
										_UpperCamelCase			:			str   			=						DDIMScheduler(
										    beta_start=0.0_0_0_8_5       ,		beta_end=0.0_1_2       ,		beta_schedule="scaled_linear"       ,		clip_sample=lowercase__       ,		set_alpha_to_one=lowercase__       ,		)
										torch.manual_seed(0					)
										_UpperCamelCase			:			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					)
										_UpperCamelCase			:			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=1_000       ,		)
										_UpperCamelCase			:			Union[str, Any]   			=						CLIPTextModel(lowercase__					)
										_UpperCamelCase			:			str   			=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

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

										_UpperCamelCase			:			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  snake_case__				(						self					:						int       ,		lowercase__					:						Optional[Any]       ,		lowercase__					:						Union[str, Any]=0					)							->Optional[int]:


										'''simple docstring'''




										if str(lowercase__					).startswith("mps"					):
																_UpperCamelCase			:			List[str]   			=						torch.manual_seed(lowercase__					)
										else:
																_UpperCamelCase			:			Optional[int]   			=						torch.Generator(device=lowercase__					).manual_seed(lowercase__					)

										_UpperCamelCase			:			Tuple   			=						2

										_UpperCamelCase			:			Optional[int]   			=						[
										    randn_tensor(
										        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,		generator=lowercase__       ,		device=torch.device(lowercase__					)       ,		),
										    randn_tensor(
										        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,		generator=lowercase__       ,		device=torch.device(lowercase__					)       ,		),
										]

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


										'''simple docstring'''




										_UpperCamelCase			:			Optional[Any]   			=						self.get_dummy_components()
										_UpperCamelCase			:			Optional[Any]   			=						self.pipeline_class(**lowercase__					)
										pipe.to(lowercase__					)

										_UpperCamelCase			:			Dict   			=						1_0.0
										_UpperCamelCase			:			List[Any]   			=						4

										_UpperCamelCase			:			Any   			=						self.get_dummy_inputs(lowercase__					)
										_UpperCamelCase			:			List[Any]   			=						steps
										_UpperCamelCase			:			Dict   			=						scale
										_UpperCamelCase			:			Any   			=						pipe(**lowercase__					)[0]

										_UpperCamelCase			:			Optional[int]   			=						self.get_dummy_inputs(lowercase__					)
										_UpperCamelCase			:			Any   			=						steps
										_UpperCamelCase			:			Dict   			=						scale
										_UpperCamelCase			:			Dict   			=						pipe(**lowercase__       ,		control_guidance_start=0.1       ,		control_guidance_end=0.2					)[0]

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

										_UpperCamelCase			:			Any   			=						self.get_dummy_inputs(lowercase__					)
										_UpperCamelCase			:			Optional[Any]   			=						steps
										_UpperCamelCase			:			Any   			=						scale
										_UpperCamelCase			:			int   			=						pipe(**lowercase__       ,		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  snake_case__				(						self					:						int					)							->Dict:


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


										'''simple docstring'''




										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)





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


										'''simple docstring'''




										self._test_inference_batch_single_identical(expected_max_diff=2e-3					)







				def  snake_case__				(						self					:						Any					)							->List[Any]:


										'''simple docstring'''




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





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







				'''simple docstring'''





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


										'''simple docstring'''




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







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


										'''simple docstring'''




										_UpperCamelCase			:			Union[str, Any]   			=						ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"					)

										_UpperCamelCase			:			str   			=						StableDiffusionControlNetImgaImgPipeline.from_pretrained(
										    "runwayml/stable-diffusion-v1-5"       ,		safety_checker=lowercase__       ,		controlnet=lowercase__					)
										pipe.enable_model_cpu_offload()
										pipe.set_progress_bar_config(disable=lowercase__					)

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

										_UpperCamelCase			:			Optional[int]   			=						pipe(
										    lowercase__       ,		lowercase__       ,		control_image=lowercase__       ,		generator=lowercase__       ,		output_type="np"       ,		num_inference_steps=50       ,		strength=0.6       ,		)

										_UpperCamelCase			:			int   			=						output.images[0]

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

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



'''simple docstring'''





from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_	:		Optional[int]			=     logging.get_logger(__name__)

lowerCAmelCase_	:		Tuple			=     {
    """huggingface/time-series-transformer-tourism-monthly""": (
        """https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"""
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}




class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							'''time_series_transformer'''
				UpperCAmelCase__							=							{
				    '''hidden_size''': '''d_model''',
				    '''num_attention_heads''': '''encoder_attention_heads''',
				    '''num_hidden_layers''': '''encoder_layers''',
				}





				def __init__(						self					:						int       ,		lowercase__					:						Optional[int] = None       ,		lowercase__					:						Optional[int] = None       ,		lowercase__					:						str = "student_t"       ,		lowercase__					:						str = "nll"       ,		lowercase__					:						int = 1       ,		lowercase__					:						List[int] = [1, 2, 3, 4, 5, 6, 7]       ,		lowercase__					:						Optional[Union[str, bool]] = "mean"       ,		lowercase__					:						int = 0       ,		lowercase__					:						int = 0       ,		lowercase__					:						int = 0       ,		lowercase__					:						int = 0       ,		lowercase__					:						Optional[List[int]] = None       ,		lowercase__					:						Optional[List[int]] = None       ,		lowercase__					:						int = 32       ,		lowercase__					:						int = 32       ,		lowercase__					:						int = 2       ,		lowercase__					:						int = 2       ,		lowercase__					:						int = 2       ,		lowercase__					:						int = 2       ,		lowercase__					:						bool = True       ,		lowercase__					:						str = "gelu"       ,		lowercase__					:						int = 64       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						int = 100       ,		lowercase__					:						float = 0.0_2       ,		lowercase__					:						Optional[int]=True       ,		**lowercase__					:						Optional[Any]       ,		)							->Optional[Any]:


										'''simple docstring'''




										_UpperCamelCase			:			int   			=						prediction_length
										_UpperCamelCase			:			Optional[Any]   			=						context_length or prediction_length
										_UpperCamelCase			:			List[str]   			=						distribution_output
										_UpperCamelCase			:			Optional[Any]   			=						loss
										_UpperCamelCase			:			Tuple   			=						input_size
										_UpperCamelCase			:			Optional[int]   			=						num_time_features
										_UpperCamelCase			:			Union[str, Any]   			=						lags_sequence
										_UpperCamelCase			:			int   			=						scaling
										_UpperCamelCase			:			Dict   			=						num_dynamic_real_features
										_UpperCamelCase			:			str   			=						num_static_real_features
										_UpperCamelCase			:			Any   			=						num_static_categorical_features
										if cardinality and num_static_categorical_features > 0:
																if len(lowercase__					) != num_static_categorical_features:
																						raise ValueError(
																						    "The cardinality should be a list of the same length as `num_static_categorical_features`"					)
																_UpperCamelCase			:			Dict   			=						cardinality
										else:
																_UpperCamelCase			:			str   			=						[0]
										if embedding_dimension and num_static_categorical_features > 0:
																if len(lowercase__					) != num_static_categorical_features:
																						raise ValueError(
																						    "The embedding dimension should be a list of the same length as `num_static_categorical_features`"					)
																_UpperCamelCase			:			int   			=						embedding_dimension
										else:
																_UpperCamelCase			:			Optional[Any]   			=						[min(50       ,		(cat + 1) // 2					) for cat in self.cardinality]
										_UpperCamelCase			:			Union[str, Any]   			=						num_parallel_samples

										# Transformer architecture configuration
										_UpperCamelCase			:			int   			=						input_size * len(lowercase__					) + self._number_of_features
										_UpperCamelCase			:			Optional[Any]   			=						d_model
										_UpperCamelCase			:			str   			=						encoder_attention_heads
										_UpperCamelCase			:			List[Any]   			=						decoder_attention_heads
										_UpperCamelCase			:			str   			=						encoder_ffn_dim
										_UpperCamelCase			:			List[str]   			=						decoder_ffn_dim
										_UpperCamelCase			:			Union[str, Any]   			=						encoder_layers
										_UpperCamelCase			:			Optional[int]   			=						decoder_layers

										_UpperCamelCase			:			int   			=						dropout
										_UpperCamelCase			:			Optional[int]   			=						attention_dropout
										_UpperCamelCase			:			int   			=						activation_dropout
										_UpperCamelCase			:			List[Any]   			=						encoder_layerdrop
										_UpperCamelCase			:			int   			=						decoder_layerdrop

										_UpperCamelCase			:			List[Any]   			=						activation_function
										_UpperCamelCase			:			Any   			=						init_std

										_UpperCamelCase			:			Optional[Any]   			=						use_cache

										super().__init__(is_encoder_decoder=lowercase__       ,		**lowercase__					)







				@property
				def  snake_case__				(						self					:						Union[str, Any]					)							->int:


										'''simple docstring'''




										return (
										    sum(self.embedding_dimension					)
										    + self.num_dynamic_real_features
										    + self.num_time_features
										    + self.num_static_real_features
										    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
										)


