import base64 import logging from typing import Callable from threading import Lock from secrets import compare_digest from io import BytesIO import asyncio import concurrent.futures import os import glob from fastapi import FastAPI, Depends, HTTPException from fastapi.security import HTTPBasic, HTTPBasicCredentials from pydantic import BaseModel, Field from PIL import Image import numpy as np from backend_lsnet.inference import process_image_from_pil try: from modules import shared from modules.call_queue import queue_lock as webui_queue_lock IN_WEBUI = True except ImportError: IN_WEBUI = False shared = type('Shared', (), {'cmd_opts': type('CmdOpts', (), {'api_auth': None})()})() webui_queue_lock = None # 设置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def get_available_checkpoints(model_name): """Get available checkpoint files for the model""" models_dir = "models/lsnet" model_dir = os.path.join(models_dir, model_name) if os.path.exists(model_dir): checkpoints = [] for ext in ['*.pth', '*.ckpt', '*.safetensors']: checkpoints.extend(glob.glob(os.path.join(model_dir, ext))) return [os.path.basename(f) for f in checkpoints] return [] def get_available_csv(model_name): """Get available CSV files for the model""" models_dir = "models/lsnet" model_dir = os.path.join(models_dir, model_name) if os.path.exists(model_dir): csv_files = glob.glob(os.path.join(model_dir, "*.csv")) return [os.path.basename(f) for f in csv_files] return [] def get_checkpoint_path(model_name, checkpoint_name): """Get full checkpoint path""" models_dir = "models/lsnet" return os.path.join(models_dir, model_name, checkpoint_name) class InferenceRequest(BaseModel): input_image: str = Field(..., description="Input image as Base64 encoded string") model_name: str = Field('Kaloscope', description="Model name (subfolder in models/lsnet/)") device: str = Field('cuda', description="Device to use") top_k: int = Field(5, ge=1, le=20, description="Number of top predictions") threshold: float = Field(0.0, ge=0.0, le=1.0, description="Probability threshold") class InferenceResponse(BaseModel): results: dict = Field(..., description="Inference results") info: str = Field(..., description="Additional information") class CancelResponse(BaseModel): info: str = Field(..., description="Cancel operation result") class Api: def __init__(self, app: FastAPI, queue_lock: Lock = None, prefix: str = "/lsnet/v1"): self.app = app self.queue_lock = queue_lock or Lock() self.prefix = prefix self.credentials = {} self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) if IN_WEBUI and shared.cmd_opts.api_auth: for auth in shared.cmd_opts.api_auth.split(","): user, password = auth.split(":") self.credentials[user] = password self.add_api_route( "infer", self.endpoint_infer, methods=["POST"], response_model=InferenceResponse, summary="Perform artist style inference", description="Classify or cluster an image using LSNet artist model." ) self.add_api_route( "cancel", self.endpoint_cancel, methods=["POST"], response_model=CancelResponse, summary="Cancel the current inference task", description="Terminates the ongoing inference task." ) def auth(self, creds: HTTPBasicCredentials = Depends(HTTPBasic())): if not self.credentials: return True if creds.username in self.credentials: if compare_digest(creds.password, self.credentials[creds.username]): return True raise HTTPException( status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"} ) def add_api_route(self, path: str, endpoint: Callable, **kwargs): path = f"{self.prefix}/{path}" if self.prefix else path if self.credentials: return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs) return self.app.add_api_route(path, endpoint, **kwargs) def decode_base64_image(self, base64_str: str) -> Image.Image: try: img_data = base64.b64decode(base64_str, validate=True) img = Image.open(BytesIO(img_data)).convert("RGB") return img except base64.binascii.Error: raise HTTPException(400, "Invalid Base64 string format") except Exception as e: raise HTTPException(400, f"Failed to decode image: {str(e)}") async def run_inference(self, image, **kwargs): """Run inference in a separate thread""" loop = asyncio.get_event_loop() try: return await loop.run_in_executor(self.executor, lambda: process_image_from_pil(image, **kwargs)) except Exception as e: logger.error(f"Inference execution failed: {str(e)}") raise async def endpoint_infer(self, req: InferenceRequest): logger.info(f"Received inference request: model_name={req.model_name}") try: with self.queue_lock: input_image = self.decode_base64_image(req.input_image) checkpoints = get_available_checkpoints(req.model_name) if not checkpoints: raise HTTPException(400, f"No checkpoints found for model {req.model_name}") checkpoint_name = checkpoints[0] # use first available checkpoint = get_checkpoint_path(req.model_name, checkpoint_name) if not os.path.exists(checkpoint): raise HTTPException(400, f"Checkpoint not found: {checkpoint}") # Prepare inference arguments csv_files = get_available_csv(req.model_name) class_csv = None if csv_files: class_csv = os.path.join("models/lsnet", req.model_name, csv_files[0]) # use first available infer_args = { "model": "lsnet_xl_artist", # 固定模型架构 "checkpoint": checkpoint, "mode": "classify", # default to classify "device": req.device, "top_k": req.top_k, "threshold": req.threshold, "class_csv": class_csv } # Run inference results = await self.run_inference(input_image, **infer_args) return InferenceResponse(results=results, info="Inference completed successfully") except Exception as e: logger.error(f"Inference failed: {str(e)}") raise HTTPException(500, f"Inference failed: {str(e)}") async def endpoint_cancel(self): # For simplicity, just return a message since inference is quick return CancelResponse(info="No active inference to cancel") def on_app_started(demo, app): """Called when the webui app starts""" queue_lock = webui_queue_lock or Lock() api = Api(app, queue_lock) logger.info("LSNet API routes added to webui")