""" API Routes FastAPI routes for image generation service. """ import logging from fastapi import APIRouter, HTTPException, Response from fastapi.responses import FileResponse import asyncio from api.models import ( GenerateRequest, JobResponse, JobStatus, HealthResponse, ModelsResponse ) from worker.queue_manager import queue_manager from worker.generator import generator from dreamtail_storage.file_manager import file_manager import config logger = logging.getLogger(__name__) router = APIRouter() @router.post("/generate", response_model=JobResponse, status_code=202) async def generate_image(request: GenerateRequest): """ Submit an image generation job. Returns immediately with a job_id. Use /status/{job_id} to check progress. """ try: # Submit job to queue job_id = await queue_manager.submit_job( prompt=request.prompt, negative_prompt=request.negative_prompt, params=request.params.dict() if request.params else {} ) job = queue_manager.get_job(job_id) return JobResponse( job_id=job_id, status=job.status, created_at=job.created_at, message=f"Job queued. Queue position: {queue_manager.get_queue_size()}" ) except asyncio.QueueFull: raise HTTPException( status_code=503, detail=f"Queue is full (max: {config.MAX_QUEUE_SIZE}). Please try again later." ) except Exception as e: logger.error(f"Error submitting job: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.get("/status/{job_id}", response_model=JobStatus) async def get_job_status(job_id: str): """ Get the status of a generation job. Returns job progress, status, and timestamps. """ job = queue_manager.get_job(job_id) if not job: raise HTTPException(status_code=404, detail="Job not found") return JobStatus( job_id=job.job_id, status=job.status, progress=job.progress, created_at=job.created_at, started_at=job.started_at, completed_at=job.completed_at, error=job.error, prompt=job.prompt ) @router.get("/result/{job_id}") async def get_result(job_id: str): """ Download the generated image for a completed job. Returns the image file as PNG. """ job = queue_manager.get_job(job_id) if not job: raise HTTPException(status_code=404, detail="Job not found") if job.status != "completed": raise HTTPException( status_code=400, detail=f"Job is {job.status}, not completed. Check /status/{job_id}" ) # Check if image file exists image_path = file_manager.get_image_path(job_id) if not image_path: raise HTTPException( status_code=404, detail="Image file not found (may have been cleaned up)" ) # Return image file return FileResponse( path=image_path, media_type="image/png", filename=f"{job_id}.png" ) @router.get("/health", response_model=HealthResponse) async def health_check(): """ Health check endpoint. Returns service status and basic statistics. """ import time from main import start_time return HealthResponse( status="healthy" if generator.model_loaded else "unhealthy", version=config.APP_VERSION, model_loaded=generator.model_loaded, queue_size=queue_manager.get_queue_size(), active_jobs=queue_manager.get_active_jobs(), uptime_seconds=time.time() - start_time ) @router.get("/models", response_model=ModelsResponse) async def get_models_info(): """ Get information about loaded models and configuration. """ model_info = generator.get_model_info() return ModelsResponse( base_model=config.SDXL_MODEL_ID, refiner_model=config.SDXL_REFINER_ID if config.USE_REFINER else None, refiner_enabled=config.USE_REFINER, device=model_info.get("device", "unknown"), fp16_enabled=config.USE_FP16 ) @router.get("/storage") async def get_storage_info(): """ Get storage statistics (admin endpoint). """ stats = file_manager.get_storage_stats() return { "total_images": stats["total_images"], "total_size_mb": round(stats["total_size_mb"], 2), "storage_path": stats["storage_path"], "retention_days": config.IMAGE_RETENTION_DAYS }