


# 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      lowercase__     (			lowercase_,					lowercase_,					lowercase_,					unittest.TestCase	):
	a_         =StableDiffusionControlNetImgaImgPipeline
	a_         =TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
	a_         =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	a_         =IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""}	)
	a_         =IMAGE_TO_IMAGE_IMAGE_PARAMS

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							torch.manual_seed(0       )
							lowerCAmelCase__       =						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__       =						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__       =						DDIMScheduler(
							    beta_start=0.00_085		,    beta_end=0.012		,    beta_schedule="scaled_linear"		,    clip_sample=__UpperCAmelCase		,    set_alpha_to_one=__UpperCAmelCase		,    )
							torch.manual_seed(0       )
							lowerCAmelCase__       =						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__       =						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		,    )
							lowerCAmelCase__       =						CLIPTextModel(__UpperCAmelCase       )
							lowerCAmelCase__       =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

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

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase=0       )->					Optional[int]:







							'''simple docstring'''





							if str(__UpperCAmelCase       ).startswith("mps"       ):
													lowerCAmelCase__       =						torch.manual_seed(__UpperCAmelCase       )
							else:
													lowerCAmelCase__       =						torch.Generator(device=__UpperCAmelCase       ).manual_seed(__UpperCAmelCase       )

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







							'''simple docstring'''





							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3       )


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







							'''simple docstring'''





							self._test_inference_batch_single_identical(expected_max_diff=2E-3       )





class      lowercase__     (			lowercase_,					lowercase_,					unittest.TestCase	):
	a_         =StableDiffusionControlNetImgaImgPipeline
	a_         =TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
	a_         =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	a_         =frozenset([]	)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							torch.manual_seed(0       )
							lowerCAmelCase__       =						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       ):
													if isinstance(__UpperCAmelCase		,    torch.nn.Convad       ):
																			torch.nn.init.normal(m.weight       )
																			m.bias.data.fill_(1.0       )

							lowerCAmelCase__       =						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       )
							lowerCAmelCase__       =						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       )
							lowerCAmelCase__       =						DDIMScheduler(
							    beta_start=0.00_085		,    beta_end=0.012		,    beta_schedule="scaled_linear"		,    clip_sample=__UpperCAmelCase		,    set_alpha_to_one=__UpperCAmelCase		,    )
							torch.manual_seed(0       )
							lowerCAmelCase__       =						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__       =						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		,    )
							lowerCAmelCase__       =						CLIPTextModel(__UpperCAmelCase       )
							lowerCAmelCase__       =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

							lowerCAmelCase__       =						MultiControlNetModel([controlneta, controlneta]       )

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

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase=0       )->					int:







							'''simple docstring'''





							if str(__UpperCAmelCase       ).startswith("mps"       ):
													lowerCAmelCase__       =						torch.manual_seed(__UpperCAmelCase       )
							else:
													lowerCAmelCase__       =						torch.Generator(device=__UpperCAmelCase       ).manual_seed(__UpperCAmelCase       )

							lowerCAmelCase__       =						2

							lowerCAmelCase__       =						[
							    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       )		,    ),
							]

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







							'''simple docstring'''





							lowerCAmelCase__       =						self.get_dummy_components()
							lowerCAmelCase__       =						self.pipeline_class(**__UpperCAmelCase       )
							pipe.to(__UpperCAmelCase       )

							lowerCAmelCase__       =						10.0
							lowerCAmelCase__       =						4

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase       )
							lowerCAmelCase__       =						steps
							lowerCAmelCase__       =						scale
							lowerCAmelCase__       =						pipe(**__UpperCAmelCase       )[0]

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

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

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase       )
							lowerCAmelCase__       =						steps
							lowerCAmelCase__       =						scale
							lowerCAmelCase__       =						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 		UpperCAmelCase  (							self       )->					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 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3       )

	def 		UpperCAmelCase  (							self       )->					str:







							'''simple docstring'''





							self._test_inference_batch_single_identical(expected_max_diff=2E-3       )


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







							'''simple docstring'''





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

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







							'''simple docstring'''





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


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







							'''simple docstring'''





							lowerCAmelCase__       =						ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"       )

							lowerCAmelCase__       =						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       )

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

							lowerCAmelCase__       =						pipe(
							    __UpperCAmelCase		,    __UpperCAmelCase		,    control_image=__UpperCAmelCase		,    generator=__UpperCAmelCase		,    output_type="np"		,    num_inference_steps=50		,    strength=0.6		,    )

							lowerCAmelCase__       =						output.images[0]

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

							lowerCAmelCase__       =						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














# limitations under the License.

# NOTE: This file is deprecated and will be removed in a future version.
# It only exists so that temporarely `from diffusers.pipelines import DiffusionPipeline` works

from .pipelines import DiffusionPipeline, ImagePipelineOutput  # noqa: F401
from .utils import deprecate


deprecate(
    '''pipelines_utils''',
    '''0.22.0''',
    '''Importing `DiffusionPipeline` or `ImagePipelineOutput` from diffusers.pipeline_utils is deprecated. Please import from diffusers.pipelines.pipeline_utils instead.''',
    standard_warn=False,
    stacklevel=3,
)






