#!/usr/bin/env python3 """Measure local and downstream impact of isolated Layer-17 Predictor calls. For every prompt, chunk 1..6 is independently evaluated with FPFF, FFPF, and FPPF while every other chunk remains FFFF. A shadow Full forward is executed at each selected Predictor step to measure hidden/flow/x0 error on the exact rollout state. The shadow output is never used by the generated trajectory. """ from __future__ import annotations import argparse import csv import json import math import os import sys import time from pathlib import Path from typing import Any def _preparse_gpu() -> str: parser = argparse.ArgumentParser(add_help=False) parser.add_argument("--gpu", default="4") args, _ = parser.parse_known_args() os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) return str(args.gpu) PHYSICAL_GPU = _preparse_gpu() import lpips import torch from omegaconf import OmegaConf REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from scripts import evaluate_single_block_fppf as base from utils.misc import set_seed from utils.wan_wrapper import WanVAEWrapper SCHEDULES = ("FPFF", "FFPF", "FPPF") EPS = 1e-8 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--gpu", default=PHYSICAL_GPU) parser.add_argument( "--config_path", type=Path, default=Path("configs/self_forcing_sid.yaml") ) parser.add_argument( "--checkpoint_path", type=Path, default=Path("checkpoints/self_forcing_dmd.pt"), ) parser.add_argument( "--dataset_root", type=Path, default=Path("outputs/predictor_offline_100_all_blocks"), ) parser.add_argument( "--sweep_dir", type=Path, default=Path("outputs/single_block_init_sweep"), ) parser.add_argument( "--reference_root", type=Path, default=Path("outputs/single_block_fppf_eval"), ) parser.add_argument( "--output_dir", type=Path, default=Path("outputs/layer17_chunk_impact_pilot"), ) parser.add_argument( "--prompt_ids", type=int, nargs="*", default=list(range(80, 90)) ) parser.add_argument("--schedules", nargs="*", choices=SCHEDULES, default=list(SCHEDULES)) parser.add_argument( "--chunks", type=int, nargs="*", default=list(range(1, base.NUM_CHUNKS)) ) parser.add_argument("--max_prompts", type=int, default=None) parser.add_argument("--metric_batch_size", type=int, default=4) parser.add_argument("--generation_seed", type=int, default=0) parser.add_argument( "--skip_lpips", action=argparse.BooleanOptionalAction, default=False ) parser.add_argument("--overwrite", action="store_true") args = parser.parse_args() if not args.prompt_ids: parser.error("At least one prompt ID is required") if any(value < 0 or value >= 100 for value in args.prompt_ids): parser.error("Prompt IDs must be in [0, 99]") if args.metric_batch_size < 1: parser.error("--metric_batch_size must be positive") if not args.schedules: parser.error("At least one schedule is required") if not args.chunks or any(chunk < 1 or chunk >= base.NUM_CHUNKS for chunk in args.chunks): parser.error("--chunks must contain values in [1, 6]") return args def resolve(path: Path) -> Path: return path.resolve() if path.is_absolute() else (REPO_ROOT / path).resolve() def rms(value: torch.Tensor) -> torch.Tensor: return value.float().square().mean().sqrt() def nrmse(prediction: torch.Tensor, target: torch.Tensor) -> float: return float(rms(prediction.float() - target.float()) / rms(target).clamp_min(EPS)) def chunk_frame_slice(chunk: int) -> slice: if chunk == 0: return slice(0, base.PIXEL_FRAMES_FIRST_CHUNK) start = base.PIXEL_FRAMES_FIRST_CHUNK + 12 * (chunk - 1) return slice(start, start + 12) def summarize_frame_range(metrics: dict[str, Any], selected: slice) -> dict[str, float]: mse_values = metrics["mse_per_frame"][selected] ssim_values = metrics["ssim_per_frame"][selected] lpips_values = metrics["lpips_per_frame"][selected] mean_mse = sum(mse_values) / len(mse_values) return { "pixel_mse": mean_mse, "psnr": -10.0 * math.log10(max(mean_mse, 1e-12)), "ssim": sum(ssim_values) / len(ssim_values), "lpips": ( sum(lpips_values) / len(lpips_values) if lpips_values else float("nan") ), } @torch.inference_mode() def generate_intervention( *, pipeline: Any, dataset_root: Path, prompt_id: int, generation_seed: int, device: torch.device, predictor: Any, source_layer: int, intervention_chunk: int, intervention_schedule: str, ) -> tuple[torch.Tensor, dict[str, Any]]: if intervention_schedule not in SCHEDULES: raise ValueError(intervention_schedule) if intervention_chunk < 1 or intervention_chunk >= base.NUM_CHUNKS: raise ValueError("Predictor intervention chunk must be 1..6") base.reset_kv_and_load_cross_cache(pipeline, dataset_root, prompt_id, device) set_seed(generation_seed) noise = torch.randn( 1, base.NUM_CHUNKS * base.FRAMES_PER_CHUNK, base.LATENT_CHANNELS, base.LATENT_HEIGHT, base.LATENT_WIDTH, dtype=torch.bfloat16, device=device, ) timesteps = pipeline.denoising_step_list.to(device=device) output_chunks: list[torch.Tensor] = [] previous_chunk_hidden: list[torch.Tensor | None] | None = None teacher = pipeline.generator.model capture = base.FinalHiddenCapture(teacher) full_calls = 0 predictor_calls = 0 shadow_full_calls = 0 local_errors: list[dict[str, float | int]] = [] started = time.perf_counter() try: for chunk in range(base.NUM_CHUNKS): noisy_input = noise[ :, chunk * base.FRAMES_PER_CHUNK : (chunk + 1) * base.FRAMES_PER_CHUNK, ] current_hidden: list[torch.Tensor | None] = [None] * base.NUM_DENOISING_STEPS denoised_pred: torch.Tensor | None = None timestep: torch.Tensor | None = None for step, current_timestep in enumerate(timesteps): timestep = torch.ones( [1, base.FRAMES_PER_CHUNK], dtype=torch.int64, device=device ) * current_timestep selected = ( chunk == intervention_chunk and intervention_schedule[step] == "P" ) if selected: anchor_hidden = current_hidden[step - 1] if anchor_hidden is None or previous_chunk_hidden is None: raise RuntimeError("Predictor inputs are unavailable") previous_hidden = previous_chunk_hidden[step] if previous_hidden is None: raise RuntimeError("Previous-chunk hidden is unavailable") history = pipeline.kv_cache1[source_layer] cross = pipeline.crossattn_cache[source_layer] pred_hidden, pred_flow, _ = base.predictor_step( predictor=predictor, teacher=teacher, noisy_input=noisy_input, timestep=timestep, anchor_hidden=anchor_hidden, previous_hidden=previous_hidden, history_cache=history, cross_cache=cross, current_start=chunk * base.TOKENS_PER_CHUNK, ) pred_x0 = pipeline.generator._convert_flow_pred_to_x0( flow_pred=pred_flow.flatten(0, 1), xt=noisy_input.flatten(0, 1), timestep=timestep.flatten(0, 1), ).unflatten(0, pred_flow.shape[:2]) # Shadow Full measures the exact counterfactual target on # this rollout state. Its x0/hidden are never accepted. capture.start() full_flow, full_x0 = pipeline.generator( noisy_image_or_video=noisy_input, conditional_dict={ "prompt_embeds": torch.zeros( 1, 1, int(teacher.text_embedding[0].in_features), dtype=torch.bfloat16, device=device, ) }, timestep=timestep, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=chunk * base.TOKENS_PER_CHUNK, ) full_hidden = capture.finish() local_errors.append( { "step": step, "timestep": float(current_timestep), "hidden_nrmse": nrmse(pred_hidden, full_hidden), "flow_nrmse": nrmse(pred_flow, full_flow), "x0_nrmse": nrmse(pred_x0, full_x0), } ) current_hidden[step] = pred_hidden denoised_pred = pred_x0 predictor_calls += 1 shadow_full_calls += 1 del full_flow, full_x0, full_hidden else: capture.start() _, denoised_pred = pipeline.generator( noisy_image_or_video=noisy_input, conditional_dict={ "prompt_embeds": torch.zeros( 1, 1, int(teacher.text_embedding[0].in_features), dtype=torch.bfloat16, device=device, ) }, timestep=timestep, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=chunk * base.TOKENS_PER_CHUNK, ) current_hidden[step] = capture.finish() full_calls += 1 if step < base.NUM_DENOISING_STEPS - 1: if denoised_pred is None: raise RuntimeError("Denoising step produced no x0") next_timestep = timesteps[step + 1] flat = denoised_pred.flatten(0, 1) noisy_input = pipeline.scheduler.add_noise( flat, torch.randn_like(flat), next_timestep * torch.ones( [base.FRAMES_PER_CHUNK], dtype=torch.long, device=device ), ).unflatten(0, denoised_pred.shape[:2]) if denoised_pred is None or timestep is None: raise RuntimeError("Chunk produced no clean latent") output_chunks.append(denoised_pred) context_timestep = torch.ones_like(timestep) * pipeline.args.context_noise pipeline.generator( noisy_image_or_video=denoised_pred, conditional_dict={ "prompt_embeds": torch.zeros( 1, 1, int(teacher.text_embedding[0].in_features), dtype=torch.bfloat16, device=device, ) }, timestep=context_timestep, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=chunk * base.TOKENS_PER_CHUNK, ) previous_chunk_hidden = current_hidden finally: capture.close() torch.cuda.synchronize() return torch.cat(output_chunks, dim=1), { "generation_time_s": time.perf_counter() - started, "full_calls": full_calls, "predictor_calls": predictor_calls, "shadow_full_calls": shadow_full_calls, "local_errors": local_errors, } def average(values: list[float]) -> float: return sum(values) / len(values) def finite_average(values: list[Any]) -> float: numeric = [ float(value) for value in values if value is not None and math.isfinite(float(value)) ] return average(numeric) if numeric else float("nan") def rankdata(values: list[float]) -> list[float]: order = sorted(range(len(values)), key=values.__getitem__) ranks = [0.0] * len(values) start = 0 while start < len(order): end = start + 1 while end < len(order) and values[order[end]] == values[order[start]]: end += 1 rank = 0.5 * (start + end - 1) for position in range(start, end): ranks[order[position]] = rank start = end return ranks def pearson(left: list[float], right: list[float]) -> float: left_mean, right_mean = average(left), average(right) left_centered = [value - left_mean for value in left] right_centered = [value - right_mean for value in right] numerator = sum(a * b for a, b in zip(left_centered, right_centered)) denominator = math.sqrt( sum(value * value for value in left_centered) * sum(value * value for value in right_centered) ) return numerator / denominator if denominator > 0 else float("nan") def spearman(left: list[float], right: list[float]) -> float: return pearson(rankdata(left), rankdata(right)) def write_csv(path: Path, rows: list[dict[str, Any]], fields: list[str]) -> None: temporary = path.with_suffix(path.suffix + ".tmp") with temporary.open("w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fields) writer.writeheader() writer.writerows(rows) os.replace(temporary, path) def group_centered_values( records: list[dict[str, Any]], field: str ) -> list[float]: """Remove schedule-by-chunk means to isolate prompt/state variation.""" groups: dict[tuple[str, int], list[float]] = {} for row in records: key = (str(row["schedule"]), int(row["chunk"])) groups.setdefault(key, []).append(float(row[field])) means = {key: average(values) for key, values in groups.items()} return [ float(row[field]) - means[(str(row["schedule"]), int(row["chunk"]))] for row in records ] def aggregate(records: list[dict[str, Any]], output_dir: Path) -> None: numeric_fields = [ "hidden_nrmse_mean", "flow_nrmse_mean", "x0_nrmse_mean", "latent_all_nrmse", "latent_current_nrmse", "latent_tail_nrmse", "psnr", "ssim", "lpips", "current_psnr", "current_ssim", "current_lpips", "tail_psnr", "tail_ssim", "tail_lpips", "generation_time_s", ] summary_rows: list[dict[str, Any]] = [] available_schedules = [ schedule for schedule in SCHEDULES if any(row["schedule"] == schedule for row in records) ] available_chunks = sorted({int(row["chunk"]) for row in records}) for schedule in available_schedules: for chunk in available_chunks: selected = [ row for row in records if row["schedule"] == schedule and row["chunk"] == chunk ] row: dict[str, Any] = { "schedule": schedule, "chunk": chunk, "num_prompts": len(selected), "alpha_eligible_linear": (base.NUM_CHUNKS - 1 - chunk) / (base.NUM_CHUNKS - 2), } for field in numeric_fields: row[field] = finite_average([item[field] for item in selected]) summary_rows.append(row) summary_fields = [ "schedule", "chunk", "num_prompts", "alpha_eligible_linear", *numeric_fields, ] write_csv(output_dir / "summary_by_schedule_chunk.csv", summary_rows, summary_fields) correlation_rows: list[dict[str, Any]] = [] for schedule in (*available_schedules, "ALL"): selected = ( records if schedule == "ALL" else [r for r in records if r["schedule"] == schedule] ) for local in ("hidden_nrmse_mean", "flow_nrmse_mean", "x0_nrmse_mean"): for downstream in ("tail_lpips", "latent_tail_nrmse", "tail_pixel_mse"): pairs = [ (float(row[local]), float(row[downstream])) for row in selected if row[local] is not None and row[downstream] is not None and math.isfinite(float(row[local])) and math.isfinite(float(row[downstream])) ] correlation_rows.append( { "schedule": schedule, "local_metric": local, "downstream_metric": downstream, "num_observations": len(pairs), "spearman": ( spearman( [pair[0] for pair in pairs], [pair[1] for pair in pairs], ) if len(pairs) >= 2 else float("nan") ), } ) correlation_fields = [ "schedule", "local_metric", "downstream_metric", "num_observations", "spearman", ] write_csv(output_dir / "local_downstream_correlations.csv", correlation_rows, correlation_fields) controlled_rows: list[dict[str, Any]] = [] for local in ("hidden_nrmse_mean", "flow_nrmse_mean", "x0_nrmse_mean"): local_residual = group_centered_values(records, local) downstream_residual = group_centered_values(records, "tail_lpips") within_cell = [] for schedule in available_schedules: for chunk in available_chunks: selected = [ row for row in records if row["schedule"] == schedule and row["chunk"] == chunk ] if len(selected) >= 2: within_cell.append( spearman( [float(row[local]) for row in selected], [float(row["tail_lpips"]) for row in selected], ) ) controlled_rows.append( { "local_metric": local, "downstream_metric": "tail_lpips", "controls": "schedule+chunk", "residual_spearman": spearman(local_residual, downstream_residual), "mean_within_cell_spearman": finite_average(within_cell), "positive_cells": sum(value > 0 for value in within_cell), "num_cells": len(within_cell), } ) controlled_fields = [ "local_metric", "downstream_metric", "controls", "residual_spearman", "mean_within_cell_spearman", "positive_cells", "num_cells", ] write_csv( output_dir / "controlled_local_downstream_correlations.csv", controlled_rows, controlled_fields, ) report = [ "# Layer-17 Predictor chunk-impact pilot", "", f"Prompts: {len(set(row['prompt_id'] for row in records))}; seed 0; " "all non-intervened chunks use FFFF.", "", "A shadow Full call measures local error at each Predictor decision, but the " "generated trajectory always consumes the Predictor output.", "", ] for schedule in available_schedules: report.extend( [ f"## {schedule}", "", "| Chunk | x0 nRMSE | Tail PSNR | Tail LPIPS | Tail latent nRMSE |", "|---:|---:|---:|---:|---:|", ] ) for row in summary_rows: if row["schedule"] != schedule: continue report.append( f"| {row['chunk']} | {row['x0_nrmse_mean']:.6f} | " f"{row['tail_psnr']:.4f} | {row['tail_lpips']:.6f} | " f"{row['latent_tail_nrmse']:.6f} |" ) report.append("") report.extend( [ "## Local-to-downstream Spearman correlations", "", "| Schedule | Local metric | Downstream metric | Spearman |", "|---|---|---|---:|", ] ) for row in correlation_rows: if row["downstream_metric"] == "tail_lpips": report.append( f"| {row['schedule']} | {row['local_metric']} | tail LPIPS | " f"{row['spearman']:.4f} |" ) report.extend( [ "", "## Correlations after controlling schedule and chunk", "", "Residual correlations remove each schedule-by-chunk mean, so they test " "whether local error explains prompt/state risk beyond the position prior.", "", "| Local metric | Residual Spearman | Mean within-cell Spearman | Positive cells |", "|---|---:|---:|---:|", ] ) for row in controlled_rows: report.append( f"| {row['local_metric']} | {row['residual_spearman']:.4f} | " f"{row['mean_within_cell_spearman']:.4f} | " f"{row['positive_cells']}/{row['num_cells']} |" ) (output_dir / "REPORT.md").write_text("\n".join(report) + "\n", encoding="utf-8") def main() -> None: args = parse_args() for name in ( "config_path", "checkpoint_path", "dataset_root", "sweep_dir", "reference_root", "output_dir", ): setattr(args, name, resolve(getattr(args, name))) args.output_dir.mkdir(parents=True, exist_ok=True) prompt_ids = sorted(set(args.prompt_ids)) if args.max_prompts is not None: prompt_ids = prompt_ids[: args.max_prompts] config = OmegaConf.merge( OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"), OmegaConf.load(args.config_path), ) device = torch.device("cuda") torch.set_grad_enabled(False) set_seed(args.generation_seed) manifest = { "status": "running", "gpu": str(args.gpu), "prompt_ids": prompt_ids, "generation_seed": args.generation_seed, "checkpoint_path": str(args.checkpoint_path), "predictor_weights": str( args.sweep_dir / "teacher_layer_17" / "predictor_final.safetensors" ), "dataset_root": str(args.dataset_root), "schedules": list(args.schedules), "chunks": list(args.chunks), "shadow_full_target": True, } base.atomic_json(args.output_dir / "manifest.json", manifest) print("[setup] loading VAE", flush=True) vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval() missing_references = [ prompt_id for prompt_id in prompt_ids if not ( args.reference_root / "ffff_reference_frames" / f"prompt_{prompt_id:04d}.safetensors" ).exists() ] if missing_references: base.prepare_reference_frames( vae=vae, dataset_root=args.dataset_root, output_dir=args.reference_root, prompt_ids=missing_references, device=device, rebuild=False, ) print("[setup] loading frozen generator and Layer-17 Predictor", flush=True) pipeline = base.build_pipeline(config, args.checkpoint_path, vae, device) experiment = base.discover_experiments( args.sweep_dir, ["teacher_layer_17"], None )[0] predictor = base.load_predictor(pipeline.generator.model, experiment, device) lpips_model = None if not args.skip_lpips: lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval() lpips_model.requires_grad_(False) records: list[dict[str, Any]] = [] total = len(prompt_ids) * len(args.schedules) * len(args.chunks) completed = 0 for prompt_id in prompt_ids: reference_latent = base.load_ffff_latent(args.dataset_root, prompt_id).to( device=device, dtype=torch.bfloat16 ) reference_u8 = base.load_reference_frames(args.reference_root, prompt_id) for schedule in args.schedules: for chunk in args.chunks: destination = ( args.output_dir / "per_intervention" / f"prompt_{prompt_id:04d}_{schedule}_chunk_{chunk:02d}.json" ) if destination.exists() and not args.overwrite: record = json.loads(destination.read_text(encoding="utf-8")) records.append(record) completed += 1 print(f"[cached] {completed}/{total} {destination.stem}", flush=True) continue started = time.perf_counter() latent, diagnostics = generate_intervention( pipeline=pipeline, dataset_root=args.dataset_root, prompt_id=prompt_id, generation_seed=args.generation_seed, device=device, predictor=predictor, source_layer=17, intervention_chunk=chunk, intervention_schedule=schedule, ) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): pixels = vae.decode_to_pixel(latent, use_cache=False) prediction_u8 = base.pixels_to_u8(pixels) frame = base.frame_metrics( reference_u8=reference_u8, prediction_u8=prediction_u8, lpips_model=lpips_model, batch_size=args.metric_batch_size, device=device, ) current_slice = chunk_frame_slice(chunk) current = summarize_frame_range(frame, current_slice) tail = summarize_frame_range(frame, slice(current_slice.start, None)) latent_chunk_start = chunk * base.FRAMES_PER_CHUNK latent_chunk_end = latent_chunk_start + base.FRAMES_PER_CHUNK errors = diagnostics.pop("local_errors") record = { "prompt_id": prompt_id, "schedule": schedule, "chunk": chunk, "predictor_steps": [int(item["step"]) for item in errors], "hidden_nrmse_mean": average( [float(item["hidden_nrmse"]) for item in errors] ), "flow_nrmse_mean": average( [float(item["flow_nrmse"]) for item in errors] ), "x0_nrmse_mean": average( [float(item["x0_nrmse"]) for item in errors] ), "local_errors": errors, "latent_all_nrmse": nrmse(latent, reference_latent), "latent_current_nrmse": nrmse( latent[:, latent_chunk_start:latent_chunk_end], reference_latent[:, latent_chunk_start:latent_chunk_end], ), "latent_tail_nrmse": nrmse( latent[:, latent_chunk_start:], reference_latent[:, latent_chunk_start:], ), "psnr": frame["psnr"], "ssim": frame["ssim"], "lpips": frame["lpips"], "current_pixel_mse": current["pixel_mse"], "current_psnr": current["psnr"], "current_ssim": current["ssim"], "current_lpips": current["lpips"], "tail_pixel_mse": tail["pixel_mse"], "tail_psnr": tail["psnr"], "tail_ssim": tail["ssim"], "tail_lpips": tail["lpips"], **diagnostics, "total_time_s": time.perf_counter() - started, } base.atomic_json(destination, record) records.append(record) completed += 1 print( f"[run] {completed}/{total} p={prompt_id} {schedule} c={chunk} " f"x0={record['x0_nrmse_mean']:.5f} " f"tail_lpips={record['tail_lpips']:.5f} " f"time={record['total_time_s']:.1f}s", flush=True, ) if hasattr(vae.model, "clear_cache"): vae.model.clear_cache() del latent, pixels, prediction_u8, frame torch.cuda.empty_cache() del reference_latent, reference_u8 fields = sorted({key for record in records for key in record if key != "local_errors"}) flattened = [{key: row.get(key) for key in fields} for row in records] write_csv(args.output_dir / "interventions.csv", flattened, fields) aggregate(records, args.output_dir) manifest["status"] = "complete" base.atomic_json(args.output_dir / "manifest.json", manifest) print(f"[complete] results -> {args.output_dir}", flush=True) if __name__ == "__main__": main()