from __future__ import annotations import os import zipfile from dataclasses import asdict, dataclass from pathlib import Path, PurePosixPath from typing import Callable, Iterable from huggingface_hub import HfApi, hf_hub_download from .config import ( DATASET_SAMPLE_ROOT, MAX_DATASET_SAMPLE_ARCHIVE_BYTES, MAX_DATASET_SAMPLE_FILE_BYTES, ) _AUDIO_EXTENSIONS = {".wav", ".flac", ".mp3", ".ogg", ".opus", ".m4a", ".aiff", ".aif"} _ARCHIVE_EXTENSIONS = {".zip"} @dataclass(frozen=True) class DatasetSampleSpec: sample_id: str label: str repo_id: str repo_type: str revision: str license_id: str attribution: str source_url: str filename: str | None = None selector: str | None = None first_download_note: str = "" @dataclass(frozen=True) class ResolvedDatasetSample: sample_id: str label: str local_path: str repo_id: str repo_type: str revision: str repository_filename: str resolved_audio_filename: str license_id: str attribution: str source_url: str archive_extracted: bool size_bytes: int def to_json(self) -> str: return json.dumps(asdict(self), ensure_ascii=False, sort_keys=True) DATASET_SAMPLE_SPECS: tuple[DatasetSampleSpec, ...] = ( DatasetSampleSpec( sample_id="sonicsets-demo-mix", label="Dataset: SonicSets ensemble mix (CC BY-NC 4.0; first download may be large)", repo_id="sonicsets-data/Stems-Evaluation-Kit", repo_type="dataset", revision="main", license_id="CC-BY-NC-4.0", attribution="SonicSets High-Fidelity Stems Evaluation Kit", source_url="https://huggingface.co/datasets/sonicsets-data/Stems-Evaluation-Kit", selector="ensemble-mix", first_download_note=( "The repository is about 333 MB in total. SESA downloads one matching audio file when possible; " "if the repository exposes only an archive, the first download can approach the archive size." ), ), DatasetSampleSpec( sample_id="legacy-symphony-01", label="Hosted fallback: Symphony of Automation 01 (small)", repo_id="set-soft/audio_separation", repo_type="model", revision="main", filename="Audio_Examples/01-The_Symphony_of_Automation.mp3", license_id="MIT repository metadata", attribution="set-soft/audio_separation hosted audio example", source_url="https://huggingface.co/set-soft/audio_separation/tree/main/Audio_Examples", first_download_note="Small hosted fallback. The repository does not separately document audio-asset provenance.", ), DatasetSampleSpec( sample_id="legacy-symphony-04", label="Hosted fallback: Symphony of Automation 04 (small)", repo_id="set-soft/audio_separation", repo_type="model", revision="main", filename="Audio_Examples/04-The_Symphony_of_Automation.mp3", license_id="MIT repository metadata", attribution="set-soft/audio_separation hosted audio example", source_url="https://huggingface.co/set-soft/audio_separation/tree/main/Audio_Examples", first_download_note="Small hosted fallback. The repository does not separately document audio-asset provenance.", ), DatasetSampleSpec( sample_id="legacy-symphony-05", label="Hosted fallback: Symphony of Automation 05 (small)", repo_id="set-soft/audio_separation", repo_type="model", revision="main", filename="Audio_Examples/05-The_Symphony_of_Automation.mp3", license_id="MIT repository metadata", attribution="set-soft/audio_separation hosted audio example", source_url="https://huggingface.co/set-soft/audio_separation/tree/main/Audio_Examples", first_download_note="Small hosted fallback. The repository does not separately document audio-asset provenance.", ), ) _SPEC_BY_ID = {item.sample_id: item for item in DATASET_SAMPLE_SPECS} DEFAULT_DATASET_SAMPLE_ID = DATASET_SAMPLE_SPECS[0].sample_id def dataset_sample_choices() -> list[tuple[str, str]]: return [(item.label, item.sample_id) for item in DATASET_SAMPLE_SPECS] def dataset_sample_markdown(sample_id: str | None) -> str: spec = _SPEC_BY_ID.get(str(sample_id or "")) if not spec: return "Select a curated hosted sample, or upload a file to override it." note = f" \n{spec.first_download_note}" if spec.first_download_note else "" return ( f"**Source:** `{spec.repo_id}` ยท **license:** `{spec.license_id}` \n" f"Attribution: {spec.attribution}.{note}" ) def _repo_token() -> str | bool: return os.environ.get("HF_TOKEN") or False def _score_audio_candidate(filename: str, selector: str | None) -> int: path = PurePosixPath(filename) if path.suffix.lower() not in _AUDIO_EXTENSIONS: return -10_000 text = str(path).lower().replace("-", "_").replace(" ", "_") score = 0 if selector == "ensemble-mix": preferred = ( ("full_ensemble_mix", 500), ("ensemble_mix", 450), ("full_mix", 400), ("mixture", 350), ("mix", 250), ("ensemble", 180), ) for token, points in preferred: if token in text: score += points for token in ("isolated", "stem", "vocal", "vocals", "bass", "drum", "guitar", "harmony"): if token in text: score -= 220 if "demo_01" in text or "demo01" in text or "/01" in text: score += 35 score -= len(path.parts) * 2 return score def _choose_audio_filename(filenames: Iterable[str], selector: str | None) -> str | None: ranked = sorted( ((-_score_audio_candidate(name, selector), len(name), name) for name in filenames), ) if not ranked: return None best_score = -ranked[0][0] if best_score <= 0: return None return ranked[0][2] def _choose_archive_filename(filenames: Iterable[str]) -> str | None: archives = [name for name in filenames if PurePosixPath(name).suffix.lower() in _ARCHIVE_EXTENSIONS] if not archives: return None return sorted(archives, key=lambda value: (len(PurePosixPath(value).parts), len(value), value))[0] def _safe_extract_zip(archive_path: Path, destination: Path) -> None: destination.mkdir(parents=True, exist_ok=True) root = destination.resolve() total_size = 0 with zipfile.ZipFile(archive_path) as archive: for member in archive.infolist(): if member.is_dir(): continue total_size += max(0, int(member.file_size)) if total_size > MAX_DATASET_SAMPLE_ARCHIVE_BYTES: raise RuntimeError("Dataset sample archive expands beyond the configured safety limit.") target = (destination / member.filename).resolve() if root not in target.parents: raise RuntimeError("Unsafe path found in dataset sample archive.") archive.extractall(destination) def _resolve_from_archive(archive_path: Path, spec: DatasetSampleSpec) -> Path: extracted = archive_path.parent / "extracted" marker = extracted / ".complete" if not marker.is_file(): _safe_extract_zip(archive_path, extracted) marker.write_text("ok\n", encoding="utf-8") candidates = [ str(path.relative_to(extracted).as_posix()) for path in extracted.rglob("*") if path.is_file() and path.suffix.lower() in _AUDIO_EXTENSIONS ] selected = _choose_audio_filename(candidates, spec.selector) if not selected: raise RuntimeError("No suitable mixture audio was found inside the dataset archive.") return extracted / selected def resolve_dataset_sample( sample_id: str, *, api: HfApi | None = None, download_fn: Callable[..., str] = hf_hub_download, ) -> ResolvedDatasetSample: spec = _SPEC_BY_ID.get(str(sample_id or "")) if not spec: raise ValueError("Unknown dataset sample selection.") local_dir = DATASET_SAMPLE_ROOT / spec.sample_id local_dir.mkdir(parents=True, exist_ok=True) repository_filename = spec.filename if repository_filename is None: client = api or HfApi(token=_repo_token()) filenames = client.list_repo_files( repo_id=spec.repo_id, repo_type=spec.repo_type, revision=spec.revision, ) repository_filename = _choose_audio_filename(filenames, spec.selector) if repository_filename is None: repository_filename = _choose_archive_filename(filenames) if repository_filename is None: raise RuntimeError("No supported audio file or ZIP archive was found in the dataset repository.") downloaded = Path( download_fn( repo_id=spec.repo_id, repo_type=spec.repo_type, filename=repository_filename, revision=spec.revision, local_dir=local_dir, token=_repo_token(), etag_timeout=15, ) ) if not downloaded.is_file() or downloaded.stat().st_size <= 0: raise RuntimeError("Dataset sample download did not produce a readable file.") if downloaded.stat().st_size > MAX_DATASET_SAMPLE_FILE_BYTES: raise RuntimeError("Dataset sample download exceeds the configured file-size limit.") archive_extracted = downloaded.suffix.lower() in _ARCHIVE_EXTENSIONS local_audio = _resolve_from_archive(downloaded, spec) if archive_extracted else downloaded if not local_audio.is_file() or local_audio.suffix.lower() not in _AUDIO_EXTENSIONS: raise RuntimeError("Resolved dataset sample is not a supported audio file.") return ResolvedDatasetSample( sample_id=spec.sample_id, label=spec.label, local_path=str(local_audio), repo_id=spec.repo_id, repo_type=spec.repo_type, revision=spec.revision, repository_filename=repository_filename, resolved_audio_filename=local_audio.name, license_id=spec.license_id, attribution=spec.attribution, source_url=spec.source_url, archive_extracted=archive_extracted, size_bytes=local_audio.stat().st_size, )