




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	):
      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 	__magic_name__(   self		):
        lowerCAmelCase__		       =     UNetaDModel.from_pretrained(
            '''diffusers/consistency-models-test'''      ,				subfolder='''test_unet'''      ,				)
        return unet

      @property
      def 	__magic_name__(   self		):
        lowerCAmelCase__		       =     UNetaDModel.from_pretrained(
            '''diffusers/consistency-models-test'''      ,				subfolder='''test_unet_class_cond'''      ,				)
        return unet

      def 	__magic_name__(   self      ,				__UpperCAmelCase=False		):
        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 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=0		):
        if str(__a		).startswith('''mps'''		):
          lowerCAmelCase__		       =     torch.manual_seed(__a		)
        else:
          lowerCAmelCase__		       =     torch.Generator(device=__a		).manual_seed(__a		)

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

        return inputs

      def 	__magic_name__(   self		):
        lowerCAmelCase__		       =     "cpu"  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__		       =     self.get_dummy_components()
        lowerCAmelCase__		       =     ConsistencyModelPipeline(**__a		)
        lowerCAmelCase__		       =     pipe.to(__a		)
        pipe.set_progress_bar_config(disable=__a		)

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

        lowerCAmelCase__		       =     image[0, -3:, -3:, -1]
        lowerCAmelCase__		       =     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 	__magic_name__(   self		):
        lowerCAmelCase__		       =     "cpu"  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__		       =     self.get_dummy_components(class_cond=__a		)
        lowerCAmelCase__		       =     ConsistencyModelPipeline(**__a		)
        lowerCAmelCase__		       =     pipe.to(__a		)
        pipe.set_progress_bar_config(disable=__a		)

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

        lowerCAmelCase__		       =     image[0, -3:, -3:, -1]
        lowerCAmelCase__		       =     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 	__magic_name__(   self		):
        lowerCAmelCase__		       =     "cpu"  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__		       =     self.get_dummy_components()
        lowerCAmelCase__		       =     ConsistencyModelPipeline(**__a		)
        lowerCAmelCase__		       =     pipe.to(__a		)
        pipe.set_progress_bar_config(disable=__a		)

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

        lowerCAmelCase__		       =     image[0, -3:, -3:, -1]
        lowerCAmelCase__		       =     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 	__magic_name__(   self		):
        lowerCAmelCase__		       =     "cpu"  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__		       =     self.get_dummy_components(class_cond=__a		)
        lowerCAmelCase__		       =     ConsistencyModelPipeline(**__a		)
        lowerCAmelCase__		       =     pipe.to(__a		)
        pipe.set_progress_bar_config(disable=__a		)

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

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

      def 	__magic_name__(   self		):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

      def 	__magic_name__(   self      ,				__UpperCAmelCase=0      ,				__UpperCAmelCase=False      ,				__UpperCAmelCase="cpu"      ,				__UpperCAmelCase=torch.floataa      ,				__UpperCAmelCase=(1, 3, 64, 64)		):
        lowerCAmelCase__		       =     torch.manual_seed(__a		)

        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=__a      ,				device=__a      ,				dtype=__a      ,				shape=__a		)
          lowerCAmelCase__		       =     latents

        return inputs

      def 	__magic_name__(   self      ,				__UpperCAmelCase=0      ,				__UpperCAmelCase="cpu"      ,				__UpperCAmelCase=torch.floataa      ,				__UpperCAmelCase=(1, 3, 64, 64)		):
        if type(__a		) == str:
          lowerCAmelCase__		       =     torch.device(__a		)
        lowerCAmelCase__		       =     torch.Generator(device=__a		).manual_seed(__a		)
        lowerCAmelCase__		       =     randn_tensor(__a      ,				generator=__a      ,				device=__a      ,				dtype=__a		)
        return latents

      def 	__magic_name__(   self		):
        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=__a      ,				scheduler=__a		)
        pipe.to(torch_device=__a		)
        pipe.set_progress_bar_config(disable=__a		)

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

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

        lowerCAmelCase__		       =     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 	__magic_name__(   self		):
        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=__a      ,				scheduler=__a		)
        pipe.to(torch_device=__a		)
        pipe.set_progress_bar_config(disable=__a		)

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

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

        lowerCAmelCase__		       =     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 	__magic_name__(   self		):
        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=__a      ,				scheduler=__a		)
        pipe.to(torch_device=__a      ,				torch_dtype=torch.floataa		)
        pipe.set_progress_bar_config(disable=__a		)

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

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

        lowerCAmelCase__		       =     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 	__magic_name__(   self		):
        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=__a      ,				scheduler=__a		)
        pipe.to(torch_device=__a      ,				torch_dtype=torch.floataa		)
        pipe.set_progress_bar_config(disable=__a		)

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

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

        lowerCAmelCase__		       =     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 abc import ABC, abstractmethod
from argparse import ArgumentParser






class  _lowerCAmelCase     (  _lowercase	):

      @staticmethod
      @abstractmethod
      def 	__magic_name__(   __UpperCAmelCase		):
        raise NotImplementedError()







      @abstractmethod
      def 	__magic_name__(   self		):
        raise NotImplementedError()

