


"""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 				SCREAMING_SNAKE_CASE__    (		lowercase       ,   unittest.TestCase    ):
       """simple docstring"""


       a  :    List[str]       =ConsistencyModelPipeline
       a  :    Optional[int]       =UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       a  :    Dict       =UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

       @property
       def 	lowercase__							(						self						):






         """simple docstring"""





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

       @property
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[Any]								=							UNetaDModel.from_pretrained(
             "diffusers/consistency-models-test"		,					subfolder="test_unet_class_cond"		,					)
         return unet

       def 	lowercase__							(						self		,					snake_case__=False						):






         """simple docstring"""





         if class_cond:
           lowerCAmelCase					:					List[str]								=							self.dummy_cond_unet
         else:
           lowerCAmelCase					:					List[Any]								=							self.dummy_uncond_unet

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

         lowerCAmelCase					:					List[str]								=							{
             "unet": unet,
             "scheduler": scheduler,
         }

         return components

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__=0						):






         """simple docstring"""





         if str(snake_case__						).startswith("mps"						):
           lowerCAmelCase					:					Dict								=							torch.manual_seed(snake_case__						)
         else:
           lowerCAmelCase					:					List[Any]								=							torch.Generator(device=snake_case__						).manual_seed(snake_case__						)

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

         return inputs

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							"cpu"  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase					:					Tuple								=							self.get_dummy_components()
         lowerCAmelCase					:					Dict								=							ConsistencyModelPipeline(**snake_case__						)
         lowerCAmelCase					:					Optional[int]								=							pipe.to(snake_case__						)
         pipe.set_progress_bar_config(disable=snake_case__						)

         lowerCAmelCase					:					List[str]								=							self.get_dummy_inputs(snake_case__						)
         lowerCAmelCase					:					Dict								=							pipe(**snake_case__						).images
         assert image.shape == (1, 32, 32, 3)

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






         """simple docstring"""





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

         lowerCAmelCase					:					str								=							self.get_dummy_inputs(snake_case__						)
         lowerCAmelCase					:					Optional[int]								=							0
         lowerCAmelCase					:					int								=							pipe(**snake_case__						).images
         assert image.shape == (1, 32, 32, 3)

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






         """simple docstring"""





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

         lowerCAmelCase					:					List[Any]								=							self.get_dummy_inputs(snake_case__						)
         lowerCAmelCase					:					Dict								=							1
         lowerCAmelCase					:					Dict								=							None
         lowerCAmelCase					:					Union[str, Any]								=							pipe(**snake_case__						).images
         assert image.shape == (1, 32, 32, 3)

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






         """simple docstring"""





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

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

         lowerCAmelCase					:					List[str]								=							image[0, -3:, -3:, -1]
         lowerCAmelCase					:					str								=							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 				SCREAMING_SNAKE_CASE__    (		unittest.TestCase    ):
       """simple docstring"""


       def 	lowercase__							(						self						):






         """simple docstring"""





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

       def 	lowercase__							(						self		,					snake_case__=0		,					snake_case__=False		,					snake_case__="cpu"		,					snake_case__=torch.floataa		,					snake_case__=(1, 3, 64, 64)						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							torch.manual_seed(snake_case__						)

         lowerCAmelCase					:					Dict								=							{
             "num_inference_steps": None,
             "timesteps": [22, 0],
             "class_labels": 0,
             "generator": generator,
             "output_type": "np",
         }

         if get_fixed_latents:
           lowerCAmelCase					:					int								=							self.get_fixed_latents(seed=snake_case__		,					device=snake_case__		,					dtype=snake_case__		,					shape=snake_case__						)
           lowerCAmelCase					:					Dict								=							latents

         return inputs

       def 	lowercase__							(						self		,					snake_case__=0		,					snake_case__="cpu"		,					snake_case__=torch.floataa		,					snake_case__=(1, 3, 64, 64)						):






         """simple docstring"""





         if type(snake_case__						) == str:
           lowerCAmelCase					:					Tuple								=							torch.device(snake_case__						)
         lowerCAmelCase					:					Optional[Any]								=							torch.Generator(device=snake_case__						).manual_seed(snake_case__						)
         lowerCAmelCase					:					List[str]								=							randn_tensor(snake_case__		,					generator=snake_case__		,					device=snake_case__		,					dtype=snake_case__						)
         return latents

       def 	lowercase__							(						self						):






         """simple docstring"""





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

         lowerCAmelCase					:					Optional[int]								=							self.get_inputs()
         lowerCAmelCase					:					List[str]								=							pipe(**snake_case__						).images
         assert image.shape == (1, 64, 64, 3)

         lowerCAmelCase					:					Dict								=							image[0, -3:, -3:, -1]

         lowerCAmelCase					:					Optional[int]								=							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 	lowercase__							(						self						):






         """simple docstring"""





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

         lowerCAmelCase					:					Any								=							self.get_inputs()
         lowerCAmelCase					:					str								=							1
         lowerCAmelCase					:					List[str]								=							None
         lowerCAmelCase					:					Tuple								=							pipe(**snake_case__						).images
         assert image.shape == (1, 64, 64, 3)

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

         lowerCAmelCase					:					List[Any]								=							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 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					str								=							UNetaDModel.from_pretrained("diffusers/consistency_models"		,					subfolder="diffusers_cd_imagenet64_l2"						)
         lowerCAmelCase					:					List[str]								=							CMStochasticIterativeScheduler(
             num_train_timesteps=40		,					sigma_min=0.002		,					sigma_max=80.0		,					)
         lowerCAmelCase					:					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__						)

         lowerCAmelCase					:					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__						):
           lowerCAmelCase					:					Dict								=							pipe(**snake_case__						).images
         assert image.shape == (1, 64, 64, 3)

         lowerCAmelCase					:					List[str]								=							image[0, -3:, -3:, -1]

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






         """simple docstring"""





         lowerCAmelCase					:					List[Any]								=							UNetaDModel.from_pretrained("diffusers/consistency_models"		,					subfolder="diffusers_cd_imagenet64_l2"						)
         lowerCAmelCase					:					Union[str, Any]								=							CMStochasticIterativeScheduler(
             num_train_timesteps=40		,					sigma_min=0.002		,					sigma_max=80.0		,					)
         lowerCAmelCase					:					Dict								=							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__						)

         lowerCAmelCase					:					Any								=							self.get_inputs(get_fixed_latents=snake_case__		,					device=snake_case__						)
         lowerCAmelCase					:					Tuple								=							1
         lowerCAmelCase					:					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__						):
           lowerCAmelCase					:					str								=							pipe(**snake_case__						).images
         assert image.shape == (1, 64, 64, 3)

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

         lowerCAmelCase					:					Union[str, Any]								=							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









"""simple docstring"""







import argparse

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_controlnet_from_original_ckpt


if __name__ == "__main__":
     lowerCAmelCase__				    =       argparse.ArgumentParser()

     parser.add_argument(
         '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
     )
     parser.add_argument(
         '''--original_config_file''',
         type=str,
         required=True,
         help='''The YAML config file corresponding to the original architecture.''',
     )
     parser.add_argument(
         '''--num_in_channels''',
         default=None,
         type=int,
         help='''The number of input channels. If `None` number of input channels will be automatically inferred.''',
     )
     parser.add_argument(
         '''--image_size''',
         default=512,
         type=int,
         help=(
             '''The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2'''
             ''' Base. Use 768 for Stable Diffusion v2.'''
         ),
     )
     parser.add_argument(
         '''--extract_ema''',
         action='''store_true''',
         help=(
             '''Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights'''
             ''' or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield'''
             ''' higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning.'''
         ),
     )
     parser.add_argument(
         '''--upcast_attention''',
         action='''store_true''',
         help=(
             '''Whether the attention computation should always be upcasted. This is necessary when running stable'''
             ''' diffusion 2.1.'''
         ),
     )
     parser.add_argument(
         '''--from_safetensors''',
         action='''store_true''',
         help='''If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.''',
     )
     parser.add_argument(
         '''--to_safetensors''',
         action='''store_true''',
         help='''Whether to store pipeline in safetensors format or not.''',
     )
     parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')
     parser.add_argument('''--device''', type=str, help='''Device to use (e.g. cpu, cuda:0, cuda:1, etc.)''')
     def  a__      (    SCREAMING_SNAKE_CASE   :   Tuple    ):


       '''simple docstring'''







       if string == "True":
         return True
       elif string == "False":
         return False
       else:
         raise ValueError(f"""could not parse string as bool {string}"""    )

     parser.add_argument(
         '''--use_linear_projection''', help='''Override for use linear projection''', required=False, type=parse_bool
     )

     parser.add_argument('''--cross_attention_dim''', help='''Override for cross attention_dim''', required=False, type=int)

     lowerCAmelCase__				    =       parser.parse_args()

     lowerCAmelCase__				    =       download_controlnet_from_original_ckpt(
         checkpoint_path=args.checkpoint_path,
         original_config_file=args.original_config_file,
         image_size=args.image_size,
         extract_ema=args.extract_ema,
         num_in_channels=args.num_in_channels,
         upcast_attention=args.upcast_attention,
         from_safetensors=args.from_safetensors,
         device=args.device,
         use_linear_projection=args.use_linear_projection,
         cross_attention_dim=args.cross_attention_dim,
     )





     controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)





