Implement AI Host functionality with model caching and job management
- Added ai_host.py to handle AI tasks with FastAPI integration. - Introduced HostState and ModelCache classes for managing client connections and model loading. - Implemented job submission and cancellation endpoints in the AI host. - Enhanced tests for AI host functionality, including client tracking and job cancellation. - Updated existing tests to accommodate changes in the operator and command structure. - Refactored runtime commands to use request-based architecture instead of command-line arguments. - Improved model loading logic to handle CUDA memory management and reuse models efficiently.
This commit is contained in:
@@ -23,10 +23,27 @@ SCENE_FILENAME = "scene.glb"
|
||||
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
|
||||
|
||||
|
||||
def write_status(path, progress, message, error=None):
|
||||
class TaskCancelled(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def check_cancelled(cancel_check):
|
||||
if cancel_check is not None and cancel_check():
|
||||
raise TaskCancelled("Task cancelled")
|
||||
|
||||
|
||||
def write_status(path, progress, message, error=None, state=None):
|
||||
payload = {"progress": progress, "message": message}
|
||||
if error:
|
||||
payload["error"] = error
|
||||
if state is None:
|
||||
if error:
|
||||
state = "failed"
|
||||
elif progress >= 1.0:
|
||||
state = "succeeded"
|
||||
else:
|
||||
state = "running"
|
||||
payload["state"] = state
|
||||
Path(path).write_text(
|
||||
json.dumps(payload, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
@@ -42,7 +59,13 @@ def collect_images(directory):
|
||||
)
|
||||
|
||||
|
||||
def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
|
||||
def extract_video_frames(
|
||||
video_path,
|
||||
output_directory,
|
||||
maximum_frames,
|
||||
cv2,
|
||||
cancel_check=None,
|
||||
):
|
||||
capture = cv2.VideoCapture(str(video_path))
|
||||
if not capture.isOpened():
|
||||
raise RuntimeError(f"Unable to open video: {video_path}")
|
||||
@@ -66,6 +89,7 @@ def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
|
||||
image_paths = []
|
||||
try:
|
||||
for output_index, frame_index in enumerate(frame_indices):
|
||||
check_cancelled(cancel_check)
|
||||
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
|
||||
success, frame = capture.read()
|
||||
if not success:
|
||||
@@ -82,7 +106,8 @@ def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
|
||||
return image_paths
|
||||
|
||||
|
||||
def prepare_input(input_path, video_frames, temporary_root):
|
||||
def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
|
||||
check_cancelled(cancel_check)
|
||||
input_path = Path(input_path)
|
||||
suffix = input_path.suffix.lower()
|
||||
if suffix in IMAGE_SUFFIXES:
|
||||
@@ -102,6 +127,7 @@ def prepare_input(input_path, video_frames, temporary_root):
|
||||
temporary_directory,
|
||||
video_frames,
|
||||
cv2,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
except Exception:
|
||||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||
@@ -203,19 +229,33 @@ def select_device(requested, torch):
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
def infer_foreground(input_path, model_dir, device, torch):
|
||||
from ben2 import BEN_Base
|
||||
def infer_foreground(
|
||||
input_path,
|
||||
model_dir,
|
||||
device,
|
||||
torch,
|
||||
model_cache=None,
|
||||
cancel_check=None,
|
||||
):
|
||||
from PIL import Image
|
||||
|
||||
input_path = Path(input_path)
|
||||
check_cancelled(cancel_check)
|
||||
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
|
||||
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
|
||||
load_started = time.perf_counter()
|
||||
model = BEN_Base.from_pretrained(model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
load_ms = (time.perf_counter() - load_started) * 1000.0
|
||||
owns_model = model_cache is None
|
||||
if owns_model:
|
||||
from ben2 import BEN_Base
|
||||
|
||||
load_started = time.perf_counter()
|
||||
model = BEN_Base.from_pretrained(model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
load_ms = (time.perf_counter() - load_started) * 1000.0
|
||||
else:
|
||||
model, load_ms = model_cache.get_ben2(model_dir, device)
|
||||
|
||||
check_cancelled(cancel_check)
|
||||
with Image.open(input_path) as opened_image:
|
||||
image = opened_image.convert("RGB")
|
||||
if device.type == "cuda":
|
||||
@@ -229,17 +269,23 @@ def infer_foreground(input_path, model_dir, device, torch):
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
inference_ms = (time.perf_counter() - inference_started) * 1000.0
|
||||
del model
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
if owns_model:
|
||||
del model
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
check_cancelled(cancel_check)
|
||||
return foreground, load_ms, inference_ms
|
||||
|
||||
|
||||
def remove_background(args):
|
||||
def remove_background(args, model_cache=None, cancel_check=None):
|
||||
total_started = time.perf_counter()
|
||||
write_status(args.status, 0.02, "Starting BEN2 Runtime")
|
||||
check_cancelled(cancel_check)
|
||||
|
||||
import torch
|
||||
if model_cache is None:
|
||||
import torch
|
||||
else:
|
||||
torch = model_cache.torch
|
||||
|
||||
output_directory = Path(args.output_dir)
|
||||
output_directory.mkdir(parents=True, exist_ok=False)
|
||||
@@ -252,8 +298,11 @@ def remove_background(args):
|
||||
args.model_dir,
|
||||
device,
|
||||
torch,
|
||||
model_cache=model_cache,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
|
||||
check_cancelled(cancel_check)
|
||||
write_status(args.status, 0.92, "Writing transparent PNG")
|
||||
output_path = output_directory / "foreground.png"
|
||||
foreground.convert("RGBA").save(
|
||||
@@ -264,6 +313,7 @@ def remove_background(args):
|
||||
if not output_path.is_file():
|
||||
raise RuntimeError("BEN2 did not write a foreground image")
|
||||
|
||||
check_cancelled(cancel_check)
|
||||
total_ms = (time.perf_counter() - total_started) * 1000.0
|
||||
message = (
|
||||
f"Background removed in {inference_ms:.0f} ms "
|
||||
@@ -288,10 +338,16 @@ def exported_paths(output_directory, export_format):
|
||||
return []
|
||||
|
||||
|
||||
def wait_for_exported_paths(output_directory, export_format, timeout=30.0):
|
||||
def wait_for_exported_paths(
|
||||
output_directory,
|
||||
export_format,
|
||||
timeout=30.0,
|
||||
cancel_check=None,
|
||||
):
|
||||
# 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
check_cancelled(cancel_check)
|
||||
result_paths = exported_paths(output_directory, export_format)
|
||||
if result_paths:
|
||||
return result_paths
|
||||
@@ -299,11 +355,15 @@ def wait_for_exported_paths(output_directory, export_format, timeout=30.0):
|
||||
return []
|
||||
|
||||
|
||||
def reconstruct(args):
|
||||
def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
write_status(args.status, 0.02, "Starting AI runtime")
|
||||
check_cancelled(cancel_check)
|
||||
|
||||
import torch
|
||||
from depth_anything_3.api import DepthAnything3
|
||||
if model_cache is None:
|
||||
import torch
|
||||
from depth_anything_3.api import DepthAnything3
|
||||
else:
|
||||
torch = model_cache.torch
|
||||
from depth_anything_3.utils.export import export
|
||||
|
||||
output_directory = Path(args.output_dir)
|
||||
@@ -313,9 +373,11 @@ def reconstruct(args):
|
||||
args.input,
|
||||
args.video_frames,
|
||||
output_directory.parent,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
|
||||
try:
|
||||
check_cancelled(cancel_check)
|
||||
device = select_device(args.device, torch)
|
||||
remove_background_model_dir = getattr(
|
||||
args,
|
||||
@@ -329,7 +391,10 @@ def reconstruct(args):
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
torch,
|
||||
model_cache=model_cache,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
check_cancelled(cancel_check)
|
||||
foreground.convert("RGBA").save(
|
||||
output_directory / "foreground.png",
|
||||
format="PNG",
|
||||
@@ -342,10 +407,14 @@ def reconstruct(args):
|
||||
f"model load {load_ms:.0f} ms"
|
||||
)
|
||||
write_status(args.status, 0.15, f"Loading model on {device.type}")
|
||||
model = DepthAnything3.from_pretrained(args.model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
if model_cache is None:
|
||||
model = DepthAnything3.from_pretrained(args.model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
else:
|
||||
model = model_cache.get_da3(args.model_dir, device)
|
||||
|
||||
check_cancelled(cancel_check)
|
||||
write_status(
|
||||
args.status,
|
||||
0.25,
|
||||
@@ -372,6 +441,7 @@ def reconstruct(args):
|
||||
use_ray_pose=bool(args.use_ray_pose),
|
||||
ref_view_strategy=args.ref_view_strategy,
|
||||
)
|
||||
check_cancelled(cancel_check)
|
||||
write_status(args.status, 0.88, f"Exporting {args.export_format}")
|
||||
export(
|
||||
prediction,
|
||||
@@ -379,15 +449,18 @@ def reconstruct(args):
|
||||
str(output_directory),
|
||||
**export_kwargs,
|
||||
)
|
||||
check_cancelled(cancel_check)
|
||||
|
||||
result_paths = wait_for_exported_paths(
|
||||
output_directory,
|
||||
args.export_format,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
if not result_paths:
|
||||
raise RuntimeError(
|
||||
f"No {args.export_format} files were exported"
|
||||
)
|
||||
check_cancelled(cancel_check)
|
||||
write_status(
|
||||
args.status,
|
||||
1.0,
|
||||
|
||||
Reference in New Issue
Block a user