

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  UpperCAmelCase__	( A__		,  unittest.TestCase ):






  """simple docstring"""


  a     =      ConsistencyModelPipeline
  a     =      UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  a     =      UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






  @property
  def    lowercase_	(							self		:	Dict      )	->							str:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained(
           '''diffusers/consistency-models-test'''			,     subfolder='''test_unet'''			,     )
       return unet






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






  def    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	str=False      )	->							str:
       if class_cond:
            SCREAMING_SNAKE_CASE__					=					self.dummy_cond_unet
       else:
            SCREAMING_SNAKE_CASE__					=					self.dummy_uncond_unet

       # Default to CM multistep sampler
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )

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

       return components






  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Union[str, Any]=0      )	->							Any:
       if str(__lowerCamelCase      ).startswith('''mps'''      ):
            SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )
       else:
            SCREAMING_SNAKE_CASE__					=					torch.Generator(device=__lowerCamelCase      ).manual_seed(__lowerCamelCase      )

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

       return inputs






  def    lowercase_	(							self		:	Optional[Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					pipe.to(__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE__					=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]      )

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






  def    lowercase_	(							self		:	Optional[Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components(class_cond=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					pipe.to(__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE__					=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]      )

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






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

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

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE__					=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]      )

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







  def    lowercase_	(							self		:	str      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components(class_cond=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					pipe.to(__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE__					=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]      )

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






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






  """simple docstring"""


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






  def    lowercase_	(							self		:	Optional[Any]			,     __lowerCamelCase		:	List[Any]=0			,     __lowerCamelCase		:	int=False			,     __lowerCamelCase		:	Dict="cpu"			,     __lowerCamelCase		:	Union[str, Any]=torch.floataa			,     __lowerCamelCase		:	Optional[Any]=(1, 3, 64, 64)      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )

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

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

       return inputs






  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	Any=0			,     __lowerCamelCase		:	List[str]="cpu"			,     __lowerCamelCase		:	Optional[Any]=torch.floataa			,     __lowerCamelCase		:	Dict=(1, 3, 64, 64)      )	->							List[str]:
       if type(__lowerCamelCase      ) == str:
            SCREAMING_SNAKE_CASE__					=					torch.device(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					torch.Generator(device=__lowerCamelCase      ).manual_seed(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					randn_tensor(__lowerCamelCase			,     generator=__lowerCamelCase			,     device=__lowerCamelCase			,     dtype=__lowerCamelCase      )
       return latents






  def    lowercase_	(							self		:	int      )	->							int:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]      )

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






  def    lowercase_	(							self		:	List[str]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]      )

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






  @require_torch_a
  def    lowercase_	(							self		:	Optional[int]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase			,     torch_dtype=torch.floataa      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]      )

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







  @require_torch_a
  def    lowercase_	(							self		:	Any      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase			,     torch_dtype=torch.floataa      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]      )

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





from typing import List

import datasets
from datasets.tasks import AudioClassification

from ..folder_based_builder import folder_based_builder


_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    datasets.utils.logging.get_logger(__name__)






class  UpperCAmelCase__	( folder_based_builder.FolderBasedBuilderConfig ):






  """simple docstring"""


  a     =      None
  a     =      None






class  UpperCAmelCase__	( folder_based_builder.FolderBasedBuilder ):






  """simple docstring"""


  a     =      datasets.Audio()
  a     =      "audio"
  a     =      AudioFolderConfig
  a     =      42  # definition at the bottom of the script
  a     =      AudioClassification(audio_column="audio"		,  label_column="label" )



_SCREAMING_SNAKE_CASE  :						str         =    [
    '''.aiff''',
    '''.au''',
    '''.avr''',
    '''.caf''',
    '''.flac''',
    '''.htk''',
    '''.svx''',
    '''.mat4''',
    '''.mat5''',
    '''.mpc2k''',
    '''.ogg''',
    '''.paf''',
    '''.pvf''',
    '''.raw''',
    '''.rf64''',
    '''.sd2''',
    '''.sds''',
    '''.ircam''',
    '''.voc''',
    '''.w64''',
    '''.wav''',
    '''.nist''',
    '''.wavex''',
    '''.wve''',
    '''.xi''',
    '''.mp3''',
    '''.opus''',
]
_SCREAMING_SNAKE_CASE  :						List[Any]         =    AUDIO_EXTENSIONS


