



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  __snake_case      (   SCREAMING_SNAKE_CASE			,unittest.TestCase):







							'''simple docstring'''






							UpperCamelCase__				:      int         =					ConsistencyModelPipeline
							UpperCamelCase__				:      str         =					UNCONDITIONAL_IMAGE_GENERATION_PARAMS
							UpperCamelCase__				:      Any         =					UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


							@property
							def _a      (							self		):
											a__				      =		UNetaDModel.from_pretrained(
											    """diffusers/consistency-models-test"""  ,  subfolder="""test_unet"""  ,  )
											return unet


							@property
							def _a      (							self		):
											a__				      =		UNetaDModel.from_pretrained(
											    """diffusers/consistency-models-test"""  ,  subfolder="""test_unet_class_cond"""  ,  )
											return unet


							def _a      (							self  ,  a_=False		):
											if class_cond:
															a__				      =		self.dummy_cond_unet
											else:
															a__				      =		self.dummy_uncond_unet

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

											a__				      =		{
											    """unet""": unet,
											    """scheduler""": scheduler,
											}

											return components


							def _a      (							self  ,  a_  ,  a_=0		):
											if str(a_		).startswith("""mps"""		):
															a__				      =		torch.manual_seed(a_		)
											else:
															a__				      =		torch.Generator(device=a_		).manual_seed(a_		)

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

											return inputs


							def _a      (							self		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		self.get_dummy_components()
											a__				      =		ConsistencyModelPipeline(**a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

											a__				      =		image[0, -3:, -3:, -1]
											a__				      =		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		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		self.get_dummy_components(class_cond=a_		)
											a__				      =		ConsistencyModelPipeline(**a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

											a__				      =		image[0, -3:, -3:, -1]
											a__				      =		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		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		self.get_dummy_components()
											a__				      =		ConsistencyModelPipeline(**a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

											a__				      =		image[0, -3:, -3:, -1]
											a__				      =		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		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		self.get_dummy_components(class_cond=a_		)
											a__				      =		ConsistencyModelPipeline(**a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

											a__				      =		image[0, -3:, -3:, -1]
											a__				      =		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  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							def _a      (							self  ,  a_=0  ,  a_=False  ,  a_="cpu"  ,  a_=torch.floataa  ,  a_=(1, 3, 64, 64)		):
											a__				      =		torch.manual_seed(a_		)

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

											if get_fixed_latents:
															a__				      =		self.get_fixed_latents(seed=a_  ,  device=a_  ,  dtype=a_  ,  shape=a_		)
															a__				      =		latents

											return inputs


							def _a      (							self  ,  a_=0  ,  a_="cpu"  ,  a_=torch.floataa  ,  a_=(1, 3, 64, 64)		):
											if type(a_		) == str:
															a__				      =		torch.device(a_		)
											a__				      =		torch.Generator(device=a_		).manual_seed(a_		)
											a__				      =		randn_tensor(a_  ,  generator=a_  ,  device=a_  ,  dtype=a_		)
											return latents


							def _a      (							self		):
											a__				      =		UNetaDModel.from_pretrained("""diffusers/consistency_models"""  ,  subfolder="""diffusers_cd_imagenet64_l2"""		)
											a__				      =		CMStochasticIterativeScheduler(
											    num_train_timesteps=40  ,  sigma_min=0.002  ,  sigma_max=80.0  ,  )
											a__				      =		ConsistencyModelPipeline(unet=a_  ,  scheduler=a_		)
											pipe.to(torch_device=a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

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

											a__				      =		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		):
											a__				      =		UNetaDModel.from_pretrained("""diffusers/consistency_models"""  ,  subfolder="""diffusers_cd_imagenet64_l2"""		)
											a__				      =		CMStochasticIterativeScheduler(
											    num_train_timesteps=40  ,  sigma_min=0.002  ,  sigma_max=80.0  ,  )
											a__				      =		ConsistencyModelPipeline(unet=a_  ,  scheduler=a_		)
											pipe.to(torch_device=a_		)
											pipe.set_progress_bar_config(disable=a_		)

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

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

											a__				      =		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		):
											a__				      =		UNetaDModel.from_pretrained("""diffusers/consistency_models"""  ,  subfolder="""diffusers_cd_imagenet64_l2"""		)
											a__				      =		CMStochasticIterativeScheduler(
											    num_train_timesteps=40  ,  sigma_min=0.002  ,  sigma_max=80.0  ,  )
											a__				      =		ConsistencyModelPipeline(unet=a_  ,  scheduler=a_		)
											pipe.to(torch_device=a_  ,  torch_dtype=torch.floataa		)
											pipe.set_progress_bar_config(disable=a_		)

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

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

											a__				      =		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		):
											a__				      =		UNetaDModel.from_pretrained("""diffusers/consistency_models"""  ,  subfolder="""diffusers_cd_imagenet64_l2"""		)
											a__				      =		CMStochasticIterativeScheduler(
											    num_train_timesteps=40  ,  sigma_min=0.002  ,  sigma_max=80.0  ,  )
											a__				      =		ConsistencyModelPipeline(unet=a_  ,  scheduler=a_		)
											pipe.to(torch_device=a_  ,  torch_dtype=torch.floataa		)
											pipe.set_progress_bar_config(disable=a_		)

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

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

											a__				      =		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






from .configuration_bert_masked import MaskedBertConfig
from .modeling_bert_masked import (
    MaskedBertForMultipleChoice,
    MaskedBertForQuestionAnswering,
    MaskedBertForSequenceClassification,
    MaskedBertForTokenClassification,
    MaskedBertModel,
)
from .modules import *

