
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 				_lowerCamelCase (      _A	,							unittest.TestCase ):





     """simple docstring"""







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






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




           """simple docstring"""







           torch.manual_seed(0  )
           UpperCamelCase__    :   List[str]				=				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						,						)
           UpperCamelCase__    :   List[str]				=				DDIMScheduler(
               beta_start=0.00085						,						beta_end=0.012						,						beta_schedule='''scaled_linear'''						,						clip_sample=__lowerCamelCase						,						set_alpha_to_one=__lowerCamelCase						,						)
           torch.manual_seed(0  )
           UpperCamelCase__    :   Optional[int]				=				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  )
           UpperCamelCase__    :   List[str]				=				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						,						)
           UpperCamelCase__    :   Any				=				ClapTextModelWithProjection(__lowerCamelCase  )
           UpperCamelCase__    :   List[str]				=				RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''						,						model_max_length=7_7  )

           UpperCamelCase__    :   List[str]				=				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						,						)

           UpperCamelCase__    :   Optional[int]				=				SpeechTaHifiGan(__lowerCamelCase  )

           UpperCamelCase__    :   Union[str, Any]				=				{
               "unet": unet,
               "scheduler": scheduler,
               "vae": vae,
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "vocoder": vocoder,
           }
           return components






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=0  )  ->   int:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__    :   int				=				"cpu"  # ensure determinism for the device-dependent torch.Generator

           UpperCamelCase__    :   Dict				=				self.get_dummy_components()
           UpperCamelCase__    :   List[str]				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   Tuple				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   Optional[int]				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   int				=				audioldm_pipe(**__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				output.audios[0]

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

           UpperCamelCase__    :   int				=				audio[:1_0]
           UpperCamelCase__    :   Dict				=				np.array(
               [-0.0050, 0.0050, -0.0060, 0.0033, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0033]  )

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






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




           """simple docstring"""







           UpperCamelCase__    :   List[str]				=				self.get_dummy_components()
           UpperCamelCase__    :   List[Any]				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   Tuple				=				audioldm_pipe.to(__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   str				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				3 * [inputs["prompt"]]

           # forward
           UpperCamelCase__    :   List[str]				=				audioldm_pipe(**__lowerCamelCase  )
           UpperCamelCase__    :   Dict				=				output.audios[0]

           UpperCamelCase__    :   Tuple				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   List[Any]				=				3 * [inputs.pop('''prompt'''  )]

           UpperCamelCase__    :   Optional[int]				=				audioldm_pipe.tokenizer(
               __lowerCamelCase						,						padding='''max_length'''						,						max_length=audioldm_pipe.tokenizer.model_max_length						,						truncation=__lowerCamelCase						,						return_tensors='''pt'''						,						)
           UpperCamelCase__    :   str				=				text_inputs["input_ids"].to(__lowerCamelCase  )

           UpperCamelCase__    :   str				=				audioldm_pipe.text_encoder(
               __lowerCamelCase						,						)
           UpperCamelCase__    :   Tuple				=				prompt_embeds.text_embeds
           # additional L_2 normalization over each hidden-state
           UpperCamelCase__    :   Optional[int]				=				F.normalize(__lowerCamelCase						,						dim=-1  )

           UpperCamelCase__    :   Union[str, Any]				=				prompt_embeds

           # forward
           UpperCamelCase__    :   List[Any]				=				audioldm_pipe(**__lowerCamelCase  )
           UpperCamelCase__    :   str				=				output.audios[0]

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






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




           """simple docstring"""







           UpperCamelCase__    :   Optional[Any]				=				self.get_dummy_components()
           UpperCamelCase__    :   List[str]				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   List[Any]				=				audioldm_pipe.to(__lowerCamelCase  )
           UpperCamelCase__    :   Optional[int]				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   Optional[Any]				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   int				=				3 * ["this is a negative prompt"]
           UpperCamelCase__    :   Dict				=				negative_prompt
           UpperCamelCase__    :   str				=				3 * [inputs["prompt"]]

           # forward
           UpperCamelCase__    :   Tuple				=				audioldm_pipe(**__lowerCamelCase  )
           UpperCamelCase__    :   str				=				output.audios[0]

           UpperCamelCase__    :   List[str]				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   Dict				=				3 * [inputs.pop('''prompt'''  )]

           UpperCamelCase__    :   Optional[int]				=				[]
           for p in [prompt, negative_prompt]:
                 UpperCamelCase__    :   List[Any]				=				audioldm_pipe.tokenizer(
                     __lowerCamelCase						,						padding='''max_length'''						,						max_length=audioldm_pipe.tokenizer.model_max_length						,						truncation=__lowerCamelCase						,						return_tensors='''pt'''						,						)
                 UpperCamelCase__    :   int				=				text_inputs["input_ids"].to(__lowerCamelCase  )

                 UpperCamelCase__    :   List[Any]				=				audioldm_pipe.text_encoder(
                     __lowerCamelCase						,						)
                 UpperCamelCase__    :   List[Any]				=				text_embeds.text_embeds
                 # additional L_2 normalization over each hidden-state
                 UpperCamelCase__    :   Dict				=				F.normalize(__lowerCamelCase						,						dim=-1  )

                 embeds.append(__lowerCamelCase  )

           UpperCamelCase__    :   List[Any]				=				embeds

           # forward
           UpperCamelCase__    :   Union[str, Any]				=				audioldm_pipe(**__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				output.audios[0]

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






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




           """simple docstring"""







           UpperCamelCase__    :   Dict				=				"cpu"  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   List[Any]				=				self.get_dummy_components()
           UpperCamelCase__    :   Optional[int]				=				PNDMScheduler(skip_prk_steps=__lowerCamelCase  )
           UpperCamelCase__    :   Any				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   List[str]				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   int				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   Any				=				"egg cracking"
           UpperCamelCase__    :   List[Any]				=				audioldm_pipe(**__lowerCamelCase						,						negative_prompt=__lowerCamelCase  )
           UpperCamelCase__    :   Any				=				output.audios[0]

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

           UpperCamelCase__    :   Any				=				audio[:1_0]
           UpperCamelCase__    :   Tuple				=				np.array(
               [-0.0051, 0.0050, -0.0060, 0.0034, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0032]  )

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






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







           UpperCamelCase__    :   Tuple				=				"cpu"  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   Optional[Any]				=				self.get_dummy_components()
           UpperCamelCase__    :   List[Any]				=				PNDMScheduler(skip_prk_steps=__lowerCamelCase  )
           UpperCamelCase__    :   Tuple				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   Union[str, Any]				=				"A hammer hitting a wooden surface"

           # test num_waveforms_per_prompt=1 (default)
           UpperCamelCase__    :   Optional[int]				=				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
           UpperCamelCase__    :   str				=				2
           UpperCamelCase__    :   List[str]				=				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
           UpperCamelCase__    :   Union[str, Any]				=				2
           UpperCamelCase__    :   Optional[Any]				=				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
           UpperCamelCase__    :   Optional[int]				=				2
           UpperCamelCase__    :   Tuple				=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[int]:




           """simple docstring"""







           UpperCamelCase__    :   Union[str, Any]				=				"cpu"  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   int				=				self.get_dummy_components()
           UpperCamelCase__    :   List[Any]				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )
           UpperCamelCase__    :   Any				=				audioldm_pipe.vocoder.config.sampling_rate

           UpperCamelCase__    :   Tuple				=				self.get_dummy_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				audioldm_pipe(audio_length_in_s=0.016						,						**__lowerCamelCase  )
           UpperCamelCase__    :   str				=				output.audios[0]

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

           UpperCamelCase__    :   Dict				=				audioldm_pipe(audio_length_in_s=0.032						,						**__lowerCamelCase  )
           UpperCamelCase__    :   Dict				=				output.audios[0]

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






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__    :   Any				=				self.get_dummy_components()
           UpperCamelCase__    :   List[str]				=				AudioLDMPipeline(**__lowerCamelCase  )
           UpperCamelCase__    :   Any				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   Tuple				=				["hey"]

           UpperCamelCase__    :   Tuple				=				audioldm_pipe(__lowerCamelCase						,						num_inference_steps=1  )
           UpperCamelCase__    :   Optional[int]				=				output.audios.shape
           assert audio_shape == (1, 2_5_6)

           UpperCamelCase__    :   Tuple				=				audioldm_pipe.vocoder.config
           config.model_in_dim *= 2
           UpperCamelCase__    :   Tuple				=				SpeechTaHifiGan(__lowerCamelCase  ).to(__lowerCamelCase  )
           UpperCamelCase__    :   Any				=				audioldm_pipe(__lowerCamelCase						,						num_inference_steps=1  )
           UpperCamelCase__    :   Any				=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   int:




           """simple docstring"""







           self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowerCamelCase  )






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




           """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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







           self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowerCamelCase  )




@slow
class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   int:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE="cpu"						,						__SCREAMING_SNAKE_CASE=torch.floataa						,						__SCREAMING_SNAKE_CASE=0  )  ->   Any:




           """simple docstring"""







           UpperCamelCase__    :   List[str]				=				torch.Generator(device=__lowerCamelCase  ).manual_seed(__lowerCamelCase  )
           UpperCamelCase__    :   Optional[Any]				=				np.random.RandomState(__lowerCamelCase  ).standard_normal((1, 8, 1_2_8, 1_6)  )
           UpperCamelCase__    :   Optional[Any]				=				torch.from_numpy(__lowerCamelCase  ).to(device=__lowerCamelCase						,						dtype=__lowerCamelCase  )
           UpperCamelCase__    :   int				=				{
               "prompt": "A hammer hitting a wooden surface",
               "latents": latents,
               "generator": generator,
               "num_inference_steps": 3,
               "guidance_scale": 2.5,
           }
           return inputs






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




           """simple docstring"""







           UpperCamelCase__    :   int				=				AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''  )
           UpperCamelCase__    :   List[Any]				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   Optional[Any]				=				self.get_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   int				=				2_5
           UpperCamelCase__    :   Dict				=				audioldm_pipe(**__lowerCamelCase  ).audios[0]

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

           UpperCamelCase__    :   str				=				audio[7_7_2_3_0:7_7_2_4_0]
           UpperCamelCase__    :   Tuple				=				np.array(
               [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]  )
           UpperCamelCase__    :   Union[str, Any]				=				np.abs(expected_slice - audio_slice  ).max()
           assert max_diff < 1e-2






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




           """simple docstring"""







           UpperCamelCase__    :   List[str]				=				AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''  )
           UpperCamelCase__    :   str				=				LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config  )
           UpperCamelCase__    :   str				=				audioldm_pipe.to(__lowerCamelCase  )
           audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

           UpperCamelCase__    :   List[str]				=				self.get_inputs(__lowerCamelCase  )
           UpperCamelCase__    :   str				=				audioldm_pipe(**__lowerCamelCase  ).audios[0]

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

           UpperCamelCase__    :   Optional[int]				=				audio[2_7_7_8_0:2_7_7_9_0]
           UpperCamelCase__    :   Any				=				np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]  )
           UpperCamelCase__    :   Union[str, Any]				=				np.abs(expected_slice - audio_slice  ).max()
           assert max_diff < 3e-2




import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


lowerCamelCase      =False
lowerCamelCase      =True
lowerCamelCase      =False


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

    parser.add_argument(
        "--repo_path",
        default=None,
        type=str,
        required=True,
        help="The config json file corresponding to the architecture.",
    )

    parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")

    lowerCamelCase      =parser.parse_args()

    lowerCamelCase      ={
        "image_size": "sample_size",
        "num_res_blocks": "layers_per_block",
        "block_channels": "block_out_channels",
        "down_blocks": "down_block_types",
        "up_blocks": "up_block_types",
        "downscale_freq_shift": "freq_shift",
        "resnet_num_groups": "norm_num_groups",
        "resnet_act_fn": "act_fn",
        "resnet_eps": "norm_eps",
        "num_head_channels": "attention_head_dim",
    }

    lowerCamelCase      ={
        "time_steps": "time_proj",
        "mid": "mid_block",
        "downsample_blocks": "down_blocks",
        "upsample_blocks": "up_blocks",
    }

    lowerCamelCase      ="" if has_file(args.repo_path, "config.json") else "unet"

    with open(os.path.join(args.repo_path, subfolder, "config.json"), "r", encoding="utf-8") as reader:
        lowerCamelCase      =reader.read()
        lowerCamelCase      =json.loads(text)

    if do_only_config:
        for key in config_parameters_to_change.keys():
            config.pop(key, None)

    if has_file(args.repo_path, "config.json"):
        lowerCamelCase      =UNetaDModel(**config)
    else:
        lowerCamelCase      =UNetaDConditionModel if "ldm-text2im-large-256" in args.repo_path else UNetaDModel
        lowerCamelCase      =class_name(**config)

    if do_only_config:
        model.save_config(os.path.join(args.repo_path, subfolder))

    lowerCamelCase      =dict(model.config)

    if do_only_renaming:
        for key, value in config_parameters_to_change.items():
            if key in config:
                lowerCamelCase      =config[key]
                del config[key]

        lowerCamelCase      =[k.replace("UNetRes", "") for k in config["down_block_types"]]
        lowerCamelCase      =[k.replace("UNetRes", "") for k in config["up_block_types"]]

    if do_only_weights:
        lowerCamelCase      =torch.load(os.path.join(args.repo_path, subfolder, "diffusion_pytorch_model.bin"))

        lowerCamelCase      ={}
        for param_key, param_value in state_dict.items():
            if param_key.endswith(".op.bias") or param_key.endswith(".op.weight"):
                continue
            lowerCamelCase      =False
            for key, new_key in key_parameters_to_change.items():
                if not has_changed and param_key.split(".")[0] == key:
                    lowerCamelCase      =param_value
                    lowerCamelCase      =True
            if not has_changed:
                lowerCamelCase      =param_value

        model.load_state_dict(new_state_dict)
        model.save_pretrained(os.path.join(args.repo_path, subfolder))


