parson/audioEditing

1import torch2from diffusers import DDIMScheduler3from diffusers import AudioLDM2Pipeline4from transformers import RobertaTokenizer, RobertaTokenizerFast5from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput6from typing import Any, Dict, List, Optional, Tuple, Union7 8 9class PipelineWrapper(torch.nn.Module):10    def __init__(self, model_id, device, double_precision=False, *args, **kwargs) -> None:11        super().__init__(*args, **kwargs)12        self.model_id = model_id13        self.device = device14        self.double_precision = double_precision15 16    def get_sigma(self, timestep) -> float:17        sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.model.scheduler.alphas_cumprod - 1)18        return sqrt_recipm1_alphas_cumprod[timestep]19 20    def load_scheduler(self):21        pass22 23    def get_fn_STFT(self):24        pass25 26    def vae_encode(self, x: torch.Tensor):27        pass28 29    def vae_decode(self, x: torch.Tensor):30        pass31 32    def decode_to_mel(self, x: torch.Tensor):33        pass34 35    def encode_text(self, prompts: List[str]) -> Tuple:36        pass37 38    def get_variance(self, timestep, prev_timestep):39        pass40 41    def get_alpha_prod_t_prev(self, prev_timestep):42        pass43 44    def unet_forward(self,45                     sample: torch.FloatTensor,46                     timestep: Union[torch.Tensor, float, int],47                     encoder_hidden_states: torch.Tensor,48                     class_labels: Optional[torch.Tensor] = None,49                     timestep_cond: Optional[torch.Tensor] = None,50                     attention_mask: Optional[torch.Tensor] = None,51                     cross_attention_kwargs: Optional[Dict[str, Any]] = None,52                     added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,53                     down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,54                     mid_block_additional_residual: Optional[torch.Tensor] = None,55                     encoder_attention_mask: Optional[torch.Tensor] = None,56                     replace_h_space: Optional[torch.Tensor] = None,57                     replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None,58                     return_dict: bool = True,59                     zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple:60 61        # By default samples have to be AT least a multiple of the overall upsampling factor.62        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).63        # However, the upsampling interpolation output size can be forced to fit any upsampling size64        # on the fly if necessary.65        default_overall_up_factor = 2**self.model.unet.num_upsamplers66 67        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`68        forward_upsample_size = False69        upsample_size = None70 71        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):72            # logger.info("Forward upsample size to force interpolation output size.")73            forward_upsample_size = True74 75        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension76        # expects mask of shape:77        #   [batch, key_tokens]78        # adds singleton query_tokens dimension:79        #   [batch,                    1, key_tokens]80        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:81        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)82        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)83        if attention_mask is not None:84            # assume that mask is expressed as:85            #   (1 = keep,      0 = discard)86            # convert mask into a bias that can be added to attention scores:87            #       (keep = +0,     discard = -10000.0)88            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.089            attention_mask = attention_mask.unsqueeze(1)90 91        # convert encoder_attention_mask to a bias the same way we do for attention_mask92        if encoder_attention_mask is not None:93            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.094            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)95 96        # 0. center input if necessary97        if self.model.unet.config.center_input_sample:98            sample = 2 * sample - 1.099 100        # 1. time101        timesteps = timestep102        if not torch.is_tensor(timesteps):103            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can104            # This would be a good case for the `match` statement (Python 3.10+)105            is_mps = sample.device.type == "mps"106            if isinstance(timestep, float):107                dtype = torch.float32 if is_mps else torch.float64108            else:109                dtype = torch.int32 if is_mps else torch.int64110            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)111        elif len(timesteps.shape) == 0:112            timesteps = timesteps[None].to(sample.device)113 114        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML115        timesteps = timesteps.expand(sample.shape[0])116 117        t_emb = self.model.unet.time_proj(timesteps)118 119        # `Timesteps` does not contain any weights and will always return f32 tensors120        # but time_embedding might actually be running in fp16. so we need to cast here.121        # there might be better ways to encapsulate this.122        t_emb = t_emb.to(dtype=sample.dtype)123 124        emb = self.model.unet.time_embedding(t_emb, timestep_cond)125 126        if self.model.unet.class_embedding is not None:127            if class_labels is None:128                raise ValueError("class_labels should be provided when num_class_embeds > 0")129 130            if self.model.unet.config.class_embed_type == "timestep":131                class_labels = self.model.unet.time_proj(class_labels)132 133                # `Timesteps` does not contain any weights and will always return f32 tensors134                # there might be better ways to encapsulate this.135                class_labels = class_labels.to(dtype=sample.dtype)136 137            class_emb = self.model.unet.class_embedding(class_labels).to(dtype=sample.dtype)138 139            if self.model.unet.config.class_embeddings_concat:140                emb = torch.cat([emb, class_emb], dim=-1)141            else:142                emb = emb + class_emb143 144        if self.model.unet.config.addition_embed_type == "text":145            aug_emb = self.model.unet.add_embedding(encoder_hidden_states)146            emb = emb + aug_emb147        elif self.model.unet.config.addition_embed_type == "text_image":148            # Kadinsky 2.1 - style149            if "image_embeds" not in added_cond_kwargs:150                raise ValueError(151                    f"{self.model.unet.__class__} has the config param `addition_embed_type` set to 'text_image' "152                    f"which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"153                )154 155            image_embs = added_cond_kwargs.get("image_embeds")156            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)157 158            aug_emb = self.model.unet.add_embedding(text_embs, image_embs)159            emb = emb + aug_emb160 161        if self.model.unet.time_embed_act is not None:162            emb = self.model.unet.time_embed_act(emb)163 164        if self.model.unet.encoder_hid_proj is not None and self.model.unet.config.encoder_hid_dim_type == "text_proj":165            encoder_hidden_states = self.model.unet.encoder_hid_proj(encoder_hidden_states)166        elif self.model.unet.encoder_hid_proj is not None and \167                self.model.unet.config.encoder_hid_dim_type == "text_image_proj":168            # Kadinsky 2.1 - style169            if "image_embeds" not in added_cond_kwargs:170                raise ValueError(171                    f"{self.model.unet.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' "172                    f"which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"173                )174 175            image_embeds = added_cond_kwargs.get("image_embeds")176            encoder_hidden_states = self.model.unet.encoder_hid_proj(encoder_hidden_states, image_embeds)177 178        # 2. pre-process179        sample = self.model.unet.conv_in(sample)180 181        # 3. down182        down_block_res_samples = (sample,)183        for downsample_block in self.model.unet.down_blocks:184            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:185                sample, res_samples = downsample_block(186                    hidden_states=sample,187                    temb=emb,188                    encoder_hidden_states=encoder_hidden_states,189                    attention_mask=attention_mask,190                    cross_attention_kwargs=cross_attention_kwargs,191                    encoder_attention_mask=encoder_attention_mask,192                )193            else:194                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)195 196            down_block_res_samples += res_samples197 198        if down_block_additional_residuals is not None:199            new_down_block_res_samples = ()200 201            for down_block_res_sample, down_block_additional_residual in zip(202                down_block_res_samples, down_block_additional_residuals203            ):204                down_block_res_sample = down_block_res_sample + down_block_additional_residual205                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)206 207            down_block_res_samples = new_down_block_res_samples208 209        # 4. mid210        if self.model.unet.mid_block is not None:211            sample = self.model.unet.mid_block(212                sample,213                emb,214                encoder_hidden_states=encoder_hidden_states,215                attention_mask=attention_mask,216                cross_attention_kwargs=cross_attention_kwargs,217                encoder_attention_mask=encoder_attention_mask,218            )219 220        # print(sample.shape)221 222        if replace_h_space is None:223            h_space = sample.clone()224        else:225            h_space = replace_h_space226            sample = replace_h_space.clone()227 228        if mid_block_additional_residual is not None:229            sample = sample + mid_block_additional_residual230 231        extracted_res_conns = {}232        # 5. up233        for i, upsample_block in enumerate(self.model.unet.up_blocks):234            is_final_block = i == len(self.model.unet.up_blocks) - 1235 236            res_samples = down_block_res_samples[-len(upsample_block.resnets):]237            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]238            if replace_skip_conns is not None and replace_skip_conns.get(i):239                res_samples = replace_skip_conns.get(i)240 241            if zero_out_resconns is not None:242                if (type(zero_out_resconns) is int and i >= (zero_out_resconns - 1)) or \243                        type(zero_out_resconns) is list and i in zero_out_resconns:244                    res_samples = [torch.zeros_like(x) for x in res_samples]245                # down_block_res_samples = [torch.zeros_like(x) for x in down_block_res_samples]246 247            extracted_res_conns[i] = res_samples248 249            # if we have not reached the final block and need to forward the250            # upsample size, we do it here251            if not is_final_block and forward_upsample_size:252                upsample_size = down_block_res_samples[-1].shape[2:]253 254            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:255                sample = upsample_block(256                    hidden_states=sample,257                    temb=emb,258                    res_hidden_states_tuple=res_samples,259                    encoder_hidden_states=encoder_hidden_states,260                    cross_attention_kwargs=cross_attention_kwargs,261                    upsample_size=upsample_size,262                    attention_mask=attention_mask,263                    encoder_attention_mask=encoder_attention_mask,264                )265            else:266                sample = upsample_block(267                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size268                )269 270        # 6. post-process271        if self.model.unet.conv_norm_out:272            sample = self.model.unet.conv_norm_out(sample)273            sample = self.model.unet.conv_act(sample)274        sample = self.model.unet.conv_out(sample)275 276        if not return_dict:277            return (sample,)278 279        return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns280 281 282class AudioLDM2Wrapper(PipelineWrapper):283    def __init__(self, *args, **kwargs) -> None:284        super().__init__(*args, **kwargs)285        if self.double_precision:286            self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, torch_dtype=torch.float64).to(self.device)287        else:288            try:289                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=True).to(self.device)290            except FileNotFoundError:291                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=False).to(self.device)292 293    def load_scheduler(self):294        # self.model.scheduler = DDIMScheduler.from_config(self.model_id, subfolder="scheduler")295        self.model.scheduler = DDIMScheduler.from_pretrained(self.model_id, subfolder="scheduler")296 297    def get_fn_STFT(self):298        from audioldm.audio import TacotronSTFT299        return TacotronSTFT(300            filter_length=1024,301            hop_length=160,302            win_length=1024,303            n_mel_channels=64,304            sampling_rate=16000,305            mel_fmin=0,306            mel_fmax=8000,307        )308 309    def vae_encode(self, x):310        # self.model.vae.disable_tiling()311        if x.shape[2] % 4:312            x = torch.nn.functional.pad(x, (0, 0, 4 - (x.shape[2] % 4), 0))313        return (self.model.vae.encode(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float()314        # return (self.encode_no_tiling(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float()315 316    def vae_decode(self, x):317        return self.model.vae.decode(1 / self.model.vae.config.scaling_factor * x).sample318 319    def decode_to_mel(self, x):320        if self.double_precision:321            tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().double()).detach()322        tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().float()).detach()323        if len(tmp.shape) == 1:324            tmp = tmp.unsqueeze(0)325        return tmp326 327    def encode_text(self, prompts: List[str]):328        tokenizers = [self.model.tokenizer, self.model.tokenizer_2]329        text_encoders = [self.model.text_encoder, self.model.text_encoder_2]330        prompt_embeds_list = []331        attention_mask_list = []332 333        for tokenizer, text_encoder in zip(tokenizers, text_encoders):334            text_inputs = tokenizer(335                prompts,336                padding="max_length" if isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast)) else True,337                max_length=tokenizer.model_max_length,338                truncation=True,339                return_tensors="pt",340            )341 342            text_input_ids = text_inputs.input_ids343            attention_mask = text_inputs.attention_mask344            untruncated_ids = tokenizer(prompts, padding="longest", return_tensors="pt").input_ids345 346            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] \347                    and not torch.equal(text_input_ids, untruncated_ids):348                removed_text = tokenizer.batch_decode(349                    untruncated_ids[:, tokenizer.model_max_length - 1: -1])350                print(f"The following part of your input was truncated because {text_encoder.config.model_type} can "351                      f"only handle sequences up to {tokenizer.model_max_length} tokens: {removed_text}"352                      )353 354            text_input_ids = text_input_ids.to(self.device)355            attention_mask = attention_mask.to(self.device)356 357            with torch.no_grad():358                if text_encoder.config.model_type == "clap":359                    prompt_embeds = text_encoder.get_text_features(360                        text_input_ids,361                        attention_mask=attention_mask,362                    )363                    # append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size)364                    prompt_embeds = prompt_embeds[:, None, :]365                    # make sure that we attend to this single hidden-state366                    attention_mask = attention_mask.new_ones((len(prompts), 1))367                else:368                    prompt_embeds = text_encoder(369                        text_input_ids,370                        attention_mask=attention_mask,371                    )372                    prompt_embeds = prompt_embeds[0]373 374            prompt_embeds_list.append(prompt_embeds)375            attention_mask_list.append(attention_mask)376 377        # print(f'prompt[0].shape: {prompt_embeds_list[0].shape}')378        # print(f'prompt[1].shape: {prompt_embeds_list[1].shape}')379        # print(f'attn[0].shape: {attention_mask_list[0].shape}')380        # print(f'attn[1].shape: {attention_mask_list[1].shape}')381 382        projection_output = self.model.projection_model(383            hidden_states=prompt_embeds_list[0],384            hidden_states_1=prompt_embeds_list[1],385            attention_mask=attention_mask_list[0],386            attention_mask_1=attention_mask_list[1],387        )388        projected_prompt_embeds = projection_output.hidden_states389        projected_attention_mask = projection_output.attention_mask390 391        generated_prompt_embeds = self.model.generate_language_model(392            projected_prompt_embeds,393            attention_mask=projected_attention_mask,394            max_new_tokens=None,395        )396 397        prompt_embeds = prompt_embeds.to(dtype=self.model.text_encoder_2.dtype, device=self.device)398        attention_mask = (399            attention_mask.to(device=self.device)400            if attention_mask is not None401            else torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=self.device)402        )403        generated_prompt_embeds = generated_prompt_embeds.to(dtype=self.model.language_model.dtype, device=self.device)404 405        return generated_prompt_embeds, prompt_embeds, attention_mask406 407    def get_variance(self, timestep, prev_timestep):408        alpha_prod_t = self.model.scheduler.alphas_cumprod[timestep]409        alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)410        beta_prod_t = 1 - alpha_prod_t411        beta_prod_t_prev = 1 - alpha_prod_t_prev412        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)413        return variance414 415    def get_alpha_prod_t_prev(self, prev_timestep):416        return self.model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 \417            else self.model.scheduler.final_alpha_cumprod418 419    def unet_forward(self,420                     sample: torch.FloatTensor,421                     timestep: Union[torch.Tensor, float, int],422                     encoder_hidden_states: torch.Tensor,423                     timestep_cond: Optional[torch.Tensor] = None,424                     class_labels: Optional[torch.Tensor] = None,425                     attention_mask: Optional[torch.Tensor] = None,426                     encoder_attention_mask: Optional[torch.Tensor] = None,427                     return_dict: bool = True,428                     cross_attention_kwargs: Optional[Dict[str, Any]] = None,429                     mid_block_additional_residual: Optional[torch.Tensor] = None,430                     replace_h_space: Optional[torch.Tensor] = None,431                     replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None,432                     zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple:433 434        # Translation435        encoder_hidden_states_1 = class_labels436        class_labels = None437        encoder_attention_mask_1 = encoder_attention_mask438        encoder_attention_mask = None439 440        # return self.model.unet(sample, timestep,441        #                        encoder_hidden_states=generated_prompt_embeds,442        #                        encoder_hidden_states_1=encoder_hidden_states_1,443        #                        encoder_attention_mask_1=encoder_attention_mask_1,444        #                        ), None, None445 446        # By default samples have to be AT least a multiple of the overall upsampling factor.447        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).448        # However, the upsampling interpolation output size can be forced to fit any upsampling size449        # on the fly if necessary.450        default_overall_up_factor = 2 ** self.model.unet.num_upsamplers451 452        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`453        forward_upsample_size = False454        upsample_size = None455 456        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):457            # print("Forward upsample size to force interpolation output size.")458            forward_upsample_size = True459 460        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension461        # expects mask of shape:462        #   [batch, key_tokens]463        # adds singleton query_tokens dimension:464        #   [batch,                    1, key_tokens]465        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:466        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)467        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)468        if attention_mask is not None:469            # assume that mask is expressed as:470            #   (1 = keep,      0 = discard)471            # convert mask into a bias that can be added to attention scores:472            #       (keep = +0,     discard = -10000.0)473            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0474            attention_mask = attention_mask.unsqueeze(1)475 476        # convert encoder_attention_mask to a bias the same way we do for attention_mask477        if encoder_attention_mask is not None:478            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0479            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)480 481        if encoder_attention_mask_1 is not None:482            encoder_attention_mask_1 = (1 - encoder_attention_mask_1.to(sample.dtype)) * -10000.0483            encoder_attention_mask_1 = encoder_attention_mask_1.unsqueeze(1)484 485        # 1. time486        timesteps = timestep487        if not torch.is_tensor(timesteps):488            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can489            # This would be a good case for the `match` statement (Python 3.10+)490            is_mps = sample.device.type == "mps"491            if isinstance(timestep, float):492                dtype = torch.float32 if is_mps else torch.float64493            else:494                dtype = torch.int32 if is_mps else torch.int64495            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)496        elif len(timesteps.shape) == 0:497            timesteps = timesteps[None].to(sample.device)498 499        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML500        timesteps = timesteps.expand(sample.shape[0])501 502        t_emb = self.model.unet.time_proj(timesteps)503 504        # `Timesteps` does not contain any weights and will always return f32 tensors505        # but time_embedding might actually be running in fp16. so we need to cast here.506        # there might be better ways to encapsulate this.507        t_emb = t_emb.to(dtype=sample.dtype)508 509        emb = self.model.unet.time_embedding(t_emb, timestep_cond)510        aug_emb = None511 512        if self.model.unet.class_embedding is not None:513            if class_labels is None:514                raise ValueError("class_labels should be provided when num_class_embeds > 0")515 516            if self.model.unet.config.class_embed_type == "timestep":517                class_labels = self.model.unet.time_proj(class_labels)518 519                # `Timesteps` does not contain any weights and will always return f32 tensors520                # there might be better ways to encapsulate this.521                class_labels = class_labels.to(dtype=sample.dtype)522 523            class_emb = self.model.unet.class_embedding(class_labels).to(dtype=sample.dtype)524 525            if self.model.unet.config.class_embeddings_concat:526                emb = torch.cat([emb, class_emb], dim=-1)527            else:528                emb = emb + class_emb529 530        emb = emb + aug_emb if aug_emb is not None else emb531 532        if self.model.unet.time_embed_act is not None:533            emb = self.model.unet.time_embed_act(emb)534 535        # 2. pre-process536        sample = self.model.unet.conv_in(sample)537 538        # 3. down539        down_block_res_samples = (sample,)540        for downsample_block in self.model.unet.down_blocks:541            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:542                sample, res_samples = downsample_block(543                    hidden_states=sample,544                    temb=emb,545                    encoder_hidden_states=encoder_hidden_states,546                    attention_mask=attention_mask,547                    cross_attention_kwargs=cross_attention_kwargs,548                    encoder_attention_mask=encoder_attention_mask,549                    encoder_hidden_states_1=encoder_hidden_states_1,550                    encoder_attention_mask_1=encoder_attention_mask_1,551                )552            else:553                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)554 555            down_block_res_samples += res_samples556 557        # 4. mid558        if self.model.unet.mid_block is not None:559            sample = self.model.unet.mid_block(560                sample,561                emb,562                encoder_hidden_states=encoder_hidden_states,563                attention_mask=attention_mask,564                cross_attention_kwargs=cross_attention_kwargs,565                encoder_attention_mask=encoder_attention_mask,566                encoder_hidden_states_1=encoder_hidden_states_1,567                encoder_attention_mask_1=encoder_attention_mask_1,568            )569 570        if replace_h_space is None:571            h_space = sample.clone()572        else:573            h_space = replace_h_space574            sample = replace_h_space.clone()575 576        if mid_block_additional_residual is not None:577            sample = sample + mid_block_additional_residual578 579        extracted_res_conns = {}580        # 5. up581        for i, upsample_block in enumerate(self.model.unet.up_blocks):582            is_final_block = i == len(self.model.unet.up_blocks) - 1583 584            res_samples = down_block_res_samples[-len(upsample_block.resnets):]585            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]586            if replace_skip_conns is not None and replace_skip_conns.get(i):587                res_samples = replace_skip_conns.get(i)588 589            if zero_out_resconns is not None:590                if (type(zero_out_resconns) is int and i >= (zero_out_resconns - 1)) or \591                        type(zero_out_resconns) is list and i in zero_out_resconns:592                    res_samples = [torch.zeros_like(x) for x in res_samples]593                # down_block_res_samples = [torch.zeros_like(x) for x in down_block_res_samples]594 595            extracted_res_conns[i] = res_samples596 597            # if we have not reached the final block and need to forward the598            # upsample size, we do it here599            if not is_final_block and forward_upsample_size:600                upsample_size = down_block_res_samples[-1].shape[2:]601 602            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:603                sample = upsample_block(604                    hidden_states=sample,605                    temb=emb,606                    res_hidden_states_tuple=res_samples,607                    encoder_hidden_states=encoder_hidden_states,608                    cross_attention_kwargs=cross_attention_kwargs,609                    upsample_size=upsample_size,610                    attention_mask=attention_mask,611                    encoder_attention_mask=encoder_attention_mask,612                    encoder_hidden_states_1=encoder_hidden_states_1,613                    encoder_attention_mask_1=encoder_attention_mask_1,614                )615            else:616                sample = upsample_block(617                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size618                )619 620        # 6. post-process621        if self.model.unet.conv_norm_out:622            sample = self.model.unet.conv_norm_out(sample)623            sample = self.model.unet.conv_act(sample)624        sample = self.model.unet.conv_out(sample)625 626        if not return_dict:627            return (sample,)628 629        return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns630 631    def forward(self, *args, **kwargs):632        return self633 634 635def load_model(model_id, device, double_precision=False):636    ldm_stable = AudioLDM2Wrapper(model_id=model_id, device=device, double_precision=double_precision)637    ldm_stable.load_scheduler()638    torch.cuda.empty_cache()639    return ldm_stable640