

"""simple docstring"""






import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()







class 	__UpperCamelCase    (    __lowerCAmelCase							,							__lowerCAmelCase							,							__lowerCAmelCase							,							unittest.TestCase						):
  _UpperCAmelCase          =		AltDiffusionPipeline
  _UpperCAmelCase          =		TEXT_TO_IMAGE_PARAMS
  _UpperCAmelCase          =		TEXT_TO_IMAGE_BATCH_PARAMS
  _UpperCAmelCase          =		TEXT_TO_IMAGE_IMAGE_PARAMS
  _UpperCAmelCase          =		TEXT_TO_IMAGE_IMAGE_PARAMS

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      torch.manual_seed(0      )
      _lowerCAmelCase       :   str  =    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    ,)
      _lowerCAmelCase       :   Tuple  =    DDIMScheduler(
          beta_start=0.0_0_0_8_5    ,beta_end=0.0_1_2    ,beta_schedule='scaled_linear'    ,clip_sample=lowerCamelCase__    ,set_alpha_to_one=lowerCamelCase__    ,)
      torch.manual_seed(0      )
      _lowerCAmelCase       :   List[str]  =    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    ,)

      # TODO: address the non-deterministic text encoder (fails for save-load tests)
      # torch.manual_seed(0)
      # text_encoder_config = RobertaSeriesConfig(
      #     hidden_size=32,
      #     project_dim=32,
      #     intermediate_size=37,
      #     layer_norm_eps=1e-05,
      #     num_attention_heads=4,
      #     num_hidden_layers=5,
      #     vocab_size=5002,
      # )
      # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

      torch.manual_seed(0      )
      _lowerCAmelCase       :   Union[str, Any]  =    CLIPTextConfig(
          bos_token_id=0    ,eos_token_id=2    ,hidden_size=32    ,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=5002    ,)
      _lowerCAmelCase       :   List[Any]  =    CLIPTextModel(lowerCamelCase__      )

      _lowerCAmelCase       :   List[Any]  =    XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'      )
      _lowerCAmelCase       :   Tuple  =    77

      _lowerCAmelCase       :   Dict  =    {
          '''unet''': unet,
          '''scheduler''': scheduler,
          '''vae''': vae,
          '''text_encoder''': text_encoder,
          '''tokenizer''': tokenizer,
          '''safety_checker''': None,
          '''feature_extractor''': None,
      }
      return components

  def __lowerCamelCase		(		self    ,_A    ,_A=0      ):
      '''simple docstring'''

      if str(lowerCamelCase__      ).startswith('mps'      ):
          _lowerCAmelCase       :   Optional[int]  =    torch.manual_seed(lowerCamelCase__      )
      else:
          _lowerCAmelCase       :   Dict  =    torch.Generator(device=lowerCamelCase__      ).manual_seed(lowerCamelCase__      )
      _lowerCAmelCase       :   Dict  =    {
          '''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 __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().test_attention_slicing_forward_pass(expected_max_diff=3E-3      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().test_inference_batch_single_identical(expected_max_diff=3E-3      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator

      _lowerCAmelCase       :   Optional[int]  =    self.get_dummy_components()
      torch.manual_seed(0      )
      _lowerCAmelCase       :   str  =    RobertaSeriesConfig(
          hidden_size=32    ,project_dim=32    ,intermediate_size=37    ,layer_norm_eps=1E-05    ,num_attention_heads=4    ,num_hidden_layers=5    ,vocab_size=5002    ,)
      # TODO: remove after fixing the non-deterministic text encoder
      _lowerCAmelCase       :   int  =    RobertaSeriesModelWithTransformation(lowerCamelCase__      )
      _lowerCAmelCase       :   Any  =    text_encoder

      _lowerCAmelCase       :   Dict  =    AltDiffusionPipeline(**lowerCamelCase__      )
      _lowerCAmelCase       :   Any  =    alt_pipe.to(lowerCamelCase__      )
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__      )

      _lowerCAmelCase       :   Dict  =    self.get_dummy_inputs(lowerCamelCase__      )
      _lowerCAmelCase       :   Optional[int]  =    '''A photo of an astronaut'''
      _lowerCAmelCase       :   int  =    alt_pipe(**lowerCamelCase__      )
      _lowerCAmelCase       :   Tuple  =    output.images
      _lowerCAmelCase       :   str  =    image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      _lowerCAmelCase       :   List[Any]  =    np.array(
          [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]      )

      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator

      _lowerCAmelCase       :   Optional[Any]  =    self.get_dummy_components()
      _lowerCAmelCase       :   Optional[Any]  =    PNDMScheduler(skip_prk_steps=lowerCamelCase__      )
      torch.manual_seed(0      )
      _lowerCAmelCase       :   Optional[int]  =    RobertaSeriesConfig(
          hidden_size=32    ,project_dim=32    ,intermediate_size=37    ,layer_norm_eps=1E-05    ,num_attention_heads=4    ,num_hidden_layers=5    ,vocab_size=5002    ,)
      # TODO: remove after fixing the non-deterministic text encoder
      _lowerCAmelCase       :   int  =    RobertaSeriesModelWithTransformation(lowerCamelCase__      )
      _lowerCAmelCase       :   str  =    text_encoder
      _lowerCAmelCase       :   str  =    AltDiffusionPipeline(**lowerCamelCase__      )
      _lowerCAmelCase       :   Optional[Any]  =    alt_pipe.to(lowerCamelCase__      )
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__      )

      _lowerCAmelCase       :   List[Any]  =    self.get_dummy_inputs(lowerCamelCase__      )
      _lowerCAmelCase       :   Tuple  =    alt_pipe(**lowerCamelCase__      )
      _lowerCAmelCase       :   int  =    output.images
      _lowerCAmelCase       :   Tuple  =    image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      _lowerCAmelCase       :   List[Any]  =    np.array(
          [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]      )

      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2







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

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

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

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'    ,safety_checker=lowerCamelCase__      )
      _lowerCAmelCase       :   Union[str, Any]  =    alt_pipe.to(lowerCamelCase__      )
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__      )

      _lowerCAmelCase       :   Any  =    '''A painting of a squirrel eating a burger'''
      _lowerCAmelCase       :   Optional[int]  =    torch.manual_seed(0      )
      _lowerCAmelCase       :   Optional[Any]  =    alt_pipe([prompt]    ,generator=lowerCamelCase__    ,guidance_scale=6.0    ,num_inference_steps=20    ,output_type='np'      )

      _lowerCAmelCase       :   Union[str, Any]  =    output.images

      _lowerCAmelCase       :   Tuple  =    image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      _lowerCAmelCase       :   int  =    np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]      )

      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    DDIMScheduler.from_pretrained('BAAI/AltDiffusion'    ,subfolder='scheduler'      )

      _lowerCAmelCase       :   Optional[int]  =    AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'    ,scheduler=lowerCamelCase__    ,safety_checker=lowerCamelCase__      )
      _lowerCAmelCase       :   Dict  =    alt_pipe.to(lowerCamelCase__      )
      alt_pipe.set_progress_bar_config(disable=lowerCamelCase__      )

      _lowerCAmelCase       :   Optional[int]  =    '''A painting of a squirrel eating a burger'''
      _lowerCAmelCase       :   Dict  =    torch.manual_seed(0      )

      _lowerCAmelCase       :   List[str]  =    alt_pipe([prompt]    ,generator=lowerCamelCase__    ,num_inference_steps=2    ,output_type='numpy'      )
      _lowerCAmelCase       :   Optional[int]  =    output.images

      _lowerCAmelCase       :   Any  =    image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      _lowerCAmelCase       :   str  =    np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]      )

      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2









"""simple docstring"""






import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import SPIECE_UNDERLINE, logging


_lowerCAmelCase						=		logging.get_logger(__name__)

_lowerCAmelCase						=		{"""vocab_file""": """spiece.model"""}

_lowerCAmelCase						=		{
    """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""",
    }
}

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

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







class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		VOCAB_FILES_NAMES
  _UpperCAmelCase          =		PRETRAINED_VOCAB_FILES_MAP
  _UpperCAmelCase          =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  _UpperCAmelCase          =		"left"

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

      _lowerCAmelCase       :   Optional[Any]  =    AddedToken(_A    ,lstrip=_A    ,rstrip=_A      ) if isinstance(_A    ,_A      ) else mask_token

      _lowerCAmelCase       :   Union[str, Any]  =    {} if sp_model_kwargs is None else sp_model_kwargs

      super().__init__(
          do_lower_case=_A    ,remove_space=_A    ,keep_accents=_A    ,bos_token=_A    ,eos_token=_A    ,unk_token=_A    ,sep_token=_A    ,pad_token=_A    ,cls_token=_A    ,mask_token=_A    ,additional_special_tokens=_A    ,sp_model_kwargs=self.sp_model_kwargs    ,**_A    ,)

      _lowerCAmelCase       :   int  =    3

      _lowerCAmelCase       :   Union[str, Any]  =    do_lower_case
      _lowerCAmelCase       :   Dict  =    remove_space
      _lowerCAmelCase       :   int  =    keep_accents
      _lowerCAmelCase       :   List[str]  =    vocab_file

      _lowerCAmelCase       :   List[Any]  =    spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(_A      )

  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return len(self.sp_model      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    {self.convert_ids_to_tokens(_A      ): i for i in range(self.vocab_size      )}
      vocab.update(self.added_tokens_encoder      )
      return vocab

  def __getstate__(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    self.__dict__.copy()
      _lowerCAmelCase       :   List[str]  =    None
      return state

  def __setstate__(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    d

      # for backward compatibility
      if not hasattr(self    ,'sp_model_kwargs'      ):
          _lowerCAmelCase       :   Union[str, Any]  =    {}

      _lowerCAmelCase       :   Optional[Any]  =    spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(self.vocab_file      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      if self.remove_space:
          _lowerCAmelCase       :   str  =    ' '.join(inputs.strip().split()      )
      else:
          _lowerCAmelCase       :   Dict  =    inputs
      _lowerCAmelCase       :   List[str]  =    outputs.replace('``'    ,'"'      ).replace('\'\''    ,'"'      )

      if not self.keep_accents:
          _lowerCAmelCase       :   Optional[Any]  =    unicodedata.normalize('NFKD'    ,_A      )
          _lowerCAmelCase       :   Dict  =    ''.join([c for c in outputs if not unicodedata.combining(_A      )]      )
      if self.do_lower_case:
          _lowerCAmelCase       :   Tuple  =    outputs.lower()

      return outputs

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    self.preprocess_text(_A      )
      _lowerCAmelCase       :   int  =    self.sp_model.encode(_A    ,out_type=_A      )
      _lowerCAmelCase       :   int  =    []
      for piece in pieces:
          if len(_A      ) > 1 and piece[-1] == str(','      ) and piece[-2].isdigit():
              _lowerCAmelCase       :   Union[str, Any]  =    self.sp_model.EncodeAsPieces(piece[:-1].replace(_A    ,''      )      )
              if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
                  if len(cur_pieces[0]      ) == 1:
                      _lowerCAmelCase       :   int  =    cur_pieces[1:]
                  else:
                      _lowerCAmelCase       :   Any  =    cur_pieces[0][1:]
              cur_pieces.append(piece[-1]      )
              new_pieces.extend(_A      )
          else:
              new_pieces.append(_A      )

      return new_pieces

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      return self.sp_model.PieceToId(_A      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      return self.sp_model.IdToPiece(_A      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    ''.join(_A      ).replace(_A    ,' '      ).strip()
      return out_string

  def __lowerCamelCase		(		self    ,_A    ,_A = False    ,_A = None    ,_A = True    ,**_A    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    kwargs.pop('use_source_tokenizer'    ,_A      )

      _lowerCAmelCase       :   Dict  =    self.convert_ids_to_tokens(_A    ,skip_special_tokens=_A      )

      # To avoid mixing byte-level and unicode for byte-level BPT
      # we need to build string separately for added tokens and byte-level tokens
      # cf. https://github.com/huggingface/transformers/issues/1133
      _lowerCAmelCase       :   Optional[Any]  =    []
      _lowerCAmelCase       :   int  =    []
      for token in filtered_tokens:
          if skip_special_tokens and token in self.all_special_ids:
              continue
          if token in self.added_tokens_encoder:
              if current_sub_text:
                  sub_texts.append(self.convert_tokens_to_string(_A      )      )
                  _lowerCAmelCase       :   Tuple  =    []
              sub_texts.append(_A      )
          else:
              current_sub_text.append(_A      )
      if current_sub_text:
          sub_texts.append(self.convert_tokens_to_string(_A      )      )

      # Mimic the behavior of the Rust tokenizer:
      # By default, there are no spaces between special tokens
      _lowerCAmelCase       :   List[Any]  =    ''.join(_A      )

      _lowerCAmelCase       :   Tuple  =    (
          clean_up_tokenization_spaces
          if clean_up_tokenization_spaces is not None
          else self.clean_up_tokenization_spaces
      )
      if clean_up_tokenization_spaces:
          _lowerCAmelCase       :   int  =    self.clean_up_tokenization(_A      )
          return clean_text
      else:
          return text

  def __lowerCamelCase		(		self    ,_A    ,_A = None      ):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    [self.sep_token_id]
      _lowerCAmelCase       :   Union[str, Any]  =    [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    ,_A    ,_A = None    ,_A = False      ):
      '''simple docstring'''

      if already_has_special_tokens:
          return super().get_special_tokens_mask(
              token_ids_a=_A    ,token_ids_a=_A    ,already_has_special_tokens=_A      )

      if token_ids_a is not None:
          return ([0] * len(_A      )) + [1] + ([0] * len(_A      )) + [1, 1]
      return ([0] * len(_A      )) + [1, 1]

  def __lowerCamelCase		(		self    ,_A    ,_A = None      ):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    [self.sep_token_id]
      _lowerCAmelCase       :   Any  =    [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    ,_A    ,_A = None      ):
      '''simple docstring'''

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

      if os.path.abspath(self.vocab_file      ) != os.path.abspath(_A      ) and os.path.isfile(self.vocab_file      ):
          copyfile(self.vocab_file    ,_A      )
      elif not os.path.isfile(self.vocab_file      ):
          with open(_A    ,'wb'      ) as fi:
              _lowerCAmelCase       :   Optional[Any]  =    self.sp_model.serialized_model_proto()
              fi.write(_A      )

      return (out_vocab_file,)





