109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
|
"""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}"}
|