

# 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       __snake_case  (    lowerCAmelCase__     , lowerCAmelCase__     , lowerCAmelCase__     , unittest.TestCase       ):
  __lowerCAmelCase     :	Tuple						 =  StableDiffusionControlNetImgaImgPipeline
  __lowerCAmelCase     :	str						 =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
  __lowerCAmelCase     :	Any						 =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  __lowerCAmelCase     :	Any						 =  IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'control_image'}       )
  __lowerCAmelCase     :	str						 =  IMAGE_TO_IMAGE_IMAGE_PARAMS

  def 						lowerCAmelCase__     (					self):
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						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)
      SCREAMING_SNAKE_CASE_									=						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)
      SCREAMING_SNAKE_CASE_									=						DDIMScheduler(
          beta_start=0.0_0_0_8_5						,		beta_end=0.0_1_2						,		beta_schedule='scaled_linear'						,		clip_sample=_A						,		set_alpha_to_one=_A						,		)
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						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)
      SCREAMING_SNAKE_CASE_									=						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						,		)
      SCREAMING_SNAKE_CASE_									=						CLIPTextModel(_A)
      SCREAMING_SNAKE_CASE_									=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

      SCREAMING_SNAKE_CASE_									=						{
          '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						,		_A						,		_A=0):
      if str(_A).startswith('mps'):
          SCREAMING_SNAKE_CASE_									=						torch.manual_seed(_A)
      else:
          SCREAMING_SNAKE_CASE_									=						torch.Generator(device=_A).manual_seed(_A)

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


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






class       __snake_case  (    lowerCAmelCase__     , lowerCAmelCase__     , unittest.TestCase       ):
  __lowerCAmelCase     :	str						 =  StableDiffusionControlNetImgaImgPipeline
  __lowerCAmelCase     :	int						 =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
  __lowerCAmelCase     :	int						 =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  __lowerCAmelCase     :	List[str]						 =  frozenset([]       )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

  def 						lowerCAmelCase__     (					self):
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						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(_A):
          if isinstance(_A						,		torch.nn.Convad):
              torch.nn.init.normal(m.weight)
              m.bias.data.fill_(1.0)

      SCREAMING_SNAKE_CASE_									=						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(_A)

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						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(_A)

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						DDIMScheduler(
          beta_start=0.0_0_0_8_5						,		beta_end=0.0_1_2						,		beta_schedule='scaled_linear'						,		clip_sample=_A						,		set_alpha_to_one=_A						,		)
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						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)
      SCREAMING_SNAKE_CASE_									=						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						,		)
      SCREAMING_SNAKE_CASE_									=						CLIPTextModel(_A)
      SCREAMING_SNAKE_CASE_									=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

      SCREAMING_SNAKE_CASE_									=						MultiControlNetModel([controlneta, controlneta])

      SCREAMING_SNAKE_CASE_									=						{
          '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						,		_A						,		_A=0):
      if str(_A).startswith('mps'):
          SCREAMING_SNAKE_CASE_									=						torch.manual_seed(_A)
      else:
          SCREAMING_SNAKE_CASE_									=						torch.Generator(device=_A).manual_seed(_A)

      SCREAMING_SNAKE_CASE_									=						2

      SCREAMING_SNAKE_CASE_									=						[
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,		generator=_A						,		device=torch.device(_A)						,		),
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,		generator=_A						,		device=torch.device(_A)						,		),
      ]

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

      SCREAMING_SNAKE_CASE_									=						1_0.0
      SCREAMING_SNAKE_CASE_									=						4

      SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
      SCREAMING_SNAKE_CASE_									=						steps
      SCREAMING_SNAKE_CASE_									=						scale
      SCREAMING_SNAKE_CASE_									=						pipe(**_A)[0]

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

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

      SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
      SCREAMING_SNAKE_CASE_									=						steps
      SCREAMING_SNAKE_CASE_									=						scale
      SCREAMING_SNAKE_CASE_									=						pipe(**_A						,		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):
      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):
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3)

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


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
      SCREAMING_SNAKE_CASE_									=						self.pipeline_class(**_A)
      pipe.to(_A)
      pipe.set_progress_bar_config(disable=_A)
      with tempfile.TemporaryDirectory() as tmpdir:
          try:
              # save_pretrained is not implemented for Multi-ControlNet
              pipe.save_pretrained(_A)
          except NotImplementedError:
              pass






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

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


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny')

      SCREAMING_SNAKE_CASE_									=						StableDiffusionControlNetImgaImgPipeline.from_pretrained(
          'runwayml/stable-diffusion-v1-5'						,		safety_checker=_A						,		controlnet=_A)
      pipe.enable_model_cpu_offload()
      pipe.set_progress_bar_config(disable=_A)

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

      SCREAMING_SNAKE_CASE_									=						pipe(
          _A						,		_A						,		control_image=_A						,		generator=_A						,		output_type='np'						,		num_inference_steps=50						,		strength=0.6						,		)

      SCREAMING_SNAKE_CASE_									=						output.images[0]

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

      SCREAMING_SNAKE_CASE_									=						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





import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM



@require_torch
class       __snake_case  :

  def __init__(					self						,		_A						,		_A=99						,		_A=13						,		_A=16						,		_A=7						,		_A=True						,		_A=True						,		_A=True						,		_A=False						,		_A=True						,		_A=2						,		_A=32						,		_A=4						,		_A=4						,		_A=30						,		_A=0						,		_A=1						,		_A=2						,		_A=None						,		):
      SCREAMING_SNAKE_CASE_									=						parent
      SCREAMING_SNAKE_CASE_									=						batch_size
      SCREAMING_SNAKE_CASE_									=						decoder_seq_length
      # For common tests
      SCREAMING_SNAKE_CASE_									=						self.decoder_seq_length
      SCREAMING_SNAKE_CASE_									=						is_training
      SCREAMING_SNAKE_CASE_									=						use_attention_mask
      SCREAMING_SNAKE_CASE_									=						use_labels

      SCREAMING_SNAKE_CASE_									=						vocab_size
      SCREAMING_SNAKE_CASE_									=						d_model
      SCREAMING_SNAKE_CASE_									=						d_model
      SCREAMING_SNAKE_CASE_									=						decoder_layers
      SCREAMING_SNAKE_CASE_									=						decoder_layers
      SCREAMING_SNAKE_CASE_									=						decoder_ffn_dim
      SCREAMING_SNAKE_CASE_									=						decoder_attention_heads
      SCREAMING_SNAKE_CASE_									=						decoder_attention_heads
      SCREAMING_SNAKE_CASE_									=						eos_token_id
      SCREAMING_SNAKE_CASE_									=						bos_token_id
      SCREAMING_SNAKE_CASE_									=						pad_token_id
      SCREAMING_SNAKE_CASE_									=						decoder_start_token_id
      SCREAMING_SNAKE_CASE_									=						use_cache
      SCREAMING_SNAKE_CASE_									=						max_position_embeddings

      SCREAMING_SNAKE_CASE_									=						None
      SCREAMING_SNAKE_CASE_									=						decoder_seq_length
      SCREAMING_SNAKE_CASE_									=						2
      SCREAMING_SNAKE_CASE_									=						1

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_attention_mask:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		vocab_size=2)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_labels:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						TrOCRConfig(
          vocab_size=self.vocab_size						,		d_model=self.d_model						,		decoder_layers=self.decoder_layers						,		decoder_ffn_dim=self.decoder_ffn_dim						,		decoder_attention_heads=self.decoder_attention_heads						,		eos_token_id=self.eos_token_id						,		bos_token_id=self.bos_token_id						,		use_cache=self.use_cache						,		pad_token_id=self.pad_token_id						,		decoder_start_token_id=self.decoder_start_token_id						,		max_position_embeddings=self.max_position_embeddings						,		)

      return (config, input_ids, attention_mask, lm_labels)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A						,		):
      SCREAMING_SNAKE_CASE_									=						True
      SCREAMING_SNAKE_CASE_									=						TrOCRDecoder(config=_A).to(_A).eval()
      SCREAMING_SNAKE_CASE_									=						input_ids[:2]

      input_ids[input_ids == 0] += 1
      # first forward pass
      SCREAMING_SNAKE_CASE_									=						model(_A						,		use_cache=_A)
      SCREAMING_SNAKE_CASE_									=						model(_A)
      SCREAMING_SNAKE_CASE_									=						model(_A						,		use_cache=_A)

      self.parent.assertTrue(len(_A) == len(_A))
      self.parent.assertTrue(len(_A) == len(_A) + 1)

      SCREAMING_SNAKE_CASE_									=						outputs['past_key_values']

      # create hypothetical next token and extent to next_input_ids
      SCREAMING_SNAKE_CASE_									=						ids_tensor((2, 1)						,		config.vocab_size - 1) + 1

      # append to next input_ids and
      SCREAMING_SNAKE_CASE_									=						torch.cat([input_ids, next_tokens]						,		dim=-1)

      SCREAMING_SNAKE_CASE_									=						model(_A)['last_hidden_state']
      SCREAMING_SNAKE_CASE_									=						model(_A						,		past_key_values=_A)['last_hidden_state']

      # select random slice
      SCREAMING_SNAKE_CASE_									=						ids_tensor((1,)						,		output_from_past.shape[-1]).item()
      SCREAMING_SNAKE_CASE_									=						output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
      SCREAMING_SNAKE_CASE_									=						output_from_past[:, 0, random_slice_idx].detach()

      # test that outputs are equal for slice
      assert torch.allclose(_A						,		_A						,		atol=1E-3)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						config_and_inputs

      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids, 'attention_mask': attention_mask}
      return config, inputs_dict






@require_torch
class       __snake_case  (    lowerCAmelCase__     , lowerCAmelCase__     , lowerCAmelCase__     , unittest.TestCase       ):
  __lowerCAmelCase     :	Tuple						 =  (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
  __lowerCAmelCase     :	Union[str, Any]						 =  (TrOCRForCausalLM,) if is_torch_available() else ()
  __lowerCAmelCase     :	str						 =  {'text-generation': TrOCRForCausalLM} if is_torch_available() else {}
  __lowerCAmelCase     :	Any						 =  True
  __lowerCAmelCase     :	str						 =  False

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						TrOCRStandaloneDecoderModelTester(self						,		is_training=_A)
      SCREAMING_SNAKE_CASE_									=						ConfigTester(self						,		config_class=_A)

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      self.config_tester.run_common_tests()

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past(*_A)

  def 						lowerCAmelCase__     (					self):
      return


  @unittest.skip('The model doesn\'t support left padding')  # and it's not used enough to be worth fixing :)
  def 						lowerCAmelCase__     (					self):
      pass


