




import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

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


enable_full_determinism()





class     SCREAMING_SNAKE_CASE__		(_a    ,					_a    ,					_a    ,					unittest.TestCase			):
     lowercase_ :      List[str]				= StableUnCLIPImgaImgPipeline
     lowercase_ :      Optional[Any]				= TEXT_GUIDED_IMAGE_VARIATION_PARAMS
     lowercase_ :      Any				= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     lowercase_ :      Any				= frozenset(
         []			)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     lowercase_ :      str				= frozenset([]			)
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      32
       lowerCAmelCase__										=      embedder_hidden_size

       # image encoding components

       lowerCAmelCase__										=      CLIPImageProcessor(crop_size=32       ,						size=32						)

       torch.manual_seed(0						)
       lowerCAmelCase__										=      CLIPVisionModelWithProjection(
           CLIPVisionConfig(
               hidden_size=__lowerCamelCase       ,						projection_dim=__lowerCamelCase       ,						num_hidden_layers=5       ,						num_attention_heads=4       ,						image_size=32       ,						intermediate_size=37       ,						patch_size=1       ,						)						)

       # regular denoising components

       torch.manual_seed(0						)
       lowerCAmelCase__										=      StableUnCLIPImageNormalizer(embedding_dim=__lowerCamelCase						)
       lowerCAmelCase__										=      DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''						)

       torch.manual_seed(0						)
       lowerCAmelCase__										=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

       torch.manual_seed(0						)
       lowerCAmelCase__										=      CLIPTextModel(
           CLIPTextConfig(
               bos_token_id=0       ,						eos_token_id=2       ,						hidden_size=__lowerCamelCase       ,						projection_dim=32       ,						intermediate_size=37       ,						layer_norm_eps=1e-05       ,						num_attention_heads=4       ,						num_hidden_layers=5       ,						pad_token_id=1       ,						vocab_size=10_00       ,						)						)

       torch.manual_seed(0						)
       lowerCAmelCase__										=      UNetaDConditionModel(
           sample_size=32       ,						in_channels=4       ,						out_channels=4       ,						down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')       ,						up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')       ,						block_out_channels=(32, 64)       ,						attention_head_dim=(2, 4)       ,						class_embed_type='''projection'''       ,						projection_class_embeddings_input_dim=embedder_projection_dim * 2       ,						cross_attention_dim=__lowerCamelCase       ,						layers_per_block=1       ,						upcast_attention=__lowerCamelCase       ,						use_linear_projection=__lowerCamelCase       ,						)

       torch.manual_seed(0						)
       lowerCAmelCase__										=      DDIMScheduler(
           beta_schedule='''scaled_linear'''       ,						beta_start=0.0_0085       ,						beta_end=0.012       ,						prediction_type='''v_prediction'''       ,						set_alpha_to_one=__lowerCamelCase       ,						steps_offset=1       ,						)

       torch.manual_seed(0						)
       lowerCAmelCase__										=      AutoencoderKL()

       lowerCAmelCase__										=      {
           # image encoding components
           '''feature_extractor''': feature_extractor,
           '''image_encoder''': image_encoder.eval(),
           # image noising components
           '''image_normalizer''': image_normalizer.eval(),
           '''image_noising_scheduler''': image_noising_scheduler,
           # regular denoising components
           '''tokenizer''': tokenizer,
           '''text_encoder''': text_encoder.eval(),
           '''unet''': unet.eval(),
           '''scheduler''': scheduler,
           '''vae''': vae.eval(),
       }

       return components
     def        A__      (						self	:		str       ,						__lowerCamelCase	:		Optional[Any]       ,						__lowerCamelCase	:		Any=0       ,						__lowerCamelCase	:		List[str]=True						):



       """simple docstring"""
       if str(__lowerCamelCase						).startswith('''mps'''						):
         lowerCAmelCase__										=      torch.manual_seed(__lowerCamelCase						)
       else:
         lowerCAmelCase__										=      torch.Generator(device=__lowerCamelCase						).manual_seed(__lowerCamelCase						)

       lowerCAmelCase__										=      floats_tensor((1, 3, 32, 32)       ,						rng=random.Random(__lowerCamelCase						)						).to(__lowerCamelCase						)

       if pil_image:
         lowerCAmelCase__										=      input_image * 0.5 + 0.5
         lowerCAmelCase__										=      input_image.clamp(0       ,						1						)
         lowerCAmelCase__										=      input_image.cpu().permute(0       ,						2       ,						3       ,						1						).float().numpy()
         lowerCAmelCase__										=      DiffusionPipeline.numpy_to_pil(__lowerCamelCase						)[0]

       return {
           "prompt": "An anime racoon running a marathon",
           "image": input_image,
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "np",
       }
     @skip_mps
     def        A__      (						self	:		Dict						):



       """simple docstring"""
       lowerCAmelCase__										=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       lowerCAmelCase__										=      self.get_dummy_components()
       lowerCAmelCase__										=      StableUnCLIPImgaImgPipeline(**__lowerCamelCase						)
       lowerCAmelCase__										=      sd_pipe.to(__lowerCamelCase						)
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase						)

       lowerCAmelCase__										=      self.get_dummy_inputs(__lowerCamelCase						)
       inputs.update({'''image_embeds''': None}						)
       lowerCAmelCase__										=      sd_pipe(**__lowerCamelCase						).images
       lowerCAmelCase__										=      image[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       lowerCAmelCase__										=      np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-3
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      torch_device in ['''cpu''', '''mps''']

       self._test_attention_slicing_forward_pass(test_max_difference=__lowerCamelCase						)
     def        A__      (						self	:		Any						):



       """simple docstring"""
       lowerCAmelCase__										=      torch_device in ['''cpu''', '''mps''']

       self._test_inference_batch_single_identical(test_max_difference=__lowerCamelCase						)



     @unittest.skipIf(
         torch_device != '''cuda''' or not is_xformers_available()       ,						reason='''XFormers attention is only available with CUDA and `xformers` installed'''       ,						)
     def        A__      (						self	:		Union[str, Any]						):



       """simple docstring"""
       self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__lowerCamelCase						)





@slow
@require_torch_gpu
class     SCREAMING_SNAKE_CASE__		(unittest.TestCase			):
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
     def        A__      (						self	:		str						):



       """simple docstring"""
       lowerCAmelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''						)

       lowerCAmelCase__										=      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'''						)

       lowerCAmelCase__										=      StableUnCLIPImgaImgPipeline.from_pretrained(
           '''fusing/stable-unclip-2-1-l-img2img'''       ,						torch_dtype=torch.floataa						)
       pipe.to(__lowerCamelCase						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)
       # stable unclip will oom when integration tests are run on a V100,
       # so turn on memory savings
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       lowerCAmelCase__										=      torch.Generator(device='''cpu'''						).manual_seed(0						)
       lowerCAmelCase__										=      pipe(__lowerCamelCase       ,						'''anime turle'''       ,						generator=__lowerCamelCase       ,						output_type='''np'''						)

       lowerCAmelCase__										=      output.images[0]

       assert image.shape == (7_68, 7_68, 3)

       assert_mean_pixel_difference(__lowerCamelCase       ,						__lowerCamelCase						)
     def        A__      (						self	:		int						):



       """simple docstring"""
       lowerCAmelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''						)

       lowerCAmelCase__										=      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'''						)

       lowerCAmelCase__										=      StableUnCLIPImgaImgPipeline.from_pretrained(
           '''fusing/stable-unclip-2-1-h-img2img'''       ,						torch_dtype=torch.floataa						)
       pipe.to(__lowerCamelCase						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)
       # stable unclip will oom when integration tests are run on a V100,
       # so turn on memory savings
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       lowerCAmelCase__										=      torch.Generator(device='''cpu'''						).manual_seed(0						)
       lowerCAmelCase__										=      pipe(__lowerCamelCase       ,						'''anime turle'''       ,						generator=__lowerCamelCase       ,						output_type='''np'''						)

       lowerCAmelCase__										=      output.images[0]

       assert image.shape == (7_68, 7_68, 3)

       assert_mean_pixel_difference(__lowerCamelCase       ,						__lowerCamelCase						)



     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       lowerCAmelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''						)

       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       lowerCAmelCase__										=      StableUnCLIPImgaImgPipeline.from_pretrained(
           '''fusing/stable-unclip-2-1-h-img2img'''       ,						torch_dtype=torch.floataa						)
       lowerCAmelCase__										=      pipe.to(__lowerCamelCase						)
       pipe.set_progress_bar_config(disable=__lowerCamelCase						)
       pipe.enable_attention_slicing()
       pipe.enable_sequential_cpu_offload()

       lowerCAmelCase__										=      pipe(
           __lowerCamelCase       ,						'''anime turtle'''       ,						num_inference_steps=2       ,						output_type='''np'''       ,						)

       lowerCAmelCase__										=      torch.cuda.max_memory_allocated()
       # make sure that less than 7 GB is allocated
       assert mem_bytes < 7 * 10**9







import os
import re
import warnings
from shutil import copyfile
from typing import List, Optional, Tuple

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


if is_sentencepiece_available():
      from .tokenization_ta import TaTokenizer
else:
      __magic_name__  :  Dict					=       None


__magic_name__  :  Optional[Any]					=       logging.get_logger(__name__)

__magic_name__  :  Optional[Any]					=       {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

__magic_name__  :  Union[str, Any]					=       {
    """vocab_file""": {
        """t5-small""": """https://huggingface.co/t5-small/resolve/main/spiece.model""",
        """t5-base""": """https://huggingface.co/t5-base/resolve/main/spiece.model""",
        """t5-large""": """https://huggingface.co/t5-large/resolve/main/spiece.model""",
        """t5-3b""": """https://huggingface.co/t5-3b/resolve/main/spiece.model""",
        """t5-11b""": """https://huggingface.co/t5-11b/resolve/main/spiece.model""",
    },
    """tokenizer_file""": {
        """t5-small""": """https://huggingface.co/t5-small/resolve/main/tokenizer.json""",
        """t5-base""": """https://huggingface.co/t5-base/resolve/main/tokenizer.json""",
        """t5-large""": """https://huggingface.co/t5-large/resolve/main/tokenizer.json""",
        """t5-3b""": """https://huggingface.co/t5-3b/resolve/main/tokenizer.json""",
        """t5-11b""": """https://huggingface.co/t5-11b/resolve/main/tokenizer.json""",
    },
}


# TODO(PVP) - this should be removed in Transformers v5
__magic_name__  :  Tuple					=       {
    """t5-small""": 5_12,
    """t5-base""": 5_12,
    """t5-large""": 5_12,
    """t5-3b""": 5_12,
    """t5-11b""": 5_12,
}





class     SCREAMING_SNAKE_CASE__		(_a			):
     lowercase_ :      Dict				= VOCAB_FILES_NAMES
     lowercase_ :      Any				= PRETRAINED_VOCAB_FILES_MAP
     lowercase_ :      Tuple				= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     lowercase_ :      Tuple				= ["input_ids", "attention_mask"]
     lowercase_ :      Any				= TaTokenizer

     lowercase_ :      List[int]				= []
     def __init__(						self	:		Tuple       ,						__lowerCamelCase	:		Tuple=None       ,						__lowerCamelCase	:		str=None       ,						__lowerCamelCase	:		Optional[int]="</s>"       ,						__lowerCamelCase	:		Optional[Any]="<unk>"       ,						__lowerCamelCase	:		int="<pad>"       ,						__lowerCamelCase	:		int=1_00       ,						__lowerCamelCase	:		List[str]=None       ,						**__lowerCamelCase	:		Tuple       ,						):



       """simple docstring"""
       # Add extra_ids to the special token list
       if extra_ids > 0 and additional_special_tokens is None:
         lowerCAmelCase__										=      [F"""<extra_id_{i}>""" for i in range(__lowerCamelCase						)]
       elif extra_ids > 0 and additional_special_tokens is not None:
         # Check that we have the right number of extra special tokens
         lowerCAmelCase__										=      len(set(filter(lambda __lowerCamelCase						: bool('''extra_id_''' in str(__lowerCamelCase						)						)       ,						__lowerCamelCase						)						)						)
         if extra_tokens != extra_ids:
           raise ValueError(
               F"""Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"""
               ''' provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids'''
               ''' tokens'''						)

       super().__init__(
           __lowerCamelCase       ,						tokenizer_file=__lowerCamelCase       ,						eos_token=__lowerCamelCase       ,						unk_token=__lowerCamelCase       ,						pad_token=__lowerCamelCase       ,						extra_ids=__lowerCamelCase       ,						additional_special_tokens=__lowerCamelCase       ,						**__lowerCamelCase       ,						)

       lowerCAmelCase__										=      vocab_file
       lowerCAmelCase__										=      False if not self.vocab_file else True
       lowerCAmelCase__										=      extra_ids
     @staticmethod
     def        A__      (						__lowerCamelCase	:		Dict       ,						__lowerCamelCase	:		List[str]       ,						__lowerCamelCase	:		Any						):



       """simple docstring"""
       if pretrained_model_name_or_path in TaTokenizerFast.max_model_input_sizes:
         lowerCAmelCase__										=      TaTokenizerFast.max_model_input_sizes[pretrained_model_name_or_path]
         if init_max_model_length is not None and init_max_model_length != max_model_length:
           return init_max_model_length
         elif init_max_model_length is None:
           warnings.warn(
               '''This tokenizer was incorrectly instantiated with a model max length of'''
               F""" {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this"""
               ''' behavior is kept to avoid breaking backwards compatibility when padding/encoding with'''
               ''' `truncation is True`.\n- Be aware that you SHOULD NOT rely on'''
               F""" {pretrained_model_name_or_path} automatically truncating your input to"""
               F""" {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences"""
               F""" longer than {deprecated_max_model_length} you can either instantiate this tokenizer with"""
               ''' `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please'''
               ''' instantiate this tokenizer with `model_max_length` set to your preferred value.'''       ,						__lowerCamelCase       ,						)

       return max_model_length
     def        A__      (						self	:		int       ,						__lowerCamelCase	:		str       ,						__lowerCamelCase	:		Optional[str] = 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						)
         logger.info(F"""Copy vocab file to {out_vocab_file}"""						)

       return (out_vocab_file,)
     def        A__      (						self	:		Optional[Any]       ,						__lowerCamelCase	:		List[int]       ,						__lowerCamelCase	:		Optional[List[int]] = None						):



       """simple docstring"""
       lowerCAmelCase__										=      token_ids_a + [self.eos_token_id]
       if token_ids_a is None:
         return self.prefix_tokens + token_ids_a
       else:
         lowerCAmelCase__										=      token_ids_a + [self.eos_token_id]
         return self.prefix_tokens + token_ids_a + token_ids_a
     def        A__      (						self	:		int       ,						__lowerCamelCase	:		List[int]       ,						__lowerCamelCase	:		Optional[List[int]] = None						):



       """simple docstring"""
       lowerCAmelCase__										=      [self.eos_token_id]

       if token_ids_a is None:
         return len(token_ids_a + eos						) * [0]
       return len(token_ids_a + eos + token_ids_a + eos						) * [0]
     def        A__      (						self	:		int						):



       """simple docstring"""
       return list(
           set(filter(lambda __lowerCamelCase						: bool(re.search(r'''<extra_id_\d+>'''       ,						__lowerCamelCase						)						) is not None       ,						self.additional_special_tokens						)						)						)



     def        A__      (						self	:		Dict						):



       """simple docstring"""
       return [self.convert_tokens_to_ids(__lowerCamelCase						) for token in self.get_sentinel_tokens()]

