



'''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 				snake_case (UpperCamelCase		,      unittest.TestCase   ):
       lowerCAmelCase__  :Optional[Any]						   =					ConsistencyModelPipeline
       lowerCAmelCase__  :List[Any]						   =					UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       lowerCAmelCase__  :str						   =					UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



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



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



       def 							_a      (				self     ,UpperCAmelCase_=False      )  ->							List[Any]:
           if class_cond:
               lowercase__		       =						self.dummy_cond_unet
           else:
               lowercase__		       =						self.dummy_uncond_unet

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

           lowercase__		       =						{
               "unet": unet,
               "scheduler": scheduler,
           }

           return components



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_=0      )  ->							Dict:
           if str(UpperCAmelCase_      ).startswith("mps"      ):
               lowercase__		       =						torch.manual_seed(UpperCAmelCase_      )
           else:
               lowercase__		       =						torch.Generator(device=UpperCAmelCase_      ).manual_seed(UpperCAmelCase_      )

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

           return inputs



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowercase__		       =						self.get_dummy_components()
           lowercase__		       =						ConsistencyModelPipeline(**UpperCAmelCase_      )
           lowercase__		       =						pipe.to(UpperCAmelCase_      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

           lowercase__		       =						image[0, -3:, -3:, -1]
           lowercase__		       =						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 							_a      (				self      )  ->							List[Any]:
           lowercase__		       =						"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowercase__		       =						self.get_dummy_components(class_cond=UpperCAmelCase_      )
           lowercase__		       =						ConsistencyModelPipeline(**UpperCAmelCase_      )
           lowercase__		       =						pipe.to(UpperCAmelCase_      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

           lowercase__		       =						image[0, -3:, -3:, -1]
           lowercase__		       =						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 							_a      (				self      )  ->							Any:
           lowercase__		       =						"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowercase__		       =						self.get_dummy_components()
           lowercase__		       =						ConsistencyModelPipeline(**UpperCAmelCase_      )
           lowercase__		       =						pipe.to(UpperCAmelCase_      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

           lowercase__		       =						image[0, -3:, -3:, -1]
           lowercase__		       =						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 							_a      (				self      )  ->							Dict:
           lowercase__		       =						"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowercase__		       =						self.get_dummy_components(class_cond=UpperCAmelCase_      )
           lowercase__		       =						ConsistencyModelPipeline(**UpperCAmelCase_      )
           lowercase__		       =						pipe.to(UpperCAmelCase_      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

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



       def 							_a      (				self      )  ->							Dict:
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



       def 							_a      (				self     ,UpperCAmelCase_=0     ,UpperCAmelCase_=False     ,UpperCAmelCase_="cpu"     ,UpperCAmelCase_=torch.floataa     ,UpperCAmelCase_=(1, 3, 64, 64)      )  ->							str:
           lowercase__		       =						torch.manual_seed(UpperCAmelCase_      )

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

           if get_fixed_latents:
               lowercase__		       =						self.get_fixed_latents(seed=UpperCAmelCase_     ,device=UpperCAmelCase_     ,dtype=UpperCAmelCase_     ,shape=UpperCAmelCase_      )
               lowercase__		       =						latents

           return inputs



       def 							_a      (				self     ,UpperCAmelCase_=0     ,UpperCAmelCase_="cpu"     ,UpperCAmelCase_=torch.floataa     ,UpperCAmelCase_=(1, 3, 64, 64)      )  ->							str:
           if type(UpperCAmelCase_      ) == str:
               lowercase__		       =						torch.device(UpperCAmelCase_      )
           lowercase__		       =						torch.Generator(device=UpperCAmelCase_      ).manual_seed(UpperCAmelCase_      )
           lowercase__		       =						randn_tensor(UpperCAmelCase_     ,generator=UpperCAmelCase_     ,device=UpperCAmelCase_     ,dtype=UpperCAmelCase_      )
           return latents



       def 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						UNetaDModel.from_pretrained("diffusers/consistency_models"     ,subfolder="diffusers_cd_imagenet64_l2"      )
           lowercase__		       =						CMStochasticIterativeScheduler(
               num_train_timesteps=40     ,sigma_min=0.0_02     ,sigma_max=80.0     ,)
           lowercase__		       =						ConsistencyModelPipeline(unet=UpperCAmelCase_     ,scheduler=UpperCAmelCase_      )
           pipe.to(torch_device=UpperCAmelCase_      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

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

           lowercase__		       =						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 							_a      (				self      )  ->							str:
           lowercase__		       =						UNetaDModel.from_pretrained("diffusers/consistency_models"     ,subfolder="diffusers_cd_imagenet64_l2"      )
           lowercase__		       =						CMStochasticIterativeScheduler(
               num_train_timesteps=40     ,sigma_min=0.0_02     ,sigma_max=80.0     ,)
           lowercase__		       =						ConsistencyModelPipeline(unet=UpperCAmelCase_     ,scheduler=UpperCAmelCase_      )
           pipe.to(torch_device=UpperCAmelCase_      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

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

           lowercase__		       =						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 							_a      (				self      )  ->							str:
           lowercase__		       =						UNetaDModel.from_pretrained("diffusers/consistency_models"     ,subfolder="diffusers_cd_imagenet64_l2"      )
           lowercase__		       =						CMStochasticIterativeScheduler(
               num_train_timesteps=40     ,sigma_min=0.0_02     ,sigma_max=80.0     ,)
           lowercase__		       =						ConsistencyModelPipeline(unet=UpperCAmelCase_     ,scheduler=UpperCAmelCase_      )
           pipe.to(torch_device=UpperCAmelCase_     ,torch_dtype=torch.floataa      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

           lowercase__		       =						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_      ):
               lowercase__		       =						pipe(**UpperCAmelCase_      ).images
           assert image.shape == (1, 64, 64, 3)

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

           lowercase__		       =						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 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						UNetaDModel.from_pretrained("diffusers/consistency_models"     ,subfolder="diffusers_cd_imagenet64_l2"      )
           lowercase__		       =						CMStochasticIterativeScheduler(
               num_train_timesteps=40     ,sigma_min=0.0_02     ,sigma_max=80.0     ,)
           lowercase__		       =						ConsistencyModelPipeline(unet=UpperCAmelCase_     ,scheduler=UpperCAmelCase_      )
           pipe.to(torch_device=UpperCAmelCase_     ,torch_dtype=torch.floataa      )
           pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

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

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

           lowercase__		       =						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 os
from typing import Dict, List, Union

import tensorflow as tf
from keras_nlp.tokenizers import BytePairTokenizer
from tensorflow_text import pad_model_inputs

from .tokenization_gpta import GPTaTokenizer


class 				snake_case (tf.keras.layers.Layer   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_ = None     ,UpperCAmelCase_ = None      )  ->							List[Any]:
           super().__init__()
           lowercase__		       =						pad_token_id
           lowercase__		       =						max_length
           lowercase__		       =						vocab
           lowercase__		       =						merges
           lowercase__		       =						BytePairTokenizer(UpperCAmelCase_     ,UpperCAmelCase_     ,sequence_length=UpperCAmelCase_      )



       @classmethod
       def 							_a      (				cls     ,UpperCAmelCase_     ,*UpperCAmelCase_     ,**UpperCAmelCase_      )  ->							List[str]:
           lowercase__		       =						[" ".join(UpperCAmelCase_      ) for m in tokenizer.bpe_ranks.keys()]
           lowercase__		       =						tokenizer.get_vocab()
           return cls(UpperCAmelCase_     ,UpperCAmelCase_     ,*UpperCAmelCase_     ,**UpperCAmelCase_      )



       @classmethod
       def 							_a      (				cls     ,UpperCAmelCase_     ,*UpperCAmelCase_     ,**UpperCAmelCase_      )  ->							Optional[Any]:
           lowercase__		       =						GPTaTokenizer.from_pretrained(UpperCAmelCase_     ,*UpperCAmelCase_     ,**UpperCAmelCase_      )
           return cls.from_tokenizer(UpperCAmelCase_     ,*UpperCAmelCase_     ,**UpperCAmelCase_      )



       @classmethod
       def 							_a      (				cls     ,UpperCAmelCase_      )  ->							Optional[Any]:
           return cls(**UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Tuple:
           return {
               "vocab": self.vocab,
               "merges": self.merges,
               "max_length": self.max_length,
               "pad_token_id": self.pad_token_id,
           }







       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_ = None      )  ->							Tuple:
           lowercase__		       =						self.tf_tokenizer(UpperCAmelCase_      )
           lowercase__		       =						tf.ones_like(UpperCAmelCase_      )

           if self.pad_token_id is not None:
               # pad the tokens up to max length
               lowercase__		       =						max_length if max_length is not None else self.max_length

               if max_length is not None:
                   lowercase__       ,       lowercase__		       =						pad_model_inputs(
                       UpperCAmelCase_     ,max_seq_length=UpperCAmelCase_     ,pad_value=self.pad_token_id      )

           return {"attention_mask": attention_mask, "input_ids": input_ids}




