

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		(	__magic_name__ ,  unittest.TestCase):
  SCREAMING_SNAKE_CASE     :   Union[str, Any]							=	ConsistencyModelPipeline
  SCREAMING_SNAKE_CASE     :   List[Any]							=	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  SCREAMING_SNAKE_CASE     :   Optional[Any]							=	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

  @property
  def 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

  @property
  def 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

  def 			UpperCAmelCase     (					self			,       _SCREAMING_SNAKE_CASE=False       ):





        """simple docstring"""







        if class_cond:
              SCREAMING_SNAKE_CASE_					:	Any     							=  self.dummy_cond_unet
        else:
              SCREAMING_SNAKE_CASE_					:	Any     							=  self.dummy_uncond_unet

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

        SCREAMING_SNAKE_CASE_					:	Any     							=  {
            'unet': unet,
            'scheduler': scheduler,
        }

        return components

  def 			UpperCAmelCase     (					self			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE=0       ):





        """simple docstring"""







        if str(__lowerCAmelCase       ).startswith('mps'       ):
              SCREAMING_SNAKE_CASE_					:	Union[str, Any]     							=  torch.manual_seed(__lowerCAmelCase       )
        else:
              SCREAMING_SNAKE_CASE_					:	Optional[Any]     							=  torch.Generator(device=__lowerCAmelCase       ).manual_seed(__lowerCAmelCase       )

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

        return inputs

  def 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

        SCREAMING_SNAKE_CASE_					:	List[Any]     							=  self.get_dummy_inputs(__lowerCAmelCase       )
        SCREAMING_SNAKE_CASE_					:	str     							=  pipe(**__lowerCAmelCase       ).images
        assert image.shape == (1, 32, 32, 3)

        SCREAMING_SNAKE_CASE_					:	Any     							=  image[0, -3:, -3:, -1]
        SCREAMING_SNAKE_CASE_					:	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 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

        SCREAMING_SNAKE_CASE_					:	Optional[Any]     							=  self.get_dummy_inputs(__lowerCAmelCase       )
        SCREAMING_SNAKE_CASE_					:	Any     							=  0
        SCREAMING_SNAKE_CASE_					:	Optional[Any]     							=  pipe(**__lowerCAmelCase       ).images
        assert image.shape == (1, 32, 32, 3)

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





        """simple docstring"""







        SCREAMING_SNAKE_CASE_					:	Any     							=  'cpu'  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE_					:	str     							=  self.get_dummy_components()
        SCREAMING_SNAKE_CASE_					:	int     							=  ConsistencyModelPipeline(**__lowerCAmelCase       )
        SCREAMING_SNAKE_CASE_					:	str     							=  pipe.to(__lowerCAmelCase       )
        pipe.set_progress_bar_config(disable=__lowerCAmelCase       )

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

        SCREAMING_SNAKE_CASE_					:	Optional[Any]     							=  image[0, -3:, -3:, -1]
        SCREAMING_SNAKE_CASE_					:	Optional[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 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

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

        SCREAMING_SNAKE_CASE_					:	Any     							=  image[0, -3:, -3:, -1]
        SCREAMING_SNAKE_CASE_					:	Optional[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

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

  def 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

  def 			UpperCAmelCase     (					self			,       _SCREAMING_SNAKE_CASE=0			,       _SCREAMING_SNAKE_CASE=False			,       _SCREAMING_SNAKE_CASE="cpu"			,       _SCREAMING_SNAKE_CASE=torch.floataa			,       _SCREAMING_SNAKE_CASE=(1, 3, 64, 64)       ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_					:	str     							=  torch.manual_seed(__lowerCAmelCase       )

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

        if get_fixed_latents:
              SCREAMING_SNAKE_CASE_					:	Optional[Any]     							=  self.get_fixed_latents(seed=__lowerCAmelCase			,       device=__lowerCAmelCase			,       dtype=__lowerCAmelCase			,       shape=__lowerCAmelCase       )
              SCREAMING_SNAKE_CASE_					:	Tuple     							=  latents

        return inputs

  def 			UpperCAmelCase     (					self			,       _SCREAMING_SNAKE_CASE=0			,       _SCREAMING_SNAKE_CASE="cpu"			,       _SCREAMING_SNAKE_CASE=torch.floataa			,       _SCREAMING_SNAKE_CASE=(1, 3, 64, 64)       ):





        """simple docstring"""







        if type(__lowerCAmelCase       ) == str:
              SCREAMING_SNAKE_CASE_					:	Tuple     							=  torch.device(__lowerCAmelCase       )
        SCREAMING_SNAKE_CASE_					:	Tuple     							=  torch.Generator(device=__lowerCAmelCase       ).manual_seed(__lowerCAmelCase       )
        SCREAMING_SNAKE_CASE_					:	Optional[int]     							=  randn_tensor(__lowerCAmelCase			,       generator=__lowerCAmelCase			,       device=__lowerCAmelCase			,       dtype=__lowerCAmelCase       )
        return latents

  def 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

        SCREAMING_SNAKE_CASE_					:	str     							=  self.get_inputs()
        SCREAMING_SNAKE_CASE_					:	Dict     							=  pipe(**__lowerCAmelCase       ).images
        assert image.shape == (1, 64, 64, 3)

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

        SCREAMING_SNAKE_CASE_					:	Union[str, Any]     							=  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 			UpperCAmelCase     (					self       ):





        """simple docstring"""







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

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

        SCREAMING_SNAKE_CASE_					:	Optional[Any]     							=  image[0, -3:, -3:, -1]

        SCREAMING_SNAKE_CASE_					:	List[str]     							=  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 			UpperCAmelCase     (					self       ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_					:	Dict     							=  UNetaDModel.from_pretrained('diffusers/consistency_models'			,       subfolder='diffusers_cd_imagenet64_l2'       )
        SCREAMING_SNAKE_CASE_					:	str     							=  CMStochasticIterativeScheduler(
            num_train_timesteps=40			,       sigma_min=0.002			,       sigma_max=80.0			,       )
        SCREAMING_SNAKE_CASE_					:	int     							=  ConsistencyModelPipeline(unet=__lowerCAmelCase			,       scheduler=__lowerCAmelCase       )
        pipe.to(torch_device=__lowerCAmelCase			,       torch_dtype=torch.floataa       )
        pipe.set_progress_bar_config(disable=__lowerCAmelCase       )

        SCREAMING_SNAKE_CASE_					:	Union[str, Any]     							=  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       ):
              SCREAMING_SNAKE_CASE_					:	str     							=  pipe(**__lowerCAmelCase       ).images
        assert image.shape == (1, 64, 64, 3)

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

        SCREAMING_SNAKE_CASE_					:	Dict     							=  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 			UpperCAmelCase     (					self       ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_					:	Tuple     							=  UNetaDModel.from_pretrained('diffusers/consistency_models'			,       subfolder='diffusers_cd_imagenet64_l2'       )
        SCREAMING_SNAKE_CASE_					:	Tuple     							=  CMStochasticIterativeScheduler(
            num_train_timesteps=40			,       sigma_min=0.002			,       sigma_max=80.0			,       )
        SCREAMING_SNAKE_CASE_					:	List[Any]     							=  ConsistencyModelPipeline(unet=__lowerCAmelCase			,       scheduler=__lowerCAmelCase       )
        pipe.to(torch_device=__lowerCAmelCase			,       torch_dtype=torch.floataa       )
        pipe.set_progress_bar_config(disable=__lowerCAmelCase       )

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

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

        SCREAMING_SNAKE_CASE_					:	Tuple     							=  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   __lowerCamelCase   (				__a  :int       )						->							Dict:






		"""simple docstring"""

		A__     =  len(__a       )
		A__     =  sum(__a       )

		A__     =  [[False for x in range(s + 1       )] for y in range(n + 1       )]

		for i in range(1       ,	n + 1       ):
				A__     =  True

		for i in range(1       ,	s + 1       ):
				A__     =  False

		for i in range(1       ,	n + 1       ):
				for j in range(1       ,	s + 1       ):
						A__     =  dp[i][j - 1]

						if arr[i - 1] <= j:
								A__     =  dp[i][j] or dp[i - 1][j - arr[i - 1]]

		for j in range(int(s / 2       )       ,	-1       ,	-1       ):
				if dp[n][j] is True:
						A__     =  s - 2 * j
						break

		return diff



