'''simple docstring'''





import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map





def       lowerCAmelCase_							(   snake_case_			:					Union[dict, list, tuple, torch.Tensor]     )       -> List[Tuple[int, ...]]:


     '''simple docstring'''




     UpperCAmelCase_			       =  []
     if isinstance(snake_case_   ,       snake_case_     ):
          for v in tree.values():
               shapes.extend(_fetch_dims(snake_case_     )     )
     elif isinstance(snake_case_   ,       (list, tuple)     ):
          for t in tree:
               shapes.extend(_fetch_dims(snake_case_     )     )
     elif isinstance(snake_case_   ,       torch.Tensor     ):
          shapes.append(tree.shape     )
     else:
          raise ValueError("Not supported"     )

     return shapes





@torch.jit.ignore
def       lowerCAmelCase_							(   snake_case_			:					int   ,       snake_case_			:					Tuple[int, ...]     )       -> Tuple[int, ...]:


     '''simple docstring'''




     UpperCAmelCase_			       =  []
     for d in reversed(snake_case_     ):
          idx.append(flat_idx % d     )
          UpperCAmelCase_			       =  flat_idx // d

     return tuple(reversed(snake_case_     )     )





@torch.jit.ignore
def       lowerCAmelCase_							(   snake_case_			:					Sequence[int]   ,       snake_case_			:					Sequence[int]   ,       snake_case_			:					Sequence[int]   ,       snake_case_			:					Optional[Sequence[bool]] = None   ,       snake_case_			:					Optional[Sequence[bool]] = None   ,       )       -> List[Tuple[slice, ...]]:


     '''simple docstring'''




     def reduce_edge_list(snake_case_			:					List[bool]     ) -> None:
          UpperCAmelCase_			       =  True
          for i in range(len(snake_case_     )     ):
               UpperCAmelCase_			       =  -1 * (i + 1)
               l[reversed_idx] &= tally
               UpperCAmelCase_			       =  l[reversed_idx]

     if start_edges is None:
          UpperCAmelCase_			       =  [s == 0 for s in start]
          reduce_edge_list(snake_case_     )
     if end_edges is None:
          UpperCAmelCase_			       =  [e == (d - 1) for e, d in zip(snake_case_   ,       snake_case_     )]
          reduce_edge_list(snake_case_     )

     # Base cases. Either start/end are empty and we're done, or the final,
     # one-dimensional tensor can be simply sliced
     if len(snake_case_     ) == 0:
          return [()]
     elif len(snake_case_     ) == 1:
          return [(slice(start[0]   ,       end[0] + 1     ),)]

     UpperCAmelCase_			       =  []
     UpperCAmelCase_			       =  []

     # Dimensions common to start and end can be selected directly
     for s, e in zip(snake_case_   ,       snake_case_     ):
          if s == e:
               path_list.append(slice(snake_case_   ,       s + 1     )     )
          else:
               break

     UpperCAmelCase_			       =  tuple(snake_case_     )
     UpperCAmelCase_			       =  len(snake_case_     )

     # start == end, and we're done
     if divergence_idx == len(snake_case_     ):
          return [path]

     def upper() -> Tuple[Tuple[slice, ...], ...]:
          assert start_edges is not None
          assert end_edges is not None

          UpperCAmelCase_			       =  start[divergence_idx]
          return tuple(
              path + (slice(snake_case_   ,       sdi + 1     ),) + s
              for s in _get_minimal_slice_set(
                  start[divergence_idx + 1 :]   ,       [d - 1 for d in dims[divergence_idx + 1 :]]   ,       dims[divergence_idx + 1 :]   ,       start_edges=start_edges[divergence_idx + 1 :]   ,       end_edges=[True for _ in end_edges[divergence_idx + 1 :]]   ,       )     )

     def lower() -> Tuple[Tuple[slice, ...], ...]:
          assert start_edges is not None
          assert end_edges is not None

          UpperCAmelCase_			       =  end[divergence_idx]
          return tuple(
              path + (slice(snake_case_   ,       edi + 1     ),) + s
              for s in _get_minimal_slice_set(
                  [0 for _ in start[divergence_idx + 1 :]]   ,       end[divergence_idx + 1 :]   ,       dims[divergence_idx + 1 :]   ,       start_edges=[True for _ in start_edges[divergence_idx + 1 :]]   ,       end_edges=end_edges[divergence_idx + 1 :]   ,       )     )

     # If both start and end are at the edges of the subtree rooted at
     # divergence_idx, we can just select the whole subtree at once
     if start_edges[divergence_idx] and end_edges[divergence_idx]:
          slices.append(path + (slice(start[divergence_idx]   ,       end[divergence_idx] + 1     ),)     )
     # If just start is at the edge, we can grab almost all of the subtree,
     # treating only the ragged bottom edge as an edge case
     elif start_edges[divergence_idx]:
          slices.append(path + (slice(start[divergence_idx]   ,       end[divergence_idx]     ),)     )
          slices.extend(lower()     )
     # Analogous to the previous case, but the top is ragged this time
     elif end_edges[divergence_idx]:
          slices.extend(upper()     )
          slices.append(path + (slice(start[divergence_idx] + 1   ,       end[divergence_idx] + 1     ),)     )
     # If both sides of the range are ragged, we need to handle both sides
     # separately. If there's contiguous meat in between them, we can index it
     # in one big chunk
     else:
          slices.extend(upper()     )
          UpperCAmelCase_			       =  end[divergence_idx] - start[divergence_idx]
          if middle_ground > 1:
               slices.append(path + (slice(start[divergence_idx] + 1   ,       end[divergence_idx]     ),)     )
          slices.extend(lower()     )

     return slices





@torch.jit.ignore
def       lowerCAmelCase_							(   snake_case_			:					torch.Tensor   ,       snake_case_			:					int   ,       snake_case_			:					int   ,       snake_case_			:					int     )       -> torch.Tensor:


     '''simple docstring'''




     UpperCAmelCase_			       =  t.shape[:no_batch_dims]
     UpperCAmelCase_			       =  list(_flat_idx_to_idx(snake_case_   ,       snake_case_     )     )
     # _get_minimal_slice_set is inclusive
     UpperCAmelCase_			       =  list(_flat_idx_to_idx(flat_end - 1   ,       snake_case_     )     )

     # Get an ordered list of slices to perform
     UpperCAmelCase_			       =  _get_minimal_slice_set(
         snake_case_   ,       snake_case_   ,       snake_case_   ,       )

     UpperCAmelCase_			       =  [t[s] for s in slices]

     return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]     ) for s in sliced_tensors]     )





def       lowerCAmelCase_							(   snake_case_			:					Callable   ,       snake_case_			:					Dict[str, Any]   ,       snake_case_			:					int   ,       snake_case_			:					int   ,       snake_case_			:					bool = False   ,       snake_case_			:					Any = None   ,       snake_case_			:					bool = False   ,       )       -> Any:


     '''simple docstring'''




     if not (len(snake_case_     ) > 0):
          raise ValueError("Must provide at least one input"     )

     UpperCAmelCase_			       =  [shape[:no_batch_dims] for shape in _fetch_dims(snake_case_     )]
     UpperCAmelCase_			       =  tuple([max(snake_case_     ) for s in zip(*snake_case_     )]     )

     def _prep_inputs(snake_case_			:					torch.Tensor     ) -> torch.Tensor:
          if not low_mem:
               if not sum(t.shape[:no_batch_dims]     ) == no_batch_dims:
                    UpperCAmelCase_			       =  t.expand(orig_batch_dims + t.shape[no_batch_dims:]     )
               UpperCAmelCase_			       =  t.reshape(-1   ,       *t.shape[no_batch_dims:]     )
          else:
               UpperCAmelCase_			       =  t.expand(orig_batch_dims + t.shape[no_batch_dims:]     )
          return t

     UpperCAmelCase_			       =  tensor_tree_map(_prep_inputs   ,       snake_case_     )
     UpperCAmelCase_			       =  None
     if _out is not None:
          UpperCAmelCase_			       =  tensor_tree_map(lambda snake_case_     : t.view([-1] + list(t.shape[no_batch_dims:]     )     )   ,       _out     )

     UpperCAmelCase_			       =  1
     for d in orig_batch_dims:
          flat_batch_dim *= d

     UpperCAmelCase_			       =  flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

     def _select_chunk(snake_case_			:					torch.Tensor     ) -> torch.Tensor:
          return t[i : i + chunk_size] if t.shape[0] != 1 else t

     UpperCAmelCase_			       =  0
     UpperCAmelCase_			       =  prepped_outputs
     for _ in range(snake_case_     ):
          # Chunk the input
          if not low_mem:
               UpperCAmelCase_			       =  _select_chunk
          else:
               UpperCAmelCase_			       =  partial(
                   _chunk_slice   ,       flat_start=snake_case_   ,       flat_end=min(snake_case_   ,       i + chunk_size     )   ,       no_batch_dims=len(snake_case_     )   ,       )

          UpperCAmelCase_			       =  tensor_tree_map(snake_case_   ,       snake_case_     )

          # Run the layer on the chunk
          UpperCAmelCase_			       =  layer(**snake_case_     )

          # Allocate space for the output
          if out is None:
               UpperCAmelCase_			       =  tensor_tree_map(lambda snake_case_     : t.new_zeros((flat_batch_dim,) + t.shape[1:]     )   ,       snake_case_     )

          # Put the chunk in its pre-allocated space
          if isinstance(snake_case_   ,       snake_case_     ):

               def assign(snake_case_			:					dict   ,       snake_case_			:					dict     ) -> None:
                    for k, v in da.items():
                         if isinstance(snake_case_   ,       snake_case_     ):
                              assign(snake_case_   ,       da[k]     )
                         else:
                              if _add_into_out:
                                   v[i : i + chunk_size] += da[k]
                              else:
                                   UpperCAmelCase_			       =  da[k]

               assign(snake_case_   ,       snake_case_     )
          elif isinstance(snake_case_   ,       snake_case_     ):
               for xa, xa in zip(snake_case_   ,       snake_case_     ):
                    if _add_into_out:
                         xa[i : i + chunk_size] += xa
                    else:
                         UpperCAmelCase_			       =  xa
          elif isinstance(snake_case_   ,       torch.Tensor     ):
               if _add_into_out:
                    out[i : i + chunk_size] += output_chunk
               else:
                    UpperCAmelCase_			       =  output_chunk
          else:
               raise ValueError("Not supported"     )

          i += chunk_size

     UpperCAmelCase_			       =  tensor_tree_map(lambda snake_case_     : t.view(orig_batch_dims + t.shape[1:]     )   ,       snake_case_     )

     return out






class 			__A :



    def __init__(self       :       Dict  ,     __a       :       int = 512  ,     ):
         UpperCAmelCase_			       =  max_chunk_size
         UpperCAmelCase_			       =  None
         UpperCAmelCase_			       =  None



    def    _lowercase		(self       :       List[Any]  ,     __a       :       Callable  ,     __a       :       tuple  ,     __a       :       int					):
         logging.info("Tuning chunk size..."					)

         if min_chunk_size >= self.max_chunk_size:
              return min_chunk_size

         UpperCAmelCase_			       =  [2**l for l in range(int(math.log(self.max_chunk_size  ,     2					)					) + 1					)]
         UpperCAmelCase_			       =  [c for c in candidates if c > min_chunk_size]
         UpperCAmelCase_			       =  [min_chunk_size] + candidates
         candidates[-1] += 4

         def test_chunk_size(__a       :       int					) -> bool:
              try:
                   with torch.no_grad():
                        fn(*__a  ,     chunk_size=__a					)
                   return True
              except RuntimeError:
                   return False

         UpperCAmelCase_			       =  0
         UpperCAmelCase_			       =  len(__a					) - 1
         while i > min_viable_chunk_size_index:
              UpperCAmelCase_			       =  test_chunk_size(candidates[i]					)
              if not viable:
                   UpperCAmelCase_			       =  (min_viable_chunk_size_index + i) // 2
              else:
                   UpperCAmelCase_			       =  i
                   UpperCAmelCase_			       =  (i + len(__a					) - 1) // 2

         return candidates[min_viable_chunk_size_index]



    def    _lowercase		(self       :       int  ,     __a       :       Iterable  ,     __a       :       Iterable					):
         UpperCAmelCase_			       =  True
         for aa, aa in zip(__a  ,     __a					):
              assert type(__a					) == type(__a					)
              if isinstance(__a  ,     (list, tuple)					):
                   consistent &= self._compare_arg_caches(__a  ,     __a					)
              elif isinstance(__a  ,     __a					):
                   UpperCAmelCase_			       =  [v for _, v in sorted(aa.items()  ,     key=lambda __a					: x[0]					)]
                   UpperCAmelCase_			       =  [v for _, v in sorted(aa.items()  ,     key=lambda __a					: x[0]					)]
                   consistent &= self._compare_arg_caches(__a  ,     __a					)
              else:
                   consistent &= aa == aa

         return consistent






    def    _lowercase		(self       :       List[str]  ,     __a       :       Callable  ,     __a       :       tuple  ,     __a       :       int  ,     ):
         UpperCAmelCase_			       =  True
         UpperCAmelCase_			       =  tree_map(lambda __a					: a.shape if isinstance(__a  ,     torch.Tensor					) else a  ,     __a  ,     __a					)
         if self.cached_arg_data is not None:
              # If args have changed shape/value, we need to re-tune
              assert len(self.cached_arg_data					) == len(__a					)
              UpperCAmelCase_			       =  self._compare_arg_caches(self.cached_arg_data  ,     __a					)
         else:
              # Otherwise, we can reuse the precomputed value
              UpperCAmelCase_			       =  False

         if not consistent:
              UpperCAmelCase_			       =  self._determine_favorable_chunk_size(
                  __a  ,     __a  ,     __a  ,     )
              UpperCAmelCase_			       =  arg_data

         assert self.cached_chunk_size is not None

         return self.cached_chunk_size






# 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	( snake_case_     ,			snake_case_     ,			snake_case_     ,			unittest.TestCase		):
	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        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					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.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					)
					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-05		,    num_attention_heads=4		,    num_hidden_layers=5		,    pad_token_id=1		,    vocab_size=1000		,    )
					lowerCamelCase_         =     CLIPTextModel(lowercase					)
					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        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase=0					)							->			int:
					if str(lowercase					).startswith("mps"					):
									lowerCamelCase_         =     torch.manual_seed(lowercase					)
					else:
									lowerCamelCase_         =     torch.Generator(device=lowercase					).manual_seed(lowercase					)

					lowerCamelCase_         =     2
					lowerCamelCase_         =     randn_tensor(
					    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)		,    generator=lowercase		,    device=torch.device(lowercase					)		,    )
					lowerCamelCase_         =     floats_tensor(control_image.shape		,    rng=random.Random(lowercase					)					).to(lowercase					)
					lowerCamelCase_         =     image.cpu().permute(0		,    2		,    3		,    1					)[0]
					lowerCamelCase_         =     Image.fromarray(np.uinta(lowercase					)					).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        SCREAMING_SNAKE_CASE_(       self					)							->			List[str]:
					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[str]:
					self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Any:
					self._test_inference_batch_single_identical(expected_max_diff=2e-3					)




class        _SCREAMING_SNAKE_CASE	( snake_case_     ,			snake_case_     ,			unittest.TestCase		):
	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        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					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(lowercase					):
									if isinstance(lowercase		,    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(lowercase					)

					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(lowercase					)

					torch.manual_seed(0					)
					lowerCamelCase_         =     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					)
					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-05		,    num_attention_heads=4		,    num_hidden_layers=5		,    pad_token_id=1		,    vocab_size=1000		,    )
					lowerCamelCase_         =     CLIPTextModel(lowercase					)
					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        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase=0					)							->			List[Any]:
					if str(lowercase					).startswith("mps"					):
									lowerCamelCase_         =     torch.manual_seed(lowercase					)
					else:
									lowerCamelCase_         =     torch.Generator(device=lowercase					).manual_seed(lowercase					)

					lowerCamelCase_         =     2

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

					lowerCamelCase_         =     floats_tensor(control_image[0].shape		,    rng=random.Random(lowercase					)					).to(lowercase					)
					lowerCamelCase_         =     image.cpu().permute(0		,    2		,    3		,    1					)[0]
					lowerCamelCase_         =     Image.fromarray(np.uinta(lowercase					)					).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        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					lowerCamelCase_         =     self.get_dummy_components()
					lowerCamelCase_         =     self.pipeline_class(**lowercase					)
					pipe.to(lowercase					)

					lowerCamelCase_         =     1_0.0
					lowerCamelCase_         =     4

					lowerCamelCase_         =     self.get_dummy_inputs(lowercase					)
					lowerCamelCase_         =     steps
					lowerCamelCase_         =     scale
					lowerCamelCase_         =     pipe(**lowercase					)[0]

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

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

					lowerCamelCase_         =     self.get_dummy_inputs(lowercase					)
					lowerCamelCase_         =     steps
					lowerCamelCase_         =     scale
					lowerCamelCase_         =     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        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					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]:
					self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)



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



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					lowerCamelCase_         =     self.get_dummy_components()
					lowerCamelCase_         =     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		):



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[int]:
					lowerCamelCase_         =     ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"					)

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

					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(
					    lowercase		,    lowercase		,    control_image=lowercase		,    generator=lowercase		,    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


