diff --git a/runtime/training/families/anima/loader.py b/runtime/training/families/anima/loader.py index 7f646f7..064070f 100644 --- a/runtime/training/families/anima/loader.py +++ b/runtime/training/families/anima/loader.py @@ -188,6 +188,12 @@ def load_anima_model(transformer_path, device, dtype, repo_root, *, num_blocks, num_heads = 36, 40 else: raise RuntimeError(f"未知的 model_channels={model_channels}") + # 层数以 checkpoint 为准:同为 2048 宽也有 28 层(2B)和 40 层(2.9B)两种。 + # 写死 28 会让 strict=False 的加载静默丢掉多出来的层。 + header_blocks = block_count_from_header(transformer_path) + if header_blocks and header_blocks != num_blocks: + logger.info("Anima DiT 层数按 checkpoint 取 %d(默认 %d)", header_blocks, num_blocks) + num_blocks = header_blocks config = dict( max_img_h=1024, max_img_w=1024, max_frames=128, @@ -209,6 +215,12 @@ def load_anima_model(transformer_path, device, dtype, repo_root, *, # 加载权重 sd = _load_safetensors_state_dict(Path(transformer_path)) + # RoPE 表(seq = arange(最大长度) 等)由 config 推出;有的导出把它们按别的 + # 最大长度存进了权重(2.9B: seq 256 vs 模型 512),交给模型自己重建。 + sd = { + k: v for k, v in sd.items() + if not ("pos_embedder." in k and k.rsplit(".", 1)[-1] in ("seq", "dim_spatial_range", "dim_temporal_range")) + } info = _load_weights_best_effort(model, sd, label="Transformer") # 如果 checkpoint 中完全没有 llm_adapter 权重,随机初始化会把 cross-attn 条件搞乱,直接禁用更安全 diff --git a/studio/services/soup.py b/studio/services/soup.py index ac8bab5..753aeca 100644 --- a/studio/services/soup.py +++ b/studio/services/soup.py @@ -160,13 +160,24 @@ def _unlink(path: Path) -> None: raise SoupError(f"Could not delete the file: {exc}") from exc +def _temp_part(directory: Path) -> Path: + """Reserve a unique `.part` file and release its handle right away. + + mkstemp leaves the file open; on Windows that open handle blocks the + writer (safetensors gets "access denied") and the cleanup unlink. + """ + fd, name = tempfile.mkstemp(dir=str(directory), suffix=".part") + os.close(fd) + return Path(name) + + def save_upload(name: str, data: bytes) -> dict[str, Any]: """Store an uploaded adapter, after checking it is a readable safetensors.""" ensure_dirs() if not data: raise SoupError("The uploaded file is empty.") path = upload_path(name) - tmp = Path(tempfile.mkstemp(dir=str(UPLOAD_DIR), suffix=".part")[1]) + tmp = _temp_part(UPLOAD_DIR) try: tmp.write_bytes(data) # Parse before accepting: a truncated upload should fail here, not two @@ -409,7 +420,7 @@ def merge( meta = _soup_metadata(base_info, paths, raw_weights, weights, method) - tmp = Path(tempfile.mkstemp(dir=str(OUTPUT_DIR), suffix=".part")[1]) + tmp = _temp_part(OUTPUT_DIR) try: save_file(merged, str(tmp), metadata=meta) os.replace(tmp, out) diff --git a/studio/services/tunnel.py b/studio/services/tunnel.py index 9de7b7d..785d077 100644 --- a/studio/services/tunnel.py +++ b/studio/services/tunnel.py @@ -463,8 +463,12 @@ def _run_until_exit(self, cmd: list[str]) -> tuple[list[str], Optional[int]]: waiting (killed after the timeout, e.g. Tailscale asking for consent).""" proc = self._spawn(cmd, dict(os.environ)) self._await(proc, lambda _line: None, until_exit=True) - code = proc.poll() - if code is None: + # Output EOF can arrive a moment before the OS records the exit code + # (seen on Windows: poll() still None ~30 ms after the CLI finished). + try: + code: Optional[int] = proc.wait(timeout=5) + except subprocess.TimeoutExpired: + code = None self._stop_process(proc) return list(self._log), code