miku-discord/bot/routes/image_generation.py

"""Image generation routes: generate, status, test-detection, view."""

import os
from fastapi import APIRouter
from fastapi.responses import FileResponse, JSONResponse
from utils.logger import get_logger

logger = get_logger('api')

router = APIRouter()


@router.post("/image/generate")
async def manual_image_generation(req: dict):
    """Manually trigger image generation for testing"""
    try:
        prompt = req.get("prompt", "").strip()
        if not prompt:
            return JSONResponse(status_code=400, content={"status": "error", "message": "Prompt is required"})

        from utils.image_generation import generate_image_with_comfyui
        image_path = await generate_image_with_comfyui(prompt)

        if image_path:
            return {"status": "ok", "message": f"Image generated successfully", "image_path": image_path}
        else:
            return JSONResponse(status_code=500, content={"status": "error", "message": "Failed to generate image"})

    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": f"Error: {e}"})


@router.get("/image/status")
async def get_image_generation_status():
    """Get status of image generation system"""
    try:
        from utils.image_generation import check_comfyui_status
        status = await check_comfyui_status()
        return {"status": "ok", **status}

    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": f"Error: {e}"})


@router.post("/image/test-detection")
async def test_image_detection(req: dict):
    """Test the natural language image detection system"""
    try:
        message = req.get("message", "").strip()
        if not message:
            return JSONResponse(status_code=400, content={"status": "error", "message": "Message is required"})

        from utils.image_generation import detect_image_request
        is_image_request, extracted_prompt = await detect_image_request(message)

        return {
            "status": "ok",
            "is_image_request": is_image_request,
            "extracted_prompt": extracted_prompt,
            "original_message": message
        }

    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": f"Error: {e}"})


@router.get("/image/view/{filename}")
async def view_generated_image(filename: str):
    """Serve generated images from ComfyUI output directory"""
    try:
        logger.debug(f"Image view request for: {filename}")

        # Try multiple possible paths for ComfyUI output
        possible_paths = [
            f"/app/ComfyUI/output/{filename}",
            f"/home/koko210Serve/ComfyUI/output/{filename}",
            f"./ComfyUI/output/{filename}",
        ]

        image_path = None
        for path in possible_paths:
            if os.path.exists(path):
                image_path = path
                logger.debug(f"Found image at: {path}")
                break
            else:
                logger.debug(f"Not found at: {path}")

        if not image_path:
            logger.warning(f"Image not found anywhere: {filename}")
            return JSONResponse(status_code=404, content={"status": "error", "message": f"Image not found: {filename}"})

        # Determine content type based on file extension
        ext = filename.lower().split('.')[-1]
        content_type = "image/png"
        if ext == "jpg" or ext == "jpeg":
            content_type = "image/jpeg"
        elif ext == "gif":
            content_type = "image/gif"
        elif ext == "webp":
            content_type = "image/webp"

        logger.info(f"Serving image: {image_path} as {content_type}")
        return FileResponse(image_path, media_type=content_type)

    except Exception as e:
        logger.error(f"Error serving image: {e}")
        return JSONResponse(status_code=500, content={"status": "error", "message": f"Error serving image: {e}"})