[Diffusion] Fix the five unit tests failing on main (#36726)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Kevin Mi
2026-08-27 21:18:32 -07:00
committed by GitHub
co-authored by Claude Fable 5
parent 2380121e9b
commit ad5a105a4a
7 changed files with 62 additions and 31 deletions
@@ -2837,14 +2837,19 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module):
kv_proj = apply_kv(torch.cat([k_bhnd, v_bhnd], dim=1))
k_proj, v_proj = torch.chunk(kv_proj, chunks=2, dim=1)
q_pre_dn = q_bhnd.permute(0, 1, 3, 2)
q_dn = q_proj.permute(0, 1, 3, 2)
k_pre_dn = k_bhnd.permute(0, 1, 3, 2)
k_dn = k_proj.permute(0, 1, 3, 2)
v_pre_dn = v_bhnd.permute(0, 1, 3, 2)
v_dn = v_proj.permute(0, 1, 3, 2)
# No RMS downscale here: full post-UCPE q/k/v feed the scan; inflation
# is computed from full post-UCPE K vs pre-UCPE K and absorbed only into
# beta.
# Same per-token RMS downscale as _cam_branch, so a single chunk with
# no carried state reduces exactly to the dense scan.
q_dn = _downscale_to_reference_rms(q_pre_dn, q_dn)
k_dn = _downscale_to_reference_rms(k_pre_dn, k_dn)
v_dn = _downscale_to_reference_rms(v_pre_dn, v_dn)
pre_ucpe_k_norm = torch.linalg.vector_norm(
k_pre_dn.float(), dim=2, keepdim=True
).clamp_min(1e-6)
@@ -2910,10 +2915,17 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module):
kv_proj = apply_kv(torch.cat([k_bhnd, v_bhnd], dim=1))
k_proj, v_proj = torch.chunk(kv_proj, chunks=2, dim=1)
# No RMS downscale here: full post-UCPE q/k/v feed SDPA directly.
q_dn = q_proj.permute(0, 1, 3, 2)
k_dn = k_proj.permute(0, 1, 3, 2)
v_dn = v_proj.permute(0, 1, 3, 2)
# Same per-token RMS downscale as _cam_branch_softmax, so cached
# chunks stay on the dense path's numerics.
q_dn = _downscale_to_reference_rms(
q_bhnd.permute(0, 1, 3, 2), q_proj.permute(0, 1, 3, 2)
)
k_dn = _downscale_to_reference_rms(
k_bhnd.permute(0, 1, 3, 2), k_proj.permute(0, 1, 3, 2)
)
v_dn = _downscale_to_reference_rms(
v_bhnd.permute(0, 1, 3, 2), v_proj.permute(0, 1, 3, 2)
)
q_in = q_dn.permute(0, 3, 1, 2).contiguous() # (B, N_cur, H, D)
k_in = k_dn.permute(0, 3, 1, 2).contiguous()
@@ -396,15 +396,26 @@ class ControlStateQueue:
self._current_item = default_item
self._current_seq_id: int | None = None
self._latest_sampled_seq_id: int | None = None
# A queue that never received a transition has no state to hold, so
# sampling reports None and the caller omits the control entirely;
# after the first transition the level stays held, default included.
self._received_any = False
def clear(self) -> None:
self._pending.clear()
self._current_item = self.default_item
self._current_seq_id = None
self._latest_sampled_seq_id = None
self._received_any = False
def push(self, transition: ControlStateTransition) -> None:
self._pending.append(transition)
self._received_any = True
def mark_received(self) -> None:
"""Arm the hold without a transition: a replacing script is itself a
signal, so a drained script keeps sampling its end state as held."""
self._received_any = True
def push_many(self, transitions: Sequence[ControlStateTransition]) -> None:
for transition in transitions:
@@ -416,6 +427,8 @@ class ControlStateQueue:
transitions = self._drain_pending()
if not transitions:
if not self._received_any:
return None
self._latest_sampled_seq_id = self._current_seq_id
return [self._copy_item(self._current_item) for _ in range(chunk_size)]
@@ -77,6 +77,9 @@ class RealtimeCameraControlState:
) -> None:
"""Replace active controls with a finite per-frame script."""
self.camera_state_queue.clear()
# the script is itself a signal: when it drains, the state holds at
# released rather than reporting no-signal
self.camera_state_queue.mark_received()
self.camera_script_queue.push_script(
[list(actions) for actions in camera_actions],
event_id=event_id,
@@ -1268,9 +1268,9 @@ ONE_GPU_XPU_CASES = _select_xpu_cases(ONE_GPU_XPU_CASE_IDS)
# Nested unit/ tests verified to pass on AMD/ROCm as-is (no code change).
# Enabled incrementally and AMD-only: the CUDA `multimodal-gen-unit-test`
# lane keeps the flat glob below. Files that still need fixes/skips are added
# in follow-up PRs. Paths are relative to the unit/ dir.
# The CUDA lane runs the whole unit/ tree recursively; AMD enables nested
# files incrementally as they are vetted, in follow-up PRs. Paths are
# relative to the unit/ dir.
_AMD_READY_NESTED_UNIT_TESTS = (
"realtime/test_causal_denoising.py",
"realtime/test_output_materialization.py",
@@ -1293,12 +1293,13 @@ def _discover_unit_tests() -> list[str]:
unit_dir = Path(__file__).resolve().parent.parent / "unit"
if not unit_dir.is_dir():
return []
# Flat unit/ tests run on every lane (unchanged). This keeps the CUDA
# `multimodal-gen-unit-test` job byte-identical.
flat = [f"../unit/{f.name}" for f in unit_dir.glob("test_*.py") if f.is_file()]
if not current_platform.is_hip():
return sorted(flat)
# AMD/ROCm additionally runs the vetted nested-subdir tests.
# pytest recurses into the directory: every unit test runs, nested
# subdirs included, and new files need no registry edit.
return ["../unit"]
# AMD/ROCm keeps the vetted set: the flat files plus the nested tests
# verified to pass on ROCm.
flat = [f"../unit/{f.name}" for f in unit_dir.glob("test_*.py") if f.is_file()]
nested = [
f"../unit/{rel}"
for rel in _AMD_READY_NESTED_UNIT_TESTS
@@ -158,7 +158,7 @@ def test_realtime_session_cache_rejects_missing_nonzero_chunk():
def test_lingbot_realtime_state_uses_control_script_and_prompt_queues():
state = lingbot_realtime.LingBotWorldRealtimeState()
assert state.sample_camera_actions(3) == [[], [], []]
assert state.sample_camera_actions(3) is None
state.receive_camera_action_script([["w"], ["a"], ["s"], ["d"]])
assert state.sample_camera_actions(3) == [["w"], ["a"], ["s"]]
assert state.sample_camera_actions(3) == [["d"], [], []]
@@ -1008,9 +1008,7 @@ def test_realtime_chunk_latent_preparation_uses_chunk_spec():
)
transformer = SimpleNamespace(
config=SimpleNamespace(
arch_config=SimpleNamespace(out_channels=16, num_frames_per_block=3)
)
config=SimpleNamespace(out_channels=16, num_frames_per_block=3)
)
stage = RealtimeChunkLatentPreparationStage.__new__(
RealtimeChunkLatentPreparationStage
@@ -14,6 +14,10 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts():
styles_css = (
repo_root / "python/sglang/multimodal_gen/apps/realtime_webui/styles.css"
).read_text()
playback_js = (
repo_root
/ "python/sglang/multimodal_gen/apps/realtime_webui/playback_controller.js"
).read_text()
assert "preset.actions" not in app_js
assert "repeatActions" not in app_js
@@ -48,9 +52,11 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts():
assert "Info" not in index_html
assert 'id="steps" type="number" value="4"' in index_html
assert 'id="guidance" type="number" value="1"' in index_html
assert "styles.css?v=realtime-sr-v38" in index_html
assert "app.js?v=realtime-sr-v38" in index_html
assert 'const DECODER_WORKER_URL = "./decoder_worker.js?v=rgb-worker-v6";' in app_js
assert "styles.css?v=realtime-record-v49" in index_html
assert "app.js?v=realtime-record-v75" in index_html
assert (
'const DECODER_WORKER_URL = "./decoder_worker.js?v=rgb-worker-v10";' in app_js
)
assert "const DEFAULT_TARGET_FPS = 25;" in app_js
assert "const DEFAULT_FRAME_INTERPOLATION_EXP = 1;" in app_js
assert "const DEFAULT_FRAME_INTERPOLATION_SCALE = 1.0;" in app_js
@@ -73,18 +79,16 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts():
assert "setPreviewScale(DEFAULT_PREVIEW_SCALE)" in app_js
assert "preview_scale" in app_js
assert "sr_scale" in app_js
assert "elapsedMs % targetMs" in app_js
assert "liveQueueFrameFloor(header, chunkFrameCount)" in app_js
assert "elapsedMs < targetMs" in playback_js
assert "queuedDecodeFrames > maxQueuedFrames" in app_js
assert (
'const REACTOR_PRESET_BASE_URL = "https://www.reactor.inc/lingbot-world-fast-v1";'
in app_js
)
assert "Dragon Dolly" in app_js
assert "no creature morphing" in app_js
assert "A static locked-off view of the back side of Plastic Beach" in app_js
assert "clouds slowly drifting behind the island" in app_js
assert "occasional shooting star" in app_js
assert "tiny distant pigeons" in app_js
assert "the Plastic Beach island stays centered" in app_js
assert "no camera descent, no push-in, no orbit" in app_js
assert "Ziggy Stardust" in app_js
assert "blue K. West sign" in app_js
assert "wet pavement reflecting a yellow streetlamp" in app_js
@@ -95,8 +99,8 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts():
assert app_js.index("Dragon Dolly") < app_js.index("Kid A")
assert "dragon-ride.jpg" in app_js
assert "stageRenderFps" not in app_js
assert 'setStatus("Receiving"' not in app_js
assert "decodeChain = decodeChain" in app_js
assert 'setStatus("Receiving", "live")' in app_js
assert "decodeQueue.push(" in app_js
assert "receiveChain" not in app_js
assert 'message.type === "chunk_stats"' in app_js
assert "chunkTotal > 0 ? numFrames / chunkTotal" in app_js
@@ -104,6 +108,5 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts():
assert ".workspace" in styles_css
assert ".preview-frame" in styles_css
assert ".preview-overlay" in styles_css
assert "@keyframes previewSweep" in styles_css
assert ".preview-scale-control" in styles_css
assert "--preview-scale" in styles_css
@@ -60,6 +60,7 @@ def test_remote_file_exists_returns_false_for_definitive_404(monkeypatch):
def test_remote_video_gt_candidates_survive_inconclusive_probe(monkeypatch):
monkeypatch.setenv(test_utils.CONSISTENCY_PLATFORM_ENV, "h100")
monkeypatch.setattr(test_utils, "_remote_file_exists", lambda url: None)
files = test_utils._find_remote_consistency_gt_files(