






"""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_ (__A       ,		unittest.TestCase  ):
	__magic_name__			   =       ConsistencyModelPipeline
	__magic_name__			   =       UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	__magic_name__			   =       UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Any		)		->    Tuple:
				UpperCAmelCase_  :				int      =      UNetaDModel.from_pretrained(
				    "diffusers/consistency-models-test"			,			subfolder="test_unet"			,			)
				return unet





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]		)		->    int:
				UpperCAmelCase_  :				Optional[Any]      =      UNetaDModel.from_pretrained(
				    "diffusers/consistency-models-test"			,			subfolder="test_unet_class_cond"			,			)
				return unet





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict			,			lowerCAmelCase_      :				int=False		)		->    List[Any]:
				if class_cond:
							UpperCAmelCase_  :				str      =      self.dummy_cond_unet
				else:
							UpperCAmelCase_  :				Optional[int]      =      self.dummy_uncond_unet

				# Default to CM multistep sampler
				UpperCAmelCase_  :				Optional[int]      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=8_0.0			,			)

				UpperCAmelCase_  :				Union[str, Any]      =      {
				    "unet": unet,
				    "scheduler": scheduler,
				}

				return components





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]			,			lowerCAmelCase_      :				str			,			lowerCAmelCase_      :				str=0		)		->    Optional[Any]:
				if str(lowerCAmelCase_		).startswith("mps"		):
							UpperCAmelCase_  :				Tuple      =      torch.manual_seed(lowerCAmelCase_		)
				else:
							UpperCAmelCase_  :				Tuple      =      torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)

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

				return inputs





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Dict:
				UpperCAmelCase_  :				Union[str, Any]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				Tuple      =      self.get_dummy_components()
				UpperCAmelCase_  :				Tuple      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[int]      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				int      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

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

				UpperCAmelCase_  :				Dict      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      0
				UpperCAmelCase_  :				str      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				List[Any]      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				Optional[Any]      =      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   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Optional[Any]:
				UpperCAmelCase_  :				List[str]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				List[Any]      =      self.get_dummy_components()
				UpperCAmelCase_  :				List[str]      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[Any]      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				Tuple      =      1
				UpperCAmelCase_  :				Dict      =      None
				UpperCAmelCase_  :				Dict      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				Any      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				List[Any]      =      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   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]		)		->    Union[str, Any]:
				UpperCAmelCase_  :				Optional[Any]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				Any      =      self.get_dummy_components(class_cond=lowerCAmelCase_		)
				UpperCAmelCase_  :				List[Any]      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				Dict      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				str      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				List[Any]      =      1
				UpperCAmelCase_  :				List[str]      =      None
				UpperCAmelCase_  :				Union[str, Any]      =      0
				UpperCAmelCase_  :				Dict      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				Optional[int]      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				Any      =      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     UpperCamelCase_ (unittest.TestCase  ):





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Union[str, Any]:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def   _SCREAMING_SNAKE_CASE   (   self      :				int			,			lowerCAmelCase_      :				Dict=0			,			lowerCAmelCase_      :				Union[str, Any]=False			,			lowerCAmelCase_      :				Any="cpu"			,			lowerCAmelCase_      :				int=torch.floataa			,			lowerCAmelCase_      :				Optional[int]=(1, 3, 64, 64)		)		->    str:
				UpperCAmelCase_  :				Optional[Any]      =      torch.manual_seed(lowerCAmelCase_		)

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

				if get_fixed_latents:
							UpperCAmelCase_  :				Dict      =      self.get_fixed_latents(seed=lowerCAmelCase_			,			device=lowerCAmelCase_			,			dtype=lowerCAmelCase_			,			shape=lowerCAmelCase_		)
							UpperCAmelCase_  :				List[Any]      =      latents

				return inputs





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]			,			lowerCAmelCase_      :				Optional[Any]=0			,			lowerCAmelCase_      :				List[Any]="cpu"			,			lowerCAmelCase_      :				Optional[Any]=torch.floataa			,			lowerCAmelCase_      :				Optional[Any]=(1, 3, 64, 64)		)		->    Tuple:
				if type(lowerCAmelCase_		) == str:
							UpperCAmelCase_  :				List[str]      =      torch.device(lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)
				UpperCAmelCase_  :				Dict      =      randn_tensor(lowerCAmelCase_			,			generator=lowerCAmelCase_			,			device=lowerCAmelCase_			,			dtype=lowerCAmelCase_		)
				return latents





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    List[Any]:
				UpperCAmelCase_  :				Optional[Any]      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				Dict      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=8_0.0			,			)
				UpperCAmelCase_  :				Optional[int]      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[int]      =      self.get_inputs()
				UpperCAmelCase_  :				List[Any]      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

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

				UpperCAmelCase_  :				str      =      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   _SCREAMING_SNAKE_CASE   (   self      :				Tuple		)		->    int:
				UpperCAmelCase_  :				List[Any]      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				Optional[int]      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=8_0.0			,			)
				UpperCAmelCase_  :				Any      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      self.get_inputs()
				UpperCAmelCase_  :				List[Any]      =      1
				UpperCAmelCase_  :				Union[str, Any]      =      None
				UpperCAmelCase_  :				Dict      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_  :				int      =      image[0, -3:, -3:, -1]

				UpperCAmelCase_  :				str      =      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   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]		)		->    Union[str, Any]:
				UpperCAmelCase_  :				Dict      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				int      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=8_0.0			,			)
				UpperCAmelCase_  :				Any      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_			,			torch_dtype=torch.floataa		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

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

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

				UpperCAmelCase_  :				Optional[int]      =      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   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    Optional[Any]:
				UpperCAmelCase_  :				List[Any]      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				Optional[int]      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=8_0.0			,			)
				UpperCAmelCase_  :				List[Any]      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_			,			torch_dtype=torch.floataa		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

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

				UpperCAmelCase_  :				Any      =      image[0, -3:, -3:, -1]

				UpperCAmelCase_  :				Tuple      =      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











"""simple docstring"""
from __future__ import annotations

import math

def        snake_case   (     A__		):

			if 1 < number < 4:
						# 2 and 3 are primes
						return True
			elif number < 2 or number % 2 == 0 or number % 3 == 0:
						# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
						return False

			# All primes number are in format of 6k +/- 1
			for i in range(5			,int(math.sqrt(A__		) + 1		)			,6		):
						if number % i == 0 or number % (i + 2) == 0:
									return False
			return True


lowerCamelCase_       			=   [num for num in range(3, 10_0001, 2) if not is_prime(num)]

def        snake_case   (     A__		):
			if not isinstance(A__			,A__		):
						raise ValueError("n must be an integer"		)
			if n <= 0:
						raise ValueError("n must be >= 0"		)

			UpperCAmelCase_  :				List[str]      =      []
			for num in range(len(A__		)		):
						UpperCAmelCase_  :				Dict      =      0
						while 2 * i * i <= odd_composites[num]:
									UpperCAmelCase_  :				List[str]      =      odd_composites[num] - 2 * i * i
									if is_prime(A__		):
												break
									i += 1
						else:
									list_nums.append(odd_composites[num]		)
									if len(A__		) == n:
												return list_nums

			return []

def        snake_case   (     ):
			return compute_nums(1		)[0]


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



