

'''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 				_snake_case					(   a_     , unittest.TestCase			):
  SCREAMING_SNAKE_CASE			:       Optional[int]				=     ConsistencyModelPipeline
  SCREAMING_SNAKE_CASE			:       Optional[int]				=     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  SCREAMING_SNAKE_CASE			:       Tuple				=     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE=False	):





      '''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   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE=0	):





      '''simple docstring'''

      if str(_SCREAMING_SNAKE_CASE	).startswith('mps'	):
          lowerCAmelCase 						=							torch.manual_seed(_SCREAMING_SNAKE_CASE	)
      else:
          lowerCAmelCase 						=							torch.Generator(device=_SCREAMING_SNAKE_CASE	).manual_seed(_SCREAMING_SNAKE_CASE	)

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

      return inputs

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

      lowerCAmelCase 						=							self.get_dummy_inputs(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

      lowerCAmelCase 						=							self.get_dummy_inputs(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							0
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

      lowerCAmelCase 						=							self.get_dummy_inputs(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							1
      lowerCAmelCase 						=							None
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

      lowerCAmelCase 						=							self.get_dummy_inputs(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							1
      lowerCAmelCase 						=							None
      lowerCAmelCase 						=							0
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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 				_snake_case					(   unittest.TestCase			):

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

  def   _SCREAMING_SNAKE_CASE				(				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'''

      lowerCAmelCase 						=							torch.manual_seed(_SCREAMING_SNAKE_CASE	)

      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=_SCREAMING_SNAKE_CASE      , device=_SCREAMING_SNAKE_CASE      , dtype=_SCREAMING_SNAKE_CASE      , shape=_SCREAMING_SNAKE_CASE	)
          lowerCAmelCase 						=							latents

      return inputs

  def   _SCREAMING_SNAKE_CASE				(				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(_SCREAMING_SNAKE_CASE	) == str:
          lowerCAmelCase 						=							torch.device(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							torch.Generator(device=_SCREAMING_SNAKE_CASE	).manual_seed(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							randn_tensor(_SCREAMING_SNAKE_CASE      , generator=_SCREAMING_SNAKE_CASE      , device=_SCREAMING_SNAKE_CASE      , dtype=_SCREAMING_SNAKE_CASE	)
      return latents

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''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=_SCREAMING_SNAKE_CASE      , scheduler=_SCREAMING_SNAKE_CASE	)
      pipe.to(torch_device=_SCREAMING_SNAKE_CASE	)
      pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.get_inputs()
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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   _SCREAMING_SNAKE_CASE				(				self	):





      '''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=_SCREAMING_SNAKE_CASE      , scheduler=_SCREAMING_SNAKE_CASE	)
      pipe.to(torch_device=_SCREAMING_SNAKE_CASE	)
      pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.get_inputs()
      lowerCAmelCase 						=							1
      lowerCAmelCase 						=							None
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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   _SCREAMING_SNAKE_CASE				(				self	):





      '''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=_SCREAMING_SNAKE_CASE      , scheduler=_SCREAMING_SNAKE_CASE	)
      pipe.to(torch_device=_SCREAMING_SNAKE_CASE      , torch_dtype=torch.floataa	)
      pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.get_inputs(get_fixed_latents=_SCREAMING_SNAKE_CASE      , device=_SCREAMING_SNAKE_CASE	)
      # Ensure usage of flash attention in torch 2.0
      with sdp_kernel(enable_flash=_SCREAMING_SNAKE_CASE      , enable_math=_SCREAMING_SNAKE_CASE      , enable_mem_efficient=_SCREAMING_SNAKE_CASE	):
          lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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   _SCREAMING_SNAKE_CASE				(				self	):





      '''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=_SCREAMING_SNAKE_CASE      , scheduler=_SCREAMING_SNAKE_CASE	)
      pipe.to(torch_device=_SCREAMING_SNAKE_CASE      , torch_dtype=torch.floataa	)
      pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.get_inputs(get_fixed_latents=_SCREAMING_SNAKE_CASE      , device=_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							1
      lowerCAmelCase 						=							None
      # Ensure usage of flash attention in torch 2.0
      with sdp_kernel(enable_flash=_SCREAMING_SNAKE_CASE      , enable_math=_SCREAMING_SNAKE_CASE      , enable_mem_efficient=_SCREAMING_SNAKE_CASE	):
          lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).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'''



from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import torch

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available



@dataclass
class 				_snake_case					(   a_			):
  SCREAMING_SNAKE_CASE			:       Union[List[np.ndarray], torch.FloatTensor]


try:
     if not (is_transformers_available() and is_torch_available()):
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
     from .pipeline_text_to_video_synth import TextToVideoSDPipeline
     from .pipeline_text_to_video_synth_imgaimg import VideoToVideoSDPipeline  # noqa: F401
     from .pipeline_text_to_video_zero import TextToVideoZeroPipeline





