





'''simple docstring'''


# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()
class       a_   (				_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			StableDiffusionControlNetImgaImgPipeline
     __A      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
     __A      =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     __A      =			IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}							)
     __A      =			IMAGE_TO_IMAGE_IMAGE_PARAMS





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        torch.manual_seed(0						)
        lowercase_			:Any    			= 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_			:str    			= 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_			:List[str]    			= DDIMScheduler(
            beta_start=0.0_00_85	,      beta_end=0.0_12	,      beta_schedule="scaled_linear"	,      clip_sample=__lowerCAmelCase	,      set_alpha_to_one=__lowerCAmelCase	,      )
        torch.manual_seed(0						)
        lowercase_			:Optional[int]    			= 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_			:List[str]    			= CLIPTextConfig(
            bos_token_id=0	,      eos_token_id=2	,      hidden_size=32	,      intermediate_size=37	,      layer_norm_eps=1e-0_5	,      num_attention_heads=4	,      num_hidden_layers=5	,      pad_token_id=1	,      vocab_size=1_000	,      )
        lowercase_			:Optional[Any]    			= CLIPTextModel(__lowerCAmelCase						)
        lowercase_			:Tuple    			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

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





     def 							lowercase__      ( self		:					List[str]	,      lowercase		:					Optional[int]	,      lowercase		:					int=0						):





        """simple docstring"""




        if str(__lowerCAmelCase						).startswith("mps"						):
           lowercase_			:str    			= torch.manual_seed(__lowerCAmelCase						)
        else:
           lowercase_			:int    			= torch.Generator(device=__lowerCAmelCase						).manual_seed(__lowerCAmelCase						)

        lowercase_			:List[str]    			= 2
        lowercase_			:Any    			= randn_tensor(
            (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)	,      generator=__lowerCAmelCase	,      device=torch.device(__lowerCAmelCase						)	,      )
        lowercase_			:Any    			= floats_tensor(control_image.shape	,      rng=random.Random(__lowerCAmelCase						)						).to(__lowerCAmelCase						)
        lowercase_			:List[str]    			= image.cpu().permute(0	,      2	,      3	,      1						)[0]
        lowercase_			:Dict    			= Image.fromarray(np.uinta(__lowerCAmelCase						)						).convert("RGB"						).resize((64, 64)						)
        lowercase_			:List[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 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        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 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3						)







     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        self._test_inference_batch_single_identical(expected_max_diff=2e-3						)







class       a_   (				_lowerCAmelCase	,						_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			StableDiffusionControlNetImgaImgPipeline
     __A      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
     __A      =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     __A      =			frozenset([]							)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        torch.manual_seed(0						)
        lowercase_			:List[Any]    			= 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		:					List[str]						):
           if isinstance(__lowerCAmelCase	,      torch.nn.Convad						):
              torch.nn.init.normal(m.weight						)
              m.bias.data.fill_(1.0						)

        lowercase_			:int    			= 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(__lowerCAmelCase						)

        torch.manual_seed(0						)
        lowercase_			:Optional[int]    			= 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(__lowerCAmelCase						)

        torch.manual_seed(0						)
        lowercase_			:Any    			= DDIMScheduler(
            beta_start=0.0_00_85	,      beta_end=0.0_12	,      beta_schedule="scaled_linear"	,      clip_sample=__lowerCAmelCase	,      set_alpha_to_one=__lowerCAmelCase	,      )
        torch.manual_seed(0						)
        lowercase_			:List[Any]    			= 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_			:List[Any]    			= CLIPTextConfig(
            bos_token_id=0	,      eos_token_id=2	,      hidden_size=32	,      intermediate_size=37	,      layer_norm_eps=1e-0_5	,      num_attention_heads=4	,      num_hidden_layers=5	,      pad_token_id=1	,      vocab_size=1_000	,      )
        lowercase_			:Tuple    			= CLIPTextModel(__lowerCAmelCase						)
        lowercase_			:List[str]    			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

        lowercase_			:Tuple    			= MultiControlNetModel([controlneta, controlneta]						)

        lowercase_			:List[Any]    			= {
            "unet": unet,
            "controlnet": controlnet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components





     def 							lowercase__      ( self		:					str	,      lowercase		:					int	,      lowercase		:					List[Any]=0						):





        """simple docstring"""




        if str(__lowerCAmelCase						).startswith("mps"						):
           lowercase_			:Union[str, Any]    			= torch.manual_seed(__lowerCAmelCase						)
        else:
           lowercase_			:Union[str, Any]    			= torch.Generator(device=__lowerCAmelCase						).manual_seed(__lowerCAmelCase						)

        lowercase_			:List[Any]    			= 2

        lowercase_			:str    			= [
            randn_tensor(
                (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)	,      generator=__lowerCAmelCase	,      device=torch.device(__lowerCAmelCase						)	,      ),
            randn_tensor(
                (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)	,      generator=__lowerCAmelCase	,      device=torch.device(__lowerCAmelCase						)	,      ),
        ]

        lowercase_			:int    			= floats_tensor(control_image[0].shape	,      rng=random.Random(__lowerCAmelCase						)						).to(__lowerCAmelCase						)
        lowercase_			:Tuple    			= image.cpu().permute(0	,      2	,      3	,      1						)[0]
        lowercase_			:List[str]    			= Image.fromarray(np.uinta(__lowerCAmelCase						)						).convert("RGB"						).resize((64, 64)						)
        lowercase_			:int    			= {
            "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 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= self.get_dummy_components()
        lowercase_			:int    			= self.pipeline_class(**__lowerCAmelCase						)
        pipe.to(__lowerCAmelCase						)

        lowercase_			:Optional[int]    			= 10.0
        lowercase_			:List[str]    			= 4

        lowercase_			:Dict    			= self.get_dummy_inputs(__lowerCAmelCase						)
        lowercase_			:Optional[Any]    			= steps
        lowercase_			:List[Any]    			= scale
        lowercase_			:List[Any]    			= pipe(**__lowerCAmelCase						)[0]

        lowercase_			:Optional[Any]    			= self.get_dummy_inputs(__lowerCAmelCase						)
        lowercase_			:str    			= steps
        lowercase_			:str    			= scale
        lowercase_			:Optional[int]    			= pipe(**__lowerCAmelCase	,      control_guidance_start=0.1	,      control_guidance_end=0.2						)[0]

        lowercase_			:Tuple    			= self.get_dummy_inputs(__lowerCAmelCase						)
        lowercase_			:Tuple    			= steps
        lowercase_			:Optional[Any]    			= scale
        lowercase_			:Optional[Any]    			= pipe(**__lowerCAmelCase	,      control_guidance_start=[0.1, 0.3]	,      control_guidance_end=[0.2, 0.7]						)[0]

        lowercase_			:List[Any]    			= self.get_dummy_inputs(__lowerCAmelCase						)
        lowercase_			:Tuple    			= steps
        lowercase_			:List[Any]    			= scale
        lowercase_			:Dict    			= pipe(**__lowerCAmelCase	,      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 							lowercase__      ( self		:					int						):





        """simple docstring"""




        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 							lowercase__      ( self		:					str						):





        """simple docstring"""




        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3						)





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        self._test_inference_batch_single_identical(expected_max_diff=2e-3						)







     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			:int    			= self.get_dummy_components()
        lowercase_			:Tuple    			= self.pipeline_class(**__lowerCAmelCase						)
        pipe.to(__lowerCAmelCase						)
        pipe.set_progress_bar_config(disable=__lowerCAmelCase						)
        with tempfile.TemporaryDirectory() as tmpdir:
           try:
              # save_pretrained is not implemented for Multi-ControlNet
              pipe.save_pretrained(__lowerCAmelCase						)
           except NotImplementedError:
              pass







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





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			:Tuple    			= ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"						)

        lowercase_			:int    			= StableDiffusionControlNetImgaImgPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5"	,      safety_checker=__lowerCAmelCase	,      controlnet=__lowerCAmelCase						)
        pipe.enable_model_cpu_offload()
        pipe.set_progress_bar_config(disable=__lowerCAmelCase						)

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

        lowercase_			:List[str]    			= pipe(
            __lowerCAmelCase	,      __lowerCAmelCase	,      control_image=__lowerCAmelCase	,      generator=__lowerCAmelCase	,      output_type="np"	,      num_inference_steps=50	,      strength=0.6	,      )

        lowercase_			:Optional[Any]    			= output.images[0]

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

        lowercase_			:List[Any]    			= 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 os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
			from .tokenization_xlnet import XLNetTokenizer
else:
			_a 						=      None


_a 						=      logging.get_logger(__name__)

_a 						=      {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

_a 						=      {
    "vocab_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/tokenizer.json",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/tokenizer.json",
    },
}

_a 						=      {
    "xlnet-base-cased": None,
    "xlnet-large-cased": None,
}

_a 						=      "▁"

# Segments (not really needed)
_a 						=      0
_a 						=      1
_a 						=      2
_a 						=      3
_a 						=      4






class 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						VOCAB_FILES_NAMES
			lowerCAmelCase_							     =						PRETRAINED_VOCAB_FILES_MAP
			lowerCAmelCase_							     =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCAmelCase_							     =						"""left"""
			lowerCAmelCase_							     =						XLNetTokenizer





			def __init__(    self       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=True       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase="<s>"       ,  __lowerCAmelCase="</s>"       ,  __lowerCAmelCase="<unk>"       ,  __lowerCAmelCase="<sep>"       ,  __lowerCAmelCase="<pad>"       ,  __lowerCAmelCase="<cls>"       ,  __lowerCAmelCase="<mask>"       ,  __lowerCAmelCase=["<eop>", "<eod>"]       ,  **__lowerCAmelCase       ,  ):



				'''simple docstring'''







				lowerCamelCase__								=		AddedToken(__lowerCAmelCase       ,  lstrip=__lowerCAmelCase       ,  rstrip=__lowerCAmelCase							) if isinstance(__lowerCAmelCase       ,  __lowerCAmelCase							) else mask_token

				super().__init__(
				    vocab_file=__lowerCAmelCase       ,  tokenizer_file=__lowerCAmelCase       ,  do_lower_case=__lowerCAmelCase       ,  remove_space=__lowerCAmelCase       ,  keep_accents=__lowerCAmelCase       ,  bos_token=__lowerCAmelCase       ,  eos_token=__lowerCAmelCase       ,  unk_token=__lowerCAmelCase       ,  sep_token=__lowerCAmelCase       ,  pad_token=__lowerCAmelCase       ,  cls_token=__lowerCAmelCase       ,  mask_token=__lowerCAmelCase       ,  additional_special_tokens=__lowerCAmelCase       ,  **__lowerCAmelCase       ,  )

				lowerCamelCase__								=		3
				lowerCamelCase__								=		do_lower_case
				lowerCamelCase__								=		remove_space
				lowerCamelCase__								=		keep_accents
				lowerCamelCase__								=		vocab_file
				lowerCamelCase__								=		False if not self.vocab_file else True





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase = None							):



				'''simple docstring'''







				lowerCamelCase__								=		[self.sep_token_id]
				lowerCamelCase__								=		[self.cls_token_id]
				if token_ids_a is None:
					return token_ids_a + sep + cls
				return token_ids_a + sep + token_ids_a + sep + cls





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase = None							):



				'''simple docstring'''







				lowerCamelCase__								=		[self.sep_token_id]
				lowerCamelCase__								=		[2]

				if token_ids_a is None:
					return len(token_ids_a + sep							) * [0] + cls_segment_id
				return len(token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1] + cls_segment_id





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase = None							):



				'''simple docstring'''







				if not self.can_save_slow_tokenizer:
					raise ValueError(
					    '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
					    '''tokenizer.'''							)

				if not os.path.isdir(__lowerCAmelCase							):
					logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
					return
				lowerCamelCase__								=		os.path.join(
				    __lowerCAmelCase       ,  (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']							)

				if os.path.abspath(self.vocab_file							) != os.path.abspath(__lowerCAmelCase							):
					copyfile(self.vocab_file       ,  __lowerCAmelCase							)

				return (out_vocab_file,)



