This commit is contained in:
2026-05-14 13:30:06 +08:00
parent 974946cee4
commit e43171521d
91 changed files with 10485 additions and 3254 deletions
+1
View File
@@ -0,0 +1 @@
"""HFUT Model Server 应用包(FastAPI 路由与服务)。"""
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from pathlib import Path
# python_server/ 根目录(app/ 的上一级)
APP_ROOT = Path(__file__).resolve().parent.parent
OUTPUT_DIR = APP_ROOT / "outputs"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+13
View File
@@ -0,0 +1,13 @@
from __future__ import annotations
from fastapi import FastAPI
from .routes import animate, depth, inpaint, meta, segment
app = FastAPI(title="HFUT Model Server", version="0.1.0")
app.include_router(meta.router)
app.include_router(depth.router)
app.include_router(segment.router)
app.include_router(inpaint.router)
app.include_router(animate.router)
+1
View File
@@ -0,0 +1 @@
"""HTTP 路由(按领域拆分)。"""
+135
View File
@@ -0,0 +1,135 @@
from __future__ import annotations
import datetime
import shutil
from typing import Any, Dict
from fastapi import APIRouter
from PIL import Image
from ..deps import OUTPUT_DIR
from ..schemas import AnimateRequest, CharacterAnimateRequest
from ..services.animation_frames import (
align_size_to_input,
composite_rgba_on_background,
list_png_frames,
parse_background_color,
remove_background_by_color,
)
from ..services.image_io import b64_to_pil_image, pil_image_to_png_b64
from ..services.predictors import get_animation_predictor
router = APIRouter(tags=["animation"])
@router.post("/animate")
def animate(req: AnimateRequest) -> Dict[str, Any]:
try:
model_name = req.model_name or "animatediff"
predictor = get_animation_predictor(model_name)
result_path = predictor(
prompt=req.prompt,
negative_prompt=req.negative_prompt or "",
num_inference_steps=req.num_inference_steps,
guidance_scale=req.guidance_scale,
width=req.width,
height=req.height,
video_length=req.video_length,
seed=req.seed,
)
out_dir = OUTPUT_DIR / "animation"
out_dir.mkdir(parents=True, exist_ok=True)
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
out_path = out_dir / f"{model_name}_{ts}.gif"
shutil.copy2(result_path, out_path)
return {"success": True, "output_path": str(out_path)}
except Exception as e:
return {"success": False, "error": str(e)}
@router.post("/animate/character_sequence")
def animate_character_sequence(req: CharacterAnimateRequest) -> Dict[str, Any]:
"""
输入透明背景角色 PNG,生成 30fps 角色动画 PNG 序列。
流程:
- 将角色合成到指定纯色背景,作为 AnimateDiff 控制图
- 以与输入相近、且对齐到 8 的分辨率生成 PNG 序列
- 对每帧按背景色做色键剔除,返回透明背景 PNG 序列给前端
"""
try:
model_name = req.model_name or "animatediff"
background_rgb = parse_background_color(req.background_color)
character = b64_to_pil_image(req.image_b64).convert("RGBA")
orig_w, orig_h = character.size
run_w, run_h = align_size_to_input(orig_w, orig_h, req.max_side)
predictor = get_animation_predictor(model_name)
out_dir = OUTPUT_DIR / "animation" / "character_sequence"
out_dir.mkdir(parents=True, exist_ok=True)
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
seq_dir = out_dir / f"{model_name}_{ts}"
seq_dir.mkdir(parents=True, exist_ok=True)
control_path = seq_dir / "control_image.png"
control_image = composite_rgba_on_background(
character,
background_rgb,
size=(run_w, run_h),
)
control_image.save(control_path)
frame_dir = predictor(
prompt=req.prompt,
negative_prompt=req.negative_prompt or "",
num_inference_steps=req.num_inference_steps,
guidance_scale=req.guidance_scale,
width=run_w,
height=run_h,
video_length=req.video_length,
seed=req.seed,
control_image_path=str(control_path),
output_format="png_sequence",
)
frame_paths = list_png_frames(frame_dir)
if not frame_paths:
raise FileNotFoundError(f"未找到 AnimateDiff 输出 PNG 序列: {frame_dir}")
frames: list[dict[str, Any]] = []
processed_dir = seq_dir / "frames"
processed_dir.mkdir(parents=True, exist_ok=True)
for idx, frame_path in enumerate(frame_paths):
frame = remove_background_by_color(
Image.open(frame_path),
background_rgb,
tolerance=req.background_tolerance,
output_size=(orig_w, orig_h),
)
out_path = processed_dir / f"{idx:04d}.png"
frame.save(out_path)
frames.append(
{
"index": idx,
"image_b64": pil_image_to_png_b64(frame),
"path": str(out_path),
}
)
return {
"success": True,
"fps": 30,
"frame_count": len(frames),
"width": orig_w,
"height": orig_h,
"background_color": req.background_color,
"output_dir": str(processed_dir),
"frames": frames,
}
except Exception as e:
return {"success": False, "error": str(e), "frames": []}
+32
View File
@@ -0,0 +1,32 @@
from __future__ import annotations
import numpy as np
from fastapi import APIRouter, HTTPException
from fastapi.responses import Response
from ..schemas import DepthRequest
from ..services.depth_io import depth_to_png16_bytes
from ..services.image_io import b64_to_pil_image
from ..services.predictors import get_depth_predictor
router = APIRouter(tags=["depth"])
@router.post("/depth")
def depth(req: DepthRequest):
"""
计算深度并直接返回二进制 PNG(16-bit 灰度)。
约束:
- 前端不传/不选模型;模型选择写死在后端 config.py
- 成功:HTTP 200 + Content-Type: image/png
- 失败:HTTP 500detail 为错误信息
"""
try:
pil = b64_to_pil_image(req.image_b64).convert("RGB")
predictor = get_depth_predictor()
depth_arr = predictor(pil) # type: ignore[misc]
png_bytes = depth_to_png16_bytes(np.asarray(depth_arr))
return Response(content=png_bytes, media_type="image/png")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) from e
+47
View File
@@ -0,0 +1,47 @@
from __future__ import annotations
from typing import Any, Dict
from fastapi import APIRouter
from ..deps import OUTPUT_DIR
from ..schemas import InpaintRequest
from ..services.image_io import b64_to_pil_image, default_half_mask, pil_image_to_png_b64
from ..services.predictors import get_inpaint_predictor
router = APIRouter(tags=["inpaint"])
@router.post("/inpaint")
def inpaint(req: InpaintRequest) -> Dict[str, Any]:
try:
model_name = req.model_name or "sdxl_inpaint"
pil = b64_to_pil_image(req.image_b64).convert("RGB")
if req.mask_b64:
mask = b64_to_pil_image(req.mask_b64).convert("L")
else:
mask = default_half_mask(pil)
predictor = get_inpaint_predictor(model_name)
out = predictor(
pil,
mask,
req.prompt or "",
req.negative_prompt or "",
strength=req.strength,
max_side=req.max_side,
)
out_dir = OUTPUT_DIR / "inpaint"
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / f"{model_name}_inpaint.png"
out.save(out_path)
return {
"success": True,
"output_path": str(out_path),
"output_image_b64": pil_image_to_png_b64(out),
}
except Exception as e:
return {"success": False, "error": str(e)}
+93
View File
@@ -0,0 +1,93 @@
from __future__ import annotations
from typing import Any, Dict
from fastapi import APIRouter
router = APIRouter(tags=["meta"])
@router.get("/models")
def get_models() -> Dict[str, Any]:
"""
返回一个兼容 Qt 前端的 schema:
{
"models": {
"depth": { "key": { "name": "..."} ... },
"segment": { ... },
"inpaint": { "key": { "name": "...", "params": [...] } ... }
}
}
"""
return {
"models": {
"depth": {
"midas": {"name": "MiDaS (default)"},
"zoedepth_n": {"name": "ZoeDepth (ZoeD_N)"},
"zoedepth_k": {"name": "ZoeDepth (ZoeD_K)"},
"zoedepth_nk": {"name": "ZoeDepth (ZoeD_NK)"},
"depth_anything_v2_vits": {"name": "Depth Anything V2 (vits)"},
"depth_anything_v2_vitb": {"name": "Depth Anything V2 (vitb)"},
"depth_anything_v2_vitl": {"name": "Depth Anything V2 (vitl)"},
"depth_anything_v2_vitg": {"name": "Depth Anything V2 (vitg)"},
"dpt_large": {"name": "DPT (large)"},
"dpt_hybrid": {"name": "DPT (hybrid)"},
"midas_dpt_beit_large_512": {"name": "MiDaS (dpt_beit_large_512)"},
"midas_dpt_swin2_large_384": {"name": "MiDaS (dpt_swin2_large_384)"},
"midas_dpt_swin2_tiny_256": {"name": "MiDaS (dpt_swin2_tiny_256)"},
"midas_dpt_levit_224": {"name": "MiDaS (dpt_levit_224)"},
},
"segment": {
"sam": {"name": "SAM (vit_h)"},
"mask2former_debug": {"name": "SAM (compat mask2former_debug)"},
"mask2former": {"name": "Mask2Former (not implemented)"},
},
"inpaint": {
"copy": {"name": "Copy (no-op)", "params": []},
"sdxl_inpaint": {
"name": "SDXL Inpaint",
"params": [
{"id": "prompt", "label": "提示词", "optional": True},
],
},
"lama": {"name": "LaMa (Erase / Remove Object)", "params": []},
"controlnet": {
"name": "ControlNet Inpaint (canny)",
"params": [
{"id": "prompt", "label": "提示词", "optional": True},
],
},
},
"animation": {
"animatediff": {
"name": "AnimateDiff (Text-to-Video)",
"params": [
{"id": "prompt", "label": "提示词", "optional": False},
{"id": "negative_prompt", "label": "负向提示词", "optional": True},
{"id": "num_inference_steps", "label": "采样步数", "optional": True},
{"id": "guidance_scale", "label": "CFG Scale", "optional": True},
{"id": "width", "label": "宽度", "optional": True},
{"id": "height", "label": "高度", "optional": True},
{"id": "video_length", "label": "帧数", "optional": True},
{"id": "seed", "label": "随机种子", "optional": True},
],
},
"character_sequence": {
"name": "Character PNG Sequence (30 FPS)",
"endpoint": "/animate/character_sequence",
"params": [
{"id": "image_b64", "label": "透明背景角色 PNG", "optional": False},
{"id": "prompt", "label": "动画提示词", "optional": False},
{"id": "background_color", "label": "纯色背景", "optional": True},
{"id": "video_length", "label": "帧数", "optional": True},
{"id": "background_tolerance", "label": "背景剔除阈值", "optional": True},
],
},
},
}
}
@router.get("/health")
def health() -> Dict[str, str]:
return {"status": "ok"}
+84
View File
@@ -0,0 +1,84 @@
from __future__ import annotations
from typing import Any, Dict
import numpy as np
from fastapi import APIRouter
from PIL import Image
from ..deps import OUTPUT_DIR
from ..schemas import SamPromptSegmentRequest, SegmentRequest
from ..services.image_io import b64_to_pil_image
from ..services.predictors import get_seg_predictor
from model.Seg.seg_loader import expand_mask, mask_to_contour_xy, run_sam_prompt
router = APIRouter(tags=["segment"])
@router.post("/segment")
def segment(req: SegmentRequest) -> Dict[str, Any]:
try:
model_name = req.model_name or "sam"
pil = b64_to_pil_image(req.image_b64).convert("RGB")
rgb = np.array(pil, dtype=np.uint8)
predictor = get_seg_predictor(model_name)
label_map = predictor(rgb).astype(np.int32)
out_dir = OUTPUT_DIR / "segment"
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / f"{model_name}_label.png"
Image.fromarray(np.clip(label_map, 0, 255).astype(np.uint8), mode="L").save(out_path)
return {"success": True, "label_path": str(out_path)}
except Exception as e:
return {"success": False, "error": str(e)}
@router.post("/segment/sam_prompt")
def segment_sam_prompt(req: SamPromptSegmentRequest) -> Dict[str, Any]:
"""
交互式 SAM:裁剪图 + 点/框提示,返回掩膜外轮廓点列(裁剪像素坐标)。
"""
try:
pil = b64_to_pil_image(req.image_b64).convert("RGB")
rgb = np.array(pil, dtype=np.uint8)
h, w = rgb.shape[0], rgb.shape[1]
if req.overlay_b64:
ov = b64_to_pil_image(req.overlay_b64)
if ov.size != (w, h):
return {
"success": False,
"error": f"overlay 尺寸 {ov.size} 与 image {w}x{h} 不一致",
"contour": [],
}
if len(req.point_coords) != len(req.point_labels):
return {
"success": False,
"error": "point_coords 与 point_labels 长度不一致",
"contour": [],
}
if len(req.point_coords) < 1:
return {"success": False, "error": "至少需要一个提示点", "contour": []}
pc = np.array(req.point_coords, dtype=np.float32)
if pc.ndim != 2 or pc.shape[1] != 2:
return {"success": False, "error": "point_coords 每项须为 [x,y]", "contour": []}
pl = np.array(req.point_labels, dtype=np.int64)
box = np.array(req.box_xyxy, dtype=np.float32)
mask = run_sam_prompt(rgb, pc, pl, box_xyxy=box)
if not np.any(mask):
return {"success": False, "error": "SAM 未产生有效掩膜", "contour": []}
mask = expand_mask(mask, int(req.expand_px))
contour = mask_to_contour_xy(mask, epsilon_px=2.0)
if len(contour) < 3:
return {"success": False, "error": "轮廓点数不足", "contour": []}
return {"success": True, "contour": contour, "error": None}
except Exception as e:
return {"success": False, "error": str(e), "contour": []}
+93
View File
@@ -0,0 +1,93 @@
from __future__ import annotations
from typing import Optional
from pydantic import BaseModel, Field
class ImageInput(BaseModel):
image_b64: str = Field(..., description="PNG/JPG 编码后的 base64(不含 data: 前缀)")
model_name: Optional[str] = Field(None, description="模型 key(来自 /models")
class DepthRequest(ImageInput):
pass
class SegmentRequest(ImageInput):
pass
class SamPromptSegmentRequest(BaseModel):
image_b64: str = Field(..., description="裁剪后的 RGB 图 base64PNG/JPG")
overlay_b64: Optional[str] = Field(
None,
description="与裁剪同尺寸的标记叠加 PNG base64(可选;当前用于校验尺寸一致)",
)
point_coords: list[list[float]] = Field(
...,
description="裁剪坐标系下的提示点 [[x,y], ...]",
)
point_labels: list[int] = Field(
...,
description="与 point_coords 等长:1=前景,0=背景",
)
box_xyxy: list[float] = Field(
...,
description="裁剪内笔画紧包围盒 [x1,y1,x2,y2](像素)",
min_length=4,
max_length=4,
)
expand_px: int = Field(
2,
ge=0,
le=20,
description="对输出 mask 做轻微膨胀的像素半径(0 表示不扩大;建议 1~3)",
)
class InpaintRequest(ImageInput):
prompt: Optional[str] = Field("", description="补全 prompt")
strength: float = Field(0.8, ge=0.0, le=1.0)
negative_prompt: Optional[str] = Field("", description="负向 prompt")
mask_b64: Optional[str] = Field(None, description="mask PNG base64(可选)")
max_side: int = Field(1024, ge=128, le=2048)
class AnimateRequest(BaseModel):
model_name: Optional[str] = Field(None, description="模型 key(来自 /models")
prompt: str = Field(..., description="文本提示词")
negative_prompt: Optional[str] = Field("", description="负向提示词")
num_inference_steps: int = Field(25, ge=1, le=200)
guidance_scale: float = Field(8.0, ge=0.0, le=30.0)
width: int = Field(512, ge=128, le=2048)
height: int = Field(512, ge=128, le=2048)
video_length: int = Field(16, ge=1, le=128)
seed: int = Field(-1, description="-1 表示随机种子")
class CharacterAnimateRequest(BaseModel):
image_b64: str = Field(..., description="透明背景角色 PNG 的 base64(不含 data: 前缀)")
model_name: Optional[str] = Field(None, description="模型 key(来自 /models")
prompt: str = Field(..., description="角色动画提示词")
negative_prompt: Optional[str] = Field("", description="负向提示词")
background_color: str = Field(
"#00FF00",
description="生成时使用的纯色背景,支持 #RRGGBB 或 R,G,B",
)
background_tolerance: int = Field(
40,
ge=0,
le=255,
description="自动剔除背景的 RGB 色差阈值,越大剔除越激进",
)
num_inference_steps: int = Field(25, ge=1, le=200)
guidance_scale: float = Field(8.0, ge=0.0, le=30.0)
video_length: int = Field(16, ge=1, le=128, description="前端请求生成的帧数")
seed: int = Field(-1, description="-1 表示随机种子")
max_side: int = Field(
768,
ge=128,
le=2048,
description="推理分辨率长边上限;会保持输入比例并对齐到 8 的倍数",
)
+1
View File
@@ -0,0 +1 @@
"""业务服务层(预测器缓存、图像编解码等)。"""
@@ -0,0 +1,115 @@
from __future__ import annotations
from collections import deque
from pathlib import Path
import numpy as np
from PIL import Image
def parse_background_color(value: str) -> tuple[int, int, int]:
text = value.strip()
if text.startswith("#"):
hex_text = text[1:]
if len(hex_text) != 6:
raise ValueError("background_color 使用 #RRGGBB 时必须是 6 位十六进制")
try:
return tuple(int(hex_text[i : i + 2], 16) for i in (0, 2, 4)) # type: ignore[return-value]
except ValueError as e:
raise ValueError(f"background_color 不是有效十六进制颜色: {value}") from e
parts = [p.strip() for p in text.split(",")]
if len(parts) != 3:
raise ValueError("background_color 必须为 #RRGGBB 或 R,G,B")
try:
rgb = tuple(int(p) for p in parts)
except ValueError as e:
raise ValueError(f"background_color 不是有效 RGB 颜色: {value}") from e
if any(c < 0 or c > 255 for c in rgb):
raise ValueError("background_color 的 RGB 分量必须在 0~255")
return rgb # type: ignore[return-value]
def align_size_to_input(width: int, height: int, max_side: int) -> tuple[int, int]:
run_w, run_h = width, height
if max(width, height) > max_side:
scale = max_side / float(max(width, height))
run_w = int(round(width * scale))
run_h = int(round(height * scale))
run_w = max(8, run_w - (run_w % 8))
run_h = max(8, run_h - (run_h % 8))
return run_w, run_h
def composite_rgba_on_background(
image: Image.Image,
background_rgb: tuple[int, int, int],
size: tuple[int, int],
) -> Image.Image:
rgba = image.convert("RGBA")
if rgba.size != size:
rgba = rgba.resize(size, resample=Image.BICUBIC)
bg = Image.new("RGB", size, background_rgb)
alpha = rgba.getchannel("A")
bg.paste(rgba.convert("RGB"), mask=alpha)
return bg
def remove_background_by_color(
image: Image.Image,
background_rgb: tuple[int, int, int],
tolerance: int,
output_size: tuple[int, int] | None = None,
) -> Image.Image:
rgba = image.convert("RGBA")
if output_size is not None and rgba.size != output_size:
rgba = rgba.resize(output_size, resample=Image.BICUBIC)
arr = np.asarray(rgba, dtype=np.uint8).copy()
rgb = arr[:, :, :3].astype(np.int16)
bg = np.array(background_rgb, dtype=np.int16).reshape(1, 1, 3)
diff = np.abs(rgb - bg)
bg_like = np.max(diff, axis=2) <= int(tolerance)
bg_mask = _edge_connected_mask(bg_like)
arr[bg_mask, 3] = 0
return Image.fromarray(arr, mode="RGBA")
def list_png_frames(frame_dir: Path) -> list[Path]:
frames = [p for p in frame_dir.iterdir() if p.is_file() and p.suffix.lower() == ".png"]
return sorted(frames, key=lambda p: p.name)
def _edge_connected_mask(mask: np.ndarray) -> np.ndarray:
h, w = mask.shape
visited = np.zeros((h, w), dtype=bool)
queue: deque[tuple[int, int]] = deque()
for x in range(w):
if mask[0, x]:
queue.append((0, x))
if mask[h - 1, x]:
queue.append((h - 1, x))
for y in range(h):
if mask[y, 0]:
queue.append((y, 0))
if mask[y, w - 1]:
queue.append((y, w - 1))
while queue:
y, x = queue.popleft()
if visited[y, x] or not mask[y, x]:
continue
visited[y, x] = True
if y > 0:
queue.append((y - 1, x))
if y + 1 < h:
queue.append((y + 1, x))
if x > 0:
queue.append((y, x - 1))
if x + 1 < w:
queue.append((y, x + 1))
return visited
+22
View File
@@ -0,0 +1,22 @@
from __future__ import annotations
import io
import numpy as np
from PIL import Image
def depth_to_png16_bytes(depth: np.ndarray) -> bytes:
depth = np.asarray(depth, dtype=np.float32)
dmin = float(depth.min())
dmax = float(depth.max())
if dmax > dmin:
norm = (depth - dmin) / (dmax - dmin)
else:
norm = np.zeros_like(depth, dtype=np.float32)
# 前后端约定:最远=255,最近=0(8-bit
u8 = ((1.0 - norm) * 255.0).clip(0, 255).astype(np.uint8)
img = Image.fromarray(u8, mode="L")
buf = io.BytesIO()
img.save(buf, format="PNG")
return buf.getvalue()
+25
View File
@@ -0,0 +1,25 @@
from __future__ import annotations
import base64
import io
from PIL import Image, ImageDraw
def b64_to_pil_image(b64: str) -> Image.Image:
raw = base64.b64decode(b64)
return Image.open(io.BytesIO(raw))
def pil_image_to_png_b64(img: Image.Image) -> str:
buf = io.BytesIO()
img.save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode("ascii")
def default_half_mask(img: Image.Image) -> Image.Image:
w, h = img.size
mask = Image.new("L", (w, h), 0)
draw = ImageDraw.Draw(mask)
draw.rectangle([w // 2, 0, w, h], fill=255)
return mask
+108
View File
@@ -0,0 +1,108 @@
from __future__ import annotations
from typing import Any
from PIL import Image
from config_loader import load_app_config, get_depth_backend_from_app
from model.Animation.animation_loader import (
UnifiedAnimationConfig,
AnimationBackend,
build_animation_predictor,
)
from model.Depth.depth_loader import UnifiedDepthConfig, DepthBackend, build_depth_predictor
from model.Inpaint.inpaint_loader import UnifiedInpaintConfig, InpaintBackend, build_inpaint_predictor
from model.Seg.seg_loader import UnifiedSegConfig, SegBackend, build_seg_predictor
_depth_predictor = None
_depth_backend: DepthBackend | None = None
_seg_cache: dict[str, Any] = {}
_inpaint_cache: dict[str, Any] = {}
_animation_cache: dict[str, Any] = {}
def ensure_depth_predictor() -> None:
global _depth_predictor, _depth_backend
if _depth_predictor is not None and _depth_backend is not None:
return
app_cfg = load_app_config()
backend_str = get_depth_backend_from_app(app_cfg)
try:
backend = DepthBackend(backend_str)
except Exception as e:
raise ValueError(f"config.py 中 depth.backend 不合法: {backend_str}") from e
_depth_predictor, _depth_backend = build_depth_predictor(UnifiedDepthConfig(backend=backend))
def get_depth_predictor():
ensure_depth_predictor()
return _depth_predictor
def get_seg_predictor(model_name: str):
if model_name == "mask2former_debug":
model_name = "sam"
if model_name in _seg_cache:
return _seg_cache[model_name]
if model_name == "sam":
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.SAM))
_seg_cache[model_name] = pred
return pred
if model_name == "mask2former":
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.MASK2FORMER))
_seg_cache[model_name] = pred
return pred
raise ValueError(f"未知 segment model_name: {model_name}")
def get_inpaint_predictor(model_name: str):
if model_name in _inpaint_cache:
return _inpaint_cache[model_name]
if model_name == "copy":
def _copy(image: Image.Image, *_args, **_kwargs) -> Image.Image:
return image.convert("RGB")
_inpaint_cache[model_name] = _copy
return _copy
if model_name == "sdxl_inpaint":
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.SDXL_INPAINT))
_inpaint_cache[model_name] = pred
return pred
if model_name == "controlnet":
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.CONTROLNET))
_inpaint_cache[model_name] = pred
return pred
if model_name == "lama":
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.LAMA))
_inpaint_cache[model_name] = pred
return pred
raise ValueError(f"未知 inpaint model_name: {model_name}")
def get_animation_predictor(model_name: str):
if model_name in _animation_cache:
return _animation_cache[model_name]
if model_name == "animatediff":
pred, _ = build_animation_predictor(
UnifiedAnimationConfig(backend=AnimationBackend.ANIMATEDIFF)
)
_animation_cache[model_name] = pred
return pred
raise ValueError(f"未知 animation model_name: {model_name}")
+41 -3
View File
@@ -5,7 +5,8 @@ from __future__ import annotations
当前支持:
- SDXL Inpaintdiffusers AutoPipelineForInpainting
- ControlNet占位,暂未统一封装
- ControlNetStable Diffusion + ControlNet Inpaint
- LaMatorchscript big-lama,本地封装)
"""
from dataclasses import dataclass
@@ -26,6 +27,7 @@ from config_loader import (
class InpaintBackend(str, Enum):
SDXL_INPAINT = "sdxl_inpaint"
CONTROLNET = "controlnet"
LAMA = "lama"
@dataclass
@@ -278,11 +280,13 @@ def _make_controlnet_predictor(_: UnifiedInpaintConfig):
image_run = image
mask_run = mask
# control image:使用 canny 边缘作为约束(最通用)
# control image:使用 Canny 边缘作为条件(更通用、更稳定)。
# 备注:当前项目的 ControlNet 权重未必是 inpaint 专用模型,
# 若直接把洞内置零作为条件,容易把输出“拉黑”。先回退到稳定方案。
rgb = np.array(image_run, dtype=np.uint8)
edges = cv2.Canny(rgb, 100, 200)
edges3 = np.stack([edges, edges, edges], axis=-1)
control_image = Image.fromarray(edges3)
control_image = Image.fromarray(edges3, mode="RGB")
if device == "cuda":
try:
@@ -312,6 +316,37 @@ def _make_controlnet_predictor(_: UnifiedInpaintConfig):
return _predict
def _make_lama_predictor(cfg: UnifiedInpaintConfig):
import torch
app_cfg = load_app_config()
device = cfg.device
if device is None:
device = app_cfg.inpaint.device
device = "cuda" if device.startswith("cuda") and torch.cuda.is_available() else "cpu"
from .lama_torchscript import LaMaTorchscript
lama = LaMaTorchscript(device=device)
def _predict(
image: Image.Image,
mask: Image.Image,
_prompt: str,
_negative_prompt: str = "",
strength: float = 0.8,
guidance_scale: float = 7.5,
num_inference_steps: int = 30,
max_side: int = 1024,
**_kwargs,
) -> Image.Image:
# LaMa 不使用 prompt/negative_prompt,也不使用 strength/steps 等扩散参数
_ = (strength, guidance_scale, num_inference_steps, max_side)
return lama(image, mask)
return _predict
def build_inpaint_predictor(
cfg: UnifiedInpaintConfig | None = None,
) -> tuple[Callable[..., Image.Image], InpaintBackend]:
@@ -326,6 +361,9 @@ def build_inpaint_predictor(
if cfg.backend == InpaintBackend.CONTROLNET:
return _make_controlnet_predictor(cfg), InpaintBackend.CONTROLNET
if cfg.backend == InpaintBackend.LAMA:
return _make_lama_predictor(cfg), InpaintBackend.LAMA
raise ValueError(f"不支持的补全后端: {cfg.backend}")
@@ -0,0 +1,105 @@
from __future__ import annotations
import os
from pathlib import Path
from typing import Optional
import numpy as np
from PIL import Image
DEFAULT_LAMA_URL = os.environ.get(
"LAMA_MODEL_URL",
"https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
)
def _default_cache_path() -> Path:
# 放在 python_server/outputs/cache 下,避免污染用户家目录
here = Path(__file__).resolve()
cache_dir = here.parents[2] / "outputs" / "cache"
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir / "big-lama.pt"
def _download(url: str, dst: Path) -> None:
import requests
dst.parent.mkdir(parents=True, exist_ok=True)
resp = requests.get(url, stream=True, timeout=120)
resp.raise_for_status()
tmp = dst.with_suffix(dst.suffix + ".tmp")
with tmp.open("wb") as f:
for chunk in resp.iter_content(chunk_size=1024 * 1024):
if not chunk:
continue
f.write(chunk)
tmp.replace(dst)
def _pad_to_mod8(img: np.ndarray) -> tuple[np.ndarray, tuple[int, int]]:
h, w = img.shape[:2]
pad_h = (8 - (h % 8)) % 8
pad_w = (8 - (w % 8)) % 8
if pad_h == 0 and pad_w == 0:
return img, (0, 0)
out = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
return out, (pad_h, pad_w)
def _unpad(img: np.ndarray, pad_hw: tuple[int, int]) -> np.ndarray:
pad_h, pad_w = pad_hw
if pad_h == 0 and pad_w == 0:
return img
h, w = img.shape[:2]
return img[: h - pad_h, : w - pad_w]
class LaMaTorchscript:
"""
轻量 LaMabig-lama.pt torchscript)推理封装。
- 输入:RGB 图 + L mask(白=需要修补)
- 输出:RGB 图(同尺寸)
"""
def __init__(self, device: str = "cpu", model_path: Optional[str] = None, model_url: Optional[str] = None):
import torch
self.device = device
self.model_url = model_url or DEFAULT_LAMA_URL
p = Path(model_path) if model_path else _default_cache_path()
if not p.exists():
_download(self.model_url, p)
self.model_path = p
self.model = torch.jit.load(str(self.model_path), map_location=device).eval()
@staticmethod
def _norm_img_u8(rgb: np.ndarray) -> np.ndarray:
# [H,W,3] uint8 -> float32 [0,1]
return (rgb.astype(np.float32) / 255.0).clip(0.0, 1.0)
def __call__(self, image: Image.Image, mask: Image.Image) -> Image.Image:
import torch
img = np.asarray(image.convert("RGB"), dtype=np.uint8)
m = np.asarray(mask.convert("L"), dtype=np.uint8)
m = (m >= 128).astype(np.float32) # 1=hole
img_f = self._norm_img_u8(img)
# pad to multiple of 8
img_pad, pad_hw = _pad_to_mod8(img_f)
m_pad, _ = _pad_to_mod8(np.repeat(m[:, :, None], 3, axis=2))
m1 = m_pad[:, :, 0] # [H,W]
# torch: [1,C,H,W]
it = torch.from_numpy(img_pad).permute(2, 0, 1).unsqueeze(0).to(self.device)
mt = torch.from_numpy(m1).unsqueeze(0).unsqueeze(0).to(self.device)
with torch.no_grad():
out = self.model(it, mt)
out_np = out[0].permute(1, 2, 0).detach().cpu().numpy()
out_np = np.clip(out_np * 255.0, 0.0, 255.0).astype(np.uint8)
out_np = _unpad(out_np, pad_hw)
return Image.fromarray(out_np, mode="RGB")
+37
View File
@@ -95,6 +95,43 @@ def run_sam_prompt(
return m
def expand_mask(mask_bool: np.ndarray, expand_px: int = 0) -> np.ndarray:
"""
轻微扩大二值 mask(用于消除边界过软/过细带来的“漏边”观感)。
- expand_px <= 0: 原样返回
- expand_px > 0: 做 dilation(膨胀)expand_px 像素级别
"""
mask_bool = np.asarray(mask_bool, dtype=bool)
if expand_px <= 0:
return mask_bool
# 优先用 OpenCV(通常更快)
try:
import cv2 # type: ignore[import]
u8 = (mask_bool.astype(np.uint8) * 255)
k = int(expand_px) * 2 + 1
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
out = cv2.dilate(u8, kernel, iterations=1)
return out > 127
except Exception:
pass
# 回退到 scipy(若存在)
try:
from scipy.ndimage import binary_dilation # type: ignore[import]
# 使用近似圆形结构元素
r = int(expand_px)
yy, xx = np.ogrid[-r : r + 1, -r : r + 1]
selem = (xx * xx + yy * yy) <= (r * r)
return binary_dilation(mask_bool, structure=selem)
except Exception:
# 再回退:无依赖时就不扩大(保持可用性)
return mask_bool
def mask_to_contour_xy(
mask_bool: np.ndarray,
epsilon_px: float = 2.0,
+82
View File
@@ -0,0 +1,82 @@
accelerate==1.13.0
annotated-doc==0.0.4
annotated-types==0.7.0
antlr4-python3-runtime==4.9.3
anyio==4.12.1
certifi==2026.2.25
charset-normalizer==3.4.5
click==8.3.1
cuda-bindings==12.9.4
cuda-pathfinder==1.4.1
diffusers==0.37.0
einops==0.8.2
fastapi==0.135.1
filelock==3.25.0
fsspec==2026.2.0
h11==0.16.0
hf-xet==1.3.2
httpcore==1.0.9
httptools==0.7.1
httpx==0.28.1
huggingface_hub==1.6.0
idna==3.11
ImageIO==2.37.3
importlib_metadata==8.7.1
Jinja2==3.1.6
markdown-it-py==4.0.0
MarkupSafe==3.0.3
mdurl==0.1.2
mpmath==1.3.0
networkx==3.6.1
numpy==2.4.3
nvidia-cublas-cu12==12.8.4.1
nvidia-cuda-cupti-cu12==12.8.90
nvidia-cuda-nvrtc-cu12==12.8.93
nvidia-cuda-runtime-cu12==12.8.90
nvidia-cudnn-cu12==9.10.2.21
nvidia-cufft-cu12==11.3.3.83
nvidia-cufile-cu12==1.13.1.3
nvidia-curand-cu12==10.3.9.90
nvidia-cusolver-cu12==11.7.3.90
nvidia-cusparse-cu12==12.5.8.93
nvidia-cusparselt-cu12==0.7.1
nvidia-nccl-cu12==2.27.5
nvidia-nvjitlink-cu12==12.8.93
nvidia-nvshmem-cu12==3.4.5
nvidia-nvtx-cu12==12.8.90
omegaconf==2.3.0
opencv-python==4.13.0.92
packaging==26.0
pillow==12.1.1
psutil==7.2.2
pydantic==2.12.5
pydantic_core==2.41.5
Pygments==2.19.2
python-dotenv==1.2.2
PyYAML==6.0.3
regex==2026.2.28
requests==2.32.5
rich==14.3.3
safetensors==0.7.0
scipy==1.17.1
setuptools==82.0.0
shellingham==1.5.4
starlette==0.52.1
sympy==1.14.0
timm==0.6.13
tokenizers==0.22.2
torch==2.10.0
torchvision==0.25.0
tqdm==4.67.3
transformers==5.3.0
triton==3.6.0
typer==0.24.1
typing-inspection==0.4.2
typing_extensions==4.15.0
urllib3==2.6.3
uvicorn==0.41.0
uvloop==0.22.1
watchfiles==1.1.1
websockets==16.0
xformers==0.0.35
zipp==3.23.0
+7 -477
View File
@@ -1,488 +1,18 @@
"""
source .venv/bin/activate # 仓库根 hfut-bishe/.venv
cd python_server && uvicorn app.main:app --host 0.0.0.0 --port 8000
"""
from __future__ import annotations
import base64
import datetime
import io
import os
import shutil
from dataclasses import asdict
from pathlib import Path
from typing import Any, Dict, Optional
import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel, Field
from PIL import Image, ImageDraw
from config_loader import load_app_config, get_depth_backend_from_app
from model.Depth.depth_loader import UnifiedDepthConfig, DepthBackend, build_depth_predictor
from model.Seg.seg_loader import (
UnifiedSegConfig,
SegBackend,
build_seg_predictor,
mask_to_contour_xy,
run_sam_prompt,
)
from model.Inpaint.inpaint_loader import UnifiedInpaintConfig, InpaintBackend, build_inpaint_predictor
from model.Animation.animation_loader import (
UnifiedAnimationConfig,
AnimationBackend,
build_animation_predictor,
)
APP_ROOT = Path(__file__).resolve().parent
OUTPUT_DIR = APP_ROOT / "outputs"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
app = FastAPI(title="HFUT Model Server", version="0.1.0")
class ImageInput(BaseModel):
image_b64: str = Field(..., description="PNG/JPG 编码后的 base64(不含 data: 前缀)")
model_name: Optional[str] = Field(None, description="模型 key(来自 /models")
class DepthRequest(ImageInput):
pass
class SegmentRequest(ImageInput):
pass
class SamPromptSegmentRequest(BaseModel):
image_b64: str = Field(..., description="裁剪后的 RGB 图 base64PNG/JPG")
overlay_b64: Optional[str] = Field(
None,
description="与裁剪同尺寸的标记叠加 PNG base64(可选;当前用于校验尺寸一致)",
)
point_coords: list[list[float]] = Field(
...,
description="裁剪坐标系下的提示点 [[x,y], ...]",
)
point_labels: list[int] = Field(
...,
description="与 point_coords 等长:1=前景,0=背景",
)
box_xyxy: list[float] = Field(
...,
description="裁剪内笔画紧包围盒 [x1,y1,x2,y2](像素)",
min_length=4,
max_length=4,
)
class InpaintRequest(ImageInput):
prompt: Optional[str] = Field("", description="补全 prompt")
strength: float = Field(0.8, ge=0.0, le=1.0)
negative_prompt: Optional[str] = Field("", description="负向 prompt")
# 可选 mask(白色区域为重绘)
mask_b64: Optional[str] = Field(None, description="mask PNG base64(可选)")
# 推理缩放上限(避免 OOM
max_side: int = Field(1024, ge=128, le=2048)
class AnimateRequest(BaseModel):
model_name: Optional[str] = Field(None, description="模型 key(来自 /models")
prompt: str = Field(..., description="文本提示词")
negative_prompt: Optional[str] = Field("", description="负向提示词")
num_inference_steps: int = Field(25, ge=1, le=200)
guidance_scale: float = Field(8.0, ge=0.0, le=30.0)
width: int = Field(512, ge=128, le=2048)
height: int = Field(512, ge=128, le=2048)
video_length: int = Field(16, ge=1, le=128)
seed: int = Field(-1, description="-1 表示随机种子")
def _b64_to_pil_image(b64: str) -> Image.Image:
raw = base64.b64decode(b64)
return Image.open(io.BytesIO(raw))
def _pil_image_to_png_b64(img: Image.Image) -> str:
buf = io.BytesIO()
img.save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode("ascii")
def _depth_to_png16_b64(depth: np.ndarray) -> str:
depth = np.asarray(depth, dtype=np.float32)
dmin = float(depth.min())
dmax = float(depth.max())
if dmax > dmin:
norm = (depth - dmin) / (dmax - dmin)
else:
norm = np.zeros_like(depth, dtype=np.float32)
u16 = (norm * 65535.0).clip(0, 65535).astype(np.uint16)
img = Image.fromarray(u16, mode="I;16")
return _pil_image_to_png_b64(img)
def _depth_to_png16_bytes(depth: np.ndarray) -> bytes:
depth = np.asarray(depth, dtype=np.float32)
dmin = float(depth.min())
dmax = float(depth.max())
if dmax > dmin:
norm = (depth - dmin) / (dmax - dmin)
else:
norm = np.zeros_like(depth, dtype=np.float32)
# 前后端约定:最远=255,最近=0(8-bit
u8 = ((1.0 - norm) * 255.0).clip(0, 255).astype(np.uint8)
img = Image.fromarray(u8, mode="L")
buf = io.BytesIO()
img.save(buf, format="PNG")
return buf.getvalue()
def _default_half_mask(img: Image.Image) -> Image.Image:
w, h = img.size
mask = Image.new("L", (w, h), 0)
draw = ImageDraw.Draw(mask)
draw.rectangle([w // 2, 0, w, h], fill=255)
return mask
# -----------------------------
# /models(给前端/GUI 使用)
# -----------------------------
@app.get("/models")
def get_models() -> Dict[str, Any]:
"""
返回一个兼容 Qt 前端的 schema:
{
"models": {
"depth": { "key": { "name": "..."} ... },
"segment": { ... },
"inpaint": { "key": { "name": "...", "params": [...] } ... }
}
}
"""
return {
"models": {
"depth": {
# 兼容旧配置默认值
"midas": {"name": "MiDaS (default)"},
"zoedepth_n": {"name": "ZoeDepth (ZoeD_N)"},
"zoedepth_k": {"name": "ZoeDepth (ZoeD_K)"},
"zoedepth_nk": {"name": "ZoeDepth (ZoeD_NK)"},
"depth_anything_v2_vits": {"name": "Depth Anything V2 (vits)"},
"depth_anything_v2_vitb": {"name": "Depth Anything V2 (vitb)"},
"depth_anything_v2_vitl": {"name": "Depth Anything V2 (vitl)"},
"depth_anything_v2_vitg": {"name": "Depth Anything V2 (vitg)"},
"dpt_large": {"name": "DPT (large)"},
"dpt_hybrid": {"name": "DPT (hybrid)"},
"midas_dpt_beit_large_512": {"name": "MiDaS (dpt_beit_large_512)"},
"midas_dpt_swin2_large_384": {"name": "MiDaS (dpt_swin2_large_384)"},
"midas_dpt_swin2_tiny_256": {"name": "MiDaS (dpt_swin2_tiny_256)"},
"midas_dpt_levit_224": {"name": "MiDaS (dpt_levit_224)"},
},
"segment": {
"sam": {"name": "SAM (vit_h)"},
# 兼容旧配置默认值
"mask2former_debug": {"name": "SAM (compat mask2former_debug)"},
"mask2former": {"name": "Mask2Former (not implemented)"},
},
"inpaint": {
# 兼容旧配置默认值:copy 表示不做补全
"copy": {"name": "Copy (no-op)", "params": []},
"sdxl_inpaint": {
"name": "SDXL Inpaint",
"params": [
{"id": "prompt", "label": "提示词", "optional": True},
],
},
"controlnet": {
"name": "ControlNet Inpaint (canny)",
"params": [
{"id": "prompt", "label": "提示词", "optional": True},
],
},
},
"animation": {
"animatediff": {
"name": "AnimateDiff (Text-to-Video)",
"params": [
{"id": "prompt", "label": "提示词", "optional": False},
{"id": "negative_prompt", "label": "负向提示词", "optional": True},
{"id": "num_inference_steps", "label": "采样步数", "optional": True},
{"id": "guidance_scale", "label": "CFG Scale", "optional": True},
{"id": "width", "label": "宽度", "optional": True},
{"id": "height", "label": "高度", "optional": True},
{"id": "video_length", "label": "帧数", "optional": True},
{"id": "seed", "label": "随机种子", "optional": True},
],
},
},
}
}
# -----------------------------
# Depth
# -----------------------------
_depth_predictor = None
_depth_backend: DepthBackend | None = None
def _ensure_depth_predictor() -> None:
global _depth_predictor, _depth_backend
if _depth_predictor is not None and _depth_backend is not None:
return
app_cfg = load_app_config()
backend_str = get_depth_backend_from_app(app_cfg)
try:
backend = DepthBackend(backend_str)
except Exception as e:
raise ValueError(f"config.py 中 depth.backend 不合法: {backend_str}") from e
_depth_predictor, _depth_backend = build_depth_predictor(UnifiedDepthConfig(backend=backend))
@app.post("/depth")
def depth(req: DepthRequest):
"""
计算深度并直接返回二进制 PNG(16-bit 灰度)。
约束:
- 前端不传/不选模型;模型选择写死在后端 config.py
- 成功:HTTP 200 + Content-Type: image/png
- 失败:HTTP 500detail 为错误信息
"""
try:
_ensure_depth_predictor()
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
depth_arr = _depth_predictor(pil) # type: ignore[misc]
png_bytes = _depth_to_png16_bytes(np.asarray(depth_arr))
return Response(content=png_bytes, media_type="image/png")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# -----------------------------
# Segment
# -----------------------------
_seg_cache: Dict[str, Any] = {}
def _get_seg_predictor(model_name: str):
if model_name in _seg_cache:
return _seg_cache[model_name]
# 兼容旧默认 key
if model_name == "mask2former_debug":
model_name = "sam"
if model_name == "sam":
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.SAM))
_seg_cache[model_name] = pred
return pred
if model_name == "mask2former":
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.MASK2FORMER))
_seg_cache[model_name] = pred
return pred
raise ValueError(f"未知 segment model_name: {model_name}")
@app.post("/segment")
def segment(req: SegmentRequest) -> Dict[str, Any]:
try:
model_name = req.model_name or "sam"
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
rgb = np.array(pil, dtype=np.uint8)
predictor = _get_seg_predictor(model_name)
label_map = predictor(rgb).astype(np.int32)
out_dir = OUTPUT_DIR / "segment"
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / f"{model_name}_label.png"
# 保存为 8-bit 灰度(若 label 超过 255 会截断;当前 SAM 通常不会太大)
Image.fromarray(np.clip(label_map, 0, 255).astype(np.uint8), mode="L").save(out_path)
return {"success": True, "label_path": str(out_path)}
except Exception as e:
return {"success": False, "error": str(e)}
@app.post("/segment/sam_prompt")
def segment_sam_prompt(req: SamPromptSegmentRequest) -> Dict[str, Any]:
"""
交互式 SAM:裁剪图 + 点/框提示,返回掩膜外轮廓点列(裁剪像素坐标)。
"""
try:
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
rgb = np.array(pil, dtype=np.uint8)
h, w = rgb.shape[0], rgb.shape[1]
if req.overlay_b64:
ov = _b64_to_pil_image(req.overlay_b64)
if ov.size != (w, h):
return {
"success": False,
"error": f"overlay 尺寸 {ov.size} 与 image {w}x{h} 不一致",
"contour": [],
}
if len(req.point_coords) != len(req.point_labels):
return {
"success": False,
"error": "point_coords 与 point_labels 长度不一致",
"contour": [],
}
if len(req.point_coords) < 1:
return {"success": False, "error": "至少需要一个提示点", "contour": []}
pc = np.array(req.point_coords, dtype=np.float32)
if pc.ndim != 2 or pc.shape[1] != 2:
return {"success": False, "error": "point_coords 每项须为 [x,y]", "contour": []}
pl = np.array(req.point_labels, dtype=np.int64)
box = np.array(req.box_xyxy, dtype=np.float32)
mask = run_sam_prompt(rgb, pc, pl, box_xyxy=box)
if not np.any(mask):
return {"success": False, "error": "SAM 未产生有效掩膜", "contour": []}
contour = mask_to_contour_xy(mask, epsilon_px=2.0)
if len(contour) < 3:
return {"success": False, "error": "轮廓点数不足", "contour": []}
return {"success": True, "contour": contour, "error": None}
except Exception as e:
return {"success": False, "error": str(e), "contour": []}
# -----------------------------
# Inpaint
# -----------------------------
_inpaint_cache: Dict[str, Any] = {}
def _get_inpaint_predictor(model_name: str):
if model_name in _inpaint_cache:
return _inpaint_cache[model_name]
if model_name == "copy":
def _copy(image: Image.Image, *_args, **_kwargs) -> Image.Image:
return image.convert("RGB")
_inpaint_cache[model_name] = _copy
return _copy
if model_name == "sdxl_inpaint":
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.SDXL_INPAINT))
_inpaint_cache[model_name] = pred
return pred
if model_name == "controlnet":
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.CONTROLNET))
_inpaint_cache[model_name] = pred
return pred
raise ValueError(f"未知 inpaint model_name: {model_name}")
@app.post("/inpaint")
def inpaint(req: InpaintRequest) -> Dict[str, Any]:
try:
model_name = req.model_name or "sdxl_inpaint"
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
if req.mask_b64:
mask = _b64_to_pil_image(req.mask_b64).convert("L")
else:
mask = _default_half_mask(pil)
predictor = _get_inpaint_predictor(model_name)
out = predictor(
pil,
mask,
req.prompt or "",
req.negative_prompt or "",
strength=req.strength,
max_side=req.max_side,
)
out_dir = OUTPUT_DIR / "inpaint"
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / f"{model_name}_inpaint.png"
out.save(out_path)
# 兼容 Qt 前端:直接返回结果图,避免前端再去读取服务器磁盘路径
return {
"success": True,
"output_path": str(out_path),
"output_image_b64": _pil_image_to_png_b64(out),
}
except Exception as e:
return {"success": False, "error": str(e)}
_animation_cache: Dict[str, Any] = {}
def _get_animation_predictor(model_name: str):
if model_name in _animation_cache:
return _animation_cache[model_name]
if model_name == "animatediff":
pred, _ = build_animation_predictor(
UnifiedAnimationConfig(backend=AnimationBackend.ANIMATEDIFF)
)
_animation_cache[model_name] = pred
return pred
raise ValueError(f"未知 animation model_name: {model_name}")
@app.post("/animate")
def animate(req: AnimateRequest) -> Dict[str, Any]:
try:
model_name = req.model_name or "animatediff"
predictor = _get_animation_predictor(model_name)
result_path = predictor(
prompt=req.prompt,
negative_prompt=req.negative_prompt or "",
num_inference_steps=req.num_inference_steps,
guidance_scale=req.guidance_scale,
width=req.width,
height=req.height,
video_length=req.video_length,
seed=req.seed,
)
out_dir = OUTPUT_DIR / "animation"
out_dir.mkdir(parents=True, exist_ok=True)
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
out_path = out_dir / f"{model_name}_{ts}.gif"
shutil.copy2(result_path, out_path)
return {"success": True, "output_path": str(out_path)}
except Exception as e:
return {"success": False, "error": str(e)}
@app.get("/health")
def health() -> Dict[str, str]:
return {"status": "ok"}
from app.main import app
__all__ = ["app"]
if __name__ == "__main__":
import uvicorn
port = int(os.environ.get("MODEL_SERVER_PORT", "8000"))
uvicorn.run(app, host="0.0.0.0", port=port)