
# 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__  (       __lowercase    ,__lowercase    ,__lowercase    ,unittest.TestCase    ):
   _SCREAMING_SNAKE_CASE      :							int	       =      StableDiffusionControlNetImgaImgPipeline
   _SCREAMING_SNAKE_CASE      :							str	       =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   _SCREAMING_SNAKE_CASE      :							List[str]	       =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   _SCREAMING_SNAKE_CASE      :							Union[str, Any]	       =      IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}    )
   _SCREAMING_SNAKE_CASE      :							List[Any]	       =      IMAGE_TO_IMAGE_IMAGE_PARAMS
   def       lowerCAmelCase			(self :					str   ):
    torch.manual_seed(0   )
    __a :  Optional[Any]						=  UNetaDConditionModel(
        block_out_channels=(3_2, 6_4)					,				layers_per_block=2					,				sample_size=3_2					,				in_channels=4					,				out_channels=4					,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')					,				cross_attention_dim=3_2					,				)
    torch.manual_seed(0   )
    __a :  Optional[Any]						=  ControlNetModel(
        block_out_channels=(3_2, 6_4)					,				layers_per_block=2					,				in_channels=4					,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,				cross_attention_dim=3_2					,				conditioning_embedding_out_channels=(1_6, 3_2)					,				)
    torch.manual_seed(0   )
    __a :  Tuple						=  DDIMScheduler(
        beta_start=0.0_0085					,				beta_end=0.012					,				beta_schedule='''scaled_linear'''					,				clip_sample=snake_case_					,				set_alpha_to_one=snake_case_					,				)
    torch.manual_seed(0   )
    __a :  Tuple						=  AutoencoderKL(
        block_out_channels=[3_2, 6_4]					,				in_channels=3					,				out_channels=3					,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,				latent_channels=4					,				)
    torch.manual_seed(0   )
    __a :  Optional[Any]						=  CLIPTextConfig(
        bos_token_id=0					,				eos_token_id=2					,				hidden_size=3_2					,				intermediate_size=3_7					,				layer_norm_eps=1E-05					,				num_attention_heads=4					,				num_hidden_layers=5					,				pad_token_id=1					,				vocab_size=1_0_0_0					,				)
    __a :  Dict						=  CLIPTextModel(snake_case_   )
    __a :  Union[str, Any]						=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

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

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







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





class 		UpperCamelCase__  (       __lowercase    ,__lowercase    ,unittest.TestCase    ):
   _SCREAMING_SNAKE_CASE      :							str	       =      StableDiffusionControlNetImgaImgPipeline
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
   _SCREAMING_SNAKE_CASE      :							Tuple	       =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   _SCREAMING_SNAKE_CASE      :							int	       =      frozenset([]    )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess
   def       lowerCAmelCase			(self :					Dict   ):
    torch.manual_seed(0   )
    __a :  Optional[int]						=  UNetaDConditionModel(
        block_out_channels=(3_2, 6_4)					,				layers_per_block=2					,				sample_size=3_2					,				in_channels=4					,				out_channels=4					,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')					,				cross_attention_dim=3_2					,				)
    torch.manual_seed(0   )

    def init_weights(snake_case_ :					List[Any]   ):
     if isinstance(snake_case_					,				torch.nn.Convad   ):
      torch.nn.init.normal(m.weight   )
      m.bias.data.fill_(1.0   )

    __a :  List[str]						=  ControlNetModel(
        block_out_channels=(3_2, 6_4)					,				layers_per_block=2					,				in_channels=4					,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,				cross_attention_dim=3_2					,				conditioning_embedding_out_channels=(1_6, 3_2)					,				)
    controlneta.controlnet_down_blocks.apply(snake_case_   )

    torch.manual_seed(0   )
    __a :  Any						=  ControlNetModel(
        block_out_channels=(3_2, 6_4)					,				layers_per_block=2					,				in_channels=4					,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,				cross_attention_dim=3_2					,				conditioning_embedding_out_channels=(1_6, 3_2)					,				)
    controlneta.controlnet_down_blocks.apply(snake_case_   )

    torch.manual_seed(0   )
    __a :  Optional[Any]						=  DDIMScheduler(
        beta_start=0.0_0085					,				beta_end=0.012					,				beta_schedule='''scaled_linear'''					,				clip_sample=snake_case_					,				set_alpha_to_one=snake_case_					,				)
    torch.manual_seed(0   )
    __a :  Optional[int]						=  AutoencoderKL(
        block_out_channels=[3_2, 6_4]					,				in_channels=3					,				out_channels=3					,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,				latent_channels=4					,				)
    torch.manual_seed(0   )
    __a :  Optional[Any]						=  CLIPTextConfig(
        bos_token_id=0					,				eos_token_id=2					,				hidden_size=3_2					,				intermediate_size=3_7					,				layer_norm_eps=1E-05					,				num_attention_heads=4					,				num_hidden_layers=5					,				pad_token_id=1					,				vocab_size=1_0_0_0					,				)
    __a :  List[str]						=  CLIPTextModel(snake_case_   )
    __a :  str						=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

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

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

    __a :  Tuple						=  2

    __a :  Union[str, Any]						=  [
        randn_tensor(
            (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)					,				generator=snake_case_					,				device=torch.device(snake_case_   )					,				),
        randn_tensor(
            (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)					,				generator=snake_case_					,				device=torch.device(snake_case_   )					,				),
    ]

    __a :  Optional[Any]						=  floats_tensor(control_image[0].shape					,				rng=random.Random(snake_case_   )   ).to(snake_case_   )
    __a :  Any						=  image.cpu().permute(0					,				2					,				3					,				1   )[0]
    __a :  Union[str, Any]						=  Image.fromarray(np.uinta(snake_case_   )   ).convert('''RGB'''   ).resize((6_4, 6_4)   )
    __a :  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       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  Union[str, Any]						=  self.get_dummy_components()
    __a :  Union[str, Any]						=  self.pipeline_class(**snake_case_   )
    pipe.to(snake_case_   )

    __a :  List[Any]						=  10.0
    __a :  Optional[int]						=  4

    __a :  Any						=  self.get_dummy_inputs(snake_case_   )
    __a :  Any						=  steps
    __a :  int						=  scale
    __a :  List[Any]						=  pipe(**snake_case_   )[0]

    __a :  Any						=  self.get_dummy_inputs(snake_case_   )
    __a :  List[Any]						=  steps
    __a :  Dict						=  scale
    __a :  Any						=  pipe(**snake_case_					,				control_guidance_start=0.1					,				control_guidance_end=0.2   )[0]

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

    __a :  List[str]						=  self.get_dummy_inputs(snake_case_   )
    __a :  Any						=  steps
    __a :  Optional[Any]						=  scale
    __a :  Tuple						=  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       lowerCAmelCase			(self :					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 :					Tuple   ):
    self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3   )
   def       lowerCAmelCase			(self :					Dict   ):
    self._test_inference_batch_single_identical(expected_max_diff=2E-3   )







   def       lowerCAmelCase			(self :					int   ):
    __a :  str						=  self.get_dummy_components()
    __a :  Optional[Any]						=  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       lowerCAmelCase			(self :					Tuple   ):
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()







   def       lowerCAmelCase			(self :					List[str]   ):
    __a :  Optional[Any]						=  ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''   )

    __a :  Tuple						=  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_   )

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

    __a :  List[Any]						=  pipe(
        snake_case_					,				snake_case_					,				control_image=snake_case_					,				generator=snake_case_					,				output_type='''np'''					,				num_inference_steps=5_0					,				strength=0.6					,				)

    __a :  List[Any]						=  output.images[0]

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

    __a :  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 pathlib import Path
from typing import List

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import get_tests_dir, is_tool_test
from transformers.tools.agent_types import AGENT_TYPE_MAPPING, AgentAudio, AgentImage, AgentText


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image


lowercase__       =['text', 'image', 'audio']




def    __UpperCamelCase     (				lowerCAmelCase__							:				List[str]  ):
 __a :  Optional[int]						=  []

 for input_type in input_types:
  if input_type == "text":
   inputs.append('''Text input'''  )
  elif input_type == "image":
   inputs.append(
       Image.open(Path(get_tests_dir('''fixtures/tests_samples/COCO'''  )  ) / '''000000039769.png'''  ).resize((5_1_2, 5_1_2)  )  )
  elif input_type == "audio":
   inputs.append(torch.ones(3_0_0_0  )  )
  elif isinstance(lowerCAmelCase__  ,   lowerCAmelCase__  ):
   inputs.append(create_inputs(lowerCAmelCase__  )  )
  else:
   raise ValueError(f"Invalid type requested: {input_type}"  )

 return inputs




def    __UpperCamelCase     (				lowerCAmelCase__							:				List  ):
 __a :  List[str]						=  []

 for output in outputs:
  if isinstance(lowerCAmelCase__  ,   (str, AgentText)  ):
   output_types.append('''text'''  )
  elif isinstance(lowerCAmelCase__  ,   (Image.Image, AgentImage)  ):
   output_types.append('''image'''  )
  elif isinstance(lowerCAmelCase__  ,   (torch.Tensor, AgentAudio)  ):
   output_types.append('''audio'''  )
  else:
   raise ValueError(f"Invalid output: {output}"  )

 return output_types



@is_tool_test
class 		UpperCamelCase__  :
   def       lowerCAmelCase			(self :					Any   ):
    self.assertTrue(hasattr(self.tool					,				'''inputs'''   )   )
    self.assertTrue(hasattr(self.tool					,				'''outputs'''   )   )

    __a :  Any						=  self.tool.inputs
    for _input in inputs:
     if isinstance(_input					,				snake_case_   ):
      for __input in _input:
       self.assertTrue(__input in authorized_types   )
     else:
      self.assertTrue(_input in authorized_types   )

    __a :  Optional[int]						=  self.tool.outputs
    for _output in outputs:
     self.assertTrue(_output in authorized_types   )
   def       lowerCAmelCase			(self :					List[Any]   ):
    __a :  Union[str, Any]						=  create_inputs(self.tool.inputs   )
    __a :  List[Any]						=  self.tool(*snake_case_   )

    # There is a single output
    if len(self.tool.outputs   ) == 1:
     __a :  Tuple						=  [outputs]

    self.assertListEqual(output_types(snake_case_   )					,				self.tool.outputs   )
   def       lowerCAmelCase			(self :					List[Any]   ):
    self.assertTrue(hasattr(self.tool					,				'''description'''   )   )
    self.assertTrue(hasattr(self.tool					,				'''default_checkpoint'''   )   )
    self.assertTrue(self.tool.description.startswith('''This is a tool that'''   )   )
   def       lowerCAmelCase			(self :					Any   ):
    __a :  Any						=  create_inputs(self.tool.inputs   )
    __a :  Union[str, Any]						=  self.tool(*snake_case_   )

    if not isinstance(snake_case_					,				snake_case_   ):
     __a :  Tuple						=  [outputs]

    self.assertEqual(len(snake_case_   )					,				len(self.tool.outputs   )   )

    for output, output_type in zip(snake_case_					,				self.tool.outputs   ):
     __a :  List[Any]						=  AGENT_TYPE_MAPPING[output_type]
     self.assertTrue(isinstance(snake_case_					,				snake_case_   )   )







   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  Any						=  create_inputs(self.tool.inputs   )

    __a :  Dict						=  []

    for _input, input_type in zip(snake_case_					,				self.tool.inputs   ):
     if isinstance(snake_case_					,				snake_case_   ):
      _inputs.append([AGENT_TYPE_MAPPING[_input_type](_input   ) for _input_type in input_type]   )
     else:
      _inputs.append(AGENT_TYPE_MAPPING[input_type](_input   )   )

        # Should not raise an error
    __a :  Optional[Any]						=  self.tool(*snake_case_   )

    if not isinstance(snake_case_					,				snake_case_   ):
     __a :  Dict						=  [outputs]

    self.assertEqual(len(snake_case_   )					,				len(self.tool.outputs   )   )
