app.py

"""
Caribbean Voices Hackathon - OWSM v3.1 Training & Inference Platform
Complete solution with entity extraction, fine-tuning, and inference

Built with Gradio v5.49.1
Optimized for Hugging Face Spaces with NVIDIA A10G GPUs (Linux/CUDA)
"""
import os
import sys
import time
import torch
from datetime import datetime

# Force flush stdout for logging
sys.stdout.reconfigure(line_buffering=True)

# Timing utility for startup profiling
_start_time = time.time()
_last_checkpoint = _start_time

def log_timing(step_name: str):
    """Log timing information for startup steps"""
    global _last_checkpoint
    current_time = time.time()
    elapsed_since_start = current_time - _start_time
    elapsed_since_last = current_time - _last_checkpoint
    timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
    print(f"[{timestamp}] ⏱️  {step_name}: +{elapsed_since_last:.3f}s (total: {elapsed_since_start:.3f}s)")
    _last_checkpoint = current_time
    return current_time

# Log startup beginning
startup_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
print(f"\n{'='*70}")
print(f"🚀 Startup began at: {startup_timestamp}")
print(f"{'='*70}\n")
log_timing("Application startup")

# All code lives in hf_space folder
# In HF Space, files are in root; locally they're in hf_space/
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
if BASE_DIR.endswith("hf_space"):
    # Local development - go up one level to find data files
    BASE_DIR = os.path.dirname(BASE_DIR)

# Import local modules
sys.path.insert(0, os.path.dirname(__file__))

# Check for Flash Attention (required for optimal performance on A10G GPUs)
def check_flash_attention():
    """Check if Flash Attention is available - required for A10G GPU optimization"""
    log_timing("Checking Flash Attention")
    try:
        import flash_attn
        print("✓ Flash Attention is installed - ESPnet will use optimized attention (A10G GPU)")
        return True
    except ImportError:
        print("⚠ Flash Attention not available - performance will be suboptimal")
        print("  This should not happen on HF Spaces with GPU. Check build logs.")
        return False

# Suppress NLTK download messages (they're informational)
import logging
logging.getLogger("nltk").setLevel(logging.ERROR)

# Check for torchcodec and provide helpful message if missing
def check_torchcodec():
    """Check if torchcodec is installed, attempt to install if missing"""
    log_timing("Checking torchcodec")
    try:
        import torchcodec
        print("✓ torchcodec is installed")
        return True
    except ImportError:
        print("⚠ torchcodec not found - attempting to install...")
        try:
            import subprocess
            # Try to install torchcodec (may not work in all environments)
            result = subprocess.run(
                [sys.executable, "-m", "pip", "install", "--quiet", "torchcodec>=0.1.0,<1.0.0"],
                capture_output=True,
                timeout=60
            )
            if result.returncode == 0:
                print("✓ torchcodec installed successfully")
                return True
            else:
                print(f"⚠ Could not install torchcodec: {result.stderr.decode()[:200]}")
                print("  Note: torchcodec should be in requirements.txt and install during build.")
                print("  Audio decoding may not work until torchcodec is available.")
                return False
        except Exception as e:
            print(f"⚠ Could not install torchcodec automatically: {e}")
            print("  Note: torchcodec should be in requirements.txt and install during Docker build.")
            print("  If audio decoding fails, ensure torchcodec is in requirements.txt")
            return False

# Check torchcodec on startup (non-blocking)
try:
    check_torchcodec()
except Exception as e:
    print(f"⚠ Error checking torchcodec: {e}")
    print("  Continuing anyway - torchcodec should be installed via requirements.txt")

# Import data loader to initialize dataset on startup
log_timing("Importing modules")
from data.loader import load_dataset_from_hf, HF_DATASET_NAME
from ui.interface import create_interface

print("Initializing Caribbean Voices OWSM Platform...")
print(f"Dataset: {HF_DATASET_NAME}")
print("Environment: Hugging Face Space (Linux/NVIDIA A10G GPU)")

# Check CUDA availability
log_timing("Checking CUDA availability")
if torch.cuda.is_available():
    print(f"✓ CUDA available: {torch.cuda.get_device_name(0)} (GPU: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB)")
else:
    print("⚠ CUDA not available - this should not happen on HF Spaces with GPU")

# Check Flash Attention status (required for optimal performance on A10G)
try:
    flash_attn_available = check_flash_attention()
    if not flash_attn_available:
        print("  Note: Flash Attention should be installed during build for A10G GPU optimization")
except Exception as e:
    print(f"⚠ Could not check Flash Attention status: {e}")

# Try loading dataset BEFORE creating interface so status is correct on load
log_timing("Loading dataset from Hugging Face")
try:
    load_dataset_from_hf()
    log_timing("Dataset loaded")
except Exception as e:
    print(f"⚠ Error loading dataset on startup: {e}")
    log_timing("Dataset loading failed")

# Create the Gradio interface
log_timing("Creating Gradio interface")
demo, css_path = create_interface()
log_timing("Gradio interface created")

# Launch app
if __name__ == "__main__":
    total_startup_time = time.time() - _start_time
    print(f"\n{'='*70}")
    print(f"✅ Startup complete in {total_startup_time:.3f} seconds")
    print(f"{'='*70}\n")
    print("Starting Gradio app...")
    log_timing("Launching Gradio server")
    
    # Gradio 6.x: CSS goes in launch() not Blocks()
    launch_kwargs = {"server_name": "0.0.0.0", "server_port": 7860}
    if css_path.exists():
        launch_kwargs["css_paths"] = [css_path]
    
    demo.launch(**launch_kwargs)