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

import sentencepiece as spm

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


lowercase_     :   Optional[Any] 				=				logging.get_logger(__name__)

lowercase_     :   Tuple 				=				'▁'

lowercase_     :   Optional[Any] 				=				{'vocab_file': 'sentencepiece.bpe.model'}

lowercase_     :   str 				=				{
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    }
}

lowercase_     :   List[str] 				=				{
    'facebook/mbart-large-50-one-to-many-mmt': 1_0_2_4,
}

# fmt: off
lowercase_     :   Union[str, Any] 				=				['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']






class 		_lowerCamelCase		(   UpperCamelCase_    ):
    __a  =      VOCAB_FILES_NAMES
    __a  =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    __a  =      PRETRAINED_VOCAB_FILES_MAP
    __a  =      ["input_ids", "attention_mask"]

    __a  =      []
    __a  =      []
    def __init__(  self , lowerCAmelCase , lowerCAmelCase=None , lowerCAmelCase=None , lowerCAmelCase="</s>" , lowerCAmelCase="</s>" , lowerCAmelCase="<s>" , lowerCAmelCase="<unk>" , lowerCAmelCase="<pad>" , lowerCAmelCase="<mask>" , lowerCAmelCase = None , **lowerCAmelCase , )       ->					None:
     # Mask token behave like a normal word, i.e. include the space before it
     SCREAMING_SNAKE_CASE__:   Dict=      AddedToken(lowerCAmelCase , lstrip=lowerCAmelCase , rstrip=lowerCAmelCase   ) if isinstance(lowerCAmelCase , lowerCAmelCase   ) else mask_token

     SCREAMING_SNAKE_CASE__:   Dict=      {} if sp_model_kwargs is None else sp_model_kwargs

     SCREAMING_SNAKE_CASE__:   Dict=      kwargs.get('''additional_special_tokens''' , []   )
     kwargs["additional_special_tokens"] += [
         code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
     ]

     super().__init__(
         src_lang=lowerCAmelCase , tgt_lang=lowerCAmelCase , eos_token=lowerCAmelCase , unk_token=lowerCAmelCase , sep_token=lowerCAmelCase , cls_token=lowerCAmelCase , pad_token=lowerCAmelCase , mask_token=lowerCAmelCase , sp_model_kwargs=self.sp_model_kwargs , **lowerCAmelCase , )

     SCREAMING_SNAKE_CASE__:   Optional[int]=      spm.SentencePieceProcessor(**self.sp_model_kwargs   )
     self.sp_model.Load(str(lowerCAmelCase   )   )
     SCREAMING_SNAKE_CASE__:   int=      vocab_file

     # Original fairseq vocab and spm vocab must be "aligned":
     # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
     # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
     # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
     # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

     # Mimic fairseq token-to-id alignment for the first 4 token
     SCREAMING_SNAKE_CASE__:   List[str]=      {'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

     # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
     SCREAMING_SNAKE_CASE__:   int=      1

     SCREAMING_SNAKE_CASE__:   Dict=      len(self.sp_model   )
     SCREAMING_SNAKE_CASE__:   int=      {
         code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(lowerCAmelCase   )
     }
     SCREAMING_SNAKE_CASE__:   int=      {v: k for k, v in self.lang_code_to_id.items()}
     SCREAMING_SNAKE_CASE__:   int=      len(self.sp_model   ) + len(self.lang_code_to_id   ) + self.fairseq_offset

     self.fairseq_tokens_to_ids.update(self.lang_code_to_id   )
     SCREAMING_SNAKE_CASE__:   List[Any]=      {v: k for k, v in self.fairseq_tokens_to_ids.items()}

     SCREAMING_SNAKE_CASE__:   Optional[Any]=      src_lang if src_lang is not None else '''en_XX'''
     SCREAMING_SNAKE_CASE__:   Tuple=      self.lang_code_to_id[self._src_lang]
     SCREAMING_SNAKE_CASE__:   List[str]=      tgt_lang
     self.set_src_lang_special_tokens(self._src_lang   )
    @property
    def 						UpperCamelCase_				(  self   )       ->					int:
     return len(self.sp_model   ) + len(self.lang_code_to_id   ) + self.fairseq_offset + 1  # Plus 1 for the mask token
    @property
    def 						UpperCamelCase_				(  self   )       ->					str:
     return self._src_lang
    @src_lang.setter
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					None:
     SCREAMING_SNAKE_CASE__:   Tuple=      new_src_lang
     self.set_src_lang_special_tokens(self._src_lang   )
    def __getstate__(  self   )       ->					Dict:
     SCREAMING_SNAKE_CASE__:   Tuple=      self.__dict__.copy()
     SCREAMING_SNAKE_CASE__:   Dict=      None
     return state
    def __setstate__(  self , lowerCAmelCase   )       ->					None:
     SCREAMING_SNAKE_CASE__:   Dict=      d

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

     SCREAMING_SNAKE_CASE__:   Tuple=      spm.SentencePieceProcessor(**self.sp_model_kwargs   )
     self.sp_model.Load(self.vocab_file   )
    def 						UpperCamelCase_				(  self   )       ->					Dict:
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      {self.convert_ids_to_tokens(lowerCAmelCase   ): i for i in range(self.vocab_size   )}
     vocab.update(self.added_tokens_encoder   )
     return vocab
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					List[str]:
     return self.sp_model.encode(lowerCAmelCase , out_type=lowerCAmelCase   )
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					int:
     if token in self.fairseq_tokens_to_ids:
      return self.fairseq_tokens_to_ids[token]
     SCREAMING_SNAKE_CASE__:   int=      self.sp_model.PieceToId(lowerCAmelCase   )

     # Need to return unknown token if the SP model returned 0
     return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					str:
     if index in self.fairseq_ids_to_tokens:
      return self.fairseq_ids_to_tokens[index]
     return self.sp_model.IdToPiece(index - self.fairseq_offset   )
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					List[Any]:
     SCREAMING_SNAKE_CASE__:   Any=      []
     SCREAMING_SNAKE_CASE__:   Any=      ''''''
     SCREAMING_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(lowerCAmelCase   ) + token
       SCREAMING_SNAKE_CASE__:   Tuple=      True
       SCREAMING_SNAKE_CASE__:   Optional[int]=      []
      else:
       current_sub_tokens.append(lowerCAmelCase   )
       SCREAMING_SNAKE_CASE__:   Tuple=      False
     out_string += self.sp_model.decode(lowerCAmelCase   )
     return out_string.strip()
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase = None   )       ->					Tuple[str]:
     if not os.path.isdir(lowerCAmelCase   ):
      logger.error(f'Vocabulary path ({save_directory}) should be a directory'   )
      return
     SCREAMING_SNAKE_CASE__:   Dict=      os.path.join(
         lowerCAmelCase , (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']   )

     if os.path.abspath(self.vocab_file   ) != os.path.abspath(lowerCAmelCase   ) and os.path.isfile(self.vocab_file   ):
      copyfile(self.vocab_file , lowerCAmelCase   )
     elif not os.path.isfile(self.vocab_file   ):
      with open(lowerCAmelCase , '''wb'''   ) as fi:
       SCREAMING_SNAKE_CASE__:   Tuple=      self.sp_model.serialized_model_proto()
       fi.write(lowerCAmelCase   )

     return (out_vocab_file,)
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase = None , lowerCAmelCase = False   )       ->					List[int]:

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

     SCREAMING_SNAKE_CASE__:   Optional[Any]=      [1] * len(self.prefix_tokens   )
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      [1] * len(self.suffix_tokens   )
     if token_ids_a is None:
      return prefix_ones + ([0] * len(lowerCAmelCase   )) + suffix_ones
     return prefix_ones + ([0] * len(lowerCAmelCase   )) + ([0] * len(lowerCAmelCase   )) + suffix_ones
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase = 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 , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , **lowerCAmelCase   )       ->					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'''   )
     SCREAMING_SNAKE_CASE__:   Dict=      src_lang
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      self(lowerCAmelCase , add_special_tokens=lowerCAmelCase , return_tensors=lowerCAmelCase , **lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      self.convert_tokens_to_ids(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      tgt_lang_id
     return inputs
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase = "en_XX" , lowerCAmelCase = None , lowerCAmelCase = "ro_RO" , **lowerCAmelCase , )       ->					BatchEncoding:
     SCREAMING_SNAKE_CASE__:   str=      src_lang
     SCREAMING_SNAKE_CASE__:   int=      tgt_lang
     return super().prepare_seqaseq_batch(lowerCAmelCase , lowerCAmelCase , **lowerCAmelCase   )
    def 						UpperCamelCase_				(  self   )       ->					Dict:
     return self.set_src_lang_special_tokens(self.src_lang   )
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     return self.set_tgt_lang_special_tokens(self.tgt_lang   )
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					None:
     SCREAMING_SNAKE_CASE__:   List[str]=      self.lang_code_to_id[src_lang]
     SCREAMING_SNAKE_CASE__:   Dict=      [self.cur_lang_code_id]
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      [self.eos_token_id]



    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					None:
     SCREAMING_SNAKE_CASE__:   Any=      self.lang_code_to_id[tgt_lang]
     SCREAMING_SNAKE_CASE__:   Tuple=      [self.cur_lang_code_id]
     SCREAMING_SNAKE_CASE__:   List[Any]=      [self.eos_token_id]


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 		_lowerCamelCase		(   UpperCamelCase_    ,     UpperCamelCase_    ,     UpperCamelCase_    ,     unittest.TestCase    ):
    __a  =      AltDiffusionPipeline
    __a  =      TEXT_TO_IMAGE_PARAMS
    __a  =      TEXT_TO_IMAGE_BATCH_PARAMS
    __a  =      TEXT_TO_IMAGE_IMAGE_PARAMS
    __a  =      TEXT_TO_IMAGE_IMAGE_PARAMS
    def 						UpperCamelCase_				(  self   )       ->					Optional[Any]:
     torch.manual_seed(0   )
     SCREAMING_SNAKE_CASE__:   List[Any]=      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 , )
     SCREAMING_SNAKE_CASE__:   Tuple=      DDIMScheduler(
         beta_start=0.00085 , beta_end=0.012 , beta_schedule='''scaled_linear''' , clip_sample=lowerCAmelCase , set_alpha_to_one=lowerCAmelCase , )
     torch.manual_seed(0   )
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      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   )
     SCREAMING_SNAKE_CASE__:   Tuple=      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 , )
     SCREAMING_SNAKE_CASE__:   List[Any]=      CLIPTextModel(lowerCAmelCase   )

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

     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      {
         '''unet''': unet,
         '''scheduler''': scheduler,
         '''vae''': vae,
         '''text_encoder''': text_encoder,
         '''tokenizer''': tokenizer,
         '''safety_checker''': None,
         '''feature_extractor''': None,
     }
     return components
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase=0   )       ->					List[Any]:
     if str(lowerCAmelCase   ).startswith('''mps'''   ):
      SCREAMING_SNAKE_CASE__:   Any=      torch.manual_seed(lowerCAmelCase   )
     else:
      SCREAMING_SNAKE_CASE__:   Union[str, Any]=      torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Any=      {
         '''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 						UpperCamelCase_				(  self   )       ->					List[str]:
     super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3   )
    def 						UpperCamelCase_				(  self   )       ->					Optional[Any]:
     SCREAMING_SNAKE_CASE__:   Optional[int]=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

     SCREAMING_SNAKE_CASE__:   Tuple=      self.get_dummy_components()
     torch.manual_seed(0   )
     SCREAMING_SNAKE_CASE__:   List[Any]=      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
     SCREAMING_SNAKE_CASE__:   List[str]=      RobertaSeriesModelWithTransformation(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   List[str]=      text_encoder

     SCREAMING_SNAKE_CASE__:   List[Any]=      AltDiffusionPipeline(**lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      alt_pipe.to(lowerCAmelCase   )
     alt_pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

     assert image.shape == (1, 64, 64, 3)
     SCREAMING_SNAKE_CASE__:   Any=      np.array(
         [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]   )

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



    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     SCREAMING_SNAKE_CASE__:   List[str]=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

     SCREAMING_SNAKE_CASE__:   Any=      self.get_dummy_components()
     SCREAMING_SNAKE_CASE__:   List[str]=      PNDMScheduler(skip_prk_steps=lowerCAmelCase   )
     torch.manual_seed(0   )
     SCREAMING_SNAKE_CASE__:   Tuple=      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
     SCREAMING_SNAKE_CASE__:   str=      RobertaSeriesModelWithTransformation(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   List[str]=      text_encoder
     SCREAMING_SNAKE_CASE__:   Any=      AltDiffusionPipeline(**lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   List[str]=      alt_pipe.to(lowerCAmelCase   )
     alt_pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

     assert image.shape == (1, 64, 64, 3)
     SCREAMING_SNAKE_CASE__:   List[Any]=      np.array(
         [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]   )

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






@slow
@require_torch_gpu
class 		_lowerCamelCase		(   unittest.TestCase    ):
    def 						UpperCamelCase_				(  self   )       ->					Union[str, Any]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()
    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     # make sure here that pndm scheduler skips prk
     SCREAMING_SNAKE_CASE__:   Any=      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion''' , safety_checker=lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      alt_pipe.to(lowerCAmelCase   )
     alt_pipe.set_progress_bar_config(disable=lowerCAmelCase   )

     SCREAMING_SNAKE_CASE__:   Dict=      '''A painting of a squirrel eating a burger'''
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      torch.manual_seed(0   )
     SCREAMING_SNAKE_CASE__:   Tuple=      alt_pipe([prompt] , generator=lowerCAmelCase , guidance_scale=6.0 , num_inference_steps=20 , output_type='''np'''   )

     SCREAMING_SNAKE_CASE__:   List[Any]=      output.images

     SCREAMING_SNAKE_CASE__:   Dict=      image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     SCREAMING_SNAKE_CASE__:   Optional[int]=      np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]   )

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



    def 						UpperCamelCase_				(  self   )       ->					Union[str, Any]:
     SCREAMING_SNAKE_CASE__:   List[Any]=      DDIMScheduler.from_pretrained('''BAAI/AltDiffusion''' , subfolder='''scheduler'''   )

     SCREAMING_SNAKE_CASE__:   List[str]=      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion''' , scheduler=lowerCAmelCase , safety_checker=lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      alt_pipe.to(lowerCAmelCase   )
     alt_pipe.set_progress_bar_config(disable=lowerCAmelCase   )

     SCREAMING_SNAKE_CASE__:   Optional[Any]=      '''A painting of a squirrel eating a burger'''
     SCREAMING_SNAKE_CASE__:   Any=      torch.manual_seed(0   )

     SCREAMING_SNAKE_CASE__:   str=      alt_pipe([prompt] , generator=lowerCAmelCase , num_inference_steps=2 , output_type='''numpy'''   )
     SCREAMING_SNAKE_CASE__:   List[str]=      output.images

     SCREAMING_SNAKE_CASE__:   Dict=      image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]   )

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

