Skip to content

Support wav2vec base models #9637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
comfyanonymous merged 5 commits into from
Sep 13, 2025
Merged

Support wav2vec base models #9637

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b64004f
Support wav2vec base models
kijai Aug 30, 2025
66e1091
trim trailing whitespace
kijai Aug 30, 2025
874a067
Merge remote-tracking branch 'upstream/master'
kijai Aug 31, 2025
ca2e7f3
Merge remote-tracking branch 'upstream/master'
kijai Sep 4, 2025
bfa45fb
Do interpolation after
kijai Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 37 additions & 4 deletions comfy/audio_encoders/audio_encoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,21 @@


class AudioEncoderModel():
def __init__(self, config):
def __init__(self, config, fps=50):
self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
self.model = Wav2Vec2Model(dtype=self.dtype, device=offload_device, operations=comfy.ops.manual_cast)
model_config = dict(config)
model_config.update({
"dtype": self.dtype,
"device": offload_device,
"operations": comfy.ops.manual_cast
})
self.model = Wav2Vec2Model(**model_config)
self.model.eval()
self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.model_sample_rate = 16000
self.fps = fps

def load_sd(self, sd):
return self.model.load_state_dict(sd, strict=False)
Expand All @@ -25,16 +32,42 @@ def get_sd(self):
def encode_audio(self, audio, sample_rate):
comfy.model_management.load_model_gpu(self.patcher)
audio = torchaudio.functional.resample(audio, sample_rate, self.model_sample_rate)
out, all_layers = self.model(audio.to(self.load_device))
out, all_layers = self.model(audio.to(self.load_device), fps=self.fps, sr=self.model_sample_rate)
outputs = {}
outputs["encoded_audio"] = out
outputs["encoded_audio_all_layers"] = all_layers
return outputs


def load_audio_encoder_from_sd(sd, prefix=""):
audio_encoder = AudioEncoderModel(None)
sd = comfy.utils.state_dict_prefix_replace(sd, {"wav2vec2.": ""})
embed_dim = sd["encoder.layer_norm.bias"].shape[0]
if embed_dim == 1024:# large
config = {
"embed_dim": 1024,
"num_heads": 16,
"num_layers": 24,
"conv_norm": True,
"conv_bias": True,
"do_normalize": True,
"do_stable_layer_norm": True
}
fps = 50
elif embed_dim == 768: # base
config = {
"embed_dim": 768,
"num_heads": 12,
"num_layers": 12,
"conv_norm": False,
"conv_bias": False,
"do_normalize": False, # chinese-wav2vec2-base has this False
"do_stable_layer_norm": False
}
fps = 25
else:
raise RuntimeError("ERROR: audio encoder file is invalid or unsupported embed_dim: {}".format(embed_dim))

audio_encoder = AudioEncoderModel(config, fps=fps)
m, u = audio_encoder.load_sd(sd)
if len(m) > 0:
logging.warning("missing audio encoder: {}".format(m))
Expand Down
98 changes: 77 additions & 21 deletions comfy/audio_encoders/wav2vec2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,49 @@ def forward(self, x):
x = self.conv(x)
return torch.nn.functional.gelu(self.layer_norm(x.transpose(-2, -1)).transpose(-2, -1))

class LayerGroupNormConv(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
super().__init__()
self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)
self.layer_norm = operations.GroupNorm(num_groups=out_channels, num_channels=out_channels, affine=True, device=device, dtype=dtype)

def forward(self, x):
x = self.conv(x)
return torch.nn.functional.gelu(self.layer_norm(x))

class ConvNoNorm(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
super().__init__()
self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)

def forward(self, x):
x = self.conv(x)
return torch.nn.functional.gelu(x)


class ConvFeatureEncoder(nn.Module):
def __init__(self, conv_dim, dtype=None, device=None, operations=None):
def __init__(self, conv_dim, conv_bias=False, conv_norm=True, dtype=None, device=None, operations=None):
super().__init__()
self.conv_layers = nn.ModuleList([
LayerNormConv(1, conv_dim, kernel_size=10, stride=5, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
])
if conv_norm:
self.conv_layers = nn.ModuleList([
LayerNormConv(1, conv_dim, kernel_size=10, stride=5, bias=True, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
])
else:
self.conv_layers = nn.ModuleList([
LayerGroupNormConv(1, conv_dim, kernel_size=10, stride=5, bias=conv_bias, device=device, dtype=dtype, operations=operations),
ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
ConvNoNorm(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
ConvNoNorm(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
])

def forward(self, x):
x = x.unsqueeze(1)
Expand Down Expand Up @@ -76,6 +106,7 @@ def __init__(
num_heads=12,
num_layers=12,
mlp_ratio=4.0,
do_stable_layer_norm=True,
dtype=None, device=None, operations=None
):
super().__init__()
Expand All @@ -86,20 +117,25 @@ def __init__(
embed_dim=embed_dim,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
do_stable_layer_norm=do_stable_layer_norm,
device=device, dtype=dtype, operations=operations
)
for _ in range(num_layers)
])

self.layer_norm = operations.LayerNorm(embed_dim, eps=1e-05, device=device, dtype=dtype)
self.do_stable_layer_norm = do_stable_layer_norm

def forward(self, x, mask=None):
x = x + self.pos_conv_embed(x)
all_x = ()
if not self.do_stable_layer_norm:
x = self.layer_norm(x)
for layer in self.layers:
all_x += (x,)
x = layer(x, mask)
x = self.layer_norm(x)
if self.do_stable_layer_norm:
x = self.layer_norm(x)
all_x += (x,)
return x, all_x

Expand Down Expand Up @@ -145,6 +181,7 @@ def __init__(
embed_dim=768,
num_heads=12,
mlp_ratio=4.0,
do_stable_layer_norm=True,
dtype=None, device=None, operations=None
):
super().__init__()
Expand All @@ -154,15 +191,19 @@ def __init__(
self.layer_norm = operations.LayerNorm(embed_dim, device=device, dtype=dtype)
self.feed_forward = FeedForward(embed_dim, mlp_ratio, device=device, dtype=dtype, operations=operations)
self.final_layer_norm = operations.LayerNorm(embed_dim, device=device, dtype=dtype)
self.do_stable_layer_norm = do_stable_layer_norm

def forward(self, x, mask=None):
residual = x
x = self.layer_norm(x)
if self.do_stable_layer_norm:
x = self.layer_norm(x)
x = self.attention(x, mask=mask)
x = residual + x

x = x + self.feed_forward(self.final_layer_norm(x))
return x
if not self.do_stable_layer_norm:
x = self.layer_norm(x)
return self.final_layer_norm(x + self.feed_forward(x))
else:
return x + self.feed_forward(self.final_layer_norm(x))


class Wav2Vec2Model(nn.Module):
Expand All @@ -174,34 +215,49 @@ def __init__(
final_dim=256,
num_heads=16,
num_layers=24,
conv_norm=True,
conv_bias=True,
do_normalize=True,
do_stable_layer_norm=True,
dtype=None, device=None, operations=None
):
super().__init__()

conv_dim = 512
self.feature_extractor = ConvFeatureEncoder(conv_dim, device=device, dtype=dtype, operations=operations)
self.feature_extractor = ConvFeatureEncoder(conv_dim, conv_norm=conv_norm, conv_bias=conv_bias, device=device, dtype=dtype, operations=operations)
self.feature_projection = FeatureProjection(conv_dim, embed_dim, device=device, dtype=dtype, operations=operations)

self.masked_spec_embed = nn.Parameter(torch.empty(embed_dim, device=device, dtype=dtype))
self.do_normalize = do_normalize

self.encoder = TransformerEncoder(
embed_dim=embed_dim,
num_heads=num_heads,
num_layers=num_layers,
do_stable_layer_norm=do_stable_layer_norm,
device=device, dtype=dtype, operations=operations
)

def forward(self, x, mask_time_indices=None, return_dict=False):

def forward(self, x, fps=50, sr=16000, mask_time_indices=None, return_dict=False):
x = torch.mean(x, dim=1)

x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7)
if self.do_normalize:
x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7)

features = self.feature_extractor(x)
features = self.feature_projection(features)

if fps != 50:

@comfyanonymous comfyanonymous Aug 30, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any reason this shouldn't be done by the node using the outputs like the s2v node? https://github.com/comfyanonymous/ComfyUI/blob/master/comfy_extras/nodes_wan.py#L892

@kijai kijai Aug 30, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm unsure about that, this was the way it was done in the initial MultiTalk implementation. I guess the reason could be efficiency as this way the model has to process less? Though the models being as light as they are, I'm not sure it matters, mostly I was trying to match how it was done in the initial MultiTalk implementation to get matching results.

I can test if it makes a difference.

@comfyanonymous comfyanonymous Aug 30, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the best is to have the encoder nodes return the raw output from the model as much as possible which makes it easier if multiple models use the output from the same encoder but with different processing.

@kijai kijai Aug 31, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Well I agree with that principle of course. Tested now and there is a slight difference in results when interpolating before the encode and when interpolating the already encoded features, with the same function as S2V, should we worry about that?

AnimateDiff_00005-audio.webm

@comfyanonymous comfyanonymous Aug 31, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's weird, why is there a difference?

@kijai kijai Aug 31, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know, but looking at the tensors there's already a slight difference in the features if the interpolation is done before feature_projection:

lerp before proj: shape=torch.Size([1, 601, 768]) min=-12.0035 max=11.7886 mean=-0.0433 std=1.6125
lerp after proj:  shape=torch.Size([1, 601, 768]) min=-11.9092 max=10.5639 mean=-0.0436 std=1.5525
Mean abs diff: 0.053896

and similarly between encoding interpolated features and interpolating encoded features:

lerp before encode: shape=torch.Size([13, 601, 768]) min=-4.3660 max=4.7442 mean=0.0022 std=0.3017
lerp after encode: shape=torch.Size([13, 601, 768]) min=-3.9754 max=5.0701 mean=0.0027 std=0.2990
Overall mean abs diff across all layers: 0.079161

It's very small, can it really be that sensitive, maybe it matters if the model is trained like that?

The original implementation is specifically overriding the Wav2vec2 -model to do the interpolation before projection and encode:
https://github.com/MeiGen-AI/MultiTalk/blob/636c774370d40214e6be1ea1ed98bc219dfbd8c4/src/audio_analysis/wav2vec2.py#L32

Another longer comparison test with better settings overall still shows a difference, but I can't really say if it's better or worse:

AnimateDiff_00002-audio.webm

Anyway there's no hurry with this PR as the native model implementation still needs to be done to have any use for this, so which way to do this can be decided later too.

@comfyanonymous comfyanonymous Sep 1, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do all the models do the same interpolation?

@kijai kijai Sep 1, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do all the models do the same interpolation?

As far as I can see, the normal Wav2Vec2Model implementation in transformers uses same forward for base and large, and it doesn't involve interpolation, so it seems it's just specifically this chinese model from Tencent that uses it, and it would explain why they use custom code for it. The MultiTalk paper didn't mention the fps used or anything, so I'm just guessing.

@kijai kijai Sep 10, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did end up leaving the interpolation later as you suggested, and while the result is slightly different, it has been working fine. If there are no other issues left it would be fine by me to merge this.

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)
batch_size, seq_len, _ = features.shape

x, all_x = self.encoder(features)

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