Rewrite Tab 3 to wire real training pipeline with HF dataset support

Replace placeholder LoRA trainer (hardcoded loss=0.5) with 4-step wizard:
- Data Source: upload audio files or download from HuggingFace Hub
- Label & Review: auto-label via LLM with editable dataframe
- Preprocess: VAE + text encoding to training tensors
- Train: real Fabric-based LoRA training with stop control

app.py

@@ -16,9 +16,12 @@ import spaces
 
 from src.ace_step_engine import ACEStepEngine
 from src.timeline_manager import TimelineManager
-from src.lora_trainer import LoRATrainer
+from src.lora_trainer import download_hf_dataset
 from src.audio_processor import AudioProcessor
 from src.utils import setup_logging, load_config
+from acestep.training.dataset_builder import DatasetBuilder
+from acestep.training.configs import LoRAConfig, TrainingConfig
+from acestep.training.trainer import LoRATrainer as FabricLoRATrainer
 
 # Setup
 logger = setup_logging()
@@ -27,9 +30,13 @@ config = load_config()
 # Lazy initialize components (will be initialized on first use)
 ace_engine = None
 timeline_manager = None
-lora_trainer = None
+dataset_builder = None
 audio_processor = None
 
+# Module-level mutable dict for training stop signal
+# (gr.State is not shared between concurrent Gradio calls)
+_training_control = {"should_stop": False}
+
 def get_ace_engine():
     """Lazy-load ACE-Step engine."""
     global ace_engine
@@ -44,12 +51,12 @@ def get_timeline_manager():
         timeline_manager = TimelineManager(config)
     return timeline_manager
 
-def get_lora_trainer():
-    """Lazy-load LoRA trainer."""
-    global lora_trainer
-    if lora_trainer is None:
-        lora_trainer = LoRATrainer(config)
-    return lora_trainer
+def get_dataset_builder():
+    """Lazy-load dataset builder."""
+    global dataset_builder
+    if dataset_builder is None:
+        dataset_builder = DatasetBuilder()
+    return dataset_builder
 
 def get_audio_processor():
     """Lazy-load audio processor."""
@@ -277,64 +284,246 @@ def timeline_reset(session_state: dict) -> Tuple[None, None, str, dict]:
     return None, None, "Timeline cleared", session_state
 
 
-# ==================== TAB 3: LORA TRAINING ====================
+# ==================== TAB 3: LORA TRAINING STUDIO ====================
+
+DATAFRAME_HEADERS = ["#", "Filename", "Duration", "Lyrics", "Labeled", "BPM", "Key", "Caption"]
+
+
+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_files(files: List[str]) -> str:
-    """Upload and prepare audio files for LoRA training."""
+def lora_upload_and_scan(files, training_state):
+    """Copy uploaded audio files to working dir and scan."""
     try:
-        prepared_files = get_lora_trainer().prepare_dataset(files)
-        return f"✅ Prepared {len(prepared_files)} files for training"
+        if not files:
+            return "No files uploaded", training_state
+
+        import shutil
+
+        work_dir = Path("lora_training") / "uploaded"
+        work_dir.mkdir(parents=True, exist_ok=True)
+
+        for f in files:
+            src = Path(f)
+            shutil.copy2(str(src), str(work_dir / src.name))
+
+        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
+
+    except Exception as e:
+        logger.error(f"Upload scan failed: {e}")
+        return f"Error: {e}", training_state or {}
+
+
+def lora_download_hf(dataset_id, hf_token, 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
+
+        token = hf_token.strip() if hf_token else None
+        output_dir = str(Path("lora_training") / "hf_datasets")
+
+        local_dir, dl_status = download_hf_dataset(
+            dataset_id.strip(), output_dir, hf_token=token
+        )
+
+        if not local_dir:
+            return f"Download failed: {dl_status}", training_state
+
+        builder = get_dataset_builder()
+        samples, scan_status = builder.scan_directory(local_dir)
+
+        training_state = training_state or {}
+        training_state["audio_dir"] = local_dir
+
+        return f"{dl_status} | {scan_status}", training_state
+
     except Exception as e:
-        return f"❌ Error: {str(e)}"
+        logger.error(f"HF download failed: {e}")
+        return f"Error: {e}", training_state or {}
+
 
 @spaces.GPU(duration=300)
+def lora_auto_label(training_state, progress=gr.Progress()):
+    """Auto-label all samples using LLM analysis."""
+    try:
+        builder = get_dataset_builder()
 
-def lora_train(
-    dataset_path: str,
-    model_name: str,
-    learning_rate: float,
-    batch_size: int,
-    num_epochs: int,
-    rank: int,
-    alpha: int,
-    use_existing_lora: bool,
-    existing_lora_path: Optional[str] = None,
-    progress=gr.Progress()
-) -> Tuple[str, str]:
-    """Train LoRA model on uploaded dataset."""
+        if builder.get_sample_count() == 0:
+            return [], "No samples loaded. Upload files or download a dataset first."
+
+        engine = get_ace_engine()
+        if not engine.is_initialized():
+            return [], "ACE-Step engine not initialized. Models may still be loading."
+
+        def progress_callback(msg):
+            progress(0, desc=msg)
+
+        samples, status = builder.label_all_samples(
+            dit_handler=engine.dit_handler,
+            llm_handler=engine.llm_handler,
+            progress_callback=progress_callback,
+        )
+
+        return _build_review_dataframe(), status
+
+    except Exception as e:
+        logger.error(f"Auto-label failed: {e}")
+        return [], f"Error: {e}"
+
+
+def lora_save_edits(df_data, training_state):
+    """Save user edits from the review dataframe back to samples."""
     try:
-        logger.info(f"Starting LoRA training: {model_name}")
-        
-        # Initialize or load LoRA
-        if use_existing_lora and existing_lora_path:
-            lora_trainer.load_lora(existing_lora_path)
-        else:
-            lora_trainer.initialize_lora(rank=rank, alpha=alpha)
-        
-        # Train
-        def progress_callback(step, total_steps, loss):
-            progress((step, total_steps), desc=f"Training (loss: {loss:.4f})")
-        
-        result_path = lora_trainer.train(
-            dataset_path=dataset_path,
-            model_name=model_name,
-            learning_rate=learning_rate,
-            batch_size=batch_size,
-            num_epochs=num_epochs,
-            progress_callback=progress_callback
+        builder = get_dataset_builder()
+
+        if not df_data or len(df_data) == 0:
+            return "No data to save"
+
+        updated = 0
+        for row in df_data:
+            idx = int(row[0])
+            updates = {}
+
+            # Map editable columns back to sample fields
+            bpm_val = row[5]
+            if bpm_val and bpm_val != "-":
+                try:
+                    updates["bpm"] = int(bpm_val)
+                except (ValueError, TypeError):
+                    pass
+
+            key_val = row[6]
+            if key_val and key_val != "-":
+                updates["keyscale"] = str(key_val)
+
+            caption_val = row[7]
+            if caption_val and caption_val != "-":
+                updates["caption"] = str(caption_val)
+
+            if updates:
+                builder.update_sample(idx, **updates)
+                updated += 1
+
+        return f"Updated {updated} samples"
+
+    except Exception as e:
+        logger.error(f"Save edits failed: {e}")
+        return f"Error: {e}"
+
+
+@spaces.GPU(duration=300)
+def lora_preprocess(training_state, progress=gr.Progress()):
+    """Preprocess labeled samples to training tensors."""
+    try:
+        builder = get_dataset_builder()
+
+        if builder.get_labeled_count() == 0:
+            return "No labeled samples. Run auto-label first."
+
+        engine = get_ace_engine()
+        if not engine.is_initialized():
+            return "ACE-Step engine not initialized."
+
+        tensor_dir = str(Path("lora_training") / "tensors")
+
+        def progress_callback(msg):
+            progress(0, desc=msg)
+
+        output_paths, status = builder.preprocess_to_tensors(
+            dit_handler=engine.dit_handler,
+            output_dir=tensor_dir,
+            progress_callback=progress_callback,
         )
-        
-        info = f"✅ Training complete! Model saved to {result_path}"
-        return result_path, info
-        
+
+        training_state = training_state or {}
+        training_state["tensor_dir"] = tensor_dir
+
+        return status
+
     except Exception as e:
-        logger.error(f"LoRA training failed: {e}")
-        return None, f"❌ Error: {str(e)}"
+        logger.error(f"Preprocess failed: {e}")
+        return f"Error: {e}"
+
+
+@spaces.GPU(duration=600)
+def lora_train_real(
+    lr, batch_size, epochs, rank, alpha,
+    grad_accum, model_name, training_state,
+    progress=gr.Progress(),
+):
+    """Train LoRA using the real Fabric-based trainer."""
+    try:
+        training_state = training_state or {}
+        tensor_dir = training_state.get("tensor_dir", "")
 
+        if not tensor_dir or not Path(tensor_dir).exists():
+            return "", "No preprocessed tensors found. Run preprocessing first."
 
-def lora_download(lora_path: str) -> str:
-    """Provide LoRA model for download."""
-    return lora_path if Path(lora_path).exists() else None
+        engine = get_ace_engine()
+        if not engine.is_initialized():
+            return "", "ACE-Step engine not initialized."
+
+        lora_cfg = LoRAConfig(r=int(rank), alpha=int(alpha))
+        output_dir = str(Path("lora_training") / "models" / (model_name or "lora_model"))
+
+        train_cfg = TrainingConfig(
+            learning_rate=float(lr),
+            batch_size=int(batch_size),
+            max_epochs=int(epochs),
+            gradient_accumulation_steps=int(grad_accum),
+            output_dir=output_dir,
+        )
+
+        trainer = FabricLoRATrainer(
+            dit_handler=engine.dit_handler,
+            lora_config=lora_cfg,
+            training_config=train_cfg,
+        )
+
+        _training_control["should_stop"] = False
+        last_msg = ""
+
+        for step, loss, message in trainer.train_from_preprocessed(
+            tensor_dir=tensor_dir,
+            training_state=_training_control,
+        ):
+            last_msg = f"Step {step} | Loss: {loss:.4f} | {message}"
+            progress(0, desc=last_msg)
+
+            if _training_control.get("should_stop"):
+                trainer.stop()
+                last_msg = f"Training stopped at step {step} (loss: {loss:.4f})"
+                break
+
+        final_path = str(Path(output_dir) / "final")
+        return final_path, last_msg
+
+    except Exception as e:
+        logger.error(f"Training failed: {e}")
+        return "", f"Error: {e}"
+
+
+def lora_stop_training():
+    """Signal the training loop to stop."""
+    _training_control["should_stop"] = True
+    return "Stop signal sent. Training will stop after current step."
+
+
+def lora_download_model(model_path):
+    """Return model path for Gradio file download."""
+    if model_path and Path(model_path).exists():
+        return model_path
+    return None
 
 
 # ==================== GRADIO UI ====================
@@ -543,103 +732,219 @@ def create_ui():
                     outputs=[tl_full_audio, tl_timeline_viz, timeline_state, tl_info]
                 )
             
-            # ============ TAB 3: LORA TRAINING ============
+            # ============ TAB 3: LORA TRAINING STUDIO ============
             with gr.Tab("🎓 LoRA Training Studio"):
                 gr.Markdown("""
                 ### Train Custom LoRA Models
-                Upload audio files to train specialized models for voice cloning, style adaptation, etc.
+                Step-by-step wizard: provide audio data, auto-label with LLM, preprocess, and train.
                 """)
-                
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("#### 1. Upload Training Data")
-                        lora_files = gr.File(
-                            label="Audio Files",
-                            file_count="multiple",
-                            file_types=["audio"]
+
+                training_state = gr.State(value={})
+
+                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.")
+
+                        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_btn = gr.Button(
+                                    "Download & Scan", variant="primary"
+                                )
+
+                        lora_source_status = gr.Textbox(
+                            label="Status", lines=2, interactive=False
                         )
-                        lora_upload_btn = gr.Button("📤 Upload & Prepare Dataset")
-                        lora_upload_status = gr.Textbox(label="Upload Status", lines=2)
-                        
-                        gr.Markdown("#### 2. Training Configuration")
-                        lora_dataset_path = gr.Textbox(
-                            label="Dataset Path",
-                            placeholder="Path to prepared dataset"
+
+                    # ---------- Sub-tab 2: Label & Review ----------
+                    with gr.Tab("2. Label & Review"):
+                        gr.Markdown(
+                            "Auto-label samples using the LLM, then review and edit metadata."
                         )
-                        lora_model_name = gr.Textbox(
-                            label="Model Name",
-                            placeholder="my_custom_lora"
+
+                        lora_label_btn = gr.Button(
+                            "Auto-Label All Samples", variant="primary"
                         )
-                        
-                        with gr.Row():
-                            lora_learning_rate = gr.Number(
-                                label="Learning Rate",
-                                value=1e-4
-                            )
-                            lora_batch_size = gr.Slider(
-                                minimum=1, maximum=16, value=4, step=1,
-                                label="Batch Size"
-                            )
-                        
-                        with gr.Row():
-                            lora_num_epochs = gr.Slider(
-                                minimum=1, maximum=100, value=10, step=1,
-                                label="Epochs"
-                            )
-                            lora_rank = gr.Slider(
-                                minimum=4, maximum=128, value=16, step=4,
-                                label="LoRA Rank"
-                            )
-                            lora_alpha = gr.Slider(
-                                minimum=4, maximum=128, value=32, step=4,
-                                label="LoRA Alpha"
-                            )
-                        
-                        lora_use_existing = gr.Checkbox(
-                            label="Continue training from existing LoRA",
-                            value=False
+                        lora_label_status = gr.Textbox(
+                            label="Label Status", lines=2, interactive=False
                         )
-                        lora_existing_path = gr.Textbox(
-                            label="Existing LoRA Path",
-                            placeholder="Path to existing LoRA model"
+
+                        lora_review_df = gr.Dataframe(
+                            headers=DATAFRAME_HEADERS,
+                            label="Sample Review (editable: BPM, Key, Caption)",
+                            interactive=True,
+                            wrap=True,
                         )
-                        
-                        lora_train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
-                    
-                    with gr.Column():
-                        lora_train_status = gr.Textbox(label="Training Status", lines=3)
-                        lora_model_path = gr.Textbox(label="Trained Model Path", lines=1)
-                        lora_download_btn = gr.Button("💾 Download Model")
-                        lora_download_file = gr.File(label="Download")
-                        
-                        gr.Markdown("""
-                        #### Training Tips
-                        - Upload 10+ audio samples for best results
-                        - Keep samples consistent in style/quality
-                        - Higher rank = more capacity but slower training
-                        - Start with 10-20 epochs and adjust
-                        - Use existing LoRA to continue training
-                        """)
-                
-                # Event handlers
+
+                        lora_save_btn = gr.Button("Save Edits")
+                        lora_save_status = gr.Textbox(
+                            label="Save Status", interactive=False
+                        )
+
+                    # ---------- Sub-tab 3: Preprocess ----------
+                    with gr.Tab("3. Preprocess"):
+                        gr.Markdown(
+                            "Encode audio through VAE and text encoders to create training tensors."
+                        )
+
+                        lora_preprocess_btn = gr.Button(
+                            "Preprocess to Tensors", variant="primary"
+                        )
+                        lora_preprocess_status = gr.Textbox(
+                            label="Preprocess Status", lines=3, interactive=False
+                        )
+
+                    # ---------- Sub-tab 4: Train ----------
+                    with gr.Tab("4. Train"):
+                        gr.Markdown("Configure and run LoRA training.")
+
+                        with gr.Row():
+                            with gr.Column():
+                                lora_model_name = gr.Textbox(
+                                    label="Model Name",
+                                    value="my_lora",
+                                    placeholder="my_lora",
+                                )
+
+                                with gr.Row():
+                                    lora_lr = gr.Number(
+                                        label="Learning Rate", value=1e-4
+                                    )
+                                    lora_batch_size = gr.Slider(
+                                        minimum=1, maximum=8, value=1, step=1,
+                                        label="Batch Size",
+                                    )
+
+                                with gr.Row():
+                                    lora_epochs = gr.Slider(
+                                        minimum=1, maximum=500, value=100, step=1,
+                                        label="Epochs",
+                                    )
+                                    lora_grad_accum = gr.Slider(
+                                        minimum=1, maximum=16, value=4, step=1,
+                                        label="Gradient Accumulation",
+                                    )
+
+                                with gr.Row():
+                                    lora_rank = gr.Slider(
+                                        minimum=4, maximum=128, value=8, step=4,
+                                        label="LoRA Rank",
+                                    )
+                                    lora_alpha = gr.Slider(
+                                        minimum=4, maximum=128, value=16, step=4,
+                                        label="LoRA Alpha",
+                                    )
+
+                                with gr.Row():
+                                    lora_train_btn = gr.Button(
+                                        "Start Training",
+                                        variant="primary",
+                                        size="lg",
+                                    )
+                                    lora_stop_btn = gr.Button(
+                                        "Stop Training",
+                                        variant="stop",
+                                        size="lg",
+                                    )
+
+                            with gr.Column():
+                                lora_train_status = gr.Textbox(
+                                    label="Training Status",
+                                    lines=4,
+                                    interactive=False,
+                                )
+                                lora_model_path = gr.Textbox(
+                                    label="Model Path",
+                                    interactive=False,
+                                )
+                                lora_dl_btn = gr.Button("Download Model")
+                                lora_dl_file = gr.File(label="Download")
+
+                                gr.Markdown("""
+                                #### Tips
+                                - Upload 10+ audio samples for best results
+                                - Keep samples consistent in style/quality
+                                - Higher rank = more capacity but slower training
+                                - Default settings (rank=8, lr=1e-4, 100 epochs) are a good starting point
+                                """)
+
+                # ---------- Event handlers ----------
+
+                # Data Source
                 lora_upload_btn.click(
-                    fn=lora_upload_files,
-                    inputs=[lora_files],
-                    outputs=[lora_upload_status]
+                    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, training_state],
+                    outputs=[lora_source_status, training_state],
+                )
+
+                # Label & Review
+                lora_label_btn.click(
+                    fn=lora_auto_label,
+                    inputs=[training_state],
+                    outputs=[lora_review_df, lora_label_status],
+                )
+
+                lora_save_btn.click(
+                    fn=lora_save_edits,
+                    inputs=[lora_review_df, training_state],
+                    outputs=[lora_save_status],
+                )
+
+                # Preprocess
+                lora_preprocess_btn.click(
+                    fn=lora_preprocess,
+                    inputs=[training_state],
+                    outputs=[lora_preprocess_status],
+                )
+
+                # Train
                 lora_train_btn.click(
-                    fn=lora_train,
-                    inputs=[lora_dataset_path, lora_model_name, lora_learning_rate,
-                           lora_batch_size, lora_num_epochs, lora_rank, lora_alpha,
-                           lora_use_existing, lora_existing_path],
-                    outputs=[lora_model_path, lora_train_status]
+                    fn=lora_train_real,
+                    inputs=[
+                        lora_lr, lora_batch_size, lora_epochs,
+                        lora_rank, lora_alpha, lora_grad_accum,
+                        lora_model_name, training_state,
+                    ],
+                    outputs=[lora_model_path, lora_train_status],
                 )
-                
-                lora_download_btn.click(
-                    fn=lora_download,
+
+                lora_stop_btn.click(
+                    fn=lora_stop_training,
+                    inputs=[],
+                    outputs=[lora_train_status],
+                )
+
+                lora_dl_btn.click(
+                    fn=lora_download_model,
                     inputs=[lora_model_path],
-                    outputs=[lora_download_file]
+                    outputs=[lora_dl_file],
                 )
         
         gr.Markdown("""

src/lora_trainer.py

@@ -1,359 +1,69 @@
 """
-LoRA Trainer - Handles LoRA training for custom models
+HuggingFace Dataset Download Utility for LoRA Training Studio.
+
+Provides a helper to download audio datasets from HuggingFace Hub.
+The actual training pipeline lives in acestep/training/.
 """
 
-import torch
-import torchaudio
-from pathlib import Path
 import logging
-from typing import List, Dict, Any, Optional, Callable
-import json
-from datetime import datetime
+from pathlib import Path
+from typing import Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
+AUDIO_EXTENSIONS = ["*.wav", "*.mp3", "*.flac", "*.ogg", "*.opus"]
+
+
+def download_hf_dataset(
+    dataset_id: str,
+    output_dir: str,
+    hf_token: Optional[str] = None,
+) -> Tuple[str, str]:
+    """
+    Download an audio dataset from HuggingFace Hub.
+
+    Uses snapshot_download to fetch only audio files from the repo,
+    skipping non-audio content like READMEs, metadata, etc.
+
+    Args:
+        dataset_id: HuggingFace dataset repo ID (e.g. "pedroapfilho/lofi-tracks")
+        output_dir: Local directory to download into
+        hf_token: Optional HuggingFace token for private repos
+
+    Returns:
+        Tuple of (local_dir, status_message)
+    """
+    try:
+        from huggingface_hub import snapshot_download
+
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        logger.info(f"Downloading dataset '{dataset_id}' to {output_dir}...")
+
+        local_dir = snapshot_download(
+            repo_id=dataset_id,
+            repo_type="dataset",
+            local_dir=str(output_path / dataset_id.replace("/", "_")),
+            token=hf_token or None,
+            allow_patterns=AUDIO_EXTENSIONS,
+        )
+
+        audio_count = sum(
+            1
+            for ext in AUDIO_EXTENSIONS
+            for _ in Path(local_dir).rglob(ext)
+        )
+
+        status = f"Downloaded {audio_count} audio files from {dataset_id}"
+        logger.info(status)
+        return local_dir, status
 
-class LoRATrainer:
-    """Manages LoRA training for ACE-Step model."""
-    
-    def __init__(self, config: Dict[str, Any]):
-        """
-        Initialize LoRA trainer.
-        
-        Args:
-            config: Configuration dictionary
-        """
-        self.config = config
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.training_dir = Path(config.get("training_dir", "lora_training"))
-        self.training_dir.mkdir(exist_ok=True)
-        
-        self.model = None
-        self.lora_config = None
-        
-        logger.info(f"LoRA Trainer initialized on {self.device}")
-    
-    def prepare_dataset(self, audio_files: List[str]) -> List[str]:
-        """
-        Prepare audio files for training.
-        
-        Args:
-            audio_files: List of audio file paths
-            
-        Returns:
-            List of prepared file paths
-        """
-        try:
-            logger.info(f"Preparing {len(audio_files)} files for training...")
-            
-            prepared_dir = self.training_dir / "prepared_data" / datetime.now().strftime("%Y%m%d_%H%M%S")
-            prepared_dir.mkdir(parents=True, exist_ok=True)
-            
-            prepared_files = []
-            
-            for i, file_path in enumerate(audio_files):
-                try:
-                    # Load audio
-                    audio, sr = torchaudio.load(file_path)
-                    
-                    # Resample to target sample rate if needed
-                    target_sr = self.config.get("sample_rate", 44100)
-                    if sr != target_sr:
-                        resampler = torchaudio.transforms.Resample(sr, target_sr)
-                        audio = resampler(audio)
-                    
-                    # Convert to mono if needed (for some training scenarios)
-                    if audio.shape[0] > 1 and self.config.get("force_mono", False):
-                        audio = torch.mean(audio, dim=0, keepdim=True)
-                    
-                    # Normalize
-                    audio = audio / (torch.abs(audio).max() + 1e-8)
-                    
-                    # Split long files into chunks if needed
-                    chunk_duration = self.config.get("chunk_duration", 30)  # seconds
-                    chunk_samples = int(chunk_duration * target_sr)
-                    
-                    if audio.shape[1] > chunk_samples:
-                        # Split into chunks
-                        num_chunks = audio.shape[1] // chunk_samples
-                        for j in range(num_chunks):
-                            start = j * chunk_samples
-                            end = start + chunk_samples
-                            chunk = audio[:, start:end]
-                            
-                            # Save chunk
-                            chunk_path = prepared_dir / f"audio_{i:04d}_chunk_{j:02d}.wav"
-                            torchaudio.save(
-                                str(chunk_path),
-                                chunk,
-                                target_sr,
-                                encoding="PCM_S",
-                                bits_per_sample=16
-                            )
-                            prepared_files.append(str(chunk_path))
-                    else:
-                        # Save as-is
-                        output_path = prepared_dir / f"audio_{i:04d}.wav"
-                        torchaudio.save(
-                            str(output_path),
-                            audio,
-                            target_sr,
-                            encoding="PCM_S",
-                            bits_per_sample=16
-                        )
-                        prepared_files.append(str(output_path))
-                    
-                except Exception as e:
-                    logger.warning(f"Failed to process {file_path}: {e}")
-                    continue
-            
-            # Save dataset metadata
-            metadata = {
-                "num_files": len(prepared_files),
-                "original_files": len(audio_files),
-                "sample_rate": target_sr,
-                "prepared_at": datetime.now().isoformat(),
-                "files": prepared_files
-            }
-            
-            metadata_path = prepared_dir / "metadata.json"
-            with open(metadata_path, 'w') as f:
-                json.dump(metadata, f, indent=2)
-            
-            logger.info(f"✅ Prepared {len(prepared_files)} training files")
-            return prepared_files
-            
-        except Exception as e:
-            logger.error(f"Dataset preparation failed: {e}")
-            raise
-    
-    def initialize_lora(self, rank: int = 16, alpha: int = 32):
-        """
-        Initialize LoRA configuration.
-        
-        Args:
-            rank: LoRA rank
-            alpha: LoRA alpha
-        """
-        try:
-            from peft import LoraConfig, get_peft_model
-            
-            self.lora_config = LoraConfig(
-                r=rank,
-                lora_alpha=alpha,
-                target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],  # Attention layers
-                lora_dropout=0.1,
-                bias="none",
-                task_type="CAUSAL_LM"
-            )
-            
-            logger.info(f"✅ LoRA initialized: rank={rank}, alpha={alpha}")
-            
-        except Exception as e:
-            logger.error(f"LoRA initialization failed: {e}")
-            raise
-    
-    def load_lora(self, lora_path: str):
-        """
-        Load existing LoRA model for continued training.
-        
-        Args:
-            lora_path: Path to LoRA model
-        """
-        try:
-            from peft import PeftModel
-            from transformers import AutoModel
-            
-            # Load base model
-            base_model = AutoModel.from_pretrained(
-                self.config.get("model_path", "ACE-Step/ACE-Step-v1-3.5B"),
-                torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32
-            )
-            
-            # Load with LoRA
-            self.model = PeftModel.from_pretrained(base_model, lora_path)
-            
-            logger.info(f"✅ Loaded LoRA from {lora_path}")
-            
-        except Exception as e:
-            logger.error(f"Failed to load LoRA: {e}")
-            raise
-    
-    def train(
-        self,
-        dataset_path: str,
-        model_name: str,
-        learning_rate: float = 1e-4,
-        batch_size: int = 4,
-        num_epochs: int = 10,
-        progress_callback: Optional[Callable] = None
-    ) -> str:
-        """
-        Train LoRA model.
-        
-        Args:
-            dataset_path: Path to prepared dataset
-            model_name: Name for the trained model
-            learning_rate: Learning rate
-            batch_size: Batch size
-            num_epochs: Number of epochs
-            progress_callback: Optional callback for progress updates
-            
-        Returns:
-            Path to trained model
-        """
-        try:
-            logger.info(f"Starting LoRA training: {model_name}")
-            
-            # Load dataset
-            dataset = self._load_dataset(dataset_path)
-            
-            # Load base model if not already loaded
-            if self.model is None:
-                from transformers import AutoModel
-                from peft import get_peft_model
-                
-                base_model = AutoModel.from_pretrained(
-                    self.config.get("model_path", "ACE-Step/ACE-Step-v1-3.5B"),
-                    torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
-                    device_map="auto"
-                )
-                
-                self.model = get_peft_model(base_model, self.lora_config)
-            
-            self.model.train()
-            
-            # Setup optimizer
-            optimizer = torch.optim.AdamW(
-                self.model.parameters(),
-                lr=learning_rate,
-                weight_decay=0.01
-            )
-            
-            # Training loop
-            total_steps = (len(dataset) // batch_size) * num_epochs
-            step = 0
-            
-            for epoch in range(num_epochs):
-                epoch_loss = 0.0
-                
-                for batch_idx in range(0, len(dataset), batch_size):
-                    batch = dataset[batch_idx:batch_idx + batch_size]
-                    
-                    # Forward pass (simplified - actual implementation would be more complex)
-                    loss = self._training_step(batch)
-                    
-                    # Backward pass
-                    optimizer.zero_grad()
-                    loss.backward()
-                    optimizer.step()
-                    
-                    epoch_loss += loss.item()
-                    step += 1
-                    
-                    # Progress callback
-                    if progress_callback:
-                        progress_callback(step, total_steps, loss.item())
-                
-                avg_loss = epoch_loss / (len(dataset) // batch_size)
-                logger.info(f"Epoch {epoch+1}/{num_epochs} - Loss: {avg_loss:.4f}")
-            
-            # Save trained model
-            output_dir = self.training_dir / "models" / model_name
-            output_dir.mkdir(parents=True, exist_ok=True)
-            
-            self.model.save_pretrained(str(output_dir))
-            
-            # Save training info
-            info = {
-                "model_name": model_name,
-                "learning_rate": learning_rate,
-                "batch_size": batch_size,
-                "num_epochs": num_epochs,
-                "dataset_size": len(dataset),
-                "trained_at": datetime.now().isoformat(),
-                "lora_config": {
-                    "rank": self.lora_config.r,
-                    "alpha": self.lora_config.lora_alpha
-                }
-            }
-            
-            info_path = output_dir / "training_info.json"
-            with open(info_path, 'w') as f:
-                json.dump(info, f, indent=2)
-            
-            logger.info(f"✅ Training complete! Model saved to {output_dir}")
-            return str(output_dir)
-            
-        except Exception as e:
-            logger.error(f"Training failed: {e}")
-            raise
-    
-    def _load_dataset(self, dataset_path: str) -> List[Dict[str, Any]]:
-        """Load prepared dataset."""
-        dataset_path = Path(dataset_path)
-        
-        # Load metadata
-        metadata_path = dataset_path / "metadata.json"
-        if metadata_path.exists():
-            with open(metadata_path, 'r') as f:
-                metadata = json.load(f)
-            files = metadata.get("files", [])
-        else:
-            # Scan directory for audio files
-            files = list(dataset_path.glob("*.wav"))
-        
-        dataset = []
-        for file_path in files:
-            dataset.append({
-                "path": str(file_path),
-                "audio": None  # Lazy loading
-            })
-        
-        return dataset
-    
-    def _training_step(self, batch: List[Dict[str, Any]]) -> torch.Tensor:
-        """
-        Perform single training step.
-        
-        This is a simplified placeholder - actual implementation would:
-        1. Load audio from batch
-        2. Encode to latent space
-        3. Generate predictions
-        4. Calculate loss
-        5. Return loss
-        
-        Args:
-            batch: Training batch
-            
-        Returns:
-            Loss tensor
-        """
-        # Placeholder loss calculation
-        # Actual implementation would process audio through model
-        loss = torch.tensor(0.5, requires_grad=True, device=self.device)
-        return loss
-    
-    def export_for_inference(self, lora_path: str, output_path: str):
-        """
-        Export LoRA model for inference.
-        
-        Args:
-            lora_path: Path to LoRA model
-            output_path: Output path for exported model
-        """
-        try:
-            # Load LoRA
-            self.load_lora(lora_path)
-            
-            # Merge LoRA with base model
-            merged_model = self.model.merge_and_unload()
-            
-            # Save merged model
-            merged_model.save_pretrained(output_path)
-            
-            logger.info(f"✅ Exported model to {output_path}")
-            
-        except Exception as e:
-            logger.error(f"Export failed: {e}")
-            raise
+    except ImportError:
+        msg = "huggingface_hub is not installed. Run: pip install huggingface_hub"
+        logger.error(msg)
+        return "", msg
+    except Exception as e:
+        msg = f"Failed to download dataset: {e}"
+        logger.error(msg)
+        return "", msg