



"""simple docstring"""


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

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
    from .tokenization_nllb import NllbTokenizer
else:
    snake_case_    					=	None


snake_case_    					=	logging.get_logger(__name__)


snake_case_    					=	{"""vocab_file""": """sentencepiece.bpe.model""", """tokenizer_file""": """tokenizer.json"""}

snake_case_    					=	{
    """vocab_file""": {
        """facebook/nllb-200-distilled-600M""": (
            """https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/sentencepiece.bpe.model"""
        ),
    },
    """tokenizer_file""": {
        """facebook/nllb-200-distilled-600M""": (
            """https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/tokenizer.json"""
        ),
    },
}

snake_case_    					=	{
    """facebook/nllb-large-en-ro""": 1024,
    """facebook/nllb-200-distilled-600M""": 1024,
}

# fmt: off
snake_case_    					=	["""ace_Arab""", """ace_Latn""", """acm_Arab""", """acq_Arab""", """aeb_Arab""", """afr_Latn""", """ajp_Arab""", """aka_Latn""", """amh_Ethi""", """apc_Arab""", """arb_Arab""", """ars_Arab""", """ary_Arab""", """arz_Arab""", """asm_Beng""", """ast_Latn""", """awa_Deva""", """ayr_Latn""", """azb_Arab""", """azj_Latn""", """bak_Cyrl""", """bam_Latn""", """ban_Latn""", """bel_Cyrl""", """bem_Latn""", """ben_Beng""", """bho_Deva""", """bjn_Arab""", """bjn_Latn""", """bod_Tibt""", """bos_Latn""", """bug_Latn""", """bul_Cyrl""", """cat_Latn""", """ceb_Latn""", """ces_Latn""", """cjk_Latn""", """ckb_Arab""", """crh_Latn""", """cym_Latn""", """dan_Latn""", """deu_Latn""", """dik_Latn""", """dyu_Latn""", """dzo_Tibt""", """ell_Grek""", """eng_Latn""", """epo_Latn""", """est_Latn""", """eus_Latn""", """ewe_Latn""", """fao_Latn""", """pes_Arab""", """fij_Latn""", """fin_Latn""", """fon_Latn""", """fra_Latn""", """fur_Latn""", """fuv_Latn""", """gla_Latn""", """gle_Latn""", """glg_Latn""", """grn_Latn""", """guj_Gujr""", """hat_Latn""", """hau_Latn""", """heb_Hebr""", """hin_Deva""", """hne_Deva""", """hrv_Latn""", """hun_Latn""", """hye_Armn""", """ibo_Latn""", """ilo_Latn""", """ind_Latn""", """isl_Latn""", """ita_Latn""", """jav_Latn""", """jpn_Jpan""", """kab_Latn""", """kac_Latn""", """kam_Latn""", """kan_Knda""", """kas_Arab""", """kas_Deva""", """kat_Geor""", """knc_Arab""", """knc_Latn""", """kaz_Cyrl""", """kbp_Latn""", """kea_Latn""", """khm_Khmr""", """kik_Latn""", """kin_Latn""", """kir_Cyrl""", """kmb_Latn""", """kon_Latn""", """kor_Hang""", """kmr_Latn""", """lao_Laoo""", """lvs_Latn""", """lij_Latn""", """lim_Latn""", """lin_Latn""", """lit_Latn""", """lmo_Latn""", """ltg_Latn""", """ltz_Latn""", """lua_Latn""", """lug_Latn""", """luo_Latn""", """lus_Latn""", """mag_Deva""", """mai_Deva""", """mal_Mlym""", """mar_Deva""", """min_Latn""", """mkd_Cyrl""", """plt_Latn""", """mlt_Latn""", """mni_Beng""", """khk_Cyrl""", """mos_Latn""", """mri_Latn""", """zsm_Latn""", """mya_Mymr""", """nld_Latn""", """nno_Latn""", """nob_Latn""", """npi_Deva""", """nso_Latn""", """nus_Latn""", """nya_Latn""", """oci_Latn""", """gaz_Latn""", """ory_Orya""", """pag_Latn""", """pan_Guru""", """pap_Latn""", """pol_Latn""", """por_Latn""", """prs_Arab""", """pbt_Arab""", """quy_Latn""", """ron_Latn""", """run_Latn""", """rus_Cyrl""", """sag_Latn""", """san_Deva""", """sat_Beng""", """scn_Latn""", """shn_Mymr""", """sin_Sinh""", """slk_Latn""", """slv_Latn""", """smo_Latn""", """sna_Latn""", """snd_Arab""", """som_Latn""", """sot_Latn""", """spa_Latn""", """als_Latn""", """srd_Latn""", """srp_Cyrl""", """ssw_Latn""", """sun_Latn""", """swe_Latn""", """swh_Latn""", """szl_Latn""", """tam_Taml""", """tat_Cyrl""", """tel_Telu""", """tgk_Cyrl""", """tgl_Latn""", """tha_Thai""", """tir_Ethi""", """taq_Latn""", """taq_Tfng""", """tpi_Latn""", """tsn_Latn""", """tso_Latn""", """tuk_Latn""", """tum_Latn""", """tur_Latn""", """twi_Latn""", """tzm_Tfng""", """uig_Arab""", """ukr_Cyrl""", """umb_Latn""", """urd_Arab""", """uzn_Latn""", """vec_Latn""", """vie_Latn""", """war_Latn""", """wol_Latn""", """xho_Latn""", """ydd_Hebr""", """yor_Latn""", """yue_Hant""", """zho_Hans""", """zho_Hant""", """zul_Latn"""]







class 				A_      (					SCREAMING_SNAKE_CASE_    ):



      """simple docstring"""
      __UpperCamelCase		    =							VOCAB_FILES_NAMES
      __UpperCamelCase		    =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __UpperCamelCase		    =							PRETRAINED_VOCAB_FILES_MAP
      __UpperCamelCase		    =							["""input_ids""", """attention_mask"""]
      __UpperCamelCase		    =							NllbTokenizer

      __UpperCamelCase		    =							[]
      __UpperCamelCase		    =							[]






      def __init__( self			:str       ,  lowercase_			:Dict=None       ,  lowercase_			:List[Any]=None       ,  lowercase_			:Optional[int]="<s>"       ,  lowercase_			:List[str]="</s>"       ,  lowercase_			:List[Any]="</s>"       ,  lowercase_			:Dict="<s>"       ,  lowercase_			:List[str]="<unk>"       ,  lowercase_			:Tuple="<pad>"       ,  lowercase_			:List[Any]="<mask>"       ,  lowercase_			:Optional[Any]=None       ,  lowercase_			:Any=None       ,  lowercase_			:List[str]=None       ,  lowercase_			:str=False       ,  **lowercase_			:str       ,  )    ->      Any:
            # Mask token behave like a normal word, i.e. include the space before it
            UpperCAmelCase					=			AddedToken(lowercase_       ,  lstrip=lowercase_       ,  rstrip=lowercase_ ) if isinstance(lowercase_       ,  lowercase_ ) else mask_token
            UpperCAmelCase					=			legacy_behaviour
            super().__init__(
                vocab_file=lowercase_       ,  tokenizer_file=lowercase_       ,  bos_token=lowercase_       ,  eos_token=lowercase_       ,  sep_token=lowercase_       ,  cls_token=lowercase_       ,  unk_token=lowercase_       ,  pad_token=lowercase_       ,  mask_token=lowercase_       ,  src_lang=lowercase_       ,  tgt_lang=lowercase_       ,  additional_special_tokens=lowercase_       ,  legacy_behaviour=lowercase_       ,  **lowercase_       ,  )

            UpperCAmelCase					=			vocab_file
            UpperCAmelCase					=			False if not self.vocab_file else True

            UpperCAmelCase					=			FAIRSEQ_LANGUAGE_CODES.copy()

            if additional_special_tokens is not None:
                  # Only add those special tokens if they are not already there.
                  _additional_special_tokens.extend(
                      [t for t in additional_special_tokens if t not in _additional_special_tokens] )

            self.add_special_tokens({'additional_special_tokens': _additional_special_tokens} )
            UpperCAmelCase					=			{
                lang_code: self.convert_tokens_to_ids(lowercase_ ) for lang_code in FAIRSEQ_LANGUAGE_CODES
            }

            UpperCAmelCase					=			src_lang if src_lang is not None else 'eng_Latn'
            UpperCAmelCase					=			self.convert_tokens_to_ids(self._src_lang )
            UpperCAmelCase					=			tgt_lang
            self.set_src_lang_special_tokens(self._src_lang )






      @property
      def      UpperCAmelCase__					( self			:int )    ->      str:
            return self._src_lang






      @src_lang.setter
      def      UpperCAmelCase__					( self			:int       ,  lowercase_			:str )    ->      None:
            UpperCAmelCase					=			new_src_lang
            self.set_src_lang_special_tokens(self._src_lang )






      def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:List[int]       ,  lowercase_			:Optional[List[int]] = None )    ->      List[int]:
            if token_ids_a is None:
                  return self.prefix_tokens + token_ids_a + self.suffix_tokens
            # We don't expect to process pairs, but leave the pair logic for API consistency
            return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens






      def      UpperCAmelCase__					( self			:Dict       ,  lowercase_			:List[int]       ,  lowercase_			:Optional[List[int]] = None )    ->      List[int]:
            UpperCAmelCase					=			[self.sep_token_id]
            UpperCAmelCase					=			[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 + sep + token_ids_a + sep ) * [0]






      def      UpperCAmelCase__					( self			:Any       ,  lowercase_			:List[str]       ,  lowercase_			:str       ,  lowercase_			:Optional[str]       ,  lowercase_			:Optional[str]       ,  **lowercase_			:Any )    ->      Dict:
            if src_lang is None or tgt_lang is None:
                  raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model' )
            UpperCAmelCase					=			src_lang
            UpperCAmelCase					=			self(lowercase_       ,  add_special_tokens=lowercase_       ,  return_tensors=lowercase_       ,  **lowercase_ )
            UpperCAmelCase					=			self.convert_tokens_to_ids(lowercase_ )
            UpperCAmelCase					=			tgt_lang_id
            return inputs






      def      UpperCAmelCase__					( self			:Any       ,  lowercase_			:List[str]       ,  lowercase_			:str = "eng_Latn"       ,  lowercase_			:Optional[List[str]] = None       ,  lowercase_			:str = "fra_Latn"       ,  **lowercase_			:Tuple       ,  )    ->      BatchEncoding:
            UpperCAmelCase					=			src_lang
            UpperCAmelCase					=			tgt_lang
            return super().prepare_seqaseq_batch(lowercase_       ,  lowercase_       ,  **lowercase_ )






      def      UpperCAmelCase__					( self			:List[str] )    ->      Optional[int]:
            return self.set_src_lang_special_tokens(self.src_lang )






      def      UpperCAmelCase__					( self			:List[str] )    ->      List[str]:
            return self.set_tgt_lang_special_tokens(self.tgt_lang )






      def      UpperCAmelCase__					( self			:Optional[Any]       ,  lowercase_			:List[str] )    ->      None:
            UpperCAmelCase					=			self.convert_tokens_to_ids(lowercase_ )

            if self.legacy_behaviour:
                  UpperCAmelCase					=			[]
                  UpperCAmelCase					=			[self.eos_token_id, self.cur_lang_code]
            else:
                  UpperCAmelCase					=			[self.cur_lang_code]
                  UpperCAmelCase					=			[self.eos_token_id]

            UpperCAmelCase					=			self.convert_ids_to_tokens(self.prefix_tokens )
            UpperCAmelCase					=			self.convert_ids_to_tokens(self.suffix_tokens )

            UpperCAmelCase					=			processors.TemplateProcessing(
                single=prefix_tokens_str + ['$A'] + suffix_tokens_str       ,  pair=prefix_tokens_str + ['$A', '$B'] + suffix_tokens_str       ,  special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str       ,  self.prefix_tokens + self.suffix_tokens ) )       ,  )






      def      UpperCAmelCase__					( self			:Any       ,  lowercase_			:str )    ->      None:
            UpperCAmelCase					=			self.convert_tokens_to_ids(lowercase_ )
            if self.legacy_behaviour:
                  UpperCAmelCase					=			[]
                  UpperCAmelCase					=			[self.eos_token_id, self.cur_lang_code]
            else:
                  UpperCAmelCase					=			[self.cur_lang_code]
                  UpperCAmelCase					=			[self.eos_token_id]

            UpperCAmelCase					=			self.convert_ids_to_tokens(self.prefix_tokens )
            UpperCAmelCase					=			self.convert_ids_to_tokens(self.suffix_tokens )

            UpperCAmelCase					=			processors.TemplateProcessing(
                single=prefix_tokens_str + ['$A'] + suffix_tokens_str       ,  pair=prefix_tokens_str + ['$A', '$B'] + suffix_tokens_str       ,  special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str       ,  self.prefix_tokens + self.suffix_tokens ) )       ,  )






      def      UpperCAmelCase__					( self			:List[str]       ,  lowercase_			:str       ,  lowercase_			:Optional[str] = None )    ->      Tuple[str]:
            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(lowercase_ ):
                  logger.error(f"""Vocabulary path ({save_directory}) should be a directory.""" )
                  return
            UpperCAmelCase					=			os.path.join(
                lowercase_       ,  (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'] )

            if os.path.abspath(self.vocab_file ) != os.path.abspath(lowercase_ ):
                  copyfile(self.vocab_file       ,  lowercase_ )

            return (out_vocab_file,)










"""simple docstring"""


import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
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 				A_      (					unittest.TestCase    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Tuple )    ->      Union[str, Any]:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						@property
						def      UpperCAmelCase__					( self			:int )    ->      Union[str, Any]:
												UpperCAmelCase					=			1
												UpperCAmelCase					=			3
												UpperCAmelCase					=			(32, 32)

												UpperCAmelCase					=			floats_tensor((batch_size, num_channels) + sizes       ,  rng=random.Random(0 ) ).to(lowercase_ )
												return image






						@property
						def      UpperCAmelCase__					( self			:List[Any] )    ->      Tuple:
												torch.manual_seed(0 )
												UpperCAmelCase					=			UNetaDConditionModel(
												    block_out_channels=(32, 32, 64)       ,  layers_per_block=2       ,  sample_size=32       ,  in_channels=7       ,  out_channels=4       ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D')       ,  up_block_types=('CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'UpBlock2D')       ,  cross_attention_dim=32       ,  attention_head_dim=8       ,  use_linear_projection=lowercase_       ,  only_cross_attention=(True, True, False)       ,  num_class_embeds=1_00       ,  )
												return model






						@property
						def      UpperCAmelCase__					( self			:str )    ->      Any:
												torch.manual_seed(0 )
												UpperCAmelCase					=			AutoencoderKL(
												    block_out_channels=[32, 32, 64]       ,  in_channels=3       ,  out_channels=3       ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D']       ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']       ,  latent_channels=4       ,  )
												return model






						@property
						def      UpperCAmelCase__					( self			:Optional[int] )    ->      List[str]:
												torch.manual_seed(0 )
												UpperCAmelCase					=			CLIPTextConfig(
												    bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=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       ,  hidden_act='gelu'       ,  projection_dim=5_12       ,  )
												return CLIPTextModel(lowercase_ )






						def      UpperCAmelCase__					( self			:str )    ->      Union[str, Any]:
												UpperCAmelCase					=			'cpu'  # ensure determinism for the device-dependent torch.Generator
												UpperCAmelCase					=			self.dummy_cond_unet_upscale
												UpperCAmelCase					=			DDPMScheduler()
												UpperCAmelCase					=			DDIMScheduler(prediction_type='v_prediction' )
												UpperCAmelCase					=			self.dummy_vae
												UpperCAmelCase					=			self.dummy_text_encoder
												UpperCAmelCase					=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

												UpperCAmelCase					=			self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1 )[0]
												UpperCAmelCase					=			Image.fromarray(np.uinta(lowercase_ ) ).convert('RGB' ).resize((64, 64) )

												# make sure here that pndm scheduler skips prk
												UpperCAmelCase					=			StableDiffusionUpscalePipeline(
												    unet=lowercase_       ,  low_res_scheduler=lowercase_       ,  scheduler=lowercase_       ,  vae=lowercase_       ,  text_encoder=lowercase_       ,  tokenizer=lowercase_       ,  max_noise_level=3_50       ,  )
												UpperCAmelCase					=			sd_pipe.to(lowercase_ )
												sd_pipe.set_progress_bar_config(disable=lowercase_ )

												UpperCAmelCase					=			'A painting of a squirrel eating a burger'
												UpperCAmelCase					=			torch.Generator(device=lowercase_ ).manual_seed(0 )
												UpperCAmelCase					=			sd_pipe(
												    [prompt]       ,  image=lowercase_       ,  generator=lowercase_       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  )

												UpperCAmelCase					=			output.images

												UpperCAmelCase					=			torch.Generator(device=lowercase_ ).manual_seed(0 )
												UpperCAmelCase					=			sd_pipe(
												    [prompt]       ,  image=lowercase_       ,  generator=lowercase_       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  return_dict=lowercase_       ,  )[0]

												UpperCAmelCase					=			image[0, -3:, -3:, -1]
												UpperCAmelCase					=			image_from_tuple[0, -3:, -3:, -1]

												UpperCAmelCase					=			low_res_image.size[0] * 4
												assert image.shape == (1, expected_height_width, expected_height_width, 3)
												UpperCAmelCase					=			np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661] )

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






						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      Any:
												UpperCAmelCase					=			'cpu'  # ensure determinism for the device-dependent torch.Generator
												UpperCAmelCase					=			self.dummy_cond_unet_upscale
												UpperCAmelCase					=			DDPMScheduler()
												UpperCAmelCase					=			DDIMScheduler(prediction_type='v_prediction' )
												UpperCAmelCase					=			self.dummy_vae
												UpperCAmelCase					=			self.dummy_text_encoder
												UpperCAmelCase					=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

												UpperCAmelCase					=			self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1 )[0]
												UpperCAmelCase					=			Image.fromarray(np.uinta(lowercase_ ) ).convert('RGB' ).resize((64, 64) )

												# make sure here that pndm scheduler skips prk
												UpperCAmelCase					=			StableDiffusionUpscalePipeline(
												    unet=lowercase_       ,  low_res_scheduler=lowercase_       ,  scheduler=lowercase_       ,  vae=lowercase_       ,  text_encoder=lowercase_       ,  tokenizer=lowercase_       ,  max_noise_level=3_50       ,  )
												UpperCAmelCase					=			sd_pipe.to(lowercase_ )
												sd_pipe.set_progress_bar_config(disable=lowercase_ )

												UpperCAmelCase					=			'A painting of a squirrel eating a burger'
												UpperCAmelCase					=			sd_pipe(
												    2 * [prompt]       ,  image=2 * [low_res_image]       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  )
												UpperCAmelCase					=			output.images
												assert image.shape[0] == 2

												UpperCAmelCase					=			torch.Generator(device=lowercase_ ).manual_seed(0 )
												UpperCAmelCase					=			sd_pipe(
												    [prompt]       ,  image=lowercase_       ,  generator=lowercase_       ,  num_images_per_prompt=2       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  )
												UpperCAmelCase					=			output.images
												assert image.shape[0] == 2






						@unittest.skipIf(torch_device != 'cuda'       ,  'This test requires a GPU' )
						def      UpperCAmelCase__					( self			:Optional[int] )    ->      Optional[Any]:
												UpperCAmelCase					=			self.dummy_cond_unet_upscale
												UpperCAmelCase					=			DDPMScheduler()
												UpperCAmelCase					=			DDIMScheduler(prediction_type='v_prediction' )
												UpperCAmelCase					=			self.dummy_vae
												UpperCAmelCase					=			self.dummy_text_encoder
												UpperCAmelCase					=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

												UpperCAmelCase					=			self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1 )[0]
												UpperCAmelCase					=			Image.fromarray(np.uinta(lowercase_ ) ).convert('RGB' ).resize((64, 64) )

												# put models in fp16, except vae as it overflows in fp16
												UpperCAmelCase					=			unet.half()
												UpperCAmelCase					=			text_encoder.half()

												# make sure here that pndm scheduler skips prk
												UpperCAmelCase					=			StableDiffusionUpscalePipeline(
												    unet=lowercase_       ,  low_res_scheduler=lowercase_       ,  scheduler=lowercase_       ,  vae=lowercase_       ,  text_encoder=lowercase_       ,  tokenizer=lowercase_       ,  max_noise_level=3_50       ,  )
												UpperCAmelCase					=			sd_pipe.to(lowercase_ )
												sd_pipe.set_progress_bar_config(disable=lowercase_ )

												UpperCAmelCase					=			'A painting of a squirrel eating a burger'
												UpperCAmelCase					=			torch.manual_seed(0 )
												UpperCAmelCase					=			sd_pipe(
												    [prompt]       ,  image=lowercase_       ,  generator=lowercase_       ,  num_inference_steps=2       ,  output_type='np'       ,  ).images

												UpperCAmelCase					=			low_res_image.size[0] * 4
												assert image.shape == (1, expected_height_width, expected_height_width, 3)







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



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Dict )    ->      str:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						def      UpperCAmelCase__					( self			:Optional[int] )    ->      int:
												UpperCAmelCase					=			load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/sd2-upscale/low_res_cat.png' )
												UpperCAmelCase					=			load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
												    '/upsampled_cat.npy' )

												UpperCAmelCase					=			'stabilityai/stable-diffusion-x4-upscaler'
												UpperCAmelCase					=			StableDiffusionUpscalePipeline.from_pretrained(lowercase_ )
												pipe.to(lowercase_ )
												pipe.set_progress_bar_config(disable=lowercase_ )
												pipe.enable_attention_slicing()

												UpperCAmelCase					=			'a cat sitting on a park bench'

												UpperCAmelCase					=			torch.manual_seed(0 )
												UpperCAmelCase					=			pipe(
												    prompt=lowercase_       ,  image=lowercase_       ,  generator=lowercase_       ,  output_type='np'       ,  )
												UpperCAmelCase					=			output.images[0]

												assert image.shape == (5_12, 5_12, 3)
												assert np.abs(expected_image - image ).max() < 1E-3






						def      UpperCAmelCase__					( self			:Union[str, Any] )    ->      Tuple:
												UpperCAmelCase					=			load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/sd2-upscale/low_res_cat.png' )
												UpperCAmelCase					=			load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
												    '/upsampled_cat_fp16.npy' )

												UpperCAmelCase					=			'stabilityai/stable-diffusion-x4-upscaler'
												UpperCAmelCase					=			StableDiffusionUpscalePipeline.from_pretrained(
												    lowercase_       ,  torch_dtype=torch.floataa       ,  )
												pipe.to(lowercase_ )
												pipe.set_progress_bar_config(disable=lowercase_ )
												pipe.enable_attention_slicing()

												UpperCAmelCase					=			'a cat sitting on a park bench'

												UpperCAmelCase					=			torch.manual_seed(0 )
												UpperCAmelCase					=			pipe(
												    prompt=lowercase_       ,  image=lowercase_       ,  generator=lowercase_       ,  output_type='np'       ,  )
												UpperCAmelCase					=			output.images[0]

												assert image.shape == (5_12, 5_12, 3)
												assert np.abs(expected_image - image ).max() < 5E-1






						def      UpperCAmelCase__					( self			:Any )    ->      List[Any]:
												torch.cuda.empty_cache()
												torch.cuda.reset_max_memory_allocated()
												torch.cuda.reset_peak_memory_stats()

												UpperCAmelCase					=			load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/sd2-upscale/low_res_cat.png' )

												UpperCAmelCase					=			'stabilityai/stable-diffusion-x4-upscaler'
												UpperCAmelCase					=			StableDiffusionUpscalePipeline.from_pretrained(
												    lowercase_       ,  torch_dtype=torch.floataa       ,  )
												pipe.to(lowercase_ )
												pipe.set_progress_bar_config(disable=lowercase_ )
												pipe.enable_attention_slicing(1 )
												pipe.enable_sequential_cpu_offload()

												UpperCAmelCase					=			'a cat sitting on a park bench'

												UpperCAmelCase					=			torch.manual_seed(0 )
												UpperCAmelCase					=			pipe(
												    prompt=lowercase_       ,  image=lowercase_       ,  generator=lowercase_       ,  num_inference_steps=5       ,  output_type='np'       ,  )

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





