hugging-apps/echo-memory
0
1#!/usr/bin/env python32"""3Infer memory-baseline runtime flags from CKPT path (aligned with train/memory_baselines_basic).4 5与训练脚本逐项对应（exp1_4_4_cam_rt_paper_style/train/memory_baselines_basic/）：6 7| 路径子串 | 训练脚本 | 推理 pipe / MEM_ARGS |8|----------|----------|----------------------|9| *_framepack_weight_only | run_framepack_baseline.sh | use_framepack_memory, decay=0.9, weight=1.0；**无** length compress |10| *_framepack_lencompress_r2 | run_framepack_lencompress_r2.sh | use_framepack_length_compress, framepack_ratio=2 |11| *_framepack_lencompress_r4 | run_framepack_lencompress_r4.sh | use_framepack_length_compress, framepack_ratio=4 |12| *_spatial_mem | run_spatial_memory_baseline.sh | use_spatial_memory, tokens=64, **legacy=False**（SpatialGridMemory，与 train 默认一致；grid=8 由 ckpt 权重形状在 loop_utils 中还原） |13| *_videossm_hybrid | run_videossm_hybrid_baseline.sh | 仅 context_frames=5；VideoSSM 结构由 loop_utils.load_pipeline_and_ckpt 按 ckpt key 推断 |14 15共同：训练均为 context_memory_frames=5 → CONTEXT_FRAMES_MEM_OVERRIDE=5。16 17压缩记忆：训练侧 K→K' 与 RT 对齐在 wan_video_new.framepack_*；推理传 pipe.use_framepack_length_compress + framepack_ratio 即可，与 train.py 写入 pipe 一致。18 19空间记忆：若 ckpt 中无 spatial_memory_module.*，则无法与训练 SpatialGridMemory 对齐；apply_memory_baseline_pipe 会降级 legacy 并打印警告（与 run_replay_loop_two_chunk 逻辑一致）。20 21Used by:22- bash: `eval "$(python3 memory_baseline_runtime.py bash-export "$CKPT")"` → MEM_ARGS array + CONTEXT_FRAMES_MEM_OVERRIDE23- Python: apply_memory_baseline_pipe(pipe, ckpt) after load_pipeline_and_ckpt24"""25from __future__ import annotations26 27import argparse28import json29import shlex30import sys31from dataclasses import dataclass32from typing import Dict, List, Optional, Sequence33 34 35@dataclass36class MemoryProfile:37 use_framepack_memory: bool = False38 context_temporal_decay: float = 1.039 context_attention_weight: float = 1.040 use_framepack_length_compress: bool = False41 framepack_ratio: int = 242 use_spatial_memory: bool = False43 use_spatial_memory_legacy: bool = False # True = old adaptive pool only (no SpatialGridMemory in ckpt)44 spatial_memory_tokens: int = 6445 spatial_memory_inject_mode: Optional[str] = None46 use_block_wise_ssm: bool = False47 use_videossm_hybrid: bool = False48 # Training uses context_memory_frames=5 for memory_baselines_basic49 context_override: Optional[int] = None50 51 52@dataclass(frozen=True)53class MemoryProfileSpec:54 profile_id: str55 ckpt_substrings: Sequence[str]56 paper_tag: str57 train_flags: Sequence[str]58 infer_flags: Sequence[str]59 eval_flags: Sequence[str]60 profile: MemoryProfile61 62 63MEMORY_PROFILE_REGISTRY: List[MemoryProfileSpec] = [64 MemoryProfileSpec(65 profile_id="context_k1",66 ckpt_substrings=("context_k1", "ctx_1", "ctx1"),67 paper_tag="raw_context",68 train_flags=("--context_memory_frames 1",),69 infer_flags=(),70 eval_flags=("--context_frames 1",),71 profile=MemoryProfile(context_override=1),72 ),73 MemoryProfileSpec(74 profile_id="context_k5",75 ckpt_substrings=("context_k5", "ctx_5", "ctx5"),76 paper_tag="raw_context",77 train_flags=("--context_memory_frames 5",),78 infer_flags=(),79 eval_flags=("--context_frames 5",),80 profile=MemoryProfile(context_override=5),81 ),82 MemoryProfileSpec(83 profile_id="context_k20",84 ckpt_substrings=("context_k20", "ctx_20", "ctx20"),85 paper_tag="raw_context",86 train_flags=("--context_memory_frames 20",),87 infer_flags=(),88 eval_flags=("--context_frames 20",),89 profile=MemoryProfile(context_override=20),90 ),91 MemoryProfileSpec(92 profile_id="framepack_weight_only",93 ckpt_substrings=(94 "memory_baselines_basic_framepack_weight_only",95 "memory_baselines_basic_abl_framepack_weight_two_chunk",96 ),97 paper_tag="arXiv:2504.12626",98 train_flags=(99 "--use_framepack_memory",100 "--context_temporal_decay 0.9",101 "--context_attention_weight 1.0",102 "--context_memory_frames 81",103 ),104 infer_flags=("--use_framepack_memory", "--context_temporal_decay 0.9", "--context_attention_weight 1.0"),105 eval_flags=("--context_frames 81",),106 profile=MemoryProfile(107 use_framepack_memory=True,108 context_temporal_decay=0.9,109 context_attention_weight=1.0,110 context_override=81,111 ),112 ),113 MemoryProfileSpec(114 profile_id="framepack_lencompress_r2",115 ckpt_substrings=(116 "memory_baselines_basic_framepack_lencompress_r2",117 "memory_baselines_basic_abl_framepack_len_r2_two_chunk",118 ),119 paper_tag="arXiv:2504.12626",120 train_flags=("--use_framepack_length_compress", "--framepack_ratio 2", "--context_memory_frames 81"),121 infer_flags=("--use_framepack_length_compress", "--framepack_ratio 2"),122 eval_flags=("--context_frames 81",),123 profile=MemoryProfile(124 use_framepack_length_compress=True,125 framepack_ratio=2,126 context_override=81,127 ),128 ),129 MemoryProfileSpec(130 profile_id="framepack_lencompress_r4",131 ckpt_substrings=(132 "memory_baselines_basic_framepack_lencompress_r4",133 "memory_baselines_basic_abl_framepack_len_r4_two_chunk",134 ),135 paper_tag="arXiv:2504.12626",136 train_flags=("--use_framepack_length_compress", "--framepack_ratio 4", "--context_memory_frames 81"),137 infer_flags=("--use_framepack_length_compress", "--framepack_ratio 4"),138 eval_flags=("--context_frames 81",),139 profile=MemoryProfile(140 use_framepack_length_compress=True,141 framepack_ratio=4,142 context_override=81,143 ),144 ),145 MemoryProfileSpec(146 profile_id="spatial_mem",147 ckpt_substrings=("memory_baselines_basic_spatial_mem", "spatial_mem"),148 paper_tag="arXiv:2506.05284",149 train_flags=("--use_spatial_memory", "--spatial_memory_tokens 64", "--context_memory_frames 5"),150 infer_flags=("--use_spatial_memory", "--spatial_memory_tokens 64"),151 eval_flags=("--context_frames 5",),152 profile=MemoryProfile(153 use_spatial_memory=True,154 use_spatial_memory_legacy=False,155 spatial_memory_tokens=64,156 context_override=5,157 ),158 ),159 MemoryProfileSpec(160 profile_id="framepack_hybrid_r2_weight_two_chunk",161 ckpt_substrings=("memory_baselines_basic_abl_framepack_hybrid_r2_weight_two_chunk",),162 paper_tag="hybrid_compression",163 train_flags=(164 "--use_framepack_memory",165 "--context_temporal_decay 0.95",166 "--context_attention_weight 1.1",167 "--use_framepack_length_compress",168 "--framepack_ratio 2",169 "--context_memory_frames 81",170 ),171 infer_flags=(172 "--use_framepack_memory",173 "--context_temporal_decay 0.95",174 "--context_attention_weight 1.1",175 "--use_framepack_length_compress",176 "--framepack_ratio 2",177 ),178 eval_flags=("--context_frames 81",),179 profile=MemoryProfile(180 use_framepack_memory=True,181 context_temporal_decay=0.95,182 context_attention_weight=1.1,183 use_framepack_length_compress=True,184 framepack_ratio=2,185 context_override=81,186 ),187 ),188 MemoryProfileSpec(189 profile_id="framepack_hybrid_r4_weight_two_chunk",190 ckpt_substrings=("memory_baselines_basic_abl_framepack_hybrid_r4_weight_two_chunk",),191 paper_tag="hybrid_compression",192 train_flags=(193 "--use_framepack_memory",194 "--context_temporal_decay 0.95",195 "--context_attention_weight 1.1",196 "--use_framepack_length_compress",197 "--framepack_ratio 4",198 "--context_memory_frames 81",199 ),200 infer_flags=(201 "--use_framepack_memory",202 "--context_temporal_decay 0.95",203 "--context_attention_weight 1.1",204 "--use_framepack_length_compress",205 "--framepack_ratio 4",206 ),207 eval_flags=("--context_frames 81",),208 profile=MemoryProfile(209 use_framepack_memory=True,210 context_temporal_decay=0.95,211 context_attention_weight=1.1,212 use_framepack_length_compress=True,213 framepack_ratio=4,214 context_override=81,215 ),216 ),217 MemoryProfileSpec(218 profile_id="spatial_concat_text_two_chunk",219 ckpt_substrings=("memory_baselines_basic_abl_spatial_concat_text_two_chunk", "spatial_concat_text_two_chunk"),220 paper_tag="arXiv:2506.05284",221 train_flags=(222 "--use_spatial_memory",223 "--spatial_memory_tokens 64",224 "--spatial_memory_inject_mode concat_text",225 "--context_memory_frames 1",226 ),227 infer_flags=("--use_spatial_memory", "--spatial_memory_tokens 64", "--spatial_memory_inject_mode concat_text"),228 eval_flags=("--context_frames 1",),229 profile=MemoryProfile(230 use_spatial_memory=True,231 use_spatial_memory_legacy=False,232 spatial_memory_tokens=64,233 spatial_memory_inject_mode="concat_text",234 context_override=1,235 ),236 ),237 MemoryProfileSpec(238 profile_id="spatial_inject_none_two_chunk",239 ckpt_substrings=("memory_baselines_basic_abl_spatial_inject_none_two_chunk", "spatial_inject_none_two_chunk"),240 paper_tag="arXiv:2506.05284",241 train_flags=(242 "--use_spatial_memory",243 "--spatial_memory_tokens 64",244 "--spatial_memory_inject_mode none",245 "--context_memory_frames 1",246 ),247 infer_flags=("--use_spatial_memory", "--spatial_memory_tokens 64", "--spatial_memory_inject_mode none"),248 eval_flags=("--context_frames 1",),249 profile=MemoryProfile(250 use_spatial_memory=True,251 use_spatial_memory_legacy=False,252 spatial_memory_tokens=64,253 spatial_memory_inject_mode="none",254 context_override=1,255 ),256 ),257 MemoryProfileSpec(258 profile_id="spatial_cross_attn_readout_two_chunk",259 ckpt_substrings=("memory_baselines_basic_abl_spatial_cross_attn_readout_two_chunk", "spatial_cross_attn_readout_two_chunk"),260 paper_tag="arXiv:2506.05284",261 train_flags=(262 "--use_spatial_memory",263 "--spatial_memory_tokens 64",264 "--spatial_memory_inject_mode cross_attn_readout",265 "--context_memory_frames 1",266 ),267 infer_flags=("--use_spatial_memory", "--spatial_memory_tokens 64", "--spatial_memory_inject_mode cross_attn_readout"),268 eval_flags=("--context_frames 1",),269 profile=MemoryProfile(270 use_spatial_memory=True,271 use_spatial_memory_legacy=False,272 spatial_memory_tokens=64,273 spatial_memory_inject_mode="cross_attn_readout",274 context_override=1,275 ),276 ),277 MemoryProfileSpec(278 profile_id="spatial_cross_attn_readout_t32_g4_two_chunk",279 ckpt_substrings=(280 "memory_baselines_basic_abl_spatial_cross_attn_readout_t32_g4_two_chunk",281 "spatial_cross_attn_readout_t32_g4_two_chunk",282 ),283 paper_tag="arXiv:2506.05284",284 train_flags=(285 "--use_spatial_memory",286 "--spatial_memory_tokens 32",287 "--spatial_memory_inject_mode cross_attn_readout",288 "--context_memory_frames 1",289 ),290 infer_flags=("--use_spatial_memory", "--spatial_memory_tokens 32", "--spatial_memory_inject_mode cross_attn_readout"),291 eval_flags=("--context_frames 1",),292 profile=MemoryProfile(293 use_spatial_memory=True,294 use_spatial_memory_legacy=False,295 spatial_memory_tokens=32,296 spatial_memory_inject_mode="cross_attn_readout",297 context_override=1,298 ),299 ),300 MemoryProfileSpec(301 profile_id="spatial_legacy_concat_t32_two_chunk",302 ckpt_substrings=("memory_baselines_basic_abl_spatial_legacy_concat_t32_two_chunk",),303 paper_tag="arXiv:2506.05284",304 train_flags=(305 "--use_spatial_memory",306 "--use_spatial_memory_legacy",307 "--spatial_memory_tokens 32",308 "--spatial_memory_inject_mode concat_text",309 "--context_memory_frames 1",310 ),311 infer_flags=(312 "--use_spatial_memory",313 "--use_spatial_memory_legacy",314 "--spatial_memory_tokens 32",315 "--spatial_memory_inject_mode concat_text",316 ),317 eval_flags=("--context_frames 1",),318 profile=MemoryProfile(319 use_spatial_memory=True,320 use_spatial_memory_legacy=True,321 spatial_memory_tokens=32,322 spatial_memory_inject_mode="concat_text",323 context_override=1,324 ),325 ),326 MemoryProfileSpec(327 profile_id="videossm_hybrid_legacy",328 ckpt_substrings=(329 "memory_baselines_basic_videossm_hybrid",330 "memory_baselines_basic_abl_videossm_hybrid_two_chunk",331 ),332 paper_tag="legacy_hybrid",333 train_flags=("--use_videossm_hybrid", "--context_memory_frames 5"),334 infer_flags=("load_pipeline_and_ckpt auto-infers VideoSSM hybrid from videossm_hybrid.* ckpt keys",),335 eval_flags=("--context_frames 5",),336 profile=MemoryProfile(use_videossm_hybrid=True, context_override=5),337 ),338 MemoryProfileSpec(339 profile_id="block_wise_ssm_two_chunk",340 ckpt_substrings=("memory_baselines_basic_abl_block_wise_ssm_two_chunk",),341 paper_tag="block_wise_recurrence",342 train_flags=("--use_block_wise_ssm", "--context_memory_frames 5"),343 infer_flags=("load_pipeline_and_ckpt auto-infers Block-wise SSM from block_wise_ssm.* ckpt keys",),344 eval_flags=("--context_frames 5",),345 profile=MemoryProfile(use_block_wise_ssm=True, context_override=5),346 ),347 MemoryProfileSpec(348 profile_id="no_memory_extra_two_chunk",349 ckpt_substrings=("memory_baselines_basic_abl_no_memory_extra_two_chunk",),350 paper_tag="anchor_only_reference",351 train_flags=("--context_memory_frames 5",),352 infer_flags=(),353 eval_flags=("--context_frames 5",),354 profile=MemoryProfile(context_override=5),355 ),356]357 358 359def profile_registry_payload() -> List[Dict[str, object]]:360 payload: List[Dict[str, object]] = []361 for spec in MEMORY_PROFILE_REGISTRY:362 payload.append(363 {364 "profile_id": spec.profile_id,365 "paper_tag": spec.paper_tag,366 "ckpt_substrings": list(spec.ckpt_substrings),367 "train_flags": list(spec.train_flags),368 "infer_flags": list(spec.infer_flags),369 "eval_flags": list(spec.eval_flags),370 "context_override": spec.profile.context_override,371 }372 )373 return payload374 375 376def infer_memory_profile_spec(ckpt: str) -> Optional[MemoryProfileSpec]:377 c = (ckpt or "").lower()378 for spec in MEMORY_PROFILE_REGISTRY:379 if any(token.lower() in c for token in spec.ckpt_substrings):380 return spec381 return None382 383 384def infer_memory_profile(ckpt: str) -> MemoryProfile:385 spec = infer_memory_profile_spec(ckpt)386 if spec is not None:387 return spec.profile388 return MemoryProfile()389 390 391def profile_to_argv(p: MemoryProfile) -> List[str]:392 out: List[str] = []393 if p.use_framepack_memory:394 out.extend(395 [396 "--use_framepack_memory",397 "--context_temporal_decay",398 str(p.context_temporal_decay),399 "--context_attention_weight",400 str(p.context_attention_weight),401 ]402 )403 if p.use_framepack_length_compress:404 out.extend(405 [406 "--use_framepack_length_compress",407 "--framepack_ratio",408 str(int(p.framepack_ratio)),409 ]410 )411 if p.use_spatial_memory:412 out.extend(413 [414 "--use_spatial_memory",415 "--spatial_memory_tokens",416 str(int(p.spatial_memory_tokens)),417 ]418 )419 if p.spatial_memory_inject_mode:420 out.extend(["--spatial_memory_inject_mode", str(p.spatial_memory_inject_mode)])421 if p.use_spatial_memory_legacy:422 out.append("--use_spatial_memory_legacy")423 if p.use_block_wise_ssm:424 out.append("--use_block_wise_ssm")425 if p.use_videossm_hybrid:426 out.append("--use_videossm_hybrid")427 return out428 429 430def bash_export(ckpt: str) -> str:431 """Print bash snippet: unset override, MEM_ARGS=(), optional CONTEXT_FRAMES_MEM_OVERRIDE, MEM_ARGS+=."""432 prof = infer_memory_profile(ckpt)433 lines = [434 "unset CONTEXT_FRAMES_MEM_OVERRIDE 2>/dev/null || true",435 "MEM_ARGS=()",436 ]437 if prof.context_override is not None:438 lines.append(f"CONTEXT_FRAMES_MEM_OVERRIDE={int(prof.context_override)}")439 for arg in profile_to_argv(prof):440 lines.append(f"MEM_ARGS+=({shlex.quote(arg)})")441 return "\n".join(lines) + "\n"442 443 444def apply_memory_baseline_pipe(pipe, ckpt: str) -> None:445 """Set pipe.* memory flags from ckpt path (no-op for non–memory-baseline paths)."""446 p = infer_memory_profile(ckpt)447 pipe.use_framepack_memory = bool(p.use_framepack_memory)448 pipe.context_temporal_decay = float(p.context_temporal_decay or 1.0)449 pipe.context_attention_weight = float(p.context_attention_weight or 1.0)450 pipe.use_framepack_length_compress = bool(p.use_framepack_length_compress)451 pipe.framepack_ratio = int(p.framepack_ratio or 2)452 pipe.use_spatial_memory = bool(p.use_spatial_memory)453 pipe.spatial_memory_tokens = int(p.spatial_memory_tokens or 64)454 if getattr(p, "spatial_memory_inject_mode", None):455 pipe.spatial_memory_inject_mode = str(p.spatial_memory_inject_mode)456 pipe.use_spatial_memory_legacy = bool(p.use_spatial_memory_legacy)457 pipe.use_block_wise_ssm = bool(p.use_block_wise_ssm)458 pipe.use_videossm_hybrid = bool(p.use_videossm_hybrid)459 # 与 run_replay_loop_two_chunk 一致：宣称 SpatialGridMemory 但 ckpt 未载入模块时，避免 silent 与训练不一致460 if (461 pipe.use_spatial_memory462 and not pipe.use_spatial_memory_legacy463 and getattr(pipe, "spatial_memory_module", None) is None464 ):465 pipe.use_spatial_memory_legacy = True466 print(467 "[memory_baseline_runtime] WARN: spatial_mem profile but spatial_memory_module missing on pipe "468 "(ckpt 缺少 spatial_memory_module.* 时与 run_spatial_memory_baseline.sh 的 SpatialGridMemory 不对齐；"469 "已降级为 adaptive pool，数值仅作参考)。",470 file=sys.stderr,471 flush=True,472 )473 474 475def effective_context_frames(ckpt: str, env_context_frames: Optional[int]) -> int:476 """For shell helpers: default CONTEXT_FRAMES from env/int, else profile override, else 1."""477 prof = infer_memory_profile(ckpt)478 if env_context_frames is not None:479 return int(env_context_frames)480 if prof.context_override is not None:481 return int(prof.context_override)482 return 1483 484 485def main() -> None:486 ap = argparse.ArgumentParser(description="Memory baseline CKPT path → MEM_ARGS / pipe flags")487 sub = ap.add_subparsers(dest="cmd", required=True)488 489 p_export = sub.add_parser("bash-export", help="Print eval-able bash for MEM_ARGS + CONTEXT_FRAMES_MEM_OVERRIDE")490 p_export.add_argument("ckpt")491 492 p_ctx = sub.add_parser("print-context", help="Print effective context_frames (int) for replay_gt-style scripts")493 p_ctx.add_argument("ckpt")494 p_ctx.add_argument(495 "--fallback",496 type=int,497 default=None,498 help="If no profile override, use this (else 1)",499 )500 sub.add_parser("list-profiles", help="Print memory profile registry JSON")501 502 args = ap.parse_args()503 if args.cmd == "bash-export":504 sys.stdout.write(bash_export(args.ckpt))505 elif args.cmd == "print-context":506 prof = infer_memory_profile(args.ckpt)507 v = prof.context_override if prof.context_override is not None else (args.fallback if args.fallback is not None else 1)508 print(int(v))509 elif args.cmd == "list-profiles":510 print(json.dumps(profile_registry_payload(), ensure_ascii=False, indent=2))511 512 513if __name__ == "__main__":514 main()515