




'''simple docstring'''



import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 						lowercase_       (					A	,    unittest.TestCase  ):


   """simple docstring"""

   lowerCamelCase_							    =		AudioLDMPipeline
   lowerCamelCase_							    =		TEXT_TO_AUDIO_PARAMS
   lowerCamelCase_							    =		TEXT_TO_AUDIO_BATCH_PARAMS
   lowerCamelCase_							    =		frozenset(
       [
           '''num_inference_steps''',
           '''num_waveforms_per_prompt''',
           '''generator''',
           '''latents''',
           '''output_type''',
           '''return_dict''',
           '''callback''',
           '''callback_steps''',
       ]  )



   def 				lowerCAmelCase_     (       self   :      Optional[int]       ):

      """simple docstring"""
      torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)    , layers_per_block=2    , sample_size=3_2    , in_channels=4    , out_channels=4    , down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")    , up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")    , cross_attention_dim=(3_2, 6_4)    , class_embed_type="simple_projection"    , projection_class_embeddings_input_dim=3_2    , class_embeddings_concat=__lowerCamelCase    , )
      _SCREAMING_SNAKE_CASE			= DDIMScheduler(
          beta_start=0.0_0_0_8_5    , beta_end=0.0_1_2    , beta_schedule="scaled_linear"    , clip_sample=__lowerCamelCase    , set_alpha_to_one=__lowerCamelCase    , )
      torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= AutoencoderKL(
          block_out_channels=[3_2, 6_4]    , in_channels=1    , out_channels=1    , down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]    , up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]    , latent_channels=4    , )
      torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= ClapTextConfig(
          bos_token_id=0    , eos_token_id=2    , hidden_size=3_2    , intermediate_size=3_7    , layer_norm_eps=1e-05    , num_attention_heads=4    , num_hidden_layers=5    , pad_token_id=1    , vocab_size=1_0_0_0    , projection_dim=3_2    , )
      _SCREAMING_SNAKE_CASE			= ClapTextModelWithProjection(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"    , model_max_length=7_7       )

      _SCREAMING_SNAKE_CASE			= SpeechTaHifiGanConfig(
          model_in_dim=8    , sampling_rate=1_6_0_0_0    , upsample_initial_channel=1_6    , upsample_rates=[2, 2]    , upsample_kernel_sizes=[4, 4]    , resblock_kernel_sizes=[3, 7]    , resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]    , normalize_before=__lowerCamelCase    , )

      _SCREAMING_SNAKE_CASE			= SpeechTaHifiGan(__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= {
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "vocoder": vocoder,
      }
      return components



   def 				lowerCAmelCase_     (       self   :      str    , __lowerCamelCase   :      Optional[Any]    , __lowerCamelCase   :      List[Any]=0       ):

      """simple docstring"""


      if str(__lowerCamelCase       ).startswith("mps"       ):
         _SCREAMING_SNAKE_CASE			= torch.manual_seed(__lowerCamelCase       )
      else:
         _SCREAMING_SNAKE_CASE			= torch.Generator(device=__lowerCamelCase       ).manual_seed(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= {
          "prompt": "A hammer hitting a wooden surface",
          "generator": generator,
          "num_inference_steps": 2,
          "guidance_scale": 6.0,
      }
      return inputs



   def 				lowerCAmelCase_     (       self   :      Optional[Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= "cpu"  # ensure determinism for the device-dependent torch.Generator

      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      assert audio.ndim == 1
      assert len(__lowerCamelCase       ) == 2_5_6

      _SCREAMING_SNAKE_CASE			= audio[:1_0]
      _SCREAMING_SNAKE_CASE			= np.array(
          [-0.0_0_5_0, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_3, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_3]       )

      assert np.abs(audio_slice - expected_slice       ).max() < 1e-2



   def 				lowerCAmelCase_     (       self   :      Dict       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= 3 * [inputs["prompt"]]

      # forward
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= 3 * [inputs.pop("prompt"       )]

      _SCREAMING_SNAKE_CASE			= audioldm_pipe.tokenizer(
          __lowerCamelCase    , padding="max_length"    , max_length=audioldm_pipe.tokenizer.model_max_length    , truncation=__lowerCamelCase    , return_tensors="pt"    , )
      _SCREAMING_SNAKE_CASE			= text_inputs["input_ids"].to(__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= audioldm_pipe.text_encoder(
          __lowerCamelCase    , )
      _SCREAMING_SNAKE_CASE			= prompt_embeds.text_embeds
      # additional L_2 normalization over each hidden-state
      _SCREAMING_SNAKE_CASE			= F.normalize(__lowerCamelCase    , dim=-1       )

      _SCREAMING_SNAKE_CASE			= prompt_embeds

      # forward
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      assert np.abs(audio_a - audio_a       ).max() < 1e-2



   def 				lowerCAmelCase_     (       self   :      Tuple       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= 3 * ["this is a negative prompt"]
      _SCREAMING_SNAKE_CASE			= negative_prompt
      _SCREAMING_SNAKE_CASE			= 3 * [inputs["prompt"]]

      # forward
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= 3 * [inputs.pop("prompt"       )]

      _SCREAMING_SNAKE_CASE			= []
      for p in [prompt, negative_prompt]:
         _SCREAMING_SNAKE_CASE			= audioldm_pipe.tokenizer(
             __lowerCamelCase    , padding="max_length"    , max_length=audioldm_pipe.tokenizer.model_max_length    , truncation=__lowerCamelCase    , return_tensors="pt"    , )
         _SCREAMING_SNAKE_CASE			= text_inputs["input_ids"].to(__lowerCamelCase       )

         _SCREAMING_SNAKE_CASE			= audioldm_pipe.text_encoder(
             __lowerCamelCase    , )
         _SCREAMING_SNAKE_CASE			= text_embeds.text_embeds
         # additional L_2 normalization over each hidden-state
         _SCREAMING_SNAKE_CASE			= F.normalize(__lowerCamelCase    , dim=-1       )

         embeds.append(__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE,			_SCREAMING_SNAKE_CASE			= embeds

      # forward
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      assert np.abs(audio_a - audio_a       ).max() < 1e-2



   def 				lowerCAmelCase_     (       self   :      Tuple       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= "cpu"  # ensure determinism for the device-dependent torch.Generator
      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= PNDMScheduler(skip_prk_steps=__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= "egg cracking"
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase    , negative_prompt=__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      assert audio.ndim == 1
      assert len(__lowerCamelCase       ) == 2_5_6

      _SCREAMING_SNAKE_CASE			= audio[:1_0]
      _SCREAMING_SNAKE_CASE			= np.array(
          [-0.0_0_5_1, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_4, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_2]       )

      assert np.abs(audio_slice - expected_slice       ).max() < 1e-2



   def 				lowerCAmelCase_     (       self   :      int       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= "cpu"  # ensure determinism for the device-dependent torch.Generator
      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= PNDMScheduler(skip_prk_steps=__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= "A hammer hitting a wooden surface"

      # test num_waveforms_per_prompt=1 (default)
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(__lowerCamelCase    , num_inference_steps=2       ).audios

      assert audios.shape == (1, 2_5_6)

      # test num_waveforms_per_prompt=1 (default) for batch of prompts
      _SCREAMING_SNAKE_CASE			= 2
      _SCREAMING_SNAKE_CASE			= audioldm_pipe([prompt] * batch_size    , num_inference_steps=2       ).audios

      assert audios.shape == (batch_size, 2_5_6)

      # test num_waveforms_per_prompt for single prompt
      _SCREAMING_SNAKE_CASE			= 2
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(__lowerCamelCase    , num_inference_steps=2    , num_waveforms_per_prompt=__lowerCamelCase       ).audios

      assert audios.shape == (num_waveforms_per_prompt, 2_5_6)

      # test num_waveforms_per_prompt for batch of prompts
      _SCREAMING_SNAKE_CASE			= 2
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(
          [prompt] * batch_size    , num_inference_steps=2    , num_waveforms_per_prompt=__lowerCamelCase       ).audios

      assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_5_6)



   def 				lowerCAmelCase_     (       self   :      List[str]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= "cpu"  # ensure determinism for the device-dependent torch.Generator
      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.vocoder.config.sampling_rate

      _SCREAMING_SNAKE_CASE			= self.get_dummy_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(audio_length_in_s=0.0_1_6    , **__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      assert audio.ndim == 1
      assert len(__lowerCamelCase       ) / vocoder_sampling_rate == 0.0_1_6

      _SCREAMING_SNAKE_CASE			= audioldm_pipe(audio_length_in_s=0.0_3_2    , **__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= output.audios[0]

      assert audio.ndim == 1
      assert len(__lowerCamelCase       ) / vocoder_sampling_rate == 0.0_3_2



   def 				lowerCAmelCase_     (       self   :      Tuple       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.get_dummy_components()
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline(**__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= ["hey"]

      _SCREAMING_SNAKE_CASE			= audioldm_pipe(__lowerCamelCase    , num_inference_steps=1       )
      _SCREAMING_SNAKE_CASE			= output.audios.shape
      assert audio_shape == (1, 2_5_6)

      _SCREAMING_SNAKE_CASE			= audioldm_pipe.vocoder.config
      config.model_in_dim *= 2
      _SCREAMING_SNAKE_CASE			= SpeechTaHifiGan(__lowerCamelCase       ).to(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(__lowerCamelCase    , num_inference_steps=1       )
      _SCREAMING_SNAKE_CASE			= output.audios.shape
      # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
      assert audio_shape == (1, 2_5_6)



   def 				lowerCAmelCase_     (       self   :      int       ):

      """simple docstring"""
      self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowerCamelCase       )



   def 				lowerCAmelCase_     (       self   :      int       ):

      """simple docstring"""
      self._test_inference_batch_single_identical(test_mean_pixel_difference=__lowerCamelCase       )



   @unittest.skipIf(
       torch_device != "cuda" or not is_xformers_available()    , reason="XFormers attention is only available with CUDA and `xformers` installed"    , )
   def 				lowerCAmelCase_     (       self   :      int       ):

      """simple docstring"""
      self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowerCamelCase       )







@slow
class 						lowercase_       (					unittest.TestCase  ):


   """simple docstring"""

   def 				lowerCAmelCase_     (       self   :      int       ):

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



   def 				lowerCAmelCase_     (       self   :      int    , __lowerCamelCase   :      int    , __lowerCamelCase   :      Union[str, Any]="cpu"    , __lowerCamelCase   :      Dict=torch.floataa    , __lowerCamelCase   :      Union[str, Any]=0       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= torch.Generator(device=__lowerCamelCase       ).manual_seed(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= np.random.RandomState(__lowerCamelCase       ).standard_normal((1, 8, 1_2_8, 1_6)       )
      _SCREAMING_SNAKE_CASE			= torch.from_numpy(__lowerCamelCase       ).to(device=__lowerCamelCase    , dtype=__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= {
          "prompt": "A hammer hitting a wooden surface",
          "latents": latents,
          "generator": generator,
          "num_inference_steps": 3,
          "guidance_scale": 2.5,
      }
      return inputs



   def 				lowerCAmelCase_     (       self   :      Tuple       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline.from_pretrained("cvssp/audioldm"       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= self.get_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= 2_5
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       ).audios[0]

      assert audio.ndim == 1
      assert len(__lowerCamelCase       ) == 8_1_9_2_0

      _SCREAMING_SNAKE_CASE			= audio[7_7_2_3_0:7_7_2_4_0]
      _SCREAMING_SNAKE_CASE			= np.array(
          [-0.4_8_8_4, -0.4_6_0_7, 0.0_0_2_3, 0.5_0_0_7, 0.5_8_9_6, 0.5_1_5_1, 0.3_8_1_3, -0.0_2_0_8, -0.3_6_8_7, -0.4_3_1_5]       )
      _SCREAMING_SNAKE_CASE			= np.abs(expected_slice - audio_slice       ).max()
      assert max_diff < 1e-2



   def 				lowerCAmelCase_     (       self   :      Optional[Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= AudioLDMPipeline.from_pretrained("cvssp/audioldm"       )
      _SCREAMING_SNAKE_CASE			= LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe.to(__lowerCamelCase       )
      audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= self.get_inputs(__lowerCamelCase       )
      _SCREAMING_SNAKE_CASE			= audioldm_pipe(**__lowerCamelCase       ).audios[0]

      assert audio.ndim == 1
      assert len(__lowerCamelCase       ) == 8_1_9_2_0

      _SCREAMING_SNAKE_CASE			= audio[2_7_7_8_0:2_7_7_9_0]
      _SCREAMING_SNAKE_CASE			= np.array([-0.2_1_3_1, -0.0_8_7_3, -0.0_1_2_4, -0.0_1_8_9, 0.0_5_6_9, 0.1_3_7_3, 0.1_8_8_3, 0.2_8_8_6, 0.3_2_9_7, 0.2_2_1_2]       )
      _SCREAMING_SNAKE_CASE			= np.abs(expected_slice - audio_slice       ).max()
      assert max_diff < 3e-2





import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


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

						parser.add_argument(
						    '--checkpoint_path', default=None, type=str, required=True, help='Path to the checkpoint to convert.'
						)
						# !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
						parser.add_argument(
						    '--original_config_file',
						    default=None,
						    type=str,
						    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(
						    '--scheduler_type',
						    default='pndm',
						    type=str,
						    help='Type of scheduler to use. Should be one of [\'pndm\', \'lms\', \'ddim\', \'euler\', \'euler-ancestral\', \'dpm\']',
						)
						parser.add_argument(
						    '--pipeline_type',
						    default=None,
						    type=str,
						    help=(
						        'The pipeline type. One of \'FrozenOpenCLIPEmbedder\', \'FrozenCLIPEmbedder\', \'PaintByExample\''
						        '. If `None` pipeline will be automatically inferred.'
						    ),
						)
						parser.add_argument(
						    '--image_size',
						    default=None,
						    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(
						    '--prediction_type',
						    default=None,
						    type=str,
						    help=(
						        'The prediction type that the model was trained on. Use \'epsilon\' for Stable Diffusion v1.X and Stable'
						        ' Diffusion v2 Base. Use \'v_prediction\' 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.)')
						parser.add_argument(
						    '--stable_unclip',
						    type=str,
						    default=None,
						    required=False,
						    help='Set if this is a stable unCLIP model. One of \'txt2img\' or \'img2img\'.',
						)
						parser.add_argument(
						    '--stable_unclip_prior',
						    type=str,
						    default=None,
						    required=False,
						    help='Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.',
						)
						parser.add_argument(
						    '--clip_stats_path',
						    type=str,
						    help='Path to the clip stats file. Only required if the stable unclip model\'s config specifies `model.params.noise_aug_config.params.clip_stats_path`.',
						    required=False,
						)
						parser.add_argument(
						    '--controlnet', action='store_true', default=None, help='Set flag if this is a controlnet checkpoint.'
						)
						parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
						parser.add_argument(
						    '--vae_path',
						    type=str,
						    default=None,
						    required=False,
						    help='Set to a path, hub id to an already converted vae to not convert it again.',
						)
						__UpperCamelCase     		=							parser.parse_args()

						__UpperCamelCase     		=							download_from_original_stable_diffusion_ckpt(
						    checkpoint_path=args.checkpoint_path,
						    original_config_file=args.original_config_file,
						    image_size=args.image_size,
						    prediction_type=args.prediction_type,
						    model_type=args.pipeline_type,
						    extract_ema=args.extract_ema,
						    scheduler_type=args.scheduler_type,
						    num_in_channels=args.num_in_channels,
						    upcast_attention=args.upcast_attention,
						    from_safetensors=args.from_safetensors,
						    device=args.device,
						    stable_unclip=args.stable_unclip,
						    stable_unclip_prior=args.stable_unclip_prior,
						    clip_stats_path=args.clip_stats_path,
						    controlnet=args.controlnet,
						    vae_path=args.vae_path,
						)

						if args.half:
												pipe.to(torch_dtype=torch.floataa)

						if args.controlnet:
												# only save the controlnet model
												pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
						else:
												pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)


