Add overlap working-memory ablation and LoRA epoch2 record

This commit is contained in:
Xin Wang
2026-08-24 09:52:56 +08:00
parent 3cc0c3da85
commit 20446c0bd8
6 changed files with 851 additions and 8 deletions

View File

@@ -496,6 +496,28 @@ class CausalSourceAV:
"""Map one global output tick to a source frame at or before it."""
global_time = grid_index / self.fps
clip_index, frame_index, source_time = self._source_frame_for_time(
global_time,
count_tail_reuse=count_tail_reuse,
)
return clip_index, frame_index, global_time, source_time
def _source_frame_for_time(
self,
global_time: float,
*,
count_tail_reuse: bool = True,
) -> tuple[int, int, float]:
"""Map an arbitrary causal timestamp to the latest observed frame."""
if global_time < -EPSILON:
raise ValueError(f"Negative video timestamp: {global_time}")
if global_time > self.required_duration + EPSILON:
raise ValueError(
f"Video timestamp {global_time:g}s exceeds required duration "
f"{self.required_duration:g}s"
)
global_time = min(max(0.0, float(global_time)), self.required_duration)
clip_index, clip, local_time = self._clip_for_time(global_time)
frame_index = int(
np.searchsorted(clip.source_frame_starts, local_time + EPSILON, side="right")
@@ -510,7 +532,14 @@ class CausalSourceAV:
)
if count_tail_reuse and local_time > float(clip.source_frame_starts[-1]) + EPSILON:
self._tail_frame_reuses += 1
return clip_index, frame_index, global_time, source_time
return clip_index, frame_index, source_time
def frame_at_or_before(self, global_time: float) -> Any:
"""Decode the latest source frame whose PTS is not later than ``time``."""
clip_index, frame_index, _ = self._source_frame_for_time(global_time)
self._request_count += 1
return self._load_frame(clip_index, frame_index)
def _load_frame(self, clip_index: int, frame_index: int) -> Any:
from PIL import Image
@@ -797,6 +826,89 @@ class CausalSourceAV:
}
class EndpointAlignedAV:
"""Working-window view whose final visual sample is exactly current time.
Complete cached blocks use the source reader's fixed global half-open grid,
for example ``[18.0, 18.5, 19.0, 19.5]`` for ``[18,20)`` at 2 FPS. A
two-second working interval ending at ``t`` instead uses
``[t-1.5, t-1.0, t-0.5, t]``. The number of visual samples therefore
stays matched while the working branch explicitly observes the current
causal frame. Audio remains the exact half-open slice ``[start,end)``.
"""
def __init__(
self,
source: CausalSourceAV,
*,
max_interval_seconds: float,
) -> None:
if max_interval_seconds <= 0:
raise ValueError("max_interval_seconds must be positive")
self.source = source
self.fps = source.fps
self.sample_rate = source.sample_rate
self.max_interval_seconds = float(max_interval_seconds)
max_frames = self.max_interval_seconds * self.fps
if not math.isclose(max_frames, round(max_frames), abs_tol=1e-6):
raise ValueError(
"working-memory seconds * video FPS must be integral, got "
f"{max_frames}"
)
self.max_frames = int(round(max_frames))
def sample_times(self, start: float, end: float) -> list[float]:
if end <= start:
raise ValueError(f"Empty video interval [{start}, {end}]")
duration = end - start
if duration > self.max_interval_seconds + EPSILON:
raise ValueError(
f"Working interval is {duration:g}s, exceeds configured "
f"{self.max_interval_seconds:g}s"
)
count = max(1, int(math.ceil(duration * self.fps - EPSILON)))
count = min(count, self.max_frames)
first = end - (count - 1) / self.fps
if first < start - EPSILON:
raise RuntimeError(
f"Endpoint-aligned grid starts before its interval: "
f"{first:g} < {start:g}"
)
return [first + index / self.fps for index in range(count)]
def frames(self, start: float, end: float) -> list[Any]:
frames = [
self.source.frame_at_or_before(timestamp)
for timestamp in self.sample_times(start, end)
]
if len(frames) % 2:
frames.append(frames[-1].copy())
self.source._odd_frame_duplications += 1
return frames
def audio_slice(self, start: float, end: float) -> np.ndarray:
return self.source.audio_slice(start, end)
def metadata(self) -> dict[str, Any]:
return {
"visual_sampling": "endpoint_aligned_latest_source_pts_not_after_tick",
"audio_interval": "half_open_[start,end)",
"working_memory_seconds": self.max_interval_seconds,
"output_video_fps": self.fps,
"full_window_frame_count": self.max_frames,
"full_window_relative_video_times": [
-(self.max_frames - 1 - index) / self.fps
for index in range(self.max_frames)
],
"includes_current_video_frame": True,
"oldest_visual_tick_replaced_by_current": True,
"timestamp_note": (
"video is endpoint-aligned while audio retains the full "
"half-open working interval"
),
}
def _manifest_session_specs(
manifest_path: Path,
session_name: str,
@@ -901,7 +1013,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser.add_argument("--output", type=Path)
parser.add_argument(
"--mode",
choices=("causal-prefix", "open-stream-replay"),
choices=(
"causal-prefix",
"open-stream-replay",
"overlap-working-memory",
),
default="open-stream-replay",
)
parser.add_argument("--model-id", default=core.MODEL_ID)
@@ -909,6 +1025,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser.add_argument("--end", type=float)
parser.add_argument("--step-seconds", type=float, default=0.2)
parser.add_argument("--native-block-seconds", type=float, default=2.0)
parser.add_argument(
"--working-memory-seconds",
type=float,
default=2.0,
help=(
"Sliding recent recap used by overlap-working-memory; its duration "
"is included in --memory-seconds rather than added on top"
),
)
parser.add_argument(
"--memory-seconds",
type=float,
@@ -1046,6 +1171,16 @@ def main(
raise ValueError("model and processor must be supplied together")
if args.verify_every < 0 or args.verify_continuous_every < 0:
raise SystemExit("verification intervals must be >= 0")
if args.working_memory_seconds <= 0:
raise SystemExit("--working-memory-seconds must be positive")
if (
args.mode == "overlap-working-memory"
and args.verify_continuous_every > 0
):
raise SystemExit(
"--verify-continuous-every is not defined for the intentional "
"overlap representation; use --verify-every"
)
if args.frame_cache_size < 0:
raise SystemExit("--frame-cache-size must be >= 0")
if args.output is None and not args.media_only:
@@ -1131,6 +1266,11 @@ def main(
return
duration = end - args.start
if args.mode == "overlap-working-memory" and args.memory_seconds <= 0:
raise SystemExit(
"overlap-working-memory uses a matched finite presentation budget; "
"pass a positive --memory-seconds (recommended: 8)"
)
effective_memory = _effective_memory_seconds(
args.memory_seconds,
duration=duration,
@@ -1167,6 +1307,8 @@ def main(
model.requires_grad_(False)
media_path_for_api = reader.clips[0].path
overlap_extractor = None
working_grid_av = None
if args.mode == "causal-prefix":
extractor = core.CausalPrefixExtractor(
model, processor, config, fixed_grid_av=reader
@@ -1179,7 +1321,7 @@ def main(
use_audio_in_video=True,
show_progress=not args.no_progress,
)
else:
elif args.mode == "open-stream-replay":
extractor = core.OpenStreamReplayExtractor(
model, processor, config, fixed_grid_av=reader
)
@@ -1193,20 +1335,51 @@ def main(
verify_continuous_every=args.verify_continuous_every,
show_progress=not args.no_progress,
)
else:
working_grid_av = EndpointAlignedAV(
reader,
max_interval_seconds=args.working_memory_seconds,
)
overlap_extractor = core.OverlappingWorkingMemoryExtractor(
model,
processor,
config,
working_memory_seconds=args.working_memory_seconds,
fixed_grid_av=reader,
working_grid_av=working_grid_av,
)
result = overlap_extractor.extract(
video_path=media_path_for_api,
audio_path=None,
start=args.start,
end=end,
use_audio_in_video=True,
verify_every=args.verify_every,
show_progress=not args.no_progress,
)
if target_steps is not None and len(result.times) != target_steps:
raise RuntimeError(
f"Extractor returned {len(result.times)} states, expected {target_steps}"
)
persistent_memory = args.memory_seconds == 0
persistent_memory = (
args.mode == "open-stream-replay" and args.memory_seconds == 0
)
payload: dict[str, Any] = {
"model_id": args.model_id,
"mode": args.mode,
"stream_representation": (
"single_causal_prefix"
if args.mode == "causal-prefix"
else "single_open_media_turn_with_2s_payload_commits"
else (
"single_open_media_turn_with_2s_payload_commits"
if args.mode == "open-stream-replay"
else (
"serialized_cached_complete_blocks_then_overlapping_"
"endpoint_aligned_recent_recap"
)
)
),
"config": asdict(config),
"requested_memory_seconds": args.memory_seconds,
@@ -1221,7 +1394,9 @@ def main(
"source_manifest_record": manifest_record,
"direct_media": reader.metadata(),
"video_sampling_grid": (
"direct_pts_causal_global_grid_half_open_with_tail_duplication"
"fixed_global_half_open_plus_endpoint_aligned_working_recap"
if args.mode == "overlap-working-memory"
else "direct_pts_causal_global_grid_half_open_with_tail_duplication"
),
"open_stream_policy": (
{
@@ -1242,6 +1417,39 @@ def main(
if args.mode == "open-stream-replay"
else None
),
"overlap_working_memory_policy": (
{
"total_presentation_seconds": config.memory_seconds,
"fixed_history_seconds": (
config.memory_seconds - args.working_memory_seconds
),
"fixed_block_seconds": config.native_block_seconds,
"working_memory_seconds": args.working_memory_seconds,
"fixed_blocks": (
overlap_extractor.fixed_block_budget
if overlap_extractor is not None
else None
),
"serialization": (
"fixed complete AV payloads followed by overlapping "
"working AV payload in the same open media turn"
),
"overlap_is_intentional": True,
"working_media": (
working_grid_av.metadata()
if working_grid_av is not None
else None
),
"layout": (
overlap_extractor.layout_records
if overlap_extractor is not None
else None
),
"backbone": "frozen",
}
if args.mode == "overlap-working-memory"
else None
),
"times": result.times,
"states": result.states,
"verification": result.verification,