104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from sam_audio_comfy.audio import (
|
|
AudioChunk,
|
|
crossfade_audio_chunks,
|
|
plan_audio_chunks,
|
|
prepare_audio,
|
|
prepare_visual_prompt,
|
|
result_to_audio,
|
|
slice_visual_frames,
|
|
)
|
|
|
|
|
|
def test_chunk_plan_covers_audio_with_requested_overlap():
|
|
assert plan_audio_chunks(21, 1, 10.0, 2.0) == [
|
|
AudioChunk(0, 10),
|
|
AudioChunk(8, 18),
|
|
AudioChunk(16, 21),
|
|
]
|
|
assert plan_audio_chunks(21, 1, 0.0, 2.0) == [AudioChunk(0, 21)]
|
|
assert plan_audio_chunks(5, 1, 10.0, 2.0) == [AudioChunk(0, 5)]
|
|
|
|
with pytest.raises(ValueError, match="shorter than chunk_duration"):
|
|
plan_audio_chunks(21, 1, 10.0, 10.0)
|
|
|
|
|
|
def test_crossfade_audio_chunks_normalizes_overlap():
|
|
first = torch.ones(1, 1, 6)
|
|
second = torch.full((1, 1, 6), 3.0)
|
|
output = crossfade_audio_chunks(
|
|
[(AudioChunk(0, 6), first), (AudioChunk(4, 10), second)], 10
|
|
)
|
|
|
|
assert torch.equal(output[..., :4], torch.ones(1, 1, 4))
|
|
assert torch.allclose(output[0, 0, 4:6], torch.tensor([5 / 3, 7 / 3]))
|
|
assert torch.equal(output[..., 6:], torch.full((1, 1, 4), 3.0))
|
|
|
|
|
|
def test_visual_frames_are_sliced_to_chunk_time_range():
|
|
frames = torch.arange(10).view(10, 1)
|
|
|
|
assert torch.equal(
|
|
slice_visual_frames(frames, AudioChunk(2, 6), 10), frames[2:6]
|
|
)
|
|
assert torch.equal(
|
|
slice_visual_frames(frames, AudioChunk(0, 10), 10), frames
|
|
)
|
|
|
|
|
|
def test_prepare_audio_preserves_batch_and_resamples():
|
|
waveform = torch.stack(
|
|
[torch.ones(2, 8), torch.full((2, 8), 0.5)], dim=0
|
|
)
|
|
rows, length = prepare_audio(
|
|
{"waveform": waveform, "sample_rate": 8}, target_sample_rate=4
|
|
)
|
|
|
|
assert len(rows) == 2
|
|
assert rows[0].shape == (2, 4)
|
|
assert rows[1].shape == (2, 4)
|
|
assert length == 4
|
|
|
|
|
|
def test_result_to_audio_crops_and_pads_mono_batch():
|
|
audio = result_to_audio(
|
|
[torch.arange(8), torch.arange(3)], sample_rate=48_000, target_length=5
|
|
)
|
|
|
|
assert audio["sample_rate"] == 48_000
|
|
assert audio["waveform"].shape == (2, 1, 5)
|
|
assert torch.equal(audio["waveform"][0, 0], torch.arange(5).float())
|
|
assert torch.equal(
|
|
audio["waveform"][1, 0], torch.tensor([0.0, 1.0, 2.0, 0.0, 0.0])
|
|
)
|
|
|
|
|
|
def test_prepare_visual_prompt_converts_layout_scale_and_repeats_mask():
|
|
images = torch.zeros(3, 4, 6, 3)
|
|
images[..., 0] = 1.0
|
|
mask = torch.tensor([[0.0, 1.0], [0.0, 0.0]])
|
|
|
|
frames, masks = prepare_visual_prompt(images, mask)
|
|
|
|
assert frames.shape == (3, 3, 4, 6)
|
|
assert frames.dtype == torch.uint8
|
|
assert frames[:, 0].min().item() == 255
|
|
assert frames[:, 1:].max().item() == 0
|
|
assert masks.shape == (3, 1, 4, 6)
|
|
assert masks.dtype == torch.bool
|
|
assert torch.equal(masks[0], masks[1])
|
|
|
|
|
|
def test_prepare_visual_prompt_rejects_mismatched_frame_counts():
|
|
images = torch.zeros(3, 4, 6, 3)
|
|
mask = torch.zeros(2, 4, 6)
|
|
|
|
try:
|
|
prepare_visual_prompt(images, mask)
|
|
except ValueError as error:
|
|
assert "2 frames" in str(error)
|
|
else:
|
|
raise AssertionError("expected a frame-count validation error")
|