






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class  _lowercase   (   unittest.TestCase						):
      def 						lowercase__					(     self   ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      @property
      def 						lowercase__					(     self   ):
       snake_case__  :    Tuple							       =1
       snake_case__  :    int							       =3
       snake_case__  :    Dict							       =(3_2, 3_2)

       snake_case__  :    Tuple							       =floats_tensor((batch_size, num_channels) + sizes				,			rng=random.Random(0   )   ).to(a   )
       return image
      @property
      def 						lowercase__					(     self   ):
       torch.manual_seed(0   )
       snake_case__  :    Dict							       =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				,			)
       return model
      @property
      def 						lowercase__					(     self   ):
       torch.manual_seed(0   )
       snake_case__  :    List[str]							       =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				,			)
       return model
      @property
      def 						lowercase__					(     self   ):
       torch.manual_seed(0   )
       snake_case__  :    List[Any]							       =RobertaSeriesConfig(
           hidden_size=3_2				,			project_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=5_0_0_6				,			)
       return RobertaSeriesModelWithTransformation(a   )
      @property
      def 						lowercase__					(     self   ):
       def extract(*a				,			**a   ):







        class  _lowercase   :
              def __init__(     self   ):
               snake_case__  :    Tuple							       =torch.ones([0]   )



              def 						lowercase__					(     self				,			a   ):
               self.pixel_values.to(a   )
               return self




        return Out()

       return extract
      def 						lowercase__					(     self   ):
       snake_case__  :    int							       ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
       snake_case__  :    Any							       =self.dummy_cond_unet
       snake_case__  :    Tuple							       =PNDMScheduler(skip_prk_steps=a   )
       snake_case__  :    Optional[int]							       =self.dummy_vae
       snake_case__  :    Dict							       =self.dummy_text_encoder
       snake_case__  :    int							       =XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""   )
       snake_case__  :    Any							       =7_7

       snake_case__  :    List[str]							       =self.dummy_image.to(a   )
       snake_case__  :    List[str]							       =init_image / 2 + 0.5

       # make sure here that pndm scheduler skips prk
       snake_case__  :    str							       =AltDiffusionImgaImgPipeline(
           unet=a				,			scheduler=a				,			vae=a				,			text_encoder=a				,			tokenizer=a				,			safety_checker=a				,			feature_extractor=self.dummy_extractor				,			)
       snake_case__  :    Optional[Any]							       =VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor				,			do_normalize=a   )
       snake_case__  :    List[str]							       =alt_pipe.to(a   )
       alt_pipe.set_progress_bar_config(disable=a   )

       snake_case__  :    Optional[Any]							       ="""A painting of a squirrel eating a burger"""
       snake_case__  :    Tuple							       =torch.Generator(device=a   ).manual_seed(0   )
       snake_case__  :    str							       =alt_pipe(
           [prompt]				,			generator=a				,			guidance_scale=6.0				,			num_inference_steps=2				,			output_type="""np"""				,			image=a				,			)

       snake_case__  :    Tuple							       =output.images

       snake_case__  :    str							       =torch.Generator(device=a   ).manual_seed(0   )
       snake_case__  :    List[str]							       =alt_pipe(
           [prompt]				,			generator=a				,			guidance_scale=6.0				,			num_inference_steps=2				,			output_type="""np"""				,			image=a				,			return_dict=a				,			)[0]

       snake_case__  :    Tuple							       =image[0, -3:, -3:, -1]
       snake_case__  :    str							       =image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 3_2, 3_2, 3)
       snake_case__  :    str							       =np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]   )

       assert np.abs(image_slice.flatten() - expected_slice   ).max() < 5e-3
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice   ).max() < 5e-3
      @unittest.skipIf(torch_device != """cuda"""				,			"""This test requires a GPU"""   )
      def 						lowercase__					(     self   ):
       snake_case__  :    Union[str, Any]							       =self.dummy_cond_unet
       snake_case__  :    List[Any]							       =PNDMScheduler(skip_prk_steps=a   )
       snake_case__  :    Optional[int]							       =self.dummy_vae
       snake_case__  :    Union[str, Any]							       =self.dummy_text_encoder
       snake_case__  :    int							       =XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""   )
       snake_case__  :    Union[str, Any]							       =7_7

       snake_case__  :    Dict							       =self.dummy_image.to(a   )

       # put models in fp16
       snake_case__  :    str							       =unet.half()
       snake_case__  :    int							       =vae.half()
       snake_case__  :    Tuple							       =bert.half()

       # make sure here that pndm scheduler skips prk
       snake_case__  :    Optional[int]							       =AltDiffusionImgaImgPipeline(
           unet=a				,			scheduler=a				,			vae=a				,			text_encoder=a				,			tokenizer=a				,			safety_checker=a				,			feature_extractor=self.dummy_extractor				,			)
       snake_case__  :    Optional[int]							       =VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor				,			do_normalize=a   )
       snake_case__  :    Tuple							       =alt_pipe.to(a   )
       alt_pipe.set_progress_bar_config(disable=a   )

       snake_case__  :    str							       ="""A painting of a squirrel eating a burger"""
       snake_case__  :    List[str]							       =torch.manual_seed(0   )
       snake_case__  :    Any							       =alt_pipe(
           [prompt]				,			generator=a				,			num_inference_steps=2				,			output_type="""np"""				,			image=a				,			).images

       assert image.shape == (1, 3_2, 3_2, 3)



      @unittest.skipIf(torch_device != """cuda"""				,			"""This test requires a GPU"""   )
      def 						lowercase__					(     self   ):
       snake_case__  :    Dict							       =load_image(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
           """/img2img/sketch-mountains-input.jpg"""   )
       # resize to resolution that is divisible by 8 but not 16 or 32
       snake_case__  :    Tuple							       =init_image.resize((7_6_0, 5_0_4)   )

       snake_case__  :    Any							       ="""BAAI/AltDiffusion"""
       snake_case__  :    Optional[int]							       =AltDiffusionImgaImgPipeline.from_pretrained(
           a				,			safety_checker=a				,			)
       pipe.to(a   )
       pipe.set_progress_bar_config(disable=a   )
       pipe.enable_attention_slicing()

       snake_case__  :    Tuple							       ="""A fantasy landscape, trending on artstation"""

       snake_case__  :    Optional[Any]							       =torch.manual_seed(0   )
       snake_case__  :    Any							       =pipe(
           prompt=a				,			image=a				,			strength=0.75				,			guidance_scale=7.5				,			generator=a				,			output_type="""np"""				,			)
       snake_case__  :    Optional[int]							       =output.images[0]

       snake_case__  :    Union[str, Any]							       =image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

       assert image.shape == (5_0_4, 7_6_0, 3)
       snake_case__  :    Optional[int]							       =np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]   )

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







@slow
@require_torch_gpu
class  _lowercase   (   unittest.TestCase						):
      def 						lowercase__					(     self   ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



      def 						lowercase__					(     self   ):
       snake_case__  :    Any							       =load_image(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
           """/img2img/sketch-mountains-input.jpg"""   )
       snake_case__  :    Optional[int]							       =init_image.resize((7_6_8, 5_1_2)   )
       snake_case__  :    Dict							       =load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"""   )

       snake_case__  :    List[Any]							       ="""BAAI/AltDiffusion"""
       snake_case__  :    str							       =AltDiffusionImgaImgPipeline.from_pretrained(
           a				,			safety_checker=a				,			)
       pipe.to(a   )
       pipe.set_progress_bar_config(disable=a   )
       pipe.enable_attention_slicing()

       snake_case__  :    Tuple							       ="""A fantasy landscape, trending on artstation"""

       snake_case__  :    Optional[int]							       =torch.manual_seed(0   )
       snake_case__  :    str							       =pipe(
           prompt=a				,			image=a				,			strength=0.75				,			guidance_scale=7.5				,			generator=a				,			output_type="""np"""				,			)
       snake_case__  :    int							       =output.images[0]

       assert image.shape == (5_1_2, 7_6_8, 3)
       # img2img is flaky across GPUs even in fp32, so using MAE here
       assert np.abs(expected_image - image   ).max() < 1e-2













import os
import re
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 logging


__lowerCamelCase	:		Dict										=   logging.get_logger(__name__)

__lowerCamelCase	:		Any										=   {"""vocab_file""": """spiece.model"""}

__lowerCamelCase	:		Optional[Any]										=   {
    """vocab_file""": {
        """google/bigbird-roberta-base""": """https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model""",
        """google/bigbird-roberta-large""": (
            """https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model"""
        ),
        """google/bigbird-base-trivia-itc""": (
            """https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model"""
        ),
    }
}

__lowerCamelCase	:		List[str]										=   {
    """google/bigbird-roberta-base""": 40_96,
    """google/bigbird-roberta-large""": 40_96,
    """google/bigbird-base-trivia-itc""": 40_96,
}







class  _lowercase   (   _A						):
      _a							:						int										=							VOCAB_FILES_NAMES
      _a							:						List[str]										=							PRETRAINED_VOCAB_FILES_MAP
      _a							:						Any										=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _a							:						Tuple										=							['input_ids', 'attention_mask']
      _a							:						List[int]										=							[]
      def __init__(     self				,			a				,			a="<unk>"				,			a="<s>"				,			a="</s>"				,			a="<pad>"				,			a="[SEP]"				,			a="[MASK]"				,			a="[CLS]"				,			a = None				,			**a				,			):
       snake_case__  :    str							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else bos_token
       snake_case__  :    Optional[int]							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else eos_token
       snake_case__  :    Any							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else unk_token
       snake_case__  :    List[str]							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else pad_token
       snake_case__  :    int							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else cls_token
       snake_case__  :    Tuple							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else sep_token

       # Mask token behave like a normal word, i.e. include the space before it
       snake_case__  :    Union[str, Any]							       =AddedToken(a				,			lstrip=a				,			rstrip=a   ) if isinstance(a				,			a   ) else mask_token

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

       super().__init__(
           bos_token=a				,			eos_token=a				,			unk_token=a				,			pad_token=a				,			sep_token=a				,			mask_token=a				,			cls_token=a				,			sp_model_kwargs=self.sp_model_kwargs				,			**a				,			)

       snake_case__  :    Optional[int]							       =vocab_file

       snake_case__  :    List[str]							       =spm.SentencePieceProcessor(**self.sp_model_kwargs   )
       self.sp_model.Load(a   )
      @property
      def 						lowercase__					(     self   ):
       return self.sp_model.get_piece_size()
      def 						lowercase__					(     self   ):
       snake_case__  :    Optional[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   ):
       snake_case__  :    List[Any]							       =self.__dict__.copy()
       snake_case__  :    Optional[int]							       =None
       return state
      def __setstate__(     self				,			a   ):
       snake_case__  :    Optional[Any]							       =d

       # for backward compatibility
       if not hasattr(self				,			"""sp_model_kwargs"""   ):
        snake_case__  :    Optional[Any]							       ={}

       snake_case__  :    Tuple							       =spm.SentencePieceProcessor(**self.sp_model_kwargs   )
       self.sp_model.Load(self.vocab_file   )
      def 						lowercase__					(     self				,			a   ):
       return self.sp_model.encode(a				,			out_type=a   )
      def 						lowercase__					(     self				,			a   ):
       return self.sp_model.piece_to_id(a   )
      def 						lowercase__					(     self				,			a   ):
       snake_case__  :    Dict							       =self.sp_model.IdToPiece(a   )
       return token
      def 						lowercase__					(     self				,			a   ):
       snake_case__  :    Optional[int]							       =[]
       snake_case__  :    Dict							       =""""""
       snake_case__  :    Dict							       =False
       for token in tokens:
        # make sure that special tokens are not decoded using sentencepiece model
        if token in self.all_special_tokens:
         if not prev_is_special:
          out_string += " "
         out_string += self.sp_model.decode(a   ) + token
         snake_case__  :    Optional[Any]							       =True
         snake_case__  :    Dict							       =[]
        else:
         current_sub_tokens.append(a   )
         snake_case__  :    List[Any]							       =False
       out_string += self.sp_model.decode(a   )
       return out_string.strip()
      def 						lowercase__					(     self				,			a				,			a = False				,			a = None				,			a = True				,			**a				,			):
       snake_case__  :    List[str]							       =kwargs.pop("""use_source_tokenizer"""				,			a   )

       snake_case__  :    int							       =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
       snake_case__  :    List[Any]							       =[]
       snake_case__  :    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   )   )
          snake_case__  :    List[Any]							       =[]
         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:
       # No space before [MASK] and [SEP]
       if spaces_between_special_tokens:
        snake_case__  :    Dict							       =re.sub(R""" (\[(MASK|SEP)\])"""				,			R"""\1"""				,			""" """.join(a   )   )
       else:
        snake_case__  :    Dict							       ="""""".join(a   )

       snake_case__  :    str							       =(
           clean_up_tokenization_spaces
           if clean_up_tokenization_spaces is not None
           else self.clean_up_tokenization_spaces
       )
       if clean_up_tokenization_spaces:
        snake_case__  :    Any							       =self.clean_up_tokenization(a   )
        return clean_text
       else:
        return text
      def 						lowercase__					(     self				,			a				,			a = None   ):
       if not os.path.isdir(a   ):
        logger.error(F"Vocabulary path ({save_directory}) should be a directory"   )
        return
       snake_case__  :    Any							       =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:
         snake_case__  :    Any							       =self.sp_model.serialized_model_proto()
         fi.write(a   )

       return (out_vocab_file,)
      def 						lowercase__					(     self				,			a				,			a = None   ):
       if token_ids_a is None:
        return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
       snake_case__  :    Union[str, Any]							       =[self.cls_token_id]
       snake_case__  :    Dict							       =[self.sep_token_id]
       return cls + token_ids_a + sep + token_ids_a + sep
      def 						lowercase__					(     self				,			a				,			a = None				,			a = False   ):
       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 None:
        return [1] + ([0] * len(a   )) + [1]
       return [1] + ([0] * len(a   )) + [1] + ([0] * len(a   )) + [1]



      def 						lowercase__					(     self				,			a				,			a = None   ):
       snake_case__  :    List[str]							       =[self.sep_token_id]
       snake_case__  :    List[Any]							       =[self.cls_token_id]
       if token_ids_a is None:
        return len(cls + token_ids_a + sep   ) * [0]
       return len(cls + token_ids_a + sep   ) * [0] + len(token_ids_a + sep   ) * [1]





