




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 ):
  _UpperCamelCase  :     Optional[Any]					=     ConsistencyModelPipeline
  _UpperCamelCase  :     int					=     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  _UpperCamelCase  :     int					=     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

  # Override required_optional_params to remove num_images_per_prompt
  _UpperCamelCase  :     List[Any]					=     frozenset(
      [
          '''num_inference_steps''',
          '''generator''',
          '''latents''',
          '''output_type''',
          '''return_dict''',
          '''callback''',
          '''callback_steps''',
      ] )






  @property
  def    __a     (    self						:      Optional[int] )						->       Tuple:






       """simple docstring"""







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






  @property
  def    __a     (    self						:      Optional[Any] )						->       Tuple:






       """simple docstring"""







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






  def    __a     (    self						:      Optional[int] ,		_A						:      Any=False )						->       Optional[Any]:






       """simple docstring"""







       if class_cond:
            lowercase   :  Dict        =					self.dummy_cond_unet
       else:
            lowercase   :  Dict        =					self.dummy_uncond_unet

       # Default to CM multistep sampler
       lowercase   :  List[str]        =					CMStochasticIterativeScheduler(
           num_train_timesteps=40 ,		sigma_min=0.002 ,		sigma_max=80.0 ,		)

       lowercase   :  Dict        =					{
           '''unet''': unet,
           '''scheduler''': scheduler,
       }

       return components






  def    __a     (    self						:      str ,		_A						:      Optional[int] ,		_A						:      Any=0 )						->       Union[str, Any]:






       """simple docstring"""







       if str(_A ).startswith('''mps''' ):
            lowercase   :  Optional[int]        =					torch.manual_seed(_A )
       else:
            lowercase   :  Dict        =					torch.Generator(device=_A ).manual_seed(_A )

       lowercase   :  Union[str, Any]        =					{
           '''batch_size''': 1,
           '''num_inference_steps''': None,
           '''timesteps''': [22, 0],
           '''generator''': generator,
           '''output_type''': '''np''',
       }

       return inputs






  def    __a     (    self						:      List[Any] )						->       Union[str, Any]:






       """simple docstring"""







       lowercase   :  Any        =					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       lowercase   :  Union[str, Any]        =					self.get_dummy_components()
       lowercase   :  Union[str, Any]        =					ConsistencyModelPipeline(**_A )
       lowercase   :  Dict        =					pipe.to(_A )
       pipe.set_progress_bar_config(disable=_A )

       lowercase   :  str        =					self.get_dummy_inputs(_A )
       lowercase   :  Optional[int]        =					pipe(**_A ).images
       assert image.shape == (1, 32, 32, 3)

       lowercase   :  List[Any]        =					image[0, -3:, -3:, -1]
       lowercase   :  Union[str, Any]        =					np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004] )

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






  def    __a     (    self						:      Optional[int] )						->       Dict:






       """simple docstring"""







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

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

       lowercase   :  Optional[Any]        =					image[0, -3:, -3:, -1]
       lowercase   :  Tuple        =					np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004] )

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






  def    __a     (    self						:      str )						->       Dict:






       """simple docstring"""







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

       lowercase   :  List[Any]        =					self.get_dummy_inputs(_A )
       lowercase   :  Union[str, Any]        =					1
       lowercase   :  Optional[int]        =					None
       lowercase   :  Dict        =					pipe(**_A ).images
       assert image.shape == (1, 32, 32, 3)

       lowercase   :  Optional[int]        =					image[0, -3:, -3:, -1]
       lowercase   :  Tuple        =					np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987] )

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






  def    __a     (    self						:      List[str] )						->       List[Any]:






       """simple docstring"""







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

       lowercase   :  Any        =					self.get_dummy_inputs(_A )
       lowercase   :  List[Any]        =					1
       lowercase   :  Tuple        =					None
       lowercase   :  List[str]        =					0
       lowercase   :  Union[str, Any]        =					pipe(**_A ).images
       assert image.shape == (1, 32, 32, 3)

       lowercase   :  Dict        =					image[0, -3:, -3:, -1]
       lowercase   :  Union[str, Any]        =					np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987] )

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


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






  def    __a     (    self						:      Optional[int] )						->       Tuple:






       """simple docstring"""







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






  def    __a     (    self						:      Optional[Any] ,		_A						:      Any=0 ,		_A						:      Optional[Any]=False ,		_A						:      Any="cpu" ,		_A						:      Union[str, Any]=torch.floataa ,		_A						:      List[Any]=(1, 3, 64, 64) )						->       str:






       """simple docstring"""







       lowercase   :  List[str]        =					torch.manual_seed(_A )

       lowercase   :  Union[str, Any]        =					{
           '''num_inference_steps''': None,
           '''timesteps''': [22, 0],
           '''class_labels''': 0,
           '''generator''': generator,
           '''output_type''': '''np''',
       }

       if get_fixed_latents:
            lowercase   :  List[Any]        =					self.get_fixed_latents(seed=_A ,		device=_A ,		dtype=_A ,		shape=_A )
            lowercase   :  Optional[Any]        =					latents

       return inputs






  def    __a     (    self						:      List[Any] ,		_A						:      int=0 ,		_A						:      Any="cpu" ,		_A						:      Any=torch.floataa ,		_A						:      Dict=(1, 3, 64, 64) )						->       Dict:






       """simple docstring"""







       if type(_A ) == str:
            lowercase   :  Tuple        =					torch.device(_A )
       lowercase   :  str        =					torch.Generator(device=_A ).manual_seed(_A )
       lowercase   :  int        =					randn_tensor(_A ,		generator=_A ,		device=_A ,		dtype=_A )
       return latents






  def    __a     (    self						:      Any )						->       List[str]:






       """simple docstring"""







       lowercase   :  Union[str, Any]        =					UNetaDModel.from_pretrained('''diffusers/consistency_models''' ,		subfolder='''diffusers_cd_imagenet64_l2''' )
       lowercase   :  Optional[int]        =					CMStochasticIterativeScheduler(
           num_train_timesteps=40 ,		sigma_min=0.002 ,		sigma_max=80.0 ,		)
       lowercase   :  str        =					ConsistencyModelPipeline(unet=_A ,		scheduler=_A )
       pipe.to(torch_device=_A )
       pipe.set_progress_bar_config(disable=_A )

       lowercase   :  List[Any]        =					self.get_inputs()
       lowercase   :  Optional[Any]        =					pipe(**_A ).images
       assert image.shape == (1, 64, 64, 3)

       lowercase   :  int        =					image[0, -3:, -3:, -1]

       lowercase   :  Optional[Any]        =					np.array([0.0_888, 0.0_881, 0.0_666, 0.0_479, 0.0_292, 0.0_195, 0.0_201, 0.0_163, 0.0_254] )

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






  def    __a     (    self						:      str )						->       Union[str, Any]:






       """simple docstring"""







       lowercase   :  int        =					UNetaDModel.from_pretrained('''diffusers/consistency_models''' ,		subfolder='''diffusers_cd_imagenet64_l2''' )
       lowercase   :  Tuple        =					CMStochasticIterativeScheduler(
           num_train_timesteps=40 ,		sigma_min=0.002 ,		sigma_max=80.0 ,		)
       lowercase   :  Tuple        =					ConsistencyModelPipeline(unet=_A ,		scheduler=_A )
       pipe.to(torch_device=_A )
       pipe.set_progress_bar_config(disable=_A )

       lowercase   :  Dict        =					self.get_inputs()
       lowercase   :  int        =					1
       lowercase   :  List[str]        =					None
       lowercase   :  int        =					pipe(**_A ).images
       assert image.shape == (1, 64, 64, 3)

       lowercase   :  Any        =					image[0, -3:, -3:, -1]

       lowercase   :  str        =					np.array([0.0_340, 0.0_152, 0.0_063, 0.0_267, 0.0_221, 0.0_107, 0.0_416, 0.0_186, 0.0_217] )

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






  @require_torch_a
  def    __a     (    self						:      Union[str, Any] )						->       List[str]:






       """simple docstring"""







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

       lowercase   :  Union[str, Any]        =					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 ):
            lowercase   :  Union[str, Any]        =					pipe(**_A ).images
       assert image.shape == (1, 64, 64, 3)

       lowercase   :  int        =					image[0, -3:, -3:, -1]

       lowercase   :  Tuple        =					np.array([0.1_875, 0.1_428, 0.1_289, 0.2_151, 0.2_092, 0.1_477, 0.1_877, 0.1_641, 0.1_353] )

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






  @require_torch_a
  def    __a     (    self						:      Optional[Any] )						->       Any:






       """simple docstring"""







       lowercase   :  Dict        =					UNetaDModel.from_pretrained('''diffusers/consistency_models''' ,		subfolder='''diffusers_cd_imagenet64_l2''' )
       lowercase   :  Union[str, Any]        =					CMStochasticIterativeScheduler(
           num_train_timesteps=40 ,		sigma_min=0.002 ,		sigma_max=80.0 ,		)
       lowercase   :  Any        =					ConsistencyModelPipeline(unet=_A ,		scheduler=_A )
       pipe.to(torch_device=_A ,		torch_dtype=torch.floataa )
       pipe.set_progress_bar_config(disable=_A )

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

       lowercase   :  Any        =					image[0, -3:, -3:, -1]

       lowercase   :  int        =					np.array([0.1_663, 0.1_948, 0.2_275, 0.1_680, 0.1_204, 0.1_245, 0.1_858, 0.1_338, 0.2_095] )

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





'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase_					=   {"configuration_fnet": ["FNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "FNetConfig"]}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   ["FNetTokenizer"]

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   ["FNetTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   [
      "FNET_PRETRAINED_MODEL_ARCHIVE_LIST",
      "FNetForMaskedLM",
      "FNetForMultipleChoice",
      "FNetForNextSentencePrediction",
      "FNetForPreTraining",
      "FNetForQuestionAnswering",
      "FNetForSequenceClassification",
      "FNetForTokenClassification",
      "FNetLayer",
      "FNetModel",
      "FNetPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_fnet import FNetTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_fnet_fast import FNetTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_fnet import (
        FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        FNetForMaskedLM,
        FNetForMultipleChoice,
        FNetForNextSentencePrediction,
        FNetForPreTraining,
        FNetForQuestionAnswering,
        FNetForSequenceClassification,
        FNetForTokenClassification,
        FNetLayer,
        FNetModel,
        FNetPreTrainedModel,
    )


else:
  import sys

  UpperCAmelCase_					=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



