course-agent-od/backend/src/api/image.py

"""
图像生成和分析API
支持文生图和图生图功能
"""

import json
from typing import List, Optional
from fastapi import APIRouter, Depends, HTTPException, status, UploadFile, File, Form
from sqlalchemy.orm import Session

from ..core.database import get_db
from ..core.security import get_current_user
from ..models.user import User
from ..models.generated_image import GeneratedImageRecord
from ..models.image_models import (
    TextToImageRequest, TextToImageResponse, GeneratedImage,
    ImageEditRequest, ImageEditResponse, ImageEditResult,
    ImageVariationRequest, ImageVariationResponse, ImageVariationResult,
    ModelInfo, TemplateInfo, StyleInfo, SizeInfo, EditModeInfo
)
from ..services.text_to_image_service import TextToImageService
from ..services.image_to_image_service import ImageToImageService

router = APIRouter(prefix="/image", tags=["图像生成"])

# 服务实例
text_to_image_service = TextToImageService()
image_to_image_service = ImageToImageService()


def _save_record(db: Session, user_id: int, image_id: str, image_type: str,
                 image_url: str, prompt: str = None, model: str = None,
                 style: str = None, size: str = None, template: str = None,
                 original_filename: str = None, mode: str = None,
                 strength: float = None, metadata: dict = None):
    """保存图像生成记录到数据库"""
    record = GeneratedImageRecord(
        user_id=user_id,
        image_id=image_id,
        image_type=image_type,
        prompt=prompt,
        model=model,
        style=style,
        size=size,
        template=template,
        image_url=image_url,
        original_filename=original_filename,
        mode=mode,
        strength=str(strength) if strength else None,
        generation_meta=json.dumps(metadata or {}, ensure_ascii=False),
    )
    db.add(record)
    db.commit()


@router.post("/text-to-image", response_model=TextToImageResponse)
async def text_to_image(
    request: TextToImageRequest,
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """文生图 - 从文字描述生成图像"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        images_data = await text_to_image_service.generate(
            prompt=request.prompt, model=request.model, template=request.template,
            style=request.style, size=request.size, num_images=request.num_images
        )

        images = [
            GeneratedImage(id=img["id"], url=img["url"], prompt=img["prompt"],
                           model=img["model"], metadata=img["metadata"])
            for img in images_data
        ]

        for img in images_data:
            _save_record(db, user.id, img["id"], "text_to_image", img["url"],
                         prompt=request.prompt, model=request.model,
                         style=request.style, size=request.size,
                         template=request.template, metadata=img.get("metadata"))

        return TextToImageResponse(images=images, total=len(images))

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"文生图失败: {str(e)}")


@router.post("/image-to-image/edit", response_model=ImageEditResponse)
async def image_edit(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    mode: str = Form("optimize"),
    mask: Optional[UploadFile] = File(None),
    strength: float = Form(0.8),
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """图生图 - 图像编辑"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        result_data = await image_to_image_service.edit_image(
            image_file=image, prompt=prompt, mode=mode, mask_file=mask, strength=strength
        )

        result = ImageEditResult(
            id=result_data["id"], url=result_data["url"],
            original_filename=result_data["original_filename"],
            edit_prompt=result_data["edit_prompt"], mode=result_data["mode"],
            metadata=result_data["metadata"]
        )

        _save_record(db, user.id, result_data["id"], "image_edit", result_data["url"],
                     prompt=prompt, model="Qwen-Image-Edit",
                     original_filename=result_data.get("original_filename"),
                     mode=mode, strength=strength, metadata=result_data.get("metadata"))

        return ImageEditResponse(result=result)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"图像编辑失败: {str(e)}")


@router.post("/image-to-image/variations", response_model=ImageVariationResponse)
async def image_variations(
    image: UploadFile = File(...),
    num_variations: int = Form(3),
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """图生图 - 生成图像变体"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        variations_data = await image_to_image_service.create_variation(
            image_file=image, num_variations=num_variations
        )

        variations = [
            ImageVariationResult(id=var["id"], url=var["url"], type=var["type"], metadata=var["metadata"])
            for var in variations_data
        ]

        for var in variations_data:
            _save_record(db, user.id, var["id"], "variation", var["url"],
                         model="Qwen-Image-Edit", metadata=var.get("metadata"))

        return ImageVariationResponse(variations=variations, total=len(variations))

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"生成图像变体失败: {str(e)}")


@router.post("/image-to-image/style-transfer", response_model=ImageEditResponse)
async def style_transfer(
    image: UploadFile = File(...),
    style_prompt: str = Form(...),
    strength: float = Form(0.8),
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """图生图 - 风格转换"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        result_data = await image_to_image_service.style_transfer(
            image_file=image, style_prompt=style_prompt, strength=strength
        )

        result = ImageEditResult(
            id=result_data["id"], url=result_data["url"],
            original_filename=result_data["original_filename"],
            edit_prompt=result_data["edit_prompt"], mode=result_data["mode"],
            metadata=result_data["metadata"]
        )

        _save_record(db, user.id, result_data["id"], "style_transfer", result_data["url"],
                     prompt=style_prompt, model="Qwen-Image-Edit",
                     original_filename=result_data.get("original_filename"),
                     mode="style_transfer", strength=strength, metadata=result_data.get("metadata"))

        return ImageEditResponse(result=result)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"风格转换失败: {str(e)}")


@router.post("/image-to-image/optimize", response_model=ImageEditResponse)
async def optimize_image(
    image: UploadFile = File(...),
    optimization_prompt: str = Form("优化图像质量，增强细节，提高清晰度"),
    strength: float = Form(0.6),
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """图生图 - 图像优化"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        result_data = await image_to_image_service.optimize_image(
            image_file=image, optimization_prompt=optimization_prompt, strength=strength
        )

        result = ImageEditResult(
            id=result_data["id"], url=result_data["url"],
            original_filename=result_data["original_filename"],
            edit_prompt=result_data["edit_prompt"], mode=result_data["mode"],
            metadata=result_data["metadata"]
        )

        _save_record(db, user.id, result_data["id"], "optimize", result_data["url"],
                     prompt=optimization_prompt, model="Qwen-Image-Edit",
                     original_filename=result_data.get("original_filename"),
                     mode="optimize", strength=strength, metadata=result_data.get("metadata"))

        return ImageEditResponse(result=result)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"图像优化失败: {str(e)}")


@router.post("/image-to-image/outpaint", response_model=ImageEditResponse)
async def outpaint_image(
    image: UploadFile = File(...),
    expansion_prompt: str = Form(...),
    strength: float = Form(0.7),
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """图生图 - 图像智能扩展"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        result_data = await image_to_image_service.outpaint_image(
            image_file=image, expansion_prompt=expansion_prompt, strength=strength
        )

        result = ImageEditResult(
            id=result_data["id"], url=result_data["url"],
            original_filename=result_data["original_filename"],
            edit_prompt=result_data["edit_prompt"], mode=result_data["mode"],
            metadata=result_data["metadata"]
        )

        _save_record(db, user.id, result_data["id"], "outpaint", result_data["url"],
                     prompt=expansion_prompt, model="Qwen-Image-Edit",
                     original_filename=result_data.get("original_filename"),
                     mode="outpaint", strength=strength, metadata=result_data.get("metadata"))

        return ImageEditResponse(result=result)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"图像扩展失败: {str(e)}")


# 配置信息端点
@router.get("/models")
async def get_available_models():
    """获取可用的模型列表"""
    models_data = text_to_image_service.get_available_models()
    return {
        "models": [
            ModelInfo(id=model["id"], name=model["name"], description=model["description"])
            for model in models_data
        ]
    }


@router.get("/templates")
async def get_available_templates():
    """获取可用的提示词模板"""
    templates_data = text_to_image_service.get_available_templates()
    return {
        "templates": [
            TemplateInfo(id=template["id"], name=template["name"], description=template["description"])
            for template in templates_data
        ]
    }


@router.get("/styles")
async def get_available_styles():
    """获取可用的风格选项"""
    styles_data = text_to_image_service.get_available_styles()
    return {
        "styles": [
            StyleInfo(id=style["id"], name=style["name"], description=style["description"])
            for style in styles_data
        ]
    }


@router.get("/sizes")
async def get_available_sizes():
    """获取可用的尺寸选项"""
    sizes_data = text_to_image_service.get_available_sizes()
    return {
        "sizes": [
            SizeInfo(id=size["id"], name=size["name"], description=size["description"])
            for size in sizes_data
        ]
    }


@router.get("/edit-modes")
async def get_edit_modes():
    """获取可用的编辑模式"""
    modes_data = image_to_image_service.get_available_modes()
    return {
        "modes": [
            EditModeInfo(id=mode["id"], name=mode["name"], description=mode["description"])
            for mode in modes_data
        ]
    }


@router.get("/style-presets")
async def get_style_presets():
    """获取预设风格选项"""
    presets_data = image_to_image_service.get_style_presets()
    return {
        "presets": [
            StyleInfo(id=preset["id"], name=preset["name"], description=preset["description"])
            for preset in presets_data
        ]
    }


@router.get("/history")
async def get_image_history(
    page: int = 1,
    page_size: int = 20,
    current_user: str = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """获取用户的图像生成历史"""
    try:
        user = db.query(User).filter(User.username == current_user).first()
        if not user:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")

        total = db.query(GeneratedImageRecord).filter(GeneratedImageRecord.user_id == user.id).count()
        records = (
            db.query(GeneratedImageRecord)
            .filter(GeneratedImageRecord.user_id == user.id)
            .order_by(GeneratedImageRecord.created_at.desc())
            .offset((page - 1) * page_size)
            .limit(page_size)
            .all()
        )

        images = []
        for r in records:
            img = {
                "id": r.image_id,
                "url": r.image_url,
                "image_type": r.image_type,
                "prompt": r.prompt,
                "model": r.model,
                "created_at": r.created_at.isoformat() if r.created_at else None,
            }
            if r.style:
                img["style"] = r.style
            if r.size:
                img["size"] = r.size
            if r.template:
                img["template"] = r.template
            if r.mode:
                img["mode"] = r.mode
            if r.original_filename:
                img["original_filename"] = r.original_filename
            images.append(img)

        return {
            "images": images,
            "total": total,
            "page": page,
            "page_size": page_size,
            "has_next": page * page_size < total,
            "has_prev": page > 1
        }

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"获取图像历史失败: {str(e)}")