
"""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 					lowerCAmelCase			(				lowerCAmelCase_     ,						unittest.TestCase					):
     """simple docstring"""







     __lowercase						:List[str]          =    ConsistencyModelPipeline
     __lowercase						:Any          =    UNCONDITIONAL_IMAGE_GENERATION_PARAMS
     __lowercase						:Any          =    UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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




     @property
     def    _lowerCAmelCase		(      self							)					->  str:







            '''simple docstring'''





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




     @property
     def    _lowerCAmelCase		(      self							)					->  str:







            '''simple docstring'''





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




     def    _lowerCAmelCase		(      self ,		UpperCamelCase__=False							)					->  Union[str, Any]:







            '''simple docstring'''





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

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

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

            return components




     def    _lowerCAmelCase		(      self ,		UpperCamelCase__ ,		UpperCamelCase__=0							)					->  Optional[int]:







            '''simple docstring'''





            if str(__lowerCamelCase							).startswith('''mps'''							):
                   lowerCamelCase_       							=		torch.manual_seed(__lowerCamelCase							)
            else:
                   lowerCamelCase_       							=		torch.Generator(device=__lowerCamelCase							).manual_seed(__lowerCamelCase							)

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

            return inputs




     def    _lowerCAmelCase		(      self							)					->  Tuple:







            '''simple docstring'''





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

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

            lowerCamelCase_       							=		image[0, -3:, -3:, -1]
            lowerCamelCase_       							=		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    _lowerCAmelCase		(      self							)					->  Dict:







            '''simple docstring'''





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

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

            lowerCamelCase_       							=		image[0, -3:, -3:, -1]
            lowerCamelCase_       							=		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    _lowerCAmelCase		(      self							)					->  int:







            '''simple docstring'''





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

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

            lowerCamelCase_       							=		image[0, -3:, -3:, -1]
            lowerCamelCase_       							=		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    _lowerCAmelCase		(      self							)					->  int:







            '''simple docstring'''





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

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

            lowerCamelCase_       							=		image[0, -3:, -3:, -1]
            lowerCamelCase_       							=		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 					lowerCAmelCase			(				unittest.TestCase					):
     """simple docstring"""







     def    _lowerCAmelCase		(      self							)					->  int:







            '''simple docstring'''





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




     def    _lowerCAmelCase		(      self ,		UpperCamelCase__=0 ,		UpperCamelCase__=False ,		UpperCamelCase__="cpu" ,		UpperCamelCase__=torch.floataa ,		UpperCamelCase__=(1, 3, 64, 64)							)					->  Any:







            '''simple docstring'''





            lowerCamelCase_       							=		torch.manual_seed(__lowerCamelCase							)

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

            if get_fixed_latents:
                   lowerCamelCase_       							=		self.get_fixed_latents(seed=__lowerCamelCase ,		device=__lowerCamelCase ,		dtype=__lowerCamelCase ,		shape=__lowerCamelCase							)
                   lowerCamelCase_       							=		latents

            return inputs




     def    _lowerCAmelCase		(      self ,		UpperCamelCase__=0 ,		UpperCamelCase__="cpu" ,		UpperCamelCase__=torch.floataa ,		UpperCamelCase__=(1, 3, 64, 64)							)					->  int:







            '''simple docstring'''





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




     def    _lowerCAmelCase		(      self							)					->  int:







            '''simple docstring'''





            lowerCamelCase_       							=		UNetaDModel.from_pretrained('''diffusers/consistency_models''' ,		subfolder='''diffusers_cd_imagenet64_l2'''							)
            lowerCamelCase_       							=		CMStochasticIterativeScheduler(
                num_train_timesteps=40 ,		sigma_min=0.002 ,		sigma_max=80.0 ,		)
            lowerCamelCase_       							=		ConsistencyModelPipeline(unet=__lowerCamelCase ,		scheduler=__lowerCamelCase							)
            pipe.to(torch_device=__lowerCamelCase							)
            pipe.set_progress_bar_config(disable=__lowerCamelCase							)

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

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

            lowerCamelCase_       							=		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    _lowerCAmelCase		(      self							)					->  Any:







            '''simple docstring'''





            lowerCamelCase_       							=		UNetaDModel.from_pretrained('''diffusers/consistency_models''' ,		subfolder='''diffusers_cd_imagenet64_l2'''							)
            lowerCamelCase_       							=		CMStochasticIterativeScheduler(
                num_train_timesteps=40 ,		sigma_min=0.002 ,		sigma_max=80.0 ,		)
            lowerCamelCase_       							=		ConsistencyModelPipeline(unet=__lowerCamelCase ,		scheduler=__lowerCamelCase							)
            pipe.to(torch_device=__lowerCamelCase							)
            pipe.set_progress_bar_config(disable=__lowerCamelCase							)

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

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

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







            '''simple docstring'''





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

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

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

            lowerCamelCase_       							=		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    _lowerCAmelCase		(      self							)					->  str:







            '''simple docstring'''





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

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

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

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





'''simple docstring'''






# This code is adapted from OpenAI's release
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py

import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import tempfile






def 		UpperCamelCase__      (		__SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE		)    ->   Dict:
      snake_case__       :   int	      =    multiprocessing.Manager()
      snake_case__       :   Any	      =    manager.list()

      snake_case__       :   List[str]	      =    multiprocessing.Process(target=__SCREAMING_SNAKE_CASE				,      args=(check_program, result, timeout)		)
      p.start()
      p.join(timeout=timeout + 1		)
      if p.is_alive():
            p.kill()

      if not result:
            result.append('timed out'		)

      return {
          "task_id": task_id,
          "passed": result[0] == "passed",
          "result": result[0],
          "completion_id": completion_id,
      }






def 		UpperCamelCase__      (		__SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE		)    ->   Tuple:
      with create_tempdir():
            # These system calls are needed when cleaning up tempdir.
            import os
            import shutil

            snake_case__       :   List[Any]	      =    shutil.rmtree
            snake_case__       :   Tuple	      =    os.rmdir
            snake_case__       :   Any	      =    os.chdir

            # Disable functionalities that can make destructive changes to the test.
            reliability_guard()

            # Run program.
            try:
                  snake_case__       :   Union[str, Any]	      =    {}
                  with swallow_io():
                        with time_limit(__SCREAMING_SNAKE_CASE		):
                              exec(__SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE		)
                  result.append('passed'		)
            except TimeoutException:
                  result.append('timed out'		)
            except BaseException as e:
                  result.append(f"failed: {e}"		)

            # Needed for cleaning up.
            snake_case__       :   List[Any]	      =    rmtree
            snake_case__       :   Optional[int]	      =    rmdir
            snake_case__       :   Dict	      =    chdir






@contextlib.contextmanager
def 		UpperCamelCase__      (		__SCREAMING_SNAKE_CASE		)    ->   List[str]:
      def signal_handler(__SCREAMING_SNAKE_CASE				,      __SCREAMING_SNAKE_CASE		):
            raise TimeoutException('Timed out!'		)

      signal.setitimer(signal.ITIMER_REAL				,      __SCREAMING_SNAKE_CASE		)
      signal.signal(signal.SIGALRM				,      __SCREAMING_SNAKE_CASE		)
      try:
            yield
      finally:
            signal.setitimer(signal.ITIMER_REAL				,      0		)






@contextlib.contextmanager
def 		UpperCamelCase__      (		)    ->   List[str]:
      snake_case__       :   Dict	      =    WriteOnlyStringIO()
      with contextlib.redirect_stdout(__SCREAMING_SNAKE_CASE		):
            with contextlib.redirect_stderr(__SCREAMING_SNAKE_CASE		):
                  with redirect_stdin(__SCREAMING_SNAKE_CASE		):
                        yield






@contextlib.contextmanager
def 		UpperCamelCase__      (		)    ->   Optional[int]:
      with tempfile.TemporaryDirectory() as dirname:
            with chdir(__SCREAMING_SNAKE_CASE		):
                  yield dirname






class        lowercase_      (						lowerCAmelCase_							):
     pass






class        lowercase_      (						io.StringIO							):





     def 			_lowerCAmelCase      (     self     :					Optional[Any]     ,							*__lowerCamelCase     :					int     ,							**__lowerCamelCase     :					Optional[int]			):
           raise OSError





     def 			_lowerCAmelCase      (     self     :					int     ,							*__lowerCamelCase     :					List[Any]     ,							**__lowerCamelCase     :					Tuple			):
           raise OSError





     def 			_lowerCAmelCase      (     self     :					Any     ,							*__lowerCamelCase     :					int     ,							**__lowerCamelCase     :					Union[str, Any]			):
           raise OSError





     def 			_lowerCAmelCase      (     self     :					Tuple     ,							*__lowerCamelCase     :					List[str]     ,							**__lowerCamelCase     :					Optional[Any]			):
           return False






class        lowercase_      (						contextlib._RedirectStream							):  # type: ignore
     A_											=			"stdin"






@contextlib.contextmanager
def 		UpperCamelCase__      (		__SCREAMING_SNAKE_CASE		)    ->   List[str]:
      if root == ".":
            yield
            return
      snake_case__       :   str	      =    os.getcwd()
      os.chdir(__SCREAMING_SNAKE_CASE		)
      try:
            yield
      except BaseException as exc:
            raise exc
      finally:
            os.chdir(__SCREAMING_SNAKE_CASE		)






def 		UpperCamelCase__      (		__SCREAMING_SNAKE_CASE=None		)    ->   Any:

      if maximum_memory_bytes is not None:
            import resource

            resource.setrlimit(resource.RLIMIT_AS				,      (maximum_memory_bytes, maximum_memory_bytes)		)
            resource.setrlimit(resource.RLIMIT_DATA				,      (maximum_memory_bytes, maximum_memory_bytes)		)
            if not platform.uname().system == "Darwin":
                  resource.setrlimit(resource.RLIMIT_STACK				,      (maximum_memory_bytes, maximum_memory_bytes)		)

      faulthandler.disable()

      import builtins

      snake_case__       :   List[Any]	      =    None
      snake_case__       :   List[str]	      =    None

      import os

      snake_case__       :   List[Any]	      =    '1'

      snake_case__       :   Union[str, Any]	      =    None
      snake_case__       :   Tuple	      =    None
      snake_case__       :   Dict	      =    None
      snake_case__       :   Tuple	      =    None
      snake_case__       :   List[str]	      =    None
      snake_case__       :   Any	      =    None
      snake_case__       :   Optional[Any]	      =    None
      snake_case__       :   Union[str, Any]	      =    None
      snake_case__       :   int	      =    None
      snake_case__       :   Optional[int]	      =    None
      snake_case__       :   Union[str, Any]	      =    None
      snake_case__       :   List[Any]	      =    None
      snake_case__       :   Optional[int]	      =    None
      snake_case__       :   Union[str, Any]	      =    None
      snake_case__       :   Tuple	      =    None
      snake_case__       :   int	      =    None
      snake_case__       :   Optional[Any]	      =    None
      snake_case__       :   Optional[int]	      =    None
      snake_case__       :   List[str]	      =    None
      snake_case__       :   List[str]	      =    None
      snake_case__       :   List[str]	      =    None
      snake_case__       :   Dict	      =    None
      snake_case__       :   Optional[Any]	      =    None
      snake_case__       :   List[Any]	      =    None
      snake_case__       :   Tuple	      =    None
      snake_case__       :   Optional[int]	      =    None
      snake_case__       :   int	      =    None

      import shutil

      snake_case__       :   List[str]	      =    None
      snake_case__       :   List[Any]	      =    None
      snake_case__       :   List[Any]	      =    None

      import subprocess

      snake_case__       :   int	      =    None  # type: ignore

      snake_case__       :   Dict	      =    None

      import sys

      snake_case__       :   Tuple	      =    None
      snake_case__       :   Tuple	      =    None
      snake_case__       :   List[str]	      =    None
      snake_case__       :   Optional[int]	      =    None
      snake_case__       :   List[Any]	      =    None
