


import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

from transformers import is_torch_available
from transformers.testing_utils import is_flaky, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


snake_case__			 =						1e-4

if is_torch_available():
  import torch

  from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
  from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder


@require_torch
class 					lowerCAmelCase_ :




      def __init__(	self						:   Tuple   ,     __A						:   Optional[int]   ,     __A						:   List[str]=16   ,     __A						:   str=13   ,     __A						:   Tuple=7   ,     __A						:   Any=14   ,     __A						:   List[str]=10   ,     __A						:   Any=19   ,     __A						:   List[Any]=5   ,     __A						:   Tuple=4   ,     __A						:   Tuple=True   ,     __A						:   Union[str, Any]=16   ,     __A						:   Optional[Any]=2   ,     __A						:   Union[str, Any]=4   ,     __A						:   Dict=4   ,     __A						:   Any="gelu"   ,     __A						:   List[Any]=0.1   ,     __A						:   Union[str, Any]=0.1   ,     __A						:   Optional[int]=[1, 2, 3, 4, 5]   ,     __A						:   int=25   ,     __A						:   Optional[int]=5   ,     )      ->List[Any]:






         """simple docstring"""






         a__ :int						  =    d_model
         a__ :Optional[Any]						  =    parent
         a__ :List[str]						  =    batch_size
         a__ :int						  =    prediction_length
         a__ :Optional[int]						  =    context_length
         a__ :Dict						  =    cardinality
         a__ :int						  =    num_time_features
         a__ :Union[str, Any]						  =    lags_sequence
         a__ :Any						  =    embedding_dimension
         a__ :Optional[int]						  =    is_training
         a__ :List[Any]						  =    hidden_size
         a__ :Dict						  =    num_hidden_layers
         a__ :Union[str, Any]						  =    num_attention_heads
         a__ :Optional[int]						  =    intermediate_size
         a__ :List[Any]						  =    hidden_act
         a__ :Optional[int]						  =    hidden_dropout_prob
         a__ :Optional[Any]						  =    attention_probs_dropout_prob

         a__ :Optional[Any]						  =    context_length
         a__ :Union[str, Any]						  =    prediction_length + label_length
         a__ :List[str]						  =    label_length

         a__ :List[str]						  =    moving_average
         a__ :Dict						  =    autocorrelation_factor




      def  _snake_case					(	self						:   Union[str, Any]						)      ->Optional[int]:






         """simple docstring"""






         return AutoformerConfig(
             d_model=self.d_model   ,     encoder_layers=self.num_hidden_layers   ,     decoder_layers=self.num_hidden_layers   ,     encoder_attention_heads=self.num_attention_heads   ,     decoder_attention_heads=self.num_attention_heads   ,     encoder_ffn_dim=self.intermediate_size   ,     decoder_ffn_dim=self.intermediate_size   ,     dropout=self.hidden_dropout_prob   ,     attention_dropout=self.attention_probs_dropout_prob   ,     prediction_length=self.prediction_length   ,     context_length=self.context_length   ,     label_length=self.label_length   ,     lags_sequence=self.lags_sequence   ,     num_time_features=self.num_time_features   ,     num_static_categorical_features=1   ,     cardinality=[self.cardinality]   ,     embedding_dimension=[self.embedding_dimension]   ,     moving_average=self.moving_average   ,     )




      def  _snake_case					(	self						:   List[Any]   ,     __A						:   Any						)      ->Union[str, Any]:






         """simple docstring"""






         a__ :List[str]						  =    config.context_length + max(config.lags_sequence						)

         a__ :List[Any]						  =    ids_tensor([self.batch_size, 1]   ,     config.cardinality[0]						)
         a__ :str						  =    floats_tensor([self.batch_size, _past_length, config.num_time_features]						)
         a__ :Optional[int]						  =    floats_tensor([self.batch_size, _past_length]						)
         a__ :List[Any]						  =    floats_tensor([self.batch_size, _past_length]						) > 0.5

         # decoder inputs
         a__ :List[str]						  =    floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]						)
         a__ :Any						  =    floats_tensor([self.batch_size, config.prediction_length]						)

         a__ :Union[str, Any]						  =    {
             "past_values": past_values,
             "static_categorical_features": static_categorical_features,
             "past_time_features": past_time_features,
             "past_observed_mask": past_observed_mask,
             "future_time_features": future_time_features,
             "future_values": future_values,
         }
         return inputs_dict




      def  _snake_case					(	self						:   Any						)      ->List[str]:






         """simple docstring"""






         a__ :Optional[int]						  =    self.get_config()
         a__ :Dict						  =    self.prepare_autoformer_inputs_dict(__A						)
         return config, inputs_dict




      def  _snake_case					(	self						:   List[Any]						)      ->Any:






         """simple docstring"""






         a__ :Optional[int]						  =    self.prepare_config_and_inputs()
         return config, inputs_dict






      def  _snake_case					(	self						:   List[str]   ,     __A						:   Optional[Any]   ,     __A						:   Any						)      ->List[Any]:






         """simple docstring"""






         a__ :Any						  =    AutoformerModel(config=__A						).to(__A						).eval()
         a__ :Tuple						  =    model(**__A						)

         a__ :int						  =    outputs.encoder_last_hidden_state
         a__ :Tuple						  =    outputs.last_hidden_state

         with tempfile.TemporaryDirectory() as tmpdirname:
            a__ :str						  =    model.get_encoder()
            encoder.save_pretrained(__A						)
            a__ :List[Any]						  =    AutoformerEncoder.from_pretrained(__A						).to(__A						)

         a__ :Optional[int]						  =    model.create_network_inputs(**__A						)
         a__ :Tuple						  =    model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]						)

         a__ :Optional[Any]						  =    torch.cat(
             (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])   ,     dim=-1   ,     )
         a__ :Optional[Any]						  =    encoder(inputs_embeds=__A						)[0]
         self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1E-3						)

         a__ :Any						  =    (
             torch.mean(transformer_inputs[:, : config.context_length, ...]   ,     dim=1						)
             .unsqueeze(1						)
             .repeat(1   ,     config.prediction_length   ,     1						)
         )
         a__ :int						  =    torch.zeros(
             [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]]   ,     device=enc_input.device   ,     )

         a__ :Optional[Any]						  =    torch.cat(
             (
                 torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)   ,     dim=1						),
                 feature[:, config.context_length - config.label_length :, ...],
             )   ,     dim=-1   ,     )
         a__ :int						  =    torch.cat(
             (
                 torch.cat((trend_input[:, -config.label_length :, ...], mean)   ,     dim=1						),
                 feature[:, config.context_length - config.label_length :, ...],
             )   ,     dim=-1   ,     )

         with tempfile.TemporaryDirectory() as tmpdirname:
            a__ :List[Any]						  =    model.get_decoder()
            decoder.save_pretrained(__A						)
            a__ :List[str]						  =    AutoformerDecoder.from_pretrained(__A						).to(__A						)

         a__ :Optional[int]						  =    decoder(
             trend=__A   ,     inputs_embeds=__A   ,     encoder_hidden_states=__A   ,     )[0]

         self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1E-3						)






@require_torch
class 					lowerCAmelCase_ (       lowercase__			,lowercase__			,unittest.TestCase):
      lowerCamelCase_         =     (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
      lowerCamelCase_         =     (AutoformerForPrediction,) if is_torch_available() else ()
      lowerCamelCase_         =     {'feature-extraction': AutoformerModel} if is_torch_available() else {}
      lowerCamelCase_         =     False
      lowerCamelCase_         =     False
      lowerCamelCase_         =     False
      lowerCamelCase_         =     False
      lowerCamelCase_         =     False
      lowerCamelCase_         =     False




      def  _snake_case					(	self						:   Optional[Any]						)      ->Tuple:






         """simple docstring"""






         a__ :List[str]						  =    AutoformerModelTester(self						)
         a__ :Tuple						  =    ConfigTester(self   ,     config_class=__A   ,     has_text_modality=__A						)




      def  _snake_case					(	self						:   Tuple						)      ->str:






         """simple docstring"""






         self.config_tester.run_common_tests()




      def  _snake_case					(	self						:   Any						)      ->Tuple:






         """simple docstring"""






         a__ :Tuple						  =    self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
            a__ :str						  =    model_class(__A						)

            with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_pretrained(__A						)
               a__ :Dict						  =    model_class.from_pretrained(__A   ,     output_loading_info=__A						)
            self.assertEqual(info["missing_keys"]   ,     []						)




      def  _snake_case					(	self						:   str						)      ->Dict:






         """simple docstring"""






         a__ :Any						  =    self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.check_encoder_decoder_model_standalone(*__A						)




      @unittest.skip(reason="Model has no tokens embeddings"						)
      def  _snake_case					(	self						:   Optional[Any]						)      ->Union[str, Any]:






         """simple docstring"""






         pass




      def  _snake_case					(	self						:   Dict						)      ->int:






         """simple docstring"""






         a__ :Optional[int]						  =    inspect.signature(getattr(__A   ,     "forward"						)						)
         # The main input is the name of the argument after `self`
         a__ :str						  =    list(model_signature.parameters.keys()						)[1]
         self.assertEqual(AutoformerModel.main_input_name   ,     __A						)




      def  _snake_case					(	self						:   Optional[Any]						)      ->Dict:






         """simple docstring"""






         a__ :Tuple						  =    self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            a__ :Optional[Any]						  =    model_class(__A						)
            a__ :List[Any]						  =    inspect.signature(model.forward						)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            a__ :Tuple						  =    [*signature.parameters.keys()]

            a__ :Optional[Any]						  =    [
                "past_values",
                "past_time_features",
                "past_observed_mask",
                "static_categorical_features",
                "static_real_features",
                "future_values",
                "future_time_features",
            ]

            if model.__class__.__name__ in ["AutoformerForPrediction"]:
               expected_arg_names.append("future_observed_mask"						)

            expected_arg_names.extend(
                [
                    "decoder_attention_mask",
                    "head_mask",
                    "decoder_head_mask",
                    "cross_attn_head_mask",
                    "encoder_outputs",
                    "past_key_values",
                    "output_hidden_states",
                    "output_attentions",
                    "use_cache",
                    "return_dict",
                ]						)

            self.assertListEqual(arg_names[: len(__A						)]   ,     __A						)




      def  _snake_case					(	self						:   Any						)      ->List[Any]:






         """simple docstring"""






         a__ :List[str]						  =    self.model_tester.prepare_config_and_inputs_for_common()
         a__ :Union[str, Any]						  =    True

         a__ :List[Any]						  =    getattr(self.model_tester   ,     "seq_length"   ,     __A						)
         a__ :Union[str, Any]						  =    getattr(self.model_tester   ,     "decoder_seq_length"   ,     __A						)
         a__ :Optional[Any]						  =    getattr(self.model_tester   ,     "encoder_seq_length"   ,     __A						)
         a__ :List[str]						  =    getattr(self.model_tester   ,     "d_model"   ,     __A						)
         a__ :str						  =    getattr(self.model_tester   ,     "num_attention_heads"   ,     __A						)
         a__ :int						  =    d_model // num_attention_heads

         for model_class in self.all_model_classes:
            a__ :Dict						  =    True
            a__ :Tuple						  =    False
            a__ :int						  =    True
            a__ :Optional[int]						  =    model_class(__A						)
            model.to(__A						)
            model.eval()
            with torch.no_grad():
               a__ :Optional[Any]						  =    model(**self._prepare_for_class(__A   ,     __A						)						)
            a__ :Dict						  =    outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
            self.assertEqual(len(__A						)   ,     self.model_tester.num_hidden_layers						)

            # check that output_attentions also work using config
            del inputs_dict["output_attentions"]
            a__ :Any						  =    True
            a__ :Union[str, Any]						  =    model_class(__A						)
            model.to(__A						)
            model.eval()
            with torch.no_grad():
               a__ :str						  =    model(**self._prepare_for_class(__A   ,     __A						)						)
            a__ :Optional[int]						  =    outputs.encoder_attentions
            self.assertEqual(len(__A						)   ,     self.model_tester.num_hidden_layers						)

            self.assertListEqual(
                list(attentions[0].shape[-3:]						)   ,     [self.model_tester.num_attention_heads, encoder_seq_length, dim]   ,     )
            a__ :Tuple						  =    len(__A						)

            a__ :int						  =    7

            if "last_hidden_state" in outputs:
               correct_outlen += 1

            if "trend" in outputs:
               correct_outlen += 1

            if "past_key_values" in outputs:
               correct_outlen += 1  # past_key_values have been returned

            if "loss" in outputs:
               correct_outlen += 1

            if "params" in outputs:
               correct_outlen += 1

            self.assertEqual(__A   ,     __A						)

            # decoder attentions
            a__ :Optional[Any]						  =    outputs.decoder_attentions
            self.assertIsInstance(__A   ,     (list, tuple)						)
            self.assertEqual(len(__A						)   ,     self.model_tester.num_hidden_layers						)
            self.assertListEqual(
                list(decoder_attentions[0].shape[-3:]						)   ,     [self.model_tester.num_attention_heads, decoder_seq_length, dim]   ,     )

            # cross attentions
            a__ :int						  =    outputs.cross_attentions
            self.assertIsInstance(__A   ,     (list, tuple)						)
            self.assertEqual(len(__A						)   ,     self.model_tester.num_hidden_layers						)
            self.assertListEqual(
                list(cross_attentions[0].shape[-3:]						)   ,     [self.model_tester.num_attention_heads, decoder_seq_length, dim]   ,     )

         # Check attention is always last and order is fine
         a__ :List[str]						  =    True
         a__ :int						  =    True
         a__ :int						  =    model_class(__A						)
         model.to(__A						)
         model.eval()
         with torch.no_grad():
            a__ :List[str]						  =    model(**self._prepare_for_class(__A   ,     __A						)						)

         self.assertEqual(out_len + 2   ,     len(__A						)						)

         a__ :Optional[int]						  =    outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

         self.assertEqual(len(__A						)   ,     self.model_tester.num_hidden_layers						)
         self.assertListEqual(
             list(self_attentions[0].shape[-3:]						)   ,     [self.model_tester.num_attention_heads, encoder_seq_length, dim]   ,     )






      @is_flaky()
      def  _snake_case					(	self						:   int						)      ->Optional[int]:






         """simple docstring"""






         super().test_retain_grad_hidden_states_attentions()






def 				lowerCamelCase__	(      a	:   Any="train-batch.pt"				)    ->	Tuple:





   """simple docstring"""
   a__ :Union[str, Any]						  =    hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch"	,  filename=_lowerCamelCase	,  repo_type="dataset"				)
   a__ :Union[str, Any]						  =    torch.load(_lowerCamelCase	,  map_location=_lowerCamelCase				)
   return batch







@require_torch
@slow
class 					lowerCAmelCase_ (       unittest.TestCase):




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






         """simple docstring"""






         a__ :Union[str, Any]						  =    AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly"						).to(__A						)
         a__ :Optional[int]						  =    prepare_batch()

         with torch.no_grad():
            a__ :Optional[int]						  =    model(
                past_values=batch["past_values"]   ,     past_time_features=batch["past_time_features"]   ,     past_observed_mask=batch["past_observed_mask"]   ,     static_categorical_features=batch["static_categorical_features"]   ,     future_values=batch["future_values"]   ,     future_time_features=batch["future_time_features"]   ,     )[0]

         a__ :int						  =    torch.Size(
             (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)						)
         self.assertEqual(output.shape   ,     __A						)

         a__ :Optional[Any]						  =    torch.tensor(
             [[0.3_593, -1.3_398, 0.6_330], [0.2_279, 1.5_396, -0.1_792], [0.0_450, 1.3_225, -0.2_335]]   ,     device=__A						)
         self.assertTrue(torch.allclose(output[0, :3, :3]   ,     __A   ,     atol=__A						)						)




      def  _snake_case					(	self						:   int						)      ->Any:






         """simple docstring"""






         a__ :Optional[int]						  =    AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"						).to(__A						)
         a__ :Optional[int]						  =    prepare_batch("val-batch.pt"						)
         with torch.no_grad():
            a__ :Optional[int]						  =    model(
                past_values=batch["past_values"]   ,     past_time_features=batch["past_time_features"]   ,     past_observed_mask=batch["past_observed_mask"]   ,     static_categorical_features=batch["static_categorical_features"]   ,     ).encoder_last_hidden_state
         a__ :Union[str, Any]						  =    torch.Size((64, model.config.context_length, model.config.d_model)						)
         self.assertEqual(output.shape   ,     __A						)

         a__ :Optional[int]						  =    torch.tensor(
             [[-0.0_734, -0.9_036, 0.8_358], [4.7_186, 2.4_113, 1.9_581], [1.7_953, 2.3_558, 1.2_970]]   ,     device=__A						)
         self.assertTrue(torch.allclose(output[0, :3, :3]   ,     __A   ,     atol=__A						)						)






      def  _snake_case					(	self						:   Any						)      ->Tuple:






         """simple docstring"""






         a__ :Tuple						  =    AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"						).to(__A						)
         a__ :List[str]						  =    prepare_batch("val-batch.pt"						)
         with torch.no_grad():
            a__ :Union[str, Any]						  =    model.generate(
                static_categorical_features=batch["static_categorical_features"]   ,     past_time_features=batch["past_time_features"]   ,     past_values=batch["past_values"]   ,     future_time_features=batch["future_time_features"]   ,     past_observed_mask=batch["past_observed_mask"]   ,     )
         a__ :Optional[Any]						  =    torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)						)
         self.assertEqual(outputs.sequences.shape   ,     __A						)

         a__ :Optional[Any]						  =    torch.tensor([3_130.6_763, 4_056.5_293, 7_053.0_786]   ,     device=__A						)
         a__ :List[str]						  =    outputs.sequences.mean(dim=1						)
         self.assertTrue(torch.allclose(mean_prediction[0, -3:]   ,     __A   ,     rtol=1E-1						)						)






import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
		snake_case__			 =						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.''',
		)
		snake_case__			 =						parser.parse_args()

		snake_case__			 =						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)


