
# 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						(		A_			,A_			,A_			,unittest.TestCase					):
   A__   :						Union[str, Any]			=	StableDiffusionControlNetImgaImgPipeline
   A__   :						Union[str, Any]			=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   A__   :						List[str]			=	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   A__   :						List[str]			=	IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}					)
   A__   :						List[str]			=	IMAGE_TO_IMAGE_IMAGE_PARAMS



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


          '''simple docstring'''


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

          snake_case			:  str     							= {
              "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						:   List[Any]       ,		snake_case__						:   Union[str, Any]       ,		snake_case__						:   Tuple=0					)						->      str:


          '''simple docstring'''


          if str(snake_case__					).startswith("mps"					):
                 snake_case			:  Tuple     							= torch.manual_seed(snake_case__					)
          else:
                 snake_case			:  Optional[Any]     							= torch.Generator(device=snake_case__					).manual_seed(snake_case__					)

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


          '''simple docstring'''


          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)



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


          '''simple docstring'''


          self._test_inference_batch_single_identical(expected_max_diff=2e-3					)






class     UpperCAmelCase						(		A_			,A_			,unittest.TestCase					):
   A__   :						Union[str, Any]			=	StableDiffusionControlNetImgaImgPipeline
   A__   :						Dict			=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   A__   :						int			=	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   A__   :						List[str]			=	frozenset([]					)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess



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


          '''simple docstring'''


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

          snake_case			:  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(snake_case__					)

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

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

          snake_case			:  Any     							= MultiControlNetModel([controlneta, controlneta]					)

          snake_case			:  Union[str, Any]     							= {
              "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						:   List[str]       ,		snake_case__						:   Tuple       ,		snake_case__						:   List[Any]=0					)						->      int:


          '''simple docstring'''


          if str(snake_case__					).startswith("mps"					):
                 snake_case			:  Tuple     							= torch.manual_seed(snake_case__					)
          else:
                 snake_case			:  List[str]     							= torch.Generator(device=snake_case__					).manual_seed(snake_case__					)

          snake_case			:  str     							= 2

          snake_case			:  int     							= [
              randn_tensor(
                  (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,		generator=snake_case__       ,		device=torch.device(snake_case__					)       ,		),
              randn_tensor(
                  (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,		generator=snake_case__       ,		device=torch.device(snake_case__					)       ,		),
          ]

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


          '''simple docstring'''


          snake_case			:  Tuple     							= self.get_dummy_components()
          snake_case			:  Optional[int]     							= self.pipeline_class(**snake_case__					)
          pipe.to(snake_case__					)

          snake_case			:  Optional[int]     							= 10.0
          snake_case			:  str     							= 4

          snake_case			:  Tuple     							= self.get_dummy_inputs(snake_case__					)
          snake_case			:  Any     							= steps
          snake_case			:  Any     							= scale
          snake_case			:  Optional[int]     							= pipe(**snake_case__					)[0]

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

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

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


          '''simple docstring'''


          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)



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


          '''simple docstring'''


          self._test_inference_batch_single_identical(expected_max_diff=2e-3					)



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


          '''simple docstring'''


          snake_case			:  Dict     							= self.get_dummy_components()
          snake_case			:  int     							= self.pipeline_class(**snake_case__					)
          pipe.to(snake_case__					)
          pipe.set_progress_bar_config(disable=snake_case__					)
          with tempfile.TemporaryDirectory() as tmpdir:
                 try:
                        # save_pretrained is not implemented for Multi-ControlNet
                        pipe.save_pretrained(snake_case__					)
                 except NotImplementedError:
                        pass






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



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


          '''simple docstring'''


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



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


          '''simple docstring'''


          snake_case			:  Any     							= ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"					)

          snake_case			:  List[str]     							= StableDiffusionControlNetImgaImgPipeline.from_pretrained(
              "runwayml/stable-diffusion-v1-5"       ,		safety_checker=snake_case__       ,		controlnet=snake_case__					)
          pipe.enable_model_cpu_offload()
          pipe.set_progress_bar_config(disable=snake_case__					)

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

          snake_case			:  Any     							= pipe(
              snake_case__       ,		snake_case__       ,		control_image=snake_case__       ,		generator=snake_case__       ,		output_type="np"       ,		num_inference_steps=50       ,		strength=0.6       ,		)

          snake_case			:  Any     							= output.images[0]

          assert image.shape == (5_12, 5_12, 3)

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







from __future__ import annotations

from collections.abc import Callable
from typing import Any, Generic, TypeVar

__lowerCamelCase				=			TypeVar("""T""")






class     UpperCAmelCase						(		Generic[T]					):



   def __init__(self						:   int       ,		snake_case__						:   list[T]       ,		snake_case__						:   Callable[[T, T], T]					)						->      None:


          '''simple docstring'''


          snake_case			:  Any | T     							= None

          snake_case			:  int     							= len(snake_case__					)
          snake_case			:  list[T]     							= [any_type for _ in range(self.N					)] + arr
          snake_case			:  str     							= fnc
          self.build()



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


          '''simple docstring'''


          for p in range(self.N - 1       ,		0       ,		-1					):
                 snake_case			:  int     							= self.fn(self.st[p * 2]       ,		self.st[p * 2 + 1]					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   Dict       ,		snake_case__						:   int       ,		snake_case__						:   T					)						->      None:


          '''simple docstring'''


          p += self.N
          snake_case			:  int     							= v
          while p > 1:
                 snake_case			:  List[str]     							= p // 2
                 snake_case			:  Tuple     							= self.fn(self.st[p * 2]       ,		self.st[p * 2 + 1]					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[Any]       ,		snake_case__						:   int       ,		snake_case__						:   int					)						->      T | None:  # noqa: E741


          '''simple docstring'''


          snake_case							,   snake_case			:  Optional[int]     							= l + self.N, r + self.N

          snake_case			:  T | None     							= None
          while l <= r:
                 if l % 2 == 1:
                        snake_case			:  List[str]     							= self.st[l] if res is None else self.fn(snake_case__       ,		self.st[l]					)
                 if r % 2 == 0:
                        snake_case			:  int     							= self.st[r] if res is None else self.fn(snake_case__       ,		self.st[r]					)
                 snake_case							,   snake_case			:  str     							= (l + 1) // 2, (r - 1) // 2
          return res


if __name__ == "__main__":
   from functools import reduce

   __lowerCamelCase				=			[1, 10, -2, 9, -3, 8, 4, -7, 5, 6, 11, -12]

   __lowerCamelCase				=			{
       0: 7,
       1: 2,
       2: 6,
       3: -14,
       4: 5,
       5: 4,
       6: 7,
       7: -10,
       8: 9,
       9: 10,
       10: 12,
       11: 1,
   }

   __lowerCamelCase				=			SegmentTree(test_array, min)
   __lowerCamelCase				=			SegmentTree(test_array, max)
   __lowerCamelCase				=			SegmentTree(test_array, lambda a, b: a + b)


   def      UpperCamelCase   (       ):
          for i in range(len(__lowerCamelCase       )       ):
                 for j in range(__lowerCamelCase							,   len(__lowerCamelCase       )       ):
                        snake_case			:  int     							= reduce(__lowerCamelCase							,   test_array[i : j + 1]       )
                        snake_case			:  Tuple     							= reduce(__lowerCamelCase							,   test_array[i : j + 1]       )
                        snake_case			:  Union[str, Any]     							= reduce(lambda __lowerCamelCase							,   __lowerCamelCase       : a + b							,   test_array[i : j + 1]       )
                        assert min_range == min_segment_tree.query(__lowerCamelCase							,   __lowerCamelCase       )
                        assert max_range == max_segment_tree.query(__lowerCamelCase							,   __lowerCamelCase       )
                        assert sum_range == sum_segment_tree.query(__lowerCamelCase							,   __lowerCamelCase       )

   test_all_segments()

   for index, value in test_updates.items():
      __lowerCamelCase				=			value
      min_segment_tree.update(index, value)
      max_segment_tree.update(index, value)
      sum_segment_tree.update(index, value)
      test_all_segments()





