



"""simple docstring"""


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 			_a (					lowerCAmelCase			,							unittest.TestCase):
			"""simple docstring"""
			UpperCamelCase__			      =	ConsistencyModelPipeline
			UpperCamelCase__			      =	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			UpperCamelCase__			      =	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



			@property
			def  lowercase__						(		self     :    List[Any]   )->Union[str, Any]:
							_UpperCAmelCase					=    UNetaDModel.from_pretrained(
							    '''diffusers/consistency-models-test'''				,   subfolder='''test_unet'''				,   )
							return unet



			@property
			def  lowercase__						(		self     :    Optional[Any]   )->Optional[int]:
							_UpperCAmelCase					=    UNetaDModel.from_pretrained(
							    '''diffusers/consistency-models-test'''				,   subfolder='''test_unet_class_cond'''				,   )
							return unet



			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    List[str]=False   )->int:
							if class_cond:
											_UpperCAmelCase					=    self.dummy_cond_unet
							else:
											_UpperCAmelCase					=    self.dummy_uncond_unet

							# Default to CM multistep sampler
							_UpperCAmelCase					=    CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,   sigma_min=0.0_0_2				,   sigma_max=8_0.0				,   )

							_UpperCAmelCase					=    {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							}

							return components



			def  lowercase__						(		self     :    Any				,   __UpperCamelCase     :    Any				,   __UpperCamelCase     :    int=0   )->List[str]:
							if str(__UpperCamelCase   ).startswith('''mps'''   ):
											_UpperCAmelCase					=    torch.manual_seed(__UpperCamelCase   )
							else:
											_UpperCAmelCase					=    torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )

							_UpperCAmelCase					=    {
							    '''batch_size''': 1,
							    '''num_inference_steps''': None,
							    '''timesteps''': [2_2, 0],
							    '''generator''': generator,
							    '''output_type''': '''np''',
							}

							return inputs



			def  lowercase__						(		self     :    Any   )->List[Any]:
							_UpperCAmelCase					=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCAmelCase					=    self.get_dummy_components()
							_UpperCAmelCase					=    ConsistencyModelPipeline(**__UpperCamelCase   )
							_UpperCAmelCase					=    pipe.to(__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]
							_UpperCAmelCase					=    np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]   )

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



			def  lowercase__						(		self     :    Union[str, Any]   )->Tuple:
							_UpperCAmelCase					=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCAmelCase					=    self.get_dummy_components(class_cond=__UpperCamelCase   )
							_UpperCAmelCase					=    ConsistencyModelPipeline(**__UpperCamelCase   )
							_UpperCAmelCase					=    pipe.to(__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    0
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]
							_UpperCAmelCase					=    np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]   )

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



			def  lowercase__						(		self     :    Tuple   )->int:
							_UpperCAmelCase					=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCAmelCase					=    self.get_dummy_components()
							_UpperCAmelCase					=    ConsistencyModelPipeline(**__UpperCamelCase   )
							_UpperCAmelCase					=    pipe.to(__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    1
							_UpperCAmelCase					=    None
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]
							_UpperCAmelCase					=    np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]   )

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



			def  lowercase__						(		self     :    int   )->int:
							_UpperCAmelCase					=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCAmelCase					=    self.get_dummy_components(class_cond=__UpperCamelCase   )
							_UpperCAmelCase					=    ConsistencyModelPipeline(**__UpperCamelCase   )
							_UpperCAmelCase					=    pipe.to(__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    1
							_UpperCAmelCase					=    None
							_UpperCAmelCase					=    0
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]
							_UpperCAmelCase					=    np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]   )

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



@slow
@require_torch_gpu
class 			_a (					unittest.TestCase):
			"""simple docstring"""



			def  lowercase__						(		self     :    Union[str, Any]   )->List[Any]:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



			def  lowercase__						(		self     :    List[str]				,   __UpperCamelCase     :    Optional[int]=0				,   __UpperCamelCase     :    Optional[Any]=False				,   __UpperCamelCase     :    Optional[int]="cpu"				,   __UpperCamelCase     :    Tuple=torch.floataa				,   __UpperCamelCase     :    Dict=(1, 3, 6_4, 6_4)   )->Optional[Any]:
							_UpperCAmelCase					=    torch.manual_seed(__UpperCamelCase   )

							_UpperCAmelCase					=    {
							    '''num_inference_steps''': None,
							    '''timesteps''': [2_2, 0],
							    '''class_labels''': 0,
							    '''generator''': generator,
							    '''output_type''': '''np''',
							}

							if get_fixed_latents:
											_UpperCAmelCase					=    self.get_fixed_latents(seed=__UpperCamelCase				,   device=__UpperCamelCase				,   dtype=__UpperCamelCase				,   shape=__UpperCamelCase   )
											_UpperCAmelCase					=    latents

							return inputs



			def  lowercase__						(		self     :    Optional[int]				,   __UpperCamelCase     :    Dict=0				,   __UpperCamelCase     :    Union[str, Any]="cpu"				,   __UpperCamelCase     :    Tuple=torch.floataa				,   __UpperCamelCase     :    str=(1, 3, 6_4, 6_4)   )->List[Any]:
							if type(__UpperCamelCase   ) == str:
											_UpperCAmelCase					=    torch.device(__UpperCamelCase   )
							_UpperCAmelCase					=    torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )
							_UpperCAmelCase					=    randn_tensor(__UpperCamelCase				,   generator=__UpperCamelCase				,   device=__UpperCamelCase				,   dtype=__UpperCamelCase   )
							return latents



			def  lowercase__						(		self     :    Any   )->int:
							_UpperCAmelCase					=    UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,   subfolder='''diffusers_cd_imagenet64_l2'''   )
							_UpperCAmelCase					=    CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,   sigma_min=0.0_0_2				,   sigma_max=8_0.0				,   )
							_UpperCAmelCase					=    ConsistencyModelPipeline(unet=__UpperCamelCase				,   scheduler=__UpperCamelCase   )
							pipe.to(torch_device=__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_inputs()
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]

							_UpperCAmelCase					=    np.array([0.0_8_8_8, 0.0_8_8_1, 0.0_6_6_6, 0.0_4_7_9, 0.0_2_9_2, 0.0_1_9_5, 0.0_2_0_1, 0.0_1_6_3, 0.0_2_5_4]   )

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



			def  lowercase__						(		self     :    Any   )->Dict:
							_UpperCAmelCase					=    UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,   subfolder='''diffusers_cd_imagenet64_l2'''   )
							_UpperCAmelCase					=    CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,   sigma_min=0.0_0_2				,   sigma_max=8_0.0				,   )
							_UpperCAmelCase					=    ConsistencyModelPipeline(unet=__UpperCamelCase				,   scheduler=__UpperCamelCase   )
							pipe.to(torch_device=__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_inputs()
							_UpperCAmelCase					=    1
							_UpperCAmelCase					=    None
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]

							_UpperCAmelCase					=    np.array([0.0_3_4_0, 0.0_1_5_2, 0.0_0_6_3, 0.0_2_6_7, 0.0_2_2_1, 0.0_1_0_7, 0.0_4_1_6, 0.0_1_8_6, 0.0_2_1_7]   )

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



			@require_torch_a
			def  lowercase__						(		self     :    str   )->List[str]:
							_UpperCAmelCase					=    UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,   subfolder='''diffusers_cd_imagenet64_l2'''   )
							_UpperCAmelCase					=    CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,   sigma_min=0.0_0_2				,   sigma_max=8_0.0				,   )
							_UpperCAmelCase					=    ConsistencyModelPipeline(unet=__UpperCamelCase				,   scheduler=__UpperCamelCase   )
							pipe.to(torch_device=__UpperCamelCase				,   torch_dtype=torch.floataa   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    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   ):
											_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]

							_UpperCAmelCase					=    np.array([0.1_8_7_5, 0.1_4_2_8, 0.1_2_8_9, 0.2_1_5_1, 0.2_0_9_2, 0.1_4_7_7, 0.1_8_7_7, 0.1_6_4_1, 0.1_3_5_3]   )

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



			@require_torch_a
			def  lowercase__						(		self     :    str   )->Union[str, Any]:
							_UpperCAmelCase					=    UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,   subfolder='''diffusers_cd_imagenet64_l2'''   )
							_UpperCAmelCase					=    CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,   sigma_min=0.0_0_2				,   sigma_max=8_0.0				,   )
							_UpperCAmelCase					=    ConsistencyModelPipeline(unet=__UpperCamelCase				,   scheduler=__UpperCamelCase   )
							pipe.to(torch_device=__UpperCamelCase				,   torch_dtype=torch.floataa   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_inputs(get_fixed_latents=__UpperCamelCase				,   device=__UpperCamelCase   )
							_UpperCAmelCase					=    1
							_UpperCAmelCase					=    None
							# Ensure usage of flash attention in torch 2.0
							with sdp_kernel(enable_flash=__UpperCamelCase				,   enable_math=__UpperCamelCase				,   enable_mem_efficient=__UpperCamelCase   ):
											_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							_UpperCAmelCase					=    image[0, -3:, -3:, -1]

							_UpperCAmelCase					=    np.array([0.1_6_6_3, 0.1_9_4_8, 0.2_2_7_5, 0.1_6_8_0, 0.1_2_0_4, 0.1_2_4_5, 0.1_8_5_8, 0.1_3_3_8, 0.2_0_9_5]   )

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








"""simple docstring"""


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

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


if is_sentencepiece_available():
					from .tokenization_fnet import FNetTokenizer
else:
					__A   :							Tuple										=   None

__A   :							Optional[Any]										=   logging.get_logger(__name__)
__A   :							Union[str, Any]										=   {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

__A   :							int										=   {
    "vocab_file": {
        "google/fnet-base": "https://huggingface.co/google/fnet-base/resolve/main/spiece.model",
        "google/fnet-large": "https://huggingface.co/google/fnet-large/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "google/fnet-base": "https://huggingface.co/google/fnet-base/resolve/main/tokenizer.json",
        "google/fnet-large": "https://huggingface.co/google/fnet-large/resolve/main/tokenizer.json",
    },
}

__A   :							Optional[int]										=   {
    "google/fnet-base": 512,
    "google/fnet-large": 512,
}

__A   :							List[Any]										=   "▁"
class 			_a (					lowerCAmelCase):
			"""simple docstring"""
			UpperCamelCase__			      =	VOCAB_FILES_NAMES
			UpperCamelCase__			      =	PRETRAINED_VOCAB_FILES_MAP
			UpperCamelCase__			      =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			UpperCamelCase__			      =	["""input_ids""", """token_type_ids"""]
			UpperCamelCase__			      =	FNetTokenizer



			def __init__(		self     :    Optional[Any]				,   __UpperCamelCase     :    Tuple=None				,   __UpperCamelCase     :    Dict=None				,   __UpperCamelCase     :    List[str]=False				,   __UpperCamelCase     :    Union[str, Any]=True				,   __UpperCamelCase     :    List[str]=True				,   __UpperCamelCase     :    List[Any]="<unk>"				,   __UpperCamelCase     :    Any="[SEP]"				,   __UpperCamelCase     :    Dict="<pad>"				,   __UpperCamelCase     :    Optional[int]="[CLS]"				,   __UpperCamelCase     :    List[Any]="[MASK]"				,   **__UpperCamelCase     :    Optional[Any]				,   )->Dict:
							# Mask token behave like a normal word, i.e. include the space before it and
							# is included in the raw text, there should be a match in a non-normalized sentence.
							_UpperCAmelCase					=    (
							    AddedToken(__UpperCamelCase				,   lstrip=__UpperCamelCase				,   rstrip=__UpperCamelCase				,   normalized=__UpperCamelCase   )
							    if isinstance(__UpperCamelCase				,   __UpperCamelCase   )
							    else mask_token
							)

							super().__init__(
							    __UpperCamelCase				,   tokenizer_file=__UpperCamelCase				,   do_lower_case=__UpperCamelCase				,   remove_space=__UpperCamelCase				,   keep_accents=__UpperCamelCase				,   unk_token=__UpperCamelCase				,   sep_token=__UpperCamelCase				,   pad_token=__UpperCamelCase				,   cls_token=__UpperCamelCase				,   mask_token=__UpperCamelCase				,   **__UpperCamelCase				,   )

							_UpperCAmelCase					=    do_lower_case
							_UpperCAmelCase					=    remove_space
							_UpperCAmelCase					=    keep_accents
							_UpperCAmelCase					=    vocab_file
							_UpperCAmelCase					=    False if not self.vocab_file else True



			def  lowercase__						(		self     :    int				,   __UpperCamelCase     :    List[int]				,   __UpperCamelCase     :    Optional[List[int]] = None   )->List[int]:
							_UpperCAmelCase					=    [self.sep_token_id]
							_UpperCAmelCase					=    [self.cls_token_id]
							if token_ids_a is None:
											return cls + token_ids_a + sep
							return cls + token_ids_a + sep + token_ids_a + sep



			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    List[int]				,   __UpperCamelCase     :    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   ) * [0] + len(token_ids_a + sep   ) * [1]



			def  lowercase__						(		self     :    List[str]				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    Optional[str] = None   )->Tuple[str]:
							if not os.path.isdir(__UpperCamelCase   ):
											logger.error(F'Vocabulary path ({save_directory}) should be a directory'   )
											return
							_UpperCAmelCase					=    os.path.join(
							    __UpperCamelCase				,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']   )

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

							return (out_vocab_file,)



