"""Image generation routes: generate, status, test-detection, view.""" import os from fastapi import APIRouter from fastapi.responses import FileResponse 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 {"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 {"status": "error", "message": "Failed to generate image"} except Exception as e: return {"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 {"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 {"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 {"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 {"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 {"status": "error", "message": f"Error serving image: {e}"}