
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import difflib
from collections.abc import Mapping
from typing import Any, Callable, Dict



class       a_   ( _a						):
   pass







def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  str		,      SCREAMING_SNAKE_CASE_  :  Dict[str, Callable]		,      SCREAMING_SNAKE_CASE_  :  Optional[int]=None		,      SCREAMING_SNAKE_CASE_  :  Union[str, Any]=False					)     ->    Tuple:
      try:
            _lowercase									=				ast.parse(SCREAMING_SNAKE_CASE_					)
      except SyntaxError as e:
            print("""The code generated by the agent is not valid.\n"""		,      SCREAMING_SNAKE_CASE_					)
            return
      if state is None:
            _lowercase									=				{}
      _lowercase									=				None
      for idx, node in enumerate(expression.body					):
            try:
                  _lowercase									=				evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
            except InterpretorError as e:
                  _lowercase									=				f"""Evaluation of the code stopped at line {idx} before the end because of the following error"""
                  if chat_mode:
                        msg += (
                            f""". Copy paste the following error message and send it back to the agent:\nI get an error: '{e}'"""
                        )
                  else:
                        msg += f""":\n{e}"""
                  print(SCREAMING_SNAKE_CASE_					)
                  break
            if line_result is not None:
                  _lowercase									=				line_result

      return result




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  ast.AST		,      SCREAMING_SNAKE_CASE_  :  Dict[str, Any]		,      SCREAMING_SNAKE_CASE_  :  Dict[str, Callable]					)     ->    List[str]:
      if isinstance(SCREAMING_SNAKE_CASE_		,      ast.Assign					):
            # Assignement -> we evaluate the assignement which should update the state
            # We return the variable assigned as it may be used to determine the final result.
            return evaluate_assign(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Call					):
            # Function call -> we return the value of the function call
            return evaluate_call(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Constant					):
            # Constant -> just return the value
            return expression.value
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Dict					):
            # Dict -> evaluate all keys and values
            _lowercase									=				[evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					) for k in expression.keys]
            _lowercase									=				[evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					) for v in expression.values]
            return dict(zip(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Expr					):
            # Expression -> evaluate the content
            return evaluate_ast(expression.value		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.For					):
            # For loop -> execute the loop
            return evaluate_for(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.FormattedValue					):
            # Formatted value (part of f-string) -> evaluate the content and return
            return evaluate_ast(expression.value		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.If					):
            # If -> execute the right branch
            return evaluate_if(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif hasattr(SCREAMING_SNAKE_CASE_		,      """Index"""					) and isinstance(SCREAMING_SNAKE_CASE_		,      ast.Index					):
            return evaluate_ast(expression.value		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.JoinedStr					):
            return "".join([str(evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)					) for v in expression.values]					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.List					):
            # List -> evaluate all elements
            return [evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					) for elt in expression.elts]
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Name					):
            # Name -> pick up the value in the state
            return evaluate_name(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Subscript					):
            # Subscript -> return the value of the indexing
            return evaluate_subscript(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      else:
            # For now we refuse anything else. Let's add things as we need them.
            raise InterpretorError(f"""{expression.__class__.__name__} is not supported."""					)




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  List[str]		,      SCREAMING_SNAKE_CASE_  :  str		,      SCREAMING_SNAKE_CASE_  :  Optional[Any]					)     ->    Optional[int]:
      _lowercase									=				assign.targets
      _lowercase									=				evaluate_ast(assign.value		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)

      if len(SCREAMING_SNAKE_CASE_					) == 1:
            _lowercase									=				result
      else:
            if len(SCREAMING_SNAKE_CASE_					) != len(SCREAMING_SNAKE_CASE_					):
                  raise InterpretorError(f"""Expected {len(SCREAMING_SNAKE_CASE_					)} values but got {len(SCREAMING_SNAKE_CASE_					)}."""					)
            for var_name, r in zip(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					):
                  _lowercase									=				r
      return result




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Dict		,      SCREAMING_SNAKE_CASE_  :  str		,      SCREAMING_SNAKE_CASE_  :  Optional[Any]					)     ->    Optional[int]:
      if not isinstance(call.func		,      ast.Name					):
            raise InterpretorError(
                f"""It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func} of """
                f"""type {type(call.func					)}."""					)
      _lowercase									=				call.func.id
      if func_name not in tools:
            raise InterpretorError(
                f"""It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func.id})."""					)

      _lowercase									=				tools[func_name]
      # Todo deal with args
      _lowercase									=				[evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					) for arg in call.args]
      _lowercase									=				{keyword.arg: evaluate_ast(keyword.value		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					) for keyword in call.keywords}
      return func(*SCREAMING_SNAKE_CASE_		,      **SCREAMING_SNAKE_CASE_					)




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Any		,      SCREAMING_SNAKE_CASE_  :  Union[str, Any]		,      SCREAMING_SNAKE_CASE_  :  str					)     ->    Optional[int]:
      _lowercase									=				evaluate_ast(subscript.slice		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      _lowercase									=				evaluate_ast(subscript.value		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      if isinstance(SCREAMING_SNAKE_CASE_		,      (list, tuple)					):
            return value[int(SCREAMING_SNAKE_CASE_					)]
      if index in value:
            return value[index]
      if isinstance(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					) and isinstance(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					):
            _lowercase									=				difflib.get_close_matches(SCREAMING_SNAKE_CASE_		,      list(value.keys()					)					)
            if len(SCREAMING_SNAKE_CASE_					) > 0:
                  return value[close_matches[0]]

      raise InterpretorError(f"""Could not index {value} with '{index}'."""					)




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Dict		,      SCREAMING_SNAKE_CASE_  :  Optional[int]		,      SCREAMING_SNAKE_CASE_  :  List[str]					)     ->    Union[str, Any]:
      if name.id in state:
            return state[name.id]
      _lowercase									=				difflib.get_close_matches(name.id		,      list(state.keys()					)					)
      if len(SCREAMING_SNAKE_CASE_					) > 0:
            return state[close_matches[0]]
      raise InterpretorError(f"""The variable `{name.id}` is not defined."""					)




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  int		,      SCREAMING_SNAKE_CASE_  :  Dict		,      SCREAMING_SNAKE_CASE_  :  int					)     ->    Union[str, Any]:
      if len(condition.ops					) > 1:
            raise InterpretorError("""Cannot evaluate conditions with multiple operators"""					)

      _lowercase									=				evaluate_ast(condition.left		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      _lowercase									=				condition.ops[0]
      _lowercase									=				evaluate_ast(condition.comparators[0]		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)

      if isinstance(SCREAMING_SNAKE_CASE_		,      ast.Eq					):
            return left == right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.NotEq					):
            return left != right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Lt					):
            return left < right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.LtE					):
            return left <= right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Gt					):
            return left > right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.GtE					):
            return left >= right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.Is					):
            return left is right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.IsNot					):
            return left is not right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.In					):
            return left in right
      elif isinstance(SCREAMING_SNAKE_CASE_		,      ast.NotIn					):
            return left not in right
      else:
            raise InterpretorError(f"""Operator not supported: {comparator}"""					)




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Union[str, Any]		,      SCREAMING_SNAKE_CASE_  :  Dict		,      SCREAMING_SNAKE_CASE_  :  Union[str, Any]					)     ->    Optional[Any]:
      _lowercase									=				None
      if evaluate_condition(if_statement.test		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					):
            for line in if_statement.body:
                  _lowercase									=				evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
                  if line_result is not None:
                        _lowercase									=				line_result
      else:
            for line in if_statement.orelse:
                  _lowercase									=				evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
                  if line_result is not None:
                        _lowercase									=				line_result
      return result




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Any		,      SCREAMING_SNAKE_CASE_  :  Optional[Any]		,      SCREAMING_SNAKE_CASE_  :  str					)     ->    Any:
      _lowercase									=				None
      _lowercase									=				evaluate_ast(for_loop.iter		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
      for counter in iterator:
            _lowercase									=				counter
            for expression in for_loop.body:
                  _lowercase									=				evaluate_ast(SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_					)
                  if line_result is not None:
                        _lowercase									=				line_result
      return result

# This model implementation is heavily based on:

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,
    StableDiffusionControlNetInpaintPipeline,
    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 (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
    TEXT_TO_IMAGE_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()



class       a_   ( _a				, _a				, _a				, unittest.TestCase						):
			a  : Optional[int]				      =       StableDiffusionControlNetInpaintPipeline
			a  : Union[str, Any]				      =       TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
			a  : List[Any]				      =       TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
			a  : int				      =       frozenset({'''control_image'''}						)  # skip `image` and `mask` for now, only test for control_image
			a  : int				      =       TEXT_TO_IMAGE_IMAGE_PARAMS






			def       UpperCamelCase_							(			self     ):
									torch.manual_seed(0     )
									_lowercase									=				UNetaDConditionModel(
									    block_out_channels=(32, 64)							, layers_per_block=2							, sample_size=32							, in_channels=9							, out_channels=4							, down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")							, up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")							, cross_attention_dim=32							, )
									torch.manual_seed(0     )
									_lowercase									=				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     )
									_lowercase									=				DDIMScheduler(
									    beta_start=0.0_00_85							, beta_end=0.0_12							, beta_schedule="""scaled_linear"""							, clip_sample=__UpperCamelCase							, set_alpha_to_one=__UpperCamelCase							, )
									torch.manual_seed(0     )
									_lowercase									=				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     )
									_lowercase									=				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							, )
									_lowercase									=				CLIPTextModel(__UpperCamelCase     )
									_lowercase									=				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

									_lowercase									=				{
									    """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     ):
									if str(__UpperCamelCase     ).startswith("""mps"""     ):
															_lowercase									=				torch.manual_seed(__UpperCamelCase     )
									else:
															_lowercase									=				torch.Generator(device=__UpperCamelCase     ).manual_seed(__UpperCamelCase     )

									_lowercase									=				2
									_lowercase									=				randn_tensor(
									    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)							, generator=__UpperCamelCase							, device=torch.device(__UpperCamelCase     )							, )
									_lowercase									=				floats_tensor((1, 3, 32, 32)							, rng=random.Random(__UpperCamelCase     )     ).to(__UpperCamelCase     )
									_lowercase									=				init_image.cpu().permute(0							, 2							, 3							, 1     )[0]

									_lowercase									=				Image.fromarray(np.uinta(__UpperCamelCase     )     ).convert("""RGB"""     ).resize((64, 64)     )
									_lowercase									=				Image.fromarray(np.uinta(init_image + 4     )     ).convert("""RGB"""     ).resize((64, 64)     )

									_lowercase									=				{
									    """prompt""": """A painting of a squirrel eating a burger""",
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """guidance_scale""": 6.0,
									    """output_type""": """numpy""",
									    """image""": image,
									    """mask_image""": mask_image,
									    """control_image""": control_image,
									}

									return inputs






			def       UpperCamelCase_							(			self     ):
									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     ):
									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3     )






			def       UpperCamelCase_							(			self     ):
									self._test_inference_batch_single_identical(expected_max_diff=2E-3     )







class       a_   ( _a						):
			a  : List[str]				      =       StableDiffusionControlNetInpaintPipeline
			a  : Any				      =       TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
			a  : int				      =       TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
			a  : Union[str, Any]				      =       frozenset([]						)






			def       UpperCamelCase_							(			self     ):
									torch.manual_seed(0     )
									_lowercase									=				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     )
									_lowercase									=				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     )
									_lowercase									=				DDIMScheduler(
									    beta_start=0.0_00_85							, beta_end=0.0_12							, beta_schedule="""scaled_linear"""							, clip_sample=__UpperCamelCase							, set_alpha_to_one=__UpperCamelCase							, )
									torch.manual_seed(0     )
									_lowercase									=				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     )
									_lowercase									=				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							, )
									_lowercase									=				CLIPTextModel(__UpperCamelCase     )
									_lowercase									=				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

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







class       a_   ( _a				, _a				, unittest.TestCase						):
			a  : str				      =       StableDiffusionControlNetInpaintPipeline
			a  : Any				      =       TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
			a  : Optional[Any]				      =       TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS






			def       UpperCamelCase_							(			self     ):
									torch.manual_seed(0     )
									_lowercase									=				UNetaDConditionModel(
									    block_out_channels=(32, 64)							, layers_per_block=2							, sample_size=32							, in_channels=9							, 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     )

									_lowercase									=				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     )
									_lowercase									=				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     )
									_lowercase									=				DDIMScheduler(
									    beta_start=0.0_00_85							, beta_end=0.0_12							, beta_schedule="""scaled_linear"""							, clip_sample=__UpperCamelCase							, set_alpha_to_one=__UpperCamelCase							, )
									torch.manual_seed(0     )
									_lowercase									=				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     )
									_lowercase									=				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							, )
									_lowercase									=				CLIPTextModel(__UpperCamelCase     )
									_lowercase									=				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

									_lowercase									=				MultiControlNetModel([controlneta, controlneta]     )

									_lowercase									=				{
									    """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     ):
									if str(__UpperCamelCase     ).startswith("""mps"""     ):
															_lowercase									=				torch.manual_seed(__UpperCamelCase     )
									else:
															_lowercase									=				torch.Generator(device=__UpperCamelCase     ).manual_seed(__UpperCamelCase     )

									_lowercase									=				2

									_lowercase									=				[
									    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     )							, ),
									]
									_lowercase									=				floats_tensor((1, 3, 32, 32)							, rng=random.Random(__UpperCamelCase     )     ).to(__UpperCamelCase     )
									_lowercase									=				init_image.cpu().permute(0							, 2							, 3							, 1     )[0]

									_lowercase									=				Image.fromarray(np.uinta(__UpperCamelCase     )     ).convert("""RGB"""     ).resize((64, 64)     )
									_lowercase									=				Image.fromarray(np.uinta(init_image + 4     )     ).convert("""RGB"""     ).resize((64, 64)     )

									_lowercase									=				{
									    """prompt""": """A painting of a squirrel eating a burger""",
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """guidance_scale""": 6.0,
									    """output_type""": """numpy""",
									    """image""": image,
									    """mask_image""": mask_image,
									    """control_image""": control_image,
									}

									return inputs






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.get_dummy_components()
									_lowercase									=				self.pipeline_class(**__UpperCamelCase     )
									pipe.to(__UpperCamelCase     )

									_lowercase									=				10.0
									_lowercase									=				4

									_lowercase									=				self.get_dummy_inputs(__UpperCamelCase     )
									_lowercase									=				steps
									_lowercase									=				scale
									_lowercase									=				pipe(**__UpperCamelCase     )[0]

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

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

									_lowercase									=				self.get_dummy_inputs(__UpperCamelCase     )
									_lowercase									=				steps
									_lowercase									=				scale
									_lowercase									=				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     ):
									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     ):
									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3     )






			def       UpperCamelCase_							(			self     ):
									self._test_inference_batch_single_identical(expected_max_diff=2E-3     )






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






			def       UpperCamelCase_							(			self     ):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""     )

									_lowercase									=				StableDiffusionControlNetInpaintPipeline.from_pretrained(
									    """runwayml/stable-diffusion-inpainting"""							, safety_checker=__UpperCamelCase							, controlnet=__UpperCamelCase     )
									pipe.enable_model_cpu_offload()
									pipe.set_progress_bar_config(disable=__UpperCamelCase     )

									_lowercase									=				torch.Generator(device="""cpu"""     ).manual_seed(0     )
									_lowercase									=				load_image(
									    """https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"""     ).resize((512, 512)     )

									_lowercase									=				load_image(
									    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
									    """/stable_diffusion_inpaint/input_bench_mask.png"""     ).resize((512, 512)     )

									_lowercase									=				"""pitch black hole"""

									_lowercase									=				load_image(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"""     ).resize((512, 512)     )

									_lowercase									=				pipe(
									    __UpperCamelCase							, image=__UpperCamelCase							, mask_image=__UpperCamelCase							, control_image=__UpperCamelCase							, generator=__UpperCamelCase							, output_type="""np"""							, num_inference_steps=3							, )

									_lowercase									=				output.images[0]

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

									_lowercase									=				load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/inpaint.npy"""     )

									assert np.abs(expected_image - image     ).max() < 9E-2






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				ControlNetModel.from_pretrained("""lllyasviel/control_v11p_sd15_inpaint"""     )

									_lowercase									=				StableDiffusionControlNetInpaintPipeline.from_pretrained(
									    """runwayml/stable-diffusion-v1-5"""							, safety_checker=__UpperCamelCase							, controlnet=__UpperCamelCase     )
									_lowercase									=				DDIMScheduler.from_config(pipe.scheduler.config     )
									pipe.enable_model_cpu_offload()
									pipe.set_progress_bar_config(disable=__UpperCamelCase     )

									_lowercase									=				torch.Generator(device="""cpu"""     ).manual_seed(33     )

									_lowercase									=				load_image(
									    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"""     )
									_lowercase									=				init_image.resize((512, 512)     )

									_lowercase									=				load_image(
									    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"""     )
									_lowercase									=				mask_image.resize((512, 512)     )

									_lowercase									=				"""a handsome man with ray-ban sunglasses"""

									def make_inpaint_condition(__UpperCamelCase							, __UpperCamelCase     ):
															_lowercase									=				np.array(image.convert("""RGB"""     )     ).astype(np.floataa     ) / 2_55.0
															_lowercase									=				np.array(image_mask.convert("""L"""     )     ).astype(np.floataa     ) / 2_55.0

															assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
															_lowercase									=				-1.0  # set as masked pixel
															_lowercase									=				np.expand_dims(__UpperCamelCase							, 0     ).transpose(0							, 3							, 1							, 2     )
															_lowercase									=				torch.from_numpy(__UpperCamelCase     )
															return image

									_lowercase									=				make_inpaint_condition(__UpperCamelCase							, __UpperCamelCase     )

									_lowercase									=				pipe(
									    __UpperCamelCase							, image=__UpperCamelCase							, mask_image=__UpperCamelCase							, control_image=__UpperCamelCase							, guidance_scale=9.0							, eta=1.0							, generator=__UpperCamelCase							, num_inference_steps=20							, output_type="""np"""							, )
									_lowercase									=				output.images[0]

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

									_lowercase									=				load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/boy_ray_ban.npy"""     )

									assert np.abs(expected_image - image     ).max() < 9E-2