





import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class        lowerCAmelCase_    (			UpperCAmelCase_							,    unittest.TestCase							):


  '''simple docstring'''
  a_							:       Optional[Any]      					=ConsistencyModelPipeline
  a_							:       List[str]      					=UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  a_							:       Optional[Any]      					=UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

  # Override required_optional_params to remove num_images_per_prompt
  a_							:       Tuple      					=frozenset(
      [
          """num_inference_steps""",
          """generator""",
          """latents""",
          """output_type""",
          """return_dict""",
          """callback""",
          """callback_steps""",
      ]							)







  @property
  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	UNetaDModel.from_pretrained(
        'diffusers/consistency-models-test'      ,				subfolder='test_unet'      ,				)
    return unet







  @property
  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	UNetaDModel.from_pretrained(
        'diffusers/consistency-models-test'      ,				subfolder='test_unet_class_cond'      ,				)
    return unet







  def UpperCamelCase_ (      self      :					Optional[int]      ,				UpperCamelCase      :					Dict=False							):


    '''simple docstring'''





    if class_cond:
      _snake_case   :				Dict														=	self.dummy_cond_unet
    else:
      _snake_case   :				Any														=	self.dummy_uncond_unet

    # Default to CM multistep sampler
    _snake_case   :				int														=	CMStochasticIterativeScheduler(
        num_train_timesteps=40      ,				sigma_min=0.0_02      ,				sigma_max=80.0      ,				)

    _snake_case   :				Optional[int]														=	{
        'unet': unet,
        'scheduler': scheduler,
    }

    return components







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


    '''simple docstring'''





    if str(UpperCamelCase							).startswith('mps'							):
      _snake_case   :				Optional[int]														=	torch.manual_seed(UpperCamelCase							)
    else:
      _snake_case   :				Union[str, Any]														=	torch.Generator(device=UpperCamelCase							).manual_seed(UpperCamelCase							)

    _snake_case   :				Tuple														=	{
        'batch_size': 1,
        'num_inference_steps': None,
        'timesteps': [22, 0],
        'generator': generator,
        'output_type': 'np',
    }

    return inputs







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	'cpu'  # ensure determinism for the device-dependent torch.Generator
    _snake_case   :				List[str]														=	self.get_dummy_components()
    _snake_case   :				Dict														=	ConsistencyModelPipeline(**UpperCamelCase							)
    _snake_case   :				List[str]														=	pipe.to(UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Dict														=	self.get_dummy_inputs(UpperCamelCase							)
    _snake_case   :				str														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 32, 32, 3)

    _snake_case   :				Optional[Any]														=	image[0, -3:, -3:, -1]
    _snake_case   :				List[str]														=	np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]							)

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







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


    '''simple docstring'''





    _snake_case   :				int														=	'cpu'  # ensure determinism for the device-dependent torch.Generator
    _snake_case   :				List[str]														=	self.get_dummy_components(class_cond=UpperCamelCase							)
    _snake_case   :				Any														=	ConsistencyModelPipeline(**UpperCamelCase							)
    _snake_case   :				Any														=	pipe.to(UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				List[Any]														=	self.get_dummy_inputs(UpperCamelCase							)
    _snake_case   :				int														=	0
    _snake_case   :				List[str]														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 32, 32, 3)

    _snake_case   :				Optional[int]														=	image[0, -3:, -3:, -1]
    _snake_case   :				Union[str, Any]														=	np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]							)

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







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


    '''simple docstring'''





    _snake_case   :				Dict														=	'cpu'  # ensure determinism for the device-dependent torch.Generator
    _snake_case   :				Optional[int]														=	self.get_dummy_components()
    _snake_case   :				int														=	ConsistencyModelPipeline(**UpperCamelCase							)
    _snake_case   :				int														=	pipe.to(UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				str														=	self.get_dummy_inputs(UpperCamelCase							)
    _snake_case   :				Tuple														=	1
    _snake_case   :				Optional[int]														=	None
    _snake_case   :				Union[str, Any]														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 32, 32, 3)

    _snake_case   :				str														=	image[0, -3:, -3:, -1]
    _snake_case   :				Tuple														=	np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]							)

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







  def UpperCamelCase_ (      self      :					str							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	'cpu'  # ensure determinism for the device-dependent torch.Generator
    _snake_case   :				Optional[int]														=	self.get_dummy_components(class_cond=UpperCamelCase							)
    _snake_case   :				Tuple														=	ConsistencyModelPipeline(**UpperCamelCase							)
    _snake_case   :				Union[str, Any]														=	pipe.to(UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Dict														=	self.get_dummy_inputs(UpperCamelCase							)
    _snake_case   :				List[Any]														=	1
    _snake_case   :				Optional[Any]														=	None
    _snake_case   :				Optional[Any]														=	0
    _snake_case   :				int														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 32, 32, 3)

    _snake_case   :				int														=	image[0, -3:, -3:, -1]
    _snake_case   :				Tuple														=	np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]							)

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




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


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()







  def UpperCamelCase_ (      self      :					List[str]      ,				UpperCamelCase      :					Tuple=0      ,				UpperCamelCase      :					List[Any]=False      ,				UpperCamelCase      :					Optional[int]="cpu"      ,				UpperCamelCase      :					List[Any]=torch.floataa      ,				UpperCamelCase      :					List[Any]=(1, 3, 64, 64)							):


    '''simple docstring'''





    _snake_case   :				str														=	torch.manual_seed(UpperCamelCase							)

    _snake_case   :				Tuple														=	{
        'num_inference_steps': None,
        'timesteps': [22, 0],
        'class_labels': 0,
        'generator': generator,
        'output_type': 'np',
    }

    if get_fixed_latents:
      _snake_case   :				Union[str, Any]														=	self.get_fixed_latents(seed=UpperCamelCase      ,				device=UpperCamelCase      ,				dtype=UpperCamelCase      ,				shape=UpperCamelCase							)
      _snake_case   :				Dict														=	latents

    return inputs







  def UpperCamelCase_ (      self      :					Union[str, Any]      ,				UpperCamelCase      :					Union[str, Any]=0      ,				UpperCamelCase      :					List[str]="cpu"      ,				UpperCamelCase      :					Optional[Any]=torch.floataa      ,				UpperCamelCase      :					Tuple=(1, 3, 64, 64)							):


    '''simple docstring'''





    if type(UpperCamelCase							) == str:
      _snake_case   :				List[Any]														=	torch.device(UpperCamelCase							)
    _snake_case   :				Any														=	torch.Generator(device=UpperCamelCase							).manual_seed(UpperCamelCase							)
    _snake_case   :				List[str]														=	randn_tensor(UpperCamelCase      ,				generator=UpperCamelCase      ,				device=UpperCamelCase      ,				dtype=UpperCamelCase							)
    return latents







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	UNetaDModel.from_pretrained('diffusers/consistency_models'      ,				subfolder='diffusers_cd_imagenet64_l2'							)
    _snake_case   :				List[str]														=	CMStochasticIterativeScheduler(
        num_train_timesteps=40      ,				sigma_min=0.0_02      ,				sigma_max=80.0      ,				)
    _snake_case   :				int														=	ConsistencyModelPipeline(unet=UpperCamelCase      ,				scheduler=UpperCamelCase							)
    pipe.to(torch_device=UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				List[str]														=	self.get_inputs()
    _snake_case   :				Union[str, Any]														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 64, 64, 3)

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

    _snake_case   :				Tuple														=	np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]							)

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







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	UNetaDModel.from_pretrained('diffusers/consistency_models'      ,				subfolder='diffusers_cd_imagenet64_l2'							)
    _snake_case   :				Optional[Any]														=	CMStochasticIterativeScheduler(
        num_train_timesteps=40      ,				sigma_min=0.0_02      ,				sigma_max=80.0      ,				)
    _snake_case   :				int														=	ConsistencyModelPipeline(unet=UpperCamelCase      ,				scheduler=UpperCamelCase							)
    pipe.to(torch_device=UpperCamelCase							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Any														=	self.get_inputs()
    _snake_case   :				List[str]														=	1
    _snake_case   :				Dict														=	None
    _snake_case   :				List[Any]														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 64, 64, 3)

    _snake_case   :				List[Any]														=	image[0, -3:, -3:, -1]

    _snake_case   :				str														=	np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]							)

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







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


    '''simple docstring'''





    _snake_case   :				List[str]														=	UNetaDModel.from_pretrained('diffusers/consistency_models'      ,				subfolder='diffusers_cd_imagenet64_l2'							)
    _snake_case   :				Union[str, Any]														=	CMStochasticIterativeScheduler(
        num_train_timesteps=40      ,				sigma_min=0.0_02      ,				sigma_max=80.0      ,				)
    _snake_case   :				int														=	ConsistencyModelPipeline(unet=UpperCamelCase      ,				scheduler=UpperCamelCase							)
    pipe.to(torch_device=UpperCamelCase      ,				torch_dtype=torch.floataa							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	self.get_inputs(get_fixed_latents=UpperCamelCase      ,				device=UpperCamelCase							)
    # Ensure usage of flash attention in torch 2.0
    with sdp_kernel(enable_flash=UpperCamelCase      ,				enable_math=UpperCamelCase      ,				enable_mem_efficient=UpperCamelCase							):
      _snake_case   :				int														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 64, 64, 3)

    _snake_case   :				str														=	image[0, -3:, -3:, -1]

    _snake_case   :				List[str]														=	np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]							)

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







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


    '''simple docstring'''





    _snake_case   :				Any														=	UNetaDModel.from_pretrained('diffusers/consistency_models'      ,				subfolder='diffusers_cd_imagenet64_l2'							)
    _snake_case   :				Optional[int]														=	CMStochasticIterativeScheduler(
        num_train_timesteps=40      ,				sigma_min=0.0_02      ,				sigma_max=80.0      ,				)
    _snake_case   :				List[str]														=	ConsistencyModelPipeline(unet=UpperCamelCase      ,				scheduler=UpperCamelCase							)
    pipe.to(torch_device=UpperCamelCase      ,				torch_dtype=torch.floataa							)
    pipe.set_progress_bar_config(disable=UpperCamelCase							)

    _snake_case   :				Tuple														=	self.get_inputs(get_fixed_latents=UpperCamelCase      ,				device=UpperCamelCase							)
    _snake_case   :				Union[str, Any]														=	1
    _snake_case   :				Optional[Any]														=	None
    # Ensure usage of flash attention in torch 2.0
    with sdp_kernel(enable_flash=UpperCamelCase      ,				enable_math=UpperCamelCase      ,				enable_mem_efficient=UpperCamelCase							):
      _snake_case   :				Optional[Any]														=	pipe(**UpperCamelCase							).images
    assert image.shape == (1, 64, 64, 3)

    _snake_case   :				Optional[int]														=	image[0, -3:, -3:, -1]

    _snake_case   :				Optional[Any]														=	np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]							)

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











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)




