





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        _lowerCAmelCase    (			UpperCAmelCase_							,    unittest.TestCase							):


  '''simple docstring'''
  a_							:       Optional[int]      					="""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""







  def UpperCamelCase_ (      self      :					Optional[Any]      ,				UpperCamelCase      :					Optional[int]=0							):


    '''simple docstring'''





    _snake_case   :				str														=	floats_tensor((1, 3, 1_28, 1_28)      ,				rng=random.Random(UpperCamelCase							)							)
    _snake_case   :				Optional[Any]														=	np.random.RandomState(UpperCamelCase							)
    _snake_case   :				Optional[int]														=	{
        '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]							):


    '''simple docstring'''





    _snake_case   :				Dict														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,				provider='CPUExecutionProvider'							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Optional[Any]														=	self.get_dummy_inputs()
    _snake_case   :				Optional[int]														=	pipe(**UpperCamelCase							).images
    _snake_case   :				List[Any]														=	image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 1_28, 1_28, 3)
    _snake_case   :				Optional[Any]														=	np.array([0.6_96_43, 0.5_84_84, 0.5_03_14, 0.5_87_60, 0.5_53_68, 0.5_96_43, 0.5_15_29, 0.4_12_17, 0.4_90_87]							)
    assert np.abs(image_slice - expected_slice							).max() < 1e-1







  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				Any														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,				provider='CPUExecutionProvider'							)
    _snake_case   :				Dict														=	PNDMScheduler.from_config(pipe.scheduler.config      ,				skip_prk_steps=UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

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

    assert image.shape == (1, 1_28, 1_28, 3)
    _snake_case   :				Union[str, Any]														=	np.array([0.6_17_37, 0.5_46_42, 0.5_31_83, 0.5_44_65, 0.5_27_42, 0.6_05_25, 0.4_99_69, 0.4_06_55, 0.4_81_54]							)

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







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,				provider='CPUExecutionProvider'							)
    _snake_case   :				str														=	LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    # warmup pass to apply optimizations
    _snake_case   :				int														=	pipe(**self.get_dummy_inputs()							)

    _snake_case   :				List[Any]														=	self.get_dummy_inputs()
    _snake_case   :				Dict														=	pipe(**UpperCamelCase							).images
    _snake_case   :				int														=	image[0, -3:, -3:, -1]

    assert image.shape == (1, 1_28, 1_28, 3)
    _snake_case   :				List[Any]														=	np.array([0.5_27_61, 0.5_99_77, 0.4_90_33, 0.4_96_19, 0.5_42_82, 0.5_03_11, 0.4_76_00, 0.4_09_18, 0.4_52_03]							)

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







  def UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    _snake_case   :				str														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,				provider='CPUExecutionProvider'							)
    _snake_case   :				Dict														=	EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

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

    assert image.shape == (1, 1_28, 1_28, 3)
    _snake_case   :				Optional[Any]														=	np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]							)

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







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,				provider='CPUExecutionProvider'							)
    _snake_case   :				Union[str, Any]														=	EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Any														=	self.get_dummy_inputs()
    _snake_case   :				List[Any]														=	pipe(**UpperCamelCase							).images
    _snake_case   :				Optional[Any]														=	image[0, -3:, -3:, -1]

    assert image.shape == (1, 1_28, 1_28, 3)
    _snake_case   :				Any														=	np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]							)

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







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				Dict														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint      ,				provider='CPUExecutionProvider'							)
    _snake_case   :				Union[str, Any]														=	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	self.get_dummy_inputs()
    _snake_case   :				Optional[int]														=	pipe(**UpperCamelCase							).images
    _snake_case   :				List[Any]														=	image[0, -3:, -3:, -1]

    assert image.shape == (1, 1_28, 1_28, 3)
    _snake_case   :				Optional[Any]														=	np.array([0.6_53_31, 0.5_82_77, 0.4_82_04, 0.5_60_59, 0.5_36_65, 0.5_62_35, 0.5_09_69, 0.4_00_09, 0.4_65_52]							)

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




@nightly
@require_onnxruntime
@require_torch_gpu
class        _lowerCAmelCase    (			unittest.TestCase							):


  '''simple docstring'''







  @property
  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    return (
        "CUDAExecutionProvider",
        {
            "gpu_mem_limit": "15000000000",  # 15GB
            "arena_extend_strategy": "kSameAsRequested",
        },
    )







  @property
  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    _snake_case   :				int														=	ort.SessionOptions()
    _snake_case   :				str														=	False
    return options







  def UpperCamelCase_ (      self      :					str							):


    '''simple docstring'''





    _snake_case   :				Dict														=	load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
        '/img2img/sketch-mountains-input.jpg'							)
    _snake_case   :				List[Any]														=	init_image.resize((7_68, 5_12)							)
    # using the PNDM scheduler by default
    _snake_case   :				Union[str, Any]														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4'      ,				revision='onnx'      ,				safety_checker=UpperCamelCase      ,				feature_extractor=UpperCamelCase      ,				provider=self.gpu_provider      ,				sess_options=self.gpu_options      ,				)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Any														=	'A fantasy landscape, trending on artstation'

    _snake_case   :				int														=	np.random.RandomState(0							)
    _snake_case   :				Optional[int]														=	pipe(
        prompt=UpperCamelCase      ,				image=UpperCamelCase      ,				strength=0.75      ,				guidance_scale=7.5      ,				num_inference_steps=10      ,				generator=UpperCamelCase      ,				output_type='np'      ,				)
    _snake_case   :				int														=	output.images
    _snake_case   :				Any														=	images[0, 2_55:2_58, 3_83:3_86, -1]

    assert images.shape == (1, 5_12, 7_68, 3)
    _snake_case   :				List[Any]														=	np.array([0.49_09, 0.50_59, 0.53_72, 0.46_23, 0.48_76, 0.50_49, 0.48_20, 0.49_56, 0.50_19]							)
    # 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]							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	load_image(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
        '/img2img/sketch-mountains-input.jpg'							)
    _snake_case   :				Any														=	init_image.resize((7_68, 5_12)							)
    _snake_case   :				Optional[int]														=	LMSDiscreteScheduler.from_pretrained(
        'runwayml/stable-diffusion-v1-5'      ,				subfolder='scheduler'      ,				revision='onnx'							)
    _snake_case   :				List[Any]														=	OnnxStableDiffusionImgaImgPipeline.from_pretrained(
        'runwayml/stable-diffusion-v1-5'      ,				revision='onnx'      ,				scheduler=UpperCamelCase      ,				safety_checker=UpperCamelCase      ,				feature_extractor=UpperCamelCase      ,				provider=self.gpu_provider      ,				sess_options=self.gpu_options      ,				)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Optional[int]														=	'A fantasy landscape, trending on artstation'

    _snake_case   :				Any														=	np.random.RandomState(0							)
    _snake_case   :				List[str]														=	pipe(
        prompt=UpperCamelCase      ,				image=UpperCamelCase      ,				strength=0.75      ,				guidance_scale=7.5      ,				num_inference_steps=20      ,				generator=UpperCamelCase      ,				output_type='np'      ,				)
    _snake_case   :				int														=	output.images
    _snake_case   :				Tuple														=	images[0, 2_55:2_58, 3_83:3_86, -1]

    assert images.shape == (1, 5_12, 7_68, 3)
    _snake_case   :				Union[str, Any]														=	np.array([0.80_43, 0.9_26, 0.95_81, 0.81_19, 0.89_54, 0.9_13, 0.72_09, 0.74_63, 0.74_31]							)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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











import argparse
import os

import transformers

from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS
from .utils import logging


logging.set_verbosity_info()

lowerCAmelCase_						   =    logging.get_logger(__name__)


lowerCAmelCase_						   =    {name: getattr(transformers, name + """Fast""") for name in SLOW_TO_FAST_CONVERTERS}


def     lowerCamelCase_	(      lowerCAmelCase:	List[Any]    ,							lowerCAmelCase:	Optional[Any]    ,							lowerCAmelCase:	Dict    ,							lowerCAmelCase:	Union[str, Any]				)->		Optional[int]:
  if tokenizer_name is not None and tokenizer_name not in TOKENIZER_CLASSES:
    raise ValueError(F"""Unrecognized tokenizer name, should be one of {list(TOKENIZER_CLASSES.keys()				)}."""				)

  if tokenizer_name is None:
    _snake_case   :				Tuple														=	TOKENIZER_CLASSES
  else:
    _snake_case   :				Union[str, Any]														=	{tokenizer_name: getattr(lowerCAmelCase    ,							tokenizer_name + 'Fast'				)}

  logger.info(F"""Loading tokenizer classes: {tokenizer_names}"""				)

  for tokenizer_name in tokenizer_names:
    _snake_case   :				Dict														=	TOKENIZER_CLASSES[tokenizer_name]

    _snake_case   :				Optional[Any]														=	True
    if checkpoint_name is None:
      _snake_case   :				Union[str, Any]														=	list(tokenizer_class.max_model_input_sizes.keys()				)
    else:
      _snake_case   :				Optional[int]														=	[checkpoint_name]

    logger.info(F"""For tokenizer {tokenizer_class.__class__.__name__} loading checkpoints: {checkpoint_names}"""				)

    for checkpoint in checkpoint_names:
      logger.info(F"""Loading {tokenizer_class.__class__.__name__} {checkpoint}"""				)

      # Load tokenizer
      _snake_case   :				str														=	tokenizer_class.from_pretrained(lowerCAmelCase    ,							force_download=lowerCAmelCase				)

      # Save fast tokenizer
      logger.info(F"""Save fast tokenizer to {dump_path} with prefix {checkpoint} add_prefix {add_prefix}"""				)

      # For organization names we create sub-directories
      if "/" in checkpoint:
        _snake_case							,					_snake_case   :				Tuple														=	checkpoint.split('/'				)
        _snake_case   :				int														=	os.path.join(lowerCAmelCase    ,							lowerCAmelCase				)
      elif add_prefix:
        _snake_case   :				Dict														=	checkpoint
        _snake_case   :				Optional[Any]														=	dump_path
      else:
        _snake_case   :				str														=	None
        _snake_case   :				Union[str, Any]														=	dump_path

      logger.info(F"""=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}"""				)

      if checkpoint in list(tokenizer.pretrained_vocab_files_map.values()				)[0]:
        _snake_case   :				Optional[Any]														=	list(tokenizer.pretrained_vocab_files_map.values()				)[0][checkpoint]
        _snake_case   :				Optional[int]														=	file_path.split(lowerCAmelCase				)[-1][0]
        if next_char == "/":
          _snake_case   :				Union[str, Any]														=	os.path.join(lowerCAmelCase    ,							lowerCAmelCase				)
          _snake_case   :				str														=	None

        logger.info(F"""=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}"""				)

      _snake_case   :				Optional[int]														=	tokenizer.save_pretrained(
          lowerCAmelCase    ,							legacy_format=lowerCAmelCase    ,							filename_prefix=lowerCAmelCase				)
      logger.info(F"""=> File names {file_names}"""				)

      for file_name in file_names:
        if not file_name.endswith('tokenizer.json'				):
          os.remove(lowerCAmelCase				)
          logger.info(F"""=> removing {file_name}"""				)


if __name__ == "__main__":
      lowerCAmelCase_						   =    argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--dump_path""", default=None, type=str, required=True, help="""Path to output generated fast tokenizer files."""
      )
      parser.add_argument(
          """--tokenizer_name""",
          default=None,
          type=str,
          help=(
              F"""Optional tokenizer type selected in the list of {list(TOKENIZER_CLASSES.keys())}. If not given, will """
              """download and convert all the checkpoints from AWS."""
          ),
      )
      parser.add_argument(
          """--checkpoint_name""",
          default=None,
          type=str,
          help="""Optional checkpoint name. If not given, will download and convert the canonical checkpoints from AWS.""",
      )
      parser.add_argument(
          """--force_download""",
          action="""store_true""",
          help="""Re-download checkpoints.""",
      )
      lowerCAmelCase_						   =    parser.parse_args()

      convert_slow_checkpoint_to_fast(args.tokenizer_name, args.checkpoint_name, args.dump_path, args.force_download)




