dreamtail/main.py

"""
DreamTail - SDXL Image Generation Service

Main application entry point.
"""

import sys
import os
from pathlib import Path

# Add app directory to Python path
sys.path.insert(0, str(Path(__file__).parent))

import asyncio
import logging
import time
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

from api.routes import router
from worker.queue_manager import queue_manager
from worker.generator import generator
from dreamtail_storage.file_manager import file_manager
from dreamtail_storage.cleanup_task import cleanup_task
import config

# Configure logging
logging.basicConfig(
    level=getattr(logging, config.LOG_LEVEL),
    format=config.LOG_FORMAT,
    datefmt=config.LOG_DATE_FORMAT
)

logger = logging.getLogger(__name__)

# Track application start time
start_time = time.time()


async def process_jobs():
    """
    Background worker that processes jobs from the queue.
    """
    logger.info("Job processor started")

    while True:
        try:
            # Get next job from queue (blocks until available)
            job_id = await queue_manager.get_next_job()

            if not job_id:
                continue

            job = queue_manager.get_job(job_id)
            if not job:
                logger.warning(f"Job {job_id} not found in queue")
                continue

            # Mark job as started
            await queue_manager.start_job(job_id)

            logger.info(f"Processing job {job_id}")

            try:
                # Progress callback
                async def update_progress(progress: int):
                    await queue_manager.update_progress(job_id, progress)

                # Resolve face image path if specified
                face_image = job.params.get("face_image")
                if face_image:
                    # Handle "vixy" shortcut for default face
                    if face_image.lower() == "vixy":
                        # Use all faces in the vixy directory
                        vixy_faces = list(config.FACE_REFERENCE_DIR.glob("*.jpg")) + \
                                    list(config.FACE_REFERENCE_DIR.glob("*.png"))
                        if vixy_faces:
                            face_image = [str(f) for f in vixy_faces]
                            logger.info(f"Using {len(face_image)} Vixy reference faces")
                        else:
                            logger.warning("No Vixy faces found, generating without face lock")
                            face_image = None
                    else:
                        # Look for specific face file
                        face_path = config.FACE_REFERENCE_DIR / face_image
                        if face_path.exists():
                            face_image = str(face_path)
                        else:
                            logger.warning(f"Face image {face_image} not found, generating without face lock")
                            face_image = None

                # Generate image
                image = await generator.generate_image(
                    prompt=job.prompt,
                    negative_prompt=job.negative_prompt,
                    width=job.params.get("width", config.DEFAULT_WIDTH),
                    height=job.params.get("height", config.DEFAULT_HEIGHT),
                    num_inference_steps=job.params.get("num_inference_steps", config.DEFAULT_STEPS),
                    guidance_scale=job.params.get("guidance_scale", config.DEFAULT_GUIDANCE_SCALE),
                    seed=job.params.get("seed"),
                    progress_callback=update_progress,
                    face_image=face_image,
                    face_strength=job.params.get("face_strength", config.DEFAULT_FACE_STRENGTH),
                )

                # Save image to disk
                result_path = await file_manager.save_image(job_id, image)

                # Mark job as completed
                await queue_manager.complete_job(job_id, result_path)

                logger.info(f"Job {job_id} completed successfully")

            except Exception as e:
                logger.error(f"Job {job_id} failed: {e}")
                await queue_manager.fail_job(job_id, str(e))

        except asyncio.CancelledError:
            logger.info("Job processor cancelled")
            break
        except Exception as e:
            logger.error(f"Error in job processor: {e}")
            await asyncio.sleep(5)  # Wait before retrying


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Lifespan context manager for startup and shutdown.
    """
    # Startup
    logger.info(f"🎨 Starting {config.APP_NAME} v{config.APP_VERSION}")
    logger.info(f"Storage: {config.STORAGE_DIR}")
    logger.info(f"Models: {config.MODELS_DIR}")

    # Initialize SDXL model
    try:
        await generator.initialize()
    except Exception as e:
        logger.error(f"Failed to initialize SDXL model: {e}")
        logger.warning("Service will start but generation will fail until model is loaded")

    # Start background tasks
    worker_task = asyncio.create_task(process_jobs())
    await cleanup_task.start()

    logger.info(f"✅ {config.APP_NAME} ready on http://{config.API_HOST}:{config.API_PORT}")

    yield

    # Shutdown
    logger.info(f"🛑 Shutting down {config.APP_NAME}...")

    # Cancel worker task
    worker_task.cancel()
    try:
        await worker_task
    except asyncio.CancelledError:
        pass

    # Stop cleanup task
    await cleanup_task.stop()

    logger.info("✅ Shutdown complete")


# Create FastAPI application
app = FastAPI(
    title=config.APP_NAME,
    version=config.APP_VERSION,
    description="SDXL Image Generation Service for Jetson AGX Orin",
    lifespan=lifespan
)

# Add CORS middleware (allow all origins for multi-client support)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include API routes
app.include_router(router)


@app.get("/")
async def root():
    """Root endpoint with service information."""
    return {
        "service": config.APP_NAME,
        "version": config.APP_VERSION,
        "status": "running",
        "model": config.SDXL_MODEL_ID,
        "endpoints": {
            "generate": "POST /generate",
            "status": "GET /status/{job_id}",
            "result": "GET /result/{job_id}",
            "health": "GET /health",
            "models": "GET /models"
        }
    }


if __name__ == "__main__":
    # Run with uvicorn
    uvicorn.run(
        "main:app",
        host=config.API_HOST,
        port=config.API_PORT,
        log_level=config.LOG_LEVEL.lower()
    )