Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from diffusers import DDIMScheduler | |
| from diffusers import AudioLDM2Pipeline | |
| from transformers import RobertaTokenizer, RobertaTokenizerFast | |
| from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput | |
| from typing import Any, Dict, List, Optional, Tuple, Union | |
| class PipelineWrapper(torch.nn.Module): | |
| def __init__(self, model_id, device, double_precision=False, *args, **kwargs) -> None: | |
| super().__init__(*args, **kwargs) | |
| self.model_id = model_id | |
| self.device = device | |
| self.double_precision = double_precision | |
| def get_sigma(self, timestep) -> float: | |
| sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.model.scheduler.alphas_cumprod - 1) | |
| return sqrt_recipm1_alphas_cumprod[timestep] | |
| def load_scheduler(self): | |
| pass | |
| def get_fn_STFT(self): | |
| pass | |
| def vae_encode(self, x: torch.Tensor): | |
| pass | |
| def vae_decode(self, x: torch.Tensor): | |
| pass | |
| def decode_to_mel(self, x: torch.Tensor): | |
| pass | |
| def encode_text(self, prompts: List[str]) -> Tuple: | |
| pass | |
| def get_variance(self, timestep, prev_timestep): | |
| pass | |
| def get_alpha_prod_t_prev(self, prev_timestep): | |
| pass | |
| def unet_forward(self, | |
| sample: torch.FloatTensor, | |
| timestep: Union[torch.Tensor, float, int], | |
| encoder_hidden_states: torch.Tensor, | |
| class_labels: Optional[torch.Tensor] = None, | |
| timestep_cond: Optional[torch.Tensor] = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
| added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, | |
| down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, | |
| mid_block_additional_residual: Optional[torch.Tensor] = None, | |
| encoder_attention_mask: Optional[torch.Tensor] = None, | |
| replace_h_space: Optional[torch.Tensor] = None, | |
| replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None, | |
| return_dict: bool = True, | |
| zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple: | |
| # By default samples have to be AT least a multiple of the overall upsampling factor. | |
| # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). | |
| # However, the upsampling interpolation output size can be forced to fit any upsampling size | |
| # on the fly if necessary. | |
| default_overall_up_factor = 2**self.model.unet.num_upsamplers | |
| # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` | |
| forward_upsample_size = False | |
| upsample_size = None | |
| if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): | |
| # logger.info("Forward upsample size to force interpolation output size.") | |
| forward_upsample_size = True | |
| # ensure attention_mask is a bias, and give it a singleton query_tokens dimension | |
| # expects mask of shape: | |
| # [batch, key_tokens] | |
| # adds singleton query_tokens dimension: | |
| # [batch, 1, key_tokens] | |
| # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: | |
| # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) | |
| # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) | |
| if attention_mask is not None: | |
| # assume that mask is expressed as: | |
| # (1 = keep, 0 = discard) | |
| # convert mask into a bias that can be added to attention scores: | |
| # (keep = +0, discard = -10000.0) | |
| attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 | |
| attention_mask = attention_mask.unsqueeze(1) | |
| # convert encoder_attention_mask to a bias the same way we do for attention_mask | |
| if encoder_attention_mask is not None: | |
| encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 | |
| encoder_attention_mask = encoder_attention_mask.unsqueeze(1) | |
| # 0. center input if necessary | |
| if self.model.unet.config.center_input_sample: | |
| sample = 2 * sample - 1.0 | |
| # 1. time | |
| timesteps = timestep | |
| if not torch.is_tensor(timesteps): | |
| # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can | |
| # This would be a good case for the `match` statement (Python 3.10+) | |
| is_mps = sample.device.type == "mps" | |
| if isinstance(timestep, float): | |
| dtype = torch.float32 if is_mps else torch.float64 | |
| else: | |
| dtype = torch.int32 if is_mps else torch.int64 | |
| timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) | |
| elif len(timesteps.shape) == 0: | |
| timesteps = timesteps[None].to(sample.device) | |
| # broadcast to batch dimension in a way that's compatible with ONNX/Core ML | |
| timesteps = timesteps.expand(sample.shape[0]) | |
| t_emb = self.model.unet.time_proj(timesteps) | |
| # `Timesteps` does not contain any weights and will always return f32 tensors | |
| # but time_embedding might actually be running in fp16. so we need to cast here. | |
| # there might be better ways to encapsulate this. | |
| t_emb = t_emb.to(dtype=sample.dtype) | |
| emb = self.model.unet.time_embedding(t_emb, timestep_cond) | |
| if self.model.unet.class_embedding is not None: | |
| if class_labels is None: | |
| raise ValueError("class_labels should be provided when num_class_embeds > 0") | |
| if self.model.unet.config.class_embed_type == "timestep": | |
| class_labels = self.model.unet.time_proj(class_labels) | |
| # `Timesteps` does not contain any weights and will always return f32 tensors | |
| # there might be better ways to encapsulate this. | |
| class_labels = class_labels.to(dtype=sample.dtype) | |
| class_emb = self.model.unet.class_embedding(class_labels).to(dtype=sample.dtype) | |
| if self.model.unet.config.class_embeddings_concat: | |
| emb = torch.cat([emb, class_emb], dim=-1) | |
| else: | |
| emb = emb + class_emb | |
| if self.model.unet.config.addition_embed_type == "text": | |
| aug_emb = self.model.unet.add_embedding(encoder_hidden_states) | |
| emb = emb + aug_emb | |
| elif self.model.unet.config.addition_embed_type == "text_image": | |
| # Kadinsky 2.1 - style | |
| if "image_embeds" not in added_cond_kwargs: | |
| raise ValueError( | |
| f"{self.model.unet.__class__} has the config param `addition_embed_type` set to 'text_image' " | |
| f"which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" | |
| ) | |
| image_embs = added_cond_kwargs.get("image_embeds") | |
| text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) | |
| aug_emb = self.model.unet.add_embedding(text_embs, image_embs) | |
| emb = emb + aug_emb | |
| if self.model.unet.time_embed_act is not None: | |
| emb = self.model.unet.time_embed_act(emb) | |
| if self.model.unet.encoder_hid_proj is not None and self.model.unet.config.encoder_hid_dim_type == "text_proj": | |
| encoder_hidden_states = self.model.unet.encoder_hid_proj(encoder_hidden_states) | |
| elif self.model.unet.encoder_hid_proj is not None and \ | |
| self.model.unet.config.encoder_hid_dim_type == "text_image_proj": | |
| # Kadinsky 2.1 - style | |
| if "image_embeds" not in added_cond_kwargs: | |
| raise ValueError( | |
| f"{self.model.unet.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' " | |
| f"which requires the keyword argument `image_embeds` to be passed in `added_conditions`" | |
| ) | |
| image_embeds = added_cond_kwargs.get("image_embeds") | |
| encoder_hidden_states = self.model.unet.encoder_hid_proj(encoder_hidden_states, image_embeds) | |
| # 2. pre-process | |
| sample = self.model.unet.conv_in(sample) | |
| # 3. down | |
| down_block_res_samples = (sample,) | |
| for downsample_block in self.model.unet.down_blocks: | |
| if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: | |
| sample, res_samples = downsample_block( | |
| hidden_states=sample, | |
| temb=emb, | |
| encoder_hidden_states=encoder_hidden_states, | |
| attention_mask=attention_mask, | |
| cross_attention_kwargs=cross_attention_kwargs, | |
| encoder_attention_mask=encoder_attention_mask, | |
| ) | |
| else: | |
| sample, res_samples = downsample_block(hidden_states=sample, temb=emb) | |
| down_block_res_samples += res_samples | |
| if down_block_additional_residuals is not None: | |
| new_down_block_res_samples = () | |
| for down_block_res_sample, down_block_additional_residual in zip( | |
| down_block_res_samples, down_block_additional_residuals | |
| ): | |
| down_block_res_sample = down_block_res_sample + down_block_additional_residual | |
| new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) | |
| down_block_res_samples = new_down_block_res_samples | |
| # 4. mid | |
| if self.model.unet.mid_block is not None: | |
| sample = self.model.unet.mid_block( | |
| sample, | |
| emb, | |
| encoder_hidden_states=encoder_hidden_states, | |
| attention_mask=attention_mask, | |
| cross_attention_kwargs=cross_attention_kwargs, | |
| encoder_attention_mask=encoder_attention_mask, | |
| ) | |
| # print(sample.shape) | |
| if replace_h_space is None: | |
| h_space = sample.clone() | |
| else: | |
| h_space = replace_h_space | |
| sample = replace_h_space.clone() | |
| if mid_block_additional_residual is not None: | |
| sample = sample + mid_block_additional_residual | |
| extracted_res_conns = {} | |
| # 5. up | |
| for i, upsample_block in enumerate(self.model.unet.up_blocks): | |
| is_final_block = i == len(self.model.unet.up_blocks) - 1 | |
| res_samples = down_block_res_samples[-len(upsample_block.resnets):] | |
| down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] | |
| if replace_skip_conns is not None and replace_skip_conns.get(i): | |
| res_samples = replace_skip_conns.get(i) | |
| if zero_out_resconns is not None: | |
| if (type(zero_out_resconns) is int and i >= (zero_out_resconns - 1)) or \ | |
| type(zero_out_resconns) is list and i in zero_out_resconns: | |
| res_samples = [torch.zeros_like(x) for x in res_samples] | |
| # down_block_res_samples = [torch.zeros_like(x) for x in down_block_res_samples] | |
| extracted_res_conns[i] = res_samples | |
| # if we have not reached the final block and need to forward the | |
| # upsample size, we do it here | |
| if not is_final_block and forward_upsample_size: | |
| upsample_size = down_block_res_samples[-1].shape[2:] | |
| if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: | |
| sample = upsample_block( | |
| hidden_states=sample, | |
| temb=emb, | |
| res_hidden_states_tuple=res_samples, | |
| encoder_hidden_states=encoder_hidden_states, | |
| cross_attention_kwargs=cross_attention_kwargs, | |
| upsample_size=upsample_size, | |
| attention_mask=attention_mask, | |
| encoder_attention_mask=encoder_attention_mask, | |
| ) | |
| else: | |
| sample = upsample_block( | |
| hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size | |
| ) | |
| # 6. post-process | |
| if self.model.unet.conv_norm_out: | |
| sample = self.model.unet.conv_norm_out(sample) | |
| sample = self.model.unet.conv_act(sample) | |
| sample = self.model.unet.conv_out(sample) | |
| if not return_dict: | |
| return (sample,) | |
| return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns | |
| class AudioLDM2Wrapper(PipelineWrapper): | |
| def __init__(self, *args, **kwargs) -> None: | |
| super().__init__(*args, **kwargs) | |
| if self.double_precision: | |
| self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, torch_dtype=torch.float64).to(self.device) | |
| else: | |
| try: | |
| self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=True).to(self.device) | |
| except FileNotFoundError: | |
| self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=False).to(self.device) | |
| def load_scheduler(self): | |
| # self.model.scheduler = DDIMScheduler.from_config(self.model_id, subfolder="scheduler") | |
| self.model.scheduler = DDIMScheduler.from_pretrained(self.model_id, subfolder="scheduler") | |
| def get_fn_STFT(self): | |
| from audioldm.audio import TacotronSTFT | |
| return TacotronSTFT( | |
| filter_length=1024, | |
| hop_length=160, | |
| win_length=1024, | |
| n_mel_channels=64, | |
| sampling_rate=16000, | |
| mel_fmin=0, | |
| mel_fmax=8000, | |
| ) | |
| def vae_encode(self, x): | |
| # self.model.vae.disable_tiling() | |
| if x.shape[2] % 4: | |
| x = torch.nn.functional.pad(x, (0, 0, 4 - (x.shape[2] % 4), 0)) | |
| return (self.model.vae.encode(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float() | |
| # return (self.encode_no_tiling(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float() | |
| def vae_decode(self, x): | |
| return self.model.vae.decode(1 / self.model.vae.config.scaling_factor * x).sample | |
| def decode_to_mel(self, x): | |
| if self.double_precision: | |
| tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().double()).detach() | |
| tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().float()).detach() | |
| if len(tmp.shape) == 1: | |
| tmp = tmp.unsqueeze(0) | |
| return tmp | |
| def encode_text(self, prompts: List[str]): | |
| tokenizers = [self.model.tokenizer, self.model.tokenizer_2] | |
| text_encoders = [self.model.text_encoder, self.model.text_encoder_2] | |
| prompt_embeds_list = [] | |
| attention_mask_list = [] | |
| for tokenizer, text_encoder in zip(tokenizers, text_encoders): | |
| text_inputs = tokenizer( | |
| prompts, | |
| padding="max_length" if isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast)) else True, | |
| max_length=tokenizer.model_max_length, | |
| truncation=True, | |
| return_tensors="pt", | |
| ) | |
| text_input_ids = text_inputs.input_ids | |
| attention_mask = text_inputs.attention_mask | |
| untruncated_ids = tokenizer(prompts, padding="longest", return_tensors="pt").input_ids | |
| if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] \ | |
| and not torch.equal(text_input_ids, untruncated_ids): | |
| removed_text = tokenizer.batch_decode( | |
| untruncated_ids[:, tokenizer.model_max_length - 1: -1]) | |
| print(f"The following part of your input was truncated because {text_encoder.config.model_type} can " | |
| f"only handle sequences up to {tokenizer.model_max_length} tokens: {removed_text}" | |
| ) | |
| text_input_ids = text_input_ids.to(self.device) | |
| attention_mask = attention_mask.to(self.device) | |
| with torch.no_grad(): | |
| if text_encoder.config.model_type == "clap": | |
| prompt_embeds = text_encoder.get_text_features( | |
| text_input_ids, | |
| attention_mask=attention_mask, | |
| ) | |
| # append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size) | |
| prompt_embeds = prompt_embeds[:, None, :] | |
| # make sure that we attend to this single hidden-state | |
| attention_mask = attention_mask.new_ones((len(prompts), 1)) | |
| else: | |
| prompt_embeds = text_encoder( | |
| text_input_ids, | |
| attention_mask=attention_mask, | |
| ) | |
| prompt_embeds = prompt_embeds[0] | |
| prompt_embeds_list.append(prompt_embeds) | |
| attention_mask_list.append(attention_mask) | |
| # print(f'prompt[0].shape: {prompt_embeds_list[0].shape}') | |
| # print(f'prompt[1].shape: {prompt_embeds_list[1].shape}') | |
| # print(f'attn[0].shape: {attention_mask_list[0].shape}') | |
| # print(f'attn[1].shape: {attention_mask_list[1].shape}') | |
| projection_output = self.model.projection_model( | |
| hidden_states=prompt_embeds_list[0], | |
| hidden_states_1=prompt_embeds_list[1], | |
| attention_mask=attention_mask_list[0], | |
| attention_mask_1=attention_mask_list[1], | |
| ) | |
| projected_prompt_embeds = projection_output.hidden_states | |
| projected_attention_mask = projection_output.attention_mask | |
| generated_prompt_embeds = self.model.generate_language_model( | |
| projected_prompt_embeds, | |
| attention_mask=projected_attention_mask, | |
| max_new_tokens=None, | |
| ) | |
| prompt_embeds = prompt_embeds.to(dtype=self.model.text_encoder_2.dtype, device=self.device) | |
| attention_mask = ( | |
| attention_mask.to(device=self.device) | |
| if attention_mask is not None | |
| else torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=self.device) | |
| ) | |
| generated_prompt_embeds = generated_prompt_embeds.to(dtype=self.model.language_model.dtype, device=self.device) | |
| return generated_prompt_embeds, prompt_embeds, attention_mask | |
| def get_variance(self, timestep, prev_timestep): | |
| alpha_prod_t = self.model.scheduler.alphas_cumprod[timestep] | |
| alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep) | |
| beta_prod_t = 1 - alpha_prod_t | |
| beta_prod_t_prev = 1 - alpha_prod_t_prev | |
| variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) | |
| return variance | |
| def get_alpha_prod_t_prev(self, prev_timestep): | |
| return self.model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 \ | |
| else self.model.scheduler.final_alpha_cumprod | |
| def unet_forward(self, | |
| sample: torch.FloatTensor, | |
| timestep: Union[torch.Tensor, float, int], | |
| encoder_hidden_states: torch.Tensor, | |
| timestep_cond: Optional[torch.Tensor] = None, | |
| class_labels: Optional[torch.Tensor] = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| encoder_attention_mask: Optional[torch.Tensor] = None, | |
| return_dict: bool = True, | |
| cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
| mid_block_additional_residual: Optional[torch.Tensor] = None, | |
| replace_h_space: Optional[torch.Tensor] = None, | |
| replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None, | |
| zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple: | |
| # Translation | |
| encoder_hidden_states_1 = class_labels | |
| class_labels = None | |
| encoder_attention_mask_1 = encoder_attention_mask | |
| encoder_attention_mask = None | |
| # return self.model.unet(sample, timestep, | |
| # encoder_hidden_states=generated_prompt_embeds, | |
| # encoder_hidden_states_1=encoder_hidden_states_1, | |
| # encoder_attention_mask_1=encoder_attention_mask_1, | |
| # ), None, None | |
| # By default samples have to be AT least a multiple of the overall upsampling factor. | |
| # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). | |
| # However, the upsampling interpolation output size can be forced to fit any upsampling size | |
| # on the fly if necessary. | |
| default_overall_up_factor = 2 ** self.model.unet.num_upsamplers | |
| # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` | |
| forward_upsample_size = False | |
| upsample_size = None | |
| if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): | |
| # print("Forward upsample size to force interpolation output size.") | |
| forward_upsample_size = True | |
| # ensure attention_mask is a bias, and give it a singleton query_tokens dimension | |
| # expects mask of shape: | |
| # [batch, key_tokens] | |
| # adds singleton query_tokens dimension: | |
| # [batch, 1, key_tokens] | |
| # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: | |
| # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) | |
| # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) | |
| if attention_mask is not None: | |
| # assume that mask is expressed as: | |
| # (1 = keep, 0 = discard) | |
| # convert mask into a bias that can be added to attention scores: | |
| # (keep = +0, discard = -10000.0) | |
| attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 | |
| attention_mask = attention_mask.unsqueeze(1) | |
| # convert encoder_attention_mask to a bias the same way we do for attention_mask | |
| if encoder_attention_mask is not None: | |
| encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 | |
| encoder_attention_mask = encoder_attention_mask.unsqueeze(1) | |
| if encoder_attention_mask_1 is not None: | |
| encoder_attention_mask_1 = (1 - encoder_attention_mask_1.to(sample.dtype)) * -10000.0 | |
| encoder_attention_mask_1 = encoder_attention_mask_1.unsqueeze(1) | |
| # 1. time | |
| timesteps = timestep | |
| if not torch.is_tensor(timesteps): | |
| # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can | |
| # This would be a good case for the `match` statement (Python 3.10+) | |
| is_mps = sample.device.type == "mps" | |
| if isinstance(timestep, float): | |
| dtype = torch.float32 if is_mps else torch.float64 | |
| else: | |
| dtype = torch.int32 if is_mps else torch.int64 | |
| timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) | |
| elif len(timesteps.shape) == 0: | |
| timesteps = timesteps[None].to(sample.device) | |
| # broadcast to batch dimension in a way that's compatible with ONNX/Core ML | |
| timesteps = timesteps.expand(sample.shape[0]) | |
| t_emb = self.model.unet.time_proj(timesteps) | |
| # `Timesteps` does not contain any weights and will always return f32 tensors | |
| # but time_embedding might actually be running in fp16. so we need to cast here. | |
| # there might be better ways to encapsulate this. | |
| t_emb = t_emb.to(dtype=sample.dtype) | |
| emb = self.model.unet.time_embedding(t_emb, timestep_cond) | |
| aug_emb = None | |
| if self.model.unet.class_embedding is not None: | |
| if class_labels is None: | |
| raise ValueError("class_labels should be provided when num_class_embeds > 0") | |
| if self.model.unet.config.class_embed_type == "timestep": | |
| class_labels = self.model.unet.time_proj(class_labels) | |
| # `Timesteps` does not contain any weights and will always return f32 tensors | |
| # there might be better ways to encapsulate this. | |
| class_labels = class_labels.to(dtype=sample.dtype) | |
| class_emb = self.model.unet.class_embedding(class_labels).to(dtype=sample.dtype) | |
| if self.model.unet.config.class_embeddings_concat: | |
| emb = torch.cat([emb, class_emb], dim=-1) | |
| else: | |
| emb = emb + class_emb | |
| emb = emb + aug_emb if aug_emb is not None else emb | |
| if self.model.unet.time_embed_act is not None: | |
| emb = self.model.unet.time_embed_act(emb) | |
| # 2. pre-process | |
| sample = self.model.unet.conv_in(sample) | |
| # 3. down | |
| down_block_res_samples = (sample,) | |
| for downsample_block in self.model.unet.down_blocks: | |
| if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: | |
| sample, res_samples = downsample_block( | |
| hidden_states=sample, | |
| temb=emb, | |
| encoder_hidden_states=encoder_hidden_states, | |
| attention_mask=attention_mask, | |
| cross_attention_kwargs=cross_attention_kwargs, | |
| encoder_attention_mask=encoder_attention_mask, | |
| encoder_hidden_states_1=encoder_hidden_states_1, | |
| encoder_attention_mask_1=encoder_attention_mask_1, | |
| ) | |
| else: | |
| sample, res_samples = downsample_block(hidden_states=sample, temb=emb) | |
| down_block_res_samples += res_samples | |
| # 4. mid | |
| if self.model.unet.mid_block is not None: | |
| sample = self.model.unet.mid_block( | |
| sample, | |
| emb, | |
| encoder_hidden_states=encoder_hidden_states, | |
| attention_mask=attention_mask, | |
| cross_attention_kwargs=cross_attention_kwargs, | |
| encoder_attention_mask=encoder_attention_mask, | |
| encoder_hidden_states_1=encoder_hidden_states_1, | |
| encoder_attention_mask_1=encoder_attention_mask_1, | |
| ) | |
| if replace_h_space is None: | |
| h_space = sample.clone() | |
| else: | |
| h_space = replace_h_space | |
| sample = replace_h_space.clone() | |
| if mid_block_additional_residual is not None: | |
| sample = sample + mid_block_additional_residual | |
| extracted_res_conns = {} | |
| # 5. up | |
| for i, upsample_block in enumerate(self.model.unet.up_blocks): | |
| is_final_block = i == len(self.model.unet.up_blocks) - 1 | |
| res_samples = down_block_res_samples[-len(upsample_block.resnets):] | |
| down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] | |
| if replace_skip_conns is not None and replace_skip_conns.get(i): | |
| res_samples = replace_skip_conns.get(i) | |
| if zero_out_resconns is not None: | |
| if (type(zero_out_resconns) is int and i >= (zero_out_resconns - 1)) or \ | |
| type(zero_out_resconns) is list and i in zero_out_resconns: | |
| res_samples = [torch.zeros_like(x) for x in res_samples] | |
| # down_block_res_samples = [torch.zeros_like(x) for x in down_block_res_samples] | |
| extracted_res_conns[i] = res_samples | |
| # if we have not reached the final block and need to forward the | |
| # upsample size, we do it here | |
| if not is_final_block and forward_upsample_size: | |
| upsample_size = down_block_res_samples[-1].shape[2:] | |
| if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: | |
| sample = upsample_block( | |
| hidden_states=sample, | |
| temb=emb, | |
| res_hidden_states_tuple=res_samples, | |
| encoder_hidden_states=encoder_hidden_states, | |
| cross_attention_kwargs=cross_attention_kwargs, | |
| upsample_size=upsample_size, | |
| attention_mask=attention_mask, | |
| encoder_attention_mask=encoder_attention_mask, | |
| encoder_hidden_states_1=encoder_hidden_states_1, | |
| encoder_attention_mask_1=encoder_attention_mask_1, | |
| ) | |
| else: | |
| sample = upsample_block( | |
| hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size | |
| ) | |
| # 6. post-process | |
| if self.model.unet.conv_norm_out: | |
| sample = self.model.unet.conv_norm_out(sample) | |
| sample = self.model.unet.conv_act(sample) | |
| sample = self.model.unet.conv_out(sample) | |
| if not return_dict: | |
| return (sample,) | |
| return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns | |
| def forward(self, *args, **kwargs): | |
| return self | |
| def load_model(model_id, device, double_precision=False): | |
| ldm_stable = AudioLDM2Wrapper(model_id=model_id, device=device, double_precision=double_precision) | |
| ldm_stable.load_scheduler() | |
| torch.cuda.empty_cache() | |
| return ldm_stable | |