


'''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 				__UpperCamelCase							(						UpperCamelCase_						, unittest.TestCase   ):
	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  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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




	@property
	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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




	def  __UpperCAmelCase				(  self		,							__a=False				):

						'''simple docstring'''





						if class_cond:
											__a     :			Optional[Any]          =    self.dummy_cond_unet
						else:
											__a     :			Dict          =    self.dummy_uncond_unet

						# Default to CM multistep sampler
						__a     :			Optional[Any]          =    CMStochasticIterativeScheduler(
						    num_train_timesteps=40		,							sigma_min=0.002		,							sigma_max=80.0		,							)

						__a     :			Optional[Any]          =    {
						    '''unet''': unet,
						    '''scheduler''': scheduler,
						}

						return components




	def  __UpperCAmelCase				(  self		,							__a		,							__a=0				):

						'''simple docstring'''





						if str(__SCREAMING_SNAKE_CASE				).startswith('mps'				):
											__a     :			Dict          =    torch.manual_seed(__SCREAMING_SNAKE_CASE				)
						else:
											__a     :			Union[str, Any]          =    torch.Generator(device=__SCREAMING_SNAKE_CASE				).manual_seed(__SCREAMING_SNAKE_CASE				)

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

						return inputs




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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

						__a     :			List[str]          =    self.get_dummy_inputs(__SCREAMING_SNAKE_CASE				)
						__a     :			List[str]          =    pipe(**__SCREAMING_SNAKE_CASE				).images
						assert image.shape == (1, 32, 32, 3)

						__a     :			Dict          =    image[0, -3:, -3:, -1]
						__a     :			Any          =    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  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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

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

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

						'''simple docstring'''





						__a     :			List[Any]          =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
						__a     :			List[str]          =    self.get_dummy_components()
						__a     :			Any          =    ConsistencyModelPipeline(**__SCREAMING_SNAKE_CASE				)
						__a     :			str          =    pipe.to(__SCREAMING_SNAKE_CASE				)
						pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE				)

						__a     :			Any          =    self.get_dummy_inputs(__SCREAMING_SNAKE_CASE				)
						__a     :			str          =    1
						__a     :			Optional[Any]          =    None
						__a     :			List[str]          =    pipe(**__SCREAMING_SNAKE_CASE				).images
						assert image.shape == (1, 32, 32, 3)

						__a     :			str          =    image[0, -3:, -3:, -1]
						__a     :			List[Any]          =    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  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			List[str]          =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
						__a     :			str          =    self.get_dummy_components(class_cond=__SCREAMING_SNAKE_CASE				)
						__a     :			Optional[int]          =    ConsistencyModelPipeline(**__SCREAMING_SNAKE_CASE				)
						__a     :			List[str]          =    pipe.to(__SCREAMING_SNAKE_CASE				)
						pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE				)

						__a     :			int          =    self.get_dummy_inputs(__SCREAMING_SNAKE_CASE				)
						__a     :			str          =    1
						__a     :			Union[str, Any]          =    None
						__a     :			Any          =    0
						__a     :			str          =    pipe(**__SCREAMING_SNAKE_CASE				).images
						assert image.shape == (1, 32, 32, 3)

						__a     :			Optional[int]          =    image[0, -3:, -3:, -1]
						__a     :			int          =    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   ):




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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




	def  __UpperCAmelCase				(  self		,							__a=0		,							__a=False		,							__a="cpu"		,							__a=torch.floataa		,							__a=(1, 3, 64, 64)				):

						'''simple docstring'''





						__a     :			Tuple          =    torch.manual_seed(__SCREAMING_SNAKE_CASE				)

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

						if get_fixed_latents:
											__a     :			List[str]          =    self.get_fixed_latents(seed=__SCREAMING_SNAKE_CASE		,							device=__SCREAMING_SNAKE_CASE		,							dtype=__SCREAMING_SNAKE_CASE		,							shape=__SCREAMING_SNAKE_CASE				)
											__a     :			Tuple          =    latents

						return inputs




	def  __UpperCAmelCase				(  self		,							__a=0		,							__a="cpu"		,							__a=torch.floataa		,							__a=(1, 3, 64, 64)				):

						'''simple docstring'''





						if type(__SCREAMING_SNAKE_CASE				) == str:
											__a     :			List[str]          =    torch.device(__SCREAMING_SNAKE_CASE				)
						__a     :			List[Any]          =    torch.Generator(device=__SCREAMING_SNAKE_CASE				).manual_seed(__SCREAMING_SNAKE_CASE				)
						__a     :			Union[str, Any]          =    randn_tensor(__SCREAMING_SNAKE_CASE		,							generator=__SCREAMING_SNAKE_CASE		,							device=__SCREAMING_SNAKE_CASE		,							dtype=__SCREAMING_SNAKE_CASE				)
						return latents




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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

						__a     :			Dict          =    self.get_inputs()
						__a     :			Dict          =    pipe(**__SCREAMING_SNAKE_CASE				).images
						assert image.shape == (1, 64, 64, 3)

						__a     :			Optional[Any]          =    image[0, -3:, -3:, -1]

						__a     :			str          =    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  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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

						__a     :			Dict          =    self.get_inputs()
						__a     :			Optional[int]          =    1
						__a     :			Tuple          =    None
						__a     :			Tuple          =    pipe(**__SCREAMING_SNAKE_CASE				).images
						assert image.shape == (1, 64, 64, 3)

						__a     :			Dict          =    image[0, -3:, -3:, -1]

						__a     :			Union[str, Any]          =    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  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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

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

						__a     :			List[Any]          =    image[0, -3:, -3:, -1]

						__a     :			Union[str, Any]          =    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  __UpperCAmelCase				(  self				):

						'''simple docstring'''





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

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

						__a     :			List[Any]          =    image[0, -3:, -3:, -1]

						__a     :			Optional[int]          =    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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCamelCase  =logging.get_logger(__name__)

lowerCamelCase  ={
    "distilbert-base-uncased": "https://huggingface.co/distilbert-base-uncased/resolve/main/config.json",
    "distilbert-base-uncased-distilled-squad": (
        "https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json"
    ),
    "distilbert-base-cased": "https://huggingface.co/distilbert-base-cased/resolve/main/config.json",
    "distilbert-base-cased-distilled-squad": (
        "https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json"
    ),
    "distilbert-base-german-cased": "https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json",
    "distilbert-base-multilingual-cased": (
        "https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json"
    ),
    "distilbert-base-uncased-finetuned-sst-2-english": (
        "https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json"
    ),
}




class 					_lowerCamelCase			(   UpperCamelCase_			):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		 =	'''distilbert'''
     SCREAMING_SNAKE_CASE_		 =	{
         '''hidden_size''': '''dim''',
         '''num_attention_heads''': '''n_heads''',
         '''num_hidden_layers''': '''n_layers''',
     }






     def __init__(				self    ,    __SCREAMING_SNAKE_CASE=3_0_5_2_2    ,    __SCREAMING_SNAKE_CASE=5_1_2    ,    __SCREAMING_SNAKE_CASE=False    ,    __SCREAMING_SNAKE_CASE=6    ,    __SCREAMING_SNAKE_CASE=1_2    ,    __SCREAMING_SNAKE_CASE=7_6_8    ,    __SCREAMING_SNAKE_CASE=4 * 7_6_8    ,    __SCREAMING_SNAKE_CASE=0.1    ,    __SCREAMING_SNAKE_CASE=0.1    ,    __SCREAMING_SNAKE_CASE="gelu"    ,    __SCREAMING_SNAKE_CASE=0.02    ,    __SCREAMING_SNAKE_CASE=0.1    ,    __SCREAMING_SNAKE_CASE=0.2    ,    __SCREAMING_SNAKE_CASE=0    ,    **__SCREAMING_SNAKE_CASE    ,    )	-> Optional[Any]:




           """simple docstring"""







           UpperCamelCase__     : List[Any]											=					vocab_size
           UpperCamelCase__     : Dict											=					max_position_embeddings
           UpperCamelCase__     : Any											=					sinusoidal_pos_embds
           UpperCamelCase__     : Dict											=					n_layers
           UpperCamelCase__     : List[Any]											=					n_heads
           UpperCamelCase__     : Dict											=					dim
           UpperCamelCase__     : Dict											=					hidden_dim
           UpperCamelCase__     : Optional[int]											=					dropout
           UpperCamelCase__     : Optional[Any]											=					attention_dropout
           UpperCamelCase__     : Tuple											=					activation
           UpperCamelCase__     : Optional[int]											=					initializer_range
           UpperCamelCase__     : Optional[int]											=					qa_dropout
           UpperCamelCase__     : str											=					seq_classif_dropout
           super().__init__(**__SCREAMING_SNAKE_CASE    ,    pad_token_id=__SCREAMING_SNAKE_CASE	)




class 					_lowerCamelCase			(   UpperCamelCase_			):





     """simple docstring"""







     @property
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Mapping[str, Mapping[int, str]]:




           """simple docstring"""







           if self.task == "multiple-choice":
                 UpperCamelCase__     : Union[str, Any]											=					{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
           else:
                 UpperCamelCase__     : int											=					{0: '''batch''', 1: '''sequence'''}
           return OrderedDict(
               [
                   ('''input_ids''', dynamic_axis),
                   ('''attention_mask''', dynamic_axis),
               ]	)


