




# 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 				UpperCamelCase_    ( UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__		,					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        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->str:





           '''simple docstring'''

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

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


    def        SCREAMING_SNAKE_CASE							(			self      :	Any					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	Optional[int]=0)						->Tuple:





           '''simple docstring'''

           if str(UpperCAmelCase__).startswith('''mps'''):
                  A__       =    torch.manual_seed(UpperCAmelCase__)
           else:
                  A__       =    torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)

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





           '''simple docstring'''

           self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3)





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





           '''simple docstring'''

           self._test_inference_batch_single_identical(expected_max_diff=2e-3)


class 				UpperCamelCase_    ( UpperCAmelCase__		,					UpperCAmelCase__		,					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        SCREAMING_SNAKE_CASE							(			self      :	int)						->Optional[int]:





           '''simple docstring'''

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

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

           A__       =    MultiControlNetModel([controlneta, controlneta])

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


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any]					, UpperCAmelCase__      :	List[Any]					, UpperCAmelCase__      :	Union[str, Any]=0)						->Dict:





           '''simple docstring'''

           if str(UpperCAmelCase__).startswith('''mps'''):
                  A__       =    torch.manual_seed(UpperCAmelCase__)
           else:
                  A__       =    torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)

           A__       =    2

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

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





           '''simple docstring'''

           A__       =    self.get_dummy_components()
           A__       =    self.pipeline_class(**UpperCAmelCase__)
           pipe.to(UpperCAmelCase__)

           A__       =    10.0
           A__       =    4

           A__       =    self.get_dummy_inputs(UpperCAmelCase__)
           A__       =    steps
           A__       =    scale
           A__       =    pipe(**UpperCAmelCase__)[0]

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

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

           A__       =    self.get_dummy_inputs(UpperCAmelCase__)
           A__       =    steps
           A__       =    scale
           A__       =    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        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->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        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->int:





           '''simple docstring'''

           self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3)


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





           '''simple docstring'''

           self._test_inference_batch_single_identical(expected_max_diff=2e-3)





    def        SCREAMING_SNAKE_CASE							(			self      :	Any)						->int:





           '''simple docstring'''

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

    '''simple docstring'''






    def        SCREAMING_SNAKE_CASE							(			self      :	List[Any])						->List[str]:





           '''simple docstring'''

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





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





           '''simple docstring'''

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

           A__       =    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__)

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

           A__       =    pipe(
               UpperCAmelCase__					, UpperCAmelCase__					, control_image=UpperCAmelCase__					, generator=UpperCAmelCase__					, output_type='''np'''					, num_inference_steps=50					, strength=0.6					, )

           A__       =    output.images[0]

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

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








from __future__ import annotations

import typing
from collections.abc import Iterable

import numpy as np

_lowerCamelCase	:		str							       =   typing.Union[Iterable[float], Iterable[int], np.ndarray]  # noqa: UP007
_lowerCamelCase	:		Tuple							       =   typing.Union[np.floataa, int, float]  # noqa: UP007







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_						)   ->       VectorOut:



       """simple docstring"""

       return np.sqrt(np.sum((np.asarray(lowercase_						) - np.asarray(lowercase_						)) ** 2						)						)







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_						)   ->       VectorOut:



       """simple docstring"""

       return sum((va - va) ** 2 for va, va in zip(lowercase_ ,				lowercase_						)						) ** (1 / 2)


if __name__ == "__main__":







  def SCREAMING_SNAKE_CASE				(       )   ->       None:



         """simple docstring"""

         from timeit import timeit

         print('''Without Numpy'''						)
         print(
             timeit(
                 '''euclidean_distance_no_np([1, 2, 3], [4, 5, 6])''' ,				number=10_000 ,				globals=globals() ,				)						)
         print('''With Numpy'''						)
         print(
             timeit(
                 '''euclidean_distance([1, 2, 3], [4, 5, 6])''' ,				number=10_000 ,				globals=globals() ,				)						)
  benchmark()


