src/utils.py

"""
Utility functions for ACE-Step application
"""

import logging
import yaml
from pathlib import Path
from typing import Dict, Any
import sys


def setup_logging(log_level: str = "INFO") -> logging.Logger:
    """
    Setup logging configuration.
    
    Args:
        log_level: Logging level
        
    Returns:
        Logger instance
    """
    # Create logs directory
    log_dir = Path("logs")
    log_dir.mkdir(exist_ok=True)
    
    # Setup logging
    logging.basicConfig(
        level=getattr(logging, log_level),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_dir / "ace_step.log"),
            logging.StreamHandler(sys.stdout)
        ]
    )
    
    logger = logging.getLogger("ace_step")
    logger.info("Logging initialized")
    
    return logger


def load_config(config_path: str = "config.yaml") -> Dict[str, Any]:
    """
    Load configuration from YAML file.
    
    Args:
        config_path: Path to config file
        
    Returns:
        Configuration dictionary
    """
    config_file = Path(config_path)
    
    if config_file.exists():
        with open(config_file, 'r') as f:
            config = yaml.safe_load(f)
    else:
        # Default configuration
        config = {
            "checkpoint_dir": "./checkpoints",
            "dit_model_path": "acestep-v15-turbo",
            "lm_model_path": "acestep-5Hz-lm-1.7B",
            "model_path": "ACE-Step/ACE-Step-v1-3.5B",
            "sample_rate": 44100,
            "output_dir": "outputs",
            "timeline_dir": "timelines",
            "training_dir": "lora_training",
            "chunk_duration": 30,
            "force_mono": False,
            "device": "auto",
            "use_flash_attention": False,
            "offload_to_cpu": False
        }
        
        # Save default config
        with open(config_file, 'w') as f:
            yaml.dump(config, f, default_flow_style=False)
    
    return config


def format_duration(seconds: float) -> str:
    """
    Format duration in seconds to human-readable string.
    
    Args:
        seconds: Duration in seconds
        
    Returns:
        Formatted string (e.g., "2:30")
    """
    minutes = int(seconds // 60)
    secs = int(seconds % 60)
    return f"{minutes}:{secs:02d}"


def validate_audio_file(file_path: str) -> bool:
    """
    Validate audio file.
    
    Args:
        file_path: Path to audio file
        
    Returns:
        True if valid, False otherwise
    """
    import torchaudio
    
    try:
        audio, sr = torchaudio.load(file_path)
        return True
    except:
        return False


def get_audio_info(file_path: str) -> Dict[str, Any]:
    """
    Get audio file information.
    
    Args:
        file_path: Path to audio file
        
    Returns:
        Dictionary with audio info
    """
    import torchaudio
    
    try:
        audio, sr = torchaudio.load(file_path)
        
        return {
            "duration": audio.shape[1] / sr,
            "sample_rate": sr,
            "channels": audio.shape[0],
            "samples": audio.shape[1],
            "format": Path(file_path).suffix
        }
    except Exception as e:
        return {"error": str(e)}


def ensure_directories():
    """Create necessary directories if they don't exist."""
    dirs = ["outputs", "timelines", "lora_training", "logs", "models"]
    
    for dir_name in dirs:
        Path(dir_name).mkdir(exist_ok=True)