






"""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__       (    UpperCAmelCase__       ,	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 _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


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




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	UNetaDModel.from_pretrained(
              '''diffusers/consistency-models-test'''		,			subfolder='''test_unet_class_cond'''		,			)
          return unet
      def _SCREAMING_SNAKE_CASE				(		self  :						str		,			lowercase_  :						Tuple=False):




          '''simple docstring'''


          if class_cond:
              SCREAMING_SNAKE_CASE_   :						str							=	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.0_02		,			sigma_max=80.0		,			)

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

          return components
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict		,			lowercase_  :						Optional[Any]		,			lowercase_  :						Dict=0):




          '''simple docstring'''


          if str(lowercase_).startswith('''mps'''):
              SCREAMING_SNAKE_CASE_   :						Dict							=	torch.manual_seed(lowercase_)
          else:
              SCREAMING_SNAKE_CASE_   :						Dict							=	torch.Generator(device=lowercase_).manual_seed(lowercase_)

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

          return inputs
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


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

          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.get_dummy_inputs(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Dict							=	pipe(**lowercase_).images
          assert image.shape == (1, 32, 32, 3)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	image[0, -3:, -3:, -1]
          SCREAMING_SNAKE_CASE_   :						Dict							=	np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


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

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

          SCREAMING_SNAKE_CASE_   :						str							=	image[0, -3:, -3:, -1]
          SCREAMING_SNAKE_CASE_   :						str							=	np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


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

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

          SCREAMING_SNAKE_CASE_   :						Tuple							=	image[0, -3:, -3:, -1]
          SCREAMING_SNAKE_CASE_   :						Any							=	np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3



      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


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

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

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	image[0, -3:, -3:, -1]
          SCREAMING_SNAKE_CASE_   :						Dict							=	np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3




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






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]		,			lowercase_  :						int=0		,			lowercase_  :						List[Any]=False		,			lowercase_  :						List[str]="cpu"		,			lowercase_  :						Tuple=torch.floataa		,			lowercase_  :						List[str]=(1, 3, 64, 64)):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	torch.manual_seed(lowercase_)

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

          if get_fixed_latents:
              SCREAMING_SNAKE_CASE_   :						Tuple							=	self.get_fixed_latents(seed=lowercase_		,			device=lowercase_		,			dtype=lowercase_		,			shape=lowercase_)
              SCREAMING_SNAKE_CASE_   :						int							=	latents

          return inputs
      def _SCREAMING_SNAKE_CASE				(		self  :						str		,			lowercase_  :						Any=0		,			lowercase_  :						Union[str, Any]="cpu"		,			lowercase_  :						Union[str, Any]=torch.floataa		,			lowercase_  :						Dict=(1, 3, 64, 64)):




          '''simple docstring'''


          if type(lowercase_) == str:
              SCREAMING_SNAKE_CASE_   :						Any							=	torch.device(lowercase_)
          SCREAMING_SNAKE_CASE_   :						Any							=	torch.Generator(device=lowercase_).manual_seed(lowercase_)
          SCREAMING_SNAKE_CASE_   :						str							=	randn_tensor(lowercase_		,			generator=lowercase_		,			device=lowercase_		,			dtype=lowercase_)
          return latents
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


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

          SCREAMING_SNAKE_CASE_   :						List[str]							=	self.get_inputs()
          SCREAMING_SNAKE_CASE_   :						List[str]							=	pipe(**lowercase_).images
          assert image.shape == (1, 64, 64, 3)

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

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Any							=	UNetaDModel.from_pretrained('''diffusers/consistency_models'''		,			subfolder='''diffusers_cd_imagenet64_l2''')
          SCREAMING_SNAKE_CASE_   :						str							=	CMStochasticIterativeScheduler(
              num_train_timesteps=40		,			sigma_min=0.0_02		,			sigma_max=80.0		,			)
          SCREAMING_SNAKE_CASE_   :						int							=	ConsistencyModelPipeline(unet=lowercase_		,			scheduler=lowercase_)
          pipe.to(torch_device=lowercase_)
          pipe.set_progress_bar_config(disable=lowercase_)

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

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

          SCREAMING_SNAKE_CASE_   :						List[str]							=	np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2
      @require_torch_a
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


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

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

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

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3



      @require_torch_a
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	UNetaDModel.from_pretrained('''diffusers/consistency_models'''		,			subfolder='''diffusers_cd_imagenet64_l2''')
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	CMStochasticIterativeScheduler(
              num_train_timesteps=40		,			sigma_min=0.0_02		,			sigma_max=80.0		,			)
          SCREAMING_SNAKE_CASE_   :						Any							=	ConsistencyModelPipeline(unet=lowercase_		,			scheduler=lowercase_)
          pipe.to(torch_device=lowercase_		,			torch_dtype=torch.floataa)
          pipe.set_progress_bar_config(disable=lowercase_)

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

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

          SCREAMING_SNAKE_CASE_   :						Dict							=	np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3










"""simple docstring"""






import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()




class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      @property
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	UNetaDModel(
              block_out_channels=(32, 64)		,			layers_per_block=2		,			sample_size=32		,			in_channels=3		,			out_channels=3		,			down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')		,			up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')		,			)
          return model



      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	self.dummy_uncond_unet
          SCREAMING_SNAKE_CASE_   :						int							=	ScoreSdeVeScheduler()

          SCREAMING_SNAKE_CASE_   :						List[str]							=	ScoreSdeVePipeline(unet=lowercase_		,			scheduler=lowercase_)
          sde_ve.to(lowercase_)
          sde_ve.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						Dict							=	sde_ve(num_inference_steps=2		,			output_type='''numpy'''		,			generator=lowercase_).images

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	sde_ve(num_inference_steps=2		,			output_type='''numpy'''		,			generator=lowercase_		,			return_dict=lowercase_)[
              0
          ]

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

          assert image.shape == (1, 32, 32, 3)
          SCREAMING_SNAKE_CASE_   :						str							=	np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2




@slow
@require_torch
class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	'''google/ncsnpp-church-256'''
          SCREAMING_SNAKE_CASE_   :						str							=	UNetaDModel.from_pretrained(lowercase_)

          SCREAMING_SNAKE_CASE_   :						Any							=	ScoreSdeVeScheduler.from_pretrained(lowercase_)

          SCREAMING_SNAKE_CASE_   :						Any							=	ScoreSdeVePipeline(unet=lowercase_		,			scheduler=lowercase_)
          sde_ve.to(lowercase_)
          sde_ve.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	sde_ve(num_inference_steps=10		,			output_type='''numpy'''		,			generator=lowercase_).images

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

          assert image.shape == (1, 256, 256, 3)

          SCREAMING_SNAKE_CASE_   :						int							=	np.array([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2


