


'''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  __magic_name__							(		lowerCAmelCase ,unittest.TestCase    ):
       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       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =UNetaDModel.from_pretrained(
              'diffusers/consistency-models-test'  ,							subfolder='test_unet'  ,							)
          return unet






       @property
       def       lowerCAmelCase		(      self)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =UNetaDModel.from_pretrained(
              'diffusers/consistency-models-test'  ,							subfolder='test_unet_class_cond'  ,							)
          return unet






       def       lowerCAmelCase		(      self  ,							snake_case=False)						->  Dict:




          '''simple docstring'''
          if class_cond:
             _UpperCAmelCase       :       List[Any]							   =self.dummy_cond_unet
          else:
             _UpperCAmelCase       :       Optional[Any]							   =self.dummy_uncond_unet

          # Default to CM multistep sampler
          _UpperCAmelCase       :       Optional[int]							   =CMStochasticIterativeScheduler(
              num_train_timesteps=4_0  ,							sigma_min=0.0_02  ,							sigma_max=80.0  ,							)

          _UpperCAmelCase       :       Optional[int]							   ={
              'unet': unet,
              'scheduler': scheduler,
          }

          return components






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case=0)						->  Dict:




          '''simple docstring'''
          if str(snake_case).startswith('mps'):
             _UpperCAmelCase       :       Union[str, Any]							   =torch.manual_seed(snake_case)
          else:
             _UpperCAmelCase       :       int							   =torch.Generator(device=snake_case).manual_seed(snake_case)

          _UpperCAmelCase       :       Tuple							   ={
              'batch_size': 1,
              'num_inference_steps': None,
              'timesteps': [2_2, 0],
              'generator': generator,
              'output_type': 'np',
          }

          return inputs






       def       lowerCAmelCase		(      self)						->  Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   ='cpu'  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase       :       str							   =self.get_dummy_components()
          _UpperCAmelCase       :       List[str]							   =ConsistencyModelPipeline(**snake_case)
          _UpperCAmelCase       :       Any							   =pipe.to(snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Any							   =self.get_dummy_inputs(snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =pipe(**snake_case).images
          assert image.shape == (1, 3_2, 3_2, 3)

          _UpperCAmelCase       :       List[str]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       int							   =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       lowerCAmelCase		(      self)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   ='cpu'  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase       :       str							   =self.get_dummy_components(class_cond=snake_case)
          _UpperCAmelCase       :       Tuple							   =ConsistencyModelPipeline(**snake_case)
          _UpperCAmelCase       :       str							   =pipe.to(snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Union[str, Any]							   =self.get_dummy_inputs(snake_case)
          _UpperCAmelCase       :       str							   =0
          _UpperCAmelCase       :       Dict							   =pipe(**snake_case).images
          assert image.shape == (1, 3_2, 3_2, 3)

          _UpperCAmelCase       :       List[str]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Any							   =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       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       List[Any]							   ='cpu'  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase       :       List[str]							   =self.get_dummy_components()
          _UpperCAmelCase       :       List[Any]							   =ConsistencyModelPipeline(**snake_case)
          _UpperCAmelCase       :       Dict							   =pipe.to(snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =self.get_dummy_inputs(snake_case)
          _UpperCAmelCase       :       str							   =1
          _UpperCAmelCase       :       str							   =None
          _UpperCAmelCase       :       Optional[Any]							   =pipe(**snake_case).images
          assert image.shape == (1, 3_2, 3_2, 3)

          _UpperCAmelCase       :       str							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       List[str]							   =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       lowerCAmelCase		(      self)						->  Any:




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

          _UpperCAmelCase       :       Any							   =self.get_dummy_inputs(snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =1
          _UpperCAmelCase       :       str							   =None
          _UpperCAmelCase       :       Optional[int]							   =0
          _UpperCAmelCase       :       str							   =pipe(**snake_case).images
          assert image.shape == (1, 3_2, 3_2, 3)

          _UpperCAmelCase       :       Optional[int]							   =image[0, -3:, -3:, -1]
          _UpperCAmelCase       :       Optional[int]							   =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  __magic_name__							(		unittest.TestCase    ):






       def       lowerCAmelCase		(      self)						->  Dict:




          '''simple docstring'''
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






       def       lowerCAmelCase		(      self  ,							snake_case=0  ,							snake_case=False  ,							snake_case="cpu"  ,							snake_case=torch.floataa  ,							snake_case=(1, 3, 6_4, 6_4))						->  Dict:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =torch.manual_seed(snake_case)

          _UpperCAmelCase       :       Union[str, Any]							   ={
              'num_inference_steps': None,
              'timesteps': [2_2, 0],
              'class_labels': 0,
              'generator': generator,
              'output_type': 'np',
          }

          if get_fixed_latents:
             _UpperCAmelCase       :       Optional[Any]							   =self.get_fixed_latents(seed=snake_case  ,							device=snake_case  ,							dtype=snake_case  ,							shape=snake_case)
             _UpperCAmelCase       :       List[Any]							   =latents

          return inputs






       def       lowerCAmelCase		(      self  ,							snake_case=0  ,							snake_case="cpu"  ,							snake_case=torch.floataa  ,							snake_case=(1, 3, 6_4, 6_4))						->  Any:




          '''simple docstring'''
          if type(snake_case) == str:
             _UpperCAmelCase       :       Optional[int]							   =torch.device(snake_case)
          _UpperCAmelCase       :       str							   =torch.Generator(device=snake_case).manual_seed(snake_case)
          _UpperCAmelCase       :       List[str]							   =randn_tensor(snake_case  ,							generator=snake_case  ,							device=snake_case  ,							dtype=snake_case)
          return latents






       def       lowerCAmelCase		(      self)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       List[Any]							   =UNetaDModel.from_pretrained('diffusers/consistency_models'  ,							subfolder='diffusers_cd_imagenet64_l2')
          _UpperCAmelCase       :       Optional[Any]							   =CMStochasticIterativeScheduler(
              num_train_timesteps=4_0  ,							sigma_min=0.0_02  ,							sigma_max=80.0  ,							)
          _UpperCAmelCase       :       Optional[int]							   =ConsistencyModelPipeline(unet=snake_case  ,							scheduler=snake_case)
          pipe.to(torch_device=snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       List[Any]							   =self.get_inputs()
          _UpperCAmelCase       :       Optional[Any]							   =pipe(**snake_case).images
          assert image.shape == (1, 6_4, 6_4, 3)

          _UpperCAmelCase       :       Union[str, Any]							   =image[0, -3:, -3:, -1]

          _UpperCAmelCase       :       Tuple							   =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       lowerCAmelCase		(      self)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       Any							   =UNetaDModel.from_pretrained('diffusers/consistency_models'  ,							subfolder='diffusers_cd_imagenet64_l2')
          _UpperCAmelCase       :       List[str]							   =CMStochasticIterativeScheduler(
              num_train_timesteps=4_0  ,							sigma_min=0.0_02  ,							sigma_max=80.0  ,							)
          _UpperCAmelCase       :       List[str]							   =ConsistencyModelPipeline(unet=snake_case  ,							scheduler=snake_case)
          pipe.to(torch_device=snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =self.get_inputs()
          _UpperCAmelCase       :       List[str]							   =1
          _UpperCAmelCase       :       Tuple							   =None
          _UpperCAmelCase       :       Tuple							   =pipe(**snake_case).images
          assert image.shape == (1, 6_4, 6_4, 3)

          _UpperCAmelCase       :       Optional[Any]							   =image[0, -3:, -3:, -1]

          _UpperCAmelCase       :       Optional[Any]							   =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       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =UNetaDModel.from_pretrained('diffusers/consistency_models'  ,							subfolder='diffusers_cd_imagenet64_l2')
          _UpperCAmelCase       :       str							   =CMStochasticIterativeScheduler(
              num_train_timesteps=4_0  ,							sigma_min=0.0_02  ,							sigma_max=80.0  ,							)
          _UpperCAmelCase       :       Any							   =ConsistencyModelPipeline(unet=snake_case  ,							scheduler=snake_case)
          pipe.to(torch_device=snake_case  ,							torch_dtype=torch.floataa)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =self.get_inputs(get_fixed_latents=snake_case  ,							device=snake_case)
          # Ensure usage of flash attention in torch 2.0
          with sdp_kernel(enable_flash=snake_case  ,							enable_math=snake_case  ,							enable_mem_efficient=snake_case):
             _UpperCAmelCase       :       str							   =pipe(**snake_case).images
          assert image.shape == (1, 6_4, 6_4, 3)

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

          _UpperCAmelCase       :       str							   =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       lowerCAmelCase		(      self)						->  Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =UNetaDModel.from_pretrained('diffusers/consistency_models'  ,							subfolder='diffusers_cd_imagenet64_l2')
          _UpperCAmelCase       :       List[str]							   =CMStochasticIterativeScheduler(
              num_train_timesteps=4_0  ,							sigma_min=0.0_02  ,							sigma_max=80.0  ,							)
          _UpperCAmelCase       :       int							   =ConsistencyModelPipeline(unet=snake_case  ,							scheduler=snake_case)
          pipe.to(torch_device=snake_case  ,							torch_dtype=torch.floataa)
          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase       :       Any							   =self.get_inputs(get_fixed_latents=snake_case  ,							device=snake_case)
          _UpperCAmelCase       :       List[str]							   =1
          _UpperCAmelCase       :       Tuple							   =None
          # Ensure usage of flash attention in torch 2.0
          with sdp_kernel(enable_flash=snake_case  ,							enable_math=snake_case  ,							enable_mem_efficient=snake_case):
             _UpperCAmelCase       :       Any							   =pipe(**snake_case).images
          assert image.shape == (1, 6_4, 6_4, 3)

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

          _UpperCAmelCase       :       Optional[Any]							   =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'''






from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowercase							=logging.get_logger(__name__)

lowercase							={
    'facebook/levit-128S': 'https://huggingface.co/facebook/levit-128S/resolve/main/config.json',
    # See all LeViT models at https://huggingface.co/models?filter=levit
}






class  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										="levit"






       def __init__(      self  ,							snake_case=2_2_4  ,							snake_case=3  ,							snake_case=3  ,							snake_case=2  ,							snake_case=1  ,							snake_case=1_6  ,							snake_case=[1_2_8, 2_5_6, 3_8_4]  ,							snake_case=[4, 8, 1_2]  ,							snake_case=[4, 4, 4]  ,							snake_case=[1_6, 1_6, 1_6]  ,							snake_case=0  ,							snake_case=[2, 2, 2]  ,							snake_case=[2, 2, 2]  ,							snake_case=0.02  ,							**snake_case  ,							)						->  Any:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       List[str]							   =image_size
          _UpperCAmelCase       :       str							   =num_channels
          _UpperCAmelCase       :       int							   =kernel_size
          _UpperCAmelCase       :       Any							   =stride
          _UpperCAmelCase       :       Tuple							   =padding
          _UpperCAmelCase       :       Optional[int]							   =hidden_sizes
          _UpperCAmelCase       :       List[str]							   =num_attention_heads
          _UpperCAmelCase       :       List[Any]							   =depths
          _UpperCAmelCase       :       Any							   =key_dim
          _UpperCAmelCase       :       List[Any]							   =drop_path_rate
          _UpperCAmelCase       :       int							   =patch_size
          _UpperCAmelCase       :       Tuple							   =attention_ratio
          _UpperCAmelCase       :       Any							   =mlp_ratio
          _UpperCAmelCase       :       Optional[Any]							   =initializer_range
          _UpperCAmelCase       :       str							   =[
              ['Subsample', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
              ['Subsample', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
          ]






class  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										=version.parse("1.11"    )






       @property
       def       lowerCAmelCase		(      self)						->  Mapping[str, Mapping[int, str]]:




          '''simple docstring'''
          return OrderedDict(
              [
                  ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
              ])






       @property
       def       lowerCAmelCase		(      self)						->  float:




          '''simple docstring'''
          return 1E-4

