





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







							"""simple docstring"""




							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  A_     ( self     :					List[str]						)				->			List[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[str]     =						UNetaDModel.from_pretrained(
														    "diffusers/consistency-models-test",			subfolder="test_unet",			)
														return unet






							@property
							def  A_     ( self     :					Optional[Any]						)				->			Tuple:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						UNetaDModel.from_pretrained(
														    "diffusers/consistency-models-test",			subfolder="test_unet_class_cond",			)
														return unet






							def  A_     ( self     :					int,			_UpperCAmelCase     :					int=False						)				->			Optional[Any]:





														"""simple docstring"""
														if class_cond:
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						self.dummy_cond_unet
														else:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.dummy_uncond_unet

														# Default to CM multistep sampler
														SCREAMING_SNAKE_CASE__       :      List[str]     =						CMStochasticIterativeScheduler(
														    num_train_timesteps=4_0,			sigma_min=0.002,			sigma_max=80.0,			)

														SCREAMING_SNAKE_CASE__       :      Dict     =						{
														    "unet": unet,
														    "scheduler": scheduler,
														}

														return components






							def  A_     ( self     :					Optional[int],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					List[str]=0						)				->			Optional[Any]:





														"""simple docstring"""
														if str(_UpperCAmelCase						).startswith("mps"						):
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						torch.manual_seed(_UpperCAmelCase						)
														else:
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      int     =						{
														    "batch_size": 1,
														    "num_inference_steps": None,
														    "timesteps": [2_2, 0],
														    "generator": generator,
														    "output_type": "np",
														}

														return inputs






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Dict     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.get_dummy_components()
														SCREAMING_SNAKE_CASE__       :      List[str]     =						ConsistencyModelPipeline(**_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						pipe.to(_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      str     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						pipe(**_UpperCAmelCase						).images
														assert image.shape == (1, 3_2, 3_2, 3)

														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						image[0, -3:, -3:, -1]
														SCREAMING_SNAKE_CASE__       :      str     =						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  A_     ( self     :					Dict						)				->			Dict:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.get_dummy_components(class_cond=_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						ConsistencyModelPipeline(**_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Any     =						pipe.to(_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Dict     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						0
														SCREAMING_SNAKE_CASE__       :      List[str]     =						pipe(**_UpperCAmelCase						).images
														assert image.shape == (1, 3_2, 3_2, 3)

														SCREAMING_SNAKE_CASE__       :      int     =						image[0, -3:, -3:, -1]
														SCREAMING_SNAKE_CASE__       :      str     =						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  A_     ( self     :					Union[str, Any]						)				->			List[str]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
														SCREAMING_SNAKE_CASE__       :      Dict     =						self.get_dummy_components()
														SCREAMING_SNAKE_CASE__       :      Any     =						ConsistencyModelPipeline(**_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      List[str]     =						pipe.to(_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Any     =						1
														SCREAMING_SNAKE_CASE__       :      int     =						None
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						pipe(**_UpperCAmelCase						).images
														assert image.shape == (1, 3_2, 3_2, 3)

														SCREAMING_SNAKE_CASE__       :      Dict     =						image[0, -3:, -3:, -1]
														SCREAMING_SNAKE_CASE__       :      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  A_     ( self     :					int						)				->			Tuple:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						"cpu"  # ensure determinism for the device-dependent torch.Generator
														SCREAMING_SNAKE_CASE__       :      Any     =						self.get_dummy_components(class_cond=_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						ConsistencyModelPipeline(**_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Dict     =						pipe.to(_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Dict     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						1
														SCREAMING_SNAKE_CASE__       :      int     =						None
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						0
														SCREAMING_SNAKE_CASE__       :      int     =						pipe(**_UpperCAmelCase						).images
														assert image.shape == (1, 3_2, 3_2, 3)

														SCREAMING_SNAKE_CASE__       :      int     =						image[0, -3:, -3:, -1]
														SCREAMING_SNAKE_CASE__       :      List[str]     =						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 						lowerCamelCase       (unittest.TestCase      ):







							"""simple docstring"""




							def  A_     ( self     :					int						)				->			Any:





														"""simple docstring"""
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()






							def  A_     ( self     :					List[str],			_UpperCAmelCase     :					Optional[int]=0,			_UpperCAmelCase     :					List[Any]=False,			_UpperCAmelCase     :					List[str]="cpu",			_UpperCAmelCase     :					List[str]=torch.floataa,			_UpperCAmelCase     :					List[Any]=(1, 3, 6_4, 6_4)						)				->			Optional[int]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						torch.manual_seed(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Any     =						{
														    "num_inference_steps": None,
														    "timesteps": [2_2, 0],
														    "class_labels": 0,
														    "generator": generator,
														    "output_type": "np",
														}

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

														return inputs






							def  A_     ( self     :					Optional[Any],			_UpperCAmelCase     :					int=0,			_UpperCAmelCase     :					Union[str, Any]="cpu",			_UpperCAmelCase     :					Optional[Any]=torch.floataa,			_UpperCAmelCase     :					Dict=(1, 3, 6_4, 6_4)						)				->			List[str]:





														"""simple docstring"""
														if type(_UpperCAmelCase						) == str:
																					SCREAMING_SNAKE_CASE__       :      int     =						torch.device(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						randn_tensor(_UpperCAmelCase,			generator=_UpperCAmelCase,			device=_UpperCAmelCase,			dtype=_UpperCAmelCase						)
														return latents






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						UNetaDModel.from_pretrained("diffusers/consistency_models",			subfolder="diffusers_cd_imagenet64_l2"						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						CMStochasticIterativeScheduler(
														    num_train_timesteps=4_0,			sigma_min=0.002,			sigma_max=80.0,			)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						ConsistencyModelPipeline(unet=_UpperCAmelCase,			scheduler=_UpperCAmelCase						)
														pipe.to(torch_device=_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      int     =						self.get_inputs()
														SCREAMING_SNAKE_CASE__       :      Any     =						pipe(**_UpperCAmelCase						).images
														assert image.shape == (1, 6_4, 6_4, 3)

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

														SCREAMING_SNAKE_CASE__       :      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  A_     ( self     :					List[str]						)				->			Optional[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						UNetaDModel.from_pretrained("diffusers/consistency_models",			subfolder="diffusers_cd_imagenet64_l2"						)
														SCREAMING_SNAKE_CASE__       :      Any     =						CMStochasticIterativeScheduler(
														    num_train_timesteps=4_0,			sigma_min=0.002,			sigma_max=80.0,			)
														SCREAMING_SNAKE_CASE__       :      Any     =						ConsistencyModelPipeline(unet=_UpperCAmelCase,			scheduler=_UpperCAmelCase						)
														pipe.to(torch_device=_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Any     =						self.get_inputs()
														SCREAMING_SNAKE_CASE__       :      Tuple     =						1
														SCREAMING_SNAKE_CASE__       :      Tuple     =						None
														SCREAMING_SNAKE_CASE__       :      int     =						pipe(**_UpperCAmelCase						).images
														assert image.shape == (1, 6_4, 6_4, 3)

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

														SCREAMING_SNAKE_CASE__       :      Dict     =						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  A_     ( self     :					Any						)				->			Optional[int]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						UNetaDModel.from_pretrained("diffusers/consistency_models",			subfolder="diffusers_cd_imagenet64_l2"						)
														SCREAMING_SNAKE_CASE__       :      List[str]     =						CMStochasticIterativeScheduler(
														    num_train_timesteps=4_0,			sigma_min=0.002,			sigma_max=80.0,			)
														SCREAMING_SNAKE_CASE__       :      str     =						ConsistencyModelPipeline(unet=_UpperCAmelCase,			scheduler=_UpperCAmelCase						)
														pipe.to(torch_device=_UpperCAmelCase,			torch_dtype=torch.floataa						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

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

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

														SCREAMING_SNAKE_CASE__       :      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  A_     ( self     :					Union[str, Any]						)				->			Any:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						UNetaDModel.from_pretrained("diffusers/consistency_models",			subfolder="diffusers_cd_imagenet64_l2"						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						CMStochasticIterativeScheduler(
														    num_train_timesteps=4_0,			sigma_min=0.002,			sigma_max=80.0,			)
														SCREAMING_SNAKE_CASE__       :      int     =						ConsistencyModelPipeline(unet=_UpperCAmelCase,			scheduler=_UpperCAmelCase						)
														pipe.to(torch_device=_UpperCAmelCase,			torch_dtype=torch.floataa						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

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

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

														SCREAMING_SNAKE_CASE__       :      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













def _a (  SCREAMING_SNAKE_CASE__      :  int = 50_00_00_00       )	->    int:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      str     =						set()
							SCREAMING_SNAKE_CASE__       :      Dict     =						int((limit - 24) ** (1 / 2)       )

							SCREAMING_SNAKE_CASE__       :      Any     =						set(range(3			,				prime_square_limit + 1			,				2       )       )
							primes.add(2       )
							for p in range(3			,				prime_square_limit + 1			,				2       ):
														if p not in primes:
																					continue
														primes.difference_update(set(range(p * p			,				prime_square_limit + 1			,				SCREAMING_SNAKE_CASE__       )       )       )

							for primea in primes:
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						primea * primea
														for primea in primes:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						primea * primea * primea
																					if square + cube >= limit - 16:
																												break
																					for primea in primes:
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						primea * primea * primea * primea
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						square + cube + tetr
																												if total >= limit:
																																			break
																												ret.add(SCREAMING_SNAKE_CASE__       )

							return len(SCREAMING_SNAKE_CASE__       )


if __name__ == "__main__":
		print(f"{solution() = }")






