


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	):
      UpperCamelCase__				:    Union[str, Any] =ConsistencyModelPipeline
      UpperCamelCase__				:    Dict =UNCONDITIONAL_IMAGE_GENERATION_PARAMS
      UpperCamelCase__				:    Dict =UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






      @property
      def 			lowerCamelCase			(   self		:	List[str]					)   ->		List[Any]:
             """simple docstring"""




             _lowerCamelCase			:						Union[str, Any]            =UNetaDModel.from_pretrained(
                 'diffusers/consistency-models-test'     ,  subfolder='test_unet'     ,  )
             return unet






      @property
      def 			lowerCamelCase			(   self		:	str					)   ->		Optional[int]:
             """simple docstring"""




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






      def 			lowerCamelCase			(   self		:	Tuple     ,  lowercase_		:	Optional[int]=False					)   ->		Dict:
             """simple docstring"""




             if class_cond:
                    _lowerCamelCase			:						Optional[int]            =self.dummy_cond_unet
             else:
                    _lowerCamelCase			:						Dict            =self.dummy_uncond_unet

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

             _lowerCamelCase			:						Tuple            ={
                 'unet': unet,
                 'scheduler': scheduler,
             }

             return components






      def 			lowerCamelCase			(   self		:	str     ,  lowercase_		:	str     ,  lowercase_		:	Optional[Any]=0					)   ->		int:
             """simple docstring"""




             if str(lowercase__					).startswith('mps'					):
                    _lowerCamelCase			:						int            =torch.manual_seed(lowercase__					)
             else:
                    _lowerCamelCase			:						Dict            =torch.Generator(device=lowercase__					).manual_seed(lowercase__					)

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

             return inputs






      def 			lowerCamelCase			(   self		:	List[Any]					)   ->		Optional[Any]:
             """simple docstring"""




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

             _lowerCamelCase			:						int            =self.get_dummy_inputs(lowercase__					)
             _lowerCamelCase			:						str            =pipe(**lowercase__					).images
             assert image.shape == (1, 32, 32, 3)

             _lowerCamelCase			:						Any            =image[0, -3:, -3:, -1]
             _lowerCamelCase			:						Union[str, 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 			lowerCamelCase			(   self		:	List[str]					)   ->		List[str]:
             """simple docstring"""




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

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

             _lowerCamelCase			:						Optional[int]            =image[0, -3:, -3:, -1]
             _lowerCamelCase			:						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 			lowerCamelCase			(   self		:	int					)   ->		Optional[int]:
             """simple docstring"""




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

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

             _lowerCamelCase			:						List[Any]            =image[0, -3:, -3:, -1]
             _lowerCamelCase			:						Union[str, 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 			lowerCamelCase			(   self		:	Any					)   ->		str:
             """simple docstring"""




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

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

             _lowerCamelCase			:						List[Any]            =image[0, -3:, -3:, -1]
             _lowerCamelCase			:						Tuple            =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 A    (  unittest.TestCase	):






      def 			lowerCamelCase			(   self		:	int					)   ->		Dict:
             """simple docstring"""




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






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




             _lowerCamelCase			:						str            =torch.manual_seed(lowercase__					)

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

             if get_fixed_latents:
                    _lowerCamelCase			:						Any            =self.get_fixed_latents(seed=lowercase__     ,  device=lowercase__     ,  dtype=lowercase__     ,  shape=lowercase__					)
                    _lowerCamelCase			:						str            =latents

             return inputs






      def 			lowerCamelCase			(   self		:	Tuple     ,  lowercase_		:	Dict=0     ,  lowercase_		:	Any="cpu"     ,  lowercase_		:	Union[str, Any]=torch.floataa     ,  lowercase_		:	Tuple=(1, 3, 64, 64)					)   ->		int:
             """simple docstring"""




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






      def 			lowerCamelCase			(   self		:	Optional[Any]					)   ->		Optional[Any]:
             """simple docstring"""




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

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

             _lowerCamelCase			:						Tuple            =image[0, -3:, -3:, -1]

             _lowerCamelCase			:						Optional[int]            =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 			lowerCamelCase			(   self		:	int					)   ->		Tuple:
             """simple docstring"""




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

             _lowerCamelCase			:						str            =self.get_inputs()
             _lowerCamelCase			:						Optional[Any]            =1
             _lowerCamelCase			:						int            =None
             _lowerCamelCase			:						Tuple            =pipe(**lowercase__					).images
             assert image.shape == (1, 64, 64, 3)

             _lowerCamelCase			:						Tuple            =image[0, -3:, -3:, -1]

             _lowerCamelCase			:						Optional[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 			lowerCamelCase			(   self		:	Union[str, Any]					)   ->		Tuple:
             """simple docstring"""




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

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

             _lowerCamelCase			:						str            =image[0, -3:, -3:, -1]

             _lowerCamelCase			:						Tuple            =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 			lowerCamelCase			(   self		:	List[str]					)   ->		Union[str, Any]:
             """simple docstring"""




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

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

             _lowerCamelCase			:						Union[str, Any]            =image[0, -3:, -3:, -1]

             _lowerCamelCase			:						List[str]            =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


_lowercase:	Optional[Any]       			=		logging.get_logger(__name__)

_lowercase:	Tuple       			=		{
    '''andreasmadsen/efficient_mlm_m0.40''': (
        '''https://huggingface.co/andreasmadsen/efficient_mlm_m0.40/resolve/main/config.json'''
    ),
}



class        lowerCamelCase__				(   UpperCAmelCase				):
			UpperCamelCase__    							="roberta-prelayernorm"





			def __init__(       self	:	int  ,  lowercase__	:	int=5_02_65  ,  lowercase__	:	str=7_68  ,  lowercase__	:	List[Any]=12  ,  lowercase__	:	Any=12  ,  lowercase__	:	Optional[Any]=30_72  ,  lowercase__	:	List[Any]="gelu"  ,  lowercase__	:	str=0.1  ,  lowercase__	:	List[Any]=0.1  ,  lowercase__	:	Tuple=5_12  ,  lowercase__	:	Optional[int]=2  ,  lowercase__	:	Dict=0.0_2  ,  lowercase__	:	Tuple=1e-12  ,  lowercase__	:	Optional[int]=1  ,  lowercase__	:	Optional[int]=0  ,  lowercase__	:	Union[str, Any]=2  ,  lowercase__	:	Any="absolute"  ,  lowercase__	:	Any=True  ,  lowercase__	:	Any=None  ,  **lowercase__	:	Optional[int]  ,  ):
										super().__init__(pad_token_id=lowercase__  ,  bos_token_id=lowercase__  ,  eos_token_id=lowercase__  ,  **lowercase__						)

										_lowerCAmelCase							=							vocab_size
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							hidden_act
										_lowerCAmelCase							=							intermediate_size
										_lowerCAmelCase							=							hidden_dropout_prob
										_lowerCAmelCase							=							attention_probs_dropout_prob
										_lowerCAmelCase							=							max_position_embeddings
										_lowerCAmelCase							=							type_vocab_size
										_lowerCAmelCase							=							initializer_range
										_lowerCAmelCase							=							layer_norm_eps
										_lowerCAmelCase							=							position_embedding_type
										_lowerCAmelCase							=							use_cache
										_lowerCAmelCase							=							classifier_dropout







class        lowerCamelCase__				(   UpperCAmelCase				):





			@property
			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										if self.task == "multiple-choice":
																	_lowerCAmelCase							=							{0: 'batch', 1: 'choice', 2: 'sequence'}
										else:
																	_lowerCAmelCase							=							{0: 'batch', 1: 'sequence'}
										return OrderedDict(
										    [
										        ('input_ids', dynamic_axis),
										        ('attention_mask', dynamic_axis),
										    ]						)






