


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       __lowerCAmelCase   (      A    ,						unittest.TestCase						):
			UpperCamelCase        =					ConsistencyModelPipeline
			UpperCamelCase        =					UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			UpperCamelCase        =					UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



			@property
			def    _lowerCamelCase	(						self       :  List[str])       ->   Union[str, Any]:



						"""simple docstring"""






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



			@property
			def    _lowerCamelCase	(						self       :  Optional[int])       ->   Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       UNetaDModel.from_pretrained(
						    'diffusers/consistency-models-test'			,		subfolder='test_unet_class_cond'			,		)
						return unet



			def    _lowerCamelCase	(						self       :  Any			,		A       :  Tuple=False)       ->   int:



						"""simple docstring"""






						if class_cond:
									_UpperCAmelCase			=       self.dummy_cond_unet
						else:
									_UpperCAmelCase			=       self.dummy_uncond_unet

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

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

						return components



			def    _lowerCamelCase	(						self       :  str			,		A       :  Optional[Any]			,		A       :  List[Any]=0)       ->   int:



						"""simple docstring"""






						if str(A).startswith('mps'):
									_UpperCAmelCase			=       torch.manual_seed(A)
						else:
									_UpperCAmelCase			=       torch.Generator(device=A).manual_seed(A)

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

						return inputs



			def    _lowerCamelCase	(						self       :  Dict)       ->   Tuple:



						"""simple docstring"""






						_UpperCAmelCase			=       'cpu'  # ensure determinism for the device-dependent torch.Generator
						_UpperCAmelCase			=       self.get_dummy_components()
						_UpperCAmelCase			=       ConsistencyModelPipeline(**A)
						_UpperCAmelCase			=       pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

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

						_UpperCAmelCase			=       image[0, -3:, -3:, -1]
						_UpperCAmelCase			=       np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4])

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



			def    _lowerCamelCase	(						self       :  Optional[Any])       ->   Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       'cpu'  # ensure determinism for the device-dependent torch.Generator
						_UpperCAmelCase			=       self.get_dummy_components(class_cond=A)
						_UpperCAmelCase			=       ConsistencyModelPipeline(**A)
						_UpperCAmelCase			=       pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

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

						_UpperCAmelCase			=       image[0, -3:, -3:, -1]
						_UpperCAmelCase			=       np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4])

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



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



						"""simple docstring"""






						_UpperCAmelCase			=       'cpu'  # ensure determinism for the device-dependent torch.Generator
						_UpperCAmelCase			=       self.get_dummy_components()
						_UpperCAmelCase			=       ConsistencyModelPipeline(**A)
						_UpperCAmelCase			=       pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

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

						_UpperCAmelCase			=       image[0, -3:, -3:, -1]
						_UpperCAmelCase			=       np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7])

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



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



						"""simple docstring"""






						_UpperCAmelCase			=       'cpu'  # ensure determinism for the device-dependent torch.Generator
						_UpperCAmelCase			=       self.get_dummy_components(class_cond=A)
						_UpperCAmelCase			=       ConsistencyModelPipeline(**A)
						_UpperCAmelCase			=       pipe.to(A)
						pipe.set_progress_bar_config(disable=A)

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

						_UpperCAmelCase			=       image[0, -3:, -3:, -1]
						_UpperCAmelCase			=       np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7])

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







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



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



						"""simple docstring"""






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



			def    _lowerCamelCase	(						self       :  Tuple			,		A       :  Any=0			,		A       :  int=False			,		A       :  Optional[int]="cpu"			,		A       :  Optional[Any]=torch.floataa			,		A       :  List[str]=(1, 3, 64, 64))       ->   List[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       torch.manual_seed(A)

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

						if get_fixed_latents:
									_UpperCAmelCase			=       self.get_fixed_latents(seed=A			,		device=A			,		dtype=A			,		shape=A)
									_UpperCAmelCase			=       latents

						return inputs



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



						"""simple docstring"""






						if type(A) == str:
									_UpperCAmelCase			=       torch.device(A)
						_UpperCAmelCase			=       torch.Generator(device=A).manual_seed(A)
						_UpperCAmelCase			=       randn_tensor(A			,		generator=A			,		device=A			,		dtype=A)
						return latents



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



						"""simple docstring"""






						_UpperCAmelCase			=       UNetaDModel.from_pretrained('diffusers/consistency_models'			,		subfolder='diffusers_cd_imagenet64_l2')
						_UpperCAmelCase			=       CMStochasticIterativeScheduler(
						    num_train_timesteps=40			,		sigma_min=0.0_0_2			,		sigma_max=80.0			,		)
						_UpperCAmelCase			=       ConsistencyModelPipeline(unet=A			,		scheduler=A)
						pipe.to(torch_device=A)
						pipe.set_progress_bar_config(disable=A)

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

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

						_UpperCAmelCase			=       np.array([0.0_8_8_8, 0.0_8_8_1, 0.0_6_6_6, 0.0_4_7_9, 0.0_2_9_2, 0.0_1_9_5, 0.0_2_0_1, 0.0_1_6_3, 0.0_2_5_4])

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



			def    _lowerCamelCase	(						self       :  Optional[int])       ->   Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase			=       UNetaDModel.from_pretrained('diffusers/consistency_models'			,		subfolder='diffusers_cd_imagenet64_l2')
						_UpperCAmelCase			=       CMStochasticIterativeScheduler(
						    num_train_timesteps=40			,		sigma_min=0.0_0_2			,		sigma_max=80.0			,		)
						_UpperCAmelCase			=       ConsistencyModelPipeline(unet=A			,		scheduler=A)
						pipe.to(torch_device=A)
						pipe.set_progress_bar_config(disable=A)

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

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

						_UpperCAmelCase			=       np.array([0.0_3_4_0, 0.0_1_5_2, 0.0_0_6_3, 0.0_2_6_7, 0.0_2_2_1, 0.0_1_0_7, 0.0_4_1_6, 0.0_1_8_6, 0.0_2_1_7])

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



			@require_torch_a
			def    _lowerCamelCase	(						self       :  Any)       ->   List[str]:



						"""simple docstring"""






						_UpperCAmelCase			=       UNetaDModel.from_pretrained('diffusers/consistency_models'			,		subfolder='diffusers_cd_imagenet64_l2')
						_UpperCAmelCase			=       CMStochasticIterativeScheduler(
						    num_train_timesteps=40			,		sigma_min=0.0_0_2			,		sigma_max=80.0			,		)
						_UpperCAmelCase			=       ConsistencyModelPipeline(unet=A			,		scheduler=A)
						pipe.to(torch_device=A			,		torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=A)

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

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

						_UpperCAmelCase			=       np.array([0.1_8_7_5, 0.1_4_2_8, 0.1_2_8_9, 0.2_1_5_1, 0.2_0_9_2, 0.1_4_7_7, 0.1_8_7_7, 0.1_6_4_1, 0.1_3_5_3])

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



			@require_torch_a
			def    _lowerCamelCase	(						self       :  str)       ->   int:



						"""simple docstring"""






						_UpperCAmelCase			=       UNetaDModel.from_pretrained('diffusers/consistency_models'			,		subfolder='diffusers_cd_imagenet64_l2')
						_UpperCAmelCase			=       CMStochasticIterativeScheduler(
						    num_train_timesteps=40			,		sigma_min=0.0_0_2			,		sigma_max=80.0			,		)
						_UpperCAmelCase			=       ConsistencyModelPipeline(unet=A			,		scheduler=A)
						pipe.to(torch_device=A			,		torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=A)

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

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

						_UpperCAmelCase			=       np.array([0.1_6_6_3, 0.1_9_4_8, 0.2_2_7_5, 0.1_6_8_0, 0.1_2_0_4, 0.1_2_4_5, 0.1_8_5_8, 0.1_3_3_8, 0.2_0_9_5])

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





from typing import TYPE_CHECKING

from ...utils import _LazyModule


UpperCAmelCase__       =	{"tokenization_bertweet": ["BertweetTokenizer"]}


if TYPE_CHECKING:
							from .tokenization_bertweet import BertweetTokenizer

else:
							import sys

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