


import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()



class 							_UpperCamelCase (						A,A,A,unittest.TestCase						):

   '''simple docstring'''







   a_				:							Optional[int]    = StableUnCLIPPipeline
   a_				:							List[Any]    = TEXT_TO_IMAGE_PARAMS
   a_				:							Union[str, Any]    = TEXT_TO_IMAGE_BATCH_PARAMS
   a_				:							int    = TEXT_TO_IMAGE_IMAGE_PARAMS
   a_				:							List[str]    = TEXT_TO_IMAGE_IMAGE_PARAMS

   # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
   a_				:							List[str]    = False

   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						3_2
         __lowerCamelCase     :      Optional[Any]       						=						embedder_hidden_size

         # prior components

         torch.manual_seed(0							)
         __lowerCamelCase     :      str       						=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

         torch.manual_seed(0							)
         __lowerCamelCase     :      str       						=						CLIPTextModelWithProjection(
             CLIPTextConfig(
                 bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=_lowerCamelCase       ,  projection_dim=_lowerCamelCase       ,  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       ,  )							)

         torch.manual_seed(0							)
         __lowerCamelCase     :      Optional[int]       						=						PriorTransformer(
             num_attention_heads=2       ,  attention_head_dim=1_2       ,  embedding_dim=_lowerCamelCase       ,  num_layers=1       ,  )

         torch.manual_seed(0							)
         __lowerCamelCase     :      Any       						=						DDPMScheduler(
             variance_type="""fixed_small_log"""       ,  prediction_type="""sample"""       ,  num_train_timesteps=1_0_0_0       ,  clip_sample=_lowerCamelCase       ,  clip_sample_range=5.0       ,  beta_schedule="""squaredcos_cap_v2"""       ,  )

         # regular denoising components

         torch.manual_seed(0							)
         __lowerCamelCase     :      int       						=						StableUnCLIPImageNormalizer(embedding_dim=_lowerCamelCase							)
         __lowerCamelCase     :      int       						=						DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""							)

         torch.manual_seed(0							)
         __lowerCamelCase     :      Any       						=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

         torch.manual_seed(0							)
         __lowerCamelCase     :      List[str]       						=						CLIPTextModel(
             CLIPTextConfig(
                 bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=_lowerCamelCase       ,  projection_dim=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       ,  )							)

         torch.manual_seed(0							)
         __lowerCamelCase     :      int       						=						UNetaDConditionModel(
             sample_size=3_2       ,  in_channels=4       ,  out_channels=4       ,  down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")       ,  up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")       ,  block_out_channels=(3_2, 6_4)       ,  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     :      Any       						=						DDIMScheduler(
             beta_schedule="""scaled_linear"""       ,  beta_start=0.00_085       ,  beta_end=0.012       ,  prediction_type="""v_prediction"""       ,  set_alpha_to_one=_lowerCamelCase       ,  steps_offset=1       ,  )

         torch.manual_seed(0							)
         __lowerCamelCase     :      Optional[Any]       						=						AutoencoderKL()

         __lowerCamelCase     :      Union[str, Any]       						=						{
             # prior components
             """prior_tokenizer""": prior_tokenizer,
             """prior_text_encoder""": prior_text_encoder,
             """prior""": prior,
             """prior_scheduler""": prior_scheduler,
             # image noising components
             """image_normalizer""": image_normalizer,
             """image_noising_scheduler""": image_noising_scheduler,
             # regular denoising components
             """tokenizer""": tokenizer,
             """text_encoder""": text_encoder,
             """unet""": unet,
             """scheduler""": scheduler,
             """vae""": vae,
         }

         return components

   def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : str=0							):






         '''simple docstring'''



         if str(_lowerCamelCase							).startswith("""mps"""							):
               __lowerCamelCase     :      Union[str, Any]       						=						torch.manual_seed(_lowerCamelCase							)
         else:
               __lowerCamelCase     :      Tuple       						=						torch.Generator(device=_lowerCamelCase							).manual_seed(_lowerCamelCase							)
         __lowerCamelCase     :      Any       						=						{
             """prompt""": """A painting of a squirrel eating a burger""",
             """generator""": generator,
             """num_inference_steps""": 2,
             """prior_num_inference_steps""": 2,
             """output_type""": """numpy""",
         }
         return inputs

   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      int       						=						torch_device == """cpu"""

         self._test_attention_slicing_forward_pass(test_max_difference=_lowerCamelCase							)


   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						torch_device in ["""cpu""", """mps"""]

         self._test_inference_batch_single_identical(test_max_difference=_lowerCamelCase							)



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

   '''simple docstring'''







   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""							)

         __lowerCamelCase     :      Optional[int]       						=						StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""       ,  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     :      str       						=						torch.Generator(device="""cpu"""							).manual_seed(0							)
         __lowerCamelCase     :      Union[str, Any]       						=						pipe("""anime turle"""       ,  generator=_lowerCamelCase       ,  output_type="""np"""							)

         __lowerCamelCase     :      Optional[int]       						=						output.images[0]

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

         assert_mean_pixel_difference(_lowerCamelCase       ,  _lowerCamelCase							)


   def   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



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

         __lowerCamelCase     :      Dict       						=						StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""       ,  torch_dtype=torch.floataa							)
         __lowerCamelCase     :      List[Any]       						=						pipe.to(_lowerCamelCase							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()

         __lowerCamelCase     :      Optional[Any]       						=						pipe(
             """anime turtle"""       ,  prior_num_inference_steps=2       ,  num_inference_steps=2       ,  output_type="""np"""       ,  )

         __lowerCamelCase     :      int       						=						torch.cuda.max_memory_allocated()
         # make sure that less than 7 GB is allocated
         assert mem_bytes < 7 * 1_0**9







import json
from typing import Iterator, List, Union

from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, trainers
from tokenizers.implementations.base_tokenizer import BaseTokenizer
from tokenizers.models import Unigram
from tokenizers.processors import TemplateProcessing



class 							_UpperCamelCase (						A						):

   '''simple docstring'''







   def __init__(  self  : str       ,  _lowerCamelCase  : str = "▁"       ,  _lowerCamelCase  : bool = True       ,  _lowerCamelCase  : Union[str, AddedToken] = "<unk>"       ,  _lowerCamelCase  : Union[str, AddedToken] = "</s>"       ,  _lowerCamelCase  : Union[str, AddedToken] = "<pad>"       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						{
             """pad""": {"""id""": 0, """token""": pad_token},
             """eos""": {"""id""": 1, """token""": eos_token},
             """unk""": {"""id""": 2, """token""": unk_token},
         }

         __lowerCamelCase     :      Optional[Any]       						=						[None] * len(self.special_tokens							)
         for token_dict in self.special_tokens.values():
               __lowerCamelCase     :      Tuple       						=						token_dict["""token"""]

         __lowerCamelCase     :      Optional[int]       						=						Tokenizer(Unigram()							)

         __lowerCamelCase     :      str       						=						normalizers.Sequence(
             [
                 normalizers.Nmt(),
                 normalizers.NFKC(),
                 normalizers.Replace(Regex(""" {2,}"""							)       ,  """ """							),
                 normalizers.Lowercase(),
             ]							)
         __lowerCamelCase     :      int       						=						pre_tokenizers.Sequence(
             [
                 pre_tokenizers.Metaspace(replacement=_lowerCamelCase       ,  add_prefix_space=_lowerCamelCase							),
                 pre_tokenizers.Digits(individual_digits=_lowerCamelCase							),
                 pre_tokenizers.Punctuation(),
             ]							)
         __lowerCamelCase     :      Tuple       						=						decoders.Metaspace(replacement=_lowerCamelCase       ,  add_prefix_space=_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						TemplateProcessing(
             single=F"""$A {self.special_tokens["eos"]["token"]}"""       ,  special_tokens=[(self.special_tokens["""eos"""]["""token"""], self.special_tokens["""eos"""]["""id"""])]       ,  )

         __lowerCamelCase     :      Tuple       						=						{
             """model""": """SentencePieceUnigram""",
             """replacement""": replacement,
             """add_prefix_space""": add_prefix_space,
         }

         super().__init__(_lowerCamelCase       ,  _lowerCamelCase							)

   def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : Union[str, List[str]]       ,  _lowerCamelCase  : int = 8_0_0_0       ,  _lowerCamelCase  : bool = True       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						trainers.UnigramTrainer(
             vocab_size=_lowerCamelCase       ,  special_tokens=self.special_tokens_list       ,  show_progress=_lowerCamelCase       ,  )

         if isinstance(_lowerCamelCase       ,  _lowerCamelCase							):
               __lowerCamelCase     :      Optional[int]       						=						[files]
         self._tokenizer.train(_lowerCamelCase       ,  trainer=_lowerCamelCase							)

         self.add_unk_id()

   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : Union[Iterator[str], Iterator[Iterator[str]]]       ,  _lowerCamelCase  : int = 8_0_0_0       ,  _lowerCamelCase  : bool = True       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						trainers.UnigramTrainer(
             vocab_size=_lowerCamelCase       ,  special_tokens=self.special_tokens_list       ,  show_progress=_lowerCamelCase       ,  )

         self._tokenizer.train_from_iterator(_lowerCamelCase       ,  trainer=_lowerCamelCase							)

         self.add_unk_id()


   def   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						json.loads(self._tokenizer.to_str()							)

         __lowerCamelCase     :      Optional[int]       						=						self.special_tokens["""unk"""]["""id"""]

         __lowerCamelCase     :      List[str]       						=						Tokenizer.from_str(json.dumps(_lowerCamelCase							)							)



