



"""simple docstring"""




import gc
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class       UpperCAmelCase_    (							_UpperCamelCase			,     unittest.TestCase		):
      __SCREAMING_SNAKE_CASE		:  int						=					LDMTextToImagePipeline
      __SCREAMING_SNAKE_CASE		:  List[str]						=					TEXT_TO_IMAGE_PARAMS - {
          'negative_prompt',
          'negative_prompt_embeds',
          'cross_attention_kwargs',
          'prompt_embeds',
      }
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					PipelineTesterMixin.required_optional_params - {
          'num_images_per_prompt',
          'callback',
          'callback_steps',
      }
      __SCREAMING_SNAKE_CASE		:  Union[str, Any]						=					TEXT_TO_IMAGE_BATCH_PARAMS
      __SCREAMING_SNAKE_CASE		:  Any						=					False



      def 		snake_case_   (	self	:      str		):
           torch.manual_seed(0		)
           _UpperCAmelCase     :      List[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	,      )
           _UpperCAmelCase     :      Optional[int]														=       DDIMScheduler(
               beta_start=0.00_085	,      beta_end=0.012	,      beta_schedule="scaled_linear"	,      clip_sample=A	,      set_alpha_to_one=A	,      )
           torch.manual_seed(0		)
           _UpperCAmelCase     :      Optional[Any]														=       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		)
           _UpperCAmelCase     :      Optional[int]														=       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	,      )
           _UpperCAmelCase     :      int														=       CLIPTextModel(A		)
           _UpperCAmelCase     :      Optional[Any]														=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

           _UpperCAmelCase     :      Any														=       {
               "unet": unet,
               "scheduler": scheduler,
               "vqvae": vae,
               "bert": text_encoder,
               "tokenizer": tokenizer,
           }
           return components



      def 		snake_case_   (	self	:      Union[str, Any]	,      A	:      int	,      A	:      Any=0		):
           if str(A		).startswith("mps"		):
                _UpperCAmelCase     :      List[Any]														=       torch.manual_seed(A		)
           else:
                _UpperCAmelCase     :      List[str]														=       torch.Generator(device=A		).manual_seed(A		)
           _UpperCAmelCase     :      Tuple														=       {
               "prompt": "A painting of a squirrel eating a burger",
               "generator": generator,
               "num_inference_steps": 2,
               "guidance_scale": 6.0,
               "output_type": "numpy",
           }
           return inputs



      def 		snake_case_   (	self	:      Dict		):
           _UpperCAmelCase     :      Any														=       "cpu"  # ensure determinism for the device-dependent torch.Generator

           _UpperCAmelCase     :      List[str]														=       self.get_dummy_components()
           _UpperCAmelCase     :      Any														=       LDMTextToImagePipeline(**A		)
           pipe.to(A		)
           pipe.set_progress_bar_config(disable=A		)

           _UpperCAmelCase     :      str														=       self.get_dummy_inputs(A		)
           _UpperCAmelCase     :      Dict														=       pipe(**A		).images
           _UpperCAmelCase     :      int														=       image[0, -3:, -3:, -1]

           assert image.shape == (1, 1_6, 1_6, 3)
           _UpperCAmelCase     :      Optional[int]														=       np.array([0.6_101, 0.6_156, 0.5_622, 0.4_895, 0.6_661, 0.3_804, 0.5_748, 0.6_136, 0.5_014]		)

           assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-3







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



      def 		snake_case_   (	self	:      List[str]		):
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



      def 		snake_case_   (	self	:      Any	,      A	:      List[Any]	,      A	:      int=torch.floataa	,      A	:      str=0		):
           _UpperCAmelCase     :      str														=       torch.manual_seed(A		)
           _UpperCAmelCase     :      Any														=       np.random.RandomState(A		).standard_normal((1, 4, 3_2, 3_2)		)
           _UpperCAmelCase     :      Optional[Any]														=       torch.from_numpy(A		).to(device=A	,      dtype=A		)
           _UpperCAmelCase     :      List[str]														=       {
               "prompt": "A painting of a squirrel eating a burger",
               "latents": latents,
               "generator": generator,
               "num_inference_steps": 3,
               "guidance_scale": 6.0,
               "output_type": "numpy",
           }
           return inputs



      def 		snake_case_   (	self	:      Union[str, Any]		):
           _UpperCAmelCase     :      Any														=       LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"		).to(A		)
           pipe.set_progress_bar_config(disable=A		)

           _UpperCAmelCase     :      Union[str, Any]														=       self.get_inputs(A		)
           _UpperCAmelCase     :      List[str]														=       pipe(**A		).images
           _UpperCAmelCase     :      int														=       image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 2_5_6, 2_5_6, 3)
           _UpperCAmelCase     :      Tuple														=       np.array([0.51_825, 0.52_850, 0.52_543, 0.54_258, 0.52_304, 0.52_569, 0.54_363, 0.55_276, 0.56_878]		)
           _UpperCAmelCase     :      List[Any]														=       np.abs(expected_slice - image_slice		).max()
           assert max_diff < 1e-3







@nightly
@require_torch_gpu
class       UpperCAmelCase_    (							unittest.TestCase		):



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



      def 		snake_case_   (	self	:      Any	,      A	:      Union[str, Any]	,      A	:      List[str]=torch.floataa	,      A	:      List[str]=0		):
           _UpperCAmelCase     :      int														=       torch.manual_seed(A		)
           _UpperCAmelCase     :      str														=       np.random.RandomState(A		).standard_normal((1, 4, 3_2, 3_2)		)
           _UpperCAmelCase     :      Any														=       torch.from_numpy(A		).to(device=A	,      dtype=A		)
           _UpperCAmelCase     :      List[str]														=       {
               "prompt": "A painting of a squirrel eating a burger",
               "latents": latents,
               "generator": generator,
               "num_inference_steps": 5_0,
               "guidance_scale": 6.0,
               "output_type": "numpy",
           }
           return inputs



      def 		snake_case_   (	self	:      int		):
           _UpperCAmelCase     :      List[Any]														=       LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"		).to(A		)
           pipe.set_progress_bar_config(disable=A		)

           _UpperCAmelCase     :      Tuple														=       self.get_inputs(A		)
           _UpperCAmelCase     :      Union[str, Any]														=       pipe(**A		).images[0]

           _UpperCAmelCase     :      Union[str, Any]														=       load_numpy(
               "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy"		)
           _UpperCAmelCase     :      str														=       np.abs(expected_image - image		).max()
           assert max_diff < 1e-3





"""simple docstring"""




import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_gpta import GPTaTokenizer


if TYPE_CHECKING:
 from transformers.pipelines.conversational import Conversation


_lowerCAmelCase							:   Any				=				logging.get_logger(__name__)

_lowerCAmelCase							:   Any				=				{"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

_lowerCAmelCase							:   List[str]				=				{
    "vocab_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/vocab.json",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/vocab.json",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/vocab.json",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/vocab.json",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/vocab.json",
    },
    "merges_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/merges.txt",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/merges.txt",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/merges.txt",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/merges.txt",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/tokenizer.json",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/tokenizer.json",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/tokenizer.json",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/tokenizer.json",
    },
}

_lowerCAmelCase							:   List[str]				=				{
    "gpt2": 10_24,
    "gpt2-medium": 10_24,
    "gpt2-large": 10_24,
    "gpt2-xl": 10_24,
    "distilgpt2": 10_24,
}






class       UpperCAmelCase_    (							_UpperCamelCase		):
      __SCREAMING_SNAKE_CASE		:  Optional[Any]						=					VOCAB_FILES_NAMES
      __SCREAMING_SNAKE_CASE		:  int						=					PRETRAINED_VOCAB_FILES_MAP
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					['input_ids', 'attention_mask']
      __SCREAMING_SNAKE_CASE		:  Dict						=					GPTaTokenizer



      def __init__(	self	:      Optional[Any]	,      A	:      Optional[int]=None	,      A	:      List[Any]=None	,      A	:      Optional[int]=None	,      A	:      Optional[Any]="<|endoftext|>"	,      A	:      Union[str, Any]="<|endoftext|>"	,      A	:      Optional[Any]="<|endoftext|>"	,      A	:      Optional[Any]=False	,      **A	:      str	,      ):
           super().__init__(
               A	,      A	,      tokenizer_file=A	,      unk_token=A	,      bos_token=A	,      eos_token=A	,      add_prefix_space=A	,      **A	,      )

           _UpperCAmelCase     :      List[Any]														=       kwargs.pop("add_bos_token"	,      A		)

           _UpperCAmelCase     :      List[str]														=       json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()		)
           if pre_tok_state.get("add_prefix_space"	,      A		) != add_prefix_space:
                _UpperCAmelCase     :      Dict														=       getattr(A	,      pre_tok_state.pop("type"		)		)
                _UpperCAmelCase     :      Optional[Any]														=       add_prefix_space
                _UpperCAmelCase     :      Optional[Any]														=       pre_tok_class(**A		)

           _UpperCAmelCase     :      List[Any]														=       add_prefix_space



      def 		snake_case_   (	self	:      List[Any]	,      *A	:      List[Any]	,      **A	:      List[str]		):
           _UpperCAmelCase     :      Union[str, Any]														=       kwargs.get("is_split_into_words"	,      A		)
           assert self.add_prefix_space or not is_split_into_words, (
               f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
               "to use it with pretokenized inputs."
           )

           return super()._batch_encode_plus(*A	,      **A		)



      def 		snake_case_   (	self	:      Optional[Any]	,      *A	:      List[Any]	,      **A	:      int		):
           _UpperCAmelCase     :      Union[str, Any]														=       kwargs.get("is_split_into_words"	,      A		)

           assert self.add_prefix_space or not is_split_into_words, (
               f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
               "to use it with pretokenized inputs."
           )

           return super()._encode_plus(*A	,      **A		)



      def 		snake_case_   (	self	:      Optional[int]	,      A	:      str	,      A	:      Optional[str] = None		):
           _UpperCAmelCase     :      List[Any]														=       self._tokenizer.model.save(A	,      name=A		)
           return tuple(A		)



      def 		snake_case_   (	self	:      int	,      A	:      "Conversation"		):
           _UpperCAmelCase     :      List[str]														=       []
           for is_user, text in conversation.iter_texts():
                input_ids.extend(self.encode(A	,      add_special_tokens=A		) + [self.eos_token_id]		)

           if len(A		) > self.model_max_length:
                _UpperCAmelCase     :      Optional[int]														=       input_ids[-self.model_max_length :]
           return input_ids
