58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import pytest
|
|
from typing import Dict, Union
|
|
from pathlib import Path
|
|
import torch
|
|
from ..core.io import from_disk_as_raw_3d
|
|
|
|
# Test data roots
|
|
DATA_ROOT = Path("ComfyUI_SignalProcessing/audio")
|
|
INPUT_FILE = DATA_ROOT / "inputs/002-orig.mp4"
|
|
INPUT_FILE = DATA_ROOT / "inputs/pf-01.mp3"
|
|
INPUT_IR_FILE = DATA_ROOT / "inputs/ir.wav"
|
|
INPUTS_ROOT = DATA_ROOT / "inputs"
|
|
INPUT_FILES = {
|
|
file.name: file.resolve() for file in INPUTS_ROOT.rglob("*") if file.is_file()
|
|
}
|
|
|
|
TestData = Dict[str, Union[Dict[str, Union[torch.Tensor, int]], Path]]
|
|
|
|
|
|
@pytest.fixture
|
|
def test_data(request) -> TestData:
|
|
|
|
test_name = request.node.name # Automatically get the current test function name
|
|
test_name = str(request.node.function.__name__)
|
|
|
|
OUTPUT_ROOT = DATA_ROOT / f"outputs/{test_name}"
|
|
INPUTS_ROOT = DATA_ROOT / "inputs"
|
|
|
|
param_values = (
|
|
request.node.callspec.params if hasattr(request.node, "callspec") else {}
|
|
)
|
|
pest_param_str = "_".join(f"{key}-{value}" for key, value in param_values.items())
|
|
|
|
audio_slice_begin_seconds: float = 60.0
|
|
audio_slice_duration_seconds: float = 120.0
|
|
|
|
# Prepare audio data
|
|
audio, sample_rate = from_disk_as_raw_3d(
|
|
str(INPUT_FILE.absolute()),
|
|
try_gpu=True,
|
|
start_seconds=audio_slice_begin_seconds,
|
|
duration_seconds=audio_slice_duration_seconds,
|
|
)
|
|
OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) # Ensure the output directory exists
|
|
|
|
audio_to_comfy: Dict[torch.Tensor, int] = {
|
|
"waveform": audio,
|
|
"sample_rate": sample_rate,
|
|
}
|
|
|
|
return {
|
|
"audio": audio_to_comfy,
|
|
"output_root": OUTPUT_ROOT,
|
|
"inputs_root": INPUTS_ROOT,
|
|
"test_name": test_name,
|
|
"pest_param_str": pest_param_str,
|
|
}
|