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




  '''simple docstring'''



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






  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->List[Any]:
        torch.manual_seed(0  )
        UpperCAmelCase_            =  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_            =  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_            =  DDIMScheduler(
            beta_start=0.0_0085  ,							beta_end=0.012  ,							beta_schedule='''scaled_linear'''  ,							clip_sample=UpperCAmelCase__  ,							set_alpha_to_one=UpperCAmelCase__  ,							)
        torch.manual_seed(0  )
        UpperCAmelCase_            =  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_            =  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  ,							)
        UpperCAmelCase_            =  CLIPTextModel(UpperCAmelCase__  )
        UpperCAmelCase_            =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

        UpperCAmelCase_            =  {
            '''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			:       str  ,							UpperCAmelCase__			:       Optional[int]  ,							UpperCAmelCase__			:       List[Any]=0  )				->str:
        if str(UpperCAmelCase__  ).startswith('''mps'''  ):
              UpperCAmelCase_            =  torch.manual_seed(UpperCAmelCase__  )
        else:
              UpperCAmelCase_            =  torch.Generator(device=UpperCAmelCase__  ).manual_seed(UpperCAmelCase__  )

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






  def  lowerCAmelCase__ (      self			:       str  )				->Union[str, Any]:
        self._test_inference_batch_single_identical(expected_max_diff=2e-3  )







class lowerCamelCase							(	lowerCamelCase     ,		lowerCamelCase     ,		unittest.TestCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	StableDiffusionControlNetImgaImgPipeline
  lowerCAmelCase__										=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
  lowerCAmelCase__										=	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowerCAmelCase__										=	frozenset([]					)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess






  def  lowerCAmelCase__ (      self			:       Dict  )				->Dict:
        torch.manual_seed(0  )
        UpperCAmelCase_            =  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__			:       List[str]  ):
              if isinstance(UpperCAmelCase__  ,							torch.nn.Convad  ):
                    torch.nn.init.normal(m.weight  )
                    m.bias.data.fill_(1.0  )

        UpperCAmelCase_            =  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  )
        UpperCAmelCase_            =  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  )
        UpperCAmelCase_            =  DDIMScheduler(
            beta_start=0.0_0085  ,							beta_end=0.012  ,							beta_schedule='''scaled_linear'''  ,							clip_sample=UpperCAmelCase__  ,							set_alpha_to_one=UpperCAmelCase__  ,							)
        torch.manual_seed(0  )
        UpperCAmelCase_            =  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_            =  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  ,							)
        UpperCAmelCase_            =  CLIPTextModel(UpperCAmelCase__  )
        UpperCAmelCase_            =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

        UpperCAmelCase_            =  MultiControlNetModel([controlneta, controlneta]  )

        UpperCAmelCase_            =  {
            '''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			:       Dict  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       Tuple=0  )				->Tuple:
        if str(UpperCAmelCase__  ).startswith('''mps'''  ):
              UpperCAmelCase_            =  torch.manual_seed(UpperCAmelCase__  )
        else:
              UpperCAmelCase_            =  torch.Generator(device=UpperCAmelCase__  ).manual_seed(UpperCAmelCase__  )

        UpperCAmelCase_            =  2

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

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

        UpperCAmelCase_            =  10.0
        UpperCAmelCase_            =  4

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

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

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

        UpperCAmelCase_            =  self.get_dummy_inputs(UpperCAmelCase__  )
        UpperCAmelCase_            =  steps
        UpperCAmelCase_            =  scale
        UpperCAmelCase_            =  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  lowerCAmelCase__ (      self			:       List[Any]  )				->Tuple:
        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			:       Optional[int]  )				->List[str]:
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3  )






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Any:
        self._test_inference_batch_single_identical(expected_max_diff=2e-3  )






  def  lowerCAmelCase__ (      self			:       int  )				->int:
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  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 lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       int  )				->Any:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->List[Any]:
        UpperCAmelCase_            =  ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''  )

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

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

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

        UpperCAmelCase_            =  output.images[0]

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

        UpperCAmelCase_            =  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 heapq import heappop, heappush

import numpy as np



def  __lowerCamelCase       (      _UpperCamelCase  :  np.ndarray							,    _UpperCamelCase  :  tuple[int, int]							,    _UpperCamelCase  :  tuple[int, int]							,    _UpperCamelCase  :  bool							,    ):


      '''simple docstring'''


      UpperCAmelCase_      ,		UpperCAmelCase_            =  grid.shape
      UpperCAmelCase_            =  [-1, 1, 0, 0]
      UpperCAmelCase_            =  [0, 0, -1, 1]
      if allow_diagonal:
            dx += [-1, -1, 1, 1]
            dy += [-1, 1, -1, 1]

      UpperCAmelCase_      ,		UpperCAmelCase_            =  [(0, source)], set()
      UpperCAmelCase_            =  np.full((rows, cols)							,    np.inf					)
      UpperCAmelCase_            =  0
      UpperCAmelCase_            =  np.empty((rows, cols)							,    dtype=_UpperCamelCase					)
      UpperCAmelCase_            =  None

      while queue:
            ((UpperCAmelCase_)      ,		(UpperCAmelCase_))            =  heappop(_UpperCamelCase					)
            if (x, y) in visited:
                  continue
            visited.add((x, y)					)

            if (x, y) == destination:
                  UpperCAmelCase_            =  []
                  while (x, y) != source:
                        path.append((x, y)					)
                        UpperCAmelCase_      ,		UpperCAmelCase_            =  predecessors[x, y]
                  path.append(_UpperCamelCase					)  # add the source manually
                  path.reverse()
                  return matrix[destination], path

            for i in range(len(_UpperCamelCase					)					):
                  UpperCAmelCase_      ,		UpperCAmelCase_            =  x + dx[i], y + dy[i]
                  if 0 <= nx < rows and 0 <= ny < cols:
                        UpperCAmelCase_            =  grid[nx][ny]
                        if next_node == 1 and matrix[nx, ny] > dist + 1:
                              heappush(_UpperCamelCase							,    (dist + 1, (nx, ny))					)
                              UpperCAmelCase_            =  dist + 1
                              UpperCAmelCase_            =  (x, y)

      return np.inf, []


if __name__ == "__main__":
      import doctest

      doctest.testmod()





