api_server.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RemoveAnything FastAPI服务
使用Nunchaku量化模型进行物体移除推理
"""

import os
import io
import base64
import time
import logging
from pathlib import Path
from typing import Optional

import torch
import numpy as np
import cv2
from PIL import Image
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# 导入推理模块
import sys
sys.path.insert(0, str(Path(__file__).parent))

# 确保可以导入src模块
from src.infer import load_weights, set_seed
from src.data.data_utils import (
    get_bbox_from_mask, expand_bbox, pad_to_square,
    box2squre, crop_back, expand_image_mask
)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(
    title="RemoveAnything API",
    description="基于FLUX Fill + LoRA的物体移除服务（Nunchaku量化加速）",
    version="1.0.0"
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局模型管理
class ModelManager:
    """单例模式管理模型"""
    _instance = None
    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        if not self._initialized:
            self.pipe = None
            self.redux = None
            self.device = None
            self.dtype = None
            self.service_id = "remove_anything_nunchaku_quant"
            self._initialized = True

    def load_models(self, args):
        """加载模型"""
        if self.pipe is not None:
            logger.info("模型已加载，跳过重复加载")
            return True

        logger.info("=" * 60)
        logger.info("开始加载RemoveAnything模型...")
        logger.info("=" * 60)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = torch.bfloat16

        try:
            # 加载模型
            self.pipe, self.redux = load_weights(args, self.device, self.dtype)

            if self.pipe is None or self.redux is None:
                logger.error("模型加载失败")
                return False

            logger.info("=" * 60)
            logger.info("✓ 模型加载成功！")
            logger.info(f"  - Device: {self.device}")
            logger.info(f"  - Dtype: {self.dtype}")
            logger.info(f"  - Nunchaku Quantization: Enabled")
            logger.info(f"  - Service ID: {self.service_id}")
            logger.info("=" * 60)
            return True

        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            import traceback
            traceback.print_exc()
            return False

    def is_ready(self):
        """检查模型是否就绪"""
        return self.pipe is not None and self.redux is not None

# 全局模型管理器
model_manager = ModelManager()

# 请求模型
class RemoveRequest(BaseModel):
    """移除请求"""
    image: str = Field(..., description="原始图像（base64编码）")
    mask: str = Field(..., description="遮罩图像（base64编码）")
    size: int = Field(512, description="处理尺寸", ge=256, le=1024)
    num_inference_steps: int = Field(20, description="推理步数", ge=10, le=100)
    seed: int = Field(42, description="随机种子")
    expansion_ratio: float = Field(2.0, description="遮罩扩展比例", ge=1.0, le=5.0)

class ServiceInfo(BaseModel):
    """服务信息"""
    service_id: str
    name: str
    description: str
    category: str
    version: str
    model_type: str
    quantization: str

# 工具函数
def base64_to_image(base64_str: str) -> Optional[Image.Image]:
    """Base64字符串转PIL图像"""
    try:
        # 处理data URI格式
        if base64_str.startswith('data:image'):
            base64_str = base64_str.split(',', 1)[1]

        image_data = base64.b64decode(base64_str)
        image = Image.open(io.BytesIO(image_data)).convert('RGB')
        return image
    except Exception as e:
        logger.error(f"Base64转图像失败: {e}")
        return None

def image_to_base64(image: Image.Image, format: str = 'PNG') -> str:
    """PIL图像转Base64字符串"""
    try:
        buffer = io.BytesIO()
        image.save(buffer, format=format)
        img_str = base64.b64encode(buffer.getvalue()).decode()
        return f"data:image/{format.lower()};base64,{img_str}"
    except Exception as e:
        logger.error(f"图像转Base64失败: {e}")
        return None

def pil_to_numpy(pil_image: Image.Image) -> np.ndarray:
    """PIL图像转numpy数组（保持RGB格式，与infer.py一致）"""
    return np.array(pil_image)

# API端点
@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "RemoveAnything API",
        "version": "1.0.0",
        "status": "running",
        "model_loaded": model_manager.is_ready()
    }

@app.get("/api/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy" if model_manager.is_ready() else "initializing",
        "model_loaded": model_manager.is_ready(),
        "device": str(model_manager.device) if model_manager.device else "unknown",
        "service_id": model_manager.service_id
    }

@app.get("/api/service/info", response_model=ServiceInfo)
async def get_service_info():
    """获取服务信息"""
    return ServiceInfo(
        service_id=model_manager.service_id,
        name="物体移除（Nunchaku量化）",
        description="基于FLUX Fill + LoRA的智能物体移除，使用Nunchaku INT4量化加速",
        category="image_inpainting",
        version="1.0.0",
        model_type="FLUX.1-Fill-dev + LoRA",
        quantization="Nunchaku INT4 (Transformer + T5)"
    )

@app.post("/api/remove")
async def remove_object(request: RemoveRequest):
    """
    物体移除API

    输入：
    - image: 原始图像（base64）
    - mask: 遮罩图像（base64，白色区域为要移除的部分）
    - size: 处理尺寸（默认512）
    - num_inference_steps: 推理步数（默认20）
    - seed: 随机种子（默认42）
    - expansion_ratio: 遮罩扩展比例（默认2.0）

    输出：
    - result_image: 处理后的图像（base64）
    - processing_time: 处理耗时（秒）
    """
    if not model_manager.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载，请稍后重试")

    start_time = time.time()

    try:
        # 1. 解码图像
        logger.info("解码输入图像...")
        source_image = base64_to_image(request.image)
        mask_image = base64_to_image(request.mask)

        if source_image is None or mask_image is None:
            raise HTTPException(status_code=400, detail="无效的图像数据")

        # 转换为numpy数组（RGB格式，与infer.py保持一致）
        source_rgb = pil_to_numpy(source_image)
        mask_rgb = pil_to_numpy(mask_image)

        # 确保尺寸一致
        if source_rgb.shape[:2] != mask_rgb.shape[:2]:
            logger.info("调整mask尺寸以匹配原图")
            mask_rgb = cv2.resize(mask_rgb, (source_rgb.shape[1], source_rgb.shape[0]))

        # 2. 准备参数
        logger.info(f"开始推理 - size={request.size}, steps={request.num_inference_steps}, seed={request.seed}")

        class Args:
            def __init__(self):
                self.size = request.size
                self.num_inference_steps = request.num_inference_steps
                self.seed = request.seed
                self.expansion_ratio = request.expansion_ratio
                self.output_dir = "/tmp"

        args = Args()
        set_seed(args.seed)

        # 3. 执行推理（不保存文件，直接返回结果）
        # 修改infer_single_image函数以返回图像而不是保存
        from src.data.data_utils import (
            get_bbox_from_mask, expand_bbox, pad_to_square,
            box2squre, crop_back, expand_image_mask
        )

        # 处理参考图像（与infer.py完全一致，使用RGB格式）
        ref_image = source_rgb.copy()
        ref_mask = (mask_rgb > 128).astype(np.uint8)[:, :, 0]

        ref_box_yyxx = get_bbox_from_mask(ref_mask)
        ref_mask_3 = np.stack([ref_mask, ref_mask, ref_mask], -1)
        masked_ref_image = ref_image * ref_mask_3 + np.ones_like(ref_image) * 255 * (1-ref_mask_3)

        y1, y2, x1, x2 = ref_box_yyxx
        masked_ref_image = masked_ref_image[y1:y2, x1:x2, :]
        ref_mask = ref_mask[y1:y2, x1:x2]

        ratio = 1.3
        masked_ref_image, ref_mask = expand_image_mask(masked_ref_image, ref_mask, ratio=ratio)
        masked_ref_image = pad_to_square(masked_ref_image, pad_value=255, random=False)

        # 扩展目标蒙版（与infer.py完全一致）
        tar_mask = (mask_rgb > 128).astype(np.uint8)[:, :, 0]
        kernel = np.ones((7, 7), np.uint8)
        iterations = 2
        tar_mask = cv2.dilate(tar_mask, kernel, iterations=iterations)

        # 处理目标图像（与infer.py完全一致）
        tar_image = source_rgb.copy()
        tar_box_yyxx = get_bbox_from_mask(tar_mask)
        tar_box_yyxx = expand_bbox(tar_mask, tar_box_yyxx, ratio=1.2)
        tar_box_yyxx_crop = expand_bbox(tar_image, tar_box_yyxx, ratio=args.expansion_ratio)
        tar_box_yyxx_crop = box2squre(tar_image, tar_box_yyxx_crop)
        y1, y2, x1, x2 = tar_box_yyxx_crop

        old_tar_image = tar_image.copy()
        tar_image = tar_image[y1:y2, x1:x2, :]
        tar_mask = tar_mask[y1:y2, x1:x2]

        H1, W1 = tar_image.shape[0], tar_image.shape[1]

        tar_mask = pad_to_square(tar_mask, pad_value=0)
        size = (args.size, args.size)
        tar_mask = cv2.resize(tar_mask, size)

        # 提取参考图像特征
        masked_ref_image = cv2.resize(masked_ref_image.astype(np.uint8), size).astype(np.uint8)

        with torch.no_grad():
            # 显式指定RGB模式，确保颜色正确
            pipe_prior_output = model_manager.redux(Image.fromarray(masked_ref_image, mode='RGB'))

        for key, value in pipe_prior_output.items():
            if isinstance(value, torch.Tensor):
                pipe_prior_output[key] = value.to(device=model_manager.device, dtype=model_manager.dtype)

        tar_image = pad_to_square(tar_image, pad_value=255)
        H2, W2 = tar_image.shape[0], tar_image.shape[1]
        tar_image = cv2.resize(tar_image, size)

        diptych_ref_tar = np.concatenate([masked_ref_image, tar_image], axis=1)
        tar_mask = np.stack([tar_mask, tar_mask, tar_mask], -1)
        mask_black = np.ones_like(tar_image) * 0
        mask_diptych = np.concatenate([mask_black, tar_mask], axis=1)

        # 确保数据类型正确并显式指定RGB模式
        diptych_ref_tar = diptych_ref_tar.astype(np.uint8)
        diptych_ref_tar = Image.fromarray(diptych_ref_tar, mode='RGB')

        mask_diptych[mask_diptych == 1] = 255
        mask_diptych = mask_diptych.astype(np.uint8)
        mask_diptych = Image.fromarray(mask_diptych, mode='RGB')

        # 推理
        generator = torch.Generator(model_manager.device).manual_seed(args.seed)

        edited_image = model_manager.pipe(
            image=diptych_ref_tar,
            mask_image=mask_diptych,
            height=mask_diptych.size[1],
            width=mask_diptych.size[0],
            max_sequence_length=512,
            generator=generator,
            num_inference_steps=args.num_inference_steps,
            **pipe_prior_output,
        ).images[0]

        # 裁剪结果
        width, height = edited_image.size
        left = width // 2
        edited_image = edited_image.crop((left, 0, width, height))

        # 放回原始图像
        edited_image_np = np.array(edited_image)
        edited_image_np = crop_back(edited_image_np, old_tar_image, np.array([H1, W1, H2, W2]), np.array(tar_box_yyxx_crop))
        edited_image_np = np.clip(edited_image_np, 0, 255).astype(np.uint8)
        # 显式指定RGB模式，确保输出颜色正确
        edited_image = Image.fromarray(edited_image_np, mode='RGB')

        # 4. 转换为base64
        result_base64 = image_to_base64(edited_image)

        if result_base64 is None:
            raise HTTPException(status_code=500, detail="结果图像编码失败")

        processing_time = time.time() - start_time
        logger.info(f"✓ 推理完成，耗时: {processing_time:.2f}秒")

        return {
            "success": True,
            "result_image": result_base64,
            "processing_time": round(processing_time, 2),
            "parameters": {
                "size": request.size,
                "num_inference_steps": request.num_inference_steps,
                "seed": request.seed,
                "expansion_ratio": request.expansion_ratio
            }
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"推理失败: {e}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"推理失败: {str(e)}")

@app.on_event("startup")
async def startup_event():
    """启动时加载模型"""
    logger.info("=" * 60)
    logger.info("RemoveAnything API Server 启动中...")
    logger.info("=" * 60)

    # 准备参数
    class Args:
        def __init__(self):
            self.flux_fill_path = "./models/FLUX.1-Fill-dev"
            self.lora_weights_path = "./models/lora"
            self.flux_redux_path = "./models/FLUX.1-Redux-dev"
            self.use_nunchaku = True
            self.nunchaku_svdq_path = "./models/FLUX.1-Fill-dev_quant/svdq-int4_r32-flux.1-fill-dev.safetensors"

    args = Args()

    # 加载模型
    success = model_manager.load_models(args)

    if not success:
        logger.error("=" * 60)
        logger.error("⚠️  模型加载失败，服务将以降级模式运行")
        logger.error("=" * 60)
    else:
        logger.info("=" * 60)
        logger.info("✓ RemoveAnything API Server 启动成功！")
        logger.info("=" * 60)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")