| import math |
|
|
| import torch |
|
|
| import node_helpers |
| from comfy_extras.nodes_wan import get_audio_embed_bucket_fps, linear_interpolation |
|
|
| WAN_AUDIO_INPUT_FPS = 50 |
| WAN_AUDIO_VIDEO_RATE = 30 |
| WAN_AUDIO_FPS = 16 |
| WAN_AUDIO_SAMPLE_RATE = 16000 |
|
|
|
|
| def _audio_feat(audio_encoder_output): |
| feat = torch.cat(audio_encoder_output["encoded_audio_all_layers"]) |
| return linear_interpolation(feat, input_fps=WAN_AUDIO_INPUT_FPS, output_fps=WAN_AUDIO_VIDEO_RATE) |
|
|
|
|
| def audio_encoder_output_video_frames(audio_encoder_output, fps=WAN_AUDIO_FPS): |
| audio_samples = audio_encoder_output.get("audio_samples") |
| if audio_samples is not None: |
| return max(1, int(round(audio_samples / float(WAN_AUDIO_SAMPLE_RATE) * fps))) |
| feat = _audio_feat(audio_encoder_output) |
| return max(1, int(round(feat.shape[1] * fps / WAN_AUDIO_VIDEO_RATE))) |
|
|
|
|
| def _permute_audio_embed_bucket(audio_embed_bucket): |
| audio_embed_bucket = audio_embed_bucket.unsqueeze(0) |
| if len(audio_embed_bucket.shape) == 3: |
| return audio_embed_bucket.permute(0, 2, 1) |
| return audio_embed_bucket.permute(0, 2, 3, 1) |
|
|
|
|
| def build_timeline_audio_embed(length, segments): |
| latent_t = ((length - 1) // 4) + 1 |
| batch_frames = latent_t * 4 |
| total_feat_frames = int(math.ceil(batch_frames * WAN_AUDIO_VIDEO_RATE / WAN_AUDIO_FPS)) |
|
|
| composite = None |
| cursor_auto = 0 |
| for segment in segments: |
| feat = _audio_feat(segment["audio_encoder_output"]) |
| if composite is None: |
| composite = torch.zeros( |
| feat.shape[0], total_feat_frames, feat.shape[2], |
| dtype=feat.dtype, device=feat.device) |
|
|
| start_frame = segment.get("start_frame", -1) |
| if start_frame is None or start_frame < 0: |
| start_frame = cursor_auto |
| else: |
| start_frame = int(start_frame) |
|
|
| start_feat = int(round(start_frame * WAN_AUDIO_VIDEO_RATE / WAN_AUDIO_FPS)) |
| copy_len = min(feat.shape[1], total_feat_frames - start_feat) |
| if copy_len > 0 and start_feat < total_feat_frames: |
| composite[:, start_feat:start_feat + copy_len, :] = feat[:, :copy_len, :] |
|
|
| cursor_auto = start_frame + audio_encoder_output_video_frames(segment["audio_encoder_output"]) |
|
|
| audio_embed_bucket, _ = get_audio_embed_bucket_fps( |
| composite, fps=WAN_AUDIO_FPS, batch_frames=batch_frames, m=0, video_rate=WAN_AUDIO_VIDEO_RATE) |
| audio_embed_bucket = _permute_audio_embed_bucket(audio_embed_bucket) |
| return audio_embed_bucket[:, :, :, :batch_frames] |
|
|
|
|
| def resolve_timeline_segment_ranges(length, segments): |
| batch_frames = (((length - 1) // 4) + 1) * 4 |
| resolved = [] |
| cursor_auto = 0 |
| for segment in segments: |
| start_frame = segment.get("start_frame", -1) |
| if start_frame is None or start_frame < 0: |
| start_frame = cursor_auto |
| else: |
| start_frame = int(start_frame) |
| end_frame = min(batch_frames, start_frame + audio_encoder_output_video_frames(segment["audio_encoder_output"])) |
| resolved.append({**segment, "start_frame": start_frame, "end_frame": end_frame}) |
| cursor_auto = end_frame |
| return resolved |
|
|
|
|
| def apply_timeline_audio_conditioning(positive, negative, length, segments): |
| audio_embed_bucket = build_timeline_audio_embed(length, segments) |
| if audio_embed_bucket is None or audio_embed_bucket.shape[3] <= 0: |
| return positive, negative |
| positive = node_helpers.conditioning_set_values(positive, {"audio_embed": audio_embed_bucket}) |
| negative = node_helpers.conditioning_set_values(negative, {"audio_embed": audio_embed_bucket * 0.0}) |
| return positive, negative |
