

"""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      __a							(_UpperCamelCase   ,	unittest.TestCase):






     '''simple docstring'''
     _SCREAMING_SNAKE_CASE						:str       = ConsistencyModelPipeline
     _SCREAMING_SNAKE_CASE						:Optional[Any]       = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
     _SCREAMING_SNAKE_CASE						:Dict       = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

     @property
     def 					_a			(						self  )  ->      Optional[int]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	UNetaDModel.from_pretrained(
               """diffusers/consistency-models-test"""	,    subfolder="""test_unet"""	,    )
           return unet

     @property
     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Dict      		=	UNetaDModel.from_pretrained(
               """diffusers/consistency-models-test"""	,    subfolder="""test_unet_class_cond"""	,    )
           return unet

     def 					_a			(						self	,    _a=False  )  ->      Tuple:


           """simple docstring"""


           if class_cond:
                 SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	self.dummy_cond_unet
           else:
                 SCREAMING_SNAKE_CASE__   :      Any      		=	self.dummy_uncond_unet

           # Default to CM multistep sampler
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	CMStochasticIterativeScheduler(
               num_train_timesteps=40	,    sigma_min=0.002	,    sigma_max=80.0	,    )

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	{
               "unet": unet,
               "scheduler": scheduler,
           }

           return components

     def 					_a			(						self	,    _a	,    _a=0  )  ->      Optional[int]:


           """simple docstring"""


           if str(__a  ).startswith("""mps"""  ):
                 SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	torch.manual_seed(__a  )
           else:
                 SCREAMING_SNAKE_CASE__   :      List[str]      		=	torch.Generator(device=__a  ).manual_seed(__a  )

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

           return inputs

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


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

           SCREAMING_SNAKE_CASE__   :      Dict      		=	self.get_dummy_inputs(__a  )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	pipe(**__a  ).images
           assert image.shape == (1, 32, 32, 3)

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	image[0, -3:, -3:, -1]
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	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 					_a			(						self  )  ->      Union[str, Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	"cpu"  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	self.get_dummy_components(class_cond=__a  )
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	ConsistencyModelPipeline(**__a  )
           SCREAMING_SNAKE_CASE__   :      Tuple      		=	pipe.to(__a  )
           pipe.set_progress_bar_config(disable=__a  )

           SCREAMING_SNAKE_CASE__   :      Tuple      		=	self.get_dummy_inputs(__a  )
           SCREAMING_SNAKE_CASE__   :      Tuple      		=	0
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	pipe(**__a  ).images
           assert image.shape == (1, 32, 32, 3)

           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	image[0, -3:, -3:, -1]
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	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 					_a			(						self  )  ->      List[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Tuple      		=	"cpu"  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE__   :      Any      		=	self.get_dummy_components()
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	ConsistencyModelPipeline(**__a  )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	pipe.to(__a  )
           pipe.set_progress_bar_config(disable=__a  )

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

           SCREAMING_SNAKE_CASE__   :      str      		=	image[0, -3:, -3:, -1]
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	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 					_a			(						self  )  ->      Union[str, Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Tuple      		=	"cpu"  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.get_dummy_components(class_cond=__a  )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	ConsistencyModelPipeline(**__a  )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	pipe.to(__a  )
           pipe.set_progress_bar_config(disable=__a  )

           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	self.get_dummy_inputs(__a  )
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	1
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	None
           SCREAMING_SNAKE_CASE__   :      int      		=	0
           SCREAMING_SNAKE_CASE__   :      int      		=	pipe(**__a  ).images
           assert image.shape == (1, 32, 32, 3)

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






     '''simple docstring'''

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


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

     def 					_a			(						self	,    _a=0	,    _a=False	,    _a="cpu"	,    _a=torch.floataa	,    _a=(1, 3, 64, 64)  )  ->      Tuple:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	torch.manual_seed(__a  )

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

           if get_fixed_latents:
                 SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.get_fixed_latents(seed=__a	,    device=__a	,    dtype=__a	,    shape=__a  )
                 SCREAMING_SNAKE_CASE__   :      List[Any]      		=	latents

           return inputs

     def 					_a			(						self	,    _a=0	,    _a="cpu"	,    _a=torch.floataa	,    _a=(1, 3, 64, 64)  )  ->      Dict:


           """simple docstring"""


           if type(__a  ) == str:
                 SCREAMING_SNAKE_CASE__   :      int      		=	torch.device(__a  )
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	torch.Generator(device=__a  ).manual_seed(__a  )
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	randn_tensor(__a	,    generator=__a	,    device=__a	,    dtype=__a  )
           return latents

     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[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.002	,    sigma_max=80.0	,    )
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	ConsistencyModelPipeline(unet=__a	,    scheduler=__a  )
           pipe.to(torch_device=__a  )
           pipe.set_progress_bar_config(disable=__a  )

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

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

           SCREAMING_SNAKE_CASE__   :      Tuple      		=	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 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	UNetaDModel.from_pretrained("""diffusers/consistency_models"""	,    subfolder="""diffusers_cd_imagenet64_l2"""  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	CMStochasticIterativeScheduler(
               num_train_timesteps=40	,    sigma_min=0.002	,    sigma_max=80.0	,    )
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	ConsistencyModelPipeline(unet=__a	,    scheduler=__a  )
           pipe.to(torch_device=__a  )
           pipe.set_progress_bar_config(disable=__a  )

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.get_inputs()
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	1
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	None
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	pipe(**__a  ).images
           assert image.shape == (1, 64, 64, 3)

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

           SCREAMING_SNAKE_CASE__   :      int      		=	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 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	UNetaDModel.from_pretrained("""diffusers/consistency_models"""	,    subfolder="""diffusers_cd_imagenet64_l2"""  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	CMStochasticIterativeScheduler(
               num_train_timesteps=40	,    sigma_min=0.002	,    sigma_max=80.0	,    )
           SCREAMING_SNAKE_CASE__   :      Any      		=	ConsistencyModelPipeline(unet=__a	,    scheduler=__a  )
           pipe.to(torch_device=__a	,    torch_dtype=torch.floataa  )
           pipe.set_progress_bar_config(disable=__a  )

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	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  ):
                 SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	pipe(**__a  ).images
           assert image.shape == (1, 64, 64, 3)

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

           SCREAMING_SNAKE_CASE__   :      Tuple      		=	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 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	UNetaDModel.from_pretrained("""diffusers/consistency_models"""	,    subfolder="""diffusers_cd_imagenet64_l2"""  )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	CMStochasticIterativeScheduler(
               num_train_timesteps=40	,    sigma_min=0.002	,    sigma_max=80.0	,    )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	ConsistencyModelPipeline(unet=__a	,    scheduler=__a  )
           pipe.to(torch_device=__a	,    torch_dtype=torch.floataa  )
           pipe.set_progress_bar_config(disable=__a  )

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

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

           SCREAMING_SNAKE_CASE__   :      Any      		=	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"""


a					:List[str]    = [
    (1_000, "M"),
    (900, "CM"),
    (500, "D"),
    (400, "CD"),
    (100, "C"),
    (90, "XC"),
    (50, "L"),
    (40, "XL"),
    (10, "X"),
    (9, "IX"),
    (5, "V"),
    (4, "IV"),
    (1, "I"),
]




def  _lowercase					(    __lowerCAmelCase							)       ->				int:
      SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	{"""I""": 1, """V""": 5, """X""": 10, """L""": 50, """C""": 100, """D""": 500, """M""": 1000}
      SCREAMING_SNAKE_CASE__   :      List[Any]      		=	0
      SCREAMING_SNAKE_CASE__   :      List[str]      		=	0
      while place < len(__lowerCAmelCase							):
            if (place + 1 < len(__lowerCAmelCase							)) and (vals[roman[place]] < vals[roman[place + 1]]):
                  total += vals[roman[place + 1]] - vals[roman[place]]
                  place += 2
            else:
                  total += vals[roman[place]]
                  place += 1
      return total




def  _lowercase					(    __lowerCAmelCase							)       ->				str:
      SCREAMING_SNAKE_CASE__   :      Any      		=	[]
      for arabic, roman in ROMAN:
            ((SCREAMING_SNAKE_CASE__)		,		(SCREAMING_SNAKE_CASE__))   :      List[str]      		=	divmod(__lowerCAmelCase  ,					__lowerCAmelCase							)
            result.append(roman * factor							)
            if number == 0:
                  break
      return "".join(__lowerCAmelCase							)


if __name__ == "__main__":
  import doctest

  doctest.testmod()




