Use HF dataset repo as source of truth for dataset.json

Push dataset.json to HF repo after every label/preprocess batch so
progress persists across ephemeral Space restarts. On download, pull
the JSON back and restore labels before scanning.

- Add upload_dataset_json_to_hf() and pull dataset.json in download
- Remove hf_token param (use HF_TOKEN env var)
- Reconcile stale .pt preprocessed flags on restart
- Simplify Data Source UI: remove upload files, HF token, resume accordion
- Add batch mode (max_count) for labeling and preprocessing
- Add preprocessed flag to AudioSample, scan dedup via existing_paths
- Make @spaces.GPU a no-op when running locally

acestep/training/dataset_builder_modules/core.py

@@ -18,3 +18,7 @@ class CoreMixin:
     def get_labeled_count(self) -> int:
         """Get the number of labeled samples."""
         return sum(1 for s in self.samples if s.labeled)
+
+    def get_preprocessed_count(self) -> int:
+        """Get the number of preprocessed samples."""
+        return sum(1 for s in self.samples if s.preprocessed)

acestep/training/dataset_builder_modules/label_all.py

@@ -14,9 +14,14 @@ class LabelAllMixin:
         transcribe_lyrics: bool = False,
         skip_metas: bool = False,
         only_unlabeled: bool = False,
+        max_count: int = 0,
         progress_callback=None,
     ) -> Tuple[List[AudioSample], str]:
-        """Label all samples in the dataset."""
+        """Label all samples in the dataset.
+
+        Args:
+            max_count: When > 0, stop after labeling this many samples (batch mode).
+        """
         if not self.samples:
             return [], "❌ No samples to label. Please scan a directory first."
 
@@ -30,6 +35,9 @@ class LabelAllMixin:
         if not samples_to_label:
             return self.samples, "✅ All samples already labeled"
 
+        batch_limit = max_count if max_count > 0 else len(samples_to_label)
+        samples_to_label = samples_to_label[:batch_limit]
+
         success_count = 0
         fail_count = 0
         total = len(samples_to_label)
@@ -53,10 +61,13 @@ class LabelAllMixin:
             else:
                 fail_count += 1
 
-        status_msg = f"✅ Labeled {success_count}/{total} samples"
+        total_labeled = sum(1 for s in self.samples if s.labeled)
+        total_samples = len(self.samples)
+        remaining = total_samples - total_labeled
+
+        status_msg = f"✅ Labeled {success_count} this batch"
         if fail_count > 0:
             status_msg += f" ({fail_count} failed)"
-        if only_unlabeled:
-            status_msg += f" (unlabeled only, {len(self.samples)} total)"
+        status_msg += f" | {total_labeled}/{total_samples} labeled total, {remaining} remaining"
 
         return self.samples, status_msg

acestep/training/dataset_builder_modules/models.py

@@ -31,6 +31,7 @@ class AudioSample:
     is_instrumental: bool = True
     custom_tag: str = ""
     labeled: bool = False
+    preprocessed: bool = False
     prompt_override: Optional[str] = None  # None=use global ratio, "caption" or "genre"
 
     def __post_init__(self):

acestep/training/dataset_builder_modules/preprocess.py

@@ -29,16 +29,30 @@ class PreprocessMixin:
         dit_handler,
         output_dir: str,
         max_duration: float = 240.0,
+        max_count: int = 0,
         progress_callback=None,
     ) -> Tuple[List[str], str]:
-        """Preprocess all labeled samples to tensor files for efficient training."""
-        debug_log_for("dataset", f"preprocess_to_tensors: output_dir='{output_dir}', max_duration={max_duration}")
+        """Preprocess labeled samples to tensor files for efficient training.
+
+        Args:
+            max_count: When > 0, stop after processing this many new samples (batch mode).
+        """
+        debug_log_for("dataset", f"preprocess_to_tensors: output_dir='{output_dir}', max_duration={max_duration}, max_count={max_count}")
         if not self.samples:
             return [], "❌ No samples to preprocess"
 
-        labeled_samples = [s for s in self.samples if s.labeled]
+        # Reset stale preprocessed flags (ephemeral .pt files may be gone after restart)
+        for s in self.samples:
+            if s.preprocessed and not os.path.exists(os.path.join(output_dir, f"{s.id}.pt")):
+                s.preprocessed = False
+
+        labeled_samples = [s for s in self.samples if s.labeled and not s.preprocessed]
         if not labeled_samples:
-            return [], "❌ No labeled samples to preprocess"
+            total_preprocessed = sum(1 for s in self.samples if s.preprocessed)
+            return [], f"✅ All labeled samples already preprocessed ({total_preprocessed} total)"
+
+        if max_count > 0:
+            labeled_samples = labeled_samples[:max_count]
 
         if dit_handler is None or dit_handler.model is None:
             return [], "❌ Model not initialized. Please initialize the service first."
@@ -190,6 +204,7 @@ class PreprocessMixin:
                 torch.save(output_data, output_path)
                 debug_end_verbose_for("dataset", f"torch.save[{i}]", t0)
                 output_paths.append(output_path)
+                sample.preprocessed = True
                 success_count += 1
 
             except Exception as e:
@@ -202,8 +217,13 @@ class PreprocessMixin:
         save_manifest(output_dir, self.metadata, output_paths)
         debug_end_verbose_for("dataset", "save_manifest", t0)
 
-        status = f"✅ Preprocessed {success_count}/{len(labeled_samples)} samples to {output_dir}"
+        total_preprocessed = sum(1 for s in self.samples if s.preprocessed)
+        total_labeled = sum(1 for s in self.samples if s.labeled)
+        remaining = total_labeled - total_preprocessed
+
+        status = f"✅ Preprocessed {success_count} new samples"
         if fail_count > 0:
             status += f" ({fail_count} failed)"
+        status += f" | {total_preprocessed}/{total_labeled} done, {remaining} remaining"
 
         return output_paths, status

acestep/training/dataset_builder_modules/scan.py

@@ -20,7 +20,7 @@ class ScanMixin:
             return [], f"❌ Not a directory: {directory}"
 
         self._current_dir = directory
-        self.samples = []
+        existing_paths = {s.audio_path for s in self.samples}
 
         audio_files = []
         for root, _, files in os.walk(directory):
@@ -29,19 +29,27 @@ class ScanMixin:
                 if ext in SUPPORTED_AUDIO_FORMATS:
                     audio_files.append(os.path.join(root, file))
 
-        if not audio_files:
+        audio_files.sort()
+
+        new_audio_files = [f for f in audio_files if f not in existing_paths]
+
+        if not new_audio_files and existing_paths:
+            return self.samples, (
+                f"✅ No new audio files in {directory} "
+                f"({len(self.samples)} samples already loaded)"
+            )
+
+        if not new_audio_files:
             return [], (
                 f"❌ No audio files found in {directory}\n"
                 f"Supported formats: {', '.join(SUPPORTED_AUDIO_FORMATS)}"
             )
 
-        audio_files.sort()
-
         csv_metadata = load_csv_metadata(directory)
         csv_count = 0
         lyrics_count = 0
 
-        for audio_path in audio_files:
+        for audio_path in new_audio_files:
             try:
                 duration = get_audio_duration(audio_path)
                 lyrics_content, has_lyrics_file = load_lyrics_file(audio_path)
@@ -78,7 +86,9 @@ class ScanMixin:
 
         self.metadata.num_samples = len(self.samples)
 
-        status = f"✅ Found {len(self.samples)} audio files in {directory}"
+        new_count = len(new_audio_files)
+        total_count = len(self.samples)
+        status = f"✅ Added {new_count} new files ({total_count} total samples)"
         if lyrics_count > 0:
             status += f"\n   📝 {lyrics_count} files have accompanying lyrics (.txt)"
         if csv_count > 0:

app.py

@@ -7,16 +7,24 @@ A comprehensive music generation system with three main interfaces:
 """
 
 import gradio as gr
+import pandas as pd
 import torch
 import numpy as np
 from pathlib import Path
 import json
 from typing import Optional, List, Tuple
-import spaces
+try:
+    import spaces
+except ImportError:
+    # Local dev — make @spaces.GPU a no-op
+    class _Spaces:
+        def GPU(self, fn=None, **kwargs):
+            return fn if fn else lambda f: f
+    spaces = _Spaces()
 
 from src.ace_step_engine import ACEStepEngine
 from src.timeline_manager import TimelineManager
-from src.lora_trainer import download_hf_dataset
+from src.lora_trainer import download_hf_dataset, upload_dataset_json_to_hf
 from src.audio_processor import AudioProcessor
 from src.utils import setup_logging, load_config
 from acestep.training.dataset_builder import DatasetBuilder
@@ -289,80 +297,99 @@ def timeline_reset(session_state: dict) -> Tuple[None, None, str, dict]:
 DATAFRAME_HEADERS = ["#", "Filename", "Duration", "Lyrics", "Labeled", "BPM", "Key", "Caption"]
 
 
+def _build_progress_summary():
+    """Build a one-line progress summary from current dataset builder state."""
+    builder = get_dataset_builder()
+    total = builder.get_sample_count()
+    labeled = builder.get_labeled_count()
+    preprocessed = builder.get_preprocessed_count()
+    remaining = total - labeled
+    return f"Total: {total} | Labeled: {labeled} | Preprocessed: {preprocessed} | Remaining: {remaining}"
+
+
 def _build_review_dataframe():
     """Build editable dataframe rows from current dataset builder state."""
     builder = get_dataset_builder()
     return builder.get_samples_dataframe_data()
 
 
-def lora_upload_and_scan(files, training_state):
-    """Copy uploaded audio files to working dir and scan."""
+def lora_download_hf(dataset_id, max_files, hf_offset, training_state):
+    """Download HuggingFace dataset batch, restore labels from HF repo, and scan."""
     try:
-        if not files:
-            return "No files uploaded", training_state
+        if not dataset_id or not dataset_id.strip():
+            return "Enter a dataset ID (e.g. username/dataset-name)", training_state, int(hf_offset or 0), _build_progress_summary()
 
-        import shutil
+        offset_val = int(hf_offset or 0)
+        max_files_val = int(max_files)
 
-        work_dir = Path("lora_training") / "uploaded"
-        work_dir.mkdir(parents=True, exist_ok=True)
+        local_dir, dl_status = download_hf_dataset(
+            dataset_id.strip(),
+            max_files=max_files_val,
+            offset=offset_val,
+        )
 
-        for f in files:
-            src = Path(f)
-            shutil.copy2(str(src), str(work_dir / src.name))
+        if not local_dir:
+            return f"Download failed: {dl_status}", training_state, offset_val, _build_progress_summary()
 
         builder = get_dataset_builder()
-        samples, status = builder.scan_directory(str(work_dir))
-
-        training_state = training_state or {}
-        training_state["audio_dir"] = str(work_dir)
 
-        return f"Scanned {len(samples)} audio files from uploads", training_state
+        # Restore labels/flags from dataset.json pulled from HF repo
+        dataset_json_path = str(Path(local_dir) / "dataset.json")
+        if Path(dataset_json_path).exists():
+            builder.load_dataset(dataset_json_path)
+            dl_status += " | Restored labels from HF repo"
 
-    except Exception as e:
-        logger.error(f"Upload scan failed: {e}")
-        return f"Error: {e}", training_state or {}
+        # Scan directory — skips already-tracked files via existing_paths check
+        samples, scan_status = builder.scan_directory(local_dir)
 
+        training_state = training_state or {}
+        training_state["audio_dir"] = local_dir
+        training_state["dataset_id"] = dataset_id.strip()
+        training_state["dataset_path"] = dataset_json_path
 
-def lora_download_hf(dataset_id, hf_token, max_files, training_state):
-    """Download HuggingFace dataset and scan for audio files."""
-    try:
-        if not dataset_id or not dataset_id.strip():
-            return "Enter a dataset ID (e.g. pedroapfilho/lofi-tracks)", training_state
+        next_offset = offset_val + max_files_val
 
-        token = hf_token.strip() if hf_token else None
+        return f"{dl_status} | {scan_status}", training_state, next_offset, _build_progress_summary()
 
-        local_dir, dl_status = download_hf_dataset(
-            dataset_id.strip(), max_files=int(max_files), hf_token=token
-        )
+    except Exception as e:
+        logger.error(f"HF download failed: {e}")
+        return f"Error: {e}", training_state or {}, int(hf_offset or 0), _build_progress_summary()
 
-        if not local_dir:
-            return f"Download failed: {dl_status}", training_state
 
+def lora_save_dataset_to_json(training_state):
+    """Explicitly save the current dataset to JSON."""
+    try:
         builder = get_dataset_builder()
-        samples, scan_status = builder.scan_directory(local_dir)
+
+        if builder.get_sample_count() == 0:
+            return "No samples to save"
 
         training_state = training_state or {}
-        training_state["audio_dir"] = local_dir
+        dataset_path = training_state.get("dataset_path")
+        if not dataset_path:
+            audio_dir = training_state.get("audio_dir", "lora_training")
+            dataset_path = str(Path(audio_dir) / "dataset.json")
+            training_state["dataset_path"] = dataset_path
 
-        return f"{dl_status} | {scan_status}", training_state
+        return builder.save_dataset(dataset_path)
 
     except Exception as e:
-        logger.error(f"HF download failed: {e}")
-        return f"Error: {e}", training_state or {}
+        logger.error(f"Save dataset failed: {e}")
+        return f"Error: {e}"
 
 
 @spaces.GPU(duration=300)
-def lora_auto_label(training_state, progress=gr.Progress()):
-    """Auto-label all samples using LLM analysis."""
+def lora_auto_label(label_batch_size, training_state, progress=gr.Progress()):
+    """Auto-label unlabeled samples in batches using LLM analysis, then auto-save."""
     try:
         builder = get_dataset_builder()
 
         if builder.get_sample_count() == 0:
-            return [], "No samples loaded. Upload files or download a dataset first."
+            return [], "No samples loaded. Upload files or download a dataset first.", training_state, _build_progress_summary()
 
         engine = get_ace_engine()
         if not engine.is_initialized():
-            return [], "ACE-Step engine not initialized. Models may still be loading."
+            return [], "ACE-Step engine not initialized. Models may still be loading.", training_state, _build_progress_summary()
 
         def progress_callback(msg):
             progress(0, desc=msg)
@@ -370,14 +397,32 @@ def lora_auto_label(training_state, progress=gr.Progress()):
         samples, status = builder.label_all_samples(
             dit_handler=engine.dit_handler,
             llm_handler=engine.llm_handler,
+            only_unlabeled=True,
+            max_count=int(label_batch_size),
             progress_callback=progress_callback,
         )
 
-        return _build_review_dataframe(), status
+        training_state = training_state or {}
+        dataset_path = training_state.get("dataset_path")
+        if not dataset_path:
+            audio_dir = training_state.get("audio_dir", "lora_training")
+            dataset_path = str(Path(audio_dir) / "dataset.json")
+            training_state["dataset_path"] = dataset_path
+
+        save_status = builder.save_dataset(dataset_path)
+        status += f"\n{save_status}"
+
+        # Sync to HF repo so labels persist across sessions
+        dataset_id = training_state.get("dataset_id")
+        if dataset_id:
+            hf_status = upload_dataset_json_to_hf(dataset_id, dataset_path)
+            status += f"\n{hf_status}"
+
+        return _build_review_dataframe(), status, training_state, _build_progress_summary()
 
     except Exception as e:
         logger.error(f"Auto-label failed: {e}")
-        return [], f"Error: {e}"
+        return [], f"Error: {e}", training_state or {}, _build_progress_summary()
 
 
 def lora_save_edits(df_data, training_state):
@@ -385,11 +430,22 @@ def lora_save_edits(df_data, training_state):
     try:
         builder = get_dataset_builder()
 
-        if not df_data or len(df_data) == 0:
+        if df_data is None:
+            return "No data to save"
+
+        if isinstance(df_data, pd.DataFrame):
+            if df_data.empty:
+                return "No data to save"
+            rows = df_data.values.tolist()
+        elif isinstance(df_data, list):
+            if len(df_data) == 0:
+                return "No data to save"
+            rows = df_data
+        else:
             return "No data to save"
 
         updated = 0
-        for row in df_data:
+        for row in rows:
             idx = int(row[0])
             updates = {}
 
@@ -421,17 +477,17 @@ def lora_save_edits(df_data, training_state):
 
 
 @spaces.GPU(duration=300)
-def lora_preprocess(training_state, progress=gr.Progress()):
-    """Preprocess labeled samples to training tensors."""
+def lora_preprocess(preprocess_batch_size, training_state, progress=gr.Progress()):
+    """Preprocess labeled samples to training tensors in batches."""
     try:
         builder = get_dataset_builder()
 
         if builder.get_labeled_count() == 0:
-            return "No labeled samples. Run auto-label first."
+            return "No labeled samples. Run auto-label first.", _build_progress_summary()
 
         engine = get_ace_engine()
         if not engine.is_initialized():
-            return "ACE-Step engine not initialized."
+            return "ACE-Step engine not initialized.", _build_progress_summary()
 
         tensor_dir = str(Path("lora_training") / "tensors")
 
@@ -441,17 +497,34 @@ def lora_preprocess(training_state, progress=gr.Progress()):
         output_paths, status = builder.preprocess_to_tensors(
             dit_handler=engine.dit_handler,
             output_dir=tensor_dir,
+            max_count=int(preprocess_batch_size),
             progress_callback=progress_callback,
         )
 
         training_state = training_state or {}
         training_state["tensor_dir"] = tensor_dir
 
-        return status
+        # Auto-save so preprocessed flags persist across sessions
+        dataset_path = training_state.get("dataset_path")
+        if not dataset_path:
+            audio_dir = training_state.get("audio_dir", "lora_training")
+            dataset_path = str(Path(audio_dir) / "dataset.json")
+            training_state["dataset_path"] = dataset_path
+
+        save_status = builder.save_dataset(dataset_path)
+        status += f"\n{save_status}"
+
+        # Sync to HF repo so preprocessed flags persist across sessions
+        dataset_id = training_state.get("dataset_id")
+        if dataset_id:
+            hf_status = upload_dataset_json_to_hf(dataset_id, dataset_path)
+            status += f"\n{hf_status}"
+
+        return status, _build_progress_summary()
 
     except Exception as e:
         logger.error(f"Preprocess failed: {e}")
-        return f"Error: {e}"
+        return f"Error: {e}", _build_progress_summary()
 
 
 @spaces.GPU(duration=600)
@@ -740,41 +813,38 @@ def create_ui():
 
                 training_state = gr.State(value={})
 
+                lora_progress = gr.Textbox(
+                    label="Progress",
+                    value="Total: 0 | Labeled: 0 | Preprocessed: 0 | Remaining: 0",
+                    interactive=False,
+                )
+
                 with gr.Tabs():
 
                     # ---------- Sub-tab 1: Data Source ----------
                     with gr.Tab("1. Data Source"):
-                        gr.Markdown("Choose one: upload audio files or download from HuggingFace.")
+                        gr.Markdown(
+                            "Download audio from a HuggingFace dataset repo. "
+                            "Labels and progress are synced back to the repo automatically."
+                        )
 
+                        lora_hf_id = gr.Textbox(
+                            label="Dataset ID",
+                            placeholder="username/dataset-name",
+                        )
                         with gr.Row():
-                            with gr.Column():
-                                gr.Markdown("#### Upload Files")
-                                lora_files = gr.File(
-                                    label="Audio Files (WAV, MP3, FLAC, OGG, OPUS)",
-                                    file_count="multiple",
-                                    file_types=["audio"],
-                                )
-                                lora_upload_btn = gr.Button(
-                                    "Upload & Scan", variant="primary"
-                                )
-
-                            with gr.Column():
-                                gr.Markdown("#### HuggingFace Dataset")
-                                lora_hf_id = gr.Textbox(
-                                    label="Dataset ID",
-                                    placeholder="pedroapfilho/lofi-tracks",
-                                )
-                                lora_hf_token = gr.Textbox(
-                                    label="HF Token (optional, for private repos)",
-                                    type="password",
-                                )
-                                lora_hf_max = gr.Slider(
-                                    minimum=1, maximum=500, value=50, step=1,
-                                    label="Max files",
-                                )
-                                lora_hf_btn = gr.Button(
-                                    "Download & Scan", variant="primary"
-                                )
+                            lora_hf_max = gr.Slider(
+                                minimum=1, maximum=500, value=50, step=1,
+                                label="Batch size",
+                            )
+                            lora_hf_offset = gr.Number(
+                                label="Offset (auto-increments)",
+                                value=0,
+                                precision=0,
+                            )
+                        lora_hf_btn = gr.Button(
+                            "Download Batch & Scan", variant="primary"
+                        )
 
                         lora_source_status = gr.Textbox(
                             label="Status", lines=2, interactive=False
@@ -786,11 +856,16 @@ def create_ui():
                             "Auto-label samples using the LLM, then review and edit metadata."
                         )
 
+                        lora_label_batch_size = gr.Slider(
+                            minimum=1, maximum=500, value=50, step=1,
+                            label="Label batch size (samples per run)",
+                        )
                         lora_label_btn = gr.Button(
-                            "Auto-Label All Samples", variant="primary"
+                            "Label Batch (+ auto-save)",
+                            variant="primary",
                         )
                         lora_label_status = gr.Textbox(
-                            label="Label Status", lines=2, interactive=False
+                            label="Label Status", lines=3, interactive=False
                         )
 
                         lora_review_df = gr.Dataframe(
@@ -800,7 +875,11 @@ def create_ui():
                             wrap=True,
                         )
 
-                        lora_save_btn = gr.Button("Save Edits")
+                        with gr.Row():
+                            lora_save_btn = gr.Button("Save Edits")
+                            lora_save_dataset_btn = gr.Button(
+                                "Save Dataset to JSON", variant="secondary"
+                            )
                         lora_save_status = gr.Textbox(
                             label="Save Status", interactive=False
                         )
@@ -811,8 +890,12 @@ def create_ui():
                             "Encode audio through VAE and text encoders to create training tensors."
                         )
 
+                        lora_preprocess_batch_size = gr.Slider(
+                            minimum=1, maximum=500, value=50, step=1,
+                            label="Preprocess batch size (samples per run)",
+                        )
                         lora_preprocess_btn = gr.Button(
-                            "Preprocess to Tensors", variant="primary"
+                            "Preprocess Batch (+ auto-save)", variant="primary"
                         )
                         lora_preprocess_status = gr.Textbox(
                             label="Preprocess Status", lines=3, interactive=False
@@ -895,23 +978,17 @@ def create_ui():
                 # ---------- Event handlers ----------
 
                 # Data Source
-                lora_upload_btn.click(
-                    fn=lora_upload_and_scan,
-                    inputs=[lora_files, training_state],
-                    outputs=[lora_source_status, training_state],
-                )
-
                 lora_hf_btn.click(
                     fn=lora_download_hf,
-                    inputs=[lora_hf_id, lora_hf_token, lora_hf_max, training_state],
-                    outputs=[lora_source_status, training_state],
+                    inputs=[lora_hf_id, lora_hf_max, lora_hf_offset, training_state],
+                    outputs=[lora_source_status, training_state, lora_hf_offset, lora_progress],
                 )
 
                 # Label & Review
                 lora_label_btn.click(
                     fn=lora_auto_label,
-                    inputs=[training_state],
-                    outputs=[lora_review_df, lora_label_status],
+                    inputs=[lora_label_batch_size, training_state],
+                    outputs=[lora_review_df, lora_label_status, training_state, lora_progress],
                 )
 
                 lora_save_btn.click(
@@ -920,11 +997,17 @@ def create_ui():
                     outputs=[lora_save_status],
                 )
 
+                lora_save_dataset_btn.click(
+                    fn=lora_save_dataset_to_json,
+                    inputs=[training_state],
+                    outputs=[lora_save_status],
+                )
+
                 # Preprocess
                 lora_preprocess_btn.click(
                     fn=lora_preprocess,
-                    inputs=[training_state],
-                    outputs=[lora_preprocess_status],
+                    inputs=[lora_preprocess_batch_size, training_state],
+                    outputs=[lora_preprocess_status, lora_progress],
                 )
 
                 # Train

src/lora_trainer.py

@@ -6,8 +6,10 @@ The actual training pipeline lives in acestep/training/.
 """
 
 import logging
+import os
+import shutil
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Tuple
 
 logger = logging.getLogger(__name__)
 
@@ -17,18 +19,15 @@ AUDIO_SUFFIXES = {".wav", ".mp3", ".flac", ".ogg", ".opus"}
 def download_hf_dataset(
     dataset_id: str,
     max_files: int = 50,
-    hf_token: Optional[str] = None,
+    offset: int = 0,
 ) -> Tuple[str, str]:
     """
     Download a subset of audio files from a HuggingFace dataset repo.
 
-    Lists repo contents first, picks the first N audio files,
-    then downloads them individually to the HF cache.
+    Also pulls dataset.json from the repo if it exists (restoring labels
+    and preprocessed flags from a previous session).
 
-    Args:
-        dataset_id: HuggingFace dataset repo ID (e.g. "pedroapfilho/lofi-tracks")
-        max_files: Maximum number of audio files to download
-        hf_token: Optional HuggingFace token for private repos
+    Uses HF_TOKEN env var for authentication.
 
     Returns:
         Tuple of (output_dir, status_message)
@@ -37,7 +36,7 @@ def download_hf_dataset(
         from huggingface_hub import HfApi, hf_hub_download
 
         api = HfApi()
-        token = hf_token or None
+        token = os.environ.get("HF_TOKEN")
 
         logger.info(f"Listing files in '{dataset_id}'...")
 
@@ -51,7 +50,7 @@ def download_hf_dataset(
         ]
 
         total_available = len(all_files)
-        selected = all_files[:max_files]
+        selected = all_files[offset:offset + max_files]
 
         if not selected:
             return "", f"No audio files found in {dataset_id}"
@@ -76,9 +75,23 @@ def download_hf_dataset(
             if not dest.exists():
                 dest.symlink_to(cached_path)
 
+        # Pull dataset.json from repo if it exists (restores previous session state)
+        try:
+            cached_json = hf_hub_download(
+                repo_id=dataset_id,
+                filename="dataset.json",
+                repo_type="dataset",
+                token=token,
+            )
+            dest_json = output_dir / "dataset.json"
+            shutil.copy2(cached_json, str(dest_json))
+            logger.info("Pulled dataset.json from HF repo")
+        except Exception:
+            logger.info("No dataset.json in HF repo (first session)")
+
         status = (
             f"Downloaded {len(selected)} of {total_available} "
-            f"audio files from {dataset_id}"
+            f"audio files from {dataset_id} (offset {offset})"
         )
         logger.info(status)
         return str(output_dir), status
@@ -91,3 +104,28 @@ def download_hf_dataset(
         msg = f"Failed to download dataset: {e}"
         logger.error(msg)
         return "", msg
+
+
+def upload_dataset_json_to_hf(dataset_id: str, json_path: str) -> str:
+    """Push dataset.json to the HF dataset repo for persistence across sessions."""
+    try:
+        from huggingface_hub import HfApi
+
+        token = os.environ.get("HF_TOKEN")
+        if not token:
+            return "HF_TOKEN not set — skipped HF sync"
+
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=json_path,
+            path_in_repo="dataset.json",
+            repo_id=dataset_id,
+            repo_type="dataset",
+            token=token,
+        )
+        return f"Synced dataset.json to {dataset_id}"
+
+    except Exception as e:
+        msg = f"HF sync failed: {e}"
+        logger.error(msg)
+        return msg