114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
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 == "flux_fill":
|
|
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.FLUX_FILL))
|
|
_inpaint_cache[model_name] = pred
|
|
return pred
|
|
|
|
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}")
|