[ci] Cherrypick perf regression skills to component time commit

[ci] Add performance baseline reseed skill and failed-run artifacts

.agents/skills/index.jsonl

@@ -7,4 +7,4 @@
 {"name": "search-related-work", "description": "Query the related work index for relevant papers, repos, or comparisons", "path": "search-related-work/SKILL.md", "status": "draft", "trust": "low"}
 {"name": "seed-ssim-references", "description": "Run a new or updated fastvideo/tests/ssim/ test on Modal, pull generated videos, and upload them to FastVideo/ssim-reference-videos so the test has a regression baseline", "path": "seed-ssim-references/SKILL.md", "status": "draft", "trust": "low"}
 {"name": "reseed-ssim-references", "description": "Re-seed (overwrite) HF reference videos for an existing fastvideo/tests/ssim/ test and a single model id on Modal L40S. Always backs up current refs first, regenerates on Modal, pauses for the user to eyeball before-vs-after, then uploads with --force scoped to --model-id. Sister skill to seed-ssim-references; use when intentional code change has invalidated existing refs", "path": "reseed-ssim-references/SKILL.md", "status": "draft", "trust": "low"}
-{"name": "add-model", "description": "Add a new model (or variant) to FastVideo: DiT + configs + pipeline + presets + registry + tests. Walks through FastVideo's single stage-based pipeline architecture with exact file paths and registration hooks.", "path": "add-model/SKILL.md", "status": "draft", "trust": "low"}
+{"name": "reseed-performance-baseline", "description": "Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, or environment-caused benchmark shift. Use when performance CI fails because metrics such as latency, throughput, component time, or peak memory changed for an accepted reason and the rolling median baseline must be advanced by replicating one reviewed shifted source result into three success=true records, or five records when explicitly requested", "path": "reseed-performance-baseline/SKILL.md", "status": "draft", "trust": "low"}

.agents/skills/reseed-performance-baseline/SKILL.md

@@ -0,0 +1,426 @@
+---
+name: reseed-performance-baseline
+description: Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, or environment-caused benchmark shift. Use when performance CI fails because metrics such as latency, throughput, component time, or peak memory changed for an accepted reason and the rolling median baseline in FastVideo/performance-tracking must be advanced by replicating one reviewed shifted source result into three success=true records, or five records when explicitly requested.
+---
+
+# Re-seed Performance Baseline
+
+## Purpose
+
+Replace or advance the rolling performance baseline for a single
+`(model_id, gpu_type)` pair in the HF dataset
+`FastVideo/performance-tracking`.
+
+Performance comparison uses the median of up to the last 5 successful records
+for the same model and GPU. Failed records are useful audit history, but they
+do not move the future baseline because `compare_baseline.py` loads records
+with `successful_only=True`.
+
+For a 5-record median, one shifted record is not enough to move the median if
+the other four records are from the old runtime. This skill therefore creates
+3 reviewed `success=true` records from one accepted shifted source result by
+default. If the user explicitly asks for a full reset, create 5 records.
+
+These replicated records are an intentional operator-approved baseline reset,
+not independent measurements. Mark them clearly with provenance fields so the
+HF history remains auditable.
+
+Use this skill when a performance test fails for an intentional and reviewed
+reason, such as a torch/runtime/container upgrade that legitimately increases
+peak memory or changes timings. This is the performance equivalent of
+`reseed-ssim-references`: backup first, scope tightly, require explicit human
+approval, then upload reviewed accepted baseline records.
+
+## When to use
+
+- A PR or main run failed the rolling performance comparison by more than the
+  allowed regression threshold, and maintainers agree the shift is caused by
+  an intentional runtime, dependency, hardware image, or benchmark environment
+  change rather than a FastVideo logic regression.
+- One shifted source result has been reviewed and accepted, and the operator
+  wants to replicate it into 3 successful records so the rolling median moves
+  immediately. Use 5 records only when the user explicitly asks to fully reset
+  the last-5 window.
+
+## When not to use
+
+- The benchmark failure might be a real code regression. Fix or investigate
+  the code path first.
+- The fixed benchmark thresholds in
+  `.buildkite/performance-benchmarks/tests/*.json` are too low. Those are a
+  separate gate from the rolling HF baseline and may need a code review change.
+- There is no clear source run, commit, and rationale. Baseline history is a
+  production signal; do not edit it without provenance.
+
+## Inputs
+
+| Parameter | Required | Description |
+|-----------|----------|-------------|
+| `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
+| `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
+| `source_result` | Yes | Path or Buildkite artifact URL for one accepted shifted performance JSON. Prefer the normalized `normalized_perf_*.json` artifact emitted by `compare_baseline.py`. |
+| `replica_count` | No | Number of success records to create from `source_result`. Default: `3`. Only use `5` if the user explicitly asks for a full reset. |
+| `intent_rationale` | Yes | One-line explanation for why the baseline shift is legitimate. This is written into provenance and should be reused in the PR. |
+
+Hardcoded defaults:
+
+- HF repo: `FastVideo/performance-tracking` (`HF_REPO_ID` override is
+  supported by the code, but use the default unless the user explicitly asks).
+- Local sync root: `/tmp/perf-tracking` or a timestamped local backup under
+  `performance_reseed_backup/`.
+- Baseline window: last 5 `success=true` records for the same
+  `(model_id, gpu_type)`.
+- Default reseed count: 3 replicated `success=true` records from one reviewed
+  source result. Explicit full-reset count: 5.
+
+## Steps
+
+### 1. Validate the target and source result
+
+If `source_result` is a Buildkite artifact URL, download it first into a
+local scratch directory such as `performance_reseed_source/` and use that
+downloaded JSON path for the rest of the workflow. If the agent cannot access
+the artifact because Buildkite authentication is missing, ask the user to
+download the artifact manually and provide the local path.
+
+Prefer the normalized Buildkite artifact emitted by `compare_baseline.py`:
+
+```text
+perf_reports/results/normalized_perf_*.json
+```
+
+That file is already in the HF tracking schema. Load it directly and confirm
+it has the expected baseline fields:
+
+```python
+import json
+
+with open(source_result, encoding="utf-8") as f:
+    record = json.load(f)
+```
+
+If only the older raw `fastvideo/tests/performance/results/perf_*.json`
+artifact is available, normalize it with `compare_baseline.py`'s shared helper
+before continuing. Run this from the repository root with
+`PYTHONPATH=fastvideo/tests/performance` so the script-local `hf_store` import
+resolves the same way it does in CI:
+
+```python
+import json
+from compare_baseline import normalize_performance_result
+
+with open(source_result, encoding="utf-8") as f:
+    record = normalize_performance_result(json.load(f))
+```
+
+The raw-to-normalized helper maps:
+
+- `model_id` comes from `benchmark_id`.
+- `gpu_type` comes from `device`.
+- `memory` comes from `max_peak_memory_mb`.
+- `latency` comes from `avg_generation_time_s`.
+- `throughput` comes from `throughput_fps`.
+- component timings come from the raw `text_encoder_time_s`, `dit_time_s`,
+  and `vae_decode_time_s` fields when present. If an older raw artifact lacks
+  those keys, they normalize to `None`; that source can still reseed latency,
+  throughput, and memory, but it cannot move component-time baselines.
+
+Stop if the normalized record's `model_id` or `gpu_type` does not match the
+requested `model_id` and `gpu_type`.
+
+The source record may have `success: false` when it came from a failed rolling
+baseline comparison. That is expected; only the reviewed reseed replicas become
+new `success: true` baseline records after explicit approval.
+
+Set `replica_count` to `3` by default. Set it to `5` only when the user
+explicitly asks to upload the same shifted source result 5 times for a full
+last-5 reset. Reject other counts unless the user gives a concrete reason.
+
+Check that `HF_API_KEY` is exported. The sync path may be public, but the
+upload path requires write access.
+
+### 1a. How to obtain `source_result` from CI
+
+The performance CI exports normalized source results for failed rolling
+baseline comparisons when `compare_baseline.py` ran. The preferred artifact
+comes from:
+
+```text
+perf_reports/results/normalized_perf_*.json
+```
+
+and is uploaded by Buildkite with the performance reports. The normal operator
+flow is:
+
+1. Open the failed Buildkite performance job.
+2. Download the `normalized_perf_*.json` artifact for the failed benchmark.
+3. Pass the local path or artifact URL as `source_result`.
+
+Do not scrape the Markdown performance summary to reconstruct the JSON. The
+normalized JSON artifact is the source of truth for reseed metrics and
+provenance. If only a raw `fastvideo/tests/performance/results/perf_*.json`
+artifact is present, normalize it with `normalize_performance_result()` before
+continuing. If no JSON artifact is present, the benchmark likely failed before
+writing results, so that run is not a valid source for baseline reseeding.
+
+### 2. Sync and back up existing HF records
+
+Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
+`compare_baseline.py` as a sync shortcut; on full main runs it can persist
+records, while this step must only fetch and back up existing history.
+
+The sync command pattern is:
+
+```bash
+export PERFORMANCE_TRACKING_ROOT="${PERFORMANCE_TRACKING_ROOT:-/tmp/perf-tracking}"
+export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
+PYTHONPATH=fastvideo/tests/performance python -c 'from hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
+```
+
+Then back up only the sanitized model directory:
+
+```bash
+SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
+TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)
+MODEL_SAFE=$(python - <<'PY'
+from fastvideo.tests.performance.hf_store import sanitize
+print(sanitize("<model_id>"))
+PY
+)
+BACKUP_DIR="performance_reseed_backup/${TIMESTAMP}_${SHORT_COMMIT}_${MODEL_SAFE}"
+mkdir -p "$BACKUP_DIR"
+cp -R "${PERFORMANCE_TRACKING_ROOT}/${MODEL_SAFE}" "$BACKUP_DIR/" 2>/dev/null || true
+```
+
+Write provenance next to the backup:
+
+```bash
+cat > "$BACKUP_DIR/PROVENANCE.txt" <<EOF
+model_id: <model_id>
+gpu_type: <gpu_type>
+source_result: <source_result>
+replica_count: <3_or_5>
+head_commit: $(git rev-parse HEAD)
+timestamp_utc: $(date -u +%FT%TZ)
+reason: <intent_rationale>
+EOF
+```
+
+If the backup has no prior records, this is not a destructive reseed; it is a
+first baseline seed. Continue, but report that baseline history was empty.
+
+### 3. Compute old baseline and candidate shift
+
+Load the last 5 successful records for the target:
+
+```python
+from fastvideo.tests.performance.hf_store import load_records_for_model
+
+records = load_records_for_model(
+    "/tmp/perf-tracking",
+    "<model_id>",
+    "<gpu_type>",
+    last_n=5,
+    successful_only=True,
+)
+```
+
+Print a small table showing the source result metrics, the replicated
+candidate median, and the old medians for:
+
+- `latency`
+- `throughput`
+- `memory`
+- `text_encoder_time_s`
+- `dit_time_s`
+- `vae_decode_time_s`
+
+Also print how many successful old records exist. Make clear:
+
+- 1 shifted record only seeds audit history and usually does not move the
+  median.
+- 3 replicated shifted records in a 5-record window move the median
+  immediately.
+- 5 replicated shifted records fully reset the rolling window to the source
+  result's runtime profile.
+- Replicated records are not independent measurements; they are an intentional
+  approved baseline reset and must be labeled that way.
+
+### 4. Confirm intent
+
+Require an explicit confirmation phrase before preparing the upload:
+
+> About to RE-SEED performance baseline for `<model_id>` on `<gpu_type>`.
+> This will upload `<N>` new `success=true` records to
+> `FastVideo/performance-tracking/<sanitize(model_id)>/`.
+>
+> Reason: `<intent_rationale>`
+> Source result: `<source_result>`
+> Replica count: `<replica_count>`
+> Note: these records replicate one reviewed measurement to force the rolling
+> median to the accepted runtime profile.
+> HEAD: `<git rev-parse --short=12 HEAD>`
+> Backup: `<BACKUP_DIR>`
+>
+> Reply `confirm performance reseed` to proceed, anything else to abort.
+
+Do not continue unless the user types exactly `confirm performance reseed`.
+
+### 5. Create the accepted seed records
+
+Create `replica_count` normalized records from the single source result. Use
+an explicit allowlist; do not copy the raw result JSON wholesale.
+
+Each record must include only these baseline fields plus the reseed provenance
+fields below:
+
+- `model_id`
+- `timestamp`
+- `commit_sha`
+- `gpu_type`
+- `latency`
+- `throughput`
+- `memory`
+- `text_encoder_time_s`
+- `dit_time_s`
+- `vae_decode_time_s`
+- `success: true`
+
+For normalized `normalized_perf_*.json` sources, these fields already exist.
+For older raw `perf_*.json` sources, map the raw fields exactly as
+`normalize_performance_result()` in `compare_baseline.py` does:
+
+| Normalized field | Raw source field |
+|------------------|------------------|
+| `model_id` | `benchmark_id` |
+| `gpu_type` | `device` |
+| `latency` | `avg_generation_time_s` |
+| `throughput` | `throughput_fps` |
+| `memory` | `max_peak_memory_mb` |
+| `text_encoder_time_s` | `text_encoder_time_s` |
+| `dit_time_s` | `dit_time_s` |
+| `vae_decode_time_s` | `vae_decode_time_s` |
+| `commit_sha` | `commit` |
+
+Do not upload raw-only fields such as `model_short_name`, `num_gpus`,
+`num_warmup_runs`, `num_measurement_runs`, `individual_times_s`,
+`individual_peak_memories_mb`, `thresholds`, or `pr_number`.
+
+Optional provenance fields are allowed and useful:
+
+- `baseline_reseed: true`
+- `baseline_reseed_reason`
+- `baseline_reseed_source_result`
+- `baseline_reseed_source_timestamp`
+- `baseline_reseed_replicated_source: true`
+- `baseline_reseed_batch_size`
+- `baseline_reseed_batch_index`
+- `baseline_reseed_operator`
+
+Use a fresh reseed timestamp for each replicated record, not the original
+source result timestamp. This is required because
+`load_records_for_model(..., last_n=5)` keeps the last records after loading
+the model directory; stale filenames/timestamps may not enter the last-5
+window and therefore may not move the median. Preserve the original source
+timestamp in `baseline_reseed_source_timestamp`.
+
+Use the existing filename convention from `_write_tracking_record()`:
+`<sanitize(timestamp)>_<sanitize(commit_sha)>.json` under the sanitized model
+directory, but include a deterministic suffix such as `_reseed_01`,
+`_reseed_02`, and `_reseed_03` before `.json` so the replicated files do not
+overwrite each other. For a 5-record full reset, continue through
+`_reseed_05`.
+
+If the source record already exists on HF with `success=false`, do not edit it
+in place unless the user explicitly asked for an audit-preserving correction.
+Prefer uploading new accepted seed records so failed history remains visible.
+
+### 6. Pause before upload
+
+Print:
+
+- Backup directory path.
+- HF paths that will receive the new records.
+- Old rolling medians.
+- Source metrics, replica count, and candidate median.
+- Rationale.
+
+Ask the user to reply exactly `upload`. Anything else aborts and leaves the
+prepared records plus backup on disk.
+
+### 7. Upload only the scoped records
+
+Use the shared storage helper so the path and repo type match CI:
+
+```python
+from fastvideo.tests.performance.hf_store import upload_record
+
+upload_record("<local_record_path>", record, strict=True)
+```
+
+Run it once per prepared record. Each upload goes to:
+
+```text
+FastVideo/performance-tracking/<sanitize(model_id)>/<record_filename>.json
+```
+
+Never bulk upload the whole tracking root. Never modify another model's
+directory in the same operation.
+
+### 8. Report outcome
+
+Report:
+
+- Uploaded HF paths.
+- Backup directory.
+- Old baseline window count and medians.
+- Source metrics, replica count, and candidate median.
+- Expected effect: 3 replicated shifted records move the 5-record median; 5
+  replicated shifted records fully reset the window to the accepted source
+  result.
+- Any separate threshold changes still needed in
+  `.buildkite/performance-benchmarks/tests/*.json`.
+
+Include the `intent_rationale` in the PR or follow-up comment so reviewers can
+distinguish an accepted baseline shift from a hidden regression.
+
+## Failure modes and handling
+
+- **`HF_API_KEY` unset.** Stop before upload. Do not create an untracked
+  process that appears to have reseeded but never reached HF.
+- **Source result does not match target.** Stop. The wrong benchmark or GPU
+  would poison a separate baseline.
+- **`replica_count` is 5 but the user did not explicitly ask for a full
+  reset.** Stop and use the default count of 3.
+- **The source result is noisy or suspicious.** Stop. Replicating one result
+  amplifies that measurement into the baseline, so it must be reviewed first.
+- **HF sync fails.** Stop for destructive reseeds. A stale or empty sync can
+  make the old baseline look missing.
+- **Candidate still violates fixed thresholds.** Report that this skill only
+  handles the rolling HF baseline; update benchmark JSON thresholds in code
+  review if maintainers accept the new absolute limit.
+- **The user aborts at either confirmation.** Leave the backup and prepared
+  records on disk. Nothing should be uploaded.
+- **A bad seed was uploaded.** Use the backup and HF history to identify the
+  uploaded file, then remove or supersede it with an explicitly reviewed
+  corrective record. Do not silently rewrite unrelated history.
+
+## References
+
+- `.agents/skills/reseed-ssim-references/SKILL.md` — safety pattern for
+  intentional baseline replacement.
+- `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
+  median comparison, and persistence rules.
+- `fastvideo/tests/performance/hf_store.py` — HF sync, record loading,
+  `sanitize()`, and `upload_record()`.
+- `fastvideo/tests/performance/test_inference_performance.py` — source result
+  JSON schema.
+- `.buildkite/performance-benchmarks/tests/*.json` — fixed absolute benchmark
+  thresholds, separate from rolling baseline comparisons.
+
+## Changelog
+
+| Date | Change |
+|------|--------|
+| 2026-05-03 | Initial version. Sister workflow to `reseed-ssim-references`, scoped to one performance `(model_id, gpu_type)` baseline seed with backup, confirmation, provenance, and `success=true` upload. |
+| 2026-05-03 | Current policy: replicate one approved shifted source result into 3 success records by default, or 5 only when explicitly requested. Add provenance marker for replicated-source reseeds. |

.buildkite/performance-benchmarks/tests/wan-t2v-1.3b.json

@@ -36,11 +36,17 @@
   "thresholds": {
     "L40S": {
       "max_generation_time_s": 34.0,
-      "max_peak_memory_mb": 11000.0
+      "max_peak_memory_mb": 11000.0,
+      "max_text_encoder_time_s": 5.0,
+      "max_dit_time_s": 10.0,
+      "max_vae_decode_time_s": 10.0
     },
     "default": {
       "max_generation_time_s": 120.0,
-      "max_peak_memory_mb": 30000.0
+      "max_peak_memory_mb": 30000.0,
+      "max_text_encoder_time_s": 5.0,
+      "max_dit_time_s": 10.0,
+      "max_vae_decode_time_s": 10.0
     }
   }
 }

.buildkite/scripts/pr_test.sh

@@ -121,6 +121,19 @@ upload_performance_artifacts() {
         fi
     }
 
+    _upload_normalized_perf_results() {
+        local found=0
+        while IFS= read -r -d '' target; do
+            found=1
+            log "Found normalized performance result: $target. Uploading to Buildkite..."
+            buildkite-agent artifact upload "$target"
+        done < <(find "$LOCAL_DIR" -path "*/results/normalized_perf_*.json" -print0)
+
+        if [ "$found" -eq 0 ]; then
+            log "No normalized performance result artifacts found. This is expected when the rolling performance comparison did not run."
+        fi
+    }
+
     _cleanup_modal_volume() {
         log "Cleaning up perf_reports/ from Modal Volume..."
         if modal volume rm hf-model-weights "perf_reports/" --recursive; then
@@ -139,6 +152,7 @@ upload_performance_artifacts() {
     _download_reports || { _cleanup_local; return 1; }
     _upload_dashboard
     _upload_perf_summary
+    _upload_normalized_perf_results
     _cleanup_modal_volume
     _cleanup_local
 }

.gitignore

@@ -92,11 +92,3 @@ preprocess_output_text/
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
-
-# Local clones of upstream repos used only for parity testing.
-/stable-audio-tools/
-/daVinci-MagiHuman/
-
-# Converted model weights (produced by scripts/checkpoint_conversion/*).
-# Tens of GB; should live on HF, not in git.
-/converted_weights/

docs/contributing/activation_trace.md

@@ -1,210 +0,0 @@
-# Activation Trace Mode
-
-!!! note
-    This page covers Extension 0 (module forward hooks), which is the implemented
-    tracing mechanism. Extensions 1-3 are design sketches for future work and are
-    **not yet implemented**.
-
-## Overview
-
-Activation trace mode is a zero-overhead-when-off, env-gated mechanism for
-dumping per-layer activation statistics during FastVideo inference. Its primary
-use case is **parity debugging across model ports**: enable tracing on both
-FastVideo and the upstream reference implementation, then `diff` the resulting
-JSONL files to find the first divergent layer.
-
-The mechanism is intentionally narrow. It doesn't replace general logging,
-profiling, or function tracing. It answers one question: "at which layer do
-FastVideo and the reference model first produce different numbers?"
-
-## When to use
-
-- Investigating numerical drift between FastVideo and an upstream reference.
-- Debugging mid-pipeline divergence (e.g., one block produces wrong output while earlier blocks match).
-- Validating that a refactor preserves bf16 noise-floor behavior across many layers.
-
-## When NOT to use
-
-| Goal | Use instead |
-|---|---|
-| General logging | `init_logger(__name__)` |
-| Per-stage timing | `FASTVIDEO_STAGE_LOGGING` |
-| Profiling kernel timings | `FASTVIDEO_TORCH_PROFILER_DIR` (see [Profiling](profiling.md)) |
-| Function-call tracing | `FASTVIDEO_TRACE_FUNCTION` (heavy) |
-
-## Quickstart
-
-```bash
-FASTVIDEO_TRACE_ACTIVATIONS=1 \
-FASTVIDEO_TRACE_LAYERS="^block\.layers\.[0-9]+$" \
-FASTVIDEO_TRACE_STATS="abs_mean,sum,max,shape" \
-FASTVIDEO_TRACE_OUTPUT="/tmp/fv_trace.jsonl" \
-python examples/inference/basic/basic_magi_human.py
-```
-
-Each line in `/tmp/fv_trace.jsonl` is a JSON record:
-
-```json
-{"module": "block.layers.0", "tensor": "out", "step": 0, "abs_mean": 1.234, "sum": -5.678, "max": 9.012, "shape": [1, 4096, 5120]}
-```
-
-## Configuration
-
-| Env var | Default | Description |
-|---|---|---|
-| `FASTVIDEO_TRACE_ACTIVATIONS` | `False` | Master toggle. When unset or false, **zero overhead** in the production hot path. |
-| `FASTVIDEO_TRACE_LAYERS` | `""` (all) | Python regex filter applied to `model.named_modules()` names. Empty string matches all modules. |
-| `FASTVIDEO_TRACE_STATS` | `"abs_mean,sum"` | Comma-separated stats to compute. Available: `abs_mean`, `sum`, `min`, `max`, `mean`, `std`, `shape`, `dtype`. |
-| `FASTVIDEO_TRACE_OUTPUT` | `"/tmp/fv_trace_<pid>.jsonl"` | Output file path. `<pid>` is replaced with the process ID at runtime. |
-| `FASTVIDEO_TRACE_STEPS` | `""` (all) | Comma-separated denoising step indices to capture. Empty string captures all steps. |
-
-## Workflow: parity-debug a model port
-
-1. Set up a tightly-controlled comparison: a parity test or a small standalone
-   script that loads both the FastVideo model and the upstream reference with
-   identical inputs and seeds.
-
-2. Run the FastVideo side with tracing on:
-
-   ```bash
-   FASTVIDEO_TRACE_ACTIVATIONS=1 \
-   FASTVIDEO_TRACE_LAYERS="<your regex>" \
-   FASTVIDEO_TRACE_OUTPUT="/tmp/fv_trace_fv.jsonl" \
-   python <fv_runner.py>
-   ```
-
-3. Run the upstream side. The upstream repo needs separate instrumentation. See
-   "Hooking the upstream side" below.
-
-4. Sort both files by `(module, step)` if needed, then diff:
-
-   ```bash
-   diff /tmp/fv_trace_fv.jsonl /tmp/fv_trace_upstream.jsonl
-   ```
-
-5. The first divergent line identifies the first layer where FastVideo and the
-   upstream produce different outputs. Start debugging there.
-
-## Architecture (Extension 0: module forward hooks)
-
-At pipeline initialization, `attach_activation_trace()` reads the env vars once.
-If `FASTVIDEO_TRACE_ACTIVATIONS` is unset or false, the function returns
-immediately and no hooks are registered. If tracing is on, it walks
-`model.named_modules()`, filters by the layer regex, and registers an
-`ActivationStatHook` on each matching module.
-
-During the forward pass, each hook fires after its module completes, computes
-the requested stats on the output tensor, and appends a JSON record to the
-output file.
-
-```
-ComposedPipelineBase
-  └─ attach_activation_trace()
-       ├─ reads env vars (once at startup)
-       ├─ if off: returns None immediately
-       └─ if on: walks named_modules()
-            └─ registers ActivationStatHook on matching modules
-                 └─ on each forward: compute stats → append JSONL
-```
-
-### Zero-overhead-when-off guarantee
-
-- The env var check happens **once at startup** inside `attach_activation_trace()`.
-- If the env var is unset or false, the function returns `None` immediately.
-- No hooks are registered. No branches are added to the production forward path.
-- The only cost when tracing is off is one env var lookup at pipeline
-  initialization, which takes under a microsecond.
-
-### Hooking the upstream side
-
-The upstream reference repo isn't part of FastVideo, so it can't read FastVideo
-env vars directly. Two options:
-
-**Option 1: Inline patch** in your local clone of the upstream repo. Add
-`register_forward_hook` calls in the same shape as `ActivationStatHook`. Clean
-up afterward with `git stash` or `git checkout HEAD -- <file>`.
-
-**Option 2: Wrapper script**. Write a small Python harness that imports the
-upstream model, walks its `named_modules()`, and attaches hooks externally.
-This is the same pattern used in
-`tests/local_tests/transformers/_debug_magi_human_block_parity.py`.
-
-The `add-model-trace` skill at `~/.config/opencode/skill/add-model-trace/`
-provides a script template for this purpose.
-
-## Future extensions (design only, not yet implemented)
-
-### Extension 1: FX/Dynamo backend graph rewrite
-
-**Granularity**: per-FX-node (every matmul, every add).
-
-**Mechanism**: a `torch.compile` backend that takes the captured `GraphModule`
-and inserts logger nodes after each op. Compiles into a separate artifact from
-the production graph.
-
-**Off semantics**: zero overhead. The production compile path is untouched.
-
-**When to add**: if you need to trace inside a `torch.compile`'d graph and
-Extension 0 is too coarse.
-
-**Build cost**: roughly 1-2 days. Reference:
-`torchao.quantization.pt2e._numeric_debugger`.
-
-### Extension 2: AST source injection at import time
-
-**Granularity**: per-line (between any two Python statements).
-
-**Mechanism**: an importlib loader hook rewrites Python source AST at module
-import time, inserting `if TRACE: dump(...)` statements. The decision is made
-once at import.
-
-**Off semantics**: zero overhead. If the env var is off at import time, source
-is loaded as-is.
-
-**When to add**: if you need per-line granularity that even FX-node-level can't
-provide. This is almost never the right choice.
-
-**Build cost**: roughly 1 week. Brittle and hard to debug.
-
-### Extension 3: `__torch_dispatch__` / `TorchDispatchMode`
-
-**Granularity**: per-op (every dispatcher call: matmul, add, view, etc.).
-
-**Mechanism**: a `TorchDispatchMode` context manager that intercepts all ops at
-the dispatcher level.
-
-**Off semantics**: zero overhead. PyTorch's dispatcher only invokes mode hooks
-when a mode is active.
-
-**When on**: significant overhead. Every op pays a Python callback cost. Triton
-kernels bypass it.
-
-**When to add**: useful for quantization or dtype debugging where module-level
-granularity isn't enough.
-
-**Build cost**: roughly 1 day. Reference:
-`torch.utils._python_dispatch.TorchDispatchMode`.
-
-## Comparison with similar tools
-
-| Tool | Pattern | FastVideo equivalent |
-|---|---|---|
-| SGLang `--debug-tensor-dump-output-folder` | env-gated forward hooks at startup | Extension 0 (this) |
-| TransformerEngine `DumpTensors` | config-driven selective dumps | Extension 0 (env-driven) |
-| HuggingFace `output_hidden_states=True` | source-level boolean gating | Not used; Extension 0 avoids model code edits |
-| torchao numeric debugger | FX pass + node-level loggers | Extension 1 (future) |
-| W&B `wandb.watch()` | runtime forward hooks (always on once registered) | Extension 0 has a similar mechanism, but gated off by default |
-
-## Implementation references
-
-- Module: `fastvideo/hooks/activation_trace.py`
-- Env vars: `fastvideo/envs.py` (`FASTVIDEO_TRACE_ACTIVATIONS` and friends)
-- Pipeline integration: `fastvideo/pipelines/composed_pipeline_base.py`
-- Tests: `fastvideo/tests/hooks/test_activation_trace.py`
-- Companion skill (for ad-hoc port investigations): `~/.config/opencode/skill/add-model-trace/`
-
-## Changelog
-
-| Date | Change |
-|---|---|
-| 2026-05-01 | Initial Extension 0 (module forward hooks) implementation. Extensions 1-3 designed but not implemented. |

examples/inference/basic/basic_magi_human.py

@@ -1,51 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Minimal user-runnable example for the daVinci-MagiHuman base AV pipeline.
-
-Produces an mp4 with both video (Wan 2.2 TI2V-5B VAE) and audio (Stable
-Audio Open 1.0 VAE, first-class FastVideo port in
-`fastvideo/models/vaes/oobleck.py`) muxed together via PyAV.
-
-Prerequisites (one-off):
-
-  # Accept terms of use on the gated HF repos with your HF_TOKEN:
-  #   - https://huggingface.co/google/t5gemma-9b-9b-ul2
-  #   - https://huggingface.co/stabilityai/stable-audio-open-1.0
-  # All four cross-variant shared components (Wan 2.2 VAE, T5-Gemma
-  # encoder + tokenizer, Stable Audio VAE) are lazy-loaded from their
-  # canonical upstream HF repos on first build, so a single ~25 GB
-  # cache is shared across every MagiHuman variant.
-
-The umbrella HF repo `FastVideo/MagiHuman-Diffusers` holds all four
-variants (base / distill / sr_540p / sr_1080p) under sibling subfolders
-and FastVideo will download just the requested subfolder. Local
-conversion via `scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py`
-is also supported.
-"""
-from fastvideo import VideoGenerator
-
-
-PROMPT = (
-    "A warm afternoon scene: a person sits on a park bench reading a book, "
-    "surrounded by softly swaying trees."
-)
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/base",
-        num_gpus=1,
-    )
-    output_path = "outputs_video/magi_human_basic/output_magi_human.mp4"
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path=output_path,
-        save_video=True,
-        # Defaults pulled from the registered preset (magi_human_base):
-        # height=256, width=448, fps=25, num_inference_steps=32, seed=42.
-        # Override here only if you have a specific QA scenario.
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_distill.py

@@ -1,53 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Minimal user-runnable example for the daVinci-MagiHuman DMD-2 distilled
-text-to-AV pipeline.
-
-Same arch as the base model (`basic_magi_human.py`) but with DMD-2 distilled
-weights: 8 denoising steps, no classifier-free guidance. ~4x faster than
-base at the same 256x480 resolution. Mirrors upstream
-`daVinci-MagiHuman/example/distill/run_T2V.sh`.
-
-Prerequisites (one-off):
-
-  # 1) Accept terms on the gated HF repos with your HF_TOKEN:
-  #    - https://huggingface.co/google/t5gemma-9b-9b-ul2
-  #    - https://huggingface.co/stabilityai/stable-audio-open-1.0
-  #    Cross-variant shared components (Wan 2.2 VAE + T5-Gemma + Stable
-  #    Audio VAE) are lazy-loaded from their canonical upstream HF repos
-  #    and shared with the base variant cache.
-  # 2) Convert the distill subfolder of GAIR/daVinci-MagiHuman:
-  python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \\
-      --source GAIR/daVinci-MagiHuman \\
-      --subfolder distill \\
-      --output converted_weights/magi_human_distill \\
-      --cast-bf16
-  # `--cast-bf16` is recommended (61 GB fp32 -> 30 GB bf16); the FV pipeline
-  # loads bf16 anyway, and the conversion keeps norms / RoPE bands fp32.
-"""
-from fastvideo import VideoGenerator
-
-
-PROMPT = (
-    "A warm afternoon scene: a person sits on a park bench reading a book, "
-    "surrounded by softly swaying trees."
-)
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/distill",
-        num_gpus=1,
-    )
-    output_path = "outputs_video/magi_human_basic/output_magi_human_distill.mp4"
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path=output_path,
-        save_video=True,
-        # Defaults pulled from the registered preset (magi_human_distill):
-        # height=256, width=480, fps=25, num_inference_steps=8, cfg=1, seed=42.
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_distill_ti2v.py

@@ -1,34 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Minimal daVinci-MagiHuman DMD-2 distilled text+image-to-AV example."""
-from fastvideo import VideoGenerator
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    MagiHumanDistillI2VConfig,
-)
-
-
-PROMPT = (
-    "A cheerful saxophonist performs a short line with expressive facial "
-    "motion, natural head movement, and synchronized audio in a small jazz club."
-)
-IMAGE_PATH = "assets/images/saxophonist.jpg"
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/distill",
-        num_gpus=1,
-        workload_type="i2v",
-        override_pipeline_cls_name="MagiHumanI2VPipeline",
-        pipeline_config=MagiHumanDistillI2VConfig(),
-    )
-    generator.generate_video(
-        prompt=PROMPT,
-        image_path=IMAGE_PATH,
-        output_path="outputs_video/magi_human_distill_ti2v/output_magi_human_distill_ti2v.mp4",
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_sr1080p.py

@@ -1,45 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run daVinci-MagiHuman SR-1080p text-to-AV in FastVideo.
-
-Build the converted repo on large local storage, then symlink it into the
-workspace:
-
-    python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \
-        --source GAIR/daVinci-MagiHuman \
-        --subfolder base \
-        --sr-source GAIR/daVinci-MagiHuman \
-        --sr-subfolder 1080p_sr \
-        --output /raid/william5lin_converted_weights/magi_human_sr_1080p \
-        --cast-bf16
-    ln -s /raid/william5lin_converted_weights/magi_human_sr_1080p \
-        converted_weights/magi_human_sr_1080p
-"""
-from fastvideo import VideoGenerator
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    MagiHumanSR1080pConfig,
-)
-
-
-PROMPT = (
-    "A warm afternoon scene: a person sits on a park bench reading a book, "
-    "surrounded by softly swaying trees."
-)
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/sr_1080p",
-        num_gpus=1,
-        override_pipeline_cls_name="MagiHumanSR1080pPipeline",
-        pipeline_config=MagiHumanSR1080pConfig(),
-    )
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path="outputs_video/magi_human_sr1080p/output_magi_human_sr1080p.mp4",
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_sr1080p_ti2v.py

@@ -1,34 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run daVinci-MagiHuman SR-1080p text+image-to-AV in FastVideo."""
-from fastvideo import VideoGenerator
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    MagiHumanSR1080pI2VConfig,
-)
-
-
-PROMPT = (
-    "A cheerful saxophonist performs a short line with expressive facial "
-    "motion, natural head movement, and synchronized audio in a small jazz club."
-)
-IMAGE_PATH = "assets/images/saxophonist.jpg"
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/sr_1080p",
-        num_gpus=1,
-        workload_type="i2v",
-        override_pipeline_cls_name="MagiHumanSR1080pI2VPipeline",
-        pipeline_config=MagiHumanSR1080pI2VConfig(),
-    )
-    generator.generate_video(
-        prompt=PROMPT,
-        image_path=IMAGE_PATH,
-        output_path="outputs_video/magi_human_sr1080p_ti2v/output_magi_human_sr1080p_ti2v.mp4",
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_sr540p.py

@@ -1,37 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run daVinci-MagiHuman SR-540p text-to-AV in FastVideo.
-
-The converted repo must contain both ``transformer/`` (base DiT) and
-``sr_transformer/`` (540p SR DiT). Build it with:
-
-    python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \
-        --source GAIR/daVinci-MagiHuman \
-        --subfolder base \
-        --sr-source GAIR/daVinci-MagiHuman \
-        --sr-subfolder 540p_sr \
-        --output converted_weights/magi_human_sr_540p
-"""
-from fastvideo import VideoGenerator
-
-
-PROMPT = (
-    "A warm afternoon scene: a person sits on a park bench reading a book, "
-    "surrounded by softly swaying trees."
-)
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/sr_540p",
-        num_gpus=1,
-    )
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path="outputs_video/magi_human_sr540p/output_magi_human_sr540p.mp4",
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_sr540p_ti2v.py

@@ -1,34 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run daVinci-MagiHuman SR-540p text+image-to-AV in FastVideo."""
-from fastvideo import VideoGenerator
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    MagiHumanSR540pI2VConfig,
-)
-
-
-PROMPT = (
-    "A cheerful saxophonist performs a short line with expressive facial "
-    "motion, natural head movement, and synchronized audio in a small jazz club."
-)
-IMAGE_PATH = "assets/images/saxophonist.jpg"
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/sr_540p",
-        num_gpus=1,
-        workload_type="i2v",
-        override_pipeline_cls_name="MagiHumanSRI2VPipeline",
-        pipeline_config=MagiHumanSR540pI2VConfig(),
-    )
-    generator.generate_video(
-        prompt=PROMPT,
-        image_path=IMAGE_PATH,
-        output_path="outputs_video/magi_human_sr540p_ti2v/output_magi_human_sr540p_ti2v.mp4",
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_magi_human_ti2v.py

@@ -1,34 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Minimal daVinci-MagiHuman base text+image-to-AV example."""
-from fastvideo import VideoGenerator
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    MagiHumanBaseI2VConfig,
-)
-
-
-PROMPT = (
-    "A cheerful saxophonist performs a short line with expressive facial "
-    "motion, natural head movement, and synchronized audio in a small jazz club."
-)
-IMAGE_PATH = "assets/images/saxophonist.jpg"
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/MagiHuman-Diffusers/base",
-        num_gpus=1,
-        workload_type="i2v",
-        override_pipeline_cls_name="MagiHumanI2VPipeline",
-        pipeline_config=MagiHumanBaseI2VConfig(),
-    )
-    generator.generate_video(
-        prompt=PROMPT,
-        image_path=IMAGE_PATH,
-        output_path="outputs_video/magi_human_ti2v/output_magi_human_ti2v.mp4",
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

fastvideo/configs/models/dits/__init__.py

@@ -5,7 +5,6 @@ from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
-from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
@@ -14,5 +13,5 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
     "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
-    "MagiHumanVideoConfig", "StableAudioConfig"
+    "StableAudioConfig"
 ]

fastvideo/configs/models/dits/magi_human.py

@@ -1,110 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Architecture / model config for the daVinci-MagiHuman DiT.
-
-The MagiHuman base DiT is a 15B-parameter single-stream transformer that
-jointly denoises video, audio, and text tokens in one flat sequence. Layout
-details verified against GAIR/daVinci-MagiHuman's base/ shards (2026-04-24).
-
-This file captures only configuration. The module implementation lives in
-fastvideo/models/dits/magi_human.py and the pipeline wiring in
-fastvideo/pipelines/basic/magi_human/.
-"""
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-
-
-def _is_block_layer(n: str, m) -> bool:
-    # Match "block.layers.<idx>" — the FSDP shard boundary for MagiHuman.
-    parts = n.split(".")
-    return (len(parts) >= 3 and parts[0] == "block" and parts[1] == "layers" and str.isdigit(parts[2]))
-
-
-@dataclass
-class MagiHumanArchConfig(DiTArchConfig):
-    """MagiHuman base DiT architecture constants.
-
-    **Scope contract:** fields here must match the `transformer/config.json`
-    emitted by `scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py`
-    1:1, and both are sourced from the upstream Python reference
-    `inference/common/config.py::ModelConfig` (the HF root `config.json`
-    is empty so the Python source is canonical). Pipeline-level knobs
-    (VAE stride, fps, num_inference_steps, CFG scales, flow_shift,
-    t5_gemma_target_length) and data-proxy knobs (coords_style,
-    frame_receptive_field, ref_audio_offset, text_offset) live on
-    `MagiHumanBaseConfig`, NOT here.
-
-    `param_names_mapping` is intentionally empty: the FastVideo implementation
-    keeps the same module tree as the reference (`adapter.*`,
-    `block.layers.<i>.*`, `final_linear_{video,audio}.*`,
-    `final_norm_{video,audio}.*`), so converted weights load directly.
-    """
-
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_block_layer])
-
-    # No renames needed — the FastVideo module mirrors the reference names.
-    param_names_mapping: dict = field(default_factory=dict)
-    reverse_param_names_mapping: dict = field(default_factory=dict)
-    lora_param_names_mapping: dict = field(default_factory=dict)
-
-    # --- transformer shape ---
-    num_layers: int = 40
-    hidden_size: int = 5120
-    head_dim: int = 128
-    num_query_groups: int = 8  # num_heads_kv (GQA)
-
-    # --- modality channels ---
-    # video_in_channels = z_dim (48) * patch_size product (1*2*2=4), so the
-    # embedder receives 192 per token. text_in_channels is T5Gemma-9B's
-    # encoder hidden size.
-    video_in_channels: int = 192
-    audio_in_channels: int = 64
-    text_in_channels: int = 3584
-
-    # --- block-level architecture switches ---
-    # Sandwich MoE: first and last 4 layers have per-modality experts
-    # (video/audio/text), middle layers share a single set of weights.
-    mm_layers: tuple[int, ...] = (0, 1, 2, 3, 36, 37, 38, 39)
-    local_attn_layers: tuple[int, ...] = ()
-    gelu7_layers: tuple[int, ...] = (0, 1, 2, 3)
-    post_norm_layers: tuple[int, ...] = ()
-    enable_attn_gating: bool = True
-    activation_type: str = "swiglu7"
-
-    # --- DiT patching (upstream `ModelConfig`-equivalent; NOT the VAE
-    #     stride, which is pipeline-level). ---
-    patch_size: tuple[int, int, int] = (1, 2, 2)
-    spatial_rope_interpolation: str = "extra"
-
-    # --- TReAD (token routing + early drop). Flattened from the upstream
-    #     nested `tread_config` dict so it round-trips through
-    #     `update_model_arch` cleanly. ---
-    tread_selection_rate: float = 0.5
-    tread_start_layer_idx: int = 2
-    tread_end_layer_idx: int = 25
-
-    # --- derived fields (populated in __post_init__) ---
-    num_attention_heads: int = 0  # hidden_size / head_dim
-    num_heads_kv: int = 0  # == num_query_groups
-    in_channels: int = 0  # mirror of video_in_channels (FastVideo contract)
-    out_channels: int = 0  # mirror of video_in_channels
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        self.num_attention_heads = self.hidden_size // self.head_dim
-        self.num_heads_kv = self.num_query_groups
-        self.in_channels = self.video_in_channels
-        self.out_channels = self.video_in_channels
-        # num_channels_latents is the VAE latent z_dim (48 for Wan 2.2 TI2V-5B).
-        # We don't declare z_dim on the arch config (it's a VAE property),
-        # but we still set num_channels_latents for the BaseDiT contract.
-        self.num_channels_latents = 48
-
-
-@dataclass
-class MagiHumanVideoConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=MagiHumanArchConfig)
-
-    prefix: str = "magi_human"

fastvideo/configs/models/encoders/__init__.py

@@ -9,11 +9,10 @@ from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1
 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig,
                                                                         StableAudioConditionerConfig)
-from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig"
+    "StableAudioConditionerConfig"
 ]

fastvideo/configs/models/encoders/t5gemma.py

@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Config for the T5-Gemma encoder used by daVinci-MagiHuman.
-
-The reference pipeline uses `transformers.models.t5gemma.T5GemmaEncoderModel`
-on `google/t5gemma-9b-9b-ul2`. That is a gated Google repository, so the
-encoder weights are not bundled inside GAIR/daVinci-MagiHuman; they are
-loaded from the T5-Gemma HF repo directly.
-
-Encoder shape (verified from google/t5gemma-9b-9b-ul2/config.json):
-    layers=42, hidden=3584, heads=16, kv_heads=8, head_dim=256,
-    intermediate=14336, rope_theta=10000.0, max_pos=8192,
-    layer_types alternate sliding_attention / full_attention.
-"""
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.encoders.base import (
-    TextEncoderArchConfig,
-    TextEncoderConfig,
-)
-
-
-def _is_t5gemma_model(n: str, m) -> bool:
-    return n.endswith("t5gemma_model") or n.endswith("_t5gemma_model")
-
-
-@dataclass
-class T5GemmaEncoderArchConfig(TextEncoderArchConfig):
-    architectures: list[str] = field(default_factory=lambda: ["T5GemmaEncoderModel"])
-
-    hidden_size: int = 3584
-    num_hidden_layers: int = 42
-    num_attention_heads: int = 16
-    num_key_value_heads: int = 8
-    head_dim: int = 256
-    intermediate_size: int = 14336
-    max_position_embeddings: int = 8192
-    rope_theta: float = 10000.0
-    vocab_size: int = 256000
-
-    # MagiHuman fixes prompt embed length at 640 via pad_or_trim.
-    text_len: int = 640
-
-    pad_token_id: int = 0
-    eos_token_id: int = 1
-
-    # Path to the upstream gated repo. When set, the FastVideo loader will
-    # pull the encoder directly via `T5GemmaEncoderModel.from_pretrained`.
-    t5gemma_model_path: str = "google/t5gemma-9b-9b-ul2"
-    t5gemma_dtype: str = "bfloat16"
-
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_t5gemma_model])
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        # WHY: upstream `t5_gemma_model.py:25` tokenizes without
-        # padding/max_length, then `prompt_process.py` pad_or_trim-s the
-        # encoded states. Keep only tensor return here so
-        # MagiHumanLatentPreparationStage can pad/trim post-encode while
-        # preserving the real original prompt length.
-        self.tokenizer_kwargs.pop("truncation", None)
-        self.tokenizer_kwargs.pop("max_length", None)
-        self.tokenizer_kwargs.pop("padding", None)
-
-
-@dataclass
-class T5GemmaEncoderConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(default_factory=T5GemmaEncoderArchConfig)
-
-    prefix: str = "t5gemma"

fastvideo/envs.py

@@ -35,11 +35,6 @@ if TYPE_CHECKING:
     FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS: int = 1
     FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS: int = 2
     FASTVIDEO_TORCH_PROFILE_REGIONS: str = ""
-    FASTVIDEO_TRACE_ACTIVATIONS: bool = False
-    FASTVIDEO_TRACE_LAYERS: str = ""
-    FASTVIDEO_TRACE_STATS: str = "abs_mean,sum"
-    FASTVIDEO_TRACE_OUTPUT: str = "/tmp/fv_trace_<pid>.jsonl"
-    FASTVIDEO_TRACE_STEPS: str = ""
     FASTVIDEO_SERVER_DEV_MODE: bool = False
     FASTVIDEO_STAGE_LOGGING: bool = False
     FASTVIDEO_HOST_IP: str = ""
@@ -257,22 +252,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_TORCH_PROFILE_REGIONS":
     lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""),
 
-    # Enable activation trace hooks if set.
-    "FASTVIDEO_TRACE_ACTIVATIONS":
-    lambda: bool(os.getenv("FASTVIDEO_TRACE_ACTIVATIONS", "0") != "0"),
-    # Regex filter for traced module names. Empty means all modules.
-    "FASTVIDEO_TRACE_LAYERS":
-    lambda: os.getenv("FASTVIDEO_TRACE_LAYERS", ""),
-    # Comma-separated activation stats to dump for each output tensor.
-    "FASTVIDEO_TRACE_STATS":
-    lambda: os.getenv("FASTVIDEO_TRACE_STATS", "abs_mean,sum"),
-    # JSONL sink path. The literal <pid> is replaced at runtime.
-    "FASTVIDEO_TRACE_OUTPUT":
-    lambda: os.getenv("FASTVIDEO_TRACE_OUTPUT", "/tmp/fv_trace_<pid>.jsonl"),
-    # Comma-separated denoise step indices. Empty means all steps.
-    "FASTVIDEO_TRACE_STEPS":
-    lambda: os.getenv("FASTVIDEO_TRACE_STEPS", ""),
-
     # If set, fastvideo will run in development mode, which will enable
     # some additional endpoints for developing and debugging,
     # e.g. `/reset_prefix_cache`

fastvideo/hooks/activation_trace.py

@@ -1,221 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Zero-overhead-when-off activation trace mode for FastVideo pipelines.
-
-Enable by setting FASTVIDEO_TRACE_ACTIVATIONS=1. When off, this module
-adds zero overhead — no hooks are registered, no branches exist in the
-production forward path. When on, registers forward hooks on modules
-whose name matches FASTVIDEO_TRACE_LAYERS, computes the requested stats
-(FASTVIDEO_TRACE_STATS) on each output tensor, and writes JSONL records
-to FASTVIDEO_TRACE_OUTPUT.
-
-Useful for parity debugging across model ports — log on both the
-FastVideo path and the upstream reference, diff the two JSONL files
-to find the first divergent layer.
-
-Example:
-
-    FASTVIDEO_TRACE_ACTIVATIONS=1 \
-    FASTVIDEO_TRACE_LAYERS="^block\\.layers\\.[0-9]+$" \
-    FASTVIDEO_TRACE_STATS="abs_mean,max,shape" \
-    FASTVIDEO_TRACE_OUTPUT="/tmp/fv_trace.jsonl" \
-    python examples/inference/basic/basic_magi_human.py
-"""
-from __future__ import annotations
-
-import json
-import os
-import re
-import threading
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Any
-from collections.abc import Callable, Iterator
-
-import torch
-from torch import nn
-
-from fastvideo import envs
-from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-_TRACE_STATE = threading.local()
-
-
-def current_step_idx() -> int | None:
-    return getattr(_TRACE_STATE, "step_idx", None)
-
-
-@contextmanager
-def trace_step(step_idx: int) -> Iterator[None]:
-    """Context manager that sets the current denoise step for trace records."""
-    prev = getattr(_TRACE_STATE, "step_idx", None)
-    _TRACE_STATE.step_idx = step_idx
-    try:
-        yield
-    finally:
-        _TRACE_STATE.step_idx = prev
-
-
-_STAT_FNS: dict[str, Callable[[torch.Tensor], Any]] = {
-    "abs_mean": lambda t: float(t.detach().float().abs().mean().item()),
-    "sum": lambda t: float(t.detach().float().sum().item()),
-    "min": lambda t: float(t.detach().float().min().item()),
-    "max": lambda t: float(t.detach().float().max().item()),
-    "mean": lambda t: float(t.detach().float().mean().item()),
-    "std": lambda t: float(t.detach().float().std().item()),
-    "shape": lambda t: list(t.shape),
-    "dtype": lambda t: str(t.dtype),
-}
-
-
-def _resolve_stats(spec: str) -> list[tuple[str, Callable[[torch.Tensor], Any]]]:
-    stats = []
-    for name in [s.strip() for s in spec.split(",") if s.strip()]:
-        stat_fn = _STAT_FNS.get(name)
-        if stat_fn is None:
-            logger.warning(
-                "FASTVIDEO_TRACE_STATS contains unknown stat %r; valid: %s",
-                name,
-                sorted(_STAT_FNS),
-            )
-            continue
-        stats.append((name, stat_fn))
-    return stats
-
-
-def _resolve_output_path(template: str) -> Path:
-    return Path(template.replace("<pid>", str(os.getpid())))
-
-
-def _parse_step_filter(spec: str) -> set[int] | None:
-    if not spec.strip():
-        return None
-    return {int(s.strip()) for s in spec.split(",") if s.strip()}
-
-
-class JsonlSink:
-    """Buffered JSONL writer with thread-safe append."""
-
-    def __init__(self, path: Path) -> None:
-        self.path = path
-        self.path.parent.mkdir(parents=True, exist_ok=True)
-        self._fh = open(self.path, "a", buffering=1)  # noqa: SIM115
-        self._lock = threading.Lock()
-        logger.info("Activation trace JSONL sink: %s", self.path)
-
-    def write(self, record: dict[str, Any]) -> None:
-        line = json.dumps(record, default=str) + "\n"
-        with self._lock:
-            self._fh.write(line)
-
-    def close(self) -> None:
-        with self._lock:
-            if not self._fh.closed:
-                self._fh.close()
-
-
-class ActivationStatHook(ForwardHook):
-    """Forward hook that emits per-tensor stats to a JSONL sink."""
-
-    def __init__(
-        self,
-        module_name: str,
-        stats: list[tuple[str, Callable[[torch.Tensor], Any]]],
-        sink: JsonlSink,
-        step_filter: set[int] | None,
-    ) -> None:
-        self.module_name = module_name
-        self.stats = stats
-        self.sink = sink
-        self.step_filter = step_filter
-
-    def name(self) -> str:
-        return "ActivationStatHook"
-
-    def post_forward(self, module: nn.Module, output: Any) -> Any:
-        step_idx = current_step_idx()
-        if self.step_filter is not None and step_idx not in self.step_filter:
-            return output
-        for tensor_label, tensor in _flatten_tensors(output):
-            record: dict[str, Any] = {
-                "module": self.module_name,
-                "tensor": tensor_label,
-                "step": step_idx,
-            }
-            for stat_name, stat_fn in self.stats:
-                try:
-                    record[stat_name] = stat_fn(tensor)
-                except Exception as exc:  # pragma: no cover - defensive logging
-                    record[stat_name] = f"<error: {exc!r}>"
-            self.sink.write(record)
-        return output
-
-
-def _flatten_tensors(obj: Any, prefix: str = "out") -> list[tuple[str, torch.Tensor]]:
-    """Yield (label, tensor) pairs from arbitrarily-nested forward outputs."""
-    if isinstance(obj, torch.Tensor):
-        return [(prefix, obj)]
-    if isinstance(obj, tuple | list):
-        out = []
-        for idx, item in enumerate(obj):
-            out.extend(_flatten_tensors(item, f"{prefix}[{idx}]"))
-        return out
-    if isinstance(obj, dict):
-        out = []
-        for key, value in obj.items():
-            out.extend(_flatten_tensors(value, f"{prefix}.{key}"))
-        return out
-    return []
-
-
-class ActivationTraceManager:
-
-    def __init__(self, managers: list[ModuleHookManager], sink: JsonlSink) -> None:
-        self.managers = managers
-        self.sink = sink
-
-    def remove_from_manager(self) -> None:
-        for manager in self.managers:
-            if manager.get_forward_hook("ActivationStatHook") is not None:
-                manager.remove_forward_hook("ActivationStatHook")
-            if not manager.forward_hooks:
-                ModuleHookManager.remove_from_manager(manager.module)
-        self.sink.close()
-
-
-def attach_activation_trace(model: nn.Module | None) -> ActivationTraceManager | None:
-    """Attach activation-stat hooks to model. Returns None if trace is off."""
-    if not envs.FASTVIDEO_TRACE_ACTIVATIONS or model is None:
-        return None
-
-    pattern_spec = envs.FASTVIDEO_TRACE_LAYERS
-    pattern = re.compile(pattern_spec) if pattern_spec else re.compile(".*")
-    stats = _resolve_stats(envs.FASTVIDEO_TRACE_STATS)
-    if not stats:
-        logger.warning("FASTVIDEO_TRACE_STATS yielded no valid stats; trace disabled.")
-        return None
-
-    sink = JsonlSink(_resolve_output_path(envs.FASTVIDEO_TRACE_OUTPUT))
-    step_filter = _parse_step_filter(envs.FASTVIDEO_TRACE_STEPS)
-    managers = []
-    for name, module in model.named_modules():
-        if not name or not pattern.search(name):
-            continue
-        manager = ModuleHookManager.get_from_or_default(module)
-        manager.append_forward_hook(ActivationStatHook(name, stats, sink, step_filter))
-        managers.append(manager)
-
-    logger.info(
-        "Activation trace attached to %d modules (pattern=%r, stats=%s)",
-        len(managers),
-        pattern_spec,
-        [stat_name for stat_name, _ in stats],
-    )
-    return ActivationTraceManager(managers, sink)
-
-
-def detach_activation_trace(mgr: ActivationTraceManager | None) -> None:
-    if mgr is not None:
-        mgr.remove_from_manager()

fastvideo/models/dits/magi_human.py

@@ -1,867 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""daVinci-MagiHuman DiT (base variant).
-
-Ported from https://github.com/GAIR-NLP/daVinci-MagiHuman
-(inference/model/dit/dit_module.py, ~950 lines in the reference).
-
-Architecture summary (verified against GAIR/daVinci-MagiHuman/base/ weights):
-
-  - 40 transformer layers, hidden 5120, head_dim 128.
-  - GQA with 40 query heads and 8 KV heads.
-  - Multi-modality "sandwich": layers 0..3 and 36..39 use 3-way modality
-    experts (video/audio/text) packed inside each linear as
-    weight[..., out * 3, in]. Middle layers share a single expert.
-  - Per-head attention gating: the QKV projection emits an extra
-    num_heads_q channels that are sigmoid-gated onto the attention output.
-  - Activation is GELU7 on layers 0..3 (non-gated, intermediate=4*hidden)
-    and SwiGLU7 elsewhere (gated, intermediate=int(hidden*4*2/3)//4*4).
-  - Position encoding is an element-wise Fourier embedding over 9-column
-    coords (t,h,w + original TxHxW + reference TxHxW), not a standard
-    1D/3D RoPE.
-  - Forward takes a flat concatenated token stream (video first, then
-    audio, then text) plus a modality map; the internal ModalityDispatcher
-    permutes by modality before each linear so per-expert chunks line up.
-
-Deviations from the "use fastvideo.layers primitives everywhere" guideline
-in the add-model skill:
-
-  - The packed-expert linears store weight as [out * num_experts, in].
-    FastVideo's ReplicatedLinear does not model this layout; we use raw
-    nn.Parameter with a small wrapper below. This is deliberate and scoped
-    to this DiT: ReplicatedLinear still handles the adapter.* embedders
-    and final_linear_{video,audio} (single-expert) here.
-  - Self-attention is full-sequence and crosses modalities inside the flat
-    concat stream; DistributedAttention assumes a clean spatial-sequence
-    layout, so for the first port we use torch SDPA. Multi-GPU sequence
-    parallelism is a follow-up.
-  - torch.compile via magi_compiler is replaced with a plain nn.Module.
-
-For the full history and shape-by-shape verification notes, see
-.claude/skills/add-model/SKILL.md and the scaffold PR description.
-"""
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass
-from enum import IntEnum
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from fastvideo.attention import LocalAttention
-from fastvideo.configs.models.dits.magi_human import (
-    MagiHumanArchConfig,
-    MagiHumanVideoConfig,
-)
-from fastvideo.layers.rotary_embedding import _apply_rotary_emb
-from fastvideo.models.dits.base import BaseDiT
-from fastvideo.platforms import AttentionBackendEnum
-
-
-# ---------------------------------------------------------------------------
-# Enums
-# ---------------------------------------------------------------------------
-
-
-class Modality(IntEnum):
-    VIDEO = 0
-    AUDIO = 1
-    TEXT = 2
-
-
-# ---------------------------------------------------------------------------
-# Activations
-# ---------------------------------------------------------------------------
-
-
-def swiglu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
-    """Gated swish-GLU with OpenAI-OSS-style limits and +1 linear bias."""
-    in_dtype = x.dtype
-    x = x.to(torch.float32)
-    x_glu, x_linear = x[..., ::2], x[..., 1::2]
-    x_glu = x_glu.clamp(max=limit)
-    x_linear = x_linear.clamp(min=-limit, max=limit)
-    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
-    return (out_glu * (x_linear + 1)).to(in_dtype)
-
-
-def gelu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
-    in_dtype = x.dtype
-    x = x.to(torch.float32).clamp(max=limit)
-    return (x * torch.sigmoid(alpha * x)).to(in_dtype)
-
-
-# ---------------------------------------------------------------------------
-# Modality dispatcher
-# ---------------------------------------------------------------------------
-
-
-class ModalityDispatcher:
-    """Permute a flat token stream so same-modality tokens are contiguous.
-
-    The DiT's multi-expert linears apply a different weight chunk per modality.
-    Instead of carrying a branch inside each Linear, we pre-permute tokens so
-    each chunk sees a contiguous slice, then un-permute before computing
-    RoPE/attention across the full sequence.
-    """
-
-    def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
-        self.modality_mapping = modality_mapping
-        self.num_modalities = num_modalities
-        self.permute_mapping = torch.argsort(modality_mapping)
-        self.inv_permute_mapping = torch.argsort(self.permute_mapping)
-        permuted = modality_mapping[self.permute_mapping]
-        self.group_size = torch.bincount(permuted, minlength=num_modalities).to(torch.int32)
-        self.group_size_cpu: list[int] = [int(x) for x in self.group_size.cpu().tolist()]
-
-    def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
-        return list(torch.split(x, self.group_size_cpu, dim=0))
-
-    def undispatch(self, *chunks: torch.Tensor) -> torch.Tensor:
-        return torch.cat(chunks, dim=0)
-
-    @staticmethod
-    def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
-        return x[permute_mapping]
-
-    @staticmethod
-    def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
-        return x[inv_permute_mapping]
-
-
-# ---------------------------------------------------------------------------
-# Norms, rotary embed
-# ---------------------------------------------------------------------------
-
-
-class MultiModalityRMSNorm(nn.Module):
-    """RMSNorm with optional per-modality scale.
-
-    When num_modality == 1, behaves identically to a standard RMSNorm with
-    weight initialized to zero (effective weight is 1 + weight, hence the
-    learnable +1 offset baked into the forward path). When num_modality > 1,
-    the weight tensor packs per-modality scales along its flat axis and the
-    dispatcher selects the right chunk per modality.
-    """
-
-    def __init__(self, dim: int, eps: float = 1e-6, num_modality: int = 1):
-        super().__init__()
-        self.dim = dim
-        self.eps = eps
-        self.num_modality = num_modality
-        # Always stored in fp32; matches the reference initialization.
-        self.weight = nn.Parameter(torch.zeros(dim * num_modality, dtype=torch.float32))
-
-    def _rms(self, x: torch.Tensor) -> torch.Tensor:
-        t = x.float()
-        return t * torch.rsqrt(t.pow(2).mean(dim=-1, keepdim=True) + self.eps)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        modality_dispatcher: Optional[ModalityDispatcher] = None,
-    ) -> torch.Tensor:
-        original_dtype = x.dtype
-        t = self._rms(x)
-        if self.num_modality == 1:
-            return (t * (self.weight + 1)).to(original_dtype)
-        assert modality_dispatcher is not None, (
-            "MultiModalityRMSNorm with num_modality>1 requires a dispatcher"
-        )
-        weight_chunks = self.weight.chunk(self.num_modality, dim=0)
-        parts = modality_dispatcher.dispatch(t)
-        for i in range(self.num_modality):
-            parts[i] = parts[i] * (weight_chunks[i] + 1)
-        return modality_dispatcher.undispatch(*parts).to(original_dtype)
-
-
-def _freq_bands(num_bands: int, temperature: float = 10000.0) -> torch.Tensor:
-    exp = torch.arange(0, num_bands, 1, dtype=torch.int64).float() / num_bands
-    return 1.0 / (temperature ** exp)
-
-
-class ElementWiseFourierEmbed(nn.Module):
-    """Element-wise Fourier embedding over 9-column coords (t, h, w, T, H, W,
-    ref_T, ref_H, ref_W). Produces a per-token positional embedding that
-    acts as the RoPE angle input for attention.
-
-    Weight: `bands` of shape `[dim // 8]` (fixed at init via freq_bands).
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        temperature: float = 10000.0,
-        dtype: torch.dtype = torch.float32,
-    ):
-        super().__init__()
-        self.dim = dim
-        bands = _freq_bands(dim // 8, temperature=temperature).to(dtype)
-        # `register_buffer` so state_dict keeps it, matching upstream naming.
-        self.register_buffer("bands", bands)
-
-    def forward(self, coords: torch.Tensor) -> torch.Tensor:
-        # coords: [L, 9] = (t, h, w, T, H, W, ref_T, ref_H, ref_W)
-        coords_xyz = coords[:, :3]
-        sizes = coords[:, 3:6]
-        refs = coords[:, 6:9]
-
-        scales = (refs - 1) / (sizes - 1)
-        scales[(refs == 1) & (sizes == 1)] = 1
-        # Center H and W (leave time uncentered).
-        centers = (sizes - 1) / 2
-        centers[:, 0] = 0
-        coords_xyz = coords_xyz - centers
-
-        proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands  # [L, 3, B]
-        sin_proj = proj.sin()
-        cos_proj = proj.cos()
-        return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)
-
-
-# ---------------------------------------------------------------------------
-# Packed-expert linear
-# ---------------------------------------------------------------------------
-
-
-class PackedExpertLinear(nn.Module):
-    """Linear where the weight is packed per-modality along the output axis.
-
-    Shapes:
-        weight: [out_features * num_experts, in_features]
-        bias:   [out_features * num_experts]  (optional)
-
-    When `num_experts == 1`, behaves exactly like `nn.Linear`. When
-    `num_experts > 1`, `forward` dispatches the input via the supplied
-    `ModalityDispatcher`, applies the per-modality weight/bias chunk, and
-    gathers the outputs in original order.
-
-    Why not use `ReplicatedLinear`? Because the packed-expert layout is not
-    what ReplicatedLinear (or any other fastvideo.layers.linear) is wired
-    for. Using raw `nn.Parameter` keeps weight loading trivial (names map
-    1:1 to the upstream checkpoint) and avoids quantization-path assumptions
-    that don't match this layout.
-    """
-
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        num_experts: int = 1,
-        bias: bool = False,
-        dtype: torch.dtype = torch.bfloat16,
-    ):
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.num_experts = num_experts
-        self.use_bias = bias
-        self.weight = nn.Parameter(
-            torch.empty(out_features * num_experts, in_features, dtype=dtype)
-        )
-        if bias:
-            self.bias = nn.Parameter(
-                torch.empty(out_features * num_experts, dtype=dtype)
-            )
-        else:
-            self.register_parameter("bias", None)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        modality_dispatcher: Optional[ModalityDispatcher] = None,
-    ) -> torch.Tensor:
-        if self.num_experts == 1:
-            return F.linear(x, self.weight, self.bias)
-        assert modality_dispatcher is not None, (
-            "PackedExpertLinear with num_experts>1 requires a dispatcher"
-        )
-        parts = modality_dispatcher.dispatch(x)
-        w_chunks = self.weight.chunk(self.num_experts, dim=0)
-        b_chunks = (
-            self.bias.chunk(self.num_experts, dim=0)
-            if self.bias is not None else [None] * self.num_experts
-        )
-        for i in range(self.num_experts):
-            parts[i] = F.linear(parts[i], w_chunks[i], b_chunks[i])
-        return modality_dispatcher.undispatch(*parts)
-
-
-# ---------------------------------------------------------------------------
-# Attention & MLP
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class AttentionSubConfig:
-    hidden_size: int
-    num_heads_q: int
-    num_heads_kv: int
-    head_dim: int
-    num_modality: int
-    enable_attn_gating: bool
-    use_local_attn: bool = False
-    frame_receptive_field: int = 11
-
-
-class MagiAttention(nn.Module):
-    """Self-attention with GQA + optional per-head sigmoid gating."""
-
-    def __init__(self, cfg: AttentionSubConfig):
-        super().__init__()
-        self.cfg = cfg
-        self.gating_size = cfg.num_heads_q if cfg.enable_attn_gating else 0
-        qkv_out = (
-            cfg.num_heads_q * cfg.head_dim
-            + 2 * cfg.num_heads_kv * cfg.head_dim
-            + self.gating_size
-        )
-        self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
-        self.linear_qkv = PackedExpertLinear(
-            cfg.hidden_size, qkv_out, num_experts=cfg.num_modality, bias=False,
-        )
-        self.linear_proj = PackedExpertLinear(
-            cfg.num_heads_q * cfg.head_dim, cfg.hidden_size,
-            num_experts=cfg.num_modality, bias=False,
-        )
-        self.q_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)
-        self.k_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)
-
-        self.q_size = cfg.num_heads_q * cfg.head_dim
-        self.kv_size = cfg.num_heads_kv * cfg.head_dim
-
-        self.attn = LocalAttention(
-            num_heads=cfg.num_heads_q,
-            head_size=cfg.head_dim,
-            num_kv_heads=cfg.num_heads_kv,
-            causal=False,
-            supported_attention_backends=(
-                AttentionBackendEnum.FLASH_ATTN,
-                AttentionBackendEnum.TORCH_SDPA,
-            ),
-        )
-
-    def configure_local_attention(
-        self,
-        *,
-        enabled: bool,
-        frame_receptive_field: int = 11,
-    ) -> None:
-        self.cfg.use_local_attn = enabled
-        self.cfg.frame_receptive_field = frame_receptive_field
-
-    def _sdpa(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-        """Run SDPA on [L, H, D] tensors and return [L, Hq, D]."""
-        if q.numel() == 0:
-            return q.new_empty(q.shape)
-        out = F.scaled_dot_product_attention(
-            q.transpose(0, 1).unsqueeze(0).contiguous(),
-            k.transpose(0, 1).unsqueeze(0).contiguous(),
-            v.transpose(0, 1).unsqueeze(0).contiguous(),
-            enable_gqa=self.cfg.num_heads_q != self.cfg.num_heads_kv,
-        )
-        return out.squeeze(0).transpose(0, 1).contiguous()
-
-    def _local_window_attention(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        *,
-        num_video_tokens: int,
-        num_frames: int,
-    ) -> torch.Tensor:
-        """Approximate upstream FFAHandler block accumulation with SDPA.
-
-        SR-1080p's reference kernel sums three independently-normalized
-        attention contributions:
-
-        * video frame queries -> local-window video keys;
-        * all video queries -> all audio+text keys;
-        * all audio+text queries -> full sequence keys.
-
-        This method mirrors that accumulator semantics with ordinary SDPA
-        slices. It is intentionally scoped to single-process inference; layers
-        without ``use_local_attn`` keep the existing full LocalAttention path.
-        """
-        if num_frames <= 0 or num_video_tokens <= 0:
-            return self._sdpa(q, k, v)
-        if num_video_tokens % num_frames != 0:
-            raise ValueError(
-                f"MagiHuman local attention expects video tokens divisible by "
-                f"frames, got {num_video_tokens=} and {num_frames=}."
-            )
-
-        token_per_frame = num_video_tokens // num_frames
-        out = torch.zeros(
-            q.shape[0],
-            self.cfg.num_heads_q,
-            self.cfg.head_dim,
-            device=q.device,
-            dtype=q.dtype,
-        )
-        rf = int(self.cfg.frame_receptive_field)
-
-        q_video = q[:num_video_tokens]
-        k_video = k[:num_video_tokens]
-        v_video = v[:num_video_tokens]
-        for frame_idx in range(num_frames):
-            q_start = frame_idx * token_per_frame
-            q_end = q_start + token_per_frame
-            k_start = max(0, (frame_idx - rf) * token_per_frame)
-            k_end = min(num_video_tokens, (frame_idx + rf + 1) * token_per_frame)
-            out[q_start:q_end] = self._sdpa(
-                q_video[q_start:q_end],
-                k_video[k_start:k_end],
-                v_video[k_start:k_end],
-            )
-
-        if num_video_tokens < q.shape[0]:
-            k_at = k[num_video_tokens:]
-            v_at = v[num_video_tokens:]
-            out[:num_video_tokens] = out[:num_video_tokens] + self._sdpa(
-                q[:num_video_tokens],
-                k_at,
-                v_at,
-            )
-            out[num_video_tokens:] = self._sdpa(
-                q[num_video_tokens:],
-                k,
-                v,
-            )
-        return out
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        rope: torch.Tensor,
-        permute_mapping: torch.Tensor,
-        inv_permute_mapping: torch.Tensor,
-        modality_dispatcher: ModalityDispatcher,
-        num_video_tokens: int | None = None,
-        num_frames: int | None = None,
-    ) -> torch.Tensor:
-        orig_dtype = self.linear_qkv.weight.dtype
-        h = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(orig_dtype)
-        qkv = self.linear_qkv(h, modality_dispatcher=modality_dispatcher).float()
-        q, k, v, g = torch.split(
-            qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=-1,
-        )
-        q = q.view(-1, self.cfg.num_heads_q, self.cfg.head_dim)
-        k = k.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
-        v = v.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
-        g = g.view(-1, self.cfg.num_heads_q, 1) if self.gating_size else None
-
-        q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
-        k = self.k_norm(k, modality_dispatcher=modality_dispatcher)
-
-        # Un-permute before RoPE + attention so positional order reflects
-        # the original (video, audio, text) concat — matches reference.
-        q = ModalityDispatcher.inv_permute(q, inv_permute_mapping)
-        k = ModalityDispatcher.inv_permute(k, inv_permute_mapping)
-        v = ModalityDispatcher.inv_permute(v, inv_permute_mapping)
-        if g is not None:
-            g = ModalityDispatcher.inv_permute(g, inv_permute_mapping)
-
-        # Element-wise Fourier embed packs sin/cos of 3 axes into a single
-        # `rope` tensor. Match reference's split:
-        #   sin_emb, cos_emb = rope.tensor_split(2, -1)
-        # Reference passes (cos_emb, sin_emb) but splits sin first — replicated
-        # exactly so weight parity holds. Partial RoPE: rope dim is
-        # 6 * (head_dim // 8) = 96 < head_dim (128), so the trailing 32
-        # head_dim positions stay unrotated, matching the reference.
-        sin_emb, cos_emb = rope.tensor_split(2, -1)
-        rot_dim = cos_emb.shape[-1] * 2
-        q_rot = _apply_rotary_emb(q[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
-        k_rot = _apply_rotary_emb(k[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
-        if rot_dim < q.shape[-1]:
-            q = torch.cat([q_rot, q[..., rot_dim:]], dim=-1)
-            k = torch.cat([k_rot, k[..., rot_dim:]], dim=-1)
-        else:
-            q, k = q_rot, k_rot
-
-        # Run SDPA via FastVideo's LocalAttention so the backend selection
-        # (SDPA / FlashAttn / SLA / SageAttn) flows through the standard
-        # configurable path. GQA is handled inside the SDPA backend via
-        # `enable_gqa=True` when num_heads_q != num_heads_kv, so we no
-        # longer need the manual `repeat_interleave` here.
-        # Attention math runs at orig_dtype (bf16 in production and in the
-        # parity test, since PackedExpertLinear's default is bf16, matching
-        # upstream BaseLinear at dit_module.py:330). The gating multiply
-        # promotes back to fp32 implicitly via PyTorch's type-promotion
-        # rules: bf16_attn_out * sigmoid(fp32_g) -> fp32, mirroring upstream
-        # dit_module.py:649.
-        q = q.to(orig_dtype)
-        k = k.to(orig_dtype)
-        v = v.to(orig_dtype)
-        if self.cfg.use_local_attn:
-            if num_video_tokens is None or num_frames is None:
-                raise ValueError("MagiHuman local attention requires video token/frame metadata.")
-            out = self._local_window_attention(
-                q,
-                k,
-                v,
-                num_video_tokens=num_video_tokens,
-                num_frames=num_frames,
-            )
-        else:
-            out = self.attn(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)).squeeze(0)
-
-        out = ModalityDispatcher.permute(out, permute_mapping)
-        if g is not None:
-            g = ModalityDispatcher.permute(g, permute_mapping)
-            out = out * torch.sigmoid(g)
-        out = out.reshape(-1, self.cfg.num_heads_q * self.cfg.head_dim).to(orig_dtype)
-        return self.linear_proj(out, modality_dispatcher=modality_dispatcher)
-
-
-@dataclass
-class MLPSubConfig:
-    hidden_size: int
-    intermediate_size: int
-    activation: str              # "swiglu7" or "gelu7"
-    num_modality: int
-    gated: bool
-
-
-class MagiMLP(nn.Module):
-    def __init__(self, cfg: MLPSubConfig):
-        super().__init__()
-        self.cfg = cfg
-        self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
-        up_out = cfg.intermediate_size * 2 if cfg.gated else cfg.intermediate_size
-        self.up_gate_proj = PackedExpertLinear(
-            cfg.hidden_size, up_out, num_experts=cfg.num_modality, bias=False,
-        )
-        self.down_proj = PackedExpertLinear(
-            cfg.intermediate_size, cfg.hidden_size,
-            num_experts=cfg.num_modality, bias=False,
-        )
-        self._act = swiglu7 if cfg.activation == "swiglu7" else gelu7
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        modality_dispatcher: ModalityDispatcher,
-    ) -> torch.Tensor:
-        orig_dtype = self.up_gate_proj.weight.dtype
-        x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(orig_dtype)
-        x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).float()
-        x = self._act(x).to(orig_dtype)
-        x = self.down_proj(x, modality_dispatcher=modality_dispatcher).float()
-        return x
-
-
-class MagiTransformerLayer(nn.Module):
-    def __init__(self, arch: MagiHumanArchConfig, layer_idx: int):
-        super().__init__()
-        num_modality = 3 if layer_idx in arch.mm_layers else 1
-        self.post_norm = layer_idx in arch.post_norm_layers
-        self.layer_idx = layer_idx
-
-        self.attention = MagiAttention(AttentionSubConfig(
-            hidden_size=arch.hidden_size,
-            num_heads_q=arch.num_attention_heads,
-            num_heads_kv=arch.num_heads_kv,
-            head_dim=arch.head_dim,
-            num_modality=num_modality,
-            enable_attn_gating=arch.enable_attn_gating,
-            use_local_attn=layer_idx in arch.local_attn_layers,
-        ))
-
-        is_gelu7 = layer_idx in arch.gelu7_layers
-        if is_gelu7:
-            intermediate = arch.hidden_size * 4
-            gated = False
-            activation = "gelu7"
-        else:
-            intermediate = (arch.hidden_size * 4 * 2 // 3) // 4 * 4
-            gated = True
-            activation = "swiglu7"
-
-        self.mlp = MagiMLP(MLPSubConfig(
-            hidden_size=arch.hidden_size,
-            intermediate_size=intermediate,
-            activation=activation,
-            num_modality=num_modality,
-            gated=gated,
-        ))
-
-        if self.post_norm:
-            self.attn_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality)
-            self.mlp_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        rope: torch.Tensor,
-        permute_mapping: torch.Tensor,
-        inv_permute_mapping: torch.Tensor,
-        modality_dispatcher: ModalityDispatcher,
-        num_video_tokens: int | None = None,
-        num_frames: int | None = None,
-    ) -> torch.Tensor:
-        attn_out = self.attention(
-            hidden_states, rope, permute_mapping, inv_permute_mapping, modality_dispatcher,
-            num_video_tokens=num_video_tokens,
-            num_frames=num_frames,
-        )
-        if self.post_norm:
-            attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
-        hidden_states = hidden_states + attn_out
-
-        mlp_out = self.mlp(hidden_states, modality_dispatcher=modality_dispatcher)
-        if self.post_norm:
-            mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher)
-        return hidden_states + mlp_out
-
-
-# ---------------------------------------------------------------------------
-# Adapter (per-modality embedders + Fourier RoPE producer)
-# ---------------------------------------------------------------------------
-
-
-class MagiAdapter(nn.Module):
-    def __init__(self, arch: MagiHumanArchConfig):
-        super().__init__()
-        # Embedders stay in fp32 to match the reference dtype exactly.
-        self.video_embedder = nn.Linear(
-            arch.video_in_channels, arch.hidden_size, bias=True, dtype=torch.float32,
-        )
-        self.text_embedder = nn.Linear(
-            arch.text_in_channels, arch.hidden_size, bias=True, dtype=torch.float32,
-        )
-        self.audio_embedder = nn.Linear(
-            arch.audio_in_channels, arch.hidden_size, bias=True, dtype=torch.float32,
-        )
-        self.rope = ElementWiseFourierEmbed(arch.head_dim)
-        # RoPE cache: coords_mapping is the same tensor object across timesteps
-        # in the denoising loop, so data_ptr()+shape+dtype+device is a fast,
-        # collision-free key that avoids recomputing the Fourier embed each step.
-        self._cached_rope: Optional[torch.Tensor] = None
-        self._cached_rope_key: Optional[tuple] = None
-
-    def _rope_cache_key(self, t: torch.Tensor) -> tuple:
-        return (t.data_ptr(), t.shape, t.dtype, t.device)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        coords_mapping: torch.Tensor,
-        video_mask: torch.Tensor,
-        audio_mask: torch.Tensor,
-        text_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        key = self._rope_cache_key(coords_mapping)
-        if key != self._cached_rope_key:
-            self._cached_rope = self.rope(coords_mapping)
-            self._cached_rope_key = key
-        rope = self._cached_rope
-        # Embedder dtypes may differ from x's dtype when FastVideo's FSDP
-        # loader casts all weights to `pipeline_config.precision` (bf16).
-        # Match the weight dtype per modality.
-        v_w = self.video_embedder.weight
-        a_w = self.audio_embedder.weight
-        t_w = self.text_embedder.weight
-        out = torch.zeros(
-            x.shape[0], self.video_embedder.out_features,
-            device=x.device, dtype=v_w.dtype,
-        )
-        out[text_mask] = self.text_embedder(
-            x[text_mask, : self.text_embedder.in_features].to(t_w.dtype)
-        ).to(out.dtype)
-        out[audio_mask] = self.audio_embedder(
-            x[audio_mask, : self.audio_embedder.in_features].to(a_w.dtype)
-        ).to(out.dtype)
-        out[video_mask] = self.video_embedder(
-            x[video_mask, : self.video_embedder.in_features].to(v_w.dtype)
-        ).to(out.dtype)
-        return out, rope
-
-
-# ---------------------------------------------------------------------------
-# Top-level DiT
-# ---------------------------------------------------------------------------
-
-
-class _TransformerBlock(nn.Module):
-    """Thin ModuleList wrapper to keep the 'block.layers.<i>' state_dict
-    naming identical to the upstream checkpoint (which uses a magi_compile
-    decorator producing `block.layers.<i>.*`)."""
-
-    def __init__(self, arch: MagiHumanArchConfig):
-        super().__init__()
-        self.layers = nn.ModuleList([
-            MagiTransformerLayer(arch, i) for i in range(arch.num_layers)
-        ])
-
-    def configure_local_attention(
-        self,
-        local_attn_layers: tuple[int, ...],
-        frame_receptive_field: int = 11,
-    ) -> None:
-        enabled_layers = set(local_attn_layers)
-        for idx, layer in enumerate(self.layers):
-            layer.attention.configure_local_attention(
-                enabled=idx in enabled_layers,
-                frame_receptive_field=frame_receptive_field,
-            )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope: torch.Tensor,
-        permute_mapping: torch.Tensor,
-        inv_permute_mapping: torch.Tensor,
-        modality_dispatcher: ModalityDispatcher,
-        num_video_tokens: int | None = None,
-        num_frames: int | None = None,
-    ) -> torch.Tensor:
-        for layer in self.layers:
-            x = layer(
-                x,
-                rope,
-                permute_mapping,
-                inv_permute_mapping,
-                modality_dispatcher,
-                num_video_tokens=num_video_tokens,
-                num_frames=num_frames,
-            )
-        return x
-
-
-_CFG = MagiHumanVideoConfig()
-
-
-class MagiHumanDiT(BaseDiT):
-    """Top-level DiT for daVinci-MagiHuman (base).
-
-    Forward signature mirrors the reference `DiTModel.forward`: it takes a
-    flat token stream, its per-token coords and modality mapping, and
-    returns per-modality outputs packed into a max-channel-width tensor.
-
-    This scaffold is single-GPU only; the `ulysses_scheduler().dispatch(...)`
-    sequence-parallel wrapping in the reference has no equivalent here yet.
-    """
-
-    # BaseDiT requires these class attrs. Source them from the config so
-    # they stay in sync with MagiHumanVideoConfig edits.
-    _fsdp_shard_conditions = _CFG._fsdp_shard_conditions
-    _compile_conditions = _CFG._compile_conditions
-    _supported_attention_backends = _CFG._supported_attention_backends
-    param_names_mapping = _CFG.param_names_mapping
-    reverse_param_names_mapping = _CFG.reverse_param_names_mapping
-    lora_param_names_mapping = _CFG.lora_param_names_mapping
-
-    def __init__(self, config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs):
-        super().__init__(config=config, hf_config=hf_config or {})
-        arch: MagiHumanArchConfig = getattr(config, "arch_config", config)
-        self.arch = arch
-
-        # BaseDiT contract instance vars.
-        self.hidden_size = arch.hidden_size
-        self.num_attention_heads = arch.num_attention_heads
-        self.num_channels_latents = arch.num_channels_latents
-
-        self.adapter = MagiAdapter(arch)
-        self.block = _TransformerBlock(arch)
-        self.final_norm_video = MultiModalityRMSNorm(arch.hidden_size)
-        self.final_norm_audio = MultiModalityRMSNorm(arch.hidden_size)
-        self.final_linear_video = nn.Linear(
-            arch.hidden_size, arch.video_in_channels, bias=False, dtype=torch.float32,
-        )
-        self.final_linear_audio = nn.Linear(
-            arch.hidden_size, arch.audio_in_channels, bias=False, dtype=torch.float32,
-        )
-        # Dispatcher + mask cache: modality_mapping is the same tensor object
-        # across all timesteps in the denoising loop; data_ptr()+shape+dtype+device
-        # is a fast, collision-free key that avoids rebuilding ModalityDispatcher
-        # (which calls argsort + bincount) on every forward call.
-        self._cached_dispatcher: Optional[ModalityDispatcher] = None
-        self._cached_video_mask: Optional[torch.Tensor] = None
-        self._cached_audio_mask: Optional[torch.Tensor] = None
-        self._cached_text_mask: Optional[torch.Tensor] = None
-        self._cached_modality_key: Optional[tuple] = None
-
-    def configure_local_attention(
-        self,
-        local_attn_layers: tuple[int, ...] | list[int],
-        frame_receptive_field: int = 11,
-    ) -> None:
-        layers = tuple(int(layer) for layer in local_attn_layers)
-        self.arch.local_attn_layers = layers
-        self.block.configure_local_attention(layers, frame_receptive_field)
-
-    def _modality_cache_key(self, t: torch.Tensor) -> tuple:
-        return (t.data_ptr(), t.shape, t.dtype, t.device)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        coords_mapping: torch.Tensor,
-        modality_mapping: torch.Tensor,
-    ) -> torch.Tensor:
-        """
-        Args:
-            x:                [L, max(V_ch, A_ch, T_ch)]
-            coords_mapping:   [L, 9]
-            modality_mapping: [L]  (int in {VIDEO, AUDIO, TEXT})
-        Returns:
-            out: [L, max(V_ch, A_ch)] with video channels in video slots and
-                 audio channels in audio slots; text slots are zero.
-        """
-        key = self._modality_cache_key(modality_mapping)
-        if key != self._cached_modality_key:
-            self._cached_dispatcher = ModalityDispatcher(modality_mapping, num_modalities=3)
-            self._cached_video_mask = modality_mapping == Modality.VIDEO
-            self._cached_audio_mask = modality_mapping == Modality.AUDIO
-            self._cached_text_mask = modality_mapping == Modality.TEXT
-            self._cached_modality_key = key
-        dispatcher = self._cached_dispatcher
-        video_mask = self._cached_video_mask
-        audio_mask = self._cached_audio_mask
-        text_mask = self._cached_text_mask
-        num_video_tokens = int(video_mask.sum().item())
-        if num_video_tokens:
-            num_frames = int(coords_mapping[:num_video_tokens, 0].max().item()) + 1
-        else:
-            num_frames = 0
-
-        x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
-        # Keep the residual stream in adapter dtype (fp32) entering the block.
-        # Upstream daVinci-MagiHuman dit_module.py:923 casts to params_dtype,
-        # which is fp32 by default; each layer's pre_norm.to(bf16) handles
-        # the bf16 internal-compute boundary, and linear_proj outputs bf16
-        # which gets promoted back to fp32 by the residual addition. Casting
-        # the residual to bf16 here degrades the cross-layer accumulator and
-        # compounds visibly over 40 layers in pipeline parity.
-        x = ModalityDispatcher.permute(x, dispatcher.permute_mapping)
-
-        x = self.block(
-            x, rope,
-            permute_mapping=dispatcher.permute_mapping,
-            inv_permute_mapping=dispatcher.inv_permute_mapping,
-            modality_dispatcher=dispatcher,
-            num_video_tokens=num_video_tokens,
-            num_frames=num_frames,
-        )
-        x = ModalityDispatcher.inv_permute(x, dispatcher.inv_permute_mapping)
-
-        x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
-        x_video = self.final_norm_video(x_video)
-        x_video = self.final_linear_video(x_video)
-
-        x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
-        x_audio = self.final_norm_audio(x_audio)
-        x_audio = self.final_linear_audio(x_audio)
-
-        max_ch = max(self.arch.video_in_channels, self.arch.audio_in_channels)
-        out = torch.zeros(x.shape[0], max_ch, device=x.device, dtype=x.dtype)
-        out[video_mask, : self.arch.video_in_channels] = x_video.to(out.dtype)
-        out[audio_mask, : self.arch.audio_in_channels] = x_audio.to(out.dtype)
-        return out
-
-
-EntryClass = MagiHumanDiT

fastvideo/models/encoders/t5gemma.py

@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""T5-Gemma encoder wrapper for daVinci-MagiHuman.
-
-MagiHuman uses `transformers.models.t5gemma.T5GemmaEncoderModel` on
-`google/t5gemma-9b-9b-ul2` (a gated Google repo). This wrapper follows the
-same lazy-loading pattern as `fastvideo/models/encoders/gemma.py`: we keep
-the HF module under `self._t5gemma_model` and exclude it from
-`named_parameters` so FastVideo's weight loader does not try to load
-encoder shards from the converted repo directory.
-
-For the base MagiHuman T2V port there are no additional connector layers
-on top — the pipeline prompt-preprocessing stage handles pad-or-trim to
-`text_len` and exposes both the padded embedding and the original length.
-"""
-from __future__ import annotations
-
-import os
-
-import torch
-from torch import nn
-
-from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig
-from fastvideo.models.encoders.base import TextEncoder
-from fastvideo.platforms import AttentionBackendEnum
-
-
-class T5GemmaEncoderModel(TextEncoder):
-    """Thin wrapper over HuggingFace's `T5GemmaEncoderModel`.
-
-    On first `forward`, the wrapper lazily instantiates the upstream encoder
-    from `t5gemma_model_path` (defaulting to `google/t5gemma-9b-9b-ul2`).
-    Afterwards, forward returns a `BaseEncoderOutput` with
-    `last_hidden_state = [B, L, 3584]` matching MagiHuman's
-    `context.half()` output.
-    """
-
-    _supported_attention_backends = (
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-    )
-
-    def __init__(self, config: TextEncoderConfig) -> None:
-        super().__init__(config)
-        arch = config.arch_config
-        self.t5gemma_model_path: str = arch.t5gemma_model_path
-        self.t5gemma_dtype: str = arch.t5gemma_dtype
-        self._t5gemma_model = None
-
-    def named_parameters(self, prefix: str = "", recurse: bool = True):
-        # The upstream encoder is loaded lazily and its parameters are
-        # managed by HF, not FastVideo's loader. Hide them from the parent
-        # module-tree traversal so Diffusers-repo weight loading does not
-        # try to match them.
-        for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
-            if name.startswith("_t5gemma_model.") or name == "_t5gemma_model":
-                continue
-            yield name, param
-
-    def _build_t5gemma_model(self, device: torch.device | None = None):
-        from transformers.models.t5gemma import T5GemmaEncoderModel as HFEncoder
-
-        path = self.t5gemma_model_path
-        if not path:
-            raise ValueError(
-                "t5gemma_model_path must be set. Expected "
-                "`google/t5gemma-9b-9b-ul2` or a local path to an "
-                "equivalent T5-Gemma encoder."
-            )
-        dtype = getattr(torch, self.t5gemma_dtype, torch.bfloat16)
-        model = HFEncoder.from_pretrained(
-            path,
-            is_encoder_decoder=False,
-            dtype=dtype,
-        )
-        if os.getenv("FASTVIDEO_ATTENTION_BACKEND") == "TORCH_SDPA":
-            if hasattr(model.config, "attn_implementation"):
-                model.config.attn_implementation = "sdpa"
-            if hasattr(model.config, "_attn_implementation"):
-                model.config._attn_implementation = "sdpa"
-        if device is not None:
-            model = model.to(device=device)
-        model.eval()
-        return model
-
-    @property
-    def t5gemma_model(self):
-        if self._t5gemma_model is None:
-            # Lazy-load on CPU if no device is known yet; `forward` will
-            # move the model to the input's device on first call.
-            self._t5gemma_model = self._build_t5gemma_model()
-        return self._t5gemma_model
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        output_hidden_states: bool | None = None,
-        **kwargs,
-    ) -> BaseEncoderOutput:
-        # Ensure the lazy-loaded encoder lives on the same device as the
-        # input; lazy-loading leaves it on CPU until the first forward.
-        ref = input_ids if input_ids is not None else inputs_embeds
-        target_device = ref.device if ref is not None else None
-        model = self.t5gemma_model
-        if target_device is not None:
-            first_param = next(model.parameters(), None)
-            if first_param is not None and first_param.device != target_device:
-                model = model.to(device=target_device)
-                self._t5gemma_model = model
-        outputs = model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            inputs_embeds=inputs_embeds,
-            output_hidden_states=bool(output_hidden_states),
-        )
-        # MagiHuman casts to fp16 at this point; keep the raw dtype here and
-        # leave precision management to the pipeline's postprocess stage.
-        return BaseEncoderOutput(
-            last_hidden_state=outputs["last_hidden_state"],
-            hidden_states=getattr(outputs, "hidden_states", None),
-            attention_mask=attention_mask,
-        )
-
-
-EntryClass = T5GemmaEncoderModel

fastvideo/models/loader/component_loader.py

@@ -78,7 +78,6 @@ class ComponentLoader(ABC):
         module_loaders = {
             "scheduler": (SchedulerLoader, "diffusers"),
             "transformer": (TransformerLoader, "diffusers"),
-            "sr_transformer": (TransformerLoader, "diffusers"),
             "transformer_2": (TransformerLoader, "diffusers"),
             "transformer_3": (TransformerLoader, "diffusers"),
             "vae": (VAELoader, "diffusers"),

fastvideo/pipelines/basic/magi_human/__init__.py

@@ -1 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0

fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py

@@ -1,417 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""MagiHuman base text-to-AV pipeline.
-
-Top-level composition for the daVinci-MagiHuman base model. Wires:
-
-    InputValidationStage -> TextEncodingStage (T5-Gemma)
-        -> MagiHumanLatentPreparationStage
-        -> MagiHumanDenoisingStage
-        -> DecodingStage (Wan 2.2 TI2V-5B VAE decode for video)
-        -> MagiHumanAudioDecodingStage (Stable Audio Open 1.0 VAE decode)
-
-The base checkpoint is a joint audio-visual generator; both the video
-and audio paths run in the denoising loop and both are decoded.
-
-`load_modules` is overridden so the four cross-variant shared components
-(text_encoder, tokenizer, audio_vae, video vae) lazy-load from their
-canonical upstream HF repos at first build time instead of being
-bundled inside every converted MagiHuman variant. This keeps each
-variant's converted repo at ~5-30 GB (transformer + scheduler +
-model_index.json) instead of ~30-55 GB, and lets all variants share
-the same ~25 GB of cached upstream weights.
-"""
-from __future__ import annotations
-
-import json
-import os
-from pathlib import Path
-from typing import Any
-
-from transformers import AutoTokenizer
-
-from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
-from fastvideo.configs.models.vaes import OobleckVAEConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.encoders.t5gemma import T5GemmaEncoderModel
-from fastvideo.models.vaes.sa_audio import SAAudioVAEModel
-from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
-    FlowUniPCMultistepScheduler, )
-from fastvideo.pipelines.basic.magi_human.stages import (
-    MagiHumanAudioDecodingStage,
-    MagiHumanDenoisingStage,
-    MagiHumanLatentPreparationStage,
-    MagiHumanReferenceImageStage,
-    MagiHumanSRDenoisingStage,
-    MagiHumanSRLatentPreparationStage,
-)
-from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.pipelines.stages import (
-    DecodingStage,
-    InputValidationStage,
-    TextEncodingStage,
-)
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-_T5GEMMA_HF_ID = "google/t5gemma-9b-9b-ul2"
-_SA_AUDIO_HF_ID = "stabilityai/stable-audio-open-1.0"
-_WAN_VAE_HF_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-
-
-def _ensure_hf_token_env() -> str | None:
-    """Surface any of the three common HF token env vars as `HF_TOKEN`.
-
-    FastVideo workers spawn child processes that inherit env; both
-    `huggingface_hub` and `transformers.AutoTokenizer.from_pretrained`
-    look at `HF_TOKEN` / `HUGGINGFACE_HUB_TOKEN` by default but not
-    `HF_API_KEY`. If only the latter is set, gated downloads fail with
-    401. Aliasing at pipeline-load time is the minimum-disruption fix.
-    """
-    for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        value = os.environ.get(src)
-        if value:
-            os.environ.setdefault("HF_TOKEN", value)
-            os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", value)
-            return value
-    return None
-
-
-class MagiHumanPipeline(ComposedPipelineBase):
-    """Base MagiHuman text-to-AV pipeline (no LoRA, no distill, no SR)."""
-
-    _required_config_modules = [
-        "text_encoder",
-        "tokenizer",
-        "vae",
-        "transformer",
-        "scheduler",
-        "audio_vae",
-    ]
-
-    def load_modules(
-        self,
-        fastvideo_args: FastVideoArgs,
-        loaded_modules: dict[str, Any] | None = None,
-    ) -> dict[str, Any]:
-        """Load the variant-specific transformer + scheduler from the
-        converted MagiHuman repo and lazy-load the four cross-variant
-        shared components from their canonical upstream HF repos:
-
-          * text_encoder, tokenizer -> ``google/t5gemma-9b-9b-ul2``
-            (gated, requires HF token with accepted terms of use)
-          * audio_vae -> ``stabilityai/stable-audio-open-1.0`` (gated)
-          * vae -> ``Wan-AI/Wan2.2-TI2V-5B-Diffusers``
-
-        Backwards-compatible with bundled converted repos: if any of
-        these subfolders is present locally and listed in
-        ``model_index.json``, the standard component loader picks it up
-        via super(). Otherwise the loader is told to skip the entry and
-        we lazy-load it here.
-        """
-        # T5-Gemma is gated: expose `HF_API_KEY` as `HF_TOKEN` if needed.
-        _ensure_hf_token_env()
-
-        # Resolve to a local cache path so we can inspect
-        # model_index.json before invoking super(). `maybe_download_model`
-        # is idempotent for local paths; super() repeats the call cheaply
-        # via `_load_config`.
-        local_path = maybe_download_model(self.model_path)
-
-        # Identify which cross-variant shared keys are bundled in the
-        # converted repo (declared in model_index.json with a non-null
-        # spec) versus absent (the umbrella scheme). Bundled keys stay
-        # in `required_config_modules` and are loaded normally by super()
-        # from `<model_path>/<key>/`. Absent keys are temporarily
-        # dropped so super() does not fail the "every required entry
-        # must appear in model_index.json" check, then lazy-loaded
-        # below.
-        model_index: dict[str, Any] = {}
-        try:
-            with open(Path(local_path) / "model_index.json") as f:
-                model_index = json.load(f)
-        except (FileNotFoundError, json.JSONDecodeError):
-            pass
-
-        def _is_bundled(key: str) -> bool:
-            spec = model_index.get(key)
-            return (isinstance(spec, list | tuple) and len(spec) >= 1 and spec[0] is not None)
-
-        deferred = []
-        for key in ("text_encoder", "tokenizer", "audio_vae", "vae"):
-            if key in self.required_config_modules and not _is_bundled(key):
-                self.required_config_modules.remove(key)
-                deferred.append(key)
-
-        try:
-            modules = super().load_modules(fastvideo_args, loaded_modules)
-        finally:
-            for key in deferred:
-                if key not in self.required_config_modules:
-                    self.required_config_modules.append(key)
-
-        # For each lazy-load key, prefer whatever super() already loaded
-        # (a bundled subfolder, or a caller-provided override merged in
-        # via `loaded_modules`). Fall back to the caller-provided
-        # `loaded_modules` entry for keys absent from model_index.json
-        # (super() never iterates those). Otherwise lazy-load from the
-        # canonical upstream HF repo.
-        def _resolve(key: str) -> bool:
-            """Return True if `modules[key]` is already populated."""
-            if modules.get(key) is not None:
-                return True
-            if loaded_modules and key in loaded_modules:
-                modules[key] = loaded_modules[key]
-                return True
-            return False
-
-        if not _resolve("text_encoder"):
-            logger.info("Building T5-Gemma text encoder (lazy-load from %s)", _T5GEMMA_HF_ID)
-            enc_config = T5GemmaEncoderConfig()
-            enc_config.arch_config.t5gemma_model_path = _T5GEMMA_HF_ID
-            modules["text_encoder"] = T5GemmaEncoderModel(enc_config)
-
-        if not _resolve("tokenizer"):
-            logger.info("Loading T5-Gemma tokenizer from %s", _T5GEMMA_HF_ID)
-            modules["tokenizer"] = AutoTokenizer.from_pretrained(_T5GEMMA_HF_ID)
-
-        if not _resolve("audio_vae"):
-            logger.info(
-                "Building Stable Audio Open 1.0 VAE (lazy-load from %s) — "
-                "requires HF terms accepted for gated repo",
-                _SA_AUDIO_HF_ID,
-            )
-            audio_config = OobleckVAEConfig()
-            audio_config.pretrained_path = _SA_AUDIO_HF_ID
-            modules["audio_vae"] = SAAudioVAEModel(audio_config)
-
-        if not _resolve("vae"):
-            modules["vae"] = self._load_video_vae(fastvideo_args)
-
-        return modules
-
-    def _load_video_vae(self, fastvideo_args: FastVideoArgs) -> Any:
-        """Resolve the video VAE: prefer a bundled ``vae/`` subfolder in
-        the converted repo (legacy), fall back to lazy-downloading the
-        Wan 2.2 TI2V-5B VAE shards from upstream.
-
-        Either way the load goes through FastVideo's standard
-        ``VAELoader`` so the result is the same FV ``AutoencoderKLWan``
-        nn.Module that the bundled path produces.
-        """
-        from fastvideo.models.loader.component_loader import VAELoader
-
-        bundled = Path(self.model_path) / "vae"
-        if bundled.is_dir() and (bundled / "config.json").is_file():
-            logger.info("Loading bundled video VAE from %s", bundled)
-            return VAELoader().load(str(bundled), fastvideo_args)
-
-        from huggingface_hub import snapshot_download
-
-        logger.info(
-            "Bundled vae/ not found at %s; lazy-loading Wan 2.2 TI2V-5B VAE from %s",
-            self.model_path,
-            _WAN_VAE_HF_ID,
-        )
-        snapshot = snapshot_download(
-            repo_id=_WAN_VAE_HF_ID,
-            allow_patterns=["vae/*"],
-        )
-        vae_dir = os.path.join(snapshot, "vae")
-        if not os.path.isdir(vae_dir):
-            raise RuntimeError(
-                f"snapshot_download returned {snapshot} but no vae/ "
-                f"subfolder was found inside it. Check that {_WAN_VAE_HF_ID} "
-                "still exposes a Diffusers-format vae/ folder.", )
-        return VAELoader().load(vae_dir, fastvideo_args)
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
-        # MagiHuman applies `flow_shift` during timestep setup; keep the
-        # scheduler constructor at its default no-op shift.
-        self.modules["scheduler"] = FlowUniPCMultistepScheduler()
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self._add_input_and_conditioning_stages(fastvideo_args)
-        self._add_base_latent_and_denoising_stages(fastvideo_args)
-        self._add_decode_stages()
-
-    def _add_input_and_conditioning_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(
-            stage_name="input_validation_stage",
-            stage=InputValidationStage(),
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self._add_reference_image_stage(fastvideo_args)
-
-    def _add_base_latent_and_denoising_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        pc = fastvideo_args.pipeline_config
-        dit_arch = pc.dit_config.arch_config
-
-        # Data-proxy + eval knobs come from the PipelineConfig (`pc`).
-        # Only DiT-architecture fields live on `dit_arch` now.
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=MagiHumanLatentPreparationStage(
-                vae_stride=tuple(pc.vae_stride),
-                z_dim=pc.z_dim,
-                patch_size=tuple(dit_arch.patch_size),
-                fps=pc.fps,
-                t5_gemma_target_length=pc.t5_gemma_target_length,
-                coords_style=pc.coords_style,
-                text_offset=pc.text_offset,
-                audio_in_channels=dit_arch.audio_in_channels,
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=MagiHumanDenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-                patch_size=tuple(dit_arch.patch_size),
-                video_in_channels=dit_arch.video_in_channels,
-                audio_in_channels=dit_arch.audio_in_channels,
-                video_txt_guidance_scale=pc.video_txt_guidance_scale,
-                audio_txt_guidance_scale=pc.audio_txt_guidance_scale,
-                cfg_number=pc.cfg_number,
-                coords_style=pc.coords_style,
-                video_guidance_high_t_threshold=pc.video_guidance_high_t_threshold,
-                video_guidance_low_t_value=pc.video_guidance_low_t_value,
-            ),
-        )
-
-    def _add_decode_stages(self) -> None:
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=DecodingStage(vae=self.get_module("vae"), pipeline=self),
-        )
-
-        self.add_stage(
-            stage_name="audio_decoding_stage",
-            stage=MagiHumanAudioDecodingStage(audio_vae=self.get_module("audio_vae"), ),
-        )
-
-    def _add_reference_image_stage(self, fastvideo_args: FastVideoArgs) -> None:
-        return
-
-
-class MagiHumanI2VPipeline(MagiHumanPipeline):
-    """MagiHuman text+image-to-AV pipeline using the T2V DiT weights."""
-
-    def _add_reference_image_stage(self, fastvideo_args: FastVideoArgs) -> None:
-        pc = fastvideo_args.pipeline_config
-        self.add_stage(
-            stage_name="reference_image_stage",
-            stage=MagiHumanReferenceImageStage(
-                vae=self.get_module("vae"),
-                vae_scale_factor=pc.vae_stride[1],
-            ),
-        )
-
-
-class MagiHumanSRPipeline(MagiHumanPipeline):
-    """Two-stage MagiHuman base + SR-540p text-to-AV pipeline."""
-
-    _required_config_modules = [
-        "text_encoder",
-        "tokenizer",
-        "vae",
-        "transformer",
-        "sr_transformer",
-        "scheduler",
-        "audio_vae",
-    ]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self._add_input_and_conditioning_stages(fastvideo_args)
-        self._add_base_latent_and_denoising_stages(fastvideo_args)
-        self._add_sr_latent_and_denoising_stages(fastvideo_args)
-        self._add_decode_stages()
-
-    def _add_sr_latent_and_denoising_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        pc = fastvideo_args.pipeline_config
-        dit_arch = pc.dit_config.arch_config
-        sr_transformer = self.get_module("sr_transformer")
-        sr_local_attn_layers = tuple(getattr(pc, "sr_local_attn_layers", ()))
-        if sr_local_attn_layers and hasattr(sr_transformer, "configure_local_attention"):
-            sr_transformer.configure_local_attention(
-                sr_local_attn_layers,
-                frame_receptive_field=pc.frame_receptive_field,
-            )
-
-        self.add_stage(
-            stage_name="sr_latent_preparation_stage",
-            stage=MagiHumanSRLatentPreparationStage(
-                vae=self.get_module("vae"),
-                vae_stride=tuple(pc.vae_stride),
-                patch_size=tuple(dit_arch.patch_size),
-                noise_value=pc.noise_value,
-                sr_audio_noise_scale=pc.sr_audio_noise_scale,
-                sr_height=pc.sr_height,
-                sr_width=pc.sr_width,
-                vae_scale_factor=pc.vae_stride[1],
-            ),
-        )
-        self.add_stage(
-            stage_name="sr_denoising_stage",
-            stage=MagiHumanSRDenoisingStage(
-                transformer=sr_transformer,
-                scheduler=self.get_module("scheduler"),
-                patch_size=tuple(dit_arch.patch_size),
-                video_in_channels=dit_arch.video_in_channels,
-                audio_in_channels=dit_arch.audio_in_channels,
-                sr_num_inference_steps=pc.sr_num_inference_steps,
-                sr_video_txt_guidance_scale=pc.sr_video_txt_guidance_scale,
-                use_cfg_trick=pc.use_cfg_trick,
-                cfg_trick_start_frame=pc.cfg_trick_start_frame,
-                cfg_trick_value=pc.cfg_trick_value,
-                cfg_number=pc.cfg_number,
-                coords_style="v1",
-            ),
-        )
-
-
-class MagiHumanSRI2VPipeline(MagiHumanSRPipeline):
-    """Two-stage MagiHuman base + SR-540p text+image-to-AV pipeline."""
-
-    def _add_reference_image_stage(self, fastvideo_args: FastVideoArgs) -> None:
-        pc = fastvideo_args.pipeline_config
-        self.add_stage(
-            stage_name="reference_image_stage",
-            stage=MagiHumanReferenceImageStage(
-                vae=self.get_module("vae"),
-                vae_scale_factor=pc.vae_stride[1],
-            ),
-        )
-
-
-class MagiHumanSR1080pPipeline(MagiHumanSRPipeline):
-    """Two-stage MagiHuman base + SR-1080p text-to-AV pipeline.
-
-    The stage chain is identical to SR-540p. The paired pipeline config enables
-    block-sparse local-window attention on 32 SR-DiT layers and requests the
-    1080p latent target.
-    """
-
-
-class MagiHumanSR1080pI2VPipeline(MagiHumanSRI2VPipeline):
-    """Two-stage MagiHuman base + SR-1080p text+image-to-AV pipeline."""
-
-
-EntryClass = [
-    MagiHumanPipeline,
-    MagiHumanI2VPipeline,
-    MagiHumanSRPipeline,
-    MagiHumanSRI2VPipeline,
-    MagiHumanSR1080pPipeline,
-    MagiHumanSR1080pI2VPipeline,
-]

fastvideo/pipelines/basic/magi_human/pipeline_configs.py

@@ -1,236 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""PipelineConfig for the daVinci-MagiHuman base text-to-AV pipeline."""
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.configs.models.dits import MagiHumanVideoConfig
-from fastvideo.configs.models.encoders import (
-    BaseEncoderOutput,
-    T5GemmaEncoderConfig,
-)
-from fastvideo.configs.models.vaes import OobleckVAEConfig, WanVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-
-
-def t5gemma_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-    """Return per-prompt last_hidden_state as a batched [B, L, D] tensor.
-
-    MagiHuman pads/trims the embedding to a fixed length in its own
-    `pad_or_trim` helper at pipeline time. Here we simply hand through
-    whatever the tokenizer produced — the latent-prep stage is responsible
-    for pad/trim so that the original context length can be preserved.
-    """
-    hidden = outputs.last_hidden_state
-    assert torch.isnan(hidden).sum() == 0
-    # Keep the shape the tokenizer emitted; the pipeline stage handles
-    # pad-or-trim to t5_gemma_target_length=640.
-    return hidden
-
-
-@dataclass
-class MagiHumanBaseConfig(PipelineConfig):
-    """Base MagiHuman text-to-AV pipeline config (prompt → video + audio).
-
-    MagiHuman's base model is a joint audio-visual generator. This config
-    wires up both the video VAE (Wan 2.2 TI2V-5B) and the audio VAE
-    (Stable Audio Open 1.0); the pipeline produces an mp4 with a muxed
-    audio track. The framework's `WorkloadType` enum has no `T2AV`
-    variant yet, so the registry entry uses `WorkloadType.T2V` as a
-    placeholder.
-    """
-
-    # DiT
-    dit_config: DiTConfig = field(default_factory=MagiHumanVideoConfig)
-    # VAE — Wan 2.2 TI2V-5B. Diffusers `vae/config.json` drives arch_config
-    # at load time, including z_dim=48 and scale_factor_temporal=4 /
-    # scale_factor_spatial=16.
-    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
-    vae_tiling: bool = False
-    vae_sp: bool = False
-
-    # Audio VAE — Stable Audio Open 1.0 (Oobleck), shared with the
-    # standalone Stable Audio pipeline. Lazy-loaded from
-    # `stabilityai/stable-audio-open-1.0` (HF gated, Apache 2.0).
-    audio_vae_config: VAEConfig = field(default_factory=OobleckVAEConfig)
-
-    # Denoising (flow-matching UniPC).
-    flow_shift: float | None = 5.0
-
-    # Text encoding
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5GemmaEncoderConfig(), ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda: (t5gemma_postprocess_text, ))
-
-    # Precisions — the DiT runs bf16 internally, the text encoder is
-    # bf16-native, and the VAE decode path benefits from fp32 for long
-    # sequences.
-    precision: str = "bf16"
-    vae_precision: str = "fp32"
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
-
-    # MagiHuman-specific defaults surfaced for the pipeline stages. These
-    # are pipeline-level knobs sourced from the upstream
-    # `EvaluationConfig` / `DataProxyConfig` (not `ModelConfig`), so they
-    # belong here, NOT on `MagiHumanArchConfig`.
-    t5_gemma_target_length: int = 640
-    fps: int = 25
-    num_inference_steps: int = 32
-    video_txt_guidance_scale: float = 5.0
-    audio_txt_guidance_scale: float = 5.0
-    cfg_number: int = 2
-
-    # VAE / data-proxy knobs (were on ArchConfig before; moved here).
-    vae_stride: tuple[int, int, int] = (4, 16, 16)
-    z_dim: int = 48
-    frame_receptive_field: int = 11
-    coords_style: str = "v2"
-    ref_audio_offset: int = 1000
-    text_offset: int = 0
-
-    # Video CFG step-dependent guidance: low-t steps use a relaxed scale.
-    # Upstream daVinci-MagiHuman/inference/pipeline/video_generate.py:426
-    # uses 5.0 for high-t and 2.0 for low-t with cutoff at t=500.
-    video_guidance_high_t_threshold: int = 500
-    video_guidance_low_t_value: float = 2.0
-
-    def __post_init__(self) -> None:
-        # Base text-to-AV does not need the VAE encoder (no reference-image
-        # conditioning). Keep decoder only to save memory.
-        self.vae_config.load_encoder = False
-        self.vae_config.load_decoder = True
-
-
-@dataclass
-class MagiHumanBaseI2VConfig(MagiHumanBaseConfig):
-    """Base MagiHuman text+image-to-AV pipeline config.
-
-    TI2V reuses the T2V DiT weights; the only pipeline-side difference is
-    that a reference image is encoded with the Wan VAE and reinserted into
-    the first video-latent frame before every denoise step.
-    """
-
-    image_conditioning: bool = True
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-
-
-@dataclass
-class MagiHumanDistillConfig(MagiHumanBaseConfig):
-    """DMD-2 distilled MagiHuman text-to-AV pipeline config.
-
-    Same arch as base (identical 331 keys, same shapes, same module tree),
-    but trained via DMD-2 for 8-step inference without classifier-free
-    guidance. Weights are stored in fp32 upstream; the conversion script's
-    `--cast-bf16` flag reduces the checkpoint to ~30 GB on disk.
-    """
-
-    num_inference_steps: int = 8
-    cfg_number: int = 1  # DMD distilled models skip CFG.
-    # Lower flow_shift matches the distilled DMD schedule; if parity later
-    # shows drift, measure against `scheduler_config.json` generated by the
-    # conversion script for the distill subfolder.
-    flow_shift: float | None = 5.0
-
-
-@dataclass
-class MagiHumanDistillI2VConfig(MagiHumanDistillConfig):
-    """DMD-2 distilled MagiHuman text+image-to-AV pipeline config."""
-
-    image_conditioning: bool = True
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-
-
-@dataclass
-class MagiHumanSR540pConfig(MagiHumanBaseConfig):
-    """Two-stage MagiHuman base + SR-540p text-to-AV pipeline config."""
-
-    noise_value: int = 220
-    sr_audio_noise_scale: float = 0.7
-    sr_num_inference_steps: int = 5
-    sr_video_txt_guidance_scale: float = 3.5
-    use_cfg_trick: bool = True
-    cfg_trick_start_frame: int = 13
-    cfg_trick_value: float = 2.0
-    # Upstream example/sr_540p uses sr_height=512, sr_width=896. Despite the
-    # marketing name, these are the VAE/patch-aligned dimensions actually run.
-    sr_height: int = 512
-    sr_width: int = 896
-    sr_local_attn_layers: tuple[int, ...] = ()
-
-
-@dataclass
-class MagiHumanSR540pI2VConfig(MagiHumanSR540pConfig):
-    """Two-stage MagiHuman base + SR-540p text+image-to-AV config."""
-
-    image_conditioning: bool = True
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-
-
-_SR_1080P_LOCAL_ATTN_LAYERS: tuple[int, ...] = (
-    0,
-    1,
-    2,
-    4,
-    5,
-    6,
-    8,
-    9,
-    10,
-    12,
-    13,
-    14,
-    16,
-    17,
-    18,
-    20,
-    21,
-    22,
-    24,
-    25,
-    26,
-    28,
-    29,
-    30,
-    32,
-    33,
-    34,
-    35,
-    36,
-    37,
-    38,
-    39,
-)
-
-
-@dataclass
-class MagiHumanSR1080pConfig(MagiHumanSR540pConfig):
-    """Two-stage MagiHuman base + SR-1080p text-to-AV pipeline config."""
-
-    sr_height: int = 1080
-    sr_width: int = 1920
-    sr_local_attn_layers: tuple[int, ...] = _SR_1080P_LOCAL_ATTN_LAYERS
-
-
-@dataclass
-class MagiHumanSR1080pI2VConfig(MagiHumanSR1080pConfig):
-    """Two-stage MagiHuman base + SR-1080p text+image-to-AV config."""
-
-    image_conditioning: bool = True
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True

fastvideo/pipelines/basic/magi_human/presets.py

@@ -1,225 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Presets for the daVinci-MagiHuman pipelines."""
-from fastvideo.api.presets import InferencePreset, PresetStageSpec
-
-# Keep this in sync with upstream MagiEvaluator.negative_prompt
-# (daVinci-MagiHuman/inference/pipeline/video_generate.py:222-224): the
-# video, audio-quality, and speech-delivery blocks all condition CFG.
-_MAGI_HUMAN_NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, "
-                               "paintings, images, static, overall gray, worst quality, low quality, JPEG "
-                               "compression residue, ugly, incomplete, extra fingers, poorly drawn hands, "
-                               "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, "
-                               "still picture, messy background, three legs, many people in the background, "
-                               "walking backwards, low quality, worst quality, poor quality, noise, background "
-                               "noise, hiss, hum, buzz, crackle, static, compression artifacts, MP3 artifacts, "
-                               "digital clipping, distortion, muffled, muddy, unclear, echo, reverb, room echo, "
-                               "over-reverberated, hollow sound, distant, washed out, harsh, shrill, piercing, "
-                               "grating, tinny, thin sound, boomy, bass-heavy, flat EQ, over-compressed, "
-                               "abrupt cut, jarring transition, sudden silence, looping artifact, music, "
-                               "instrumental, sirens, alarms, crowd noise, unrelated sound effects, chaotic, "
-                               "disorganized, messy, cheap sound, emotionless, flat delivery, deadpan, lifeless, "
-                               "apathetic, robotic, mechanical, monotone, flat intonation, undynamic, boring, "
-                               "reading from a script, AI voice, synthetic, text-to-speech, TTS, insincere, "
-                               "fake emotion, exaggerated, overly dramatic, melodramatic, cheesy, cringey, "
-                               "hesitant, unconfident, tired, weak voice, stuttering, stammering, mumbling, "
-                               "slurred speech, mispronounced, bad articulation, lisp, vocal fry, creaky voice, "
-                               "mouth clicks, lip smacks, wet mouth sounds, heavy breathing, audible inhales, "
-                               "plosives, p-pops, coughing, clearing throat, sneezing, speaking too fast, rushed, "
-                               "speaking too slow, dragged out, unnatural pauses, awkward silence, choppy, "
-                               "disjointed, multiple speakers, two voices, background talking, out of tune, "
-                               "off-key, autotune artifacts")
-
-_DENOISE_STAGE = PresetStageSpec(
-    name="denoise",
-    kind="denoising",
-    description="Joint video+audio UniPC flow-matching denoise pass.",
-    allowed_overrides=frozenset({
-        "num_inference_steps",
-        "guidance_scale",
-    }),
-)
-
-MAGI_HUMAN_BASE = InferencePreset(
-    name="magi_human_base",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman base text-to-AV at 256x480, 4s @ 25 fps. "
-                 "Produces an mp4 with muxed audio + video. workload_type "
-                 "is `t2v` because the framework enum has no `t2av` variant yet."),
-    workload_type="t2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        # Upstream pipeline.py:61-64 defaults br_width=480, br_height=272,
-        # and video_generate.py:254-261 snaps height to 256 while width stays
-        # 480, so the rendered default is 256x480.
-        "width": 480,
-        # num_frames is derived by the pipeline as `seconds*fps + 1`; we
-        # surface it here for APIs that expect a concrete default.
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 5.0,  # used as video_txt_guidance_scale
-        "num_inference_steps": 32,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_DISTILL = InferencePreset(
-    name="magi_human_distill",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman DMD-2 distilled text-to-AV at 256x480, 4s @ "
-                 "25 fps. 8-step inference, no classifier-free guidance. Produces "
-                 "an mp4 with muxed audio + video."),
-    workload_type="t2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        # DMD: cfg=1 at the pipeline level. guidance_scale is kept at 1.0
-        # for interop; the DenoisingStage ignores it when cfg_number=1.
-        "guidance_scale": 1.0,
-        "num_inference_steps": 8,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_BASE_TI2V = InferencePreset(
-    name="magi_human_base_ti2v",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman base text+image-to-AV at 256x480, 4s @ 25 fps. "
-                 "The reference image is VAE-encoded and pinned to the first "
-                 "video latent frame at each denoise step."),
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 32,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_DISTILL_TI2V = InferencePreset(
-    name="magi_human_distill_ti2v",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman DMD-2 distilled text+image-to-AV at 256x480, "
-                 "4s @ 25 fps. 8-step inference, no CFG."),
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 1.0,
-        "num_inference_steps": 8,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_SR_540P = InferencePreset(
-    name="magi_human_sr_540p",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman two-stage base + SR-540p text-to-AV. "
-                 "Base pass runs at 256x480; SR pass refines to upstream's "
-                 "aligned 512x896 output with muxed audio."),
-    workload_type="t2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 32,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_SR_540P_TI2V = InferencePreset(
-    name="magi_human_sr_540p_ti2v",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman two-stage base + SR-540p text+image-to-AV. "
-                 "The reference image is encoded at base resolution and then "
-                 "re-encoded at SR resolution before the SR denoise pass."),
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 32,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_SR_1080P = InferencePreset(
-    name="magi_human_sr_1080p",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman two-stage base + SR-1080p text-to-AV. "
-                 "The SR DiT uses upstream local-window attention in 32 of "
-                 "40 layers and refines to 1080p-class output."),
-    workload_type="t2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 32,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-MAGI_HUMAN_SR_1080P_TI2V = InferencePreset(
-    name="magi_human_sr_1080p_ti2v",
-    version=1,
-    model_family="magi_human",
-    description=("daVinci-MagiHuman two-stage base + SR-1080p text+image-to-AV. "
-                 "The SR DiT uses upstream local-window attention in 32 of "
-                 "40 layers; the reference image is re-encoded at SR resolution."),
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "seed": 42,
-        "height": 256,
-        "width": 480,
-        "num_frames": 101,
-        "fps": 25,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 32,
-        "negative_prompt": _MAGI_HUMAN_NEGATIVE_PROMPT,
-    },
-)
-
-ALL_PRESETS = (
-    MAGI_HUMAN_BASE,
-    MAGI_HUMAN_DISTILL,
-    MAGI_HUMAN_BASE_TI2V,
-    MAGI_HUMAN_DISTILL_TI2V,
-    MAGI_HUMAN_SR_540P,
-    MAGI_HUMAN_SR_540P_TI2V,
-    MAGI_HUMAN_SR_1080P,
-    MAGI_HUMAN_SR_1080P_TI2V,
-)

fastvideo/pipelines/basic/magi_human/stages/__init__.py

@@ -1,16 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo.pipelines.basic.magi_human.stages.audio_decoding import MagiHumanAudioDecodingStage
-from fastvideo.pipelines.basic.magi_human.stages.denoising import MagiHumanDenoisingStage
-from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import MagiHumanLatentPreparationStage
-from fastvideo.pipelines.basic.magi_human.stages.reference_image import MagiHumanReferenceImageStage
-from fastvideo.pipelines.basic.magi_human.stages.sr_denoising import MagiHumanSRDenoisingStage
-from fastvideo.pipelines.basic.magi_human.stages.sr_latent_preparation import MagiHumanSRLatentPreparationStage
-
-__all__ = [
-    "MagiHumanAudioDecodingStage",
-    "MagiHumanDenoisingStage",
-    "MagiHumanLatentPreparationStage",
-    "MagiHumanReferenceImageStage",
-    "MagiHumanSRDenoisingStage",
-    "MagiHumanSRLatentPreparationStage",
-]

fastvideo/pipelines/basic/magi_human/stages/audio_decoding.py

@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Audio decoding stage for daVinci-MagiHuman.
-
-Takes the denoised audio latent that `MagiHumanDenoisingStage` leaves
-on `batch.audio_latents` and decodes it to a waveform using the
-Stable Audio Open 1.0 VAE. Mirrors the upstream post-process path
-(see `MagiEvaluator.post_process` in
-daVinci-MagiHuman/inference/pipeline/video_generate.py:503):
-
-    latent_audio.squeeze(0)                 # (L, C_latent)
-    audio = self.audio_vae.decode(latent_audio.T)   # (1, audio_ch, samples)
-    audio = audio.squeeze(0).T.cpu().numpy()        # (samples, audio_ch)
-    audio = resample_audio_sinc(audio, _UPSTREAM_AUDIO_TIME_STRETCH)
-
-The stage stores the resampled waveform on `batch.extra["audio"]`
-(shape `[samples, audio_channels]`) and the sample rate on
-`batch.extra["audio_sample_rate"]`. FastVideo's `VideoGenerator._mux_audio`
-then reads those, writes a temp wav, and muxes it into the output mp4
-via PyAV — same plumbing LTX-2 and Stable Audio use.
-"""
-from __future__ import annotations
-
-import numpy as np
-import torch
-from scipy.signal import resample as _scipy_resample
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-# 441/512 is daVinci-MagiHuman's audio time-stretch ratio that aligns
-# the 44.1 kHz Stable-Audio output with the 25-fps video frame rate.
-# See daVinci-MagiHuman/inference/pipeline/video_generate.py:516.
-_UPSTREAM_AUDIO_TIME_STRETCH = 441.0 / 512.0
-
-# Stable Audio Open 1.0 native sample rate (per stabilityai/stable-audio-open-1.0
-# model card and fastvideo/configs/models/vaes/oobleck.py::OobleckVAEArchConfig.sampling_rate).
-_SA_AUDIO_OPEN_SAMPLE_RATE = 44100
-
-
-def _resample_sinc(audio: np.ndarray, time_stretching: float) -> np.ndarray:
-    """Resample the audio to ``new_length = int(L * time_stretching)`` samples.
-
-    Mirrors upstream ``video_process.resample_audio_sinc`` which calls
-    ``scipy.signal.resample`` (FFT-based polyphase resampling that
-    approximates ideal sinc interpolation). This avoids the
-    high-frequency aliasing and roll-off that ``F.interpolate(mode='linear')``
-    would introduce on a 25 fps × ~5 s wav (`scipy` is already a direct
-    fastvideo dep, so this is dependency-free relative to the previous
-    implementation).
-    """
-    if time_stretching == 1.0:
-        return audio
-    new_length = int(audio.shape[0] * time_stretching)
-    resampled = _scipy_resample(audio.astype(np.float32), new_length, axis=0)
-    return np.asarray(resampled, dtype=np.float32)
-
-
-class MagiHumanAudioDecodingStage(PipelineStage):
-    """Decode `batch.audio_latents` to a waveform using Stable Audio's VAE.
-
-    The VAE is loaded lazily by `SAAudioVAEModel.sa_audio_vae_model` — the
-    first call triggers a snapshot_download (requires HF token + accepted
-    terms on stabilityai/stable-audio-open-1.0).
-    """
-
-    def __init__(
-        self,
-        audio_vae,
-        time_stretching: float = _UPSTREAM_AUDIO_TIME_STRETCH,
-    ) -> None:
-        super().__init__()
-        self.audio_vae = audio_vae
-        self.time_stretching = time_stretching
-
-    def verify_input(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def verify_output(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    @torch.inference_mode()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        latent_audio = getattr(batch, "audio_latents", None)
-        if latent_audio is None:
-            # Joint AV: missing audio latents means the denoising stage broke.
-            raise ValueError("MagiHumanAudioDecodingStage requires batch.audio_latents to be set. "
-                             "Did the denoising stage produce them? Joint AV pipeline expects "
-                             "both video and audio latents from MagiHumanDenoisingStage.")
-
-        # Upstream shape: `[B, L, C_latent]` from the DiT; AutoencoderOobleck
-        # expects `[B, C_latent, L]`. MagiEvaluator.post_process does
-        # `latent_audio.squeeze(0); audio_vae.decode(latent_audio.T)`
-        # (which yields `[C_latent, L]`, implicit batch=1). We keep the
-        # batch dim and transpose L<->C.
-        latent_bcl = latent_audio.permute(0, 2, 1).contiguous()
-
-        # Decode: [B, C_latent, L] -> [B, audio_channels, samples]
-        audio_out = self.audio_vae.decode(latent_bcl)
-
-        audio_np = audio_out.squeeze(0).T.float().cpu().numpy()
-        audio_np = _resample_sinc(audio_np, self.time_stretching)
-
-        # Conform to FastVideo convention: VideoGenerator._mux_audio
-        # reads these two keys and muxes via PyAV.
-        if batch.extra is None:
-            batch.extra = {}
-        batch.extra["audio"] = audio_np
-        batch.extra["audio_sample_rate"] = int(getattr(self.audio_vae, "sampling_rate", _SA_AUDIO_OPEN_SAMPLE_RATE))
-        return batch

fastvideo/pipelines/basic/magi_human/stages/denoising.py

@@ -1,228 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Joint-modality denoising stage for daVinci-MagiHuman base text-to-AV.
-
-Runs the FlowUniPC denoise loop with CFG=2 over video + audio latents
-jointly. Text embeddings are already pad-or-trimmed to `t5_gemma_target_length`
-by `MagiHumanLatentPreparationStage`; the original context lengths are
-stashed on the batch as `magi_original_text_lens` / `magi_original_neg_text_lens`.
-"""
-from __future__ import annotations
-
-import copy
-
-import torch
-from tqdm import tqdm
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.hooks.activation_trace import trace_step
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-    StaticPackedInputs,
-    assemble_packed_inputs,
-    build_static_packed_inputs,
-    unpack_tokens,
-)
-
-
-def _dit_forward(
-    dit,
-    video_latent: torch.Tensor,
-    audio_feat_len: int,
-    txt_feat: torch.Tensor,
-    txt_feat_len: int,
-    static_packed: StaticPackedInputs,
-    coords_style: str,
-    video_in_channels: int,
-    audio_in_channels: int,
-    patch_size: tuple[int, int, int],
-) -> tuple[torch.Tensor, torch.Tensor]:
-    x, coords, mm = assemble_packed_inputs(
-        static=static_packed,
-        txt_feat=txt_feat,
-        txt_feat_len=txt_feat_len,
-        coords_style=coords_style,
-    )
-    video_token_num = static_packed.video_token_num
-    out = dit(x, coords, mm)
-    return unpack_tokens(
-        out,
-        video_token_num=video_token_num,
-        audio_feat_len=audio_feat_len,
-        video_in_channels=video_in_channels,
-        audio_in_channels=audio_in_channels,
-        latent_shape=tuple(video_latent.shape),
-        patch_size=patch_size,
-    )
-
-
-def _overwrite_first_frame(
-    video_latent: torch.Tensor,
-    image_latent: torch.Tensor | None,
-) -> torch.Tensor:
-    if image_latent is not None:
-        video_latent[:, :, :1] = image_latent.to(
-            device=video_latent.device,
-            dtype=video_latent.dtype,
-        )[:, :, :1]
-    return video_latent
-
-
-class MagiHumanDenoisingStage(PipelineStage):
-    """UniPC-flow joint denoising with CFG=2 over (video, audio) latents."""
-
-    def __init__(
-        self,
-        transformer,
-        scheduler,
-        patch_size: tuple[int, int, int] = (1, 2, 2),
-        video_in_channels: int = 192,
-        audio_in_channels: int = 64,
-        video_txt_guidance_scale: float = 5.0,
-        audio_txt_guidance_scale: float = 5.0,
-        cfg_number: int = 2,
-        coords_style: str = "v2",
-        video_guidance_high_t_threshold: int = 500,
-        video_guidance_low_t_value: float = 2.0,
-    ) -> None:
-        super().__init__()
-        self.transformer = transformer
-        self.scheduler = scheduler
-        self.patch_size = patch_size
-        self.video_in_channels = video_in_channels
-        self.audio_in_channels = audio_in_channels
-        self.video_txt_guidance_scale = video_txt_guidance_scale
-        self.audio_txt_guidance_scale = audio_txt_guidance_scale
-        self.cfg_number = cfg_number
-        self.coords_style = coords_style
-        self.video_guidance_high_t_threshold = video_guidance_high_t_threshold
-        self.video_guidance_low_t_value = video_guidance_low_t_value
-
-    def verify_input(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def verify_output(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    @torch.inference_mode()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        device = batch.latents.device
-        shift = fastvideo_args.pipeline_config.flow_shift
-        # Video and audio use independent FlowUniPC state (upstream
-        # inference/pipeline/video_generate.py:404-407 instantiates two
-        # separate schedulers). Sharing one scheduler causes the
-        # `model_outputs` buffer for the video step to pollute the audio
-        # step's diff calculation (different shapes -> broadcast error).
-        video_scheduler = copy.deepcopy(self.scheduler)
-        audio_scheduler = copy.deepcopy(self.scheduler)
-        video_scheduler.set_timesteps(
-            batch.num_inference_steps,
-            device=device,
-            shift=shift,
-        )
-        audio_scheduler.set_timesteps(
-            batch.num_inference_steps,
-            device=device,
-            shift=shift,
-        )
-        timesteps = video_scheduler.timesteps
-
-        video_latent = batch.latents
-        audio_latent = batch.audio_latents
-        image_latent = getattr(batch, "image_latent", None)
-
-        # Expect [1, L, 3584] text embeds plus a list of original lengths.
-        txt_feat = batch.prompt_embeds[0]
-        txt_feat_len = int(batch.magi_original_text_lens[0])
-
-        neg_txt_feat: torch.Tensor | None = None
-        neg_txt_feat_len: int = 0
-        if self.cfg_number == 2:
-            neg_list = batch.negative_prompt_embeds or []
-            if not neg_list:
-                raise ValueError("CFG=2 requires negative prompt embeddings; got None. "
-                                 "Did the prompt encoding stage run?")
-            else:
-                neg_txt_feat = neg_list[0]
-                neg_txt_feat_len = int(batch.magi_original_neg_text_lens[0])
-
-        audio_feat_len = int(audio_latent.shape[1])
-
-        disable_tqdm = not getattr(fastvideo_args, "log_level_progress", True)
-        for idx, t in enumerate(tqdm(timesteps, disable=disable_tqdm)):
-            video_latent = _overwrite_first_frame(video_latent, image_latent)
-            # Precompute packed video+audio tokens after any TI2V first-frame
-            # overwrite. Text varies per cond/uncond call and is attached in
-            # _dit_forward via assemble_packed_inputs.
-            static_packed = build_static_packed_inputs(
-                video_latent=video_latent,
-                audio_latent=audio_latent,
-                audio_feat_len=audio_feat_len,
-                patch_size=self.patch_size,
-                coords_style=self.coords_style,
-                layout=getattr(batch, "magi_static_packed_layout", None),
-            )
-            with trace_step(idx), set_forward_context(
-                    current_timestep=int(t.item()) if torch.is_tensor(t) else int(t),
-                    attn_metadata=None,
-            ):
-                v_cond_video, v_cond_audio = _dit_forward(
-                    self.transformer,
-                    video_latent=video_latent,
-                    audio_feat_len=audio_feat_len,
-                    txt_feat=txt_feat,
-                    txt_feat_len=txt_feat_len,
-                    static_packed=static_packed,
-                    coords_style=self.coords_style,
-                    video_in_channels=self.video_in_channels,
-                    audio_in_channels=self.audio_in_channels,
-                    patch_size=self.patch_size,
-                )
-
-                if self.cfg_number == 2:
-                    v_uncond_video, v_uncond_audio = _dit_forward(
-                        self.transformer,
-                        video_latent=video_latent,
-                        audio_feat_len=audio_feat_len,
-                        txt_feat=neg_txt_feat,
-                        txt_feat_len=neg_txt_feat_len,
-                        static_packed=static_packed,
-                        coords_style=self.coords_style,
-                        video_in_channels=self.video_in_channels,
-                        audio_in_channels=self.audio_in_channels,
-                        patch_size=self.patch_size,
-                    )
-                else:
-                    v_uncond_video = None
-                    v_uncond_audio = None
-
-            if self.cfg_number == 2:
-                video_guidance = (self.video_txt_guidance_scale
-                                  if t > self.video_guidance_high_t_threshold else self.video_guidance_low_t_value)
-                assert v_uncond_video is not None and v_uncond_audio is not None
-                v_video = v_uncond_video + video_guidance * (v_cond_video - v_uncond_video)
-                v_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio)
-            else:
-                v_video = v_cond_video
-                v_audio = v_cond_audio
-
-            # Independent scheduler state per modality (see comment above).
-            video_latent = video_scheduler.step(
-                v_video,
-                t,
-                video_latent,
-                return_dict=False,
-            )[0]
-            audio_latent = audio_scheduler.step(
-                v_audio,
-                t,
-                audio_latent,
-                return_dict=False,
-            )[0]
-
-        video_latent = _overwrite_first_frame(video_latent, image_latent)
-        batch.latents = video_latent
-        batch.audio_latents = audio_latent
-        return batch

fastvideo/pipelines/basic/magi_human/stages/latent_preparation.py

@@ -1,590 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Latent preparation stage for daVinci-MagiHuman base text-to-AV.
-
-Produces:
-  - random video latent of shape `[1, z_dim, latent_T, latent_H, latent_W]`,
-  - random audio latent of shape `[1, num_frames, 64]` (the DiT jointly
-    denoises both modalities),
-  - padded T5-Gemma text embedding (target length 640) plus the original
-    (pre-pad) context length, which the UniPC + CFG loop needs so the
-    unconditional path sees the same padded length.
-
-Also stakes out the per-token coords / modality map that the DiT consumes
-(replicates the reference `MagiDataProxy.process_input`).
-"""
-from __future__ import annotations
-
-from typing import Literal
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-# Matches inference/common/sequence_schema.py in the reference.
-MODALITY_VIDEO = 0
-MODALITY_AUDIO = 1
-MODALITY_TEXT = 2
-
-# Audio temporal compression ratio: 1 audio frame → 1/4 latent frame.
-# Mirrors data_proxy.py:206 `(audio_feat_len - 1) // 4 + 1` where 4 is
-# the audio VAE's temporal stride (same as vae_stride[0] for video).
-_AUDIO_TEMPORAL_COMPRESSION = 4
-
-# v1 text-coord reference shape: (T=2, H=1, W=1).
-# Mirrors data_proxy.py:202 `ref_feat_shape=(2, 1, 1)` for coords_style=="v1".
-_V1_TEXT_REF_SHAPE: tuple[int, int, int] = (2, 1, 1)
-
-
-def _build_coords(
-    shape: tuple[int, int, int],
-    ref_feat_shape: tuple[int, int, int],
-    offset_thw: tuple[int, int, int] = (0, 0, 0),
-    device: torch.device | None = None,
-    dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    if device is None:
-        device = torch.device("cpu")
-    ori_t, ori_h, ori_w = shape
-    ref_t, ref_h, ref_w = ref_feat_shape
-    offset_t, offset_h, offset_w = offset_thw
-    time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t
-    h_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h
-    w_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w
-    tg, hg, wg = torch.meshgrid(time_rng, h_rng, w_rng, indexing="ij")
-    coords = torch.stack([tg, hg, wg], dim=-1).reshape(-1, 3)
-    meta = torch.tensor(
-        [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w],
-        device=device,
-        dtype=dtype,
-    ).expand(coords.size(0), -1)
-    return torch.cat([coords, meta], dim=-1)
-
-
-def _pad_or_trim_dim1(t: torch.Tensor, target: int) -> tuple[torch.Tensor, int]:
-    """Pad-or-trim along dim 1. Returns (new_tensor, original_length)."""
-    current = t.size(1)
-    if current < target:
-        pad = [0, 0, 0, target - current]
-        return F.pad(t, pad, "constant", 0.0), current
-    return t[:, :target], target
-
-
-def _img2tokens(x_t: torch.Tensor, t_patch: int, patch: int) -> torch.Tensor:
-    """Pack a video latent [B, C, T, H, W] -> [B, L, C * t_patch * patch^2].
-
-    Per-token feature ordering is channel-major ``(C pT pH pW)``: the DiT's
-    ``video_embedder`` weight was trained on the layout produced by
-    upstream's grouped-conv ``UnfoldNd`` packer (channel slowest, patch
-    elements fastest). Spatial-major ``(pT pH pW C)`` silently permutes the
-    in-features and produces noise output. Asymmetric with
-    ``unpack_tokens`` which uses ``(pT pH pW C)`` to match
-    ``final_linear_video``'s trained output layout.
-    """
-    B, C, T, H, W = x_t.shape
-    assert T % t_patch == 0 and H % patch == 0 and W % patch == 0, (
-        f"Latent dims {T,H,W} must divide ({t_patch}, {patch}, {patch})")
-    return rearrange(
-        x_t,
-        "B C (T pT) (H pH) (W pW) -> B (T H W) (C pT pH pW)",
-        pT=t_patch,
-        pH=patch,
-        pW=patch,
-    ).contiguous()
-
-
-class MagiHumanLatentPreparationStage(PipelineStage):
-    """Prepare latents, coords, modality maps, and padded text embed."""
-
-    def __init__(
-        self,
-        vae_stride: tuple[int, int, int] = (4, 16, 16),
-        z_dim: int = 48,
-        patch_size: tuple[int, int, int] = (1, 2, 2),
-        fps: int = 25,
-        t5_gemma_target_length: int = 640,
-        coords_style: Literal["v1", "v2"] = "v2",
-        text_offset: int = 0,
-        audio_in_channels: int = 64,
-    ) -> None:
-        super().__init__()
-        self.vae_stride = vae_stride
-        self.z_dim = z_dim
-        self.patch_size = patch_size
-        self.fps = fps
-        self.t5_gemma_target_length = t5_gemma_target_length
-        self.coords_style = coords_style
-        self.text_offset = text_offset
-        self.audio_in_channels = audio_in_channels
-
-    def verify_input(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def verify_output(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        fps = self.fps
-        # Prefer the caller-provided `batch.num_frames` (the standard
-        # SamplingParam knob — production preset and SSIM tests both set
-        # it). Fall back to `batch.num_seconds * fps + 1` when num_frames
-        # is unset or the image-default sentinel (1). This matches
-        # upstream MagiDataProxy.process_input which derives `num_frames
-        # = seconds * fps + 1` and rejects values that don't satisfy
-        # `(num_frames - 1) % vae_temporal_stride == 0`.
-        requested_num_frames = int(getattr(batch, "num_frames", None) or 0)
-        if requested_num_frames > 1:
-            num_frames = requested_num_frames
-        else:
-            seconds = int(getattr(batch, "num_seconds", None) or 4)
-            num_frames = seconds * fps + 1
-        latent_T = (num_frames - 1) // 4 + 1
-
-        # Match upstream pipeline.py:61-64 + video_generate.py:254-261:
-        # the requested 272p height snaps to 256, while width stays 480.
-        br_h = int(batch.height) if batch.height else 256
-        br_w = int(batch.width) if batch.width else 480
-        pT, pH, pW = self.patch_size
-        vt, vh, vw = self.vae_stride
-        # Snap to patch granularity (matches reference).
-        latent_H = (br_h // vh // pH) * pH
-        latent_W = (br_w // vw // pW) * pW
-        actual_H = latent_H * vh
-        actual_W = latent_W * vw
-        batch.height = actual_H
-        batch.width = actual_W
-
-        generator = torch.Generator(device=device)
-        if batch.seed is not None:
-            generator.manual_seed(int(batch.seed))
-
-        # Video latent: [1, z_dim, latent_T, latent_H, latent_W]
-        video_latent = torch.randn(
-            (1, self.z_dim, latent_T, latent_H, latent_W),
-            generator=generator,
-            device=device,
-            dtype=torch.float32,
-        )
-        image_latent = getattr(batch, "image_latent", None)
-        if image_latent is not None:
-            video_latent[:, :, :1] = image_latent.to(
-                device=video_latent.device,
-                dtype=video_latent.dtype,
-            )[:, :, :1]
-        # Audio latent: [1, num_frames, audio_in_channels]
-        audio_latent = torch.randn(
-            (1, num_frames, self.audio_in_channels),
-            generator=generator,
-            device=device,
-            dtype=torch.float32,
-        )
-
-        # Prompt embeds: the upstream TextEncodingStage already ran. It
-        # produced a list of [1, L, D] tensors per prompt. Pad/trim each
-        # to the target length and store the original length so the DiT
-        # stage can build the correct modality-map slices.
-        padded_prompt_embeds: list[torch.Tensor] = []
-        padded_prompt_lens: list[int] = []
-        for embed in batch.prompt_embeds:
-            # embed: [1, L, 3584]
-            padded, original = _pad_or_trim_dim1(
-                embed.to(torch.float32),
-                target=self.t5_gemma_target_length,
-            )
-            padded_prompt_embeds.append(padded)
-            padded_prompt_lens.append(original)
-        batch.prompt_embeds = padded_prompt_embeds
-        # Stash the original text length list on the batch for the denoise
-        # stage — FastVideo's ForwardBatch doesn't have a first-class field
-        # for this so we attach it.
-        batch.magi_original_text_lens = padded_prompt_lens
-
-        # Matching negative prompts.
-        if batch.negative_prompt_embeds is not None and batch.negative_prompt_embeds:
-            padded_neg: list[torch.Tensor] = []
-            padded_neg_lens: list[int] = []
-            for embed in batch.negative_prompt_embeds:
-                padded, original = _pad_or_trim_dim1(
-                    embed.to(torch.float32),
-                    target=self.t5_gemma_target_length,
-                )
-                padded_neg.append(padded)
-                padded_neg_lens.append(original)
-            batch.negative_prompt_embeds = padded_neg
-            batch.magi_original_neg_text_lens = padded_neg_lens
-
-        batch.latents = video_latent
-        batch.audio_latents = audio_latent
-        batch.num_frames = num_frames
-        batch.magi_latent_T = latent_T
-        batch.magi_latent_H = latent_H
-        batch.magi_latent_W = latent_W
-        # Precompute the step-invariant packed layout (coords / modality
-        # maps / channel-padding width) once; the denoise loop reuses it
-        # every step instead of rebuilding meshgrids on each call.
-        batch.magi_static_packed_layout = precompute_static_packed_layout(
-            latent_shape=tuple(video_latent.shape),  # type: ignore[arg-type]
-            audio_feat_len=int(audio_latent.shape[1]),
-            z_dim=self.z_dim,
-            audio_in_channels=self.audio_in_channels,
-            patch_size=self.patch_size,
-            coords_style=self.coords_style,
-            device=video_latent.device,
-        )
-        return batch
-
-
-class StaticPackedInputs:
-    """Step-invariant packed inputs: video+audio tokens, coords, modality map.
-
-    Computed once before the denoise loop; reused for every cond/uncond call.
-    Text tokens are NOT included here because cond/uncond have different lengths.
-    """
-
-    __slots__ = (
-        "video_tokens",
-        "audio_tokens",
-        "video_coords",
-        "audio_coords",
-        "video_mm",
-        "audio_mm",
-        "video_token_num",
-        "audio_feat_len",
-        "max_ch",
-    )
-
-    def __init__(
-        self,
-        video_tokens: torch.Tensor,
-        audio_tokens: torch.Tensor,
-        video_coords: torch.Tensor,
-        audio_coords: torch.Tensor,
-        video_mm: torch.Tensor,
-        audio_mm: torch.Tensor,
-        max_ch: int,
-    ) -> None:
-        self.video_tokens = video_tokens
-        self.audio_tokens = audio_tokens
-        self.video_coords = video_coords
-        self.audio_coords = audio_coords
-        self.video_mm = video_mm
-        self.audio_mm = audio_mm
-        self.video_token_num = video_tokens.size(0)
-        self.audio_feat_len = audio_tokens.size(0)
-        self.max_ch = max_ch
-
-
-class StaticPackedLayout:
-    """Step- and value-invariant portion of the static packed inputs.
-
-    Coords, modality maps, and the channel-padding width depend only on the
-    latent shape, audio length, channel widths, and patch sizes — all fixed
-    for a single generation. Precompute once before the denoise loop and
-    reuse on every step. Only the per-step token tensors must be rebuilt.
-    """
-
-    __slots__ = (
-        "video_coords",
-        "audio_coords",
-        "video_mm",
-        "audio_mm",
-        "max_ch",
-        "video_token_num",
-        "audio_feat_len",
-    )
-
-    def __init__(
-        self,
-        video_coords: torch.Tensor,
-        audio_coords: torch.Tensor,
-        video_mm: torch.Tensor,
-        audio_mm: torch.Tensor,
-        max_ch: int,
-        video_token_num: int,
-        audio_feat_len: int,
-    ) -> None:
-        self.video_coords = video_coords
-        self.audio_coords = audio_coords
-        self.video_mm = video_mm
-        self.audio_mm = audio_mm
-        self.max_ch = max_ch
-        self.video_token_num = video_token_num
-        self.audio_feat_len = audio_feat_len
-
-
-def precompute_static_packed_layout(
-    latent_shape: tuple[int, int, int, int, int],
-    audio_feat_len: int,
-    z_dim: int,
-    audio_in_channels: int,
-    patch_size: tuple[int, int, int],
-    coords_style: Literal["v1", "v2"] = "v2",
-    device: torch.device | None = None,
-    dtype: torch.dtype = torch.float32,
-) -> StaticPackedLayout:
-    """Precompute the invariant fields used by ``build_static_packed_inputs``.
-
-    Arguments are derived from configs and latent shape — none depend on
-    the current denoising-step values. Call this once in the latent
-    preparation stage (or any pre-loop site) and pass the result via the
-    ``layout=`` arg of ``build_static_packed_inputs`` to skip the
-    meshgrid/full() work on every step.
-    """
-    pT, pH, pW = patch_size
-    _, _, T, H, W = latent_shape
-    if device is None:
-        device = torch.device("cpu")
-
-    video_token_num = (T // pT) * (H // pH) * (W // pW)
-    # `_img2tokens` packs to channel `z_dim * pT * pH * pW`; audio tokens
-    # are `audio_in_channels` wide — both are config constants.
-    max_ch = max(z_dim * pT * pH * pW, audio_in_channels)
-
-    video_ref_shape = (T // pT, H // pH, W // pW)
-    video_coords = _build_coords(
-        shape=video_ref_shape,
-        ref_feat_shape=video_ref_shape,
-        device=device,
-        dtype=dtype,
-    )
-
-    if coords_style == "v2":
-        audio_ref_t = (audio_feat_len - 1) // _AUDIO_TEMPORAL_COMPRESSION + 1
-        audio_coords = _build_coords(
-            shape=(audio_feat_len, 1, 1),
-            ref_feat_shape=(audio_ref_t // pT, 1, 1),
-            device=device,
-            dtype=dtype,
-        )
-    else:
-        audio_coords = _build_coords(
-            shape=(audio_feat_len, 1, 1),
-            ref_feat_shape=(T // pT, 1, 1),
-            device=device,
-            dtype=dtype,
-        )
-
-    video_mm = torch.full((video_token_num, ), MODALITY_VIDEO, dtype=torch.int64, device=device)
-    audio_mm = torch.full((audio_feat_len, ), MODALITY_AUDIO, dtype=torch.int64, device=device)
-
-    return StaticPackedLayout(
-        video_coords=video_coords,
-        audio_coords=audio_coords,
-        video_mm=video_mm,
-        audio_mm=audio_mm,
-        max_ch=max_ch,
-        video_token_num=video_token_num,
-        audio_feat_len=audio_feat_len,
-    )
-
-
-def build_static_packed_inputs(
-    video_latent: torch.Tensor,
-    audio_latent: torch.Tensor,
-    audio_feat_len: int,
-    patch_size: tuple[int, int, int],
-    coords_style: Literal["v1", "v2"] = "v2",
-    layout: StaticPackedLayout | None = None,
-) -> StaticPackedInputs:
-    """Build the step-invariant portion of the packed token stream.
-
-    Returns video+audio tokens (padded to a common channel width), their
-    coords, and their modality slices. Text is excluded because cond/uncond
-    differ in length; call assemble_packed_inputs to attach text per call.
-
-    Mirrors SingleData.token_sequence / coords_mapping / modality_mapping in
-    inference/pipeline/data_proxy.py, minus the text portion.
-
-    When ``layout`` is provided, coords / modality maps / max_ch are taken
-    from the precomputed values and only the per-step token tensors are
-    rebuilt; this is the hot-path call from the denoising loop. When
-    ``layout`` is None the function recomputes everything from scratch
-    (e.g. for one-shot tests via ``build_packed_inputs``).
-    """
-    pT, pH, pW = patch_size
-    assert video_latent.size(0) == 1, "batch size 1 required for MagiHuman base"
-
-    video_tokens = _img2tokens(video_latent, t_patch=pT, patch=pH)[0]
-    audio_tokens = audio_latent[0, :audio_feat_len].contiguous()
-
-    if layout is not None:
-        max_ch = layout.max_ch
-        video_tokens = F.pad(video_tokens, (0, max_ch - video_tokens.size(-1)))
-        audio_tokens = F.pad(audio_tokens, (0, max_ch - audio_tokens.size(-1)))
-        return StaticPackedInputs(
-            video_tokens=video_tokens,
-            audio_tokens=audio_tokens,
-            video_coords=layout.video_coords,
-            audio_coords=layout.audio_coords,
-            video_mm=layout.video_mm,
-            audio_mm=layout.audio_mm,
-            max_ch=max_ch,
-        )
-
-    # Slow path: rebuild every invariant from scratch. Kept for the
-    # ``build_packed_inputs`` one-shot wrapper used by tests/parity helpers.
-    _, z_dim, T, H, W = video_latent.shape
-
-    max_ch = max(video_tokens.size(-1), audio_tokens.size(-1))
-    video_tokens = F.pad(video_tokens, (0, max_ch - video_tokens.size(-1)))
-    audio_tokens = F.pad(audio_tokens, (0, max_ch - audio_tokens.size(-1)))
-
-    device = video_tokens.device
-    dtype = video_tokens.dtype
-    video_token_num = video_tokens.size(0)
-
-    video_mm = torch.full((video_token_num, ), MODALITY_VIDEO, dtype=torch.int64, device=device)
-    audio_mm = torch.full((audio_feat_len, ), MODALITY_AUDIO, dtype=torch.int64, device=device)
-
-    video_ref_shape = (T // pT, H // pH, W // pW)
-    video_coords = _build_coords(
-        shape=video_ref_shape,
-        ref_feat_shape=video_ref_shape,
-        device=device,
-        dtype=dtype,
-    )
-
-    if coords_style == "v2":
-        audio_ref_t = (audio_feat_len - 1) // _AUDIO_TEMPORAL_COMPRESSION + 1
-        audio_coords = _build_coords(
-            shape=(audio_feat_len, 1, 1),
-            ref_feat_shape=(audio_ref_t // pT, 1, 1),
-            device=device,
-            dtype=dtype,
-        )
-    else:
-        audio_coords = _build_coords(
-            shape=(audio_feat_len, 1, 1),
-            ref_feat_shape=(T // pT, 1, 1),
-            device=device,
-            dtype=dtype,
-        )
-
-    return StaticPackedInputs(
-        video_tokens=video_tokens,
-        audio_tokens=audio_tokens,
-        video_coords=video_coords,
-        audio_coords=audio_coords,
-        video_mm=video_mm,
-        audio_mm=audio_mm,
-        max_ch=max_ch,
-    )
-
-
-def assemble_packed_inputs(
-    static: StaticPackedInputs,
-    txt_feat: torch.Tensor,
-    txt_feat_len: int,
-    coords_style: Literal["v1", "v2"] = "v2",
-    text_offset: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Attach per-call text tokens to the precomputed static packed inputs.
-
-    Returns (token_seq, coords, modality_map) ready for the DiT.
-    """
-    text_tokens = txt_feat[0, :txt_feat_len].contiguous()
-    max_ch = max(static.max_ch, text_tokens.size(-1))
-
-    video_tokens = F.pad(static.video_tokens, (0, max_ch - static.video_tokens.size(-1)))
-    audio_tokens = F.pad(static.audio_tokens, (0, max_ch - static.audio_tokens.size(-1)))
-    text_tokens = F.pad(text_tokens, (0, max_ch - text_tokens.size(-1)))
-    token_seq = torch.cat([video_tokens, audio_tokens, text_tokens], dim=0)
-
-    device = token_seq.device
-    dtype = token_seq.dtype
-    text_mm = torch.full((txt_feat_len, ), MODALITY_TEXT, dtype=torch.int64, device=device)
-    mm = torch.cat([static.video_mm, static.audio_mm, text_mm], dim=0)
-
-    if coords_style == "v2":
-        text_coords = _build_coords(
-            shape=(txt_feat_len, 1, 1),
-            ref_feat_shape=(1, 1, 1),
-            offset_thw=(-txt_feat_len, 0, 0),
-            device=device,
-            dtype=dtype,
-        )
-    else:
-        text_coords = _build_coords(
-            shape=(txt_feat_len, 1, 1),
-            ref_feat_shape=_V1_TEXT_REF_SHAPE,
-            offset_thw=(text_offset, 0, 0),
-            device=device,
-            dtype=dtype,
-        )
-
-    coords = torch.cat([static.video_coords, static.audio_coords, text_coords], dim=0)
-    return token_seq, coords, mm
-
-
-def build_packed_inputs(
-    video_latent: torch.Tensor,
-    audio_latent: torch.Tensor,
-    audio_feat_len: int,
-    txt_feat: torch.Tensor,
-    txt_feat_len: int,
-    patch_size: tuple[int, int, int],
-    coords_style: Literal["v1", "v2"] = "v2",
-    text_offset: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Build the full packed token stream in one call (backwards-compat wrapper).
-
-    Equivalent to assemble_packed_inputs(build_static_packed_inputs(...), ...).
-    Prefer calling the two helpers separately when the static portion can be
-    reused across multiple calls (e.g. cond/uncond in the denoise loop).
-    """
-    static = build_static_packed_inputs(
-        video_latent=video_latent,
-        audio_latent=audio_latent,
-        audio_feat_len=audio_feat_len,
-        patch_size=patch_size,
-        coords_style=coords_style,
-    )
-    return assemble_packed_inputs(
-        static=static,
-        txt_feat=txt_feat,
-        txt_feat_len=txt_feat_len,
-        coords_style=coords_style,
-        text_offset=text_offset,
-    )
-
-
-def unpack_tokens(
-    output: torch.Tensor,  # [L, max(V_ch, A_ch)]
-    video_token_num: int,
-    audio_feat_len: int,
-    video_in_channels: int,
-    audio_in_channels: int,
-    latent_shape: tuple[int, int, int, int, int],  # [1, z_dim, T, H, W]
-    patch_size: tuple[int, int, int],
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Inverse of `build_packed_inputs` for the DiT output.
-
-    Splits the flat output back into a video latent (un-patched into
-    B C T H W) and an audio latent (B, L, 64).
-    """
-    pT, pH, pW = patch_size
-    _, z_dim, T, H, W = latent_shape
-    tH, tW = H // pH, W // pW
-
-    video_flat = output[:video_token_num, :video_in_channels]
-    video_latent = rearrange(
-        video_flat,
-        "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)",
-        H=tH,
-        W=tW,
-        pT=pT,
-        pH=pH,
-        pW=pW,
-    ).contiguous().unsqueeze(0)
-
-    audio_latent = output[
-        video_token_num:video_token_num + audio_feat_len,
-        :audio_in_channels,
-    ].unsqueeze(0)
-
-    return video_latent, audio_latent

fastvideo/pipelines/basic/magi_human/stages/reference_image.py

@@ -1,101 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Reference-image encoding for MagiHuman TI2V.
-
-The upstream daVinci-MagiHuman TI2V path encodes the user image through the
-Wan VAE and overwrites the first denoising latent frame with that clean latent
-at every step. This stage mirrors `MagiEvaluator.encode_image` and stashes the
-normalized latent on `batch.image_latent` for the latent-prep and denoise stages.
-"""
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-from diffusers.utils import load_image
-from diffusers.video_processor import VideoProcessor
-from PIL import Image
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-
-def _resizecrop(image: Image.Image, height: int, width: int) -> Image.Image:
-    """Mirror upstream `resizecrop`: center-crop to target aspect ratio."""
-    current_width, current_height = image.size
-    if current_width == width and current_height == height:
-        return image
-    if current_height / current_width > height / width:
-        new_width = int(current_width)
-        new_height = int(new_width * height / width)
-    else:
-        new_height = int(current_height)
-        new_width = int(new_height * width / height)
-    left = (current_width - new_width) / 2
-    top = (current_height - new_height) / 2
-    right = (current_width + new_width) / 2
-    bottom = (current_height + new_height) / 2
-    return image.crop((left, top, right, bottom))
-
-
-class MagiHumanReferenceImageStage(PipelineStage):
-    """Encode a TI2V reference image into the first-frame video latent."""
-
-    def __init__(self, vae: Any, vae_scale_factor: int = 16) -> None:
-        super().__init__()
-        self.vae = vae
-        self.video_processor = VideoProcessor(vae_scale_factor=vae_scale_factor)
-
-    def verify_input(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def verify_output(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    @torch.inference_mode()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        image = getattr(batch, "image", None) or batch.pil_image
-        if image is None and batch.image_path is not None:
-            image = load_image(batch.image_path)
-        if image is None:
-            raise ValueError("MagiHuman TI2V requires `image_path` or `pil_image`.")
-        if not isinstance(image, Image.Image):
-            raise TypeError(f"MagiHuman TI2V expects a PIL image or image path, got {type(image)}")
-        if batch.height is None or batch.width is None:
-            raise ValueError("MagiHuman TI2V requires concrete height and width before image encoding.")
-
-        height = int(batch.height)
-        width = int(batch.width)
-        device = get_local_torch_device()
-
-        image = _resizecrop(image.convert("RGB"), height, width)
-        image_tensor = self.video_processor.preprocess(
-            image,
-            height=height,
-            width=width,
-        ).to(device=device, dtype=torch.float32)
-        image_tensor = image_tensor.unsqueeze(2)
-
-        self.vae = self.vae.to(device)
-        encoded = self.vae.encode(image_tensor)
-        image_latent = encoded.mean if hasattr(encoded, "mean") else encoded
-
-        # FastVideo's Wan VAE returns unnormalized posterior means; upstream
-        # `WanVAE.encode` applies `(mu - mean) / std` before returning.
-        shift_factor = getattr(self.vae, "shift_factor", None)
-        if shift_factor is not None:
-            if isinstance(shift_factor, torch.Tensor):
-                image_latent = image_latent - shift_factor.to(image_latent.device, image_latent.dtype)
-            else:
-                image_latent = image_latent - shift_factor
-        scaling_factor = getattr(self.vae, "scaling_factor", None)
-        if scaling_factor is not None:
-            if isinstance(scaling_factor, torch.Tensor):
-                image_latent = image_latent * scaling_factor.to(image_latent.device, image_latent.dtype)
-            else:
-                image_latent = image_latent * scaling_factor
-
-        batch.image_latent = image_latent.to(torch.float32)
-        return batch

fastvideo/pipelines/basic/magi_human/stages/sr_denoising.py

@@ -1,156 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""SR video-only denoising stage for daVinci-MagiHuman SR-540p."""
-from __future__ import annotations
-
-import copy
-
-import torch
-from tqdm import tqdm
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.hooks.activation_trace import trace_step
-from fastvideo.pipelines.basic.magi_human.stages.denoising import (
-    _dit_forward,
-    _overwrite_first_frame,
-)
-from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-    build_static_packed_inputs, )
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-
-class MagiHumanSRDenoisingStage(PipelineStage):
-    """Denoise only the SR video latent; audio passes through unchanged."""
-
-    def __init__(
-        self,
-        transformer,
-        scheduler,
-        patch_size: tuple[int, int, int] = (1, 2, 2),
-        video_in_channels: int = 192,
-        audio_in_channels: int = 64,
-        sr_num_inference_steps: int = 5,
-        sr_video_txt_guidance_scale: float = 3.5,
-        use_cfg_trick: bool = True,
-        cfg_trick_start_frame: int = 13,
-        cfg_trick_value: float = 2.0,
-        cfg_number: int = 2,
-        coords_style: str = "v1",
-    ) -> None:
-        super().__init__()
-        self.transformer = transformer
-        self.scheduler = scheduler
-        self.patch_size = patch_size
-        self.video_in_channels = video_in_channels
-        self.audio_in_channels = audio_in_channels
-        self.sr_num_inference_steps = sr_num_inference_steps
-        self.sr_video_txt_guidance_scale = sr_video_txt_guidance_scale
-        self.use_cfg_trick = use_cfg_trick
-        self.cfg_trick_start_frame = cfg_trick_start_frame
-        self.cfg_trick_value = cfg_trick_value
-        self.cfg_number = cfg_number
-        self.coords_style = coords_style
-
-    def verify_input(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def verify_output(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    @torch.inference_mode()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        device = batch.latents.device
-        shift = fastvideo_args.pipeline_config.flow_shift
-        video_scheduler = copy.deepcopy(self.scheduler)
-        video_scheduler.set_timesteps(
-            self.sr_num_inference_steps,
-            device=device,
-            shift=shift,
-        )
-
-        video_latent = batch.latents
-        audio_latent = batch.audio_latents
-        audio_feat_len = int(audio_latent.shape[1])
-        image_latent = getattr(batch, "image_latent", None)
-
-        txt_feat = batch.prompt_embeds[0]
-        txt_feat_len = int(batch.magi_original_text_lens[0])
-
-        neg_txt_feat: torch.Tensor | None = None
-        neg_txt_feat_len = 0
-        if self.cfg_number == 2:
-            neg_list = batch.negative_prompt_embeds or []
-            if not neg_list:
-                raise ValueError("SR CFG=2 requires negative prompt embeddings.")
-            neg_txt_feat = neg_list[0]
-            neg_txt_feat_len = int(batch.magi_original_neg_text_lens[0])
-
-        latent_length = video_latent.shape[2]
-        guidance = torch.tensor(
-            self.sr_video_txt_guidance_scale,
-            device=device,
-            dtype=video_latent.dtype,
-        ).expand(1, 1, latent_length, 1, 1).clone()
-        if self.use_cfg_trick:
-            guidance[:, :, :self.cfg_trick_start_frame] = min(
-                self.cfg_trick_value,
-                self.sr_video_txt_guidance_scale,
-            )
-
-        disable_tqdm = not getattr(fastvideo_args, "log_level_progress", True)
-        for idx, t in enumerate(tqdm(video_scheduler.timesteps, disable=disable_tqdm)):
-            video_latent = _overwrite_first_frame(video_latent, image_latent)
-            static_packed = build_static_packed_inputs(
-                video_latent=video_latent,
-                audio_latent=audio_latent,
-                audio_feat_len=audio_feat_len,
-                patch_size=self.patch_size,
-                coords_style=self.coords_style,
-                layout=getattr(batch, "magi_static_packed_layout", None),
-            )
-            with trace_step(idx), set_forward_context(
-                    current_timestep=int(t.item()) if torch.is_tensor(t) else int(t),
-                    attn_metadata=None,
-            ):
-                v_cond_video, _ = _dit_forward(
-                    self.transformer,
-                    video_latent=video_latent,
-                    audio_feat_len=audio_feat_len,
-                    txt_feat=txt_feat,
-                    txt_feat_len=txt_feat_len,
-                    static_packed=static_packed,
-                    coords_style=self.coords_style,
-                    video_in_channels=self.video_in_channels,
-                    audio_in_channels=self.audio_in_channels,
-                    patch_size=self.patch_size,
-                )
-                if self.cfg_number == 2:
-                    assert neg_txt_feat is not None
-                    v_uncond_video, _ = _dit_forward(
-                        self.transformer,
-                        video_latent=video_latent,
-                        audio_feat_len=audio_feat_len,
-                        txt_feat=neg_txt_feat,
-                        txt_feat_len=neg_txt_feat_len,
-                        static_packed=static_packed,
-                        coords_style=self.coords_style,
-                        video_in_channels=self.video_in_channels,
-                        audio_in_channels=self.audio_in_channels,
-                        patch_size=self.patch_size,
-                    )
-                    v_video = v_uncond_video + guidance * (v_cond_video - v_uncond_video)
-                else:
-                    v_video = v_cond_video
-
-            video_latent = video_scheduler.step(
-                v_video,
-                t,
-                video_latent,
-                return_dict=False,
-            )[0]
-
-        batch.latents = _overwrite_first_frame(video_latent, image_latent)
-        batch.audio_latents = audio_latent
-        return batch

fastvideo/pipelines/basic/magi_human/stages/sr_latent_preparation.py

@@ -1,219 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Super-resolution latent preparation for daVinci-MagiHuman SR-540p."""
-from __future__ import annotations
-
-from functools import partial
-from typing import Any
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from diffusers.utils import load_image
-from diffusers.video_processor import VideoProcessor
-from PIL import Image
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.basic.magi_human.stages.reference_image import (
-    _resizecrop, )
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-
-class ZeroSNRDDPMDiscretization:
-    """Upstream ZeroSNR schedule used to corrupt interpolated SR latents."""
-
-    def __init__(
-        self,
-        linear_start: float = 0.00085,
-        linear_end: float = 0.0120,
-        num_timesteps: int = 1000,
-        shift_scale: float = 1.0,
-        keep_start: bool = False,
-        post_shift: bool = False,
-    ) -> None:
-        if keep_start and not post_shift:
-            linear_start = linear_start / (shift_scale + (1 - shift_scale) * linear_start)
-        self.num_timesteps = num_timesteps
-        betas = torch.linspace(
-            linear_start**0.5,
-            linear_end**0.5,
-            num_timesteps,
-            dtype=torch.float64,
-        )**2
-        alphas = 1.0 - betas.numpy()
-        self.alphas_cumprod = np.cumprod(alphas, axis=0)
-        self.post_shift = post_shift
-        self.shift_scale = shift_scale
-
-        if not post_shift:
-            self.alphas_cumprod = self.alphas_cumprod / (shift_scale + (1 - shift_scale) * self.alphas_cumprod)
-
-    def __call__(
-        self,
-        n: int,
-        do_append_zero: bool = True,
-        device: str | torch.device = "cpu",
-        flip: bool = False,
-    ) -> torch.Tensor:
-        sigmas = self.get_sigmas(n, device=device)
-        if do_append_zero:
-            sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
-        return torch.flip(sigmas, (0, )) if flip else sigmas
-
-    def get_sigmas(
-        self,
-        n: int,
-        device: str | torch.device = "cpu",
-    ) -> torch.Tensor:
-        if n < self.num_timesteps:
-            timesteps = np.linspace(
-                self.num_timesteps - 1,
-                0,
-                n,
-                endpoint=False,
-            ).astype(int)[::-1]
-            alphas_cumprod = self.alphas_cumprod[timesteps]
-        elif n == self.num_timesteps:
-            alphas_cumprod = self.alphas_cumprod
-        else:
-            raise ValueError(f"n must be <= {self.num_timesteps}, got {n}")
-
-        to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
-        alphas_cumprod_sqrt = to_torch(alphas_cumprod).sqrt()
-        alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone()
-        alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone()
-
-        alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T
-        alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 - alphas_cumprod_sqrt_T)
-
-        if self.post_shift:
-            alphas_cumprod_sqrt = (alphas_cumprod_sqrt**2 / (self.shift_scale +
-                                                             (1 - self.shift_scale) * alphas_cumprod_sqrt**2))**0.5
-        return torch.flip(alphas_cumprod_sqrt, (0, ))
-
-
-class MagiHumanSRLatentPreparationStage(PipelineStage):
-    """Upsample base latents, add SR noise, and refresh SR conditioning."""
-
-    def __init__(
-        self,
-        vae: Any,
-        vae_stride: tuple[int, int, int] = (4, 16, 16),
-        patch_size: tuple[int, int, int] = (1, 2, 2),
-        noise_value: int = 220,
-        sr_audio_noise_scale: float = 0.7,
-        sr_height: int = 512,
-        sr_width: int = 896,
-        vae_scale_factor: int = 16,
-    ) -> None:
-        super().__init__()
-        self.vae = vae
-        self.vae_stride = vae_stride
-        self.patch_size = patch_size
-        self.noise_value = noise_value
-        self.sr_audio_noise_scale = sr_audio_noise_scale
-        self.sr_height = sr_height
-        self.sr_width = sr_width
-        self.sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True)
-        self.video_processor = VideoProcessor(vae_scale_factor=vae_scale_factor)
-
-    def verify_input(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    def verify_output(self, batch, fastvideo_args):
-        return VerificationResult()
-
-    @torch.inference_mode()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        device = batch.latents.device
-        _, _, latent_t, _, _ = batch.latents.shape
-        _, vh, vw = self.vae_stride
-        _, pH, pW = self.patch_size
-        latent_h = (self.sr_height // vh // pH) * pH
-        latent_w = (self.sr_width // vw // pW) * pW
-        actual_h = latent_h * vh
-        actual_w = latent_w * vw
-
-        latent_video = F.interpolate(
-            batch.latents,
-            size=(latent_t, latent_h, latent_w),
-            mode="trilinear",
-            align_corners=True,
-        )
-        if self.noise_value != 0:
-            noise = torch.randn_like(latent_video, device=device)
-            sigma = self.sigmas.to(device)[self.noise_value]
-            latent_video = latent_video * sigma + noise * (1 - sigma**2)**0.5
-
-        batch.latents = latent_video
-        batch.audio_latents = (
-            torch.randn_like(batch.audio_latents, device=batch.audio_latents.device) * self.sr_audio_noise_scale +
-            batch.audio_latents * (1 - self.sr_audio_noise_scale))
-        batch.height = actual_h
-        batch.width = actual_w
-        batch.magi_latent_T = latent_t
-        batch.magi_latent_H = latent_h
-        batch.magi_latent_W = latent_w
-        # Invalidate the static packed layout precomputed by the base
-        # latent prep stage: SR upsamples `batch.latents` to a larger
-        # spatial grid, which changes video_token_num / video_coords /
-        # video_mm. The SR denoising loop's
-        # `getattr(batch, "magi_static_packed_layout", None)` will then
-        # fall back to the slow path of `build_static_packed_inputs`,
-        # which rebuilds those fields from the new latent shape. SR
-        # only does ~5 denoising steps so the meshgrid recompute cost
-        # is negligible relative to SR-DiT forward.
-        batch.magi_static_packed_layout = None
-
-        if getattr(batch, "image_latent", None) is not None:
-            batch.image_latent = self._encode_image(batch, actual_h, actual_w)
-        return batch
-
-    def _encode_image(
-        self,
-        batch: ForwardBatch,
-        height: int,
-        width: int,
-    ) -> torch.Tensor:
-        image = getattr(batch, "image", None) or batch.pil_image
-        if image is None and batch.image_path is not None:
-            image = load_image(batch.image_path)
-        if image is None:
-            raise ValueError("MagiHuman SR TI2V requires an image for SR re-encoding.")
-        if not isinstance(image, Image.Image):
-            raise TypeError(f"Expected PIL image or image path, got {type(image)}")
-
-        device = get_local_torch_device()
-        image = _resizecrop(image.convert("RGB"), height, width)
-        image_tensor = self.video_processor.preprocess(
-            image,
-            height=height,
-            width=width,
-        ).to(device=device, dtype=torch.float32)
-        image_tensor = image_tensor.unsqueeze(2)
-
-        self.vae = self.vae.to(device)
-        encoded = self.vae.encode(image_tensor)
-        image_latent = encoded.mean if hasattr(encoded, "mean") else encoded
-
-        shift_factor = getattr(self.vae, "shift_factor", None)
-        if shift_factor is not None:
-            if isinstance(shift_factor, torch.Tensor):
-                image_latent = image_latent - shift_factor.to(
-                    image_latent.device,
-                    image_latent.dtype,
-                )
-            else:
-                image_latent = image_latent - shift_factor
-        scaling_factor = getattr(self.vae, "scaling_factor", None)
-        if scaling_factor is not None:
-            if isinstance(scaling_factor, torch.Tensor):
-                image_latent = image_latent * scaling_factor.to(
-                    image_latent.device,
-                    image_latent.dtype,
-                )
-            else:
-                image_latent = image_latent * scaling_factor
-        return image_latent.to(torch.float32)

fastvideo/pipelines/composed_pipeline_base.py

@@ -16,7 +16,6 @@ from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, get_world_group)
 from fastvideo.distributed.communication_op import (warmup_sequence_parallel_communication)
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.hooks.activation_trace import attach_activation_trace, detach_activation_trace
 from fastvideo.logger import init_logger
 from fastvideo.profiler import get_or_create_profiler
 from fastvideo.models.loader.component_loader import PipelineComponentLoader
@@ -63,7 +62,6 @@ class ComposedPipelineBase(ABC):
         self.model_path: str = model_path
         self._stages: list[PipelineStage] = []
         self._stage_name_mapping: dict[str, PipelineStage] = {}
-        self._trace_mgr = None
 
         if required_config_modules is not None:
             self._required_config_modules = required_config_modules
@@ -185,8 +183,6 @@ class ComposedPipelineBase(ABC):
                 )
                 logger.info("Torch Compile enabled for DiT")
 
-        self._trace_mgr = attach_activation_trace(self.modules.get("transformer"))
-
         if not self.fastvideo_args.training_mode:
             logger.info("Creating pipeline stages...")
             self.create_pipeline_stages(self.fastvideo_args)
@@ -459,10 +455,3 @@ class ComposedPipelineBase(ABC):
 
     def train(self) -> None:
         raise NotImplementedError("if training_mode is True, the pipeline must implement this method")
-
-    def close(self) -> None:
-        detach_activation_trace(getattr(self, "_trace_mgr", None))
-        self._trace_mgr = None
-
-    def __del__(self):
-        self.close()

fastvideo/registry.py

@@ -27,16 +27,6 @@ from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    MagiHumanBaseConfig,
-    MagiHumanBaseI2VConfig,
-    MagiHumanDistillConfig,
-    MagiHumanDistillI2VConfig,
-    MagiHumanSR1080pConfig,
-    MagiHumanSR1080pI2VConfig,
-    MagiHumanSR540pConfig,
-    MagiHumanSR540pI2VConfig,
-)
 from fastvideo.configs.pipelines.turbodiffusion import (
     TurboDiffusionI2V_A14B_Config,
     TurboDiffusionT2V_14B_Config,
@@ -294,146 +284,6 @@ def _register_configs() -> None:
         default_preset="stable_audio_open_small",
     )
 
-    # daVinci-MagiHuman SR-1080p (two-stage base + local-window SR text-to-AV).
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanSR1080pConfig,
-        workload_types=(WorkloadType.T2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-SR-1080p-Diffusers",
-            "FastVideo/MagiHuman-Diffusers/sr_1080p",
-        ],
-        model_detectors=[
-            lambda path:
-            (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower()) and
-             ("sr_1080p" in path.lower() or "sr-1080p" in path.lower() or "1080p_sr" in path.lower() or "sr1080p" in
-              path.lower()) and "ti2v" not in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_sr_1080p",
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanSR1080pI2VConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-SR-1080p-TI2V-Diffusers",
-        ],
-        model_detectors=[
-            lambda path:
-            (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower()) and
-             ("sr_1080p" in path.lower() or "sr-1080p" in path.lower() or "1080p_sr" in path.lower() or "sr1080p" in
-              path.lower()) and "ti2v" in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_sr_1080p_ti2v",
-    )
-
-    # daVinci-MagiHuman SR-540p (two-stage base + SR text-to-AV).
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanSR540pConfig,
-        workload_types=(WorkloadType.T2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-SR-540p-Diffusers",
-            "FastVideo/MagiHuman-Diffusers/sr_540p",
-        ],
-        model_detectors=[
-            lambda path:
-            (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower()) and
-             ("sr_540p" in path.lower() or "sr-540p" in path.lower() or "540p_sr" in path.lower() or "srpipeline" in
-              path.lower()) and "1080" not in path.lower() and "ti2v" not in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_sr_540p",
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanSR540pI2VConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-SR-540p-TI2V-Diffusers",
-        ],
-        model_detectors=[
-            lambda path:
-            (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower()) and
-             ("sr_540p" in path.lower() or "sr-540p" in path.lower() or "540p_sr" in path.lower() or "srpipeline" in
-              path.lower()) and "1080" not in path.lower() and "ti2v" in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_sr_540p_ti2v",
-    )
-
-    # daVinci-MagiHuman (base text-to-AV).
-    # NOTE: WorkloadType has no T2AV variant yet; using T2V as the
-    # placeholder until the enum is extended (same as Stable Audio).
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanBaseConfig,
-        workload_types=(WorkloadType.T2V, ),
-        hf_model_paths=[
-            "GAIR/daVinci-MagiHuman",
-            "FastVideo/MagiHuman-Base-Diffusers",
-            "FastVideo/MagiHuman-Diffusers/base",
-        ],
-        model_detectors=[
-            lambda path:
-            (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower()
-              ) and "distill" not in path.lower() and "ti2v" not in path.lower() and "sr_540p" not in path.lower() and
-             "sr-540p" not in path.lower() and "540p_sr" not in path.lower() and "sr_1080p" not in path.lower() and
-             "sr-1080p" not in path.lower() and "1080p_sr" not in path.lower() and "srpipeline" not in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_base",
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanBaseI2VConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-Base-TI2V-Diffusers",
-        ],
-        model_detectors=[
-            lambda path:
-            (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower()
-              ) and "ti2v" in path.lower() and "distill" not in path.lower() and "sr_540p" not in path.lower() and
-             "sr-540p" not in path.lower() and "540p_sr" not in path.lower() and "sr_1080p" not in path.lower() and
-             "sr-1080p" not in path.lower() and "1080p_sr" not in path.lower() and "srpipeline" not in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_base_ti2v",
-    )
-    # daVinci-MagiHuman (DMD-2 distilled text-to-AV)
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanDistillConfig,
-        workload_types=(WorkloadType.T2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-Distilled-Diffusers",
-            "FastVideo/MagiHuman-Diffusers/distill",
-        ],
-        model_detectors=[
-            lambda path: (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower())
-                          and "distill" in path.lower() and "ti2v" not in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_distill",
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=MagiHumanDistillI2VConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/MagiHuman-Distilled-TI2V-Diffusers",
-        ],
-        model_detectors=[
-            lambda path: (("magihuman" in path.lower() or "magi_human" in path.lower() or "magi-human" in path.lower())
-                          and "ti2v" in path.lower() and "distill" in path.lower()),
-        ],
-        model_family="magi_human",
-        default_preset="magi_human_distill_ti2v",
-    )
-
     # Hunyuan 1.5 (specific)
     register_configs(
         sampling_param_cls=None,
@@ -974,8 +824,6 @@ def _register_presets() -> None:
         ALL_PRESETS as LONGCAT_PRESETS, )
     from fastvideo.pipelines.basic.ltx2.presets import (
         ALL_PRESETS as LTX2_PRESETS, )
-    from fastvideo.pipelines.basic.magi_human.presets import (
-        ALL_PRESETS as MAGI_HUMAN_PRESETS, )
     from fastvideo.pipelines.basic.matrixgame.presets import (
         ALL_PRESETS as MATRIXGAME_PRESETS, )
     from fastvideo.pipelines.basic.sd35.presets import (
@@ -997,7 +845,6 @@ def _register_presets() -> None:
         LINGBOTWORLD_PRESETS,
         LONGCAT_PRESETS,
         LTX2_PRESETS,
-        MAGI_HUMAN_PRESETS,
         MATRIXGAME_PRESETS,
         SD35_PRESETS,
         STABLE_AUDIO_PRESETS,

fastvideo/tests/hooks/test_activation_trace.py

@@ -1,157 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import json
-from pathlib import Path
-
-import torch
-from torch import nn
-
-from fastvideo.hooks.activation_trace import (
-    attach_activation_trace,
-    detach_activation_trace,
-    trace_step,
-)
-from fastvideo.hooks.hooks import ModuleHookManager
-
-
-class ToyModel(nn.Module):
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.block = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
-        self.other = nn.Linear(2, 2)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.other(self.block(x))
-
-
-class TupleLayer(nn.Module):
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.proj = nn.Linear(2, 2)
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        out = self.proj(x)
-        return out, out + 1
-
-
-class TupleOutputModel(nn.Module):
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.tuple = TupleLayer()
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        return self.tuple(x)
-
-
-def _read_jsonl(path: Path) -> list[dict]:
-    return [json.loads(line) for line in path.read_text().splitlines()]
-
-
-def test_attach_activation_trace_off_returns_none(monkeypatch) -> None:
-    monkeypatch.delenv("FASTVIDEO_TRACE_ACTIVATIONS", raising=False)
-    model = ToyModel()
-
-    manager = attach_activation_trace(model)
-
-    assert manager is None
-    assert len(model._forward_hooks) == 0
-    assert ModuleHookManager.get_from(model.block[0]) is None
-
-
-def test_attach_activation_trace_on_respects_layer_filter(
-    monkeypatch,
-    tmp_path,
-) -> None:
-    monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1")
-    monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0.*")
-    monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(tmp_path / "trace.jsonl"))
-    model = ToyModel()
-
-    manager = attach_activation_trace(model)
-
-    try:
-        assert manager is not None
-        assert ModuleHookManager.get_from(model.block[0]) is not None
-        assert ModuleHookManager.get_from(model.block[1]) is None
-        assert ModuleHookManager.get_from(model.other) is None
-    finally:
-        detach_activation_trace(manager)
-
-
-def test_activation_trace_writes_configured_stats(monkeypatch, tmp_path) -> None:
-    path = tmp_path / "trace.jsonl"
-    monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1")
-    monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0")
-    monkeypatch.setenv("FASTVIDEO_TRACE_STATS", "abs_mean,sum,shape,dtype")
-    monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path))
-    model = ToyModel()
-    manager = attach_activation_trace(model)
-
-    try:
-        with trace_step(3):
-            model(torch.ones(1, 2))
-    finally:
-        detach_activation_trace(manager)
-
-    records = _read_jsonl(path)
-    assert len(records) == 1
-    record = records[0]
-    assert record["module"] == "block.0"
-    assert record["tensor"] == "out"
-    assert record["step"] == 3
-    assert {"abs_mean", "sum", "shape", "dtype"}.issubset(record)
-    assert record["shape"] == [1, 2]
-    assert record["dtype"] == "torch.float32"
-
-
-def test_activation_trace_step_filter(monkeypatch, tmp_path) -> None:
-    path = tmp_path / "trace.jsonl"
-    monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1")
-    monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0")
-    monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path))
-    monkeypatch.setenv("FASTVIDEO_TRACE_STEPS", "0,2")
-    model = ToyModel()
-    manager = attach_activation_trace(model)
-
-    try:
-        for step_idx in range(4):
-            with trace_step(step_idx):
-                model(torch.ones(1, 2))
-    finally:
-        detach_activation_trace(manager)
-
-    assert [record["step"] for record in _read_jsonl(path)] == [0, 2]
-
-
-def test_activation_trace_flattens_tuple_outputs(monkeypatch, tmp_path) -> None:
-    path = tmp_path / "trace.jsonl"
-    monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1")
-    monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", "tuple$")
-    monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path))
-    model = TupleOutputModel()
-    manager = attach_activation_trace(model)
-
-    try:
-        model(torch.ones(1, 2))
-    finally:
-        detach_activation_trace(manager)
-
-    records = _read_jsonl(path)
-    assert [record["tensor"] for record in records] == ["out[0]", "out[1]"]
-
-
-def test_detach_activation_trace_removes_hooks(monkeypatch, tmp_path) -> None:
-    monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1")
-    monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0")
-    monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(tmp_path / "trace.jsonl"))
-    model = ToyModel()
-    manager = attach_activation_trace(model)
-
-    assert ModuleHookManager.get_from(model.block[0]) is not None
-
-    detach_activation_trace(manager)
-
-    assert ModuleHookManager.get_from(model.block[0]) is None

fastvideo/tests/modal/pr_test.py

@@ -240,17 +240,23 @@ def run_lora_extraction_tests():
               ],
               volumes={"/root/data": model_vol})
 def run_performance_tests():
-    # Dashboard runs after compare_baseline regardless of regression result so
-    # the trend view is always available when investigating a failed gate.
+    # compare_baseline.py emits normalized_perf_*.json artifacts for manual
+    # performance-baseline reseeds when the rolling comparison runs.
     run_test(
         "export HF_HOME='/root/data/.cache' && "
         "export PERFORMANCE_TRACKING_ROOT='/tmp/perf-tracking' && "
         "hf auth login --token $HF_API_KEY && "
-        "pytest ./fastvideo/tests/performance -vs && "
-        "{ python ./fastvideo/tests/performance/compare_baseline.py; "
+        "pytest ./fastvideo/tests/performance -vs; "
+        "PYTEST_RC=$?; "
+        "PERF_RC=0; "
+        "if [ $PYTEST_RC -eq 0 ]; then "
+        "python ./fastvideo/tests/performance/compare_baseline.py; "
         "PERF_RC=$?; "
+        "fi; "
         "python ./fastvideo/tests/performance/dashboard.py || true; "
-        "exit $PERF_RC; }")
+        "FINAL_RC=$PYTEST_RC; "
+        "if [ $FINAL_RC -eq 0 ]; then FINAL_RC=$PERF_RC; fi; "
+        "exit $FINAL_RC")
 
 
 @app.function(gpu="L40S:1",

fastvideo/tests/performance/compare_baseline.py

@@ -36,7 +36,23 @@ TRACKING_ROOT = os.environ.get(
     "PERFORMANCE_TRACKING_ROOT",
     "/tmp/perf-tracking",
 )
+PERF_REPORTS_DIR = os.environ.get("PERF_REPORTS_DIR", "/root/data/perf_reports")
 MAX_REGRESSION = float(os.environ.get("PERF_MAX_REGRESSION", "0.05"))
+METRICS = (
+    ("latency", "Latency", 3),
+    ("throughput", "Throughput", 3),
+    ("memory", "Memory", 1),
+    ("text_encoder_time_s", "Text Enc", 3),
+    ("dit_time_s", "DiT", 3),
+    ("vae_decode_time_s", "VAE Decode", 3),
+)
+LOWER_IS_BETTER_METRICS = {
+    "latency",
+    "memory",
+    "text_encoder_time_s",
+    "dit_time_s",
+    "vae_decode_time_s",
+}
 
 
 def _should_persist_tracking() -> bool:
@@ -44,7 +60,6 @@ def _should_persist_tracking() -> bool:
     branch = os.environ.get("BUILDKITE_BRANCH", "")
     return test_scope == "full" and branch == "main"
 
-
 def _load_current_results() -> list[dict[str, Any]]:
     pattern = os.path.join(RESULTS_DIR, "perf_*.json")
     records: list[dict[str, Any]] = []
@@ -54,7 +69,14 @@ def _load_current_results() -> list[dict[str, Any]]:
     return records
 
 
-def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
+def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
+    """Normalize a raw perf_*.json result into the HF tracking schema.
+
+    The Buildkite artifact intentionally keeps the raw benchmark output from
+    test_inference_performance.py. Baseline comparison, main-branch persistence,
+    and manual baseline reseeds should all use this mapping so the stored HF
+    records do not drift from the artifact schema.
+    """
     benchmark_id = result.get("benchmark_id", "unknown")
     model_id = benchmark_id
 
@@ -66,6 +88,9 @@ def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
     latency = safe_float(result.get("avg_generation_time_s"))
     throughput = safe_float(result.get("throughput_fps"))
     memory = safe_float(result.get("max_peak_memory_mb"))
+    text_encoder_time = safe_float(result.get("text_encoder_time_s"))
+    dit_time = safe_float(result.get("dit_time_s"))
+    vae_decode_time = safe_float(result.get("vae_decode_time_s"))
 
     return {
         "model_id": model_id,
@@ -75,10 +100,17 @@ def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
         "latency": latency,
         "throughput": throughput,
         "memory": memory,
+        "text_encoder_time_s": text_encoder_time,
+        "dit_time_s": dit_time,
+        "vae_decode_time_s": vae_decode_time,
         "success": True,
     }
 
 
+def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
+    return normalize_performance_result(result)
+
+
 def _write_tracking_record(record: dict[str, Any]) -> str:
     model_dir = os.path.join(TRACKING_ROOT, sanitize(record["model_id"]))
     os.makedirs(model_dir, exist_ok=True)
@@ -93,6 +125,24 @@ def _write_tracking_record(record: dict[str, Any]) -> str:
     return out_path
 
 
+def _write_normalized_artifact(record: dict[str, Any]) -> None:
+    try:
+        results_dir = os.path.join(PERF_REPORTS_DIR, "results")
+        os.makedirs(results_dir, exist_ok=True)
+        timestamp = sanitize(record["timestamp"])
+        model_id = sanitize(record["model_id"])
+        commit = sanitize(record["commit_sha"] or "unknown")
+        path = os.path.join(
+            results_dir,
+            f"normalized_perf_{model_id}_{timestamp}_{commit}.json",
+        )
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(record, f, indent=2)
+        print(f"Normalized performance result written to {path}")
+    except Exception as e:
+        print(f"Failed to write normalized performance result artifact: {e}")
+
+
 def _baseline_metric(records: list[dict[str, Any]], key: str) -> float | None:
     values = [safe_float(r.get(key)) for r in records]
     values = [v for v in values if v is not None]
@@ -108,7 +158,8 @@ def _check_regressions(
 ) -> list[str]:
     failures: list[str] = []
 
-    for metric in ("latency", "memory"):
+    for metric in ("latency", "memory", "text_encoder_time_s", "dit_time_s",
+                   "vae_decode_time_s"):
         baseline = _baseline_metric(baseline_records, metric)
         curr = safe_float(current.get(metric))
         if baseline is None or curr is None or baseline <= 0:
@@ -141,7 +192,7 @@ def _metric_delta_percent(
     if curr is None or baseline is None or baseline <= 0:
         return None
 
-    if metric in ("latency", "memory"):
+    if metric in LOWER_IS_BETTER_METRICS:
         return (curr - baseline) / baseline * 100.0
     if metric == "throughput":
         return (baseline - curr) / baseline * 100.0
@@ -161,28 +212,27 @@ def _build_summary_row(
 ) -> dict[str, Any]:
     """Format a single benchmark result as a row for the Markdown table."""
 
-    latency_base = _baseline_metric(baseline_records, "latency")
-    throughput_base = _baseline_metric(baseline_records, "throughput")
-    memory_base = _baseline_metric(baseline_records, "memory")
+    metric_values: dict[str, dict[str, float | None]] = {}
+    regressions: list[float] = []
+    for metric, _label, _precision in METRICS:
+        curr = safe_float(record.get(metric))
+        baseline = _baseline_metric(baseline_records, metric)
+        regression = _metric_delta_percent(metric, record, baseline_records)
+        metric_values[metric] = {
+            "curr": curr,
+            "base": baseline,
+            "regression_pct": regression,
+        }
+        if regression is not None:
+            regressions.append(regression)
 
-    # Calculate percentages for the 'Worst Regression' column
-    latency_reg = _metric_delta_percent("latency", record, baseline_records)
-    throughput_reg = _metric_delta_percent("throughput", record, baseline_records)
-    memory_reg = _metric_delta_percent("memory", record, baseline_records)
-
-    regressions = [v for v in (latency_reg, throughput_reg, memory_reg) if v is not None]
     worst_regression_pct = max(regressions) if regressions else None
 
     return {
         "model_id": record["model_id"],
         "gpu_type": record["gpu_type"],
         "baseline_n": len(baseline_records),
-        "latency_curr": safe_float(record.get("latency")),
-        "latency_base": latency_base,
-        "throughput_curr": safe_float(record.get("throughput")),
-        "throughput_base": throughput_base,
-        "memory_curr": safe_float(record.get("memory")),
-        "memory_base": memory_base,
+        "metrics": metric_values,
         "worst_regression_pct": worst_regression_pct,
         "failed": has_failed,
     }
@@ -199,24 +249,24 @@ def _build_markdown_summary(
         "",
         ("| Model | GPU | Baseline N | Latency (curr/base) | "
          "Throughput (curr/base) | Memory (curr/base) | "
-         "Worst Regression | Status |"),
-        "|---|---|---:|---|---|---|---:|---|",
+         "Text Enc (curr/base) | DiT (curr/base) | "
+         "VAE Decode (curr/base) | Worst Regression | Status |"),
+        "|---|---|---:|---|---|---|---|---|---|---:|---|",
     ]
 
     for row in summary_rows:
-        latency = (f"{_compact_value(row['latency_curr'])} / "
-                   f"{_compact_value(row['latency_base'])}")
-        throughput = (f"{_compact_value(row['throughput_curr'])} / "
-                      f"{_compact_value(row['throughput_base'])}")
-        memory = (f"{_compact_value(row['memory_curr'], 1)} / "
-                  f"{_compact_value(row['memory_base'], 1)}")
+        metric_cells = []
+        for metric, _label, precision in METRICS:
+            values = row["metrics"][metric]
+            metric_cells.append(f"{_compact_value(values['curr'], precision)} / "
+                                f"{_compact_value(values['base'], precision)}")
 
         worst_reg = ("n/a" if row["worst_regression_pct"] is None else f"{row['worst_regression_pct']:.1f}%")
         status = "FAIL" if row["failed"] else "PASS"
 
         lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
                      f"{row['baseline_n']} | "
-                     f"{latency} | {throughput} | {memory} | "
+                     f"{' | '.join(metric_cells)} | "
                      f"{worst_reg} | {status} |")
 
     return "\n".join(lines) + "\n"
@@ -233,11 +283,10 @@ def _emit_markdown_summary(markdown: str, commit_sha: str) -> None:
 
     # 2. Write to Modal volume for Buildkite to pick up in post-run hook
     try:
-        perf_reports_dir = "/root/data/perf_reports"
-        os.makedirs(perf_reports_dir, exist_ok=True)
+        os.makedirs(PERF_REPORTS_DIR, exist_ok=True)
         short_sha = commit_sha[:7] if commit_sha else "unknown"
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        report_path = os.path.join(perf_reports_dir, f"perf_{short_sha}_{timestamp}.md")
+        report_path = os.path.join(PERF_REPORTS_DIR, f"perf_{short_sha}_{timestamp}.md")
         with open(report_path, "w", encoding="utf-8") as f:
             f.write(markdown + "\n")
         print(f"Performance report written to {report_path}")
@@ -286,6 +335,8 @@ def main() -> int:
             record["success"] = not failures
             all_failures.extend(failures)
 
+        _write_normalized_artifact(record)
+
         # Strict upload: a silent failure would freeze the rolling baseline.
         if persist_tracking:
             current_path = _write_tracking_record(record)

fastvideo/tests/performance/dashboard.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
+from html import escape
 from datetime import datetime
 
 import plotly.express as px
@@ -7,6 +8,15 @@ import pandas as pd
 
 from hf_store import sync_from_hf, load_as_dataframe
 
+METRICS = (
+    "latency",
+    "throughput",
+    "memory",
+    "text_encoder_time_s",
+    "dit_time_s",
+    "vae_decode_time_s",
+)
+
 # -----------------------------
 # 1. Grouping
 # -----------------------------
@@ -19,15 +29,36 @@ def group_data(df: pd.DataFrame):
 # -----------------------------
 # 2. Plot builder
 # -----------------------------
-def build_plots(df: pd.DataFrame) -> list:
+def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
     figs = []
+    skipped_metrics: list[dict[str, object]] = []
 
     for (model_id, gpu_type), g in group_data(df):
         g = g.sort_values("timestamp")
 
         # One chart per metric so the y-axes aren't on wildly different scales
-        for metric in ("latency", "throughput", "memory"):
-            if g[metric].isna().all():
+        for metric in METRICS:
+            if metric not in g.columns:
+                skipped_metrics.append({
+                    "model_id": model_id,
+                    "gpu_type": gpu_type,
+                    "metric": metric,
+                    "reason": "column missing from loaded records",
+                    "records": len(g),
+                    "non_null": 0,
+                })
+                continue
+
+            non_null = int(g[metric].notna().sum())
+            if non_null == 0:
+                skipped_metrics.append({
+                    "model_id": model_id,
+                    "gpu_type": gpu_type,
+                    "metric": metric,
+                    "reason": "no non-null values in loaded records",
+                    "records": len(g),
+                    "non_null": non_null,
+                })
                 continue
 
             fig = px.line(
@@ -41,20 +72,57 @@ def build_plots(df: pd.DataFrame) -> list:
             )
             figs.append(fig)
 
-    return figs
+    return figs, skipped_metrics
+
+
+def render_skipped_metrics(skipped_metrics: list[dict[str, object]]) -> str:
+    if not skipped_metrics:
+        return ""
+
+    rows = [
+        "<h3>Skipped Metric Plots</h3>",
+        "<table>",
+        ("<thead><tr><th>Model</th><th>GPU</th><th>Metric</th>"
+         "<th>Records</th><th>Non-null</th><th>Reason</th></tr></thead>"),
+        "<tbody>",
+    ]
+    for item in skipped_metrics:
+        rows.append(
+            "<tr>"
+            f"<td>{escape(str(item['model_id']))}</td>"
+            f"<td>{escape(str(item['gpu_type']))}</td>"
+            f"<td>{escape(str(item['metric']))}</td>"
+            f"<td>{item['records']}</td>"
+            f"<td>{item['non_null']}</td>"
+            f"<td>{escape(str(item['reason']))}</td>"
+            "</tr>"
+        )
+    rows.extend(["</tbody>", "</table>"])
+    return "\n".join(rows)
 
 # -----------------------------
 # 3. Render HTML dashboard
 # -----------------------------
-def render_html(figs: list, days: int) -> str:
+def render_html(figs: list, skipped_metrics: list[dict[str, object]],
+                days: int) -> str:
     html_parts = [
         "<html>",
         "<head><meta charset='utf-8'>",
-        "<style>body { font-family: sans-serif; margin: 2rem; }</style>",
+        ("<style>"
+         "body { font-family: sans-serif; margin: 2rem; }"
+         "table { border-collapse: collapse; margin: 1rem 0 2rem; }"
+         "th, td { border: 1px solid #ddd; padding: 0.4rem 0.6rem; "
+         "text-align: left; }"
+         "th { background: #f5f5f5; }"
+         "</style>"),
         "</head><body>",
         f"<h2>Performance Dashboard (last {days} days)</h2>",
     ]
 
+    skipped_html = render_skipped_metrics(skipped_metrics)
+    if skipped_html:
+        html_parts.append(skipped_html)
+
     for fig in figs:
         html_parts.append(fig.to_html(full_html=False, include_plotlyjs="cdn"))
 
@@ -67,7 +135,7 @@ def render_html(figs: list, days: int) -> str:
 def main() -> None:
     days = int(os.environ.get("DASHBOARD_DAYS", "30"))
 
-    local_dir = sync_from_hf("/tmp/perf-tracking")
+    local_dir = sync_from_hf("/tmp/perf-tracking", reuse_existing=True)
     df = load_as_dataframe(local_dir, days=days)
 
     if df.empty:
@@ -79,8 +147,15 @@ def main() -> None:
           f"{df['gpu_type'].nunique()} GPU type(s), "
           f"date range: {df['timestamp'].min()} → {df['timestamp'].max()}")
 
-    figs = build_plots(df)
-    html = render_html(figs, days)
+    figs, skipped_metrics = build_plots(df)
+    if skipped_metrics:
+        print("Skipped metric plots:")
+        for item in skipped_metrics:
+            print(f"  - {item['model_id']} | {item['gpu_type']} | "
+                  f"{item['metric']}: {item['reason']} "
+                  f"({item['non_null']}/{item['records']} non-null)")
+    print(f"Generated {len(figs)} metric plot(s)")
+    html = render_html(figs, skipped_metrics, days)
 
     commit_sha = os.environ.get("BUILDKITE_COMMIT", "unknown")[:7]
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -97,4 +172,4 @@ def main() -> None:
     print(f"Dashboard generated: {output_file}")
 
 if __name__ == "__main__":
-    main()
+    main()

fastvideo/tests/performance/hf_store.py

@@ -25,6 +25,7 @@ from huggingface_hub import HfApi, snapshot_download
 
 HF_REPO_ID: str = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking")
 HF_TOKEN: str | None = os.environ.get("HF_API_KEY")
+SYNC_MARKER = ".hf_sync_complete"
 
 # ---------------------------------------------------------------------------
 # Low-level helpers
@@ -51,7 +52,16 @@ def safe_float(value: Any) -> float | None:
 # ---------------------------------------------------------------------------
 
 
-def sync_from_hf(local_dir: str, *, strict: bool = False) -> str:
+def _sync_marker_path(local_dir: str) -> str:
+    return os.path.join(local_dir, SYNC_MARKER)
+
+
+def sync_from_hf(
+    local_dir: str,
+    *,
+    strict: bool = False,
+    reuse_existing: bool = False,
+) -> str:
     """Download the HF dataset repo snapshot to *local_dir*.
 
     Returns *local_dir* so callers can chain: ``load_records(sync_from_hf(...))``.
@@ -61,7 +71,19 @@ def sync_from_hf(local_dir: str, *, strict: bool = False) -> str:
     unavailable. Callers that depend on the sync for correctness (e.g. the
     main-branch baseline writer) must pass ``strict=True`` so that misconfig
     or transient HF errors fail loud rather than silently reset the baseline.
+
+    When ``reuse_existing=True``, a previous successful sync in ``local_dir``
+    is reused. This avoids duplicate HF snapshot checks when compare and
+    dashboard scripts run sequentially in the same CI job.
     """
+    marker_path = _sync_marker_path(local_dir)
+    if reuse_existing and os.path.exists(marker_path):
+        print(f"hf_store: reusing existing sync at {local_dir}")
+        return local_dir
+
+    if not reuse_existing and os.path.exists(marker_path):
+        os.remove(marker_path)
+
     if not HF_REPO_ID:
         msg = "hf_store: HF_REPO_ID not set"
         if strict:
@@ -78,6 +100,12 @@ def sync_from_hf(local_dir: str, *, strict: bool = False) -> str:
             token=HF_TOKEN,
             allow_patterns="*.json",
         )
+        os.makedirs(local_dir, exist_ok=True)
+        with open(marker_path, "w", encoding="utf-8") as marker:
+            json.dump({
+                "repo_id": HF_REPO_ID,
+                "synced_at": datetime.now(timezone.utc).isoformat(),
+            }, marker)
     except Exception as exc:
         if strict:
             raise
@@ -223,14 +251,18 @@ def load_records_for_model(
 # DataFrame helpers (dashboard / analytics consumers)
 # ---------------------------------------------------------------------------
 
-_NUMERIC_COLS = ("latency", "throughput", "memory")
+_NUMERIC_COLS = (
+    "latency", "throughput", "memory",
+    "text_encoder_time_s", "dit_time_s", "vae_decode_time_s",
+)
 
 
 def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     """Apply standard type coercions to a raw records DataFrame.
 
     - Parses ``timestamp`` to UTC-aware datetime.
-    - Coerces ``latency``, ``throughput``, ``memory`` to float.
+    - Coerces ``latency``, ``throughput``, ``memory``, ``text_encoder_time_s``,
+      ``dit_time_s``, ``vae_decode_time_s`` to float.
     - Adds a ``config_id`` column (first 7 chars of ``commit_sha``).
 
     Returns the mutated DataFrame (also modifies in place for efficiency).

fastvideo/tests/performance/test_inference_performance.py

@@ -10,17 +10,34 @@ import glob
 import json
 import os
 import time
+from collections.abc import Mapping
 from datetime import datetime, timezone
 
 import torch
 import pytest
 
+import fastvideo.envs as envs
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
 from fastvideo.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)
 
+STAGE_METRIC_MAP: dict[str, str] = {
+    "TextEncoderStage": "text_encoder_time_s",
+    "TextEncodingStage": "text_encoder_time_s",
+    "CLIPTextEncoderStage": "text_encoder_time_s",
+    "T5TextEncoderStage": "text_encoder_time_s",
+    "DiTStage": "dit_time_s",
+    "DenoisingStage": "dit_time_s",
+    "DmdDenoisingStage": "dit_time_s",
+    "TransformerStage": "dit_time_s",
+    "VAEDecodeStage": "vae_decode_time_s",
+    "VAEDecoderStage": "vae_decode_time_s",
+    "DecodingStage": "vae_decode_time_s",
+    "DecodeStage": "vae_decode_time_s",
+}
+
 # -- Config discovery -------------------------------------------------------
 
 _BENCHMARKS_DIR = os.path.join(
@@ -73,16 +90,53 @@ def _shutdown_executor(generator):
         generator.executor.shutdown()
 
 
+def _extract_component_times(result: dict) -> dict[str, float | None]:
+    component_times: dict[str, float | None] = {
+        "text_encoder_time_s": None,
+        "dit_time_s": None,
+        "vae_decode_time_s": None,
+    }
+    logging_info = result.get("logging_info")
+    if logging_info is None:
+        return component_times
+    if isinstance(logging_info, Mapping):
+        stages: dict = logging_info.get("stages", {}) or {}
+    else:
+        stages: dict = getattr(logging_info, "stages", {}) or {}
+    if not stages:
+        return component_times
+    logger.info("Discovered pipeline stages: %s", list(stages.keys()))
+    for stage_name, stage_data in stages.items():
+        metric_key = STAGE_METRIC_MAP.get(stage_name)
+        if metric_key is None:
+            logger.debug("Unmapped stage '%s' (%.3fs)",
+                         stage_name,
+                         stage_data.get("execution_time", 0))
+            continue
+        elapsed = stage_data.get("execution_time")
+        if elapsed is None:
+            continue
+        existing = component_times[metric_key]
+        component_times[metric_key] = (
+            elapsed if existing is None else existing + elapsed)
+    return component_times
+
+
+def _avg_component(all_component_times: list[dict], key: str) -> float | None:
+    vals = [r[key] for r in all_component_times if r.get(key) is not None]
+    return round(sum(vals) / len(vals), 3) if vals else None
+
+
 def _run_generation(generator, prompt, generation_kwargs):
-    """Run a single generation, return (elapsed_s, peak_memory_mb)."""
+    """Run a single generation, return (elapsed_s, peak_memory_mb, component_times)."""
     torch.cuda.synchronize()
     start = time.perf_counter()
     result = generator.generate_video(prompt, **generation_kwargs)
     torch.cuda.synchronize()
     elapsed = time.perf_counter() - start
     peak_memory_mb = result.get("peak_memory_mb", 0.0) or 0.0
-    return elapsed, peak_memory_mb
-
+    component_times = _extract_component_times(result)
+    return elapsed, peak_memory_mb, component_times
 
 def _write_results(results):
     """Write JSON results to the results directory."""
@@ -102,15 +156,7 @@ def _write_results(results):
 
 # -- Test -------------------------------------------------------------------
 
-
-@pytest.mark.parametrize(
-    "cfg",
-    _BENCHMARK_CONFIGS,
-    ids=[c["benchmark_id"] for c in _BENCHMARK_CONFIGS],
-)
-def test_inference_performance(cfg):
-    """Measure generation latency and peak GPU memory,
-    assert against device-aware thresholds."""
+def _run_benchmark(cfg):
     run_config = cfg.get("run_config", {})
     required_gpus = run_config.get("required_gpus", 1)
     available = torch.cuda.device_count()
@@ -152,12 +198,18 @@ def test_inference_performance(cfg):
 
         times = []
         peak_memories = []
+        all_component_times = []
         for i in range(num_measure):
             logger.info("Measurement run %d/%d", i + 1, num_measure)
-            elapsed, peak_mb = _run_generation(generator, prompt, gen_kwargs)
+            elapsed, peak_mb, component_times = _run_generation(
+                generator,
+                prompt,
+                gen_kwargs,
+            )
             logger.info("  Time: %.2fs, Peak memory: %.0fMB", elapsed, peak_mb)
             times.append(elapsed)
             peak_memories.append(peak_mb)
+            all_component_times.append(component_times)
     finally:
         _shutdown_executor(generator)
 
@@ -186,6 +238,11 @@ def test_inference_performance(cfg):
         "commit": os.environ.get("BUILDKITE_COMMIT", ""),
         "pr_number": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
         "timestamp": datetime.now(timezone.utc).isoformat(),
+        "text_encoder_time_s": _avg_component(all_component_times,
+                                              "text_encoder_time_s"),
+        "dit_time_s": _avg_component(all_component_times, "dit_time_s"),
+        "vae_decode_time_s": _avg_component(all_component_times,
+                                            "vae_decode_time_s"),
     }
 
     logger.info(
@@ -203,3 +260,41 @@ def test_inference_performance(cfg):
     assert max_peak_memory <= max_mem, (
         f"Peak memory {max_peak_memory:.0f}MB exceeds "
         f"threshold {max_mem:.0f}MB for {device_name}")
+
+    component_thresholds = {
+        "text_encoder_time_s": thresholds.get("max_text_encoder_time_s"),
+        "dit_time_s": thresholds.get("max_dit_time_s"),
+        "vae_decode_time_s": thresholds.get("max_vae_decode_time_s"),
+    }
+    for metric, max_val in component_thresholds.items():
+        if max_val is None:
+            continue
+        actual = results[metric]
+        if actual is not None:
+            assert actual <= max_val, (
+                f"{metric} {actual:.3f}s exceeds threshold {max_val:.3f}s "
+                f"for {device_name}")
+
+
+@pytest.mark.parametrize(
+    "cfg",
+    _BENCHMARK_CONFIGS,
+    ids=[c["benchmark_id"] for c in _BENCHMARK_CONFIGS],
+)
+def test_inference_performance(cfg):
+    """Measure generation latency, peak GPU memory, and component-level timings
+    (text encoder, DiT, VAE decode). Assert each against device-aware thresholds.
+    """
+
+    original_env = os.environ.get("FASTVIDEO_STAGE_LOGGING")
+    original_getter = envs.environment_variables["FASTVIDEO_STAGE_LOGGING"]
+    os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
+    envs.environment_variables["FASTVIDEO_STAGE_LOGGING"] = lambda: True
+    try:
+        _run_benchmark(cfg)
+    finally:
+        if original_env is None:
+            os.environ.pop("FASTVIDEO_STAGE_LOGGING", None)
+        else:
+            os.environ["FASTVIDEO_STAGE_LOGGING"] = original_env
+        envs.environment_variables["FASTVIDEO_STAGE_LOGGING"] = original_getter

fastvideo/tests/ssim/test_magi_human_similarity.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""SSIM-based similarity test for daVinci-MagiHuman base text-to-AV.
-
-Reference videos for this test are seeded separately via the
-`.agents/skills/seed-ssim-references/` skill on Modal L40S and uploaded
-to `FastVideo/ssim-reference-videos`. Until refs exist, the first run
-will fail downloading; run the seed skill once and commit the URLs.
-
-Resolution + steps kept small enough for a CI budget; the full-quality
-variant falls back to the registered preset defaults.
-"""
-import os
-
-import pytest
-
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.logger import init_logger
-from fastvideo.tests.ssim.inference_similarity_utils import (
-    resolve_inference_device_reference_folder,
-    run_text_to_video_similarity_test,
-)
-
-logger = init_logger(__name__)
-
-# 15B DiT + T5-Gemma 9B + Wan VAE + Stable-Audio VAE doesn't fit on a
-# single L40S (44 GB). Shard across 2 ranks via FSDP.
-REQUIRED_GPUS = 2
-
-device_reference_folder = resolve_inference_device_reference_folder(logger)
-
-# Umbrella HF repo holds all four variants under sibling subfolders;
-# `maybe_download_model` parses "org/repo/subfolder" and only fetches
-# the selected subfolder. Override via `MAGI_HUMAN_MODEL_PATH` to point
-# at a local converted_weights/ dir.
-_MAGI_HUMAN_MODEL_PATH = os.getenv(
-    "MAGI_HUMAN_MODEL_PATH",
-    "FastVideo/MagiHuman-Diffusers/base",
-)
-
-MAGI_HUMAN_BASE_PARAMS = {
-    "num_gpus": 2,
-    "model_path": _MAGI_HUMAN_MODEL_PATH,
-    # height/width/guidance_scale/seed/fps mirror the registered
-    # `magi_human_base` preset defaults (see
-    # `fastvideo/pipelines/basic/magi_human/presets.py::MAGI_HUMAN_BASE`)
-    # so the SSIM test exercises the same code path as
-    # `examples/inference/basic/basic_magi_human.py`. Only the budget
-    # knobs (num_frames, num_inference_steps, sp_size) differ for CI fit.
-    "height": 256,
-    "width": 480,
-    "num_frames": 26,                # seconds=1 at fps=25 + 1; preset = 101
-    "num_inference_steps": 8,        # CI budget; preset = 32
-    "guidance_scale": 5.0,
-    "seed": 42,
-    "sp_size": 2,
-    "tp_size": 1,
-    "fps": 25,
-}
-
-try:
-    _MAGI_HUMAN_FULL_DEFAULTS = SamplingParam.from_pretrained(_MAGI_HUMAN_MODEL_PATH)
-    MAGI_HUMAN_FULL_PARAMS = {
-        "num_gpus": MAGI_HUMAN_BASE_PARAMS["num_gpus"],
-        "model_path": MAGI_HUMAN_BASE_PARAMS["model_path"],
-        "height": _MAGI_HUMAN_FULL_DEFAULTS.height,
-        "width": _MAGI_HUMAN_FULL_DEFAULTS.width,
-        "num_frames": _MAGI_HUMAN_FULL_DEFAULTS.num_frames,
-        "num_inference_steps": _MAGI_HUMAN_FULL_DEFAULTS.num_inference_steps,
-        "guidance_scale": _MAGI_HUMAN_FULL_DEFAULTS.guidance_scale,
-        "seed": _MAGI_HUMAN_FULL_DEFAULTS.seed,
-        "sp_size": MAGI_HUMAN_BASE_PARAMS["sp_size"],
-        "tp_size": MAGI_HUMAN_BASE_PARAMS["tp_size"],
-        "fps": _MAGI_HUMAN_FULL_DEFAULTS.fps,
-    }
-except Exception:
-    # Model not registered / accessible on this machine — fall back to the
-    # quick params as the full-quality map too; the test will skip anyway
-    # when the model path is unavailable.
-    MAGI_HUMAN_FULL_PARAMS = MAGI_HUMAN_BASE_PARAMS
-
-
-MAGI_HUMAN_MODEL_TO_PARAMS = {
-    "MagiHuman-Base-Diffusers": MAGI_HUMAN_BASE_PARAMS,
-}
-FULL_QUALITY_MAGI_HUMAN_MODEL_TO_PARAMS = {
-    "MagiHuman-Base-Diffusers": MAGI_HUMAN_FULL_PARAMS,
-}
-
-MAGI_HUMAN_TEST_PROMPTS = [
-    "A person sitting by a window, softly lit by afternoon sun, waving at "
-    "the camera with a gentle smile.",
-]
-
-
-@pytest.mark.parametrize("prompt", MAGI_HUMAN_TEST_PROMPTS)
-@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
-@pytest.mark.parametrize("model_id", list(MAGI_HUMAN_MODEL_TO_PARAMS.keys()))
-def test_magi_human_base_inference_similarity(
-    prompt: str,
-    attention_backend_name: str,
-    model_id: str,
-) -> None:
-    run_text_to_video_similarity_test(
-        logger=logger,
-        script_dir=os.path.dirname(os.path.abspath(__file__)),
-        device_reference_folder=device_reference_folder,
-        prompt=prompt,
-        attention_backend_name=attention_backend_name,
-        model_id=model_id,
-        default_params_map=MAGI_HUMAN_MODEL_TO_PARAMS,
-        full_quality_params_map=FULL_QUALITY_MAGI_HUMAN_MODEL_TO_PARAMS,
-        min_acceptable_ssim=0.60,
-    )

fastvideo/tests/stages/test_magi_human_sr_latent_preparation.py

@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression test: SR latent prep must invalidate the static-packed layout.
-
-The base latent prep stage (`MagiHumanLatentPreparationStage`) precomputes
-``batch.magi_static_packed_layout`` for the BASE-resolution latent and stashes
-it on the batch so the base denoising loop can reuse it across all denoising
-steps (C4 perf optimization, commit 4190c720).
-
-The SR latent prep stage (`MagiHumanSRLatentPreparationStage`) upsamples
-``batch.latents`` to a much larger spatial grid (e.g. 256x480 -> 512x896 for
-SR-540p), which changes the layout's video_token_num / video_coords / video_mm.
-Without invalidating the layout, the SR denoising loop reuses the stale
-base-sized layout and crashes inside ``MagiHumanDiT.adapter`` with::
-
-    IndexError: The shape of the mask [3243] at index 0 does not match
-    the shape of the indexed tensor [11771, 3584] at index 0
-
-See git f1eeb630 for the fix and a commit-message-level explanation.
-
-This is a pure logic test — no GPU, no model load, no upstream daVinci-MagiHuman
-clone needed. It runs in the default CI suite.
-"""
-from __future__ import annotations
-
-import torch
-
-from fastvideo.pipelines.basic.magi_human.stages.sr_latent_preparation import (
-    MagiHumanSRLatentPreparationStage,
-    ZeroSNRDDPMDiscretization,
-)
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-
-def _make_stage() -> MagiHumanSRLatentPreparationStage:
-    """Bypass __init__: only set the fields the T2V forward() path reads."""
-    stage = MagiHumanSRLatentPreparationStage.__new__(
-        MagiHumanSRLatentPreparationStage)
-    # vae + video_processor are only used by `_encode_image` (TI2V path);
-    # T2V skips that branch when `batch.image_latent is None`.
-    stage.vae = None
-    stage.video_processor = None
-    stage.vae_stride = (4, 16, 16)
-    stage.patch_size = (1, 2, 2)
-    # `noise_value=0` skips the sigma noise-injection branch — keeps the
-    # test deterministic and avoids depending on torch.randn.
-    stage.noise_value = 0
-    stage.sr_audio_noise_scale = 0.7
-    stage.sr_height = 512
-    stage.sr_width = 896
-    stage.sigmas = ZeroSNRDDPMDiscretization()(
-        1000, do_append_zero=False, flip=True)
-    return stage
-
-
-def test_sr_latent_prep_invalidates_static_packed_layout():
-    stage = _make_stage()
-
-    base_latent = torch.randn(1, 48, 7, 16, 30, dtype=torch.float32)
-    audio = torch.randn(1, 26, 64, dtype=torch.float32)
-
-    batch = ForwardBatch(data_type="video")
-    batch.latents = base_latent
-    batch.audio_latents = audio
-
-    sentinel = object()
-    batch.magi_static_packed_layout = sentinel  # type: ignore[attr-defined]
-
-    out = stage.forward(batch, fastvideo_args=None)  # type: ignore[arg-type]
-
-    assert out is batch
-    # Sanity: SR actually upsampled to a different spatial grid.
-    assert out.latents.shape[-1] != base_latent.shape[-1]
-    assert out.latents.shape[-2] != base_latent.shape[-2]
-    # The bug-fix invariant: the stale base-sized layout is gone, so the
-    # SR denoising loop's `getattr(batch, "magi_static_packed_layout", None)`
-    # falls back to None and `build_static_packed_inputs` rebuilds the
-    # layout from the new SR-sized latent.
-    assert getattr(out, "magi_static_packed_layout", "<missing>") is None

fastvideo/utils.py

@@ -496,23 +496,14 @@ def import_pynvml():
 def maybe_download_model(model_name_or_path: str, local_dir: str | None = None, download: bool = True) -> str:
     """
     Check if the model path is a Hugging Face Hub model ID and download it if needed.
-
-    Supports an "umbrella" repo layout where a single HF repo holds multiple
-    pipeline variants under sibling subfolders. If the input is shaped as
-    ``org/repo/subfolder`` (i.e. a non-existent local path with 3+ slash-
-    separated components and at least one segment that does not look like a
-    posix-absolute path), treat the first two components as the HF repo id
-    and the remainder as a subfolder; only the subfolder's blobs are
-    downloaded, and the returned local path points inside that subfolder.
-
+    
     Args:
-        model_name_or_path: Local path, Hugging Face Hub model ID, or
-            ``org/repo/subfolder`` umbrella-repo reference.
+        model_name_or_path: Local path or Hugging Face Hub model ID
         local_dir: Local directory to save the model
         download: Whether to download the model from Hugging Face Hub
-
+        
     Returns:
-        Local path to the model (or to the subfolder inside the snapshot).
+        Local path to the model
     """
 
     # If the path exists locally, return it
@@ -520,32 +511,8 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None,
         logger.info("Model already exists locally at %s", model_name_or_path)
         return model_name_or_path
 
-    # Detect the umbrella-repo "org/repo/subfolder[/nested]" form. HF Hub
-    # repo ids are exactly two components ("org/name"); anything more is
-    # always a subfolder reference. Local absolute paths are excluded by
-    # the os.path.exists check above and by the leading-slash test below.
-    repo_id = model_name_or_path
-    subfolder: str | None = None
-    parts = model_name_or_path.split("/")
-    if (len(parts) >= 3 and not model_name_or_path.startswith("/") and not model_name_or_path.startswith(".")
-            and "" not in parts):
-        repo_id = "/".join(parts[:2])
-        subfolder = "/".join(parts[2:])
-
     # Otherwise, assume it's a HF Hub model ID and try to download it
     try:
-        if subfolder is not None:
-            logger.info("Downloading umbrella-repo subfolder %s/%s from HF Hub...", repo_id, subfolder)
-            with get_lock(model_name_or_path):
-                local_path = snapshot_download(
-                    repo_id=repo_id,
-                    allow_patterns=[f"{subfolder}/**"],
-                    local_dir=local_dir,
-                )
-            local_path = os.path.join(local_path, subfolder)
-            logger.info("Downloaded subfolder to %s", local_path)
-            return str(local_path)
-
         logger.info("Downloading model snapshot from HF Hub for %s...", model_name_or_path)
         with get_lock(model_name_or_path):
             local_path = snapshot_download(repo_id=model_name_or_path,
@@ -600,38 +567,19 @@ def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
         raise ValueError(f"Model directory {model_path} does not contain model_index.json. "
                          "Only Hugging Face diffusers format is supported.")
 
-    # Load the config first so directory checks below can be conditional
-    # on what model_index.json actually declares.
-    with open(config_path) as f:
-        config = json.load(f)
-
-    # transformer/ is mandatory for every supported pipeline; the variant-
-    # specific DiT weights live there.
+    # Check for transformer and vae directories
     transformer_dir = os.path.join(model_path, "transformer")
+    vae_dir = os.path.join(model_path, "vae")
+
     if not os.path.exists(transformer_dir):
         raise ValueError(f"Model directory {model_path} does not contain a transformer/ directory.")
 
-    # Other components (vae, text_encoder, audio_vae, tokenizer, ...) are
-    # only required to live in a local subfolder if model_index.json
-    # actually lists them. Pipelines that lazy-load shared components
-    # from upstream HF repos (e.g. MagiHuman lazy-loading the Wan VAE,
-    # T5-Gemma, Stable Audio) emit a model_index.json that omits those
-    # keys, and the pipeline subclass handles the load at module-build
-    # time. Enforce only the "declared but missing" mismatch.
-    _OPTIONAL_COMPONENT_DIRS = (
-        "vae",
-        "text_encoder",
-        "tokenizer",
-        "audio_vae",
-        "scheduler",
-        "image_encoder",
-    )
-    for key in _OPTIONAL_COMPONENT_DIRS:
-        if key in config:
-            subdir = os.path.join(model_path, key)
-            if not os.path.exists(subdir):
-                raise ValueError(f"Model directory {model_path} declares `{key}` in "
-                                 f"model_index.json but is missing the {key}/ subfolder.")
+    if not os.path.exists(vae_dir):
+        raise ValueError(f"Model directory {model_path} does not contain a vae/ directory.")
+
+    # Load the config
+    with open(config_path) as f:
+        config = json.load(f)
 
     # Verify diffusers version exists
     if "_diffusers_version" not in config:

pyproject.toml

@@ -171,8 +171,6 @@ follow_imports = "silent"
 
 [tool.codespell]
 skip = "./data,./wandb,ui/package-lock.json"
-# "TReAD" is daVinci-MagiHuman's acronym (Token Routing and Early Drop).
-ignore-words-list = "TReAD,tread"
 
 [tool.ruff]
 # Allow lines to be as long as 120.

scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py

@@ -1,564 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Convert daVinci-MagiHuman (GAIR-NLP) weights to a Diffusers-format repo.
-
-MagiHuman publishes weights in a raw layout on HuggingFace
-(https://huggingface.co/GAIR/daVinci-MagiHuman). The layout is:
-
-    base/                       <- DiT safetensors (sharded)
-    distill/                    <- distilled DiT (out of scope for the base port)
-    540p_sr/, 1080p_sr/         <- super-resolution DiTs (out of scope)
-    turbo_vae/                  <- optional fast VAE decoder (out of scope for first cut)
-
-The base DiT uses Wan-AI/Wan2.2-TI2V-5B's VAE and google/t5gemma-9b-9b-ul2's
-encoder at inference time; neither is bundled upstream.
-
-This converter takes the raw MagiHuman base DiT and emits a Diffusers-style
-directory so `VideoGenerator.from_pretrained(...)` can load it standalone:
-
-    <output>/
-        model_index.json
-        transformer/
-            config.json
-            diffusion_pytorch_model-00001-of-00N.safetensors (+ index)
-        scheduler/
-            scheduler_config.json                    (FlowUniPC default)
-        vae/                                         (optional; --bundle-vae)
-        audio_vae/                                   (optional; --bundle-audio-vae)
-        text_encoder/, tokenizer/                    (optional; --bundle-text-encoder)
-
-By default the converted repo is MINIMAL: only `transformer/`,
-`scheduler/`, and `model_index.json` are emitted (~5-30 GB depending on
-variant). The four cross-variant shared components — Wan VAE, Stable
-Audio VAE, T5-Gemma encoder, and tokenizer — are lazy-loaded by
-`MagiHumanPipeline.load_modules` from their canonical upstream HF repos
-on first build, so all MagiHuman variants share a single ~25 GB cache
-of upstream weights. Pass the `--bundle-*` flags only if you want to
-ship a self-contained snapshot.
-
-The DiT key names pass through unchanged — the FastVideo `MagiHumanDiT` module
-mirrors the reference module tree (`adapter.*`, `block.layers.*`, `final_*`),
-so no regex remapping is needed. The conversion is effectively a reshard +
-Diffusers wrapper.
-
-    Example (minimal artifact, ~5-30 GB):
-    python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \\
-        --source GAIR/daVinci-MagiHuman \\
-        --subfolder base \\
-        --output converted_weights/magi_human_base
-
-Example (self-contained SR-540p artifact with base + SR DiTs):
-    python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \
-        --source GAIR/daVinci-MagiHuman \
-        --subfolder base \
-        --sr-source GAIR/daVinci-MagiHuman \
-        --sr-subfolder 540p_sr \
-        --output converted_weights/magi_human_sr_540p
-
-Example (self-contained snapshot with shared components bundled):
-    python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \\
-        --source GAIR/daVinci-MagiHuman \\
-        --subfolder base \\
-        --output converted_weights/magi_human_base \\
-        --bundle-vae --bundle-audio-vae --bundle-text-encoder
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import shutil
-from collections import OrderedDict
-from pathlib import Path
-
-import torch
-from huggingface_hub import hf_hub_download, snapshot_download
-from safetensors.torch import load_file, save_file
-
-
-# MagiHuman base arch — keys must be valid `MagiHumanArchConfig` fields.
-# FastVideo's `TransformerLoader.load` calls `ArchConfig.update_model_arch`
-# with this dict (minus `_class_name`, `_diffusers_version`) and rejects
-# any key that isn't a declared field. Pipeline-level knobs (steps, CFG,
-# guidance scales, flow_shift) live on `MagiHumanBaseConfig` and do NOT
-# belong here — they'd silently shadow the ArchConfig loader otherwise.
-MAGI_HUMAN_BASE_ARCH: dict = {
-    "_class_name": "MagiHumanDiT",
-    "_diffusers_version": "0.33.0",
-    # Transformer shape (upstream `ModelConfig`, `inference/common/config.py`).
-    "num_layers": 40,
-    "hidden_size": 5120,
-    "head_dim": 128,
-    "num_query_groups": 8,
-    # Modality channels.
-    "video_in_channels": 192,            # 48 (VAE z_dim) * patch_size product 1*2*2
-    "audio_in_channels": 64,
-    "text_in_channels": 3584,            # T5Gemma-9B encoder hidden size
-    # Block-level switches.
-    "mm_layers": [0, 1, 2, 3, 36, 37, 38, 39],
-    "local_attn_layers": [],
-    "gelu7_layers": [0, 1, 2, 3],
-    "post_norm_layers": [],
-    "enable_attn_gating": True,
-    "activation_type": "swiglu7",
-    # DiT patching / positional.
-    "patch_size": [1, 2, 2],
-    "spatial_rope_interpolation": "extra",
-    # TReAD (flattened; upstream nests as `tread_config`).
-    "tread_selection_rate": 0.5,
-    "tread_start_layer_idx": 2,
-    "tread_end_layer_idx": 25,
-}
-
-
-SCHEDULER_CONFIG: dict = {
-    "_class_name": "FlowUniPCMultistepScheduler",
-    "_diffusers_version": "0.33.0",
-    "num_train_timesteps": 1000,
-    "solver_order": 2,
-    "prediction_type": "flow_prediction",
-    "shift": 5.0,
-    "predict_x0": True,
-    "solver_type": "bh2",
-    "lower_order_final": True,
-    "disable_corrector": [],
-    "flow_shift": 5.0,
-}
-
-
-MAX_SHARD_BYTES = 5 * 1024 * 1024 * 1024  # 5 GB shards, matches HF defaults
-
-
-def _download_dit_shards(source: Path | str, subfolder: str = "base") -> list[Path]:
-    """Return local paths to all safetensors shards for the DiT."""
-    source = str(source)
-    if os.path.isdir(source):
-        shard_dir = Path(source) / subfolder
-        shards = sorted(shard_dir.glob("*.safetensors"))
-        if not shards:
-            raise FileNotFoundError(f"No safetensors under {shard_dir}")
-        return shards
-
-    # Remote HF repo — pull just the base subfolder.
-    local_dir = snapshot_download(
-        repo_id=source,
-        allow_patterns=[f"{subfolder}/*.safetensors", f"{subfolder}/*.json"],
-    )
-    shard_dir = Path(local_dir) / subfolder
-    return sorted(shard_dir.glob("*.safetensors"))
-
-
-def _load_all_shards(
-    shards: list[Path],
-    cast_bf16: bool = False,
-) -> "OrderedDict[str, torch.Tensor]":
-    """Load all safetensors shards into a single state dict.
-
-    When `cast_bf16` is True, fp32 tensors whose names match the transformer
-    core (attention / mlp / final_linear_* / adapter.{video,text,audio}_embedder)
-    are cast to bfloat16. fp32 is preserved for norms, rope bands, and any
-    other tensor where precision matters. This is the right default for
-    the distill checkpoint, which upstream ships as fp32 master weights
-    (61 GB) — casting yields a 30 GB Diffusers artifact that matches the
-    base checkpoint format.
-    """
-    # Tensors that must stay in float32 regardless of cast_bf16. These are
-    # the dtypes that appear as fp32 in the BASE checkpoint, which is the
-    # ground-truth shape of a "runtime-loadable" MagiHuman repo. The list
-    # includes:
-    #   - all RMSNorm weights (norms always run fp32 in upstream
-    #     MultiModalityRMSNorm and FV's mirror)
-    #   - the rope band buffer
-    #   - the adapter embedders (video/text/audio: weight + bias) which
-    #     upstream's Adapter declares as `dtype=torch.float32` and FV's
-    #     MagiAdapter mirrors at `magi_human.py:519-527`
-    #   - the final_linear_{video,audio} heads which upstream/FV both
-    #     declare as `dtype=torch.float32` (`magi_human.py:645-648`,
-    #     `dit_module.py:896-900`)
-    # Forgetting any of these makes `--cast-bf16` lossy for the distill
-    # checkpoint (which ships everything as fp32) and produces parity
-    # drift vs upstream that base does not exhibit (because base already
-    # ships with the right mixed-dtype layout).
-    _FP32_KEEP_SUFFIXES = (
-        ".pre_norm.weight",
-        ".q_norm.weight",
-        ".k_norm.weight",
-        ".attn_post_norm.weight",
-        ".mlp_post_norm.weight",
-        "final_norm_video.weight",
-        "final_norm_audio.weight",
-        "final_linear_video.weight",
-        "final_linear_audio.weight",
-        "adapter.video_embedder.weight",
-        "adapter.video_embedder.bias",
-        "adapter.text_embedder.weight",
-        "adapter.text_embedder.bias",
-        "adapter.audio_embedder.weight",
-        "adapter.audio_embedder.bias",
-        "adapter.rope.bands",
-    )
-    _FP32_KEEP_FULL = {"adapter.rope.bands"}
-
-    def _keep_fp32(k: str) -> bool:
-        if k in _FP32_KEEP_FULL:
-            return True
-        return any(k.endswith(s) for s in _FP32_KEEP_SUFFIXES)
-
-    state: OrderedDict[str, torch.Tensor] = OrderedDict()
-    for shard in shards:
-        piece = load_file(str(shard))
-        for k, v in piece.items():
-            if k in state:
-                raise RuntimeError(f"Duplicate key across shards: {k}")
-            if cast_bf16 and v.dtype == torch.float32 and not _keep_fp32(k):
-                v = v.to(torch.bfloat16)
-            state[k] = v
-        print(f"  loaded {shard.name} ({len(piece)} tensors)")
-    return state
-
-
-def _validate_state(state: dict[str, torch.Tensor]) -> None:
-    """Sanity-check required top-level modules are present."""
-    required_prefixes = (
-        "adapter.video_embedder.",
-        "adapter.text_embedder.",
-        "adapter.audio_embedder.",
-        "adapter.rope.bands",
-        "final_norm_video.",
-        "final_norm_audio.",
-        "final_linear_video.",
-        "final_linear_audio.",
-    )
-    for pref in required_prefixes:
-        if not any(k.startswith(pref) for k in state):
-            raise RuntimeError(f"Missing expected key prefix: {pref}")
-    # Layer count
-    layer_ids = {int(k.split(".")[2]) for k in state if k.startswith("block.layers.")}
-    if layer_ids != set(range(40)):
-        raise RuntimeError(f"Expected layers 0..39, got {sorted(layer_ids)}")
-
-
-def _shard_state_dict(
-    state: dict[str, torch.Tensor],
-    max_bytes: int = MAX_SHARD_BYTES,
-) -> tuple[list[dict[str, torch.Tensor]], dict[str, str]]:
-    """Greedy shard-packing: produce N shards of <= max_bytes, plus index."""
-    shards: list[dict[str, torch.Tensor]] = []
-    index: dict[str, str] = {}
-    cur: dict[str, torch.Tensor] = {}
-    cur_bytes = 0
-    shard_idx = 0
-    total = len(state)
-    for k, v in state.items():
-        t_bytes = v.numel() * v.element_size()
-        if cur and cur_bytes + t_bytes > max_bytes:
-            shards.append(cur)
-            cur = {}
-            cur_bytes = 0
-            shard_idx += 1
-        cur[k] = v
-        cur_bytes += t_bytes
-    if cur:
-        shards.append(cur)
-    n = len(shards)
-    for i, shard in enumerate(shards, start=1):
-        shard_name = f"diffusion_pytorch_model-{i:05d}-of-{n:05d}.safetensors"
-        for k in shard:
-            index[k] = shard_name
-    assert sum(len(s) for s in shards) == total
-    return shards, index
-
-
-def _write_transformer(
-    out_dir: Path,
-    state: dict[str, torch.Tensor],
-    arch: dict,
-    subdir: str = "transformer",
-) -> None:
-    transformer_dir = out_dir / subdir
-    transformer_dir.mkdir(parents=True, exist_ok=True)
-
-    shards, weight_map = _shard_state_dict(state)
-    n = len(shards)
-    total_bytes = sum(v.numel() * v.element_size() for v in state.values())
-    for i, shard in enumerate(shards, start=1):
-        shard_name = f"diffusion_pytorch_model-{i:05d}-of-{n:05d}.safetensors"
-        save_file(shard, str(transformer_dir / shard_name))
-        print(f"  wrote {shard_name} ({len(shard)} tensors)")
-
-    index = {"metadata": {"total_size": total_bytes}, "weight_map": weight_map}
-    with (transformer_dir / "diffusion_pytorch_model.safetensors.index.json").open("w") as f:
-        json.dump(index, f, indent=2)
-        f.write("\n")
-
-    with (transformer_dir / "config.json").open("w") as f:
-        json.dump(arch, f, indent=2)
-        f.write("\n")
-    print(f"  wrote {subdir}/config.json ({len(arch)} keys)")
-
-
-def _write_scheduler(out_dir: Path) -> None:
-    scheduler_dir = out_dir / "scheduler"
-    scheduler_dir.mkdir(parents=True, exist_ok=True)
-    with (scheduler_dir / "scheduler_config.json").open("w") as f:
-        json.dump(SCHEDULER_CONFIG, f, indent=2)
-        f.write("\n")
-    print(f"  wrote scheduler/scheduler_config.json")
-
-
-def _write_model_index(
-    out_dir: Path,
-    bundle_vae: bool,
-    bundle_text: bool,
-    bundle_audio_vae: bool = False,
-    include_sr_transformer: bool = False,
-    sr_subfolder: str = "540p_sr",
-) -> None:
-    pipeline_class = "MagiHumanPipeline"
-    if include_sr_transformer:
-        pipeline_class = (
-            "MagiHumanSR1080pPipeline"
-            if sr_subfolder == "1080p_sr" else "MagiHumanSRPipeline"
-        )
-    index = {
-        "_class_name": pipeline_class,
-        "_diffusers_version": "0.33.0",
-        "transformer": ["diffusers", "MagiHumanDiT"],
-        "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"],
-    }
-    if include_sr_transformer:
-        index["sr_transformer"] = ["diffusers", "MagiHumanDiT"]
-    if bundle_vae:
-        index["vae"] = ["diffusers", "AutoencoderKLWan"]
-    if bundle_audio_vae:
-        index["audio_vae"] = ["diffusers", "AutoencoderOobleck"]
-    if bundle_text:
-        index["text_encoder"] = ["transformers", "T5GemmaEncoderModel"]
-        index["tokenizer"] = ["transformers", "GemmaTokenizer"]
-    with (out_dir / "model_index.json").open("w") as f:
-        json.dump(index, f, indent=2)
-        f.write("\n")
-    print(f"  wrote model_index.json")
-
-
-def _bundle_wan_vae(out_dir: Path, source_repo: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers") -> None:
-    """Download the Wan 2.2 TI2V 5B VAE component into <out_dir>/vae/.
-
-    The `-Diffusers` variant has the canonical `vae/config.json` +
-    `vae/diffusion_pytorch_model.safetensors` layout. The plain
-    `Wan-AI/Wan2.2-TI2V-5B` repo ships the VAE as a single `.pth` at the
-    root, which is not `from_pretrained`-friendly.
-    """
-    print(f"  fetching VAE from {source_repo} ...")
-    local = snapshot_download(
-        repo_id=source_repo,
-        allow_patterns=["vae/*"],
-    )
-    src_vae = Path(local) / "vae"
-    if not src_vae.exists():
-        raise FileNotFoundError(f"No vae/ subdir in {source_repo}")
-    dst_vae = out_dir / "vae"
-    if dst_vae.exists():
-        shutil.rmtree(dst_vae)
-    shutil.copytree(src_vae, dst_vae)
-    print(f"  copied {src_vae} -> {dst_vae}")
-
-
-def _bundle_sa_audio_vae(out_dir: Path, source_repo: str = "stabilityai/stable-audio-open-1.0") -> None:
-    """Download the Stable Audio Open 1.0 VAE component into <out_dir>/audio_vae/.
-
-    Stability ships the VAE at `vae/config.json` +
-    `vae/diffusion_pytorch_model.safetensors` inside the main repo, so
-    the bundle is just a copy of that subdir. The repo is gated — the
-    caller's HF token must have accepted terms on
-    https://huggingface.co/stabilityai/stable-audio-open-1.0.
-    """
-    print(f"  fetching audio VAE from {source_repo} (gated) ...")
-    token = (
-        os.environ.get("HF_TOKEN")
-        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
-        or os.environ.get("HF_API_KEY")
-    )
-    local = snapshot_download(
-        repo_id=source_repo, token=token, allow_patterns=["vae/*"],
-    )
-    src = Path(local) / "vae"
-    if not src.exists():
-        raise FileNotFoundError(f"No vae/ subdir in {source_repo}")
-    dst = out_dir / "audio_vae"
-    if dst.exists():
-        shutil.rmtree(dst)
-    shutil.copytree(src, dst)
-    print(f"  copied {src} -> {dst}")
-
-
-def _bundle_text_encoder(out_dir: Path, source_repo: str = "google/t5gemma-9b-9b-ul2") -> None:
-    """Download the T5Gemma encoder + tokenizer.
-
-    T5Gemma is a Google gated repo; this step requires a write-scoped token with
-    accepted terms of use for the repo.
-    """
-    print(f"  fetching text encoder from {source_repo} (gated) ...")
-    token = (
-        os.environ.get("HF_TOKEN")
-        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
-        or os.environ.get("HF_API_KEY")
-    )
-    local = snapshot_download(
-        repo_id=source_repo,
-        token=token,
-        allow_patterns=[
-            "*.json",
-            "*.model",
-            "*.safetensors",
-            "*.safetensors.index.json",
-        ],
-    )
-    # Encoder-only bundling: keep tokenizer at the root and encoder weights
-    # under text_encoder/. HF's T5GemmaEncoderModel.from_pretrained(<dir>) on
-    # the whole repo works, but we split to match Diffusers layout.
-    src = Path(local)
-    dst_encoder = out_dir / "text_encoder"
-    dst_tokenizer = out_dir / "tokenizer"
-    if dst_encoder.exists():
-        shutil.rmtree(dst_encoder)
-    if dst_tokenizer.exists():
-        shutil.rmtree(dst_tokenizer)
-    dst_encoder.mkdir(parents=True, exist_ok=True)
-    dst_tokenizer.mkdir(parents=True, exist_ok=True)
-
-    for fname in src.iterdir():
-        if fname.name in {"tokenizer.model", "tokenizer.json", "tokenizer_config.json",
-                          "special_tokens_map.json", "spiece.model"}:
-            shutil.copy(fname, dst_tokenizer / fname.name)
-        else:
-            shutil.copy(fname, dst_encoder / fname.name)
-    print(f"  staged text_encoder and tokenizer from {src}")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
-    parser.add_argument(
-        "--source",
-        default="GAIR/daVinci-MagiHuman",
-        help="HF repo id or local directory containing base/*.safetensors shards.",
-    )
-    parser.add_argument(
-        "--subfolder",
-        default="base",
-        choices=["base", "distill", "540p_sr", "1080p_sr"],
-        help="Which MagiHuman variant to convert (scope of this skill: base).",
-    )
-    parser.add_argument(
-        "--output",
-        required=True,
-        help="Destination directory for the Diffusers-format repo.",
-    )
-    parser.add_argument(
-        "--bundle-vae",
-        action="store_true",
-        help="Download Wan-AI/Wan2.2-TI2V-5B VAE into <output>/vae/.",
-    )
-    parser.add_argument(
-        "--cast-bf16",
-        action="store_true",
-        help=(
-            "Cast fp32 DiT weights to bfloat16 on save. Recommended for the "
-            "distill subfolder (61 GB fp32 upstream -> 30 GB bf16 artifact). "
-            "Keeps norms, RoPE bands, and other precision-sensitive tensors "
-            "in fp32."
-        ),
-    )
-    parser.add_argument(
-        "--bundle-text-encoder",
-        action="store_true",
-        help="Download google/t5gemma-9b-9b-ul2 into <output>/text_encoder/ and tokenizer/. "
-             "Requires a write-scoped HF token with accepted terms of use.",
-    )
-    parser.add_argument(
-        "--bundle-audio-vae",
-        action="store_true",
-        help="Download stabilityai/stable-audio-open-1.0 VAE into <output>/audio_vae/. "
-             "Requires HF terms accepted for the Stability AI gated repo.",
-    )
-    parser.add_argument(
-        "--sr-source",
-        default=None,
-        help="Optional HF repo id or local directory containing SR DiT shards. When set, writes <output>/sr_transformer/.",
-    )
-    parser.add_argument(
-        "--sr-subfolder",
-        default="540p_sr",
-        choices=["540p_sr", "1080p_sr"],
-        help="SR source subfolder to convert into <output>/sr_transformer/.",
-    )
-    args = parser.parse_args()
-
-    out_dir = Path(args.output)
-    out_dir.mkdir(parents=True, exist_ok=True)
-
-    print(f"-> DiT shards from {args.source}/{args.subfolder}")
-    shards = _download_dit_shards(args.source, subfolder=args.subfolder)
-    print(f"  found {len(shards)} shard(s)")
-
-    print(f"-> loading DiT state dict (cast_bf16={args.cast_bf16})")
-    state = _load_all_shards(shards, cast_bf16=args.cast_bf16)
-    print(f"  total keys: {len(state)}")
-    _validate_state(state)
-    print(f"  state dict validation passed")
-
-    print(f"-> writing {out_dir}/transformer/")
-    _write_transformer(out_dir, state, MAGI_HUMAN_BASE_ARCH)
-
-    include_sr_transformer = args.sr_source is not None
-    if include_sr_transformer:
-        print(f"-> SR DiT shards from {args.sr_source}/{args.sr_subfolder}")
-        sr_shards = _download_dit_shards(args.sr_source, subfolder=args.sr_subfolder)
-        print(f"  found {len(sr_shards)} SR shard(s)")
-        print(f"-> loading SR DiT state dict (cast_bf16={args.cast_bf16})")
-        sr_state = _load_all_shards(sr_shards, cast_bf16=args.cast_bf16)
-        print(f"  total SR keys: {len(sr_state)}")
-        _validate_state(sr_state)
-        print("  SR state dict validation passed")
-        print(f"-> writing {out_dir}/sr_transformer/")
-        _write_transformer(
-            out_dir,
-            sr_state,
-            MAGI_HUMAN_BASE_ARCH,
-            subdir="sr_transformer",
-        )
-
-    print(f"-> writing {out_dir}/scheduler/")
-    _write_scheduler(out_dir)
-
-    if args.bundle_vae:
-        print(f"-> bundling video VAE (Wan 2.2 TI2V-5B)")
-        _bundle_wan_vae(out_dir)
-
-    if args.bundle_audio_vae:
-        print(f"-> bundling audio VAE (Stable Audio Open 1.0)")
-        _bundle_sa_audio_vae(out_dir)
-
-    if args.bundle_text_encoder:
-        print(f"-> bundling text encoder")
-        _bundle_text_encoder(out_dir)
-
-    print(f"-> writing model_index.json")
-    _write_model_index(
-        out_dir,
-        bundle_vae=args.bundle_vae,
-        bundle_text=args.bundle_text_encoder,
-        bundle_audio_vae=args.bundle_audio_vae,
-        include_sr_transformer=include_sr_transformer,
-        sr_subfolder=args.sr_subfolder,
-    )
-
-    print(f"\nDone. Output at: {out_dir}")
-    if not args.bundle_vae:
-        print("  (remember to fetch Wan-AI/Wan2.2-TI2V-5B VAE separately or re-run with --bundle-vae)")
-    if not args.bundle_text_encoder:
-        print("  (remember to fetch google/t5gemma-9b-9b-ul2 separately or re-run with --bundle-text-encoder)")
-
-
-if __name__ == "__main__":
-    main()

scripts/checkpoint_conversion/push_magi_human_to_hf.py

@@ -1,117 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Push a converted daVinci-MagiHuman Diffusers-format directory to the Hub.
-
-This is a thin wrapper around `huggingface_hub.create_repo` + `upload_folder`,
-dedicated to the MagiHuman upload flow. It does NOT modify `create_hf_repo.py`
-(which is LTX-2-oriented and rewrites component weights inside an existing
-Diffusers repo).
-
-Example (one-shot per variant):
-    python scripts/checkpoint_conversion/push_magi_human_to_hf.py \\
-        --local-dir converted_weights/magi_human_base \\
-        --repo-id FastVideo/MagiHuman-Base-Diffusers \\
-        --public
-
-    python scripts/checkpoint_conversion/push_magi_human_to_hf.py \\
-        --local-dir converted_weights/magi_human_distill \\
-        --repo-id FastVideo/MagiHuman-Distilled-Diffusers \\
-        --public
-
-After upload, the local directory can be deleted — the HF repo is the
-source of truth. `VideoGenerator.from_pretrained("FastVideo/...")` pulls
-shards on demand.
-"""
-from __future__ import annotations
-
-import argparse
-import os
-import sys
-from pathlib import Path
-
-from huggingface_hub import HfApi, create_repo, upload_folder
-
-
-def _validate_local_dir(local_dir: Path) -> None:
-    required = ["model_index.json", "transformer"]
-    missing = [r for r in required if not (local_dir / r).exists()]
-    if missing:
-        sys.exit(
-            f"Error: {local_dir} is missing {missing}. Run "
-            f"convert_magi_human_to_diffusers.py first."
-        )
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
-    parser.add_argument("--local-dir", required=True, help="Path to the converted Diffusers directory.")
-    parser.add_argument("--repo-id", required=True, help="Target HF repo id, e.g. FastVideo/MagiHuman-Base-Diffusers.")
-    parser.add_argument(
-        "--public",
-        action="store_true",
-        help="Create the repo as public (default: private). Mutually exclusive with --private.",
-    )
-    parser.add_argument(
-        "--private",
-        action="store_true",
-        help="Create the repo as private. Default when neither --public nor --private is set.",
-    )
-    parser.add_argument(
-        "--commit-message",
-        default="Initial upload of daVinci-MagiHuman Diffusers-format conversion.",
-    )
-    parser.add_argument(
-        "--dry-run",
-        action="store_true",
-        help="Describe what would happen without creating a repo or uploading.",
-    )
-    args = parser.parse_args()
-
-    if args.public and args.private:
-        sys.exit("Error: --public and --private are mutually exclusive.")
-    private = args.private or not args.public
-
-    local_dir = Path(args.local_dir).resolve()
-    _validate_local_dir(local_dir)
-
-    token = (
-        os.environ.get("HF_TOKEN")
-        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
-        or os.environ.get("HF_API_KEY")
-    )
-    if not token:
-        sys.exit(
-            "Error: no HF token in env (set HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY)."
-        )
-
-    api = HfApi()
-    me = api.whoami(token=token)
-    print(f"token user: {me.get('name')}")
-    print(f"source:     {local_dir}")
-    print(f"target:     {args.repo_id}")
-    print(f"visibility: {'private' if private else 'public'}")
-    if args.dry_run:
-        print("(dry run — not creating or uploading)")
-        return
-
-    print(f"-> create_repo (exist_ok=True)")
-    create_repo(
-        repo_id=args.repo_id,
-        token=token,
-        private=private,
-        exist_ok=True,
-        repo_type="model",
-    )
-
-    print(f"-> upload_folder (this can take a while for 30 GB)")
-    upload_folder(
-        repo_id=args.repo_id,
-        folder_path=str(local_dir),
-        token=token,
-        commit_message=args.commit_message,
-        repo_type="model",
-    )
-    print(f"Done. https://huggingface.co/{args.repo_id}")
-
-
-if __name__ == "__main__":
-    main()

tests/local_tests/helpers/magi_human_upstream.py

@@ -1,334 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Stubs to make upstream daVinci-MagiHuman code importable in-process.
-
-The upstream DiT (daVinci-MagiHuman/inference/model/dit/dit_module.py) hard-
-imports SandAI's internal `magi_compiler` + a distributed-runtime init
-that requires `torchrun`. Neither is available in a single-process
-parity test. This module installs the minimum stubs to let the upstream
-DiT load and run on a single GPU with cp_world_size == 1 (which makes
-Ulysses's scatter/gather a no-op).
-
-Use:
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs, load_upstream_dit,
-    )
-    install_stubs()
-    model = load_upstream_dit(base_shard_dir, device=torch.device("cuda"))
-"""
-from __future__ import annotations
-
-import importlib
-import sys
-import types
-from pathlib import Path
-from typing import Callable
-
-
-# ---------------------------------------------------------------------------
-# magi_compiler stubs — identity decorators.
-# ---------------------------------------------------------------------------
-
-
-def _install_magi_compiler_stub() -> None:
-    """Stub magi_compiler + register the `torch.ops.infra.*` ops upstream
-    calls via `torch.ops`.
-
-    Upstream code decorates plain Python fns with
-    `@magi_register_custom_op(name="infra::flash_attn_func", ...)` and
-    then calls them as `torch.ops.infra.flash_attn_func(...)`. Our stub
-    decorator has to both (a) preserve the decorated fn for direct call
-    sites and (b) register the fn under the advertised torch.ops
-    namespace so `torch.ops.infra.*` resolves.
-
-    For parity testing we route `infra::flash_attn_func` through
-    `F.scaled_dot_product_attention`, matching the FastVideo DiT's
-    kernel choice so drift measured in this test is architectural,
-    not kernel-dependent.
-    """
-    if "magi_compiler" in sys.modules:
-        return
-
-    import torch
-    import torch.nn.functional as F
-
-    pkg = types.ModuleType("magi_compiler")
-
-    def magi_compile(config_patch=None):
-        def decorator(cls_or_fn):
-            return cls_or_fn
-        return decorator
-
-    # Create one Library per (namespace, schema) pair. Track by namespace
-    # so we don't define the same op twice on re-import.
-    _libs: dict[str, torch.library.Library] = {}
-    _defined: set[tuple[str, str]] = set()
-
-    def _sdpa_flash_attn_func(q, k, v):
-        # Upstream shape: [batch=1, L, H, D]. SDPA expects [B, H, L, D]
-        # and no native GQA; expand K/V to match Q heads.
-        num_heads_q = q.shape[2]
-        num_heads_kv = k.shape[2]
-        if num_heads_q != num_heads_kv:
-            assert num_heads_q % num_heads_kv == 0
-            repeat = num_heads_q // num_heads_kv
-            k = k.repeat_interleave(repeat, dim=2)
-            v = v.repeat_interleave(repeat, dim=2)
-        q2 = q.transpose(1, 2).contiguous()
-        k2 = k.transpose(1, 2).contiguous()
-        v2 = v.transpose(1, 2).contiguous()
-        out = F.scaled_dot_product_attention(q2, k2, v2)
-        return out.transpose(1, 2).contiguous()
-
-    def _sdpa_segments(q, k, v, q_ranges, k_ranges):
-        # Upstream flex op shape: [L, H, D]. FFA accumulates each block's
-        # independently normalized attention output into the destination query
-        # slice. This SDPA fallback mirrors the accumulator semantics for
-        # SR-1080p parity tests without requiring SandAI's MagiAttention wheel.
-        out = torch.zeros(
-            q.shape[0],
-            q.shape[1],
-            q.shape[2],
-            dtype=q.dtype,
-            device=q.device,
-        )
-        num_heads_q = q.shape[1]
-        num_heads_kv = k.shape[1]
-        for q_range, k_range in zip(q_ranges.tolist(), k_ranges.tolist()):
-            qs, qe = int(q_range[0]), int(q_range[1])
-            ks, ke = int(k_range[0]), int(k_range[1])
-            q_block = q[qs:qe]
-            k_block = k[ks:ke]
-            v_block = v[ks:ke]
-            if num_heads_q != num_heads_kv:
-                assert num_heads_q % num_heads_kv == 0
-                repeat = num_heads_q // num_heads_kv
-                k_block = k_block.repeat_interleave(repeat, dim=1)
-                v_block = v_block.repeat_interleave(repeat, dim=1)
-            block_out = F.scaled_dot_product_attention(
-                q_block.transpose(0, 1).unsqueeze(0).contiguous(),
-                k_block.transpose(0, 1).unsqueeze(0).contiguous(),
-                v_block.transpose(0, 1).unsqueeze(0).contiguous(),
-            )
-            out[qs:qe] += block_out.squeeze(0).transpose(0, 1).contiguous()
-        lse = torch.empty((q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device)
-        return out, lse
-
-    def magi_register_custom_op(name=None, mutates_args=(), infer_output_meta_fn=None, is_subgraph_boundary=False, **kwargs):
-        def decorator(fn):
-            if not name:
-                return fn
-            namespace, op_name = name.split("::", 1)
-            if namespace not in _libs:
-                _libs[namespace] = torch.library.Library(namespace, "FRAGMENT")
-            if (namespace, op_name) in _defined:
-                # Already registered in a previous test run — reuse.
-                return fn
-            # Route known ops through SDPA; leave unknown ones as direct fn.
-            returns = "(Tensor, Tensor)" if op_name == "flex_flash_attn_func" else "Tensor"
-            schema_name = f"{op_name}({_infer_schema(fn)}) -> {returns}"
-            try:
-                _libs[namespace].define(schema_name)
-            except Exception:
-                pass
-            if op_name == "flash_attn_func":
-                torch.library.impl(
-                    _libs[namespace], op_name, "CUDA"
-                )(_sdpa_flash_attn_func)
-                torch.library.impl(
-                    _libs[namespace], op_name, "CPU"
-                )(_sdpa_flash_attn_func)
-            elif op_name == "flex_flash_attn_func":
-                torch.library.impl(
-                    _libs[namespace], op_name, "CUDA"
-                )(_sdpa_segments)
-            else:
-                # For ops we don't care about (compile-only wrappers), the
-                # Python fn path inside the module body is used directly —
-                # we just need `torch.ops.<ns>.<op>` to exist so module-
-                # load-time attribute lookups succeed.
-                torch.library.impl(
-                    _libs[namespace], op_name, "CUDA"
-                )(fn)
-            _defined.add((namespace, op_name))
-            return fn
-        return decorator
-
-    pkg.magi_compile = magi_compile
-    sys.modules["magi_compiler"] = pkg
-
-    api = types.ModuleType("magi_compiler.api")
-    api.magi_register_custom_op = magi_register_custom_op
-    sys.modules["magi_compiler.api"] = api
-    pkg.api = api
-
-    config_mod = types.ModuleType("magi_compiler.config")
-
-    class CompileConfig:
-        class offload_config:  # pragma: no cover - pass-through
-            gpu_resident_weight_ratio = 1.0
-    config_mod.CompileConfig = CompileConfig
-    sys.modules["magi_compiler.config"] = config_mod
-    pkg.config = config_mod
-
-
-def _infer_schema(fn) -> str:
-    """Return a minimal torch.library schema string for the given fn.
-
-    For our stub we just need *something* parseable; all real ops we
-    care about take `(q, k, v)` or `(q, k, v, q_ranges, k_ranges)` or
-    variants. Use generic `Tensor a, Tensor b, ...` arg names.
-    """
-    import inspect
-    sig = inspect.signature(fn)
-    parts = []
-    for i, name in enumerate(sig.parameters):
-        arg_name = name if name.isidentifier() else f"a{i}"
-        parts.append(f"Tensor {arg_name}")
-    return ", ".join(parts)
-
-
-# ---------------------------------------------------------------------------
-# distributed / CP stubs — single-GPU, cp_world_size == 1.
-# ---------------------------------------------------------------------------
-
-
-def _install_distributed_stubs() -> None:
-    """Monkey-patch upstream distributed + parallelism modules for cp=1."""
-    # inference.infra.distributed.*
-    # The real module requires NCCL / parallel_state to be initialized
-    # from torchrun; here we short-circuit the handful of getters the DiT
-    # actually calls.
-    import inference.infra.distributed as dist_mod
-
-    dist_mod.get_cp_world_size = lambda: 1
-    dist_mod.get_cp_group = lambda: None
-    dist_mod.get_cp_rank = lambda: 0
-    dist_mod.get_tp_rank = lambda: 0
-    dist_mod.get_pp_rank = lambda: 0
-
-    # inference.infra.parallelism.*
-    # At cp_world_size=1, scatter/gather are trivially no-ops.
-    import inference.infra.parallelism.gather_scatter_primitive as gs
-
-    def _scatter_noop(x, cp_split_sizes, group=None):
-        return x
-
-    def _gather_noop(x, cp_split_sizes, group=None):
-        return x
-
-    gs.scatter_to_context_parallel_region = _scatter_noop
-    gs.gather_from_context_parallel_region = _gather_noop
-
-    # Re-import ulysses_scheduler with patched scatter/gather in place.
-    import inference.infra.parallelism.ulysses_scheduler as us
-    us.scatter_to_context_parallel_region = _scatter_noop
-    us.gather_from_context_parallel_region = _gather_noop
-    us.get_cp_world_size = lambda: 1
-    us.get_cp_group = lambda: None
-
-    # all-to-all primitives used by flash_attn_with_cp. At cp=1 they are
-    # entered only as a no-op path (the `if cp_world_size > 1` branch is
-    # skipped), so no stubs needed there.
-
-
-def install_stubs() -> None:
-    """Install all stubs. Idempotent."""
-    _install_magi_compiler_stub()
-    repo_root = Path(__file__).resolve().parents[3]
-    upstream = repo_root / "daVinci-MagiHuman"
-    path_s = str(upstream)
-    if path_s not in sys.path:
-        sys.path.insert(0, path_s)
-    # Reload inference.* after sys.path mutation so it picks up the real
-    # upstream package (not a stale one).
-    for name in list(sys.modules):
-        if name == "inference" or name.startswith("inference."):
-            del sys.modules[name]
-    import inference  # noqa: F401
-    _install_distributed_stubs()
-
-
-# ---------------------------------------------------------------------------
-# Upstream DiTModel loader — instantiate + load base shards.
-# ---------------------------------------------------------------------------
-
-
-def _base_arch_dict() -> dict:
-    """Return the upstream `ModelConfig`-equivalent dict for the base variant.
-    Matches `inference/common/config.py::ModelConfig` defaults for base.
-    """
-    import torch
-    return dict(
-        num_layers=40,
-        hidden_size=5120,
-        head_dim=128,
-        num_query_groups=8,
-        video_in_channels=48 * 4,
-        audio_in_channels=64,
-        text_in_channels=3584,
-        checkpoint_qk_layernorm_rope=False,
-        params_dtype=torch.float32,
-        tread_config=dict(
-            selection_rate=0.5, start_layer_idx=2, end_layer_idx=25,
-        ),
-        mm_layers=[0, 1, 2, 3, 36, 37, 38, 39],
-        local_attn_layers=[],
-        enable_attn_gating=True,
-        activation_type="swiglu7",
-        gelu7_layers=[0, 1, 2, 3],
-        # derived
-        num_heads_q=40,
-        num_heads_kv=8,
-        post_norm_layers=[],
-    )
-
-
-def load_upstream_dit(base_shard_dir, device=None, dtype=None, local_attn_layers=None):
-    """Instantiate upstream `DiTModel` and load the base shards into it.
-
-    Args:
-        base_shard_dir: path to `base/` (contains `model-0000*-of-00007.safetensors`
-            and `model.safetensors.index.json`).
-        device: torch device (default cuda if available).
-        dtype: dtype cast (default: leave checkpoint dtypes as-is).
-
-    Returns:
-        An upstream `DiTModel` in `.eval()` mode with weights loaded.
-    """
-    import glob
-    import json
-    import types as _types
-
-    import torch
-    from safetensors.torch import load_file
-
-    from inference.common.config import ModelConfig  # upstream pydantic class
-    from inference.model.dit.dit_module import DiTModel
-
-    arch_dict = _base_arch_dict()
-    if local_attn_layers is not None:
-        arch_dict["local_attn_layers"] = list(local_attn_layers)
-    # ModelConfig is a pydantic BaseModel — build via kwargs.
-    model_config = ModelConfig(**arch_dict)
-
-    model = DiTModel(model_config=model_config)
-
-    # Load all base shards into a single state dict.
-    base_shard_dir = Path(base_shard_dir)
-    shard_paths = sorted(base_shard_dir.glob("*.safetensors"))
-    state = {}
-    for p in shard_paths:
-        state.update(load_file(str(p)))
-
-    missing, unexpected = model.load_state_dict(state, strict=False)
-    if missing:
-        raise RuntimeError(f"Upstream DiT missing {len(missing)} keys: {missing[:5]}")
-    if unexpected:
-        raise RuntimeError(f"Upstream DiT unexpected {len(unexpected)} keys: {unexpected[:5]}")
-
-    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device=device)
-    if dtype is not None:
-        model = model.to(dtype=dtype)
-    model.eval()
-    return model

tests/local_tests/magi-human.md

@@ -1,625 +0,0 @@
-# Local daVinci-MagiHuman Tests
-
-End-to-end parity tests for the daVinci-MagiHuman joint text-to-audio-video
-pipeline. MagiHuman is a 15B-parameter DiT that denoises video and audio
-latents in a single loop, producing synchronized video and audio from a text
-prompt. The video path uses the Wan 2.2 TI2V-5B VAE (decoder only), the audio
-path uses the Stable Audio Open 1.0 `OobleckVAE` (shared with the standalone
-Stable Audio pipeline), and text conditioning comes from a T5-Gemma 9B UL2
-encoder. The base variant runs 32-step FlowUniPC with CFG=2; the distill
-variant runs 8 steps with CFG=1. Reference implementation:
-[GAIR-NLP/daVinci-MagiHuman](https://github.com/GAIR-NLP/daVinci-MagiHuman).
-These tests compare FastVideo against the published weights and the upstream
-reference, so they're skipped in CI and run locally on a single GPU.
-
-## Setup
-
-### 1. Hugging Face access
-
-MagiHuman depends on four gated repos. Accept the terms at each URL once, then
-export your token:
-
-| Repo | Terms URL |
-|---|---|
-| `GAIR/daVinci-MagiHuman` | https://huggingface.co/GAIR/daVinci-MagiHuman |
-| `google/t5gemma-9b-9b-ul2` | https://huggingface.co/google/t5gemma-9b-9b-ul2 |
-| `stabilityai/stable-audio-open-1.0` | https://huggingface.co/stabilityai/stable-audio-open-1.0 |
-| `Wan-AI/Wan2.2-TI2V-5B` | https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B |
-
-```bash
-export HF_TOKEN=hf_...
-# any of HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY works
-```
-
-The pipeline's `_ensure_hf_token_env` helper (in
-`fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py`) aliases all
-three names to `HF_TOKEN` and `HUGGINGFACE_HUB_TOKEN` at load time, so
-whichever variable you set will be picked up. Tests skip cleanly with a
-helpful message if no token is found.
-
-### 2. Optional inference dependencies
-
-The pipeline uses the default FastVideo attention backend. No extra packages
-are required for basic inference. If you want the T5-Gemma wrapper to use
-PyTorch SDPA instead of Flash Attention, set:
-
-```bash
-export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-```
-
-The `T5GemmaEncoderModel` wrapper in
-`fastvideo/models/encoders/t5gemma.py` reads this variable and patches
-`model.config.attn_implementation` accordingly before the first forward pass.
-
-### 3. Clone the upstream reference repo
-
-The DiT parity test (`test_magi_human_parity.py`) and the pipeline parity test
-(`test_magi_human_pipeline_parity.py`) import directly from the upstream
-`daVinci-MagiHuman` package. Clone it under the repo root and add it to your
-personal ignore list:
-
-```bash
-cd <FastVideo repo root>
-git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git
-echo "/daVinci-MagiHuman/" >> .git/info/exclude   # personal ignore
-```
-
-Tests that need the clone skip cleanly if the directory is absent. The VAE
-parity tests and the smoke test do not need the upstream clone.
-
-### 4. Convert weights
-
-Run the conversion script once to produce a Diffusers-layout checkpoint. The
-`--bundle-vae`, `--bundle-audio-vae`, and `--bundle-text-encoder` flags copy
-the Wan VAE, Oobleck audio VAE, and T5-Gemma encoder into the output directory
-so the pipeline can load everything from a single path:
-
-```bash
-python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \
-    --source GAIR/daVinci-MagiHuman \
-    --output converted_weights/magi_human_base \
-    --bundle-vae \
-    --bundle-audio-vae \
-    --bundle-text-encoder
-```
-
-Disk budget: roughly 30 GB for the base checkpoint. The distill variant is a
-similar size; add `--cast-bf16` to halve the transformer shards if storage is
-tight.
-
-The tests look for the converted path in `MAGI_HUMAN_DIFFUSERS_PATH` (see
-§8 Troubleshooting). If that variable is unset, they fall back to
-`converted_weights/magi_human_base` relative to the repo root.
-
-### 5. (Optional) Pre-warm the model cache
-
-The first parity-test run downloads the T5-Gemma encoder (~18 GB), the Wan VAE
-(~2 GB), and the Stable Audio Open VAE (~1 GB) if they aren't already cached.
-To avoid the download blocking your first test run, fetch them ahead of time:
-
-```bash
-python -c "
-from huggingface_hub import snapshot_download
-snapshot_download('google/t5gemma-9b-9b-ul2')
-snapshot_download('Wan-AI/Wan2.2-TI2V-5B')
-snapshot_download('stabilityai/stable-audio-open-1.0')
-"
-```
-
-## Running the tests
-
-All MagiHuman local tests in one shot:
-
-```bash
-pytest tests/local_tests/magi_human/test_magi_human_parity.py \
-       tests/local_tests/magi_human/test_magi_human_t5gemma_parity.py \
-       tests/local_tests/magi_human/test_magi_human_sa_audio_parity.py \
-       tests/local_tests/magi_human/test_magi_human_sa_audio_official_parity.py \
-       tests/local_tests/magi_human/test_magi_human_vae_parity.py \
-       tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py \
-       tests/local_tests/magi_human/test_magi_human_pipeline_parity.py \
-       fastvideo/tests/ssim/test_magi_human_similarity.py \
-       -v -s
-```
-
-Add `-s` to print per-test diff numbers (shape / abs_mean / max diff / drift).
-
-### What each test covers
-
-**`test_magi_human_parity.py`** — DiT component parity. Loads the
-`MagiHumanTransformer3DModel` from the converted checkpoint and the upstream
-reference DiT from the `daVinci-MagiHuman` clone, feeds identical latent
-inputs, and checks that output tensors match within tolerance. Requires both
-the upstream clone and the converted weights.
-
-**`test_magi_human_t5gemma_parity.py`** — T5-Gemma encoder wrapper parity.
-Compares `fastvideo.models.encoders.t5gemma.T5GemmaEncoderModel` against a
-direct HuggingFace `T5GemmaEncoderModel.from_pretrained` call on the same
-checkpoint. Verifies that the FastVideo wrapper's lazy-load path and
-`named_parameters` exclusion don't alter the encoder's output embeddings.
-
-**`test_magi_human_sa_audio_parity.py`** — Stable Audio Open VAE wrapper
-parity. Compares the FastVideo `OobleckVAE` (shared with the standalone Stable
-Audio pipeline) against HuggingFace Diffusers' `AutoencoderOobleck` on the
-`stabilityai/stable-audio-open-1.0` weights. Encode + decode + round-trip;
-expected to be bit-identical in fp32.
-
-**`test_magi_human_sa_audio_official_parity.py`** — Stable Audio Open VAE
-parity vs the official daVinci-MagiHuman integration layer. Compares FastVideo's
-`SAAudioVAEModel` against the upstream `SAAudioFeatureExtractor.decode()` path
-from the `daVinci-MagiHuman` clone. Catches drift between FastVideo's full SA
-wrapper and the official repo's custom Stable-Audio module. Requires the upstream
-clone and the `stabilityai/stable-audio-open-1.0` gated repo. Expected to be
-bit-exact (diff=0) in fp32.
-
-**`test_magi_human_vae_parity.py`** — Wan video VAE parity. Compares the
-FastVideo Wan VAE decoder against the upstream `Wan2_2_VAE` on
-`Wan-AI/Wan2.2-TI2V-5B` weights. Decoder-only path (MagiHuman never encodes
-video at inference time).
-
-**`test_magi_human_pipeline_smoke.py`** — Preflight and smoke. Imports the
-pipeline, resolves the registry entries (`magi_human_base`,
-`magi_human_distill`), checks preset wiring, and verifies the pipeline can
-instantiate without a GPU. CPU-only; no model weights required beyond the
-converted path.
-
-**`test_magi_human_pipeline_parity.py`** — End-to-end joint AV latent parity.
-Runs a short denoising loop through the full pipeline and compares the final
-video and audio latents against the upstream reference pipeline. Requires the
-upstream clone, the converted weights, and a GPU.
-
-**`test_magi_human_similarity.py`** — Video SSIM regression (CI-runnable).
-Generates a short clip from a fixed prompt and seed, then compares frame-level
-SSIM against reference videos stored in the `FastVideo/ssim-reference-videos`
-HF dataset. The test skips cleanly until reference videos are seeded (see §7
-Open questions).
-
-### Reproducing a single test
-
-Each test file is independent. Run one:
-
-```bash
-pytest tests/local_tests/magi_human/test_magi_human_pipeline_parity.py -v -s
-```
-
-## Phase 11 status
-
-Branch tip `eeef855b` (rebased onto `origin/main` `c77a76c6`), Wave 1+4 changes applied (uncommitted working tree), NVIDIA B200. Wave 2-3 numerical-alignment investigation completed 2026-05-01; see §Numerical-alignment investigation below.
-
-| Test | Status | Diff numbers | Notes |
-|---|---|---|---|
-| `tests/local_tests/magi_human/test_magi_human_t5gemma_parity.py::test_magi_human_t5gemma_wrapper_parity` | PASS | exact (`assert_close(atol=1e-3, rtol=1e-3)`) | gated repo, requires HF token |
-| `tests/local_tests/magi_human/test_magi_human_parity.py::test_magi_human_dit_parity` | FAIL | video diff_max=0.057, diff_mean=0.008; audio diff_max=0.034, diff_mean=0.008; text exact (diff_max=0) | Tightened to `atol=0.03, rtol=0.01` (Wave 1). Bf16-noise-floor; per-layer drift ~1e-3 accumulates over 40 layers. Root cause of OQ-6 compounding. See §Numerical-alignment investigation. |
-| `tests/local_tests/magi_human/test_magi_human_vae_parity.py::test_magi_human_vae_decode_parity` | PASS | diff_max=8e-4, diff_mean=4.9e-5 | Wan VAE. Deferred to `atol=1e-3, rtol=1e-3` per OQ-7 (Wave 4). Tighten to `atol=1e-4` once Wan VAE op-order fix lands. |
-| `tests/local_tests/magi_human/test_magi_human_sa_audio_parity.py::test_magi_human_sa_audio_vae_decode_parity` | PASS | exact (`assert_close(atol=1e-5, rtol=1e-5)`, machine epsilon) | gated repo, requires HF token; uses main's shared `OobleckVAE` + `SAAudioVAEModel` wrapper |
-| `tests/local_tests/magi_human/test_magi_human_sa_audio_official_parity.py::test_magi_human_sa_audio_official_decode_parity` | PASS | `atol=1e-5, rtol=1e-5`, diff_max=0, diff_mean=0 (bit-exact) | Wave 7. Compares FV `SAAudioVAEModel` vs upstream `SAAudioFeatureExtractor.decode()`. Confirms OQ-6 is NOT in audio VAE. Requires upstream clone + gated SA repo. |
-| `tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py::test_magi_human_typed_surface_preflight` | PASS | CPU-only key/preset checks, exact key set equality, 331 keys | no skip conditions met locally |
-| `tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py::test_magi_human_pipeline_smoke` | PASS | shape-only; 2 inference steps, output shape `[B,C,T,H,W]` validated | wallclock ~50s |
-| `tests/local_tests/magi_human/test_magi_human_pipeline_parity.py::test_magi_human_pipeline_latent_parity` | FAIL | video: diff_max=6.69, diff_mean=0.47; audio: diff_max=3.45, diff_mean=1.01 | Wave 7.5: now uses real preset prompts via T5-Gemma. Wave 8 production fixes don't move parity numbers (both sides use same encoder). Residual drift is bf16+CFG amplification floor; tracked as OQ-6 RESOLVED-PRODUCTION. |
-| `fastvideo/tests/ssim/test_magi_human_similarity.py::test_magi_human_base_inference_similarity` | DEFERRED | n/a | Reference videos not yet seeded to `FastVideo/ssim-reference-videos` HF repo; tracked as OQ-2. Requires Modal L40S seeding via `seed-ssim-references` skill. |
-| _(debug)_ | INFO | Per-side layer logs: `/tmp/opencode/magi_dit_up_layers.log`, `/tmp/opencode/magi_dit_fv_layers.log` | Added in Wave 1 to `_debug_magi_human_block_parity.py`. See `add-model-trace` skill at `~/.config/opencode/skill/add-model-trace/`. |
-| `fastvideo/tests/hooks/test_activation_trace.py::*` | PASS | 6 tests covering off/on/filter/stats/step-filter/cleanup | Wave 9 activation trace infrastructure |
-
-_Last verified: 2026-05-01 (Wave 10 dtype refactor on rebased branch @ 3caeaad1; tests now under `tests/local_tests/magi_human/`)_
-
-## Design notes
-
-### Cross-variant shared component lazy-loading
-
-The four MagiHuman variants (`base`, `distill`, `sr_540p`, `sr_1080p`) ship four
-shared components — Wan 2.2 TI2V-5B VAE, T5-Gemma encoder + tokenizer, and
-Stable Audio Open 1.0 VAE — that together account for ~25 GB of weights. To
-avoid duplicating these in every converted variant repo,
-`MagiHumanPipeline.load_modules` lazy-loads all four from their canonical
-upstream HF repos at first build time:
-
-| Component | Upstream HF repo | Gated? |
-|---|---|---|
-| `text_encoder`, `tokenizer` | `google/t5gemma-9b-9b-ul2` | yes |
-| `audio_vae` | `stabilityai/stable-audio-open-1.0` | yes |
-| `vae` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | no |
-
-A converted MagiHuman variant repo therefore only needs to ship
-`transformer/`, `scheduler/`, and `model_index.json` (~5 GB for base bf16,
-~30 GB for distill bf16). Bundling the shared components is still supported
-via the conversion script's `--bundle-vae` / `--bundle-audio-vae` /
-`--bundle-text-encoder` flags but is no longer the default.
-
-The verification helper in `fastvideo/utils.py:verify_model_config_and_directory`
-treats the contents of `model_index.json` as authoritative for which component
-subfolders must exist locally; pipelines that emit a minimal `model_index.json`
-(omitting `vae`, `text_encoder`, etc.) pass verification, while pipelines that
-DO declare a component must still ship its subfolder.
-
-### Umbrella-repo subfolder syntax
-
-`fastvideo/utils.py:maybe_download_model` recognises an "umbrella" repo layout
-where a single HF repo holds multiple variants under sibling subfolders:
-
-```
-FastVideo/MagiHuman-Diffusers/
-├── base/{model_index.json, transformer/, scheduler/}
-├── distill/{...}
-├── sr_540p/{...}
-└── sr_1080p/{...}
-```
-
-Pass `org/repo/subfolder` as the model path; the loader downloads only that
-subfolder's blobs and points the pipeline at the local subfolder snapshot:
-
-```python
-generator = VideoGenerator.from_pretrained("FastVideo/MagiHuman-Diffusers/base")
-```
-
-The detection heuristic is purely structural: HF Hub repo ids are always two
-slash-separated components (`org/name`); a path with three or more components
-that does not exist locally and is not posix-absolute or relative-prefixed is
-treated as an umbrella reference. Backwards-compatible with the existing
-single-repo-per-variant layout (`FastVideo/MagiHuman-Base-Diffusers`).
-
-### T5-Gemma lazy-load exception
-
-`fastvideo/models/encoders/gemma.py:10` establishes the FastVideo precedent for
-gated foundation-model encoders: the HF model class
-(`Gemma3ForConditionalGeneration`) is imported at module top-level, and the
-actual weights are loaded lazily via `from_pretrained` inside a property or
-method.
-
-`fastvideo/models/encoders/t5gemma.py:60` follows the same pattern but is
-strictly more conservative: the HF class (`T5GemmaEncoderModel`) is imported
-inside `_build_t5gemma_model` rather than at module top-level. This avoids an
-import-time failure if `transformers.models.t5gemma` isn't available in the
-environment. The `named_parameters` override on the same class hides the
-upstream encoder from FastVideo's weight loader so the converted repo directory
-isn't scanned for T5-Gemma shards.
-
-This is the established FastVideo pattern for gated foundation-model encoders
-not yet ported natively. It is not a workaround; it's the documented approach.
-
-**Native T5-Gemma port — TRACKED FOLLOW-UP.** A future native port is
-desirable for full Phase 11 hard-rule compliance (no HF model-class imports in
-production runtime code). Scope estimate is multi-week: Gemma decoder blocks,
-T5 encoder cross-attention, RMS norm, RoPE, and tokenizer wiring all need
-native FastVideo implementations. This is tracked here until claimed by a
-follow-up PR.
-
-### Audio quality regression deferral
-
-`tests/local_tests/stable-audio.md` sets the precedent: the Stable Audio Open
-1.0 port ships local parity tests, a smoke test, and self-consistency checks
-for inpainting and audio-to-audio variation, with no `fastvideo/tests/audio/`
-quality regression test.
-
-MagiHuman's audio path is covered by `test_magi_human_pipeline_parity.py`
-(joint AV latent comparison against the upstream reference) and the basic
-example mp4 spot-check (`examples/inference/basic/basic_magi_human.py`). A
-mel-spectrogram L1 or multi-resolution STFT regression test is listed as a
-follow-up if audio drift becomes a concern in practice.
-
-### Pipeline parity tolerance budget (1-step / CFG=2)
-
-Drift is dominated by CFG amplification of single-DiT bf16 mismatch. The
-single-DiT diff_mean is ~0.008 (per `test_magi_human_dit_parity`); CFG mixes
-`v = v_uncond + 5*(v_cond - v_uncond)`, so independent bf16 errors in
-cond/uncond paths compound by ~5x, giving an expected pipeline diff_mean of
-~0.04. Observed is 0.069. `diff_max` is the noisiest statistic for bf16+CFG
-(a single fma quantization can blow it up); `atol=0.40` accommodates that.
-
-Two ratio guards catch real structural bugs:
-
-- **`abs_mean` drift < 1%** (gross-bug catcher: scheduler state leak, dropped
-  modality, CFG sign flip)
-- **`diff_mean / ref_abs` < 4%** (systematic per-element bias guard)
-
-All three guards currently pass with margin: video abs_mean rel=0.36%, audio
-abs_mean rel=0.33%; video diff_mean/ref=3.07%, audio diff_mean/ref=2.66%.
-
-The test uses `num_inference_steps=1, cfg_number=2, guidance=5.0`. Per Oracle
-analysis in this PR's review notes, this is expected bf16+CFG behavior, not a
-structural bug.
-
-## Numerical-alignment investigation (2026-05-01)
-
-Wave 2-3 investigation into why the 4-step pipeline parity fails and whether the
-DiT parity failure at `atol=0.03` indicates a real bug.
-
-### Methodology
-
-TDD-style: tighten tolerances to surface real drift, run, drill into the
-largest contributor, bisect to confirm pre-existence, then rule out hypotheses
-one by one.
-
-1. **Wave 1 (bug-surfacing changes):** Tightened DiT parity from `atol=0.1` to
-   `atol=0.03, rtol=0.01`. Tightened Wan VAE parity from `atol=5e-2` to
-   `atol=1e-4` (later deferred to `atol=1e-3` per OQ-7). Bumped pipeline parity
-   `num_inference_steps` from 1 to 4. Fixed `_find_base_shard_dir` with
-   `snapshot_download` fallback (resolves OQ-4). Added per-side layer log files
-   to `_debug_magi_human_block_parity.py`. Created new `add-model-trace` skill
-   in user dotfiles.
-
-2. **Wave 2 (run and measure):** DiT parity fails at new `atol=0.03`
-   (diff_max=0.057, diff_mean=0.008). Wan VAE parity fails at `atol=1e-4`
-   (diff_max=8e-4). 4-step pipeline parity fails with video diff_mean=1.30 vs
-   1-step 0.069, a ratio of 18.85x (expected ~4x linear). Per-block drift never
-   exceeds 0.5% threshold; cumulative peaks at MM layers (blocks 0-3 and 36-39,
-   matching `mm_layers=[0,1,2,3,36,37,38,39]`).
-
-3. **Wave 3 (drill and bisect):** Tested PackedExpertLinear hypothesis via A/B
-   patch. Bisected compounding bug to original commit. Drilled into Block[02]
-   MM-layer MLP `down_proj` amplification. Verified expert chunk ordering
-   bit-exact.
-
-### Key findings
-
-| Finding | Result | Evidence |
-|---|---|---|
-| PackedExpertLinear routing bug | **REJECTED** | A/B with `MAGI_DEBUG_PATCH_LINEAR=1` (mirrors upstream `_BF16ComputeLinear`) showed zero change in drift |
-| Wave 1 commits caused compounding | **REJECTED** | `git revert` bisect: 4-step diff_mean=1.20 with reverts vs 1.30 with Wave 1; bug pre-exists in commit 620aaf41 |
-| Expert chunk ordering mismatch | **REJECTED** | Direct FV `PackedExpertLinear` vs upstream `NativeMoELinear` test: diff=0 (bit-exact) |
-| Wan VAE op-order drift | **CONFIRMED** | FV uses `z * std + mean`; upstream uses `z / (1/std) + mean`. Bitwise non-equivalent. Shared Wan-family bug (OQ-7). |
-| MM-layer MLP `down_proj` amplification | **NORMAL** | Block[02] input drift 0.0005 → output drift 0.022 = 44x amplification. Normal sensitivity for a 15360x20480 matrix; not a routing bug. |
-| Per-forward DiT drift | **BF16 NOISE FLOOR** | diff_max=0.057 from cumulative ~1e-3 per-layer over 40 layers. Consistent with random-walk bf16 accumulation. |
-
-### Root-cause hypothesis
-
-Per-forward DiT drift is bf16 noise, not a structural bug. Diffusion sampling
-amplifies per-step bf16 perturbations geometrically over the denoise loop (a
-known ill-conditioned-ODE phenomenon). The 18.85x compounding ratio at 4 steps
-vs the expected 4x linear ratio confirms geometric amplification. The "blurry
-abstract" output at 32 steps (OQ-5) is the downstream symptom.
-
-Wave 3 ruled out all discrete implementation bugs: PackedExpertLinear routing,
-expert chunk ordering, and the conversion script are all bit-exact. The
-remaining candidates are dtype boundary mismatches around sensitive MM-layer ops
-(pre-norm, attention, MLP activation) where upstream may cast to fp32 and FV
-stays in bf16.
-
-### Wave 7 (2026-05-01): CFG + negative prompt investigation
-
-Findings:
-- **CFG math identical**: FV `v = uncond + g * (cond - uncond)` matches upstream at `denoising.py:178-181` ↔ `video_generate.py:426,456-457`. Video has `t > 500` cutoff (`5.0 → 2.0`); audio has none. Both sides apply the same formula.
-- **Scheduler args identical for T2AV base path**: `step(model_output, t, sample, return_dict=False)[0]`. Audio-skip modes (`is_a2v`/SR) are not exercised in base.
-- **Audio decode path bit-exact vs official**: New parity test [`test_magi_human_sa_audio_official_parity.py`] passes at machine-eps (diff=0). FV's `SAAudioVAEModel` is identical to upstream `SAAudioFeatureExtractor.decode()`. Confirms OQ-6 is NOT in audio VAE.
-- **Production root cause identified**: FV's preset `_MAGI_HUMAN_NEGATIVE_PROMPT` was missing the audio-quality + speech-delivery blocks present in upstream `video_generate.py:222-224`. Fix applied at `presets.py`. Audio CFG amplifies the missing-block delta 5x → consistent with observed step-1 audio amplification of ~3x.
-- **Hardening**: Replaced silent zero-fallback in `denoising.py:127-135` with `ValueError`. Missing negative embeds at CFG=2 is a real bug, not silent-success.
-
-Caveat — parity test path bypasses preset prompts: `test_magi_human_pipeline_parity.py:291-298` uses random `txt_feat` and `neg_txt_feat` (identical on both sides), so the negative-prompt fix does NOT change parity numbers. Production inference (basic example) DOES use the preset and benefits from the fix.
-
-Reframed OQ-6 root cause:
-- Production-facing "blurry abstract" output: caused by incomplete negative prompt (audio CFG didn't have the right negatives). FIXED in this commit.
-- Parity-test 4-step compounding (1.196 mean): separate phenomenon — inherent FlowUniPC multistep scheduler amplification of per-call bf16 noise (~2x per DiT call expansively, 8 calls = ~256x). NOT a code bug; would require fp32 sensitive ops or a different scheduler to materially change.
-
-### Wave 8 (2026-05-01): broader CFG/preset/fallback audit + targeted fixes
-
-Audit found 4 more HARMFUL FV-vs-upstream divergences in addition to the negative-prompt incompleteness fixed in Wave 7:
-
-| # | Item | Severity | Status |
-|---|---|---|---|
-| 1 | T5-Gemma tokenizer pre-pads to 640 BEFORE encoding (pad-token hidden states pollute DiT input; magi_original_text_lens lies about real length) | HARMFUL | FIXED — `t5gemma.py:57-64` no longer passes `truncation`/`padding`/`max_length`; pad/trim handled post-encode by `MagiHumanLatentPreparationStage._pad_or_trim_dim1` |
-| 3 | Default resolution 448x256 vs upstream's 480x272 (snapped to 256). Production users got different aspect ratio than upstream | HARMFUL | FIXED — `presets.py:50-84` and `latent_preparation.py:130-133` now use `480x256` |
-| 10 | Audio decoding silently returned no audio if `batch.audio_latents` missing (joint AV makes this a real bug) | AMBIGUOUS→HARMFUL | FIXED — `audio_decoding.py:90-96` now raises `ValueError` |
-| 7 (stale) | Parity test FV scheduler helper claimed "double-shift" | (false alarm) | Already fixed in Wave 1A; audit was reading stale state |
-
-Other items from audit (BENIGN or out-of-scope for base T2AV): distill DDIM shortcut (cfg_number=1 path), Turbo VAE default (out-of-scope), A2V branch (out-of-scope), text_offset propagation (BENIGN for default v2 coords), frame_receptive_field (BENIGN for base local_attn_layers=[]), seed fallback (AMBIGUOUS edge case).
-
-**Parity-test test edit (Wave 7.5)**: pipeline parity test now encodes real preset prompts via T5-Gemma (`test_magi_human_pipeline_parity.py:59-153, 388-395`) instead of random tensors. Validates that production-facing preset values flow through the test path.
-
-**Critical caveat — parity numbers DON'T move with these fixes**: The parity test uses the SAME encoder/decoder/tokenizer on both FV and upstream sides. So fixing tokenizer-side pre-padding doesn't change FV-vs-upstream parity (both sides got the same wrong → now both get the same right). Wave 8 fixes are real PRODUCTION improvements (actual user inference now matches upstream's tokenization, resolution, and joint-AV invariants) but the residual ~0.47 (video) / ~1.0 (audio) drift in 4-step pipeline parity is the inherent bf16+CFG amplification floor through the multistep FlowUniPC scheduler.
-
-Per-test parity numbers post-Wave-8:
-| Test | Status | diff_max | diff_mean |
-|---|---|---:|---:|
-| DiT parity (single forward) | FAIL @ atol=0.03 | 0.057 | 0.0053 |
-| T5-Gemma parity | PASS | 0.0 | 0.0 |
-| Wan VAE parity (loose per OQ-7) | PASS @ atol=1e-3 | 8e-4 | 5e-5 |
-| SA Audio VAE parity | PASS | 0.0 | 0.0 |
-| SA official parity (NEW Wave 7) | PASS | 0.0 | 0.0 |
-| Pipeline parity (real prompts, 4-step) | FAIL @ atol=0.40 | video 6.69 / audio 3.45 | video 0.47 / audio 1.01 |
-
-OQ-6 status update:
-- **Production-facing root causes**: ALL identified and FIXED — incomplete neg prompt (Wave 7), tokenizer pre-padding (Wave 8 #1), resolution defaults (Wave 8 #3), silent fallbacks (Wave 7 + Wave 8 #10).
-- **Parity-test compounding**: bf16+CFG inherent amplification floor. Cannot be improved without fp32 sensitive ops or a less-amplifying scheduler. Tracked as `RESOLVED-PRODUCTION` for OQ-6 with a separate `OPEN-IF-NEEDED` follow-up for fp32 path investigation.
-
-### Wave 9 (2026-05-01): activation trace infrastructure
-
-Built Extension 0 of FastVideo's activation trace mode at `fastvideo/hooks/activation_trace.py` (env-gated zero-overhead module forward hooks). Designed for parity-debug across model ports — enable on both FastVideo's and upstream's path, diff resulting JSONL files to find first divergent layer.
-
-Key design properties:
-- `FASTVIDEO_TRACE_ACTIVATIONS=1` master toggle. Off = single env var lookup at startup, no hooks ever registered.
-- `FASTVIDEO_TRACE_LAYERS=<regex>` selective filter.
-- `FASTVIDEO_TRACE_STATS=abs_mean,sum,max,...` configurable per-tensor stats.
-- `FASTVIDEO_TRACE_STEPS=0,1,5` step-indexed dumps via `trace_step(idx)` context manager.
-- Output: JSONL records to `FASTVIDEO_TRACE_OUTPUT` path.
-
-E2E smoke confirmed: 28,864 records generated against the magi-human pipeline.
-
-Documentation at `docs/contributing/activation_trace.md`. Future Extensions 1-3 (FX/AST/dispatch) designed but not implemented.
-
-Companion skill at `~/.config/opencode/skill/add-model-trace/` (template for one-off ad-hoc port investigations) is unchanged.
-
-### Wave 10 (2026-05-01): WanVideo-pattern dtype refactor
-
-Removed all 7 hardcoded `.to(torch.bfloat16)` casts in `fastvideo/models/dits/magi_human.py`. These were verbatim copies of upstream `daVinci-MagiHuman/inference/model/dit/dit_module.py` (lines 619, 507, 650, 694, 696). FV now follows the canonical FastVideo dtype pattern exemplified by `fastvideo/models/dits/wanvideo.py`: model dtype is **loader-owned** via `pipeline_config.dit_precision` → `default_dtype` in `component_loader.py`. Inside DiT forward, `orig_dtype = self.linear_qkv.weight.dtype` (or equivalent) is captured and used for output preservation; no hardcoded model-dtype casts remain. The top-level block-input cast (formerly `x.to(torch.bfloat16)`) is now `x.to(<loader-owned dtype>)`.
-
-Refactored sites:
-- attention pre_norm output (line ~360)
-- q/k/v post-RoPE casts (lines ~399-401)
-- attention output (line ~411)
-- MLP pre_norm + activation casts (lines ~444-447)
-- top-level block-input cast (line ~684)
-
-Production behavior unchanged: bf16 parity numbers identical to baseline (`diff_max=0.057, diff_mean=0.005`). Loader's `dit_precision="bf16"` default → all params/inputs bf16 → `orig_dtype = bf16` → outputs preserved as bf16 → same as before.
-
-fp32 parity now works end-to-end on the FV side (model is dtype-agnostic in forward), but the parity test against upstream still shows bf16-noise residual drift (post-refactor: `diff_max=0.061, diff_mean=0.0068`; pre-refactor was `0.082 / 0.0079`, ~1.2x improvement). The remaining drift is from upstream `dit_module.py` itself — upstream still hardcodes `.to(torch.bfloat16)` in its forward, so even in an fp32 parity run, upstream's intermediate tensors are bf16. **Fully fp32-clean parity would require either patching the local upstream clone OR using a build of upstream where the hardcoded casts are also config-driven.**
-
-OQ-9 (NEW): upstream `daVinci-MagiHuman/inference/model/dit/dit_module.py` has hardcoded `.to(torch.bfloat16)` casts at lines 619, 507, 650, 694, 696. For full fp32 parity validation, these would need to be patched in the local clone OR a flag added upstream. Tracked as low-priority follow-up; affects only fp32 parity testing, not production.
-
-### Wave 14 (2026-05-02): upstream E2E coherent vs FV E2E noise (REAL bug confirmed)
-
-Ran the upstream `daVinci-MagiHuman` pipeline end-to-end with the same prompt + seed (42) + steps (32) + resolution (480x256) used by `examples/inference/basic/basic_magi_human.py`. Required installing `magi_compiler` from the local subdir, `alias_free_torch`, and downgrading `diffusers` per upstream's pinned version.
-
-**Result**: upstream produces a **coherent** video — young woman in a pink shirt reading a red book on a park bench surrounded by green trees, matching the prompt. Reference at `/tmp/opencode/upstream_magi_base_4s_480x256.mp4` (frames at `/tmp/opencode/upstream_frame_*.png`). FV produces **pure colorful-blob noise** at the same configuration (`outputs_video/magi_human_basic/output_magi_human_*.mp4`).
-
-**This invalidates the Wave 13 "structural / no real bug" verdict** for OQ-6 and reopens it. The bug is in code that production exercises but the parity test bypasses — parity test still passes (~0.5% per-step drift on (2,6,6) tiny synthetic latents) yet production produces noise on real (26,16,30) latents with real text encoding.
-
-#### Falsified candidates so far
-
-1. **T5-Gemma fp16 cast (Candidate A)**. Upstream `t5_gemma_model.py:24-27` casts `outputs["last_hidden_state"].half()` (bf16→fp16) before pad/trim → fp32; FV keeps bf16 → fp32 (`fastvideo/pipelines/basic/magi_human/pipeline_configs.py:t5gemma_postprocess_text`). Parity test bypasses this because it uses FV's encoder for both upstream and FV sides (`tests/local_tests/magi_human/test_magi_human_pipeline_parity.py:118-147`). Applied `outputs.last_hidden_state.to(torch.float16)` in the postprocess function and reran the basic example → **still pure noise**, visually identical to before. Reverted.
-
-2. **Local-window video→video attention (Candidate B)**. Upstream `MagiDataProxy.process_input` returns 5 args including `local_attn_handler` (`daVinci-MagiHuman/inference/pipeline/data_proxy.py:319-382`); FV's `MagiHumanDiT.forward` only takes `(x, coords, mm)` and uses full SDPA. **Verified to be inert for the base model**: upstream `local_attn_layers` config defaults to `[]` for the base BR pipeline (`daVinci-MagiHuman/inference/common/config.py:71`); only the SR_1080 pipeline sets non-empty layer indices (lines 229-241). Base-model upstream uses `flash_attn_with_cp` (full attention) at `dit_module.py:644-645`, equivalent to FV's full SDPA.
-
-#### Root cause + fix (Oracle, 2026-05-02)
-
-**Bug**: FV's `_img2tokens` packed video latents as **spatial-major** `(pT pH pW C)` (channels innermost) at `fastvideo/pipelines/basic/magi_human/stages/latent_preparation.py:84`. Upstream's `MagiDataProxy.process_input` uses `UnfoldNd(...)` at `daVinci-MagiHuman/inference/pipeline/data_proxy.py:287-317`, which is implemented via a grouped convolution (`groups=in_channels`) that reshapes to `(batch, in_channels * kernel_size_numel, -1)` (`unfoldNd/unfold.py:66`) — i.e. **channel-major** `(C pT pH pW)` (channels slowest). The DiT's `video_embedder` (`Linear(192, 5120)`) was trained on the channel-major layout. Spatial-major input silently permutes the in-features of every video token, scrambling the entire feature representation and producing pure noise.
-
-**Why parity test passed**: `test_magi_human_pipeline_parity.py:222` imports FV's `build_packed_inputs` for the upstream side too, so both sides ate the same FV-spatial-major tokens and agreed on equally-wrong inputs. Production faces real DiT weights and breaks.
-
-**Fix**: One-character rearrange-string change in `_img2tokens`:
-```diff
-- "B C (T pT) (H pH) (W pW) -> B (T H W) (pT pH pW C)"
-+ "B C (T pT) (H pH) (W pW) -> B (T H W) (C pT pH pW)"
-```
-
-`unpack_tokens` keeps spatial-major `(pT pH pW C)` because the DiT's `final_linear_video` was trained to emit that layout, mirroring upstream's `SingleData.depack_token_sequence` at `data_proxy.py:220-228`.
-
-**Validation**: Reran `examples/inference/basic/basic_magi_human.py` at the standard 480x256 / 32 step / seed 42 prompt. Output is **coherent video** matching the prompt — woman in teal sweater on a wooden park bench reading a book, green trees, sunny park scene. Output mp4 size dropped from ~932 KB (incompressible noise) to ~222 KB (coherent video). Frame samples at `/tmp/opencode/channelmajor_frame_*.png`.
-
-#### Wave 14 follow-up (2026-05-02): re-running parity exposed dtype-boundary divergences
-
-After fixing the channel-major bug, both DiT and pipeline parity tests started failing with much larger diffs than the pre-fix baseline (DiT diff_max=0.56 vs old "0.057"; pipeline video diff_mean=0.89 vs old 0.47). The pre-fix "0.057" baseline turned out to be a *garbage-in-garbage-out cancellation*: with both sides processing scrambled tokens, the kernel-level differences (TORCH_SDPA vs flash_attn) happened to converge on noise-equilibrium output. Once the inputs were correct, the underlying dtype-boundary divergences from upstream became visible.
-
-Three additional fixes brought parity to bit-exact:
-
-1. **Attention dtype boundary mirrors upstream**: FV now hardcodes the bf16 cast for SDPA inputs (matching `daVinci-MagiHuman/inference/model/dit/dit_module.py:508` `flash_attn_with_cp` which `q.to(bf16), k.to(bf16), v.to(bf16)` regardless of weight dtype). The attention output is upcast to fp32 before the per-head gating multiply (matching upstream's `bf16 * fp32` promotion at `dit_module.py:649`), and the gated result is cast to bf16 only for `linear_proj`. Wave 10's "dtype-agnostic" `orig_dtype` cast at the SDPA call was silently running fp32 attention whenever weights happened to be fp32 (e.g., parity-test load path). Fix in `fastvideo/models/dits/magi_human.py:MagiAttention.forward`.
-
-2. **fp32 residual stream**: removed the `x.to(linear_qkv.weight.dtype)` cast at `MagiHumanDiT.forward` (was line 689). Upstream casts to `params_dtype` which defaults to fp32, so the residual stream stays fp32 across all 40 layers — internal compute still bf16, but the cross-layer accumulator is fp32. FV's bf16 residual was compounding ~6-7 bits of mantissa loss per layer × 40 layers = visible parity drift. Fix in `fastvideo/models/dits/magi_human.py:MagiHumanDiT.forward`.
-
-3. **Pipeline parity test scheduler single-shift**: `_build_fastvideo_schedulers` was still constructing `FlowUniPCMultistepScheduler(shift=shift)` and then calling `set_timesteps(... shift=shift)` (double-shift), but production was migrated to single-shift in Wave 11 (`magi_human_pipeline.py:146-149` + `denoising.py:105-116`). The test helper had a stale docstring. Fix in `tests/local_tests/magi_human/test_magi_human_pipeline_parity.py:_build_fastvideo_schedulers`.
-
-**Final parity numbers** (8 of 8 tests passing, 7 of 8 bit-exact):
-
-| Test | diff_max | diff_mean |
-|---|---|---|
-| `test_magi_human_dit_parity` | 0.0 | 0.0 |
-| `test_magi_human_t5gemma_parity` | 0.0 | 0.0 |
-| `test_magi_human_sa_audio_parity` | 0.0 | 0.0 |
-| `test_magi_human_sa_audio_official_parity` | 0.0 | 0.0 |
-| `test_magi_human_vae_parity` | 8.0e-4 | 4.9e-5 |
-| `test_magi_human_pipeline_latent_parity` | 0.0 | 0.0 |
-| `test_magi_human_pipeline_smoke` (2 cases) | passes | passes |
-
-Production E2E re-validated post-fix: still produces coherent video at the standard 480x256 / 32 step / seed 42 prompt; runtime unchanged (23.5s).
-
-**OQ-6 RESOLVED** (Wave 14, full resolution including dtype boundaries).
-
-#### Parity test fidelity follow-up (separate issue)
-
-The pipeline parity test should be updated to use upstream's *real* `MagiDataProxy.process_input` for the upstream side (instead of importing FV's `build_packed_inputs`), so it can catch this class of "both sides use FV's helper, both consume scrambled tokens, parity passes" bypass in the future. Tracked as OQ-11.
-
-### Potential mitigations (not investigated this session)
-
-- Run sensitive ops (MM-layer pre-norm, attention) in fp32 instead of bf16.
-- Match upstream's exact dtype boundaries around MLP activation (verify FV does
-  the same fp32 cast upstream does in `_BF16ComputeLinear`).
-- Use a more numerically stable scheduler (FlowUniPC may have known issues at
-  certain step counts).
-- Per-modality `up_gate_proj` drill to find the first diverging activation.
-
-### Per-side layer logs and drill methodology
-
-Layer-by-layer traces are written to:
-
-- `/tmp/opencode/magi_dit_up_layers.log` (upstream reference)
-- `/tmp/opencode/magi_dit_fv_layers.log` (FastVideo)
-
-These are produced by `tests/local_tests/magi_human/_debug_magi_human_block_parity.py`
-via forward hooks registered on each transformer block. The `add-model-trace`
-skill at `~/.config/opencode/skill/add-model-trace/` generalizes this
-methodology for future ports: forward-hook + monkey-patch + git-stash-cleanup
-with hard rules around no-source-residue cleanup.
-
-## Open questions / blockers
-
-| ID | Item | Status |
-|---|---|---|
-| OQ-1 | **Native T5-Gemma port.** Full Phase 11 compliance requires a native FastVideo T5-Gemma implementation with no HF model-class imports in production code. Multi-week scope. | TRACKED FOLLOW-UP |
-| OQ-2 | **SSIM reference videos not seeded.** `fastvideo/tests/ssim/test_magi_human_similarity.py` skips cleanly until reference videos are uploaded to `FastVideo/ssim-reference-videos` on HF via the `seed-ssim-references` skill on Modal L40S. | TRACKED FOLLOW-UP |
-| OQ-3 | **Audio quality regression metric.** Mel-spectrogram L1 / multi-resolution STFT regression deferred per `tests/local_tests/stable-audio.md` precedent. | DEFERRED |
-| OQ-4 | **`_find_base_shard_dir` is fragile across HF-cache configurations.** Wave 1 fixed the loader with `snapshot_download(repo_id, allow_patterns=['base/*.safetensors'])` fallback in 3 files. `MAGI_HUMAN_BASE_SHARD_DIR` still works as an override but is no longer required. | RESOLVED |
-| OQ-5 | **Basic-example output mp4 visual quality is impressionistic at 256x448.** Root cause identified: OQ-6 (pre-existing compounding bf16 drift over the 32-step denoise loop). Wave 2-3 investigation confirmed the 4-step pipeline parity shows 18.85x compounding ratio vs expected 4x linear. See OQ-6 for full details and mitigation candidates. | RESOLVED-ROOT-CAUSE-IDENTIFIED (see OQ-6) |
-| OQ-6 | **Video patch packing was spatial-major instead of channel-major.** Wave 14 (2026-05-02) ran upstream E2E and got coherent output; FV produced pure noise at same config. Oracle triage identified the bug in `_img2tokens` rearrange order: FV used `(pT pH pW C)` (spatial-major) but the DiT's `video_embedder` Linear weight was trained on the channel-major `(C pT pH pW)` layout that upstream's `UnfoldNd` (grouped-conv reshape, `unfoldNd/unfold.py:66`) produces. The pipeline parity test imported FV's `build_packed_inputs` for both sides at `test_magi_human_pipeline_parity.py:222`, so it consumed equally-permuted tokens on both sides and reported agreement on garbage. Fixed in `latent_preparation.py:_img2tokens` by changing the einops pattern from `(pT pH pW C)` to `(C pT pH pW)`. Validated end-to-end: `examples/inference/basic/basic_magi_human.py` now produces coherent video matching the prompt (woman on park bench reading a book, green trees). Earlier waves' production-side fixes (negative prompt, tokenizer padding, resolution defaults, silent-audio fallback) all still stand. | RESOLVED — Wave 14 |
-| OQ-11 | **Pipeline parity test imports FV's `build_packed_inputs` for the upstream side.** `tests/local_tests/magi_human/test_magi_human_pipeline_parity.py:222` calls FV's packer for both sides instead of upstream's real `MagiDataProxy.process_input`. This let the channel-major-vs-spatial-major bug (OQ-6, Wave 14) sit silent for weeks because both sides agreed on the wrong layout. Update the parity test to drive the upstream side through `MagiDataProxy.process_input` so future packing-layout regressions are caught at parity time, not at production E2E. | TRACKED FOLLOW-UP |
-| OQ-7 | **Wan VAE shared fp32 op-order drift (MEDIUM PRIORITY).** FV uses `z * std + mean` at decode normalization; upstream uses `z / (1/std) + mean`. Bitwise non-equivalent in fp32. Affects all Wan-family pipelines (`fastvideo/configs/pipelines/wan.py`, `turbodiffusion.py`, `longcat.py`, magi-human). Magi VAE test loosened to `atol=1e-3, rtol=1e-3` (Wave 4) to defer. Tighten back to `atol=1e-4` once the Wan VAE op-order fix lands. Fix should be validated against Wan2.1, Wan2.2, and magi-human. Estimated 0.5-1 day to fix and validate. | TRACKED FOLLOW-UP |
-| OQ-9 | **Upstream `dit_module.py` hardcoded bf16 casts block full fp32 parity validation.** `daVinci-MagiHuman/inference/model/dit/dit_module.py` has hardcoded `.to(torch.bfloat16)` casts at lines 619, 507, 650, 694, 696. FV's DiT forward is now dtype-agnostic (Wave 10), but parity tests against upstream still show bf16-noise residual drift in fp32 runs because upstream's intermediate tensors are bf16. Full fp32-clean parity would require patching the local upstream clone or adding a dtype-config flag upstream. Affects only fp32 parity testing, not production. | TRACKED FOLLOW-UP (LOW PRIORITY) |
-
-## Troubleshooting
-
-**`RuntimeError: Upstream DiT missing 331 keys` despite shards being present.**
-This happens when the upstream base shards are downloaded into one HF cache
-(e.g. `~/.cache/huggingface/hub/`) but `_find_base_shard_dir` resolves the
-snapshot via a different cache path (e.g. `/raid/huggingface/hub/...`) where
-only `model.safetensors.index.json` is present, not the 7 shard files.
-
-**Workaround**: explicitly set `MAGI_HUMAN_BASE_SHARD_DIR` to the snapshot dir
-that actually contains the `model-0000*-of-00007.safetensors` shards:
-
-```bash
-export MAGI_HUMAN_BASE_SHARD_DIR=~/.cache/huggingface/hub/models--GAIR--daVinci-MagiHuman/snapshots/<sha>/base
-```
-
-Tracked as open question **OQ-4** for a more robust loader.
-
-**`401 Unauthorized` on any gated repo.** Check `echo $HF_TOKEN` and confirm
-you've accepted the model terms at each URL listed in §1. The four repos have
-separate terms pages; accepting one doesn't cover the others.
-
-- T5-Gemma: https://huggingface.co/google/t5gemma-9b-9b-ul2
-- Stable Audio Open: https://huggingface.co/stabilityai/stable-audio-open-1.0
-- Wan 2.2 TI2V-5B: https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B
-- daVinci-MagiHuman: https://huggingface.co/GAIR/daVinci-MagiHuman
-
-**Override the base shard directory.** If you have the raw MagiHuman shards
-at a non-default path, point the tests at them:
-
-```bash
-export MAGI_HUMAN_BASE_SHARD_DIR=/path/to/raw/shards
-```
-
-**Override the converted weights path.** If you ran the conversion script with
-a custom `--output` path, tell the tests where to find it:
-
-```bash
-export MAGI_HUMAN_DIFFUSERS_PATH=/path/to/converted_weights/magi_human_base
-```
-
-**Missing `daVinci-MagiHuman/` clone.** Tests that need the upstream reference
-(`test_magi_human_parity.py`, `test_magi_human_pipeline_parity.py`) skip
-cleanly with a message pointing to the clone command in §3. The VAE parity
-tests and the smoke test don't need the clone.
-
-**OOM during DiT load.** The base DiT loads in bf16 by default. If you're
-tight on VRAM, use `--cast-bf16` during conversion to ensure the transformer
-shards are stored in bf16 rather than fp32. The distill variant is the same
-size; both fit on a single 80 GB GPU.
-
-**Wall-clock blew up past 10 min.** The first run downloads T5-Gemma (~18 GB),
-the Wan VAE, and the Stable Audio VAE if they aren't cached. See the pre-warm
-step in §5.
-
-## Adding new parity tests for this family
-
-`tests/local_tests/helpers/magi_human_upstream.py` contains shared reference
-loaders for the upstream DiT, VAE, and pipeline. Use these as the starting
-point for any new parity test rather than duplicating the load logic.
-
-The `_debug_magi_human_block_parity.py` and `_debug_magi_human_weight_diff.py`
-scripts in `tests/local_tests/magi_human/` are scratch tools for divergence
-investigation. They are NOT pytest tests and must NOT be promoted to formal
-tests. Run them directly with `python` when you need to inspect per-block diffs
-or weight mismatches during a parity-debug session.
-
-If you need to chase per-layer divergence on a future add-model port, see the
-`add-model-trace` skill at `~/.config/opencode/skill/add-model-trace/`.
-Generalized from `tests/local_tests/magi_human/_debug_magi_human_block_parity.py`
-(the worked magi example), it provides a forward-hook + monkey-patch +
-git-stash-cleanup methodology with hard rules around no-source-residue cleanup.

tests/local_tests/magi_human/_debug_magi_human_block_parity.py

@@ -1,358 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-block divergence debugger for MagiHumanDiT vs upstream DiTModel.
-
-Not a pytest test (filename starts with `_`). Run directly:
-
-    python tests/local_tests/transformers/_debug_magi_human_block_parity.py
-
-Mirrors the inputs / loader of `test_magi_human_dit_parity` but adds
-forward hooks on:
-
-  * `model.adapter` (post-embedding)
-  * each `model.block.layers[i]` (per-block output, 40 blocks)
-  * model output (post-final-norms)
-
-Logs (idx, label, abs_mean, sum) for both sides side-by-side, and
-prints the first block where |abs_mean diff| or |sum diff| exceeds a
-threshold so we know where to drill in.
-"""
-from __future__ import annotations
-
-import gc
-import glob
-import os
-import sys
-from pathlib import Path
-
-import torch
-
-# Match the parity test: FA on both sides.
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-sys.path.insert(0, str(REPO_ROOT))
-
-
-def _find_base_shard_dir() -> Path | None:
-    override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR")
-    if override:
-        p = Path(override)
-        return p if p.is_dir() else None
-    try:
-        from huggingface_hub import hf_hub_download
-        idx = hf_hub_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            filename="base/model.safetensors.index.json",
-        )
-        return Path(idx).parent
-    except Exception:
-        return None
-
-
-def _stat(name: str, t: torch.Tensor) -> dict:
-    f = t.detach().float()
-    return {
-        "name": name,
-        "shape": tuple(t.shape),
-        "abs_mean": f.abs().mean().item(),
-        "sum": f.sum().item(),
-        "min": f.min().item(),
-        "max": f.max().item(),
-    }
-
-
-def _attach_block_hooks(model, label: str, log: list[dict],
-                        tensors: dict[str, torch.Tensor] | None = None,
-                        drill_layer: int | None = None):
-    """Attach forward hooks to adapter + each block.layers[i].
-
-    If ``drill_layer`` is set, also hooks the submodules of
-    ``block.layers[drill_layer]`` (attention, mlp, attn_post_norm,
-    mlp_post_norm if present), letting us pinpoint which submodule
-    introduces the first measurable drift.
-    """
-    handles = []
-
-    def _hook(name):
-        def fn(_module, _inputs, outputs):
-            t = outputs[0] if isinstance(outputs, tuple) else outputs
-            if not torch.is_tensor(t):
-                return
-            log.append({"side": label, **_stat(name, t)})
-            if tensors is not None:
-                tensors[name] = t.detach().float().cpu()
-        return fn
-
-    def _pre_hook(name):
-        def fn(_module, inputs):
-            t = inputs[0] if isinstance(inputs, tuple) else inputs
-            if not torch.is_tensor(t):
-                return
-            label_in = f"{name}<in>"
-            log.append({"side": label, **_stat(label_in, t)})
-            if tensors is not None:
-                tensors[label_in] = t.detach().float().cpu()
-        return fn
-
-    handles.append(model.adapter.register_forward_hook(_hook("adapter")))
-    for i, layer in enumerate(model.block.layers):
-        handles.append(layer.register_forward_hook(_hook(f"block[{i:02d}]")))
-        if drill_layer is not None and i == drill_layer:
-            tag = f"L{i:02d}"
-            handles.append(layer.attention.register_forward_hook(
-                _hook(f"{tag}.attention")))
-            handles.append(layer.mlp.pre_norm.register_forward_hook(
-                _hook(f"{tag}.mlp.pre_norm")))
-            handles.append(layer.mlp.up_gate_proj.register_forward_hook(
-                _hook(f"{tag}.mlp.up_gate_proj")))
-            # Pre-hook on down_proj captures the post-activation tensor
-            # (the activation func is a free function, not a module, so
-            # we observe its output by intercepting down_proj's input).
-            handles.append(layer.mlp.down_proj.register_forward_pre_hook(
-                _pre_hook(f"{tag}.mlp.down_proj")))
-            handles.append(layer.mlp.down_proj.register_forward_hook(
-                _hook(f"{tag}.mlp.down_proj")))
-            handles.append(layer.mlp.register_forward_hook(
-                _hook(f"{tag}.mlp")))
-            if hasattr(layer, "attn_post_norm"):
-                handles.append(layer.attn_post_norm.register_forward_hook(
-                    _hook(f"{tag}.attn_post_norm")))
-            if hasattr(layer, "mlp_post_norm"):
-                handles.append(layer.mlp_post_norm.register_forward_hook(
-                    _hook(f"{tag}.mlp_post_norm")))
-    return handles
-
-
-def main() -> None:
-    if not torch.cuda.is_available():
-        print("Need CUDA. Skipping.")
-        return
-
-    upstream_src = REPO_ROOT / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        print(f"daVinci-MagiHuman/ not present under {REPO_ROOT}.")
-        return
-
-    base_shard_dir = _find_base_shard_dir()
-    if base_shard_dir is None or not base_shard_dir.is_dir():
-        print("Upstream base/ shards missing.")
-        return
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_DIFFUSERS_PATH",
-        REPO_ROOT / "converted_weights" / "magi_human_base",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    if not transformer_dir.is_dir():
-        print(f"Converted transformer dir missing at {transformer_dir}")
-        return
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs, load_upstream_dit,
-    )
-    install_stubs()
-
-    # Optional: monkey-patch PackedExpertLinear.forward to mirror upstream's
-    # explicit-cast torch.matmul pattern (`_BF16ComputeLinear.apply`).
-    # Toggled via env var so the experiment is reproducible.
-    if os.getenv("MAGI_DEBUG_PATCH_LINEAR") == "1":
-        from fastvideo.models.dits import magi_human as _mh
-
-        def _patched_forward(self, x, modality_dispatcher=None):
-            def _bf16_linear(inp, w, b):
-                inp_c = inp.to(torch.bfloat16)
-                w_c = w.to(torch.bfloat16)
-                out = torch.matmul(inp_c, w_c.t())
-                if b is not None:
-                    out = out + b.to(torch.bfloat16)
-                return out.to(inp.dtype)
-
-            if self.num_experts == 1:
-                return _bf16_linear(x, self.weight, self.bias)
-            assert modality_dispatcher is not None
-            parts = modality_dispatcher.dispatch(x)
-            w_chunks = self.weight.chunk(self.num_experts, dim=0)
-            b_chunks = (
-                self.bias.chunk(self.num_experts, dim=0)
-                if self.bias is not None else [None] * self.num_experts
-            )
-            for i in range(self.num_experts):
-                parts[i] = _bf16_linear(parts[i], w_chunks[i], b_chunks[i])
-            return modality_dispatcher.undispatch(*parts)
-
-        _mh.PackedExpertLinear.forward = _patched_forward
-        print("[debug] Patched PackedExpertLinear.forward to mirror "
-              "upstream's _BF16ComputeLinear pattern.")
-
-    device = torch.device("cuda:0")
-    torch.manual_seed(0)
-
-    z_dim = 48
-    pT, pH, pW = 1, 2, 2
-    lat_T, lat_H, lat_W = 2, 6, 6
-    video_latent = torch.randn((1, z_dim, lat_T, lat_H, lat_W), dtype=torch.float32, device=device)
-    num_video = (lat_T // pT) * (lat_H // pH) * (lat_W // pW)
-    num_audio = 4
-    num_text = 8
-    audio_latent = torch.randn((1, num_audio, 64), dtype=torch.float32, device=device)
-    text_feat = torch.randn((1, num_text, 3584), dtype=torch.float32, device=device)
-
-    from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-        build_packed_inputs,
-    )
-    x, coords, mm = build_packed_inputs(
-        video_latent=video_latent,
-        audio_latent=audio_latent,
-        audio_feat_len=num_audio,
-        txt_feat=text_feat,
-        txt_feat_len=num_text,
-        patch_size=(pT, pH, pW),
-        coords_style="v2",
-    )
-    total_tokens = x.shape[0]
-
-    # --- Upstream ---
-    print("Loading upstream DiTModel...")
-    upstream = load_upstream_dit(base_shard_dir, device=device, dtype=None)
-    from inference.common import VarlenHandler
-    cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device)
-    varlen = VarlenHandler(
-        cu_seqlens_q=cu, cu_seqlens_k=cu,
-        max_seqlen_q=total_tokens, max_seqlen_k=total_tokens,
-    )
-    drill_layer = int(os.getenv("MAGI_DEBUG_DRILL_LAYER", "0"))
-    up_log: list[dict] = []
-    up_tensors: dict[str, torch.Tensor] = {}
-    _attach_block_hooks(upstream, "up", up_log, tensors=up_tensors, drill_layer=drill_layer)
-    print("Running upstream forward (with hooks)...")
-    with torch.inference_mode():
-        ref_out = upstream(
-            x=x.clone(), coords_mapping=coords.clone(),
-            modality_mapping=mm.clone(),
-            varlen_handler=varlen, local_attn_handler=None,
-        ).detach().float().cpu()
-    del upstream
-    gc.collect(); torch.cuda.empty_cache()
-
-    # --- FastVideo ---
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-    print("Loading FastVideo MagiHumanDiT...")
-    fv = MagiHumanDiT(MagiHumanVideoConfig())
-    state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        state.update(load_file(shard))
-    fv.load_state_dict(state, strict=False)
-    fv = fv.to(device).eval()
-    fv_log: list[dict] = []
-    fv_tensors: dict[str, torch.Tensor] = {}
-    _attach_block_hooks(fv, "fv", fv_log, tensors=fv_tensors, drill_layer=drill_layer)
-    print("Running FastVideo forward (with hooks)...")
-    with torch.inference_mode():
-        fv_out = fv(x.clone(), coords.clone(), mm.clone()).detach().float().cpu()
-
-    # --- Side-by-side per-block comparison ---
-    # Group by name; each name should appear once on each side.
-    by_name: dict[str, dict] = {}
-    for entry in up_log + fv_log:
-        d = by_name.setdefault(entry["name"], {})
-        d[entry["side"]] = entry
-
-    print()
-    print(f"{'name':<14} {'up_shape':<22} {'up_absmean':>12} {'fv_absmean':>12} "
-          f"{'absmean_diff':>14} {'rel%':>8} {'up_sum':>14} {'fv_sum':>14} {'sum_diff':>12}")
-    print("-" * 145)
-
-    first_div_idx = None
-    rel_threshold = 0.005  # 0.5% drift in abs_mean per block
-
-    # Print in canonical order: adapter, drilled-layer submodules
-    # (intermixed with their parent block), then remaining blocks.
-    def _sort_key(n: str):
-        if n == "adapter":
-            return (0, "")
-        if n.startswith(f"L{drill_layer:02d}."):
-            # Submodule snapshots — sort to appear right before
-            # block[NN] so they read as "what fed into block[NN]'s
-            # output". Order: attention, attn_post_norm, mlp, mlp_post_norm.
-            sub_order = {
-                "attention": 0,
-                "attn_post_norm": 1,
-                "mlp.pre_norm": 2,
-                "mlp.up_gate_proj": 3,
-                "mlp.down_proj": 4,
-                "mlp": 5,
-                "mlp_post_norm": 6,
-            }.get(n.split(".", 1)[1], 9)
-            return (1, f"block[{drill_layer:02d}]", sub_order)
-        if n.startswith("block["):
-            return (1, n, 99)
-        return (2, n, 0)
-
-    Path("/tmp/opencode").mkdir(parents=True, exist_ok=True)
-    up_log_path = Path("/tmp/opencode/magi_dit_up_layers.log")
-    fv_log_path = Path("/tmp/opencode/magi_dit_fv_layers.log")
-
-    def _log_lines(entries: list[dict]) -> list[str]:
-        lines = []
-        for entry in sorted(entries, key=lambda e: _sort_key(e["name"])):
-            lines.append(
-                f"{entry['name']}\t{entry['shape']}\t{entry['abs_mean']:.6f}\t"
-                f"{entry['sum']:.6f}\t{entry['min']:.6f}\t{entry['max']:.6f}"
-            )
-        return lines
-
-    up_log_path.write_text("\n".join(_log_lines(up_log)) + "\n")
-    fv_log_path.write_text("\n".join(_log_lines(fv_log)) + "\n")
-
-    ordered_names = sorted(by_name.keys(), key=_sort_key)
-    for name in ordered_names:
-        d = by_name[name]
-        up = d.get("up")
-        fv = d.get("fv")
-        if up is None or fv is None:
-            continue
-        am_diff = abs(up["abs_mean"] - fv["abs_mean"])
-        am_rel = am_diff / max(up["abs_mean"], 1e-9)
-        sum_diff = abs(up["sum"] - fv["sum"])
-        flag = ""
-        if name.startswith("block[") and am_rel > rel_threshold:
-            flag = " <<< DIVERGE"
-            if first_div_idx is None:
-                first_div_idx = int(name[len("block["):-1])
-        print(f"{name:<14} {str(up['shape']):<22} {up['abs_mean']:>12.6f} {fv['abs_mean']:>12.6f} "
-              f"{am_diff:>14.6f} {am_rel*100:>7.3f}% {up['sum']:>14.4f} {fv['sum']:>14.4f} {sum_diff:>12.4f}{flag}")
-
-    print()
-    if first_div_idx is not None:
-        print(f"First block exceeding {rel_threshold*100:.2f}% abs_mean rel drift: block[{first_div_idx:02d}]")
-    else:
-        print(f"No block exceeded {rel_threshold*100:.2f}% — divergence is amortized across blocks.")
-
-    print("[debug] Per-side logs: /tmp/opencode/magi_dit_up_layers.log + /tmp/opencode/magi_dit_fv_layers.log  (diff with: diff /tmp/opencode/magi_dit_up_layers.log /tmp/opencode/magi_dit_fv_layers.log)")
-
-    # Final output diff
-    diff = (ref_out - fv_out).abs()
-    print()
-    print(f"Final  ref_abs={ref_out.abs().mean():.6f}  fv_abs={fv_out.abs().mean():.6f}  "
-          f"diff_max={diff.max():.6f}  diff_mean={diff.mean():.6f}")
-
-    # Element-wise diff stats for drilled submodules.
-    print()
-    print(f"Element-wise diffs for drilled L{drill_layer:02d} submodules:")
-    print(f"{'name':<30} {'shape':<22} {'diff_max':>12} {'diff_mean':>12} {'diff_rel%':>10}")
-    print("-" * 95)
-    common_names = set(up_tensors.keys()) & set(fv_tensors.keys())
-    for name in sorted(common_names):
-        a, b = up_tensors[name], fv_tensors[name]
-        if a.shape != b.shape:
-            continue
-        d = (a - b).abs()
-        ref_abs = a.abs().mean().item()
-        rel = (d.mean().item() / max(ref_abs, 1e-9)) * 100
-        print(f"{name:<30} {str(tuple(a.shape)):<22} {d.max().item():>12.6f} {d.mean().item():>12.6f} {rel:>9.4f}%")
-
-
-if __name__ == "__main__":
-    main()

tests/local_tests/magi_human/_debug_magi_human_weight_diff.py

@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Verify weights are bit-exact between upstream and FastVideo paths.
-
-If they're not, the per-block parity drift could come from weight
-mismatches (conversion-script truncation, bf16-cast-then-load, etc.)
-rather than op-ordering. Run before concluding "bf16 noise".
-"""
-from __future__ import annotations
-
-import glob
-import os
-import sys
-from pathlib import Path
-
-import torch
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-sys.path.insert(0, str(REPO_ROOT))
-
-
-def _find_base_shard_dir() -> Path | None:
-    """Return the local path to GAIR/daVinci-MagiHuman/base/ with shards present, or None."""
-    override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR")
-    if override:
-        p = Path(override)
-        return p if p.is_dir() else None
-    try:
-        from huggingface_hub import snapshot_download
-        snap = snapshot_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            allow_patterns=[
-                "base/*.safetensors",
-                "base/model.safetensors.index.json",
-            ],
-        )
-        candidate = Path(snap) / "base"
-        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
-            return candidate
-        return None
-    except Exception:
-        return None
-
-
-def main() -> None:
-    if not torch.cuda.is_available():
-        print("Need CUDA.")
-        return
-
-    upstream_src = REPO_ROOT / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        print("daVinci-MagiHuman/ missing.")
-        return
-
-    base_shard_dir = _find_base_shard_dir()
-    if base_shard_dir is None:
-        print("GAIR/daVinci-MagiHuman base shards not available locally.")
-        return
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_DIFFUSERS_PATH",
-        REPO_ROOT / "converted_weights" / "magi_human_base",
-    ))
-    transformer_dir = converted_dir / "transformer"
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs, load_upstream_dit,
-    )
-    install_stubs()
-
-    device = torch.device("cuda:0")
-
-    # Load both, dumping weight tensors to dicts for comparison.
-    print("Loading upstream...")
-    up = load_upstream_dit(base_shard_dir, device=device, dtype=None)
-    up_state = {k: v.detach().cpu() for k, v in up.state_dict().items()}
-    del up
-    import gc
-
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    print("Loading FastVideo...")
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-    fv = MagiHumanDiT(MagiHumanVideoConfig())
-    state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        state.update(load_file(shard))
-    fv.load_state_dict(state, strict=False)
-    fv_state = {k: v.detach().cpu() for k, v in fv.state_dict().items()}
-
-    # Compare overlapping keys.
-    up_keys = set(up_state.keys())
-    fv_keys = set(fv_state.keys())
-    only_up = up_keys - fv_keys
-    only_fv = fv_keys - up_keys
-    common = up_keys & fv_keys
-    print(f"Keys: common={len(common)}, only_upstream={len(only_up)}, only_fastvideo={len(only_fv)}")
-    if only_up:
-        print(f"  Only upstream (sample): {sorted(only_up)[:5]}")
-    if only_fv:
-        print(f"  Only fv (sample): {sorted(only_fv)[:5]}")
-
-    bit_exact = 0
-    diff_keys = []
-    shape_mismatch = []
-    dtype_mismatch = []
-    for k in sorted(common):
-        a, b = up_state[k], fv_state[k]
-        if a.shape != b.shape:
-            shape_mismatch.append((k, tuple(a.shape), tuple(b.shape)))
-            continue
-        if a.dtype != b.dtype:
-            dtype_mismatch.append((k, a.dtype, b.dtype))
-        d = (a.float() - b.float()).abs()
-        max_d = d.max().item()
-        if max_d == 0.0:
-            bit_exact += 1
-        else:
-            diff_keys.append((k, max_d, d.mean().item(), tuple(a.shape), str(a.dtype)))
-
-    print(f"\nWeight comparison ({len(common)} keys):")
-    print(f"  bit-exact: {bit_exact}")
-    print(f"  with diff: {len(diff_keys)}")
-    print(f"  shape mismatch: {len(shape_mismatch)}")
-    print(f"  dtype mismatch: {len(dtype_mismatch)}")
-
-    if dtype_mismatch:
-        print("\nDtype mismatches:")
-        for k, da, db in dtype_mismatch[:10]:
-            print(f"  {k}: up={da} fv={db}")
-
-    if shape_mismatch:
-        print("\nShape mismatches:")
-        for k, sa, sb in shape_mismatch[:10]:
-            print(f"  {k}: up={sa} fv={sb}")
-
-    if diff_keys:
-        print("\nTop weight diffs (max-diff sorted):")
-        diff_keys.sort(key=lambda x: -x[1])
-        for k, max_d, mean_d, shape, dtype in diff_keys[:15]:
-            print(f"  {dtype} {str(shape):<40} max={max_d:.6e} mean={mean_d:.6e}  {k}")
-
-
-if __name__ == "__main__":
-    main()

tests/local_tests/magi_human/test_magi_human_distill_parity.py

@@ -1,245 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DiT parity test for the daVinci-MagiHuman DMD-2 distilled checkpoint.
-
-The distill variant has the SAME architecture as the base model (same 40
-layers, same hidden_size, same mm_layers / gelu7_layers / local_attn_layers,
-same head_dim and num_query_groups; see
-`daVinci-MagiHuman/inference/common/config.py:ModelConfig`). Only the
-weights differ: distill is trained for 8-step DMD-2 inference without CFG.
-
-This test mirrors `test_magi_human_parity.py::test_magi_human_dit_parity`
-exactly, just pointing at the `distill/` subfolder of GAIR/daVinci-MagiHuman
-and the matching `converted_weights/magi_human_distill/`.
-
-Skips cleanly when:
-  * `daVinci-MagiHuman/` clone is absent
-  * GAIR/daVinci-MagiHuman distill shards are not locally available
-  * Converted distill weights have not been produced yet
-  * CUDA is unavailable
-"""
-from __future__ import annotations
-
-import gc
-import glob
-import os
-from pathlib import Path
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-from fastvideo.forward_context import set_forward_context
-
-
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-
-def _find_distill_shard_dir() -> Path | None:
-    """Return the local path to GAIR/daVinci-MagiHuman/distill/ shards or None."""
-    override = os.getenv("MAGI_HUMAN_DISTILL_SHARD_DIR")
-    if override:
-        p = Path(override)
-        return p if p.is_dir() else None
-    try:
-        from huggingface_hub import snapshot_download
-        snap = snapshot_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            allow_patterns=[
-                "distill/*.safetensors",
-                "distill/model.safetensors.index.json",
-            ],
-        )
-        candidate = Path(snap) / "distill"
-        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
-            return candidate
-        return None
-    except Exception:
-        return None
-
-
-def _cleanup_gpu() -> None:
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman distill DiT parity requires CUDA.",
-)
-def test_magi_human_distill_dit_parity():
-    repo_root = Path(__file__).resolve().parents[3]
-    upstream_src = repo_root / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        pytest.skip(
-            "Upstream daVinci-MagiHuman/ clone missing. Run "
-            "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`"
-        )
-
-    distill_shard_dir = _find_distill_shard_dir()
-    if distill_shard_dir is None or not distill_shard_dir.is_dir():
-        pytest.skip(
-            "GAIR/daVinci-MagiHuman distill/ shards not available locally. "
-            "Set MAGI_HUMAN_DISTILL_SHARD_DIR or run the conversion once to "
-            "populate the HF cache."
-        )
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_DISTILL_DIFFUSERS_PATH",
-        repo_root / "converted_weights" / "magi_human_distill",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    if not transformer_dir.is_dir():
-        pytest.skip(
-            f"Converted distill transformer dir missing at {transformer_dir}. Run "
-            f"scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py "
-            f"--subfolder distill --cast-bf16 first."
-        )
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs,
-        load_upstream_dit,
-    )
-    install_stubs()
-
-    device = torch.device("cuda:0")
-    torch.manual_seed(0)
-
-    z_dim = 48
-    pT, pH, pW = 1, 2, 2
-    lat_T, lat_H, lat_W = 2, 6, 6
-    video_latent = torch.randn(
-        (1, z_dim, lat_T, lat_H, lat_W),
-        dtype=torch.float32, device=device,
-    )
-    num_video_tokens = (lat_T // pT) * (lat_H // pH) * (lat_W // pW)
-    num_audio_tokens = 4
-    num_text_tokens = 8
-    audio_latent = torch.randn(
-        (1, num_audio_tokens, 64),
-        dtype=torch.float32, device=device,
-    )
-    text_feat = torch.randn(
-        (1, num_text_tokens, 3584),
-        dtype=torch.float32, device=device,
-    )
-
-    from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-        build_packed_inputs,
-    )
-    from fastvideo.models.dits.magi_human import Modality  # noqa: F401
-
-    x, coords, mm = build_packed_inputs(
-        video_latent=video_latent,
-        audio_latent=audio_latent,
-        audio_feat_len=num_audio_tokens,
-        txt_feat=text_feat,
-        txt_feat_len=num_text_tokens,
-        patch_size=(pT, pH, pW),
-        coords_style="v2",
-    )
-    assert x.shape[0] == num_video_tokens + num_audio_tokens + num_text_tokens
-
-    total_tokens = x.shape[0]
-
-    # Distill arch is identical to base; load_upstream_dit's _base_arch_dict
-    # describes both since they share num_layers / hidden_size / mm_layers
-    # / etc. Just point at the distill shards.
-    print("Loading upstream distill DiTModel from distill shards...")
-    upstream_model = load_upstream_dit(
-        distill_shard_dir,
-        device=device,
-        dtype=None,
-    )
-
-    from inference.common import VarlenHandler
-    cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device)
-    varlen = VarlenHandler(
-        cu_seqlens_q=cu,
-        cu_seqlens_k=cu,
-        max_seqlen_q=total_tokens,
-        max_seqlen_k=total_tokens,
-    )
-
-    print("Running upstream distill forward...")
-    with torch.inference_mode():
-        ref_out = upstream_model(
-            x=x.clone(),
-            coords_mapping=coords.clone(),
-            modality_mapping=mm.clone(),
-            varlen_handler=varlen,
-            local_attn_handler=None,
-        ).detach().float().cpu()
-
-    del upstream_model
-    _cleanup_gpu()
-
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-    print("Loading FastVideo MagiHumanDiT from converted distill transformer/...")
-    fv_cfg = MagiHumanVideoConfig()
-    fv_model = MagiHumanDiT(fv_cfg)
-
-    fv_state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        fv_state.update(load_file(shard))
-    missing, unexpected = fv_model.load_state_dict(fv_state, strict=False)
-    assert not missing, f"FastVideo distill DiT missing {len(missing)} keys: {missing[:5]}"
-    assert not unexpected, f"FastVideo distill DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
-
-    fv_model = fv_model.to(device=device)
-    fv_model.eval()
-
-    print("Running FastVideo distill forward...")
-    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
-        fv_out = fv_model(x.clone(), coords.clone(), mm.clone()).detach().float().cpu()
-
-    print(
-        f"ref sum={ref_out.sum().item():.4f} "
-        f"abs_mean={ref_out.abs().mean().item():.4f} "
-        f"shape={tuple(ref_out.shape)}"
-    )
-    print(
-        f"fv  sum={fv_out.sum().item():.4f} "
-        f"abs_mean={fv_out.abs().mean().item():.4f} "
-        f"shape={tuple(fv_out.shape)}"
-    )
-    diff = (ref_out - fv_out).abs()
-    print(
-        f"diff max={diff.max().item():.6f} "
-        f"mean={diff.mean().item():.6f} "
-        f"median={diff.median().item():.6f}"
-    )
-
-    ref_video = ref_out[:num_video_tokens]
-    fv_video = fv_out[:num_video_tokens]
-    ref_audio = ref_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
-    fv_audio = fv_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
-    ref_text = ref_out[num_video_tokens + num_audio_tokens:]
-    fv_text = fv_out[num_video_tokens + num_audio_tokens:]
-    video_diff = (ref_video - fv_video).abs()
-    audio_diff = (ref_audio - fv_audio).abs()
-    text_diff = (ref_text - fv_text).abs()
-    print(
-        f"video  ref_abs={ref_video.abs().mean():.4f} "
-        f"diff_max={video_diff.max():.4f} diff_mean={video_diff.mean():.4f}"
-    )
-    print(
-        f"audio  ref_abs={ref_audio.abs().mean():.4f} "
-        f"diff_max={audio_diff.max():.4f} diff_mean={audio_diff.mean():.4f}"
-    )
-    print(
-        f"text   ref_abs={ref_text.abs().mean():.4f} "
-        f"diff_max={text_diff.max():.4f} diff_mean={text_diff.mean():.4f}"
-    )
-
-    assert ref_out.shape == fv_out.shape
-    assert_close(fv_text, ref_text, atol=1e-6, rtol=1e-6)
-    # Same tolerance as base DiT parity (post-Wave 14b dtype-boundary fixes,
-    # both DiTs are bit-exact via shared upstream BaseLinear bf16 path).
-    assert_close(fv_out, ref_out, atol=0.03, rtol=0.01)
-    ref_abs = ref_out.abs().mean().item()
-    fv_abs = fv_out.abs().mean().item()
-    rel = abs(ref_abs - fv_abs) / max(ref_abs, 1e-6)
-    assert rel < 0.05, f"abs_mean drift {rel:.3%} > 5%"

tests/local_tests/magi_human/test_magi_human_parity.py

@@ -1,281 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Numerical parity test: FastVideo MagiHumanDiT vs upstream DiTModel.
-
-Loads both models from the **same converted base checkpoint** and runs them
-on identical small inputs. Asserts closeness on the joint video+audio
-output tensor.
-
-What this catches (that the preflight test does NOT):
-  - Silent weight-name mismatches that `strict=False` loading would hide.
-  - Wrong modality-expert chunking inside `PackedExpertLinear`.
-  - RoPE sin/cos ordering flipped.
-  - Per-head gating dtype / split order.
-  - swiglu7 / gelu7 off-by-one on the `+1` linear bias.
-
-Skips cleanly when:
-  - `daVinci-MagiHuman/` clone is absent (no upstream source).
-  - GAIR/daVinci-MagiHuman base shards are not available locally.
-  - CUDA is unavailable.
-
-Tolerance: `atol=5e-3, rtol=5e-3` on bf16 forward paths. The FastVideo
-attention path uses `F.scaled_dot_product_attention` while upstream uses
-`flash_attn_func`; both accumulate in bf16 but via different kernels, so
-small drift is expected and bounded.
-"""
-from __future__ import annotations
-
-import gc
-import os
-from pathlib import Path
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-from fastvideo.forward_context import set_forward_context
-
-
-# Force TORCH_SDPA for FastVideo so the attention kernel is deterministic.
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-
-def _find_base_shard_dir() -> Path | None:
-    """Return the local path to GAIR/daVinci-MagiHuman/base/ with shards present, or None."""
-    override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR")
-    if override:
-        p = Path(override)
-        return p if p.is_dir() else None
-    try:
-        from huggingface_hub import snapshot_download
-        snap = snapshot_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            allow_patterns=[
-                "base/*.safetensors",
-                "base/model.safetensors.index.json",
-            ],
-        )
-        candidate = Path(snap) / "base"
-        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
-            return candidate
-        return None
-    except Exception:
-        return None
-
-
-def _cleanup_gpu() -> None:
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman DiT parity requires CUDA.",
-)
-def test_magi_human_dit_parity():
-    repo_root = Path(__file__).resolve().parents[3]
-    upstream_src = repo_root / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        pytest.skip(
-            "Upstream daVinci-MagiHuman/ clone missing. Run "
-            "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`"
-        )
-
-    base_shard_dir = _find_base_shard_dir()
-    if base_shard_dir is None or not base_shard_dir.is_dir():
-        pytest.skip(
-            "GAIR/daVinci-MagiHuman base/ shards not available locally. "
-            "Set MAGI_HUMAN_BASE_SHARD_DIR or run the conversion once to "
-            "populate the HF cache."
-        )
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_DIFFUSERS_PATH",
-        repo_root / "converted_weights" / "magi_human_base",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    if not transformer_dir.is_dir():
-        pytest.skip(
-            f"Converted transformer dir missing at {transformer_dir}. Run "
-            f"scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py first."
-        )
-
-    # Add upstream to sys.path and install compiler/distributed stubs.
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs,
-        load_upstream_dit,
-    )
-    install_stubs()
-
-    # --- Shared inputs (deliberately small) ---
-    device = torch.device("cuda:0")
-    torch.manual_seed(0)
-
-    # Mirror MagiDataProxy.process_input for a tiny frame:
-    # video_latent: [1, z_dim, T, H, W], T=2, H=6, W=6, z_dim=48
-    # -> video tokens: (T/pT)*(H/pH)*(W/pW) with patch=(1,2,2) = 2*3*3 = 18
-    # audio tokens: 4
-    # text tokens: 8
-    # max channel width = 192 (video)
-    z_dim = 48
-    pT, pH, pW = 1, 2, 2
-    lat_T, lat_H, lat_W = 2, 6, 6
-    video_latent = torch.randn(
-        (1, z_dim, lat_T, lat_H, lat_W),
-        dtype=torch.float32, device=device,
-    )
-    num_video_tokens = (lat_T // pT) * (lat_H // pH) * (lat_W // pW)  # 18
-    num_audio_tokens = 4
-    num_text_tokens = 8
-    audio_latent = torch.randn(
-        (1, num_audio_tokens, 64),
-        dtype=torch.float32, device=device,
-    )
-    text_feat = torch.randn(
-        (1, num_text_tokens, 3584),
-        dtype=torch.float32, device=device,
-    )
-
-    # --- Build the packed inputs the DiT consumes ---
-    from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-        build_packed_inputs,
-    )
-    from fastvideo.models.dits.magi_human import Modality  # noqa: F401
-
-    x, coords, mm = build_packed_inputs(
-        video_latent=video_latent,
-        audio_latent=audio_latent,
-        audio_feat_len=num_audio_tokens,
-        txt_feat=text_feat,
-        txt_feat_len=num_text_tokens,
-        patch_size=(pT, pH, pW),
-        coords_style="v2",
-    )
-    assert x.shape[0] == num_video_tokens + num_audio_tokens + num_text_tokens
-
-    total_tokens = x.shape[0]
-
-    # --- Load upstream DiT first (so we know weights round-trip cleanly).
-    #     Upstream reads raw base/ shards; we keep it in bf16 for speed
-    #     and because that matches the FastVideo side after FSDP load.
-    print("Loading upstream DiTModel from base shards...")
-    upstream_model = load_upstream_dit(
-        base_shard_dir,
-        device=device,
-        dtype=None,  # keep checkpoint dtypes (fp32 for norms, bf16 for matmuls)
-    )
-
-    # --- VarlenHandler for upstream (batch=1, total_tokens).
-    from inference.common import VarlenHandler
-    cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device)
-    varlen = VarlenHandler(
-        cu_seqlens_q=cu,
-        cu_seqlens_k=cu,
-        max_seqlen_q=total_tokens,
-        max_seqlen_k=total_tokens,
-    )
-
-    # --- Forward upstream and capture output. ---
-    print("Running upstream forward...")
-    with torch.inference_mode():
-        ref_out = upstream_model(
-            x=x.clone(),
-            coords_mapping=coords.clone(),
-            modality_mapping=mm.clone(),
-            varlen_handler=varlen,
-            local_attn_handler=None,  # local_attn_layers=[] for base
-        ).detach().float().cpu()
-
-    # Free upstream model before loading FastVideo (saves ~30 GB on GPU).
-    del upstream_model
-    _cleanup_gpu()
-
-    # --- Load FastVideo MagiHumanDiT from the converted transformer/ ---
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-    import glob
-    print("Loading FastVideo MagiHumanDiT from converted transformer/...")
-    fv_cfg = MagiHumanVideoConfig()
-    fv_model = MagiHumanDiT(fv_cfg)
-
-    fv_state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        fv_state.update(load_file(shard))
-    missing, unexpected = fv_model.load_state_dict(fv_state, strict=False)
-    assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}"
-    assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
-
-    fv_model = fv_model.to(device=device)
-    fv_model.eval()
-
-    print("Running FastVideo forward...")
-    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
-        fv_out = fv_model(x.clone(), coords.clone(), mm.clone()).detach().float().cpu()
-
-    # Global stats
-    print(
-        f"ref sum={ref_out.sum().item():.4f} "
-        f"abs_mean={ref_out.abs().mean().item():.4f} "
-        f"shape={tuple(ref_out.shape)}"
-    )
-    print(
-        f"fv  sum={fv_out.sum().item():.4f} "
-        f"abs_mean={fv_out.abs().mean().item():.4f} "
-        f"shape={tuple(fv_out.shape)}"
-    )
-    diff = (ref_out - fv_out).abs()
-    print(
-        f"diff max={diff.max().item():.6f} "
-        f"mean={diff.mean().item():.6f} "
-        f"median={diff.median().item():.6f}"
-    )
-
-    # Per-modality diagnostic (video, audio, text). Text rows are zero-
-    # padded on both sides; video and audio should carry comparable
-    # abs_mean.
-    ref_video = ref_out[:num_video_tokens]
-    fv_video = fv_out[:num_video_tokens]
-    ref_audio = ref_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
-    fv_audio = fv_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
-    ref_text = ref_out[num_video_tokens + num_audio_tokens:]
-    fv_text = fv_out[num_video_tokens + num_audio_tokens:]
-    video_diff = (ref_video - fv_video).abs()
-    audio_diff = (ref_audio - fv_audio).abs()
-    text_diff = (ref_text - fv_text).abs()
-    print(
-        f"video  ref_abs={ref_video.abs().mean():.4f} "
-        f"diff_max={video_diff.max():.4f} diff_mean={video_diff.mean():.4f}"
-    )
-    print(
-        f"audio  ref_abs={ref_audio.abs().mean():.4f} "
-        f"diff_max={audio_diff.max():.4f} diff_mean={audio_diff.mean():.4f}"
-    )
-    print(
-        f"text   ref_abs={ref_text.abs().mean():.4f} "
-        f"diff_max={text_diff.max():.4f} diff_mean={text_diff.mean():.4f}"
-    )
-
-    # --- Assertions ---
-    assert ref_out.shape == fv_out.shape, (
-        f"shape mismatch: ref={ref_out.shape} fv={fv_out.shape}"
-    )
-
-    # Text rows are zero-padded on both sides — must match exactly.
-    assert_close(fv_text, ref_text, atol=1e-6, rtol=1e-6)
-
-    # Video + audio: bf16 single-forward DiT noise floor is ~1e-3 to
-    # 5e-3 per element. atol=0.03 catches gross structural bugs
-    # (permutation flips, sign inversions, wrong modality dispatch,
-    # missing sub-layers) while leaving 6-10x margin over actual bf16
-    # noise. Observed diff_max=0.057 will FAIL — that is the bug
-    # surfacing and is the intended spec for downstream root-cause
-    # investigation.
-    assert_close(fv_out, ref_out, atol=0.03, rtol=0.01)
-
-    # Sanity: mean magnitudes should match within 5%. A gross bug
-    # (e.g. dropping a modality branch) would show up here.
-    ref_abs = ref_out.abs().mean().item()
-    fv_abs = fv_out.abs().mean().item()
-    rel = abs(ref_abs - fv_abs) / max(ref_abs, 1e-6)
-    assert rel < 0.05, f"abs_mean drift {rel:.3%} > 5% — possible structural bug"

tests/local_tests/magi_human/test_magi_human_pipeline_parity.py

@@ -1,532 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""End-to-end latent parity test for the daVinci-MagiHuman base text-to-AV pipeline.
-
-Runs the joint video+audio FlowUniPC denoise loop with CFG=2 on both:
-  - FastVideo MagiHumanDiT loaded from `converted_weights/magi_human_base/transformer/`.
-  - Upstream daVinci-MagiHuman DiTModel loaded from the HF `base/` shards,
-    via the `magi_compiler` / distributed stubs in
-    `tests/local_tests/helpers/magi_human_upstream.py`.
-
-Both sides use the **same** `FlowUniPCMultistepScheduler` (FastVideo's
-implementation), identical latent / text inputs, identical scheduler
-state, and SDPA-routed attention — so drift here is purely the
-compound of per-call DiT parity drift through the denoise loop + CFG
-mixing amplification.
-
-What this catches (that the component-level DiT parity does NOT):
-  - Scheduler integration mistakes (state leaks between video/audio
-    schedulers, wrong shift, wrong `step()` args).
-  - CFG math errors (guidance scale switchover at t=500, per-modality
-    guidance scale wiring, unconditional-path text padding).
-  - Latent-preparation / token-unpacking drift between my
-    `build_packed_inputs` / `unpack_tokens` and the upstream
-    `MagiDataProxy` equivalents.
-  - Compounding behavior: 1% per-call DiT drift compounding through
-    `num_steps * cfg_number` calls.
-
-Skips when:
-  - `daVinci-MagiHuman/` clone or GAIR/daVinci-MagiHuman base shards
-    are not available locally.
-  - Converted transformer weights are missing (run the conversion
-    script first).
-  - CUDA is unavailable.
-
-Tolerance: `atol=0.35, rtol=0.05` on bf16 denoise-loop latents. The
-atol absorbs the observed worst-element drift (~0.31 on a signal of
-abs_mean ~2.4 — bf16 + CFG amplification + UniPC accumulation). The
-tight rtol still flags gross structural bugs (sign flip, scheduler
-state leak, modality branch drop). If tighter parity is wanted,
-chase the per-call drift first (see the DiT component parity test).
-"""
-from __future__ import annotations
-
-import gc
-import glob
-import os
-from pathlib import Path
-
-import pytest
-import torch
-import torch.nn.functional as F
-from torch.testing import assert_close
-
-
-# Force SDPA on both sides so the attention kernel is shared.
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29519")
-
-_T5GEMMA_ID = os.getenv("MAGI_HUMAN_T5GEMMA_ID", "google/t5gemma-9b-9b-ul2")
-_T5_GEMMA_TARGET_LENGTH = 640
-_SAMPLE_PROMPT = (
-    "A warm afternoon scene: a person sits on a park bench reading a book, "
-    "surrounded by softly swaying trees."
-)
-
-
-def _hf_token() -> str | None:
-    for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        token = os.environ.get(key)
-        if token:
-            return token
-    return None
-
-
-def _can_access_t5gemma() -> bool:
-    token = _hf_token()
-    if token is None:
-        return False
-    try:
-        from huggingface_hub import hf_hub_download
-        hf_hub_download(
-            repo_id=_T5GEMMA_ID,
-            filename="config.json",
-            token=token,
-        )
-        return True
-    except Exception:
-        return False
-
-
-def _pad_or_trim_dim1(t: torch.Tensor, target: int) -> tuple[torch.Tensor, int]:
-    """Mirror MagiHumanLatentPreparationStage's text pad-or-trim."""
-    current = t.size(1)
-    if current < target:
-        pad = [0, 0, 0, target - current]
-        return F.pad(t, pad, "constant", 0.0), current
-    return t[:, :target], target
-
-
-def _encode_magi_human_prompt_pair(device: torch.device):
-    """Encode the production preset prompt pair once via T5-Gemma."""
-    if not _can_access_t5gemma():
-        pytest.skip(
-            f"{_T5GEMMA_ID} not accessible — gated Google repo; set "
-            "HF_TOKEN / HF_API_KEY and accept the terms of use."
-        )
-
-    for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        token = os.environ.get(src)
-        if token:
-            os.environ.setdefault("HF_TOKEN", token)
-            os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", token)
-            break
-
-    try:
-        from transformers import AutoTokenizer
-
-        from fastvideo.configs.models.encoders.t5gemma import (
-            T5GemmaEncoderConfig,
-        )
-        from fastvideo.models.encoders.t5gemma import (
-            T5GemmaEncoderModel,
-        )
-        from fastvideo.pipelines.basic.magi_human.presets import (
-            _MAGI_HUMAN_NEGATIVE_PROMPT,
-        )
-    except Exception as exc:
-        pytest.skip(f"T5-Gemma prompt encoding dependencies unavailable: {exc}")
-
-    tokenizer = AutoTokenizer.from_pretrained(_T5GEMMA_ID)
-    enc_config = T5GemmaEncoderConfig()
-    enc_config.arch_config.t5gemma_model_path = _T5GEMMA_ID
-    encoder = T5GemmaEncoderModel(enc_config)
-
-    def encode(text: str, text_encoder=encoder) -> tuple[torch.Tensor, int]:
-        inputs = tokenizer(
-            [text],
-            return_tensors="pt",
-            padding=True,
-            truncation=False,
-        ).to(device)
-        with torch.inference_mode():
-            hidden = text_encoder(
-                input_ids=inputs["input_ids"],
-                attention_mask=inputs.get("attention_mask"),
-            ).last_hidden_state
-        return _pad_or_trim_dim1(hidden.to(torch.float32), _T5_GEMMA_TARGET_LENGTH)
-
-    txt_feat, txt_feat_len = encode(_SAMPLE_PROMPT)
-    neg_txt_feat, neg_txt_feat_len = encode(_MAGI_HUMAN_NEGATIVE_PROMPT)
-    del encoder
-    _cleanup_gpu()
-    return txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len
-
-
-def _find_base_shard_dir() -> Path | None:
-    """Return the local path to GAIR/daVinci-MagiHuman/base/ with shards present, or None."""
-    override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR")
-    if override:
-        p = Path(override)
-        return p if p.is_dir() else None
-    try:
-        from huggingface_hub import snapshot_download
-        snap = snapshot_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            allow_patterns=[
-                "base/*.safetensors",
-                "base/model.safetensors.index.json",
-            ],
-        )
-        candidate = Path(snap) / "base"
-        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
-            return candidate
-        return None
-    except Exception:
-        return None
-
-
-def _cleanup_gpu() -> None:
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-
-def _dit_forward_fv(
-    dit, video_latent, audio_latent, audio_feat_len,
-    txt_feat, txt_feat_len, patch_size, coords_style,
-    video_in_channels, audio_in_channels,
-):
-    """One FastVideo DiT call — same as MagiHumanDenoisingStage._dit_forward."""
-    from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-        build_packed_inputs, unpack_tokens,
-    )
-    x, coords, mm = build_packed_inputs(
-        video_latent=video_latent, audio_latent=audio_latent,
-        audio_feat_len=audio_feat_len, txt_feat=txt_feat,
-        txt_feat_len=txt_feat_len, patch_size=patch_size,
-        coords_style=coords_style,
-    )
-    video_token_num = x.shape[0] - audio_feat_len - txt_feat_len
-    out = dit(x, coords, mm)
-    return unpack_tokens(
-        out, video_token_num=video_token_num,
-        audio_feat_len=audio_feat_len,
-        video_in_channels=video_in_channels,
-        audio_in_channels=audio_in_channels,
-        latent_shape=tuple(video_latent.shape),
-        patch_size=patch_size,
-    )
-
-
-def _dit_forward_upstream(
-    dit, video_latent, audio_latent, audio_feat_len,
-    txt_feat, txt_feat_len, patch_size, coords_style,
-    video_in_channels, audio_in_channels,
-):
-    """One upstream DiT call — identical input construction + output
-    unpacking to the FastVideo path. The only thing that differs is
-    the DiT module and the extra `varlen_handler` / `local_attn_handler`
-    kwargs the upstream expects.
-    """
-    from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-        build_packed_inputs, unpack_tokens,
-    )
-    from inference.common import VarlenHandler
-    x, coords, mm = build_packed_inputs(
-        video_latent=video_latent, audio_latent=audio_latent,
-        audio_feat_len=audio_feat_len, txt_feat=txt_feat,
-        txt_feat_len=txt_feat_len, patch_size=patch_size,
-        coords_style=coords_style,
-    )
-    video_token_num = x.shape[0] - audio_feat_len - txt_feat_len
-    total = x.shape[0]
-    cu = torch.tensor([0, total], dtype=torch.int32, device=x.device)
-    varlen = VarlenHandler(
-        cu_seqlens_q=cu, cu_seqlens_k=cu,
-        max_seqlen_q=total, max_seqlen_k=total,
-    )
-    out = dit(
-        x=x, coords_mapping=coords, modality_mapping=mm,
-        varlen_handler=varlen, local_attn_handler=None,
-    )
-    return unpack_tokens(
-        out, video_token_num=video_token_num,
-        audio_feat_len=audio_feat_len,
-        video_in_channels=video_in_channels,
-        audio_in_channels=audio_in_channels,
-        latent_shape=tuple(video_latent.shape),
-        patch_size=patch_size,
-    )
-
-
-def _build_fastvideo_schedulers(shift: float, num_inference_steps: int, device):
-    """Mirror current FastVideo production at `magi_human_pipeline.py:146-149`
-    and `denoising.py:105-116`: default scheduler constructor (`shift=1`,
-    no-op) followed by `set_timesteps(..., shift=shift)` so the temporal
-    shift is applied exactly once. The earlier double-shift pattern was
-    reverted with the Wave 11 single-shift fix; if both __init__ and
-    set_timesteps applied non-trivial shift, the schedule would diverge
-    from upstream.
-    """
-    from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
-        FlowUniPCMultistepScheduler,
-    )
-    video_sched = FlowUniPCMultistepScheduler()
-    audio_sched = FlowUniPCMultistepScheduler()
-    video_sched.set_timesteps(num_inference_steps, device=device, shift=shift)
-    audio_sched.set_timesteps(num_inference_steps, device=device, shift=shift)
-    return video_sched, audio_sched
-
-
-def _build_upstream_schedulers(shift: float, num_inference_steps: int, device):
-    """Construct schedulers the way the official `MagiEvaluator.eval_with_text`
-    does (`daVinci-MagiHuman/inference/pipeline/video_generate.py:404-407`):
-    `FlowUniPCMultistepScheduler()` with default shift=1.0 in __init__
-    (no-op), then `set_timesteps(num_inference_steps, device, shift=self.shift)`
-    applies shift exactly once. Uses FastVideo's scheduler class for
-    the orchestration (algorithmically identical to the upstream copy
-    of the same Diffusers-derived class) but matches the upstream's
-    *call pattern*.
-    """
-    from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
-        FlowUniPCMultistepScheduler,
-    )
-    video_sched = FlowUniPCMultistepScheduler()
-    audio_sched = FlowUniPCMultistepScheduler()
-    video_sched.set_timesteps(num_inference_steps, device=device, shift=shift)
-    audio_sched.set_timesteps(num_inference_steps, device=device, shift=shift)
-    return video_sched, audio_sched
-
-
-def _run_denoise_loop(
-    dit, dit_forward_fn, video_latent, audio_latent,
-    txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len,
-    *, video_sched, audio_sched, cfg_number,
-    video_txt_guidance_scale, audio_txt_guidance_scale,
-    patch_size, coords_style, video_in_channels, audio_in_channels,
-    image_latent=None,
-):
-    """Joint video+audio FlowUniPC denoise. The schedulers are passed
-    in pre-constructed so each side can mirror its production scheduler
-    init pattern (see `_build_*_schedulers`).
-    """
-    from fastvideo.forward_context import set_forward_context
-    audio_feat_len = int(audio_latent.shape[1])
-
-    with torch.inference_mode():
-        for idx, t in enumerate(video_sched.timesteps):
-            if image_latent is not None:
-                video_latent[:, :, :1] = image_latent.to(
-                    device=video_latent.device,
-                    dtype=video_latent.dtype,
-                )[:, :, :1]
-            t_int = int(t.item()) if torch.is_tensor(t) else int(t)
-            with set_forward_context(current_timestep=t_int, attn_metadata=None):
-                v_cond_video, v_cond_audio = dit_forward_fn(
-                    dit, video_latent, audio_latent, audio_feat_len,
-                    txt_feat, txt_feat_len, patch_size, coords_style,
-                    video_in_channels, audio_in_channels,
-                )
-                if cfg_number == 2:
-                    v_uncond_video, v_uncond_audio = dit_forward_fn(
-                        dit, video_latent, audio_latent, audio_feat_len,
-                        neg_txt_feat, neg_txt_feat_len, patch_size, coords_style,
-                        video_in_channels, audio_in_channels,
-                    )
-                    # Upstream's video-guidance drop-at-t<=500 trick.
-                    video_guidance = (
-                        video_txt_guidance_scale if t > 500 else 2.0
-                    )
-                    v_video = v_uncond_video + video_guidance * (
-                        v_cond_video - v_uncond_video
-                    )
-                    v_audio = v_uncond_audio + audio_txt_guidance_scale * (
-                        v_cond_audio - v_uncond_audio
-                    )
-                else:
-                    v_video = v_cond_video
-                    v_audio = v_cond_audio
-
-            video_latent = video_sched.step(
-                v_video, t, video_latent, return_dict=False,
-            )[0]
-            audio_latent = audio_sched.step(
-                v_audio, t, audio_latent, return_dict=False,
-            )[0]
-        if image_latent is not None:
-            video_latent[:, :, :1] = image_latent.to(
-                device=video_latent.device,
-                dtype=video_latent.dtype,
-            )[:, :, :1]
-    return video_latent, audio_latent
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman pipeline parity requires CUDA.",
-)
-def test_magi_human_pipeline_latent_parity():
-    repo_root = Path(__file__).resolve().parents[3]
-    upstream_src = repo_root / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        pytest.skip(
-            "Upstream daVinci-MagiHuman/ clone missing. Run "
-            "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`"
-        )
-
-    base_shard_dir = _find_base_shard_dir()
-    if base_shard_dir is None or not base_shard_dir.is_dir():
-        pytest.skip(
-            "GAIR/daVinci-MagiHuman base/ shards not available locally."
-        )
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_DIFFUSERS_PATH",
-        repo_root / "converted_weights" / "magi_human_base",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    if not transformer_dir.is_dir():
-        pytest.skip(f"Converted transformer dir missing at {transformer_dir}")
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs, load_upstream_dit,
-    )
-    install_stubs()
-
-    # --- Shared pipeline inputs ---
-    device = torch.device("cuda:0")
-    torch.manual_seed(0)
-
-    # Deliberately tiny so 2 * CFG=2 = 4 DiT calls per side fit in
-    # CI/dev runtime budget.
-    z_dim = 48
-    patch_size = (1, 2, 2)
-    lat_T, lat_H, lat_W = 2, 6, 6
-    video_latent = torch.randn(
-        (1, z_dim, lat_T, lat_H, lat_W),
-        dtype=torch.float32, device=device,
-    )
-    audio_latent = torch.randn(
-        (1, 4, 64), dtype=torch.float32, device=device,
-    )
-    # Production-facing text embeddings: encode the example prompt and the
-    # preset negative prompt via T5-Gemma once, then feed the identical cached
-    # tensors to upstream and FastVideo. This keeps the DiT comparison focused
-    # while still validating prompt/preset content such as the full
-    # three-block MagiHuman negative prompt.
-    txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len = (
-        _encode_magi_human_prompt_pair(device)
-    )
-
-    num_inference_steps = 4        # 4 steps × CFG=2 = 8 DiT calls / side; surfaces compounding drift that 1-step hides
-    shift = 5.0
-    common_kwargs = dict(
-        cfg_number=2,
-        video_txt_guidance_scale=5.0,
-        audio_txt_guidance_scale=5.0,
-        patch_size=patch_size,
-        coords_style="v2",
-        video_in_channels=192,
-        audio_in_channels=64,
-    )
-
-    # --- Upstream side first (so we can free it before loading FastVideo). ---
-    # Upstream uses single-shift scheduler init (matches MagiEvaluator).
-    up_video_sched, up_audio_sched = _build_upstream_schedulers(
-        shift=shift, num_inference_steps=num_inference_steps, device=device,
-    )
-    print("Loading upstream DiTModel from base shards...")
-    upstream_dit = load_upstream_dit(base_shard_dir, device=device, dtype=None)
-    print("Running upstream denoise loop...")
-    ref_video, ref_audio = _run_denoise_loop(
-        upstream_dit, _dit_forward_upstream,
-        video_latent.clone(), audio_latent.clone(),
-        txt_feat.clone(), txt_feat_len,
-        neg_txt_feat.clone(), neg_txt_feat_len,
-        video_sched=up_video_sched, audio_sched=up_audio_sched,
-        **common_kwargs,
-    )
-    ref_video = ref_video.detach().float().cpu()
-    ref_audio = ref_audio.detach().float().cpu()
-    del upstream_dit
-    _cleanup_gpu()
-
-    # --- FastVideo side ---
-    # FastVideo uses double-shift scheduler init (matches
-    # `MagiHumanDenoisingStage` in production: shift in __init__ via
-    # `magi_human_pipeline.initialize_pipeline` AND in set_timesteps).
-    fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers(
-        shift=shift, num_inference_steps=num_inference_steps, device=device,
-    )
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-
-    print("Loading FastVideo MagiHumanDiT from converted transformer/...")
-    fv_cfg = MagiHumanVideoConfig()
-    fv_dit = MagiHumanDiT(fv_cfg)
-    fv_state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        fv_state.update(load_file(shard))
-    missing, unexpected = fv_dit.load_state_dict(fv_state, strict=False)
-    assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}"
-    assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
-    fv_dit = fv_dit.to(device=device)
-    fv_dit.eval()
-
-    print("Running FastVideo denoise loop...")
-    fv_video, fv_audio = _run_denoise_loop(
-        fv_dit, _dit_forward_fv,
-        video_latent.clone(), audio_latent.clone(),
-        txt_feat.clone(), txt_feat_len,
-        neg_txt_feat.clone(), neg_txt_feat_len,
-        video_sched=fv_video_sched, audio_sched=fv_audio_sched,
-        **common_kwargs,
-    )
-    fv_video = fv_video.detach().float().cpu()
-    fv_audio = fv_audio.detach().float().cpu()
-
-    # --- Report + assertions ---
-    v_diff = (ref_video - fv_video).abs()
-    a_diff = (ref_audio - fv_audio).abs()
-    print(
-        f"video  ref_abs={ref_video.abs().mean().item():.4f} "
-        f"fv_abs={fv_video.abs().mean().item():.4f} "
-        f"diff_max={v_diff.max().item():.4f} "
-        f"diff_mean={v_diff.mean().item():.4f} "
-        f"diff_median={v_diff.median().item():.4f}"
-    )
-    print(
-        f"audio  ref_abs={ref_audio.abs().mean().item():.4f} "
-        f"fv_abs={fv_audio.abs().mean().item():.4f} "
-        f"diff_max={a_diff.max().item():.4f} "
-        f"diff_mean={a_diff.mean().item():.4f} "
-        f"diff_median={a_diff.median().item():.4f}"
-    )
-
-    assert ref_video.shape == fv_video.shape
-    assert ref_audio.shape == fv_audio.shape
-
-    # Tolerance budget for 1-step / CFG=2 (bf16 DiT + bf16 CFG mix):
-    #   * Single-DiT bf16 drift: diff_mean ~0.008 on `abs ~ 1.0`
-    #     (see DiT component parity, `test_magi_human_dit_parity`).
-    #   * CFG mixes `v = v_uncond + guidance * (v_cond - v_uncond)`
-    #     with guidance=5; cond and uncond drift independently in bf16,
-    #     so the post-CFG `diff_mean` scales by ~guidance (~5x).
-    #   * One FlowUniPC scheduler step passes that through unchanged.
-    # `diff_max` is the noisiest statistic for bf16 transformer parity
-    # (a single fma quantization can blow it up). Use it only as a loose
-    # guard. The two ratio guards below catch real structural bugs:
-    # `abs_mean` drift signals scale errors / dropped branches, and
-    # `diff_mean / ref_abs` signals systematic per-element bias far
-    # beyond what bf16+CFG noise can produce.
-    assert_close(fv_video, ref_video, atol=0.40, rtol=0.05)
-    assert_close(fv_audio, ref_audio, atol=0.40, rtol=0.05)
-
-    # Global-magnitude guard — tightest single assertion. A gross bug
-    # (scheduler state leak, dropped modality branch, CFG sign flip)
-    # would shift `abs_mean` far beyond the bf16+CFG noise floor.
-    ref_v_abs = ref_video.abs().mean().item()
-    ref_a_abs = ref_audio.abs().mean().item()
-    rel_v = abs(ref_v_abs - fv_video.abs().mean().item()) / max(ref_v_abs, 1e-6)
-    rel_a = abs(ref_a_abs - fv_audio.abs().mean().item()) / max(ref_a_abs, 1e-6)
-    assert rel_v < 0.01, f"video abs_mean drift {rel_v:.2%} > 1%"
-    assert rel_a < 0.01, f"audio abs_mean drift {rel_a:.2%} > 1%"
-
-    # Per-element mean-bias guard — catches systematic shift that
-    # `abs_mean` misses (e.g. equal-magnitude flip across many elements).
-    mean_rel_v = v_diff.mean().item() / max(ref_v_abs, 1e-6)
-    mean_rel_a = a_diff.mean().item() / max(ref_a_abs, 1e-6)
-    assert mean_rel_v < 0.04, f"video mean_diff/ref_abs {mean_rel_v:.2%} > 4%"
-    assert mean_rel_a < 0.04, f"audio mean_diff/ref_abs {mean_rel_a:.2%} > 4%"

tests/local_tests/magi_human/test_magi_human_pipeline_smoke.py

@@ -1,246 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Smoke / preflight tests for the daVinci-MagiHuman base text-to-AV pipeline.
-
-Two tests:
-
-  * `test_magi_human_typed_surface_preflight` — pure-Python, no GPU, no
-    weights. Verifies that the scaffold is importable, that the preset
-    registers cleanly, and that the DiT module tree matches the upstream
-    HuggingFace checkpoint shape-for-shape on `meta` device. This is what
-    CI should run on every PR.
-
-  * `test_magi_human_pipeline_smoke` — end-to-end pipeline construction +
-    a tiny generate_video call, gated on local converted-weights paths.
-    Skips cleanly when weights are missing.
-"""
-from __future__ import annotations
-
-import json
-import os
-from pathlib import Path
-
-import pytest
-import torch
-
-
-def test_magi_human_typed_surface_preflight() -> None:
-    """Import + registry + module-tree surface check.
-
-    Covers regressions that would otherwise only surface on a GPU host:
-    preset drop from ALL_PRESETS, renamed modules, registry mis-wiring,
-    or DiT module-tree drift from the upstream checkpoint.
-    """
-    import fastvideo.registry  # noqa: F401 — triggers preset registration
-    from fastvideo.api.presets import get_preset, get_presets_for_family
-    from fastvideo.configs.models.dits.magi_human import (
-        MagiHumanArchConfig,
-        MagiHumanVideoConfig,
-    )
-    from fastvideo.configs.models.encoders.t5gemma import (
-        T5GemmaEncoderArchConfig,
-        T5GemmaEncoderConfig,
-    )
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from fastvideo.pipelines.basic.magi_human.magi_human_pipeline import (  # noqa: F401
-        MagiHumanI2VPipeline,
-        MagiHumanPipeline,
-        MagiHumanSRI2VPipeline,
-        MagiHumanSRPipeline,
-    )
-    from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-        MagiHumanBaseConfig,
-        MagiHumanBaseI2VConfig,
-        MagiHumanDistillI2VConfig,
-        MagiHumanSR540pConfig,
-        MagiHumanSR540pI2VConfig,
-    )
-    from fastvideo.pipelines.basic.magi_human.stages import (  # noqa: F401
-        MagiHumanDenoisingStage,
-        MagiHumanLatentPreparationStage,
-        MagiHumanReferenceImageStage,
-        MagiHumanSRDenoisingStage,
-        MagiHumanSRLatentPreparationStage,
-    )
-
-    # Presets are registered under the expected family.
-    names = {p.name for p in get_presets_for_family("magi_human")}
-    assert names == {
-        "magi_human_base",
-        "magi_human_distill",
-        "magi_human_base_ti2v",
-        "magi_human_distill_ti2v",
-        "magi_human_sr_540p",
-        "magi_human_sr_540p_ti2v",
-        "magi_human_sr_1080p",
-        "magi_human_sr_1080p_ti2v",
-    }
-
-    base_preset = get_preset("magi_human_base", "magi_human")
-    assert base_preset.workload_type == "t2v"
-    assert base_preset.defaults["num_inference_steps"] == 32
-    assert base_preset.defaults["fps"] == 25
-
-    distill_preset = get_preset("magi_human_distill", "magi_human")
-    assert distill_preset.workload_type == "t2v"
-    assert distill_preset.defaults["num_inference_steps"] == 8
-    assert distill_preset.defaults["guidance_scale"] == 1.0
-
-    base_ti2v_preset = get_preset("magi_human_base_ti2v", "magi_human")
-    assert base_ti2v_preset.workload_type == "i2v"
-    assert base_ti2v_preset.defaults["num_inference_steps"] == 32
-
-    distill_ti2v_preset = get_preset("magi_human_distill_ti2v", "magi_human")
-    assert distill_ti2v_preset.workload_type == "i2v"
-    assert distill_ti2v_preset.defaults["num_inference_steps"] == 8
-
-    sr_preset = get_preset("magi_human_sr_540p", "magi_human")
-    assert sr_preset.workload_type == "t2v"
-    assert sr_preset.defaults["num_inference_steps"] == 32
-
-    sr_ti2v_preset = get_preset("magi_human_sr_540p_ti2v", "magi_human")
-    assert sr_ti2v_preset.workload_type == "i2v"
-    assert sr_ti2v_preset.defaults["num_inference_steps"] == 32
-
-    # Distill pipeline config: same arch as base, CFG=1, 8 steps.
-    from fastvideo.pipelines.basic.magi_human.pipeline_configs import MagiHumanDistillConfig
-    distill_pc = MagiHumanDistillConfig()
-    assert distill_pc.num_inference_steps == 8
-    assert distill_pc.cfg_number == 1
-    assert distill_pc.dit_config.arch_config.num_layers == 40
-
-    base_i2v_pc = MagiHumanBaseI2VConfig()
-    assert base_i2v_pc.image_conditioning is True
-    assert base_i2v_pc.vae_config.load_encoder is True
-    assert base_i2v_pc.vae_config.load_decoder is True
-
-    distill_i2v_pc = MagiHumanDistillI2VConfig()
-    assert distill_i2v_pc.num_inference_steps == 8
-    assert distill_i2v_pc.cfg_number == 1
-    assert distill_i2v_pc.image_conditioning is True
-    assert distill_i2v_pc.vae_config.load_encoder is True
-
-    sr_pc = MagiHumanSR540pConfig()
-    assert sr_pc.num_inference_steps == 32
-    assert sr_pc.sr_num_inference_steps == 5
-    assert sr_pc.noise_value == 220
-    assert sr_pc.sr_audio_noise_scale == 0.7
-    assert sr_pc.sr_video_txt_guidance_scale == 3.5
-    assert sr_pc.sr_height == 512
-    assert sr_pc.sr_width == 896
-
-    sr_i2v_pc = MagiHumanSR540pI2VConfig()
-    assert sr_i2v_pc.image_conditioning is True
-    assert sr_i2v_pc.vae_config.load_encoder is True
-
-    # Config constructs with the documented defaults.
-    pc = MagiHumanBaseConfig()
-    assert pc.flow_shift == 5.0
-    assert pc.cfg_number == 2
-    assert pc.num_inference_steps == 32
-    assert pc.dit_config.arch_config.num_layers == 40
-    assert pc.dit_config.arch_config.hidden_size == 5120
-    assert pc.dit_config.arch_config.num_attention_heads == 40
-    assert pc.dit_config.arch_config.num_heads_kv == 8
-    assert pc.dit_config.arch_config.mm_layers == (0, 1, 2, 3, 36, 37, 38, 39)
-    assert pc.text_encoder_configs[0].arch_config.hidden_size == 3584
-
-    # The DiT module tree matches the upstream HF base/ checkpoint
-    # shape-for-shape. This is checkpoint-loading parity, not numerical
-    # parity — but a regression here means loaded weights won't align.
-    dit_cfg = MagiHumanVideoConfig()
-    with torch.device("meta"):
-        dit = MagiHumanDiT(dit_cfg)
-    fv_shapes = {n: tuple(p.shape) for n, p in dit.state_dict().items()}
-
-    index_path = _hf_index_path_or_none()
-    if index_path is None:
-        pytest.skip("HF repo unavailable (no network / no token) — "
-                    "skipping cross-check against GAIR/daVinci-MagiHuman.")
-    with open(index_path) as f:
-        wmap = json.load(f)["weight_map"]
-    hf_keys = set(wmap.keys())
-    fv_keys = set(fv_shapes.keys())
-
-    missing_in_fv = sorted(hf_keys - fv_keys)
-    extra_in_fv = sorted(fv_keys - hf_keys)
-    assert not missing_in_fv, f"fastvideo missing keys: {missing_in_fv[:5]}"
-    assert not extra_in_fv, f"fastvideo extra keys: {extra_in_fv[:5]}"
-    assert len(fv_keys) == 331
-
-
-def _hf_index_path_or_none() -> str | None:
-    """Return a local path to the base/ index.json, or None if unavailable."""
-    try:
-        from huggingface_hub import hf_hub_download
-    except ImportError:
-        return None
-    try:
-        return hf_hub_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            filename="base/model.safetensors.index.json",
-        )
-    except Exception:
-        return None
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman pipeline smoke test requires CUDA.",
-)
-def test_magi_human_pipeline_smoke() -> None:
-    """End-to-end smoke: build the pipeline and run a tiny denoise.
-
-    Skips cleanly when the converted-weights directory is not present.
-    """
-    diffusers_path = os.getenv(
-        "MAGI_HUMAN_DIFFUSERS_PATH",
-        "converted_weights/magi_human_base",
-    )
-    if not os.path.isdir(diffusers_path):
-        pytest.skip(
-            f"Missing converted MagiHuman repo at {diffusers_path}. "
-            f"Run scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py "
-            f"first."
-        )
-    if not os.path.isfile(os.path.join(diffusers_path, "model_index.json")):
-        pytest.skip(f"Missing model_index.json in {diffusers_path}")
-
-    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-    from fastvideo import VideoGenerator
-
-    # Small shapes to keep the smoke test cheap.
-    prompt = "A cheerful person waving at the camera in a well-lit room."
-    seed = 42
-    height = 256
-    width = 448
-    num_frames = 13   # seconds=1, 12fps for smoke; the pipeline derives it
-    fps = 12.0
-    steps = 2
-
-    generator = VideoGenerator.from_pretrained(
-        diffusers_path,
-        num_gpus=1,
-        use_fsdp_inference=False,
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=False,
-    )
-    try:
-        result = generator.generate_video(
-            prompt=prompt,
-            output_path="outputs_video/magi_human_smoke",
-            save_video=False,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            fps=fps,
-            num_inference_steps=steps,
-            seed=seed,
-        )
-    finally:
-        generator.shutdown()
-
-    samples = result["samples"]
-    assert samples.ndim == 5, f"expected [B,C,T,H,W], got {samples.shape}"
-    assert samples.shape[0] == 1

tests/local_tests/magi_human/test_magi_human_sa_audio_official_parity.py

@@ -1,162 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Parity test: FastVideo MagiHuman Stable-Audio wrapper vs the official
-daVinci-MagiHuman Stable-Audio usage path.
-
-The sibling `test_magi_human_sa_audio_parity.py` compares FastVideo's
-`SAAudioVAEModel` against Diffusers `AutoencoderOobleck`, which validates the
-low-level VAE weights/API. This test instead follows the official
-daVinci-MagiHuman integration layer:
-
-* `inference.model.sa_audio.SAAudioFeatureExtractor` is constructed from the
-  full Stable Audio checkpoint (`model_config.json` + `model.safetensors`).
-* The official loader rebuilds its local `AudioAutoencoder` from
-  `model.pretransform.config` and filters `pretransform.model.*` weights.
-* The official decode entry point is `SAAudioFeatureExtractor.decode(latents)`,
-  which calls `vae_model.decode(latents)` directly. There is no latent
-  mean/std normalization or reference-audio injection inside this decode layer.
-* Pipeline post-processing is outside the SA module: `MagiEvaluator` transposes
-  `[B, L, C] -> [C, L]` before decode, then transposes waveform samples and
-  applies `resample_audio_sinc(..., 441 / 512)`.
-
-This catches drift between FastVideo's full SA wrapper path and the official
-repo's custom Stable-Audio wrapper/module, not just the bare Diffusers VAE.
-
-Skips when:
-  * CUDA is unavailable.
-  * `daVinci-MagiHuman/` is not checked out under the repo root.
-  * `stabilityai/stable-audio-open-1.0` is inaccessible (gated; user must have
-    accepted terms and set HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY).
-"""
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-
-_SA_AUDIO_ID = "stabilityai/stable-audio-open-1.0"
-
-
-def _repo_root() -> Path:
-    return Path(__file__).resolve().parents[3]
-
-
-def _upstream_root() -> Path:
-    return _repo_root() / "daVinci-MagiHuman"
-
-
-def _hf_token():
-    for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        value = os.environ.get(key)
-        if value:
-            return value
-    return None
-
-
-def _can_access() -> bool:
-    token = _hf_token()
-    if token is None:
-        return False
-    try:
-        from huggingface_hub import hf_hub_download
-
-        hf_hub_download(
-            repo_id=_SA_AUDIO_ID,
-            filename="model_config.json",
-            token=token,
-        )
-        return True
-    except Exception:
-        return False
-
-
-def _stable_audio_snapshot() -> str:
-    from huggingface_hub import snapshot_download
-
-    return snapshot_download(
-        repo_id=_SA_AUDIO_ID,
-        token=_hf_token(),
-        allow_patterns=["model_config.json", "model.safetensors"],
-    )
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman official Stable-Audio parity requires CUDA.",
-)
-@pytest.mark.skipif(
-    not _upstream_root().exists(),
-    reason="daVinci-MagiHuman checkout is required under the repo root.",
-)
-@pytest.mark.skipif(
-    not _can_access(),
-    reason=(
-        f"{_SA_AUDIO_ID} not accessible — gated Stability AI repo; set "
-        "HF_TOKEN / HF_API_KEY and accept the terms on "
-        f"https://huggingface.co/{_SA_AUDIO_ID}."
-    ),
-)
-def test_magi_human_sa_audio_official_decode_parity():
-    # Make sure both HF helpers and FastVideo's loader see the same token alias.
-    for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        value = os.environ.get(src)
-        if value:
-            os.environ.setdefault("HF_TOKEN", value)
-            os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", value)
-            break
-
-    device = torch.device("cuda:0")
-
-    # --- Official daVinci-MagiHuman path: custom SAAudioFeatureExtractor
-    #     rebuilds AudioAutoencoder and filters `pretransform.model.*` from
-    #     the full Stable Audio checkpoint.
-    from tests.local_tests.helpers.magi_human_upstream import install_stubs
-
-    install_stubs()
-    from inference.model.sa_audio import SAAudioFeatureExtractor
-
-    upstream_vae = SAAudioFeatureExtractor(
-        device=device,
-        model_path=_stable_audio_snapshot(),
-    )
-
-    # --- FastVideo MagiHuman wrapper path: lazy loader around the native
-    #     OobleckVAE port, exactly what the MagiHuman pipeline constructs.
-    from fastvideo.configs.models.vaes import OobleckVAEConfig
-    from fastvideo.models.vaes.sa_audio import SAAudioVAEModel
-
-    fv_config = OobleckVAEConfig()
-    fv_config.pretrained_path = _SA_AUDIO_ID
-    fv_config.pretrained_dtype = "float32"
-    fv_vae = SAAudioVAEModel(fv_config)
-
-    torch.manual_seed(0)
-    latent = torch.randn(
-        (1, fv_config.arch_config.decoder_input_channels, 8),
-        dtype=torch.float32,
-        device=device,
-    )
-
-    with torch.inference_mode():
-        upstream_out = upstream_vae.decode(latent).detach().float().cpu()
-        fv_out = fv_vae.decode(latent).detach().float().cpu()
-
-    print(
-        f"upstream shape={tuple(upstream_out.shape)} "
-        f"abs_mean={upstream_out.abs().mean().item():.6f} "
-        f"range=[{upstream_out.min().item():.4f}, "
-        f"{upstream_out.max().item():.4f}]"
-    )
-    print(
-        f"fv       shape={tuple(fv_out.shape)} "
-        f"abs_mean={fv_out.abs().mean().item():.6f} "
-        f"range=[{fv_out.min().item():.4f}, {fv_out.max().item():.4f}]"
-    )
-    diff = (upstream_out - fv_out).abs()
-    print(f"diff max={diff.max().item():.6e} mean={diff.mean().item():.6e}")
-
-    assert upstream_out.shape == fv_out.shape
-    assert_close(fv_out, upstream_out, atol=1e-5, rtol=1e-5)

tests/local_tests/magi_human/test_magi_human_sa_audio_parity.py

@@ -1,123 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Parity test: MagiHuman's audio-VAE path (FastVideo `SAAudioVAEModel`
-lazy-loader around the native `OobleckVAE` port, shared with the
-standalone Stable Audio pipeline) vs `diffusers.AutoencoderOobleck.
-from_pretrained(...)` on the Stable Audio Open 1.0 VAE.
-
-Companion to `tests/local_tests/vaes/test_oobleck_vae_parity.py`, which
-already validates `OobleckVAE` itself; this test exercises the wrapper
-layer that MagiHuman uses (lazy load, device migration, decode output
-unwrap) so wrapper-level regressions don't slip past the underlying-VAE
-parity test.
-
-Skips when:
-  * CUDA is unavailable (VAE is 156M params, small enough for CPU but
-    we keep the test GPU-only to match the pipeline's runtime).
-  * `stabilityai/stable-audio-open-1.0` is inaccessible (gated; user
-    must have accepted terms on the HF repo page).
-"""
-from __future__ import annotations
-
-import os
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-
-_SA_AUDIO_ID = "stabilityai/stable-audio-open-1.0"
-
-
-def _hf_token():
-    for k in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        v = os.environ.get(k)
-        if v:
-            return v
-    return None
-
-
-def _can_access() -> bool:
-    token = _hf_token()
-    if token is None:
-        return False
-    try:
-        from huggingface_hub import hf_hub_download
-        hf_hub_download(
-            repo_id=_SA_AUDIO_ID, filename="vae/config.json", token=token,
-        )
-        return True
-    except Exception:
-        return False
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman Stable-Audio VAE parity requires CUDA.",
-)
-@pytest.mark.skipif(
-    not _can_access(),
-    reason=(f"{_SA_AUDIO_ID} not accessible — gated Stability AI repo; "
-            "set HF_TOKEN / HF_API_KEY and accept the terms on "
-            f"https://huggingface.co/{_SA_AUDIO_ID}."),
-)
-def test_magi_human_sa_audio_vae_decode_parity():
-    # Make sure HF_TOKEN is the alias the Diffusers loader actually reads.
-    for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        v = os.environ.get(src)
-        if v:
-            os.environ.setdefault("HF_TOKEN", v)
-            os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", v)
-            break
-
-    device = torch.device("cuda:0")
-
-    # --- Reference: direct HF Diffusers call, same as upstream's
-    #     SAAudioFeatureExtractor but via the Diffusers Oobleck port. ---
-    from diffusers import AutoencoderOobleck
-    ref_vae = AutoencoderOobleck.from_pretrained(
-        _SA_AUDIO_ID, subfolder="vae", torch_dtype=torch.float32,
-    ).to(device).eval()
-
-    # --- FastVideo wrapper path (shared with the standalone Stable
-    #     Audio pipeline that landed in main: `OobleckVAEConfig` +
-    #     `SAAudioVAEModel` lazy-loader around the first-class
-    #     `OobleckVAE` port). ---
-    from fastvideo.configs.models.vaes import OobleckVAEConfig
-    from fastvideo.models.vaes.sa_audio import SAAudioVAEModel
-    fv_config = OobleckVAEConfig()
-    fv_config.pretrained_path = _SA_AUDIO_ID
-    # The default `pretrained_dtype="float16"` matches official stable-
-    # audio-tools, but this parity test runs the reference path in fp32
-    # — so override here.
-    fv_config.pretrained_dtype = "float32"
-    fv_vae = SAAudioVAEModel(fv_config)
-
-    # --- Tiny shared latent ---
-    torch.manual_seed(0)
-    # decoder_input_channels=64, latent length ~8 frames for a quick test.
-    latent = torch.randn(
-        (1, fv_config.arch_config.decoder_input_channels, 8),
-        dtype=torch.float32, device=device,
-    )
-
-    with torch.inference_mode():
-        ref_out = ref_vae.decode(latent).sample.detach().float().cpu()
-        fv_out = fv_vae.decode(latent).detach().float().cpu()
-
-    print(
-        f"ref shape={tuple(ref_out.shape)} "
-        f"abs_mean={ref_out.abs().mean().item():.6f} "
-        f"range=[{ref_out.min().item():.4f}, {ref_out.max().item():.4f}]"
-    )
-    print(
-        f"fv  shape={tuple(fv_out.shape)} "
-        f"abs_mean={fv_out.abs().mean().item():.6f} "
-        f"range=[{fv_out.min().item():.4f}, {fv_out.max().item():.4f}]"
-    )
-    diff = (ref_out - fv_out).abs()
-    print(f"diff max={diff.max().item():.6e} mean={diff.mean().item():.6e}")
-
-    assert ref_out.shape == fv_out.shape
-    # Both sides call the same HF class on the same weights in fp32 —
-    # should agree to machine epsilon.
-    assert_close(fv_out, ref_out, atol=1e-5, rtol=1e-5)

tests/local_tests/magi_human/test_magi_human_sr1080p_pipeline_parity.py

@@ -1,340 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""SR-1080p local-window latent-loop parity for daVinci-MagiHuman.
-
-This mirrors the SR-540p two-stage parity test but enables upstream's
-SR2_1080 local-attention layer set on the SR DiT. The reference side uses the
-test helper's SDPA implementation of FFAHandler's segmented accumulator, so the
-assertion is a kernel-noise tolerance rather than bit-exact.
-"""
-from __future__ import annotations
-
-import glob
-import os
-from pathlib import Path
-
-import pytest
-import torch
-import torch.nn.functional as F
-from torch.testing import assert_close
-
-from fastvideo.pipelines.basic.magi_human.pipeline_configs import (
-    _SR_1080P_LOCAL_ATTN_LAYERS,
-)
-from tests.local_tests.magi_human.test_magi_human_pipeline_parity import (
-    _build_fastvideo_schedulers,
-    _build_upstream_schedulers,
-    _cleanup_gpu,
-    _dit_forward_fv,
-    _dit_forward_upstream,
-    _find_base_shard_dir,
-    _run_denoise_loop,
-)
-from tests.local_tests.magi_human.test_magi_human_sr540p_pipeline_parity import (
-    _load_fv_dit,
-    _prepare_sr_latents,
-    _run_sr_denoise_loop,
-)
-
-
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29522")
-
-
-def _find_sr1080p_shard_dir() -> Path | None:
-    override = os.getenv("MAGI_HUMAN_SR1080P_SHARD_DIR")
-    if override:
-        path = Path(override)
-        return path if path.is_dir() else None
-    try:
-        from huggingface_hub import snapshot_download
-        snap = snapshot_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            allow_patterns=[
-                "1080p_sr/*.safetensors",
-                "1080p_sr/model.safetensors.index.json",
-            ],
-        )
-        candidate = Path(snap) / "1080p_sr"
-        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
-            return candidate
-        return None
-    except Exception:
-        return None
-
-
-def _dit_forward_upstream_local(
-    dit,
-    video_latent,
-    audio_latent,
-    audio_feat_len,
-    txt_feat,
-    txt_feat_len,
-    patch_size,
-    coords_style,
-    video_in_channels,
-    audio_in_channels,
-):
-    from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
-        build_packed_inputs,
-        unpack_tokens,
-    )
-    from inference.common import VarlenHandler
-    from inference.pipeline.data_proxy import calc_local_attn_ffa_handler
-
-    x, coords, mm = build_packed_inputs(
-        video_latent=video_latent,
-        audio_latent=audio_latent,
-        audio_feat_len=audio_feat_len,
-        txt_feat=txt_feat,
-        txt_feat_len=txt_feat_len,
-        patch_size=patch_size,
-        coords_style=coords_style,
-    )
-    video_token_num = x.shape[0] - audio_feat_len - txt_feat_len
-    total = x.shape[0]
-    cu = torch.tensor([0, total], dtype=torch.int32, device=x.device)
-    varlen = VarlenHandler(
-        cu_seqlens_q=cu,
-        cu_seqlens_k=cu,
-        max_seqlen_q=total,
-        max_seqlen_k=total,
-    )
-    local_attn = calc_local_attn_ffa_handler(
-        video_token_num,
-        audio_feat_len + txt_feat_len,
-        video_latent.shape[2] // patch_size[0],
-        11,
-    )
-    out = dit(
-        x=x,
-        coords_mapping=coords,
-        modality_mapping=mm,
-        varlen_handler=varlen,
-        local_attn_handler=local_attn,
-    )
-    return unpack_tokens(
-        out,
-        video_token_num=video_token_num,
-        audio_feat_len=audio_feat_len,
-        video_in_channels=video_in_channels,
-        audio_in_channels=audio_in_channels,
-        latent_shape=tuple(video_latent.shape),
-        patch_size=patch_size,
-    )
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman SR-1080p pipeline parity requires CUDA.",
-)
-@pytest.mark.parametrize("use_image", [False, True], ids=["t2v", "ti2v"])
-def test_magi_human_sr1080p_pipeline_latent_parity(use_image: bool):
-    repo_root = Path(__file__).resolve().parents[3]
-    if not (repo_root / "daVinci-MagiHuman").exists():
-        pytest.skip("Upstream daVinci-MagiHuman/ clone missing.")
-
-    base_shard_dir = _find_base_shard_dir()
-    sr_shard_dir = _find_sr1080p_shard_dir()
-    if base_shard_dir is None or not base_shard_dir.is_dir():
-        pytest.skip("GAIR/daVinci-MagiHuman base/ shards not available locally.")
-    if sr_shard_dir is None or not sr_shard_dir.is_dir():
-        pytest.skip("GAIR/daVinci-MagiHuman 1080p_sr/ shards not available locally.")
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_SR1080P_DIFFUSERS_PATH",
-        repo_root / "converted_weights" / "magi_human_sr_1080p",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    sr_transformer_dir = converted_dir / "sr_transformer"
-    if not transformer_dir.is_dir():
-        pytest.skip(f"Converted base transformer dir missing at {transformer_dir}")
-    if not sr_transformer_dir.is_dir():
-        pytest.skip(f"Converted SR transformer dir missing at {sr_transformer_dir}")
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs,
-        load_upstream_dit,
-    )
-    install_stubs()
-
-    device = torch.device("cuda:0")
-    torch.manual_seed(1080)
-    z_dim = 48
-    patch_size = (1, 2, 2)
-    base_lat_T, base_lat_H, base_lat_W = 24, 4, 4
-    sr_lat_H, sr_lat_W = 6, 8
-    video_latent = torch.randn(
-        (1, z_dim, base_lat_T, base_lat_H, base_lat_W),
-        dtype=torch.float32,
-        device=device,
-    )
-    audio_latent = torch.randn((1, 32, 64), dtype=torch.float32, device=device)
-    base_image_latent = None
-    sr_image_latent = None
-    if use_image:
-        base_image_latent = torch.randn(
-            (1, z_dim, 1, base_lat_H, base_lat_W),
-            dtype=torch.float32,
-            device=device,
-        )
-        sr_image_latent = F.interpolate(
-            base_image_latent,
-            size=(1, sr_lat_H, sr_lat_W),
-            mode="trilinear",
-            align_corners=True,
-        )
-
-    txt_feat_len = 7
-    neg_txt_feat_len = 11
-    txt_feat = torch.randn((1, 640, 3584), dtype=torch.float32, device=device)
-    neg_txt_feat = torch.randn((1, 640, 3584), dtype=torch.float32, device=device)
-    base_steps = 4
-    sr_steps = 2
-    shift = 5.0
-    base_kwargs = dict(
-        cfg_number=2,
-        video_txt_guidance_scale=5.0,
-        audio_txt_guidance_scale=5.0,
-        patch_size=patch_size,
-        coords_style="v2",
-        video_in_channels=192,
-        audio_in_channels=64,
-        image_latent=base_image_latent,
-    )
-    sr_kwargs = dict(
-        patch_size=patch_size,
-        coords_style="v1",
-        video_in_channels=192,
-        audio_in_channels=64,
-        image_latent=sr_image_latent,
-    )
-
-    up_video_sched, up_audio_sched = _build_upstream_schedulers(
-        shift=shift,
-        num_inference_steps=base_steps,
-        device=device,
-    )
-    upstream_base = load_upstream_dit(base_shard_dir, device=device, dtype=None)
-    ref_base_video, ref_base_audio = _run_denoise_loop(
-        upstream_base,
-        _dit_forward_upstream,
-        video_latent.clone(),
-        audio_latent.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=up_video_sched,
-        audio_sched=up_audio_sched,
-        **base_kwargs,
-    )
-    del upstream_base
-    _cleanup_gpu()
-
-    torch.manual_seed(1081)
-    ref_sr_video_in, ref_sr_audio_in = _prepare_sr_latents(
-        ref_base_video,
-        ref_base_audio,
-        latent_h=sr_lat_H,
-        latent_w=sr_lat_W,
-        noise_value=220,
-    )
-    up_sr_video_sched, _ = _build_upstream_schedulers(
-        shift=shift,
-        num_inference_steps=sr_steps,
-        device=device,
-    )
-    upstream_sr = load_upstream_dit(
-        sr_shard_dir,
-        device=device,
-        dtype=None,
-        local_attn_layers=_SR_1080P_LOCAL_ATTN_LAYERS,
-    )
-    ref_video, ref_audio = _run_sr_denoise_loop(
-        upstream_sr,
-        _dit_forward_upstream_local,
-        ref_sr_video_in.clone(),
-        ref_sr_audio_in.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=up_sr_video_sched,
-        **sr_kwargs,
-    )
-    ref_video = ref_video.detach().float().cpu()
-    ref_audio = ref_audio.detach().float().cpu()
-    del upstream_sr
-    _cleanup_gpu()
-
-    fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers(
-        shift=shift,
-        num_inference_steps=base_steps,
-        device=device,
-    )
-    fv_base = _load_fv_dit(transformer_dir, device)
-    fv_base_video, fv_base_audio = _run_denoise_loop(
-        fv_base,
-        _dit_forward_fv,
-        video_latent.clone(),
-        audio_latent.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=fv_video_sched,
-        audio_sched=fv_audio_sched,
-        **base_kwargs,
-    )
-    del fv_base
-    _cleanup_gpu()
-
-    torch.manual_seed(1081)
-    fv_sr_video_in, fv_sr_audio_in = _prepare_sr_latents(
-        fv_base_video,
-        fv_base_audio,
-        latent_h=sr_lat_H,
-        latent_w=sr_lat_W,
-        noise_value=220,
-    )
-    fv_sr_video_sched, _ = _build_fastvideo_schedulers(
-        shift=shift,
-        num_inference_steps=sr_steps,
-        device=device,
-    )
-    fv_sr = _load_fv_dit(sr_transformer_dir, device)
-    fv_sr.configure_local_attention(_SR_1080P_LOCAL_ATTN_LAYERS, frame_receptive_field=11)
-    fv_video, fv_audio = _run_sr_denoise_loop(
-        fv_sr,
-        _dit_forward_fv,
-        fv_sr_video_in.clone(),
-        fv_sr_audio_in.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=fv_sr_video_sched,
-        **sr_kwargs,
-    )
-    fv_video = fv_video.detach().float().cpu()
-    fv_audio = fv_audio.detach().float().cpu()
-
-    v_diff = (ref_video - fv_video).abs()
-    a_diff = (ref_audio - fv_audio).abs()
-    print(
-        f"sr1080p {('ti2v' if use_image else 't2v')} "
-        f"video diff_max={v_diff.max().item():.4f} diff_mean={v_diff.mean().item():.4f}"
-    )
-    print(
-        f"sr1080p {('ti2v' if use_image else 't2v')} "
-        f"audio diff_max={a_diff.max().item():.4f} diff_mean={a_diff.mean().item():.4f}"
-    )
-
-    assert ref_video.shape == fv_video.shape
-    assert ref_audio.shape == fv_audio.shape
-    assert v_diff.max().item() < 0.05
-    assert_close(fv_audio, ref_audio, atol=0.0, rtol=0.0)
-    if use_image:
-        assert_close(fv_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0)
-        assert_close(ref_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0)

tests/local_tests/magi_human/test_magi_human_sr540p_pipeline_parity.py

@@ -1,400 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Two-stage SR-540p latent-loop parity for daVinci-MagiHuman."""
-from __future__ import annotations
-
-import glob
-import os
-from pathlib import Path
-
-import pytest
-import torch
-import torch.nn.functional as F
-from torch.testing import assert_close
-
-from fastvideo.pipelines.basic.magi_human.stages.sr_latent_preparation import (
-    ZeroSNRDDPMDiscretization,
-)
-from tests.local_tests.magi_human.test_magi_human_pipeline_parity import (
-    _build_fastvideo_schedulers,
-    _build_upstream_schedulers,
-    _cleanup_gpu,
-    _dit_forward_fv,
-    _dit_forward_upstream,
-    _find_base_shard_dir,
-    _run_denoise_loop,
-)
-
-
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29521")
-
-
-def _find_sr540p_shard_dir() -> Path | None:
-    override = os.getenv("MAGI_HUMAN_SR540P_SHARD_DIR")
-    if override:
-        path = Path(override)
-        return path if path.is_dir() else None
-    try:
-        from huggingface_hub import snapshot_download
-        snap = snapshot_download(
-            repo_id="GAIR/daVinci-MagiHuman",
-            allow_patterns=[
-                "540p_sr/*.safetensors",
-                "540p_sr/model.safetensors.index.json",
-            ],
-        )
-        candidate = Path(snap) / "540p_sr"
-        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
-            return candidate
-        return None
-    except Exception:
-        return None
-
-
-def _prepare_sr_latents(
-    br_video: torch.Tensor,
-    br_audio: torch.Tensor,
-    *,
-    latent_h: int,
-    latent_w: int,
-    noise_value: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    latent_video = F.interpolate(
-        br_video,
-        size=(br_video.shape[2], latent_h, latent_w),
-        mode="trilinear",
-        align_corners=True,
-    )
-    if noise_value != 0:
-        noise = torch.randn_like(latent_video, device=latent_video.device)
-        sigmas = ZeroSNRDDPMDiscretization()(
-            1000,
-            do_append_zero=False,
-            flip=True,
-            device=latent_video.device,
-        )
-        sigma = sigmas[noise_value]
-        latent_video = latent_video * sigma + noise * (1 - sigma**2)**0.5
-    sr_audio = torch.randn_like(br_audio, device=br_audio.device) * 0.7 + br_audio * 0.3
-    return latent_video, sr_audio
-
-
-def _run_sr_denoise_loop(
-    dit,
-    dit_forward_fn,
-    video_latent,
-    audio_latent,
-    txt_feat,
-    txt_feat_len,
-    neg_txt_feat,
-    neg_txt_feat_len,
-    *,
-    video_sched,
-    patch_size,
-    coords_style,
-    video_in_channels,
-    audio_in_channels,
-    image_latent=None,
-):
-    from fastvideo.forward_context import set_forward_context
-
-    audio_feat_len = int(audio_latent.shape[1])
-    latent_length = video_latent.shape[2]
-    guidance = torch.tensor(3.5, device=video_latent.device).expand(
-        1,
-        1,
-        latent_length,
-        1,
-        1,
-    ).clone()
-    guidance[:, :, :13] = 2.0
-
-    with torch.inference_mode():
-        for t in video_sched.timesteps:
-            if image_latent is not None:
-                video_latent[:, :, :1] = image_latent.to(
-                    device=video_latent.device,
-                    dtype=video_latent.dtype,
-                )[:, :, :1]
-            with set_forward_context(
-                current_timestep=int(t.item()) if torch.is_tensor(t) else int(t),
-                attn_metadata=None,
-            ):
-                v_cond_video, _ = dit_forward_fn(
-                    dit,
-                    video_latent,
-                    audio_latent,
-                    audio_feat_len,
-                    txt_feat,
-                    txt_feat_len,
-                    patch_size,
-                    coords_style,
-                    video_in_channels,
-                    audio_in_channels,
-                )
-                v_uncond_video, _ = dit_forward_fn(
-                    dit,
-                    video_latent,
-                    audio_latent,
-                    audio_feat_len,
-                    neg_txt_feat,
-                    neg_txt_feat_len,
-                    patch_size,
-                    coords_style,
-                    video_in_channels,
-                    audio_in_channels,
-                )
-                v_video = v_uncond_video + guidance * (v_cond_video - v_uncond_video)
-            video_latent = video_sched.step(
-                v_video,
-                t,
-                video_latent,
-                return_dict=False,
-            )[0]
-        if image_latent is not None:
-            video_latent[:, :, :1] = image_latent.to(
-                device=video_latent.device,
-                dtype=video_latent.dtype,
-            )[:, :, :1]
-    return video_latent, audio_latent
-
-
-def _load_fv_dit(transformer_dir: Path, device: torch.device):
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-
-    dit = MagiHumanDiT(MagiHumanVideoConfig())
-    state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        state.update(load_file(shard))
-    missing, unexpected = dit.load_state_dict(state, strict=False)
-    assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}"
-    assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
-    return dit.to(device=device).eval()
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman SR-540p pipeline parity requires CUDA.",
-)
-@pytest.mark.parametrize("use_image", [False, True], ids=["t2v", "ti2v"])
-def test_magi_human_sr540p_pipeline_latent_parity(use_image: bool):
-    repo_root = Path(__file__).resolve().parents[3]
-    if not (repo_root / "daVinci-MagiHuman").exists():
-        pytest.skip("Upstream daVinci-MagiHuman/ clone missing.")
-
-    base_shard_dir = _find_base_shard_dir()
-    sr_shard_dir = _find_sr540p_shard_dir()
-    if base_shard_dir is None or not base_shard_dir.is_dir():
-        pytest.skip("GAIR/daVinci-MagiHuman base/ shards not available locally.")
-    if sr_shard_dir is None or not sr_shard_dir.is_dir():
-        pytest.skip("GAIR/daVinci-MagiHuman 540p_sr/ shards not available locally.")
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_SR540P_DIFFUSERS_PATH",
-        repo_root / "converted_weights" / "magi_human_sr_540p",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    sr_transformer_dir = converted_dir / "sr_transformer"
-    if not transformer_dir.is_dir():
-        pytest.skip(f"Converted base transformer dir missing at {transformer_dir}")
-    if not sr_transformer_dir.is_dir():
-        pytest.skip(f"Converted SR transformer dir missing at {sr_transformer_dir}")
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs,
-        load_upstream_dit,
-    )
-    install_stubs()
-
-    device = torch.device("cuda:0")
-    torch.manual_seed(540)
-    z_dim = 48
-    patch_size = (1, 2, 2)
-    base_lat_T, base_lat_H, base_lat_W = 2, 6, 6
-    sr_lat_H, sr_lat_W = 8, 10
-    video_latent = torch.randn(
-        (1, z_dim, base_lat_T, base_lat_H, base_lat_W),
-        dtype=torch.float32,
-        device=device,
-    )
-    audio_latent = torch.randn((1, 4, 64), dtype=torch.float32, device=device)
-    base_image_latent = None
-    sr_image_latent = None
-    if use_image:
-        base_image_latent = torch.randn(
-            (1, z_dim, 1, base_lat_H, base_lat_W),
-            dtype=torch.float32,
-            device=device,
-        )
-        sr_image_latent = F.interpolate(
-            base_image_latent,
-            size=(1, sr_lat_H, sr_lat_W),
-            mode="trilinear",
-            align_corners=True,
-        )
-
-    txt_feat_len = 7
-    neg_txt_feat_len = 11
-    txt_feat = torch.randn(
-        (1, 640, 3584),
-        dtype=torch.float32,
-        device=device,
-    )
-    neg_txt_feat = torch.randn(
-        (1, 640, 3584),
-        dtype=torch.float32,
-        device=device,
-    )
-    base_steps = 4
-    sr_steps = 2
-    shift = 5.0
-    base_kwargs = dict(
-        cfg_number=2,
-        video_txt_guidance_scale=5.0,
-        audio_txt_guidance_scale=5.0,
-        patch_size=patch_size,
-        coords_style="v2",
-        video_in_channels=192,
-        audio_in_channels=64,
-        image_latent=base_image_latent,
-    )
-    sr_kwargs = dict(
-        patch_size=patch_size,
-        coords_style="v1",
-        video_in_channels=192,
-        audio_in_channels=64,
-        image_latent=sr_image_latent,
-    )
-
-    up_video_sched, up_audio_sched = _build_upstream_schedulers(
-        shift=shift,
-        num_inference_steps=base_steps,
-        device=device,
-    )
-    upstream_base = load_upstream_dit(base_shard_dir, device=device, dtype=None)
-    ref_base_video, ref_base_audio = _run_denoise_loop(
-        upstream_base,
-        _dit_forward_upstream,
-        video_latent.clone(),
-        audio_latent.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=up_video_sched,
-        audio_sched=up_audio_sched,
-        **base_kwargs,
-    )
-    del upstream_base
-    _cleanup_gpu()
-
-    torch.manual_seed(541)
-    ref_sr_video_in, ref_sr_audio_in = _prepare_sr_latents(
-        ref_base_video,
-        ref_base_audio,
-        latent_h=sr_lat_H,
-        latent_w=sr_lat_W,
-        noise_value=220,
-    )
-    up_sr_video_sched, _ = _build_upstream_schedulers(
-        shift=shift,
-        num_inference_steps=sr_steps,
-        device=device,
-    )
-    upstream_sr = load_upstream_dit(sr_shard_dir, device=device, dtype=None)
-    ref_video, ref_audio = _run_sr_denoise_loop(
-        upstream_sr,
-        _dit_forward_upstream,
-        ref_sr_video_in.clone(),
-        ref_sr_audio_in.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=up_sr_video_sched,
-        **sr_kwargs,
-    )
-    ref_video = ref_video.detach().float().cpu()
-    ref_audio = ref_audio.detach().float().cpu()
-    del upstream_sr
-    _cleanup_gpu()
-
-    fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers(
-        shift=shift,
-        num_inference_steps=base_steps,
-        device=device,
-    )
-    fv_base = _load_fv_dit(transformer_dir, device)
-    fv_base_video, fv_base_audio = _run_denoise_loop(
-        fv_base,
-        _dit_forward_fv,
-        video_latent.clone(),
-        audio_latent.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=fv_video_sched,
-        audio_sched=fv_audio_sched,
-        **base_kwargs,
-    )
-    del fv_base
-    _cleanup_gpu()
-
-    torch.manual_seed(541)
-    fv_sr_video_in, fv_sr_audio_in = _prepare_sr_latents(
-        fv_base_video,
-        fv_base_audio,
-        latent_h=sr_lat_H,
-        latent_w=sr_lat_W,
-        noise_value=220,
-    )
-    fv_sr_video_sched, _ = _build_fastvideo_schedulers(
-        shift=shift,
-        num_inference_steps=sr_steps,
-        device=device,
-    )
-    fv_sr = _load_fv_dit(sr_transformer_dir, device)
-    fv_video, fv_audio = _run_sr_denoise_loop(
-        fv_sr,
-        _dit_forward_fv,
-        fv_sr_video_in.clone(),
-        fv_sr_audio_in.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=fv_sr_video_sched,
-        **sr_kwargs,
-    )
-    fv_video = fv_video.detach().float().cpu()
-    fv_audio = fv_audio.detach().float().cpu()
-
-    v_diff = (ref_video - fv_video).abs()
-    a_diff = (ref_audio - fv_audio).abs()
-    print(
-        f"sr540p {('ti2v' if use_image else 't2v')} "
-        f"video diff_max={v_diff.max().item():.4f} diff_mean={v_diff.mean().item():.4f}"
-    )
-    print(
-        f"sr540p {('ti2v' if use_image else 't2v')} "
-        f"audio diff_max={a_diff.max().item():.4f} diff_mean={a_diff.mean().item():.4f}"
-    )
-
-    assert ref_video.shape == fv_video.shape
-    assert ref_audio.shape == fv_audio.shape
-    assert_close(fv_audio, ref_audio, atol=0.0, rtol=0.0)
-    assert_close(fv_video, ref_video, atol=0.0, rtol=0.0)
-    if use_image:
-        assert_close(fv_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0)
-        assert_close(ref_video[:, :, :1], sr_image_latent.detach().cpu(), atol=0.0, rtol=0.0)
-
-    ref_v_abs = ref_video.abs().mean().item()
-    ref_a_abs = ref_audio.abs().mean().item()
-    assert abs(ref_v_abs - fv_video.abs().mean().item()) / max(ref_v_abs, 1e-6) < 0.02
-    assert abs(ref_a_abs - fv_audio.abs().mean().item()) / max(ref_a_abs, 1e-6) < 0.02
-    assert v_diff.mean().item() / max(ref_v_abs, 1e-6) < 0.06
-    assert a_diff.mean().item() / max(ref_a_abs, 1e-6) < 0.04

tests/local_tests/magi_human/test_magi_human_t5gemma_parity.py

@@ -1,130 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Parity test: FastVideo T5GemmaEncoderModel vs direct HF
-`T5GemmaEncoderModel.from_pretrained(...)`.
-
-FastVideo's wrapper is intentionally thin — it lazy-loads the same HF
-class on the same gated repo (`google/t5gemma-9b-9b-ul2`) that the
-upstream MagiHuman pipeline uses (see
-daVinci-MagiHuman/inference/model/t5_gemma/t5_gemma_model.py). This
-test guards against future regressions in the wrapper (e.g. accidental
-mutation of `last_hidden_state`, wrong dtype cast, forgetting to pass
-attention_mask) by comparing wrapper forward output against a direct HF
-forward on the same model.
-
-Skips when the T5-Gemma repo isn't accessible (gated — requires user's
-HF token with accepted terms of use).
-"""
-from __future__ import annotations
-
-import os
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-
-_T5GEMMA_ID = "google/t5gemma-9b-9b-ul2"
-
-
-def _hf_token():
-    for k in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        v = os.environ.get(k)
-        if v:
-            return v
-    return None
-
-
-def _can_access_t5gemma() -> bool:
-    token = _hf_token()
-    if token is None:
-        return False
-    try:
-        from huggingface_hub import hf_hub_download
-        hf_hub_download(
-            repo_id=_T5GEMMA_ID, filename="config.json", token=token,
-        )
-        return True
-    except Exception:
-        return False
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman T5-Gemma parity requires CUDA (encoder is 9B params).",
-)
-@pytest.mark.skipif(
-    not _can_access_t5gemma(),
-    reason=(f"{_T5GEMMA_ID} not accessible — gated Google repo; set "
-            f"HF_TOKEN / HF_API_KEY and accept the terms of use."),
-)
-def test_magi_human_t5gemma_wrapper_parity():
-    # Alias any of the three token env vars to HF_TOKEN (what transformers
-    # reads) before constructing models.
-    for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"):
-        v = os.environ.get(src)
-        if v:
-            os.environ.setdefault("HF_TOKEN", v)
-            os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", v)
-            break
-
-    device = torch.device("cuda:0")
-
-    # --- Upstream / direct HF path (matches the reference pipeline's
-    #     `T5GemmaEncoder` wrapper exactly: see
-    #     daVinci-MagiHuman/inference/model/t5_gemma/t5_gemma_model.py) ---
-    from transformers import AutoTokenizer
-    from transformers.models.t5gemma import T5GemmaEncoderModel as HFEncoder
-
-    tokenizer = AutoTokenizer.from_pretrained(_T5GEMMA_ID)
-    ref_model = HFEncoder.from_pretrained(
-        _T5GEMMA_ID, is_encoder_decoder=False, dtype=torch.bfloat16,
-    ).to(device).eval()
-
-    # --- FastVideo wrapper path ---
-    from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
-    from fastvideo.models.encoders.t5gemma import T5GemmaEncoderModel as FVEncoder
-
-    fv_config = T5GemmaEncoderConfig()
-    fv_config.arch_config.t5gemma_model_path = _T5GEMMA_ID
-    fv_model = FVEncoder(fv_config)
-
-    # --- Identical input ---
-    prompt = (
-        "A warm afternoon scene: a person sits on a park bench reading "
-        "a book, surrounded by softly swaying trees."
-    )
-    inputs = tokenizer(
-        [prompt], return_tensors="pt", padding=True, truncation=False,
-    ).to(device)
-
-    with torch.inference_mode():
-        ref_out = ref_model(**inputs)
-        ref_hidden = ref_out["last_hidden_state"].detach().float().cpu()
-
-        # FastVideo wrapper: forward through the adapter; it lazy-loads the
-        # encoder on first call and moves it to the input's device.
-        fv_out = fv_model(
-            input_ids=inputs["input_ids"],
-            attention_mask=inputs.get("attention_mask"),
-        )
-        fv_hidden = fv_out.last_hidden_state.detach().float().cpu()
-
-    print(
-        f"ref_hidden shape={tuple(ref_hidden.shape)} "
-        f"abs_mean={ref_hidden.abs().mean().item():.6f}"
-    )
-    print(
-        f"fv_hidden  shape={tuple(fv_hidden.shape)} "
-        f"abs_mean={fv_hidden.abs().mean().item():.6f}"
-    )
-    diff = (ref_hidden - fv_hidden).abs()
-    print(
-        f"diff max={diff.max().item():.6e} "
-        f"mean={diff.mean().item():.6e}"
-    )
-
-    assert ref_hidden.shape == fv_hidden.shape
-    # Both sides run the exact same HF model on the exact same inputs;
-    # drift is bounded by nondeterminism in SDPA + bf16 matmul. This
-    # should be <= 1e-3 end-to-end.
-    assert_close(fv_hidden, ref_hidden, atol=1e-3, rtol=1e-3)

tests/local_tests/magi_human/test_magi_human_ti2v_pipeline_parity.py

@@ -1,186 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""TI2V latent-loop parity for daVinci-MagiHuman.
-
-This mirrors the base MagiHuman pipeline parity test but enables the upstream
-`latent_image is not None` branch: the clean image latent is copied into
-`latent_video[:, :, :1]` before every DiT call and once more after denoising.
-"""
-from __future__ import annotations
-
-import glob
-import os
-from pathlib import Path
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-from tests.local_tests.magi_human.test_magi_human_pipeline_parity import (
-    _build_fastvideo_schedulers,
-    _build_upstream_schedulers,
-    _cleanup_gpu,
-    _dit_forward_fv,
-    _dit_forward_upstream,
-    _encode_magi_human_prompt_pair,
-    _find_base_shard_dir,
-    _run_denoise_loop,
-)
-
-
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29520")
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="MagiHuman TI2V pipeline parity requires CUDA.",
-)
-def test_magi_human_ti2v_pipeline_latent_parity():
-    repo_root = Path(__file__).resolve().parents[3]
-    upstream_src = repo_root / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        pytest.skip("Upstream daVinci-MagiHuman/ clone missing.")
-
-    base_shard_dir = _find_base_shard_dir()
-    if base_shard_dir is None or not base_shard_dir.is_dir():
-        pytest.skip("GAIR/daVinci-MagiHuman base/ shards not available locally.")
-
-    converted_dir = Path(os.getenv(
-        "MAGI_HUMAN_DIFFUSERS_PATH",
-        repo_root / "converted_weights" / "magi_human_base",
-    ))
-    transformer_dir = converted_dir / "transformer"
-    if not transformer_dir.is_dir():
-        pytest.skip(f"Converted transformer dir missing at {transformer_dir}")
-
-    from tests.local_tests.helpers.magi_human_upstream import (
-        install_stubs,
-        load_upstream_dit,
-    )
-    install_stubs()
-
-    device = torch.device("cuda:0")
-    torch.manual_seed(123)
-
-    z_dim = 48
-    patch_size = (1, 2, 2)
-    lat_T, lat_H, lat_W = 2, 6, 6
-    video_latent = torch.randn(
-        (1, z_dim, lat_T, lat_H, lat_W),
-        dtype=torch.float32,
-        device=device,
-    )
-    audio_latent = torch.randn((1, 4, 64), dtype=torch.float32, device=device)
-    image_latent = torch.randn(
-        (1, z_dim, 1, lat_H, lat_W),
-        dtype=torch.float32,
-        device=device,
-    )
-    txt_feat, txt_feat_len, neg_txt_feat, neg_txt_feat_len = (
-        _encode_magi_human_prompt_pair(device)
-    )
-
-    num_inference_steps = 4
-    shift = 5.0
-    common_kwargs = dict(
-        cfg_number=2,
-        video_txt_guidance_scale=5.0,
-        audio_txt_guidance_scale=5.0,
-        patch_size=patch_size,
-        coords_style="v2",
-        video_in_channels=192,
-        audio_in_channels=64,
-        image_latent=image_latent,
-    )
-
-    up_video_sched, up_audio_sched = _build_upstream_schedulers(
-        shift=shift,
-        num_inference_steps=num_inference_steps,
-        device=device,
-    )
-    print("Loading upstream DiTModel from base shards...")
-    upstream_dit = load_upstream_dit(base_shard_dir, device=device, dtype=None)
-    print("Running upstream TI2V denoise loop...")
-    ref_video, ref_audio = _run_denoise_loop(
-        upstream_dit,
-        _dit_forward_upstream,
-        video_latent.clone(),
-        audio_latent.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=up_video_sched,
-        audio_sched=up_audio_sched,
-        **common_kwargs,
-    )
-    ref_video = ref_video.detach().float().cpu()
-    ref_audio = ref_audio.detach().float().cpu()
-    del upstream_dit
-    _cleanup_gpu()
-
-    fv_video_sched, fv_audio_sched = _build_fastvideo_schedulers(
-        shift=shift,
-        num_inference_steps=num_inference_steps,
-        device=device,
-    )
-    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
-    from fastvideo.models.dits.magi_human import MagiHumanDiT
-    from safetensors.torch import load_file
-
-    print("Loading FastVideo MagiHumanDiT from converted transformer/...")
-    fv_cfg = MagiHumanVideoConfig()
-    fv_dit = MagiHumanDiT(fv_cfg)
-    fv_state = {}
-    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
-        fv_state.update(load_file(shard))
-    missing, unexpected = fv_dit.load_state_dict(fv_state, strict=False)
-    assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}"
-    assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
-    fv_dit = fv_dit.to(device=device)
-    fv_dit.eval()
-
-    print("Running FastVideo TI2V denoise loop...")
-    fv_video, fv_audio = _run_denoise_loop(
-        fv_dit,
-        _dit_forward_fv,
-        video_latent.clone(),
-        audio_latent.clone(),
-        txt_feat.clone(),
-        txt_feat_len,
-        neg_txt_feat.clone(),
-        neg_txt_feat_len,
-        video_sched=fv_video_sched,
-        audio_sched=fv_audio_sched,
-        **common_kwargs,
-    )
-    fv_video = fv_video.detach().float().cpu()
-    fv_audio = fv_audio.detach().float().cpu()
-
-    v_diff = (ref_video - fv_video).abs()
-    a_diff = (ref_audio - fv_audio).abs()
-    print(
-        f"ti2v video diff_max={v_diff.max().item():.4f} "
-        f"diff_mean={v_diff.mean().item():.4f}"
-    )
-    print(
-        f"ti2v audio diff_max={a_diff.max().item():.4f} "
-        f"diff_mean={a_diff.mean().item():.4f}"
-    )
-
-    assert ref_video.shape == fv_video.shape
-    assert ref_audio.shape == fv_audio.shape
-    assert_close(fv_video, ref_video, atol=0.40, rtol=0.05)
-    assert_close(fv_audio, ref_audio, atol=0.40, rtol=0.05)
-    assert_close(fv_video[:, :, :1], image_latent.detach().cpu(), atol=0.0, rtol=0.0)
-    assert_close(ref_video[:, :, :1], image_latent.detach().cpu(), atol=0.0, rtol=0.0)
-
-    ref_v_abs = ref_video.abs().mean().item()
-    ref_a_abs = ref_audio.abs().mean().item()
-    rel_v = abs(ref_v_abs - fv_video.abs().mean().item()) / max(ref_v_abs, 1e-6)
-    rel_a = abs(ref_a_abs - fv_audio.abs().mean().item()) / max(ref_a_abs, 1e-6)
-    assert rel_v < 0.01, f"video abs_mean drift {rel_v:.2%} > 1%"
-    assert rel_a < 0.01, f"audio abs_mean drift {rel_a:.2%} > 1%"
-    assert v_diff.mean().item() / max(ref_v_abs, 1e-6) < 0.04
-    assert a_diff.mean().item() / max(ref_a_abs, 1e-6) < 0.04

tests/local_tests/magi_human/test_magi_human_vae_parity.py

@@ -1,180 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Parity test: FastVideo `AutoencoderKLWan` vs upstream `Wan2_2_VAE`.
-
-MagiHuman uses the Wan 2.2 TI2V-5B VAE. The two implementations
-compared here are:
-
-  * Upstream (SandAI port) — `inference/model/vae2_2/vae2_2_module.py::Wan2_2_VAE`
-    loaded from `Wan-AI/Wan2.2-TI2V-5B/Wan2.2_VAE.pth` (the official .pth
-    inside the daVinci-MagiHuman repo). This is the reference.
-  * FastVideo — `fastvideo.models.vaes.wanvae.AutoencoderKLWan` (the
-    class registered as `EntryClass` and resolved by the VAE component
-    loader at runtime; this is what `MagiHumanBaseConfig.vae_config`
-    materializes when the magi pipeline runs). Weights are loaded from
-    a Diffusers-format `vae/` subdir (`config.json` +
-    `diffusion_pytorch_model.safetensors`).
-
-This test decodes the same random latent through both and asserts the
-decoded videos are close. Catches regressions in:
-  - FastVideo's `AutoencoderKLWan` weight load / scale / shift handling.
-  - Any deviation in `latents_mean` / `latents_std` baked into the
-    Diffusers-format config vs the upstream constants.
-
-Skips when:
-  - CUDA is unavailable.
-  - The .pth is not locally available (requires ~2.8 GB download).
-  - The converted MagiHuman Diffusers repo (or any `Wan-AI/*-Diffusers`
-    repo with a `vae/` subdir) is not available locally.
-"""
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="VAE parity test requires CUDA.",
-)
-def test_magi_human_vae_decode_parity():
-    repo_root = Path(__file__).resolve().parents[3]
-    upstream_src = repo_root / "daVinci-MagiHuman"
-    if not upstream_src.exists():
-        pytest.skip(
-            "Upstream daVinci-MagiHuman/ clone missing — no Wan2_2_VAE source."
-        )
-
-    fv_vae_dir = Path(os.getenv(
-        "MAGI_HUMAN_VAE_DIR",
-        repo_root / "converted_weights" / "magi_human_base" / "vae",
-    ))
-    if not (fv_vae_dir / "config.json").is_file():
-        pytest.skip(f"FastVideo VAE dir missing at {fv_vae_dir}")
-
-    # Upstream Wan2_2_VAE needs the raw .pth shipped by Wan-AI/Wan2.2-TI2V-5B
-    # (NOT the -Diffusers variant; that one has safetensors, not .pth).
-    try:
-        from huggingface_hub import hf_hub_download
-        pth_path = hf_hub_download(
-            repo_id="Wan-AI/Wan2.2-TI2V-5B", filename="Wan2.2_VAE.pth",
-        )
-    except Exception as exc:
-        pytest.skip(f"Wan2.2_VAE.pth not available: {exc}")
-
-    # Push upstream + install compiler stubs (the VAE module itself doesn't
-    # need magi_compiler, but `inference.*` imports pull in siblings that do).
-    from tests.local_tests.helpers.magi_human_upstream import install_stubs
-    install_stubs()
-
-    device = torch.device("cuda:0")
-    torch.manual_seed(0)
-
-    # Tiny latent so the test stays well inside GPU memory budget.
-    # z_dim=48, T=1, H=4, W=4 -> VAE decodes to [1, 3, 1 (or 1+4*0), 64, 64]
-    z = torch.randn((1, 48, 1, 4, 4), dtype=torch.float32, device=device)
-
-    # --- Upstream decode ---
-    from inference.model.vae2_2 import Wan2_2_VAE
-    up_vae = Wan2_2_VAE(
-        vae_pth=pth_path,
-        device=device,
-        dtype=torch.float32,
-    )
-    with torch.inference_mode():
-        # Wan2_2_VAE.decode expects a (C, T, H, W) latent (no batch dim);
-        # see inference/pipeline/video_generate.py:494 — `self.vae.decode(latent.squeeze(0).to(self.dtype), ...)`.
-        up_out = up_vae.decode(z[0]).detach().float().cpu()
-
-    del up_vae
-    import gc; gc.collect(); torch.cuda.empty_cache()
-
-    # --- FastVideo decode ---
-    # Upstream `Wan2_2_VAE.decode(z)` internally normalizes via
-    # `(z - latents_mean) / latents_std` before feeding the decoder
-    # (see `scale = [mean, 1.0/std]` and the _video_vae.decode call).
-    # FastVideo's `AutoencoderKLWan.decode(z)` expects the input to
-    # ALREADY be in "decoder-input space" (the normalization is the
-    # caller's job — `DecodingStage` applies it). So we mirror the
-    # upstream transform here before calling decode.
-    import glob
-
-    from safetensors.torch import load_file as safetensors_load_file
-
-    from fastvideo.configs.models.vaes import WanVAEConfig
-    from fastvideo.models.loader.component_loader import get_diffusers_config
-    from fastvideo.models.vaes.wanvae import AutoencoderKLWan
-
-    diffusers_cfg = get_diffusers_config(model=str(fv_vae_dir))
-    diffusers_cfg.pop("_class_name", None)
-    diffusers_cfg.pop("_name_or_path", None)
-    fv_config = WanVAEConfig()
-    fv_config.load_encoder = False
-    fv_config.load_decoder = True
-    fv_config.update_model_arch(diffusers_cfg)
-    fv_vae = AutoencoderKLWan(fv_config).to(device=device, dtype=torch.float32)
-
-    # Mirror the VAE component loader: glob `*.safetensors`, merge, load
-    # non-strictly so any unused buffers (per_channel_statistics, etc.)
-    # don't fail the load.
-    sf_files = glob.glob(os.path.join(str(fv_vae_dir), "*.safetensors"))
-    assert sf_files, f"No safetensors files in {fv_vae_dir}"
-    state = {}
-    for sf in sf_files:
-        state.update(safetensors_load_file(sf))
-    fv_vae.load_state_dict(state, strict=False)
-    fv_vae.eval()
-
-    # Upstream's inner `_video_vae.decode(z, scale)` (line 874-877 of
-    # inference/model/vae2_2/vae2_2_module.py) does:
-    #     z = z / scale[1] + scale[0]    # where scale = [mean, 1/std]
-    #     = z * std + mean
-    # FastVideo's `AutoencoderKLWan.decode` expects the pre-denormalized
-    # latent — apply the same transform externally to feed both paths
-    # equivalently.
-    latents_mean = torch.tensor(
-        fv_config.arch_config.latents_mean, dtype=torch.float32, device=device,
-    )
-    latents_std = torch.tensor(
-        fv_config.arch_config.latents_std, dtype=torch.float32, device=device,
-    )
-    z_denormalized = z * latents_std.view(1, -1, 1, 1, 1) + latents_mean.view(1, -1, 1, 1, 1)
-    with torch.inference_mode():
-        fv_out_tensor = fv_vae.decode(z_denormalized)
-        fv_out = fv_out_tensor.detach().float().cpu()
-
-    # Both sides should return a video tensor of shape [..., C, T_dec, H_dec, W_dec].
-    # Normalize shapes for comparison — upstream returns a list per-video or a
-    # single tensor depending on CP group; we just squeeze batch dims.
-    def _squeeze(t):
-        while t.ndim > 4 and t.shape[0] == 1:
-            t = t[0]
-        return t
-
-    up_s = _squeeze(up_out)
-    fv_s = _squeeze(fv_out)
-    print(
-        f"up shape={tuple(up_s.shape)} abs_mean={up_s.abs().mean().item():.4f} "
-        f"range=[{up_s.min().item():.4f}, {up_s.max().item():.4f}]"
-    )
-    print(
-        f"fv shape={tuple(fv_s.shape)} abs_mean={fv_s.abs().mean().item():.4f} "
-        f"range=[{fv_s.min().item():.4f}, {fv_s.max().item():.4f}]"
-    )
-
-    # Wan VAE has a known fp32 op-ordering drift of ~8e-4 caused by
-    # `z * std + mean` (FV) vs `z / (1/std) + mean` (upstream) at decode
-    # normalization. This is a SHARED Wan-family bug, not magi-specific.
-    # Tracked as OQ-7 in tests/local_tests/magi-human.md; tighten to
-    # atol=1e-4 once the Wan VAE op-order fix lands.
-    assert up_s.shape == fv_s.shape, (
-        f"shape mismatch: up={up_s.shape} fv={fv_s.shape}"
-    )
-    diff = (up_s - fv_s).abs()
-    print(
-        f"diff max={diff.max().item():.6f} mean={diff.mean().item():.6f}"
-    )
-    assert_close(fv_s, up_s, atol=1e-3, rtol=1e-3)