

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






  """simple docstring"""


  a     =      ConsistencyModelPipeline
  a     =      UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  a     =      UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






  @property
  def    lowercase_	(							self		:	Dict      )	->							str:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained(
           '''diffusers/consistency-models-test'''			,     subfolder='''test_unet'''			,     )
       return unet






  @property
  def    lowercase_	(							self		:	int      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained(
           '''diffusers/consistency-models-test'''			,     subfolder='''test_unet_class_cond'''			,     )
       return unet






  def    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	str=False      )	->							str:
       if class_cond:
            SCREAMING_SNAKE_CASE__					=					self.dummy_cond_unet
       else:
            SCREAMING_SNAKE_CASE__					=					self.dummy_uncond_unet

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

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

       return components






  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Union[str, Any]=0      )	->							Any:
       if str(__lowerCamelCase      ).startswith('''mps'''      ):
            SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )
       else:
            SCREAMING_SNAKE_CASE__					=					torch.Generator(device=__lowerCamelCase      ).manual_seed(__lowerCamelCase      )

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

       return inputs






  def    lowercase_	(							self		:	Optional[Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					pipe.to(__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

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

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

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

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

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

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






  """simple docstring"""


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






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

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

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

       return inputs






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






  def    lowercase_	(							self		:	int      )	->							int:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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    lowercase_	(							self		:	List[str]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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    lowercase_	(							self		:	Optional[int]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase			,     torch_dtype=torch.floataa      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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    lowercase_	(							self		:	Any      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,     subfolder='''diffusers_cd_imagenet64_l2'''      )
       SCREAMING_SNAKE_CASE__					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40			,     sigma_min=0.002			,     sigma_max=80.0			,     )
       SCREAMING_SNAKE_CASE__					=					ConsistencyModelPipeline(unet=__lowerCamelCase			,     scheduler=__lowerCamelCase      )
       pipe.to(torch_device=__lowerCamelCase			,     torch_dtype=torch.floataa      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_inputs(get_fixed_latents=__lowerCamelCase			,     device=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					1
       SCREAMING_SNAKE_CASE__					=					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__					=					pipe(**__lowerCamelCase      ).images
       assert image.shape == (1, 64, 64, 3)

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

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





from importlib import import_module

from .logging import get_logger


_SCREAMING_SNAKE_CASE  :						Optional[int]         =    get_logger(__name__)






class  UpperCAmelCase__	:






  """simple docstring"""


  def __init__(							self		:	Optional[int]			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	Any=None      )	->							int:
       SCREAMING_SNAKE_CASE__					=					attrs or []
       if module is not None:
            for key in module.__dict__:
                 if key in attrs or not key.startswith('''__'''      ):
                      setattr(self			,     __lowerCamelCase			,     getattr(__lowerCamelCase			,     __lowerCamelCase      )      )
       SCREAMING_SNAKE_CASE__					=					module._original_module if isinstance(__lowerCamelCase			,     _PatchedModuleObj      ) else module






class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      []






  def __init__(							self		:	int			,     __lowerCamelCase		:	Union[str, Any]			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	List[Any]=None      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					obj
       SCREAMING_SNAKE_CASE__					=					target
       SCREAMING_SNAKE_CASE__					=					new
       SCREAMING_SNAKE_CASE__					=					target.split('''.'''      )[0]
       SCREAMING_SNAKE_CASE__					=					{}
       SCREAMING_SNAKE_CASE__					=					attrs or []






  def __enter__(							self		:	int      )	->							Tuple:
       *SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					self.target.split('''.'''      )

       # Patch modules:
       # it's used to patch attributes of submodules like "os.path.join";
       # in this case we need to patch "os" and "os.path"

       for i in range(len(__lowerCamelCase      )      ):
            try:
                 SCREAMING_SNAKE_CASE__					=					import_module('''.'''.join(submodules[: i + 1]      )      )
            except ModuleNotFoundError:
                 continue
            # We iterate over all the globals in self.obj in case we find "os" or "os.path"
            for attr in self.obj.__dir__():
                 SCREAMING_SNAKE_CASE__					=					getattr(self.obj			,     __lowerCamelCase      )
                 # We don't check for the name of the global, but rather if its value *is* "os" or "os.path".
                 # This allows to patch renamed modules like "from os import path as ospath".
                 if obj_attr is submodule or (
                     (isinstance(__lowerCamelCase			,     _PatchedModuleObj      ) and obj_attr._original_module is submodule)
                 ):
                      SCREAMING_SNAKE_CASE__					=					obj_attr
                      # patch at top level
                      setattr(self.obj			,     __lowerCamelCase			,     _PatchedModuleObj(__lowerCamelCase			,     attrs=self.attrs      )      )
                      SCREAMING_SNAKE_CASE__					=					getattr(self.obj			,     __lowerCamelCase      )
                      # construct lower levels patches
                      for key in submodules[i + 1 :]:
                           setattr(__lowerCamelCase			,     __lowerCamelCase			,     _PatchedModuleObj(getattr(__lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase      )			,     attrs=self.attrs      )      )
                           SCREAMING_SNAKE_CASE__					=					getattr(__lowerCamelCase			,     __lowerCamelCase      )
                      # finally set the target attribute
                      setattr(__lowerCamelCase			,     __lowerCamelCase			,     self.new      )

        # Patch attribute itself:
        # it's used for builtins like "open",
        # and also to patch "os.path.join" we may also need to patch "join"
        # itself if it was imported as "from os.path import join".

       if submodules:  # if it's an attribute of a submodule like "os.path.join"
            try:
                 SCREAMING_SNAKE_CASE__					=					getattr(import_module('''.'''.join(__lowerCamelCase      )      )			,     __lowerCamelCase      )
            except (AttributeError, ModuleNotFoundError):
                 return
            # We iterate over all the globals in self.obj in case we find "os.path.join"
            for attr in self.obj.__dir__():
                 # We don't check for the name of the global, but rather if its value *is* "os.path.join".
                 # This allows to patch renamed attributes like "from os.path import join as pjoin".
                 if getattr(self.obj			,     __lowerCamelCase      ) is attr_value:
                      SCREAMING_SNAKE_CASE__					=					getattr(self.obj			,     __lowerCamelCase      )
                      setattr(self.obj			,     __lowerCamelCase			,     self.new      )
       elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
            SCREAMING_SNAKE_CASE__					=					globals()['''__builtins__'''][target_attr]
            setattr(self.obj			,     __lowerCamelCase			,     self.new      )
       else:
            raise RuntimeError(f'''Tried to patch attribute {target_attr} instead of a submodule.'''      )






  def __exit__(							self		:	Optional[Any]			,     *__lowerCamelCase		:	Tuple      )	->							List[str]:
       for attr in list(self.original      ):
            setattr(self.obj			,     __lowerCamelCase			,     self.original.pop(__lowerCamelCase      )      )






  def    lowercase_	(							self		:	List[Any]      )	->							Optional[Any]:
       self.__enter__()
       self._active_patches.append(self      )







  def    lowercase_	(							self		:	int      )	->							int:
       try:
            self._active_patches.remove(self      )
       except ValueError:
            # If the patch hasn't been started this will fail
            return None

       return self.__exit__()


