mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-26 12:17:57 +08:00
Do interpolation after
This commit is contained in:
parent
ca2e7f3a6c
commit
bfa45fb0e6
@ -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))
|
||||||
|
|||||||
@ -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)
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user