audiocraft-ui/main.py

#!/usr/bin/env python3
"""Main entry point for AudioCraft Studio."""

import asyncio
import logging
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

from config.settings import get_settings
from src.core.gpu_manager import GPUMemoryManager
from src.core.model_registry import ModelRegistry
from src.services.generation_service import GenerationService
from src.services.batch_processor import BatchProcessor
from src.services.project_service import ProjectService
from src.storage.database import Database
from src.ui.app import create_app


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("audiocraft.log"),
    ],
)
logger = logging.getLogger(__name__)


async def initialize_services():
    """Initialize all application services."""
    settings = get_settings()

    # Initialize database
    logger.info("Initializing database...")
    db = Database(settings.database_path)
    await db.initialize()

    # Initialize GPU manager
    logger.info("Initializing GPU manager...")
    gpu_manager = GPUMemoryManager(
        device_id=0,
        comfyui_reserve_bytes=int(settings.comfyui_reserve_gb * 1024**3),
    )

    # Initialize model registry
    logger.info("Initializing model registry...")
    model_registry = ModelRegistry(
        gpu_manager=gpu_manager,
        max_loaded=settings.max_loaded_models,
        idle_timeout_seconds=settings.idle_unload_minutes * 60,
    )

    # Initialize services
    logger.info("Initializing services...")
    generation_service = GenerationService(
        model_registry=model_registry,
        gpu_manager=gpu_manager,
        output_dir=settings.output_dir,
    )

    batch_processor = BatchProcessor(
        generation_service=generation_service,
        max_queue_size=settings.max_queue_size,
    )

    project_service = ProjectService(
        db=db,
        output_dir=settings.output_dir,
    )

    return {
        "db": db,
        "gpu_manager": gpu_manager,
        "model_registry": model_registry,
        "generation_service": generation_service,
        "batch_processor": batch_processor,
        "project_service": project_service,
    }


def main():
    """Main entry point."""
    settings = get_settings()

    logger.info("=" * 60)
    logger.info("AudioCraft Studio")
    logger.info("=" * 60)
    logger.info(f"Host: {settings.host}")
    logger.info(f"Gradio Port: {settings.gradio_port}")
    logger.info(f"API Port: {settings.api_port}")
    logger.info(f"Output Dir: {settings.output_dir}")
    logger.info("=" * 60)

    # Initialize services
    logger.info("Initializing services...")

    try:
        services = asyncio.run(initialize_services())
    except Exception as e:
        logger.error(f"Failed to initialize services: {e}")
        logger.warning("Starting in demo mode without backend services")
        services = {}

    # Create and launch app
    logger.info("Creating Gradio application...")
    app = create_app(
        generation_service=services.get("generation_service"),
        batch_processor=services.get("batch_processor"),
        project_service=services.get("project_service"),
        gpu_manager=services.get("gpu_manager"),
        model_registry=services.get("model_registry"),
    )

    # Start batch processor if available
    batch_processor = services.get("batch_processor")
    if batch_processor:
        logger.info("Starting batch processor...")
        asyncio.run(batch_processor.start())

    # Launch the app
    logger.info("Launching application...")
    try:
        app.launch(
            server_name=settings.host,
            server_port=settings.gradio_port,
            share=False,
            show_api=settings.api_enabled,
        )
    except KeyboardInterrupt:
        logger.info("Shutting down...")
    finally:
        # Cleanup
        if batch_processor:
            asyncio.run(batch_processor.stop())
        if "db" in services:
            asyncio.run(services["db"].close())

        logger.info("Shutdown complete")


if __name__ == "__main__":
    main()