'''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__ (  UpperCAmelCase__				,  unittest.TestCase					):
     lowerCamelCase    :	Union[str, Any]				=ConsistencyModelPipeline
     lowerCamelCase    :	Dict				=UNCONDITIONAL_IMAGE_GENERATION_PARAMS
     lowerCamelCase    :	Optional[Any]				=UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






     @property
     def   SCREAMING_SNAKE_CASE__  (					self						:						Union[str, Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    UNetaDModel.from_pretrained(
             '''diffusers/consistency-models-test'''    ,							subfolder='''test_unet'''    ,							)
         return unet






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





         """simple docstring"""



         __lowerCamelCase            =    UNetaDModel.from_pretrained(
             '''diffusers/consistency-models-test'''    ,							subfolder='''test_unet_class_cond'''    ,							)
         return unet






     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[int]    ,							a						:						List[Any]=False       ):





         """simple docstring"""



         if class_cond:
             __lowerCamelCase            =    self.dummy_cond_unet
         else:
             __lowerCamelCase            =    self.dummy_uncond_unet

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

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

         return components






     def   SCREAMING_SNAKE_CASE__  (					self						:						Any    ,							a						:						str    ,							a						:						int=0       ):





         """simple docstring"""



         if str(a       ).startswith('''mps'''       ):
             __lowerCamelCase            =    torch.manual_seed(a       )
         else:
             __lowerCamelCase            =    torch.Generator(device=a       ).manual_seed(a       )

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

         return inputs






     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[int]       ):





         """simple docstring"""



         __lowerCamelCase            =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase            =    self.get_dummy_components()
         __lowerCamelCase            =    ConsistencyModelPipeline(**a       )
         __lowerCamelCase            =    pipe.to(a       )
         pipe.set_progress_bar_config(disable=a       )

         __lowerCamelCase            =    self.get_dummy_inputs(a       )
         __lowerCamelCase            =    pipe(**a       ).images
         assert image.shape == (1, 32, 32, 3)

         __lowerCamelCase            =    image[0, -3:, -3:, -1]
         __lowerCamelCase            =    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   SCREAMING_SNAKE_CASE__  (					self						:						int       ):





         """simple docstring"""



         __lowerCamelCase            =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase            =    self.get_dummy_components(class_cond=a       )
         __lowerCamelCase            =    ConsistencyModelPipeline(**a       )
         __lowerCamelCase            =    pipe.to(a       )
         pipe.set_progress_bar_config(disable=a       )

         __lowerCamelCase            =    self.get_dummy_inputs(a       )
         __lowerCamelCase            =    0
         __lowerCamelCase            =    pipe(**a       ).images
         assert image.shape == (1, 32, 32, 3)

         __lowerCamelCase            =    image[0, -3:, -3:, -1]
         __lowerCamelCase            =    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   SCREAMING_SNAKE_CASE__  (					self						:						Tuple       ):





         """simple docstring"""



         __lowerCamelCase            =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase            =    self.get_dummy_components()
         __lowerCamelCase            =    ConsistencyModelPipeline(**a       )
         __lowerCamelCase            =    pipe.to(a       )
         pipe.set_progress_bar_config(disable=a       )

         __lowerCamelCase            =    self.get_dummy_inputs(a       )
         __lowerCamelCase            =    1
         __lowerCamelCase            =    None
         __lowerCamelCase            =    pipe(**a       ).images
         assert image.shape == (1, 32, 32, 3)

         __lowerCamelCase            =    image[0, -3:, -3:, -1]
         __lowerCamelCase            =    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   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase            =    self.get_dummy_components(class_cond=a       )
         __lowerCamelCase            =    ConsistencyModelPipeline(**a       )
         __lowerCamelCase            =    pipe.to(a       )
         pipe.set_progress_bar_config(disable=a       )

         __lowerCamelCase            =    self.get_dummy_inputs(a       )
         __lowerCamelCase            =    1
         __lowerCamelCase            =    None
         __lowerCamelCase            =    0
         __lowerCamelCase            =    pipe(**a       ).images
         assert image.shape == (1, 32, 32, 3)

         __lowerCamelCase            =    image[0, -3:, -3:, -1]
         __lowerCamelCase            =    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        a__ (  unittest.TestCase					):






     def   SCREAMING_SNAKE_CASE__  (					self						:						int       ):





         """simple docstring"""



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






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





         """simple docstring"""



         __lowerCamelCase            =    torch.manual_seed(a       )

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

         if get_fixed_latents:
             __lowerCamelCase            =    self.get_fixed_latents(seed=a    ,							device=a    ,							dtype=a    ,							shape=a       )
             __lowerCamelCase            =    latents

         return inputs






     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]    ,							a						:						Any=0    ,							a						:						List[str]="cpu"    ,							a						:						int=torch.floataa    ,							a						:						Dict=(1, 3, 64, 64)       ):





         """simple docstring"""



         if type(a       ) == str:
             __lowerCamelCase            =    torch.device(a       )
         __lowerCamelCase            =    torch.Generator(device=a       ).manual_seed(a       )
         __lowerCamelCase            =    randn_tensor(a    ,							generator=a    ,							device=a    ,							dtype=a       )
         return latents






     def   SCREAMING_SNAKE_CASE__  (					self						:						Tuple       ):





         """simple docstring"""



         __lowerCamelCase            =    UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,							subfolder='''diffusers_cd_imagenet64_l2'''       )
         __lowerCamelCase            =    CMStochasticIterativeScheduler(
             num_train_timesteps=40    ,							sigma_min=0.0_02    ,							sigma_max=80.0    ,							)
         __lowerCamelCase            =    ConsistencyModelPipeline(unet=a    ,							scheduler=a       )
         pipe.to(torch_device=a       )
         pipe.set_progress_bar_config(disable=a       )

         __lowerCamelCase            =    self.get_inputs()
         __lowerCamelCase            =    pipe(**a       ).images
         assert image.shape == (1, 64, 64, 3)

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

         __lowerCamelCase            =    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   SCREAMING_SNAKE_CASE__  (					self						:						Dict       ):





         """simple docstring"""



         __lowerCamelCase            =    UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,							subfolder='''diffusers_cd_imagenet64_l2'''       )
         __lowerCamelCase            =    CMStochasticIterativeScheduler(
             num_train_timesteps=40    ,							sigma_min=0.0_02    ,							sigma_max=80.0    ,							)
         __lowerCamelCase            =    ConsistencyModelPipeline(unet=a    ,							scheduler=a       )
         pipe.to(torch_device=a       )
         pipe.set_progress_bar_config(disable=a       )

         __lowerCamelCase            =    self.get_inputs()
         __lowerCamelCase            =    1
         __lowerCamelCase            =    None
         __lowerCamelCase            =    pipe(**a       ).images
         assert image.shape == (1, 64, 64, 3)

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

         __lowerCamelCase            =    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   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,							subfolder='''diffusers_cd_imagenet64_l2'''       )
         __lowerCamelCase            =    CMStochasticIterativeScheduler(
             num_train_timesteps=40    ,							sigma_min=0.0_02    ,							sigma_max=80.0    ,							)
         __lowerCamelCase            =    ConsistencyModelPipeline(unet=a    ,							scheduler=a       )
         pipe.to(torch_device=a    ,							torch_dtype=torch.floataa       )
         pipe.set_progress_bar_config(disable=a       )

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

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

         __lowerCamelCase            =    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   SCREAMING_SNAKE_CASE__  (					self						:						Tuple       ):





         """simple docstring"""



         __lowerCamelCase            =    UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,							subfolder='''diffusers_cd_imagenet64_l2'''       )
         __lowerCamelCase            =    CMStochasticIterativeScheduler(
             num_train_timesteps=40    ,							sigma_min=0.0_02    ,							sigma_max=80.0    ,							)
         __lowerCamelCase            =    ConsistencyModelPipeline(unet=a    ,							scheduler=a       )
         pipe.to(torch_device=a    ,							torch_dtype=torch.floataa       )
         pipe.set_progress_bar_config(disable=a       )

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

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

         __lowerCamelCase            =    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

'''simple docstring'''




import json
import os
from typing import Optional, Tuple

import regex as re

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


__UpperCAmelCase					=logging.get_logger(__name__)

__UpperCAmelCase					={
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

__UpperCAmelCase					={
    "vocab_file": {"ctrl": "https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-vocab.json"},
    "merges_file": {"ctrl": "https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-merges.txt"},
}

__UpperCAmelCase					={
    "ctrl": 2_5_6,
}

__UpperCAmelCase					={
    "Pregnancy": 1_6_8_6_2_9,
    "Christianity": 7_6_7_5,
    "Explain": 1_0_6_4_2_3,
    "Fitness": 6_3_4_4_0,
    "Saving": 6_3_1_6_3,
    "Ask": 2_7_1_7_1,
    "Ass": 9_5_9_8_5,
    "Joke": 1_6_3_5_0_9,
    "Questions": 4_5_6_2_2,
    "Thoughts": 4_9_6_0_5,
    "Retail": 5_2_3_4_2,
    "Feminism": 1_6_4_3_3_8,
    "Writing": 1_1_9_9_2,
    "Atheism": 1_9_2_2_6_3,
    "Netflix": 4_8_6_1_6,
    "Computing": 3_9_6_3_9,
    "Opinion": 4_3_2_1_3,
    "Alone": 4_4_9_6_7,
    "Funny": 5_8_9_1_7,
    "Gaming": 4_0_3_5_8,
    "Human": 4_0_8_8,
    "India": 1_3_3_1,
    "Joker": 7_7_1_3_8,
    "Diet": 3_6_2_0_6,
    "Legal": 1_1_8_5_9,
    "Norman": 4_9_3_9,
    "Tip": 7_2_6_8_9,
    "Weight": 5_2_3_4_3,
    "Movies": 4_6_2_7_3,
    "Running": 2_3_4_2_5,
    "Science": 2_0_9_0,
    "Horror": 3_7_7_9_3,
    "Confession": 6_0_5_7_2,
    "Finance": 1_2_2_5_0,
    "Politics": 1_6_3_6_0,
    "Scary": 1_9_1_9_8_5,
    "Support": 1_2_6_5_4,
    "Technologies": 3_2_5_1_6,
    "Teenage": 6_6_1_6_0,
    "Event": 3_2_7_6_9,
    "Learned": 6_7_4_6_0,
    "Notion": 1_8_2_7_7_0,
    "Wikipedia": 3_7_5_8_3,
    "Books": 6_6_6_5,
    "Extract": 7_6_0_5_0,
    "Confessions": 1_0_2_7_0_1,
    "Conspiracy": 7_5_9_3_2,
    "Links": 6_3_6_7_4,
    "Narcissus": 1_5_0_4_2_5,
    "Relationship": 5_4_7_6_6,
    "Relationships": 1_3_4_7_9_6,
    "Reviews": 4_1_6_7_1,
    "News": 4_2_5_6,
    "Translation": 2_6_8_2_0,
    "multilingual": 1_2_8_4_0_6,
}






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Union[str, Any]:
    __lowerCamelCase            =    set()
    __lowerCamelCase            =    word[0]
    for char in word[1:]:
        pairs.add((prev_char, char)							)
        __lowerCamelCase            =    char

    __lowerCamelCase            =    set(UpperCamelCase__							)
    return pairs




class        a__ (  UpperCAmelCase__					):
     lowerCamelCase    :	str				=VOCAB_FILES_NAMES
     lowerCamelCase    :	Optional[int]				=PRETRAINED_VOCAB_FILES_MAP
     lowerCamelCase    :	Tuple				=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     lowerCamelCase    :	Union[str, Any]				=CONTROL_CODES






     def __init__(					self						:						Optional[Any]    ,							a						:						str    ,							a						:						List[str]    ,							a						:						Optional[int]="<unk>"    ,							**a						:						Any       ):





         """simple docstring"""



         super().__init__(unk_token=a    ,							**a       )

         with open(a    ,							encoding='''utf-8'''       ) as vocab_handle:
             __lowerCamelCase            =    json.load(a       )
         __lowerCamelCase            =    {v: k for k, v in self.encoder.items()}
         with open(a    ,							encoding='''utf-8'''       ) as merges_handle:
             __lowerCamelCase            =    merges_handle.read().split('''\n'''       )[1:-1]
         __lowerCamelCase            =    [tuple(merge.split()       ) for merge in merges]
         __lowerCamelCase            =    dict(zip(a    ,							range(len(a       )       )       )       )
         __lowerCamelCase            =    {}






     @property
     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]       ):





         """simple docstring"""



         return len(self.encoder       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						int       ):





         """simple docstring"""



         return dict(self.encoder    ,							**self.added_tokens_encoder       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[Any]    ,							a						:						Dict       ):





         """simple docstring"""



         if token in self.cache:
             return self.cache[token]
         __lowerCamelCase            =    tuple(a       )
         __lowerCamelCase            =    tuple(list(word[:-1]       ) + [word[-1] + '''</w>''']       )
         __lowerCamelCase            =    get_pairs(a       )

         if not pairs:
             return token

         while True:
             __lowerCamelCase            =    min(a    ,							key=lambda a       : self.bpe_ranks.get(a    ,							float('''inf'''       )       )       )
             if bigram not in self.bpe_ranks:
                 break
             __lowerCamelCase      ,    __lowerCamelCase            =    bigram
             __lowerCamelCase            =    []
             __lowerCamelCase            =    0
             while i < len(a       ):
                 try:
                     __lowerCamelCase            =    word.index(a    ,							a       )
                 except ValueError:
                     new_word.extend(word[i:]       )
                     break
                 else:
                     new_word.extend(word[i:j]       )
                     __lowerCamelCase            =    j

                 if word[i] == first and i < len(a       ) - 1 and word[i + 1] == second:
                     new_word.append(first + second       )
                     i += 2
                 else:
                     new_word.append(word[i]       )
                     i += 1
             __lowerCamelCase            =    tuple(a       )
             __lowerCamelCase            =    new_word
             if len(a       ) == 1:
                 break
             else:
                 __lowerCamelCase            =    get_pairs(a       )
         __lowerCamelCase            =    '''@@ '''.join(a       )
         __lowerCamelCase            =    word[:-4]
         __lowerCamelCase            =    word
         return word






     def   SCREAMING_SNAKE_CASE__  (					self						:						str    ,							a						:						List[str]       ):





         """simple docstring"""



         __lowerCamelCase            =    []

         __lowerCamelCase            =    re.findall(R'''\S+\n?'''    ,							a       )

         for token in words:
             split_tokens.extend(list(self.bpe(a       ).split(''' '''       )       )       )
         return split_tokens






     def   SCREAMING_SNAKE_CASE__  (					self						:						Any    ,							a						:						Dict       ):





         """simple docstring"""



         return self.encoder.get(a    ,							self.encoder.get(self.unk_token       )       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[Any]    ,							a						:						Optional[Any]       ):





         """simple docstring"""



         return self.decoder.get(a    ,							self.unk_token       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]    ,							a						:						Dict       ):





         """simple docstring"""



         __lowerCamelCase            =    ''' '''.join(a       ).replace('''@@ '''    ,							''''''       ).strip()
         return out_string






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[Any]    ,							a						:						str    ,							a						:						Optional[str] = None       ):





         """simple docstring"""



         if not os.path.isdir(a       ):
             logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""       )
             return
         __lowerCamelCase            =    os.path.join(
             a    ,							(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']       )
         __lowerCamelCase            =    os.path.join(
             a    ,							(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']       )

         with open(a    ,							'''w'''    ,							encoding='''utf-8'''       ) as f:
             f.write(json.dumps(self.encoder    ,							indent=2    ,							sort_keys=a    ,							ensure_ascii=a       ) + '''\n'''       )

         __lowerCamelCase            =    0
         with open(a    ,							'''w'''    ,							encoding='''utf-8'''       ) as writer:
             writer.write('''#version: 0.2\n'''       )
             for bpe_tokens, token_index in sorted(self.bpe_ranks.items()    ,							key=lambda a       : kv[1]       ):
                 if index != token_index:
                     logger.warning(
                         f"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
                         ''' Please check that the tokenizer is not corrupted!'''       )
                     __lowerCamelCase            =    token_index
                 writer.write(''' '''.join(a       ) + '''\n'''       )
                 index += 1

         return vocab_file, merge_file

     # def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
     #     filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
     #     tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
     #     tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
     #     return ''.join(tokens_generated_so_far)
