add

acestep/ui/gradio/events/training/youtube_download_registry.py

@@ -0,0 +1,82 @@
+"""Per-folder record of YouTube video IDs already downloaded by the bulk importer.
+
+The Dataset Builder's "Download All" flow stores each downloaded video's
+``v=`` key (the YouTube video ID) in a small JSON file inside the target
+folder, mapping ID -> saved mp3 filename. On the next run, links whose ID is
+already recorded (and whose mp3 still exists) are skipped instead of being
+downloaded again.
+"""
+
+import json
+import os
+import re
+from typing import Optional
+from urllib.parse import parse_qs, urlparse
+
+from loguru import logger
+
+REGISTRY_FILENAME = ".youtube_downloaded.json"
+
+_VIDEO_ID_RE = re.compile(r"^[A-Za-z0-9_-]{11}$")
+_PATH_ID_PREFIXES = ("/shorts/", "/live/", "/embed/")
+
+
+def extract_video_id(url: str) -> Optional[str]:
+    """Return the YouTube video ID (the ``v=`` key) of ``url``, or ``None``.
+
+    Handles ``watch?v=``, ``youtu.be/<id>``, ``/shorts/<id>``, ``/live/<id>``
+    and ``/embed/<id>`` links; extra query params (``list=``, ``t=``...) are
+    ignored.
+    """
+    parsed = urlparse((url or "").strip())
+    host = (parsed.hostname or "").lower()
+    candidate = None
+    if host.endswith("youtu.be"):
+        candidate = parsed.path.lstrip("/").split("/")[0]
+    elif host.endswith("youtube.com"):
+        candidate = (parse_qs(parsed.query).get("v") or [None])[0]
+        if candidate is None:
+            for prefix in _PATH_ID_PREFIXES:
+                if parsed.path.startswith(prefix):
+                    candidate = parsed.path[len(prefix):].split("/")[0]
+                    break
+    if candidate and _VIDEO_ID_RE.match(candidate):
+        return candidate
+    return None
+
+
+def _registry_path(target_dir: str) -> str:
+    """Return the registry file path inside ``target_dir``."""
+    return os.path.join(target_dir, REGISTRY_FILENAME)
+
+
+def load_registry(target_dir: str) -> dict[str, str]:
+    """Load the ``{video_id: mp3_filename}`` map for ``target_dir``.
+
+    Returns an empty dict when the file is missing or unreadable.
+    """
+    path = _registry_path(target_dir)
+    if not os.path.isfile(path):
+        return {}
+    try:
+        with open(path, encoding="utf-8") as fh:
+            data = json.load(fh)
+    except (OSError, ValueError) as exc:
+        logger.warning(f"[youtube_registry] Ignoring unreadable registry {path}: {exc}")
+        return {}
+    if not isinstance(data, dict):
+        return {}
+    return {str(k): str(v) for k, v in data.items()}
+
+
+def save_registry(target_dir: str, registry: dict[str, str]) -> None:
+    """Write the ``{video_id: mp3_filename}`` map to ``target_dir``."""
+    path = _registry_path(target_dir)
+    with open(path, "w", encoding="utf-8") as fh:
+        json.dump(registry, fh, ensure_ascii=False, indent=2, sort_keys=True)
+
+
+def is_already_downloaded(registry: dict[str, str], video_id: str, target_dir: str) -> bool:
+    """Return whether ``video_id`` is recorded and its mp3 still exists in ``target_dir``."""
+    filename = registry.get(video_id)
+    return bool(filename) and os.path.isfile(os.path.join(target_dir, filename))

acestep/ui/gradio/events/training/youtube_import.py

@@ -19,10 +19,39 @@ from loguru import logger
 
 from acestep.youtube_download import YoutubeDownloadError, download_youtube_audio
 from .training_utils import LORA_TRAINING_AUDIO_DIR
+from .youtube_download_registry import (
+    extract_video_id,
+    is_already_downloaded,
+    load_registry,
+    save_registry,
+)
 
 _URL_SPLIT_RE = re.compile(r"[,\n]+")
 
 
+def _save_registry_safely(target_dir: str, registry: dict[str, str]) -> None:
+    """Persist the downloaded-ID registry, logging (not raising) on write failure."""
+    try:
+        save_registry(target_dir, registry)
+    except OSError as exc:
+        logger.warning(f"[youtube_import] Could not save download registry: {exc}")
+
+
+def _format_index_summary(
+    downloaded_idx: list[int], skipped_idx: list[int], failed_idx: list[int]
+) -> str:
+    """Format the 1-based link numbers of each outcome, e.g. ``Downloaded: #1, #4``."""
+
+    def _fmt(indices: list[int]) -> str:
+        return ", ".join(f"#{i}" for i in indices) or "-"
+
+    return (
+        f"✅ Downloaded: {_fmt(downloaded_idx)}\n"
+        f"⏭️ Already downloaded: {_fmt(skipped_idx)}\n"
+        f"⚠️ Failed: {_fmt(failed_idx)}"
+    )
+
+
 def _split_urls(urls_text: str) -> list[str]:
     """Split a paste of one or more YouTube links, comma- and/or newline-separated."""
     raw = (urls_text or "").strip()
@@ -67,25 +96,43 @@ def download_youtube_batch(urls_text: str, target_dir: str):
     log_lines = [f"\U0001f680 Downloading {len(urls)} link(s) into: {target_dir}"]
     yield "\n".join(log_lines)
 
-    ok_count = 0
+    registry = load_registry(target_dir)
+    downloaded_idx: list[int] = []
+    skipped_idx: list[int] = []
+    failed_idx: list[int] = []
     for idx, url in enumerate(urls, start=1):
         prefix = f"[{idx}/{len(urls)}]"
+        video_id = extract_video_id(url)
+        if video_id and is_already_downloaded(registry, video_id, target_dir):
+            skipped_idx.append(idx)
+            log_lines.append(f"{prefix} ⏭️ Already downloaded ({video_id}): {registry[video_id]}")
+            yield "\n".join(log_lines)
+            continue
+
         log_lines.append(f"{prefix} \U0001f4e5 Downloading: {url}")
         yield "\n".join(log_lines)
         try:
             saved_path = download_youtube_audio(url, output_dir=target_dir)
         except YoutubeDownloadError as exc:
             logger.warning(f"[youtube_import] Download failed for {url}: {exc}")
+            failed_idx.append(idx)
             log_lines.append(f"{prefix} ⚠️ Skipped (error): {exc}")
             yield "\n".join(log_lines)
             continue
 
-        ok_count += 1
+        downloaded_idx.append(idx)
+        if video_id:
+            registry[video_id] = os.path.basename(saved_path)
+            _save_registry_safely(target_dir, registry)
         log_lines.append(f"{prefix} ✅ Saved: {os.path.basename(saved_path)}")
         yield "\n".join(log_lines)
 
+    summary = _format_index_summary(downloaded_idx, skipped_idx, failed_idx)
+    logger.info(f"[youtube_import] Batch finished in {target_dir}:\n{summary}")
     log_lines.append(
-        f"\n\U0001f3c1 Done: {ok_count}/{len(urls)} link(s) downloaded successfully.\n"
+        f"\n\U0001f3c1 Done: {len(downloaded_idx)}/{len(urls)} link(s) downloaded successfully, "
+        f"{len(skipped_idx)} already downloaded (skipped), {len(failed_idx)} failed.\n"
+        f"{summary}\n"
         f"Now click \"Scan New Directory\" above with this folder to build the dataset:\n{target_dir}"
     )
     yield "\n".join(log_lines)

acestep/ui/gradio/events/training/youtube_import_test.py

@@ -0,0 +1,94 @@
+"""Unit tests for the bulk YouTube importer's skip-already-downloaded behavior."""
+
+import json
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from acestep.ui.gradio.events.training import youtube_import
+from acestep.ui.gradio.events.training.youtube_download_registry import (
+    REGISTRY_FILENAME,
+    extract_video_id,
+    load_registry,
+)
+
+
+class ExtractVideoIdTests(unittest.TestCase):
+    """Verify the ``v=`` key is pulled from the common YouTube URL shapes."""
+
+    def test_watch_url_with_playlist_params(self):
+        """Extra list/start_radio params should not affect the extracted ID."""
+        url = "https://www.youtube.com/watch?v=JFLLHYWkIxc&list=RDEMMKSMwpeSGs9aNG657TmCfA&start_radio=1"
+        self.assertEqual(extract_video_id(url), "JFLLHYWkIxc")
+
+    def test_short_and_shorts_urls(self):
+        """youtu.be and /shorts/ links should both yield the ID."""
+        self.assertEqual(extract_video_id("https://youtu.be/JFLLHYWkIxc?t=5"), "JFLLHYWkIxc")
+        self.assertEqual(extract_video_id("https://youtube.com/shorts/JFLLHYWkIxc"), "JFLLHYWkIxc")
+
+    def test_non_youtube_or_invalid_returns_none(self):
+        """Non-YouTube hosts and malformed IDs should return None."""
+        self.assertIsNone(extract_video_id("https://vimeo.com/watch?v=JFLLHYWkIxc"))
+        self.assertIsNone(extract_video_id("https://www.youtube.com/watch?v=short"))
+        self.assertIsNone(extract_video_id(""))
+
+
+class DownloadYoutubeBatchSkipTests(unittest.TestCase):
+    """Verify recorded IDs are skipped and new downloads are recorded."""
+
+    def setUp(self):
+        """Create a scratch target folder per test."""
+        self._tmp = tempfile.TemporaryDirectory()
+        self.target_dir = self._tmp.name
+
+    def tearDown(self):
+        """Remove the scratch folder."""
+        self._tmp.cleanup()
+
+    def _fake_download(self, url, output_dir):
+        """Write a dummy mp3 named after the video ID and return its path."""
+        path = os.path.join(output_dir, f"song-{extract_video_id(url)}.mp3")
+        with open(path, "wb") as fh:
+            fh.write(b"mp3")
+        return path
+
+    def _run(self, urls_text):
+        """Run the batch with a mocked downloader, returning (mock, final log)."""
+        with patch.object(
+            youtube_import, "download_youtube_audio", side_effect=self._fake_download
+        ) as mock_dl:
+            log = list(youtube_import.download_youtube_batch(urls_text, self.target_dir))[-1]
+        return mock_dl, log
+
+    def test_new_download_is_recorded(self):
+        """A successful download should store its ID -> filename in the registry."""
+        mock_dl, _ = self._run("https://www.youtube.com/watch?v=JFLLHYWkIxc")
+        self.assertEqual(mock_dl.call_count, 1)
+        self.assertEqual(load_registry(self.target_dir), {"JFLLHYWkIxc": "song-JFLLHYWkIxc.mp3"})
+
+    def test_second_run_and_duplicates_are_skipped(self):
+        """Already-downloaded IDs (earlier run or same paste) should not be re-downloaded."""
+        self._run("https://www.youtube.com/watch?v=JFLLHYWkIxc")
+        urls = (
+            "https://youtu.be/JFLLHYWkIxc,"
+            "https://www.youtube.com/watch?v=OuRZl8f3FeM\n"
+            "https://www.youtube.com/watch?v=OuRZl8f3FeM&list=RDx"
+        )
+        mock_dl, log = self._run(urls)
+        self.assertEqual(mock_dl.call_count, 1)
+        self.assertIn("2 already downloaded", log)
+        self.assertIn("✅ Downloaded: #2", log)
+        self.assertIn("⏭️ Already downloaded: #1, #3", log)
+        self.assertIn("⚠️ Failed: -", log)
+
+    def test_redownloads_when_recorded_file_is_missing(self):
+        """A recorded ID whose mp3 was deleted should be downloaded again."""
+        with open(os.path.join(self.target_dir, REGISTRY_FILENAME), "w", encoding="utf-8") as fh:
+            json.dump({"JFLLHYWkIxc": "deleted.mp3"}, fh)
+        mock_dl, _ = self._run("https://www.youtube.com/watch?v=JFLLHYWkIxc")
+        self.assertEqual(mock_dl.call_count, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()