


"""simple docstring"""


import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
   import onnxruntime as ort

class 				_SCREAMING_SNAKE_CASE				(							__UpperCAmelCase     ,   unittest.TestCase  ):
   """simple docstring"""



   _a   :       Dict									=    '''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''


   def 					UpperCAmelCase__(						self ,     lowerCamelCase__=0						)      ->							Tuple:
    lowercase__       :       int				=						floats_tensor((1, 3, 128, 128) ,     rng=random.Random(lowerCamelCase__						)						)
    lowercase__       :       Any				=						np.random.RandomState(lowerCamelCase__						)
    lowercase__       :       str				=						{
        """prompt""": """A painting of a squirrel eating a burger""",
        """image""": image,
        """generator""": generator,
        """num_inference_steps""": 3,
        """strength""": 0.75,
        """guidance_scale""": 7.5,
        """output_type""": """numpy""",
    }
    return inputs


   def 					UpperCAmelCase__(						self						)      ->							Optional[Any]:
    lowercase__       :       List[Any]				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint ,     provider="""CPUExecutionProvider"""						)
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

    lowercase__       :       Optional[Any]				=						self.get_dummy_inputs()
    lowercase__       :       str				=						pipe(**lowerCamelCase__						).images
    lowercase__       :       int				=						image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 128, 128, 3)
    lowercase__       :       Any				=						np.array([0.6_9643, 0.5_8484, 0.5_0314, 0.5_8760, 0.5_5368, 0.5_9643, 0.5_1529, 0.4_1217, 0.4_9087]						)
    assert np.abs(image_slice - expected_slice						).max() < 1E-1


   def 					UpperCAmelCase__(						self						)      ->							Optional[int]:
    lowercase__       :       int				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint ,     provider="""CPUExecutionProvider"""						)
    lowercase__       :       Dict				=						PNDMScheduler.from_config(pipe.scheduler.config ,     skip_prk_steps=lowerCamelCase__						)
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

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

    assert image.shape == (1, 128, 128, 3)
    lowercase__       :       Dict				=						np.array([0.6_1737, 0.5_4642, 0.5_3183, 0.5_4465, 0.5_2742, 0.6_0525, 0.4_9969, 0.4_0655, 0.4_8154]						)

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


   def 					UpperCAmelCase__(						self						)      ->							Tuple:
    lowercase__       :       Union[str, Any]				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint ,     provider="""CPUExecutionProvider"""						)
    lowercase__       :       int				=						LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

    # warmup pass to apply optimizations
    lowercase__       :       Optional[int]				=						pipe(**self.get_dummy_inputs()						)

    lowercase__       :       Optional[int]				=						self.get_dummy_inputs()
    lowercase__       :       Dict				=						pipe(**lowerCamelCase__						).images
    lowercase__       :       Optional[int]				=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 128, 128, 3)
    lowercase__       :       str				=						np.array([0.5_2761, 0.5_9977, 0.4_9033, 0.4_9619, 0.5_4282, 0.5_0311, 0.4_7600, 0.4_0918, 0.4_5203]						)

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


   def 					UpperCAmelCase__(						self						)      ->							str:
    lowercase__       :       List[str]				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint ,     provider="""CPUExecutionProvider"""						)
    lowercase__       :       int				=						EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

    lowercase__       :       Tuple				=						self.get_dummy_inputs()
    lowercase__       :       Union[str, Any]				=						pipe(**lowerCamelCase__						).images
    lowercase__       :       int				=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 128, 128, 3)
    lowercase__       :       List[Any]				=						np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]						)

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


   def 					UpperCAmelCase__(						self						)      ->							Optional[int]:
    lowercase__       :       Any				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint ,     provider="""CPUExecutionProvider"""						)
    lowercase__       :       Optional[int]				=						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

    lowercase__       :       Tuple				=						self.get_dummy_inputs()
    lowercase__       :       Dict				=						pipe(**lowerCamelCase__						).images
    lowercase__       :       Dict				=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 128, 128, 3)
    lowercase__       :       Union[str, Any]				=						np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]						)

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







   def 					UpperCAmelCase__(						self						)      ->							str:
    lowercase__       :       Dict				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint ,     provider="""CPUExecutionProvider"""						)
    lowercase__       :       List[Any]				=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

    lowercase__       :       Any				=						self.get_dummy_inputs()
    lowercase__       :       Optional[Any]				=						pipe(**lowerCamelCase__						).images
    lowercase__       :       str				=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 128, 128, 3)
    lowercase__       :       Any				=						np.array([0.6_5331, 0.5_8277, 0.4_8204, 0.5_6059, 0.5_3665, 0.5_6235, 0.5_0969, 0.4_0009, 0.4_6552]						)

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

@nightly
@require_onnxruntime
@require_torch_gpu
class 				_SCREAMING_SNAKE_CASE				(							unittest.TestCase  ):
   """simple docstring"""



   @property
   def 					UpperCAmelCase__(						self						)      ->							Dict:
    return (
        "CUDAExecutionProvider",
        {
            "gpu_mem_limit": "15000000000",  # 15GB
            "arena_extend_strategy": "kSameAsRequested",
        },
    )


   @property
   def 					UpperCAmelCase__(						self						)      ->							Optional[Any]:
    lowercase__       :       List[str]				=						ort.SessionOptions()
    lowercase__       :       Optional[int]				=						False
    return options


   def 					UpperCAmelCase__(						self						)      ->							int:
    lowercase__       :       Optional[int]				=						load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
        """/img2img/sketch-mountains-input.jpg"""						)
    lowercase__       :       Dict				=						init_image.resize((768, 512)						)
    # using the PNDM scheduler by default
    lowercase__       :       int				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(
        """CompVis/stable-diffusion-v1-4""" ,     revision="""onnx""" ,     safety_checker=lowerCamelCase__ ,     feature_extractor=lowerCamelCase__ ,     provider=self.gpu_provider ,     sess_options=self.gpu_options ,     )
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

    lowercase__       :       List[Any]				=						"""A fantasy landscape, trending on artstation"""

    lowercase__       :       Any				=						np.random.RandomState(0						)
    lowercase__       :       List[str]				=						pipe(
        prompt=lowerCamelCase__ ,     image=lowerCamelCase__ ,     strength=0.75 ,     guidance_scale=7.5 ,     num_inference_steps=10 ,     generator=lowerCamelCase__ ,     output_type="""np""" ,     )
    lowercase__       :       List[Any]				=						output.images
    lowercase__       :       Tuple				=						images[0, 255:258, 383:386, -1]

    assert images.shape == (1, 512, 768, 3)
    lowercase__       :       Any				=						np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]						)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







   def 					UpperCAmelCase__(						self						)      ->							Optional[Any]:
    lowercase__       :       Tuple				=						load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
        """/img2img/sketch-mountains-input.jpg"""						)
    lowercase__       :       Optional[int]				=						init_image.resize((768, 512)						)
    lowercase__       :       List[str]				=						LMSDiscreteScheduler.from_pretrained(
        """runwayml/stable-diffusion-v1-5""" ,     subfolder="""scheduler""" ,     revision="""onnx"""						)
    lowercase__       :       Optional[Any]				=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(
        """runwayml/stable-diffusion-v1-5""" ,     revision="""onnx""" ,     scheduler=lowerCamelCase__ ,     safety_checker=lowerCamelCase__ ,     feature_extractor=lowerCamelCase__ ,     provider=self.gpu_provider ,     sess_options=self.gpu_options ,     )
    pipe.set_progress_bar_config(disable=lowerCamelCase__						)

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

    lowercase__       :       Tuple				=						np.random.RandomState(0						)
    lowercase__       :       List[str]				=						pipe(
        prompt=lowerCamelCase__ ,     image=lowerCamelCase__ ,     strength=0.75 ,     guidance_scale=7.5 ,     num_inference_steps=20 ,     generator=lowerCamelCase__ ,     output_type="""np""" ,     )
    lowercase__       :       Any				=						output.images
    lowercase__       :       Tuple				=						images[0, 255:258, 383:386, -1]

    assert images.shape == (1, 512, 768, 3)
    lowercase__       :       Union[str, Any]				=						np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]						)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



"""simple docstring"""


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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__snake_case    					=					logging.get_logger(__name__)

__snake_case    					=					'▁'

__snake_case    					=					{'vocab_file': 'sentencepiece.bpe.model'}

__snake_case    					=					{
    'vocab_file': {
        'facebook/xglm-564M': 'https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model',
    }
}

__snake_case    					=					{
    'facebook/xglm-564M': 2048,
}

class 				_SCREAMING_SNAKE_CASE				(							__UpperCAmelCase  ):
   """simple docstring"""



   _a   :       Optional[int]									=    VOCAB_FILES_NAMES
   _a   :       Union[str, Any]									=    PRETRAINED_VOCAB_FILES_MAP
   _a   :       Dict									=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _a   :       Any									=    ['''input_ids''', '''attention_mask''']


   def __init__(						self ,     lowerCamelCase__ ,     lowerCamelCase__="<s>" ,     lowerCamelCase__="</s>" ,     lowerCamelCase__="</s>" ,     lowerCamelCase__="<s>" ,     lowerCamelCase__="<unk>" ,     lowerCamelCase__="<pad>" ,     lowerCamelCase__ = None ,     **lowerCamelCase__ ,     )      ->							None:
    lowercase__       :       List[Any]				=						{} if sp_model_kwargs is None else sp_model_kwargs

    # Compatibility with the original tokenizer
    lowercase__       :       Any				=						7
    lowercase__       :       Dict				=						[F'''<madeupword{i}>''' for i in range(self.num_madeup_words						)]

    lowercase__       :       Tuple				=						kwargs.get("""additional_special_tokens""" ,     []						)
    kwargs["additional_special_tokens"] += [
        word for word in madeup_words if word not in kwargs["additional_special_tokens"]
    ]

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

    lowercase__       :       Optional[Any]				=						spm.SentencePieceProcessor(**self.sp_model_kwargs						)
    self.sp_model.Load(str(lowerCamelCase__						)						)
    lowercase__       :       Optional[Any]				=						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'

    # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
    lowercase__       :       Optional[Any]				=						1

    # Mimic fairseq token-to-id alignment for the first 4 token
    lowercase__       :       Any				=						{"""<s>""": 0, """<pad>""": 1, """</s>""": 2, """<unk>""": 3}

    lowercase__       :       Optional[int]				=						len(self.sp_model						)
    lowercase__       :       Union[str, Any]				=						{F'''<madeupword{i}>''': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words						)}
    self.fairseq_tokens_to_ids.update(lowerCamelCase__						)

    lowercase__       :       List[str]				=						{v: k for k, v in self.fairseq_tokens_to_ids.items()}


   def __getstate__(						self						)      ->							Dict:
    lowercase__       :       Any				=						self.__dict__.copy()
    lowercase__       :       List[str]				=						None
    lowercase__       :       str				=						self.sp_model.serialized_model_proto()
    return state


   def __setstate__(						self ,     lowerCamelCase__						)      ->							Optional[int]:
    lowercase__       :       int				=						d

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

    lowercase__       :       Optional[int]				=						spm.SentencePieceProcessor(**self.sp_model_kwargs						)
    self.sp_model.LoadFromSerializedProto(self.sp_model_proto						)


   def 					UpperCAmelCase__(						self ,     lowerCamelCase__ ,     lowerCamelCase__ = None						)      ->							List[int]:

    if token_ids_a is None:
     return [self.sep_token_id] + token_ids_a
    lowercase__       :       List[str]				=						[self.sep_token_id]
    return sep + token_ids_a + sep + sep + token_ids_a


   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__						)

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


   def 					UpperCAmelCase__(						self ,     lowerCamelCase__ ,     lowerCamelCase__ = None						)      ->							List[int]:
    lowercase__       :       str				=						[self.sep_token_id]

    if token_ids_a is None:
     return len(sep + token_ids_a						) * [0]
    return len(sep + token_ids_a + sep + sep + token_ids_a						) * [0]


   @property
   def 					UpperCAmelCase__(						self						)      ->							Union[str, Any]:
    return len(self.sp_model						) + self.fairseq_offset + self.num_madeup_words


   def 					UpperCAmelCase__(						self						)      ->							int:
    lowercase__       :       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]
    lowercase__       :       List[Any]				=						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__						)      ->							int:
    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__						)      ->							int:
    lowercase__       :       Dict				=						"""""".join(lowerCamelCase__						).replace(lowerCamelCase__ ,     """ """						).strip()
    return out_string







   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
    lowercase__       :       int				=						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:
      lowercase__       :       Union[str, Any]				=						self.sp_model.serialized_model_proto()
      fi.write(lowerCamelCase__						)

    return (out_vocab_file,)