(self, request, context)
| 306 | return backend_pb2.Reply(message=b"OK") |
| 307 | |
| 308 | def LoadModel(self, request, context): |
| 309 | try: |
| 310 | self.options = _parse_options(list(getattr(request, "Options", []) or [])) |
| 311 | model_path = getattr(request, "ModelPath", None) or "" |
| 312 | model_name = (request.Model or "").strip() |
| 313 | model_file = (getattr(request, "ModelFile", None) or "").strip() |
| 314 | |
| 315 | # Model dir: where we store checkpoints (always under LocalAI models path, never backend dir) |
| 316 | if model_path and model_name: |
| 317 | model_dir = os.path.join(model_path, model_name) |
| 318 | elif model_file: |
| 319 | model_dir = model_file |
| 320 | else: |
| 321 | model_dir = os.path.abspath(model_name or ".") |
| 322 | self.model_dir = model_dir |
| 323 | self.checkpoint_dir = os.path.join(model_dir, "checkpoints") |
| 324 | self.project_root = model_dir |
| 325 | self.model_path = os.path.join(self.checkpoint_dir, model_name or os.path.basename(model_dir.rstrip("/\\"))) |
| 326 | |
| 327 | config_path = model_name or os.path.basename(model_dir.rstrip("/\\")) |
| 328 | os.makedirs(self.checkpoint_dir, exist_ok=True) |
| 329 | |
| 330 | self.dit_handler = AceStepHandler() |
| 331 | # Patch handler so it uses our model dir instead of site-packages/checkpoints |
| 332 | self.dit_handler._get_project_root = lambda: self.project_root |
| 333 | device = self.options.get("device", "auto") |
| 334 | use_flash = self.options.get("use_flash_attention", True) |
| 335 | if isinstance(use_flash, str): |
| 336 | use_flash = str(use_flash).lower() in ("1", "true", "yes") |
| 337 | offload = self.options.get("offload_to_cpu", False) |
| 338 | if isinstance(offload, str): |
| 339 | offload = str(offload).lower() in ("1", "true", "yes") |
| 340 | status_msg, ok = self.dit_handler.initialize_service( |
| 341 | project_root=self.project_root, |
| 342 | config_path=config_path, |
| 343 | device=device, |
| 344 | use_flash_attention=use_flash, |
| 345 | compile_model=False, |
| 346 | offload_to_cpu=offload, |
| 347 | offload_dit_to_cpu=bool(self.options.get("offload_dit_to_cpu", False)), |
| 348 | ) |
| 349 | if not ok: |
| 350 | return backend_pb2.Result(success=False, message=f"DiT init failed: {status_msg}") |
| 351 | |
| 352 | self.llm_handler = None |
| 353 | if self.options.get("init_lm", True): |
| 354 | lm_model = self.options.get("lm_model_path", "acestep-5Hz-lm-0.6B") |
| 355 | |
| 356 | # Ensure LM model is downloaded before initializing |
| 357 | try: |
| 358 | from pathlib import Path |
| 359 | lm_success, lm_msg = ensure_lm_model( |
| 360 | model_name=lm_model, |
| 361 | checkpoints_dir=Path(self.checkpoint_dir), |
| 362 | prefer_source=None, # Auto-detect HuggingFace vs ModelScope |
| 363 | ) |
| 364 | if not lm_success: |
| 365 | print(f"[ace-step] Warning: LM model download failed: {lm_msg}", file=sys.stderr) |
nothing calls this directly
no test coverage detected