Do interpolation after

This commit is contained in:
kijai 2025-09-04 20:08:26 +03:00
parent ca2e7f3a6c
commit bfa45fb0e6
2 changed files with 4 additions and 18 deletions

View File

@ -7,7 +7,7 @@ import torchaudio
class AudioEncoderModel(): class AudioEncoderModel():
def __init__(self, config, fps=50): def __init__(self, config):
self.load_device = comfy.model_management.text_encoder_device() self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device() offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device) self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
@ -21,7 +21,6 @@ class AudioEncoderModel():
self.model.eval() self.model.eval()
self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device) self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.model_sample_rate = 16000 self.model_sample_rate = 16000
self.fps = fps
def load_sd(self, sd): def load_sd(self, sd):
return self.model.load_state_dict(sd, strict=False) return self.model.load_state_dict(sd, strict=False)
@ -32,7 +31,7 @@ class AudioEncoderModel():
def encode_audio(self, audio, sample_rate): def encode_audio(self, audio, sample_rate):
comfy.model_management.load_model_gpu(self.patcher) comfy.model_management.load_model_gpu(self.patcher)
audio = torchaudio.functional.resample(audio, sample_rate, self.model_sample_rate) audio = torchaudio.functional.resample(audio, sample_rate, self.model_sample_rate)
out, all_layers = self.model(audio.to(self.load_device), fps=self.fps, sr=self.model_sample_rate) out, all_layers = self.model(audio.to(self.load_device), sr=self.model_sample_rate)
outputs = {} outputs = {}
outputs["encoded_audio"] = out outputs["encoded_audio"] = out
outputs["encoded_audio_all_layers"] = all_layers outputs["encoded_audio_all_layers"] = all_layers
@ -52,7 +51,6 @@ def load_audio_encoder_from_sd(sd, prefix=""):
"do_normalize": True, "do_normalize": True,
"do_stable_layer_norm": True "do_stable_layer_norm": True
} }
fps = 50
elif embed_dim == 768: # base elif embed_dim == 768: # base
config = { config = {
"embed_dim": 768, "embed_dim": 768,
@ -63,11 +61,10 @@ def load_audio_encoder_from_sd(sd, prefix=""):
"do_normalize": False, # chinese-wav2vec2-base has this False "do_normalize": False, # chinese-wav2vec2-base has this False
"do_stable_layer_norm": False "do_stable_layer_norm": False
} }
fps = 25
else: else:
raise RuntimeError("ERROR: audio encoder file is invalid or unsupported embed_dim: {}".format(embed_dim)) raise RuntimeError("ERROR: audio encoder file is invalid or unsupported embed_dim: {}".format(embed_dim))
audio_encoder = AudioEncoderModel(config, fps=fps) audio_encoder = AudioEncoderModel(config)
m, u = audio_encoder.load_sd(sd) m, u = audio_encoder.load_sd(sd)
if len(m) > 0: if len(m) > 0:
logging.warning("missing audio encoder: {}".format(m)) logging.warning("missing audio encoder: {}".format(m))

View File

@ -238,26 +238,15 @@ class Wav2Vec2Model(nn.Module):
device=device, dtype=dtype, operations=operations device=device, dtype=dtype, operations=operations
) )
def forward(self, x, fps=50, sr=16000, mask_time_indices=None, return_dict=False): def forward(self, x, sr=16000, mask_time_indices=None, return_dict=False):
x = torch.mean(x, dim=1) x = torch.mean(x, dim=1)
if self.do_normalize: if self.do_normalize:
x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7) x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7)
features = self.feature_extractor(x) features = self.feature_extractor(x)
if fps != 50:
audio_duration = x.shape[1] / sr
target_seq_len = int(audio_duration * fps)
features = self.linear_interpolation(features, target_seq_len)
features = self.feature_projection(features) features = self.feature_projection(features)
batch_size, seq_len, _ = features.shape batch_size, seq_len, _ = features.shape
x, all_x = self.encoder(features) x, all_x = self.encoder(features)
return x, all_x return x, all_x
def linear_interpolation(self, features, target_seq_len):
features = features.transpose(1, 2)
output_features = torch.nn.functional.interpolate(features, size=target_seq_len, align_corners=True, mode='linear')
return output_features.transpose(1, 2)