[fix] dreamverse: steering session-flow, mock parity, and init-image validation

Steering protocol/UI fixes that interlock via one contract change — the
frontend now sends initial_rollout_prompt_id (its own prompt id for the
opening scene) in session-init and reset payloads, and both backend pre-seed
sites use it, so prompt lifecycle events (enhancing/ready/fallback) finally
match the frontend's records instead of no-oping against a backend uuid.

- controller: suppress prompt_sources_blocked while a submission is queued or
  enhancing (fixes the blank player at typed-prompt session start); segment
  prompt logging is now opt-in via DREAMVERSE_SEGMENT_PROMPT_LOG, written off
  the event loop, warning on failure (drops the hardcoded author-machine path
  and bare except).
- web: 'Generating next scene' overlay driven by explicit generation state
  instead of waitingForSegmentPrompt prop edges, so auto_extension_updated
  while idle can't strand it over the video; steering scene history is no
  longer truncated by the 24-event prompt feed cap and keeps stable numbers
  for 30+ scene sessions.
- mock server: manual_continuation_mode parity (no rewrite-flow wait, honors
  initial_rollout_prompt_id, no segment cap in manual mode) so the GPU-less
  dev backend works with the steering-only frontend.
- init image: client-side type/size validation on picker and paste paths
  (png/jpeg/webp, 15MB) with a visible error, and ws_max_size raised to 32MiB
  on both servers so a legitimate ~15MB image reaches the backend's own
  validation instead of tripping uvicorn's 16MiB frame cap.

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

@@ -1,20 +1,23 @@
 ---
 name: reseed-performance-baseline
-description: Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, or environment-caused benchmark shift using one or more reviewed normalized performance JSONs. 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 from a consistent batch of reviewed source results. The workflow backs up existing history under /tmp, validates all source JSONs for the same (model_id, gpu_type), rejects internally inconsistent source batches, uploads one success=true reseed record per accepted source JSON, and offers to clean local temp state after a successful upload.
+description: Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, environment-caused benchmark shift, or reviewed v2 calibration using one or more reviewed normalized performance JSONs. 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, or when a new v2 exact comparable identity needs its first approved baseline. The workflow backs up existing history under /tmp, validates all source JSONs for the same legacy (model_id, gpu_type) target or the same v2 exact identity, rejects internally inconsistent source batches, uploads one success=true baseline record per accepted source JSON, and offers to clean local temp state after a successful upload.
 ---
 
 # 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`.
+Replace or advance the rolling performance baseline in the HF dataset
+`FastVideo/performance-tracking`. Legacy targets are scoped by
+`(model_id, gpu_type)`. V2 targets are scoped by exact comparable identity:
+`workload_id`, `variant_id`, `benchmark_version`, `hardware_profile_id`,
+`software_profile_id`, and `recipe_fingerprint`.
 
-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`.
+Performance comparison uses the median of up to the last 5 successful,
+baseline-eligible records for the same target. Failed or calibration-only
+records are useful audit history, but they do not move the future baseline
+because `compare_baseline.py` loads records with `successful_only=True` and
+`baseline_eligible_only=True`.
 
 This skill now reseeds from a reviewed batch of one or more source performance
 JSONs. It uploads one new `success=true` record per accepted source JSON; it
@@ -22,11 +25,13 @@ does not blindly replicate one measurement into 3 or 5 records. The effective
 reseed size is therefore dynamic and equals the number of provided, validated,
 internally consistent source JSONs.
 
-If the operator provides fewer than 3 records, call out that the last-5 rolling
-median may not move immediately. If the operator provides 3 consistent shifted
-records, the rolling median usually moves immediately. If the operator provides
-5 consistent shifted records, the last-5 window is effectively reset to the new
-runtime profile.
+For baseline shifts with existing history, if the operator provides fewer than
+3 records, call out that the last-5 rolling median may not move immediately. If
+the operator provides 3 consistent shifted records, the rolling median usually
+moves immediately. If the operator provides 5 consistent shifted records, the
+last-5 window is effectively reset to the new runtime profile. For the first
+approved v2 baseline of a new exact identity, one reviewed calibration seed is
+enough for the next comparable run to leave `CALIBRATION_NEEDED`.
 
 These records are intentional operator-approved baseline resets, not ordinary
 independent main-branch persistence. Mark them clearly with provenance fields
@@ -66,8 +71,8 @@ approval, then upload reviewed accepted baseline records.
 
 | 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. |
+| `model_id` | Legacy required; v2 inferred | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. For v2 records, use the `model_id` from each source artifact only as the upload directory; comparison is by exact identity. |
+| `gpu_type` | Legacy required; v2 inferred | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. V2 hardware matching uses `hardware_profile_id`; preserve `gpu_type` as display metadata. |
 | `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
 | `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
 | `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. |
@@ -78,10 +83,14 @@ Hardcoded defaults:
   supported by the code, but use the default unless the user explicitly asks).
 - Local sync root: `/tmp/perf-tracking` (`PERFORMANCE_TRACKING_ROOT` override
   is supported).
+- Prepared-record staging root: `/tmp/performance_reseed_prepared`
+  (`PERFORMANCE_RESEED_STAGING_ROOT` override is supported). Keep it separate
+  and non-nested from the sync root.
 - Backup root: `/tmp/performance_reseed_backup`.
 - Download scratch root for source artifact URLs: `/tmp/performance_reseed_source`.
-- Baseline window: last 5 `success=true` records for the same
-  `(model_id, gpu_type)`.
+- Baseline window: last 5 `success=true`, `baseline_eligible=true` records
+  for the same legacy `(model_id, gpu_type)` target or the same v2 exact
+  comparable identity.
 - Reseed count: dynamic. Upload exactly one accepted seed record per validated
   source JSON.
 
@@ -115,12 +124,24 @@ with open(source_result, encoding="utf-8") as f:
     record = json.load(f)
 ```
 
-Stop if any normalized record's `model_id` or `gpu_type` does not match the
-requested `model_id` and `gpu_type`.
+Classify the source batch before continuing:
+
+- **Legacy source records** have no v2 exact identity fields. Stop if any
+  normalized record's `model_id` or `gpu_type` does not match the requested
+  `model_id` and `gpu_type`.
+- **V2 source records** have exact identity fields. Stop unless every source
+  record has all six comparable identity fields and they are identical across
+  the batch: `workload_id`, `variant_id`, `benchmark_version`,
+  `hardware_profile_id`, `software_profile_id`, and `recipe_fingerprint`.
+  Do not fall back to legacy `(model_id, gpu_type)` matching for v2 records.
 
 The source records may have `success: false` when they came from failed
 rolling baseline comparisons. That is expected; only the reviewed reseed
 records become new `success: true` baseline records after explicit approval.
+For a first v2 baseline seed, the source records must instead be successful
+scheduled-main full-suite `CALIBRATION_NEEDED` normalized artifacts. Reject PR,
+local, direct-run, non-main-branch, or non-full-suite calibration artifacts as
+seed sources.
 
 Sort validated source records by their original `timestamp` ascending before
 preparing the seed records. If a source timestamp is missing or unparsable,
@@ -194,7 +215,7 @@ export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
 python -c 'from fastvideo.performance.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 under `/tmp`:
+For legacy records, back up the sanitized model directory under `/tmp`:
 
 ```bash
 SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
@@ -209,6 +230,16 @@ mkdir -p "$BACKUP_DIR"
 cp -R "${PERFORMANCE_TRACKING_ROOT}/${MODEL_SAFE}" "$BACKUP_DIR/" 2>/dev/null || true
 ```
 
+For v2 records, back up the full local tracking root after sync. Exact identity
+lookup scans across model directories, so a benchmark rename may have relevant
+history outside the current source artifact's `model_id` directory:
+
+```bash
+BACKUP_DIR="/tmp/performance_reseed_backup/${TIMESTAMP}_${SHORT_COMMIT}_v2_exact_identity"
+mkdir -p "$BACKUP_DIR"
+cp -R "${PERFORMANCE_TRACKING_ROOT}" "$BACKUP_DIR/tracking-root"
+```
+
 Write provenance next to the backup:
 
 ```bash
@@ -231,7 +262,9 @@ 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:
+Load the last 5 successful baseline records for the target.
+
+For legacy targets:
 
 ```python
 from fastvideo.performance.hf_store import load_records_for_model
@@ -242,6 +275,28 @@ records = load_records_for_model(
     "<gpu_type>",
     last_n=5,
     successful_only=True,
+    baseline_eligible_only=True,
+)
+```
+
+For v2 exact-identity targets:
+
+```python
+from fastvideo.performance.hf_store import load_records_for_identity
+
+records = load_records_for_identity(
+    "/tmp/perf-tracking",
+    {
+        "workload_id": "<workload_id>",
+        "variant_id": "<variant_id>",
+        "benchmark_version": "<benchmark_version>",
+        "hardware_profile_id": "<hardware_profile_id>",
+        "software_profile_id": "<software_profile_id>",
+        "recipe_fingerprint": "<recipe_fingerprint>",
+    },
+    last_n=5,
+    successful_only=True,
+    baseline_eligible_only=True,
 )
 ```
 
@@ -257,7 +312,8 @@ medians after appending the proposed seed records, and source batch spread for:
 
 Also print how many successful old records exist. Make clear:
 
-- 1 seed record usually does not move a last-5 median by itself.
+- 1 seed record usually does not move an existing last-5 median by itself, but
+  it is enough to establish the first v2 baseline for a new exact identity.
 - 3 consistent seed records usually move the last-5 median immediately.
 - 5 consistent seed records effectively reset the last-5 window.
 - The records are intentional approved baseline resets and must be labeled
@@ -267,10 +323,10 @@ Also print how many successful old records exist. Make clear:
 
 Require an explicit confirmation phrase before preparing the upload:
 
-> About to RE-SEED performance baseline for `<model_id>` on `<gpu_type>`.
+> About to RE-SEED performance baseline for `<target description>`.
 > This will upload `<N>` new `success=true` records to
-> `FastVideo/performance-tracking/<sanitize(model_id)>/`, one per accepted
-> source JSON.
+> `FastVideo/performance-tracking/<sanitize(model_id)>/` or the source
+> artifact's v2 model directory, one per accepted source JSON.
 >
 > Reason: `<intent_rationale>`
 > Source results: `<source_results>`
@@ -288,25 +344,63 @@ Do not continue unless the user types exactly `confirm performance reseed`.
 
 ### 5. Create the accepted seed records
 
-Create one seed record from each normalized source result. Do not copy the
+Create one seed record from each normalized source result.
+
+For first v2 baseline seeds, use the scoped utility. It validates exact
+identity, requires successful scheduled-main full-suite `CALIBRATION_NEEDED`
+source artifacts, preserves the normalized v2 identity and metadata fields,
+and writes seed records with `success=true`, `baseline_eligible=true`, and
+`comparison_status=PASS`:
+
+```bash
+python fastvideo/tests/performance/seed_baseline.py \
+  --source-result <normalized_perf_1.json> \
+  --source-result <normalized_perf_2.json> \
+  --intent-rationale "<intent_rationale>" \
+  --max-intra-batch-regression 0.05 \
+  --tracking-root "${PERFORMANCE_TRACKING_ROOT}" \
+  --staging-root "${PERFORMANCE_RESEED_STAGING_ROOT:-/tmp/performance_reseed_prepared}"
+```
+
+The utility is prepare-only and intentionally has no upload option. Upload the
+scoped records only after the separate confirmation in step 6.
+
+The utility validates against an isolated fresh HF snapshot and leaves
+`PERFORMANCE_TRACKING_ROOT` untouched; that argument only proves the staging
+root is separate from the operator's tracking mirror. Before writing, it stops
+if the exact identity already has a successful baseline-eligible record or if
+the workload/variant/version already trusts another recipe. It atomically
+reserves the exact identity and writes a digest-protected upload manifest bound
+to the current HF endpoint, repository id, and repository type. Keep the
+prepared records, manifest, source files, and reservation unchanged until the
+operation is uploaded or explicitly cleaned up.
+
+If the prepared seed records look correct, upload only those scoped records in
+step 7. Do not rerun the utility with a different source list after approval.
+
+For legacy reseeds or accepted v2 baseline shifts from regression artifacts,
+create one seed record from each normalized source result. Do not copy the
 source JSON wholesale.
 
 Infer the baseline field allowlist from all existing HF records for the target
-`(model_id, gpu_type)` after syncing, including both `success=true` and
-`success=false` records. Use the union of non-provenance keys present in those
-target records, preserving only fields that also exist in the normalized
-source record or are explicitly set by the reseed workflow. Always include
-`model_id`, `timestamp`, and `success` because the upload path and baseline
-loader depend on them. Always set `timestamp` to a fresh reseed timestamp and
-`success` to `true`. Do not include unrelated source-only fields that are
-absent from existing HF records.
+after syncing, including both `success=true` and `success=false` records. For
+legacy targets the target is `(model_id, gpu_type)`. For v2 baseline-shift
+reseeds the target is the exact comparable identity. Use the union of
+non-provenance keys present in those target records, preserving only fields
+that also exist in the normalized source record or are explicitly set by the
+reseed workflow. Always include `model_id`, `timestamp`, `success`,
+`baseline_eligible`, and `comparison_status` because the upload path and
+baseline loader depend on them. Always set `timestamp` to a fresh reseed
+timestamp, `success` to `true`, `baseline_eligible` to `true`, and
+`comparison_status` to `PASS`. Do not include unrelated source-only fields
+that are absent from existing HF records.
 
 Exclude existing provenance or operator metadata from the inferred baseline
 field allowlist. At minimum, exclude keys prefixed with `baseline_reseed` and
 any fields known to be local-only audit metadata.
 
-If there are no previous HF records for the target model/GPU, fall back to this
-default baseline field list:
+If there are no previous HF records for the target, fall back to this default
+baseline field list:
 
 - `model_id`
 - `timestamp`
@@ -319,6 +413,22 @@ default baseline field list:
 - `dit_time_s`
 - `vae_decode_time_s`
 - `success`
+- `baseline_eligible`
+- `comparison_status`
+
+For v2 baseline-shift reseeds with no previous HF records for the exact
+identity, also preserve:
+
+- `workload_id`
+- `variant_id`
+- `benchmark_version`
+- `recipe_fingerprint`
+- `hardware_profile_id`
+- `software_profile_id`
+- `recipe`
+- `hardware_profile`
+- `software_profile`
+- `software_comparison_profile`
 
 Do not upload extra fields from the source artifact.
 
@@ -334,6 +444,22 @@ Optional provenance fields are allowed and useful:
 - `baseline_reseed_operator`
 - `baseline_reseed_max_intra_batch_regression`
 
+The v2 calibration seed utility writes analogous first-seed provenance:
+
+- `baseline_seed: true`
+- `baseline_seed_reason`
+- `baseline_seed_source_result`
+- `baseline_seed_source_status`
+- `baseline_seed_source_timestamp`
+- `baseline_seed_source_success`
+- `baseline_seed_source_run_source`
+- `baseline_seed_source_branch`
+- `baseline_seed_source_test_scope`
+- `baseline_seed_source_pr_number`
+- `baseline_seed_batch_size`
+- `baseline_seed_batch_index`
+- `baseline_seed_operator`
+
 Use a fresh reseed timestamp for each seed record, not the original source
 result timestamp. This is required because
 `load_records_for_model(..., last_n=5)` keeps the last records after loading
@@ -356,7 +482,8 @@ Prefer uploading new accepted seed records so failed history remains visible.
 Print:
 
 - Backup directory path under `/tmp`.
-- Prepared local record paths under `PERFORMANCE_TRACKING_ROOT`.
+- Prepared local record paths under `PERFORMANCE_RESEED_STAGING_ROOT`.
+- Prepared upload-manifest path under the identity reservation.
 - HF paths that will receive the new records.
 - Old rolling medians.
 - Source batch medians, source batch spread, reseed count, and candidate
@@ -368,22 +495,36 @@ 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:
+For a first v2 calibration seed, use the manifest uploader after the user
+replies exactly `upload`:
 
-```python
-from fastvideo.performance.hf_store import upload_record
-
-upload_record("<local_record_path>", record, strict=True)
+```bash
+python -c 'from fastvideo.tests.performance.seed_baseline import upload_prepared_seed_manifest; print(upload_prepared_seed_manifest("<prepared_manifest>"))'
 ```
 
-Run it once per prepared record. Each upload goes to:
+The uploader verifies the source and prepared-record digests, pins and scans
+the current HF revision, rechecks exact-identity and recipe-cohort conflicts,
+and writes the entire batch in one commit whose `parent_commit` must still be
+current. A concurrent Hub update makes the commit fail. Do not retry
+automatically: preserve staging, refresh/review remote state, and request a new
+explicit `upload` after the conflict is understood. Each record 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.
+Never call `upload_record()` once per first-seed record: that can partially
+land the batch and has no compare-and-swap guard.
+
+For a legacy reseed or an accepted v2 baseline shift, the first-seed manifest
+validator does not apply because an eligible baseline already exists. Upload
+only the individually reviewed records prepared in step 5 with the shared
+`upload_record(local_path, record, strict=True)` helper. Stop on the first
+failure and report exactly which records reached HF; do not silently rerun or
+replicate the remainder.
+
+Never bulk upload the tracking or staging root, and never modify another
+model's directory in the same operation.
 
 ### 8. Report outcome and offer cleanup
 
@@ -405,9 +546,14 @@ distinguish an accepted baseline shift from a hidden regression.
 After the upload is verified, ask whether the user wants to clear temporary
 local state. Explain what each directory is for:
 
-- `PERFORMANCE_TRACKING_ROOT`, usually `/tmp/perf-tracking`: local synced
-  mirror of `FastVideo/performance-tracking` plus the prepared local seed
-  records used for scoped upload.
+- `PERFORMANCE_TRACKING_ROOT`, usually `/tmp/perf-tracking`: read-only local
+  synced mirror used for operator review and reporting. First-v2 preparation
+  independently proves remote state from a fresh temporary HF snapshot.
+- `PERFORMANCE_RESEED_STAGING_ROOT`, usually
+  `/tmp/performance_reseed_prepared`: prepared local seed records used for the
+  scoped upload, plus the identity reservation and digest manifest. Keeping
+  this separate prevents aborted preparations from appearing in later
+  baseline reads.
 - `/tmp/performance_reseed_backup/<...>`: local backup of the target model's
   pre-reseed HF history plus `PROVENANCE.txt`, kept so a bad reseed can be
   audited or corrected.
@@ -417,14 +563,18 @@ local state. Explain what each directory is for:
 Ask:
 
 > Reseed succeeded. Do you want me to delete the local temp tracking mirror,
-> source downloads, and reseed backup under `/tmp`? These files are local
-> safety/audit artifacts only; HF already has the uploaded records.
+> this reseed's prepared staging records, source downloads, and reseed backup
+> under `/tmp`? These files are local safety/audit artifacts only; HF already
+> has the uploaded records.
 >
 > Reply `cleanup reseed temp` to delete them, anything else to keep them.
 
 Do not delete anything unless the user replies exactly
 `cleanup reseed temp`. If cleanup is requested, remove only the specific
-directories created for this reseed. Never remove unrelated `/tmp` contents.
+directories and prepared record paths created for this reseed. Do not remove
+the shared staging root when it contains other records. Remove this operation's
+identity reservation only with its prepared records and manifest, and never
+remove unrelated `/tmp` contents.
 
 ## Failure modes and handling
 
@@ -436,19 +586,34 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
   against the source batch median by more than `max_intra_batch_regression`.
   Ask for cleaner sources or a reviewed explanation before continuing.
 - **Too few source records to move the median.** Continue only after making
-  clear that one or two records may not immediately move the last-5 median.
+  clear that one or two records may not immediately move an existing last-5
+  median. This warning does not block a first v2 calibration seed for an exact
+  identity with no eligible baseline yet.
 - **The source results are noisy or suspicious.** Stop. Reseeding amplifies
   those measurements into the baseline, so they must be reviewed first.
 - **HF sync fails.** Stop for destructive reseeds. A stale or empty sync can
   make the old baseline look missing.
+- **The exact v2 identity already has an eligible baseline.** Stop. The
+  `CALIBRATION_NEEDED` artifact is stale; use the reviewed baseline-shift path
+  instead of the first-seed utility.
+- **The workload/variant/version trusts another recipe.** Stop. The source is
+  stale relative to the current recipe cohort and must not bypass
+  `RECIPE_MISMATCH` by creating a second trusted recipe.
+- **The staging root already has a prepared seed for the exact identity.**
+  Stop and reuse, upload, or explicitly clean that preparation. Do not prepare
+  another copy of the same measurement.
+- **The conditional Hub commit loses its parent race.** Stop without retrying.
+  Keep the preparation, refresh and review the new remote state, then request
+  a new explicit `upload` only if the seed is still valid.
 - **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.
-- **The user declines cleanup.** Keep `/tmp/perf-tracking`, the source
-  download directory if any, and `/tmp/performance_reseed_backup/<...>` in
-  place for audit/debugging.
+- **The user declines cleanup.** Keep `/tmp/perf-tracking`, the prepared seed
+  records under `/tmp/performance_reseed_prepared`, the source download
+  directory if any, and `/tmp/performance_reseed_backup/<...>` in place for
+  audit/debugging.
 - **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.
@@ -459,8 +624,9 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
   intentional baseline replacement.
 - `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
   median comparison, and persistence rules.
-- `fastvideo/performance/hf_store.py` — HF sync, record loading,
-  `sanitize()`, and `upload_record()`.
+- `fastvideo/performance/hf_store.py` — HF sync and record loading helpers.
+- `fastvideo/tests/performance/seed_baseline.py` — first-seed preparation,
+  staging reservation, manifest validation, and conditional batch upload.
 - `fastvideo/tests/performance/test_inference_performance.py` — source result
   JSON schema.
 - `.buildkite/performance-benchmarks/tests/*.json` — fixed absolute benchmark
@@ -473,3 +639,4 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
 | 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 | Previous 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. Superseded by the 2026-05-08 dynamic multi-source policy. |
 | 2026-05-08 | Replace fixed 3/5 replication with dynamic multi-source reseeding: upload one seed record per reviewed source JSON, validate intra-batch consistency, move backup/source scratch under `/tmp`, and ask whether to clean temp state after successful upload. |
+| 2026-07-13 | Keep first-v2-seed preparation outside the canonical mirror, reserve staging identities atomically, reject stale or replayed calibration seeds, and upload reviewed manifests with a single parent-guarded Hub commit. |

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

@@ -3,7 +3,7 @@
   "config_schema_version": 2,
   "workload_id": "wan-t2v",
   "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "benchmark_version": 3,
   "description": "Wan2.1 T2V 1.3B inference performance",
   "model": {
     "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",

.buildkite/pipeline.yml

@@ -436,7 +436,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 25m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: LoRA Training Tests"
                 env:
                   - TEST_TYPE=training_lora

.buildkite/scripts/pr_test.sh

@@ -76,7 +76,7 @@ EFFECTIVE_PR=${BUILDKITE_PULL_REQUEST:-false}
 if [ "$EFFECTIVE_PR" = "false" ] && [ -n "${PR_NUMBER:-}" ]; then
     EFFECTIVE_PR=$PR_NUMBER
 fi
-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} BUILDKITE_BUILD_URL=${BUILDKITE_BUILD_URL:-} BUILDKITE_BUILD_ID=${BUILDKITE_BUILD_ID:-} BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID:-} IMAGE_VERSION=$IMAGE_VERSION"
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} BUILDKITE_SOURCE=${BUILDKITE_SOURCE:-} TEST_SCOPE=${TEST_SCOPE:-} BUILDKITE_BUILD_URL=${BUILDKITE_BUILD_URL:-} BUILDKITE_BUILD_ID=${BUILDKITE_BUILD_ID:-} BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID:-} IMAGE_VERSION=$IMAGE_VERSION"
 
 POST_RUN_HOOK=""
 

.gitignore

@@ -72,8 +72,7 @@ docs/distillation/examples/
 # Python pickle files
 *.pkl
 
-# Reference videos
-!fastvideo/tests/ssim/reference_videos/**/*.mp4
+# Reference videos (negations must come after the catch-all on line below)
 
 # Static images
 !docs/assets/images/**/*.png
@@ -127,6 +126,8 @@ apps/dreamverse/web/.env.production.local
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
+!fastvideo/tests/ssim/reference_videos/**/*.mp4
+!fastvideo/tests/ssim/reference_videos/**/*.png
 
 # Editor logs and local Python version pins (accidentally committed)
 *.nvimlog

apps/dreamverse/dreamverse/config.py

@@ -67,7 +67,7 @@ MODEL_REGISTRY = {
     },
 }
 
-DEFAULT_MODEL_ID = "fast-ltx2"
+DEFAULT_MODEL_ID = "fast-ltx23"
 
 ACTIVE_MODEL_ID = (os.getenv("DREAMVERSE_MODEL_ID", "").strip() or DEFAULT_MODEL_ID)
 if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
@@ -76,14 +76,11 @@ if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
 # Active model configuration
 MODEL_CONFIG = MODEL_REGISTRY[ACTIVE_MODEL_ID]
 
-# Generation limits
-SESSION_TIMEOUT_SECONDS = 300
-
 # Frame settings
 NUM_FRAMES = 121
 FRAME_HEIGHT = 1088
 FRAME_WIDTH = 1920
-NUM_INFERENCE_STEPS = 5
+NUM_INFERENCE_STEPS = 6
 JPEG_QUALITY = 100
 BATCH_SIZE = 3
 
@@ -168,6 +165,9 @@ def _optional_env(*names: str) -> str | None:
     return None
 
 
+# Generation limits
+SESSION_TIMEOUT_SECONDS = _env_int("DREAMVERSE_SESSION_TIMEOUT_SECONDS", 300)
+
 DEVTOOLS_ENABLED = _env_bool("FASTVIDEO_ENABLE_DEVTOOLS", False)
 PROMPT_SAFETY_ENABLED = _env_bool("FASTVIDEO_ENABLE_PROMPT_SAFETY", False)
 DREAMVERSE_MAX_AUTOTUNE = _env_bool("DREAMVERSE_MAX_AUTOTUNE", True)

apps/dreamverse/dreamverse/gpu_pool.py

@@ -1023,7 +1023,11 @@ def get_available_gpus() -> list[int]:
     """Get list of available GPU IDs from environment or auto-detect."""
     cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
     if cuda_visible:
-        visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
+        try:
+            visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
+        except ValueError as exc:
+            raise RuntimeError("CUDA_VISIBLE_DEVICES must be a comma-separated list of integer GPU "
+                               f"indices (got {cuda_visible!r}); GPU UUIDs are not supported.") from exc
         return _limit_gpu_ids(visible_gpu_ids)
 
     # Auto-detect available GPUs

apps/dreamverse/dreamverse/main.py

@@ -206,7 +206,9 @@ def cli() -> None:
     args = parser.parse_args()
 
     _install_heartbeat_log_filter()
-    uvicorn.run(app, host=args.host, port=args.port)
+    # A 15MB init image (session_init_image.MAX_SESSION_INIT_IMAGE_BYTES) is ~20MB
+    # as a base64 ws message, above uvicorn's default 16MiB frame cap.
+    uvicorn.run(app, host=args.host, port=args.port, ws_max_size=32 * 1024 * 1024)
 
 
 if __name__ == "__main__":

apps/dreamverse/dreamverse/mock_server.py

@@ -30,11 +30,11 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 
 from dreamverse._deps import require_dreamverse_runtime_deps
-from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP
+from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP, SESSION_TIMEOUT_SECONDS
 from dreamverse.session_init_image import cleanup_session_init_image, persist_session_init_image
+from dreamverse.utils import _resolve_generation_segment_cap
 
 LATENCY_MS = 200
-SESSION_TIMEOUT_SECONDS = 300
 MOCK_FRAME_WIDTH = 640
 MOCK_FRAME_HEIGHT = 352
 MOCK_FPS = 24
@@ -334,6 +334,7 @@ async def websocket_endpoint(websocket: WebSocket):
         auto_extension_enabled = bool(init_data.get("auto_extension_enabled", False))
         loop_generation_enabled = bool(init_data.get("loop_generation_enabled", False))
         single_clip_mode = bool(init_data.get("single_clip_mode", False))
+        manual_continuation_mode = bool(init_data.get("manual_continuation_mode", False))
         generation_paused = False
 
         if init_type == "session_init_v2":
@@ -344,7 +345,8 @@ async def websocket_endpoint(websocket: WebSocket):
             incoming_prompts = []
 
         curated_prompts = [prompt.strip() for prompt in incoming_prompts if isinstance(prompt, str) and prompt.strip()]
-        generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
+        generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
+                                 and len(curated_prompts) == 0)
 
         try:
             session_init_image = persist_session_init_image(init_data.get("initial_image"))
@@ -373,7 +375,10 @@ async def websocket_endpoint(websocket: WebSocket):
         prompt_sources_blocked = False
         pending_seed_reset = False
         pending_seed_reset_reason = ""
-        pending_simple_submission: PromptSubmission | None = None
+        pending_simple_submission: PromptSubmission | None = (PromptSubmission(
+            prompt_id=str(init_data.get("initial_rollout_prompt_id") or uuid.uuid4()),
+            raw_prompt=initial_rollout_prompt,
+        ) if manual_continuation_mode and initial_rollout_prompt else None)
         single_clip_waiting_for_request = False
         rollout_waiting_for_rewrite = False
         initial_rollout_waiting_for_rewrite = generation_paused
@@ -393,14 +398,26 @@ async def websocket_endpoint(websocket: WebSocket):
 
         async def send_stream_start(seed_reason: str) -> None:
             await ws_send_json({
-                "type": "ltx2_stream_start",
-                "total_segments": len(curated_prompts),
-                "preset_id": preset_id,
-                "stream_mode": "av_fmp4",
-                "live_mode": True,
-                "loop_generation_enabled": loop_generation_enabled,
-                "loop_iteration": loop_iteration,
-                "generation_segment_cap": 0,
+                "type":
+                "ltx2_stream_start",
+                "total_segments":
+                len(curated_prompts),
+                "preset_id":
+                preset_id,
+                "stream_mode":
+                "av_fmp4",
+                "live_mode":
+                True,
+                "loop_generation_enabled":
+                loop_generation_enabled,
+                "loop_iteration":
+                loop_iteration,
+                "generation_segment_cap":
+                _resolve_generation_segment_cap(
+                    single_clip_mode=single_clip_mode,
+                    cap=GENERATION_SEGMENT_CAP,
+                    manual_continuation_mode=manual_continuation_mode,
+                ),
             })
             if seed_reason == "init":
                 await ws_send_json({
@@ -493,6 +510,7 @@ async def websocket_endpoint(websocket: WebSocket):
             nonlocal auto_extension_enabled
             nonlocal loop_generation_enabled
             nonlocal single_clip_mode
+            nonlocal manual_continuation_mode
             nonlocal generation_paused
             nonlocal seed_prompt_memory
             nonlocal curated_prompts
@@ -537,6 +555,7 @@ async def websocket_endpoint(websocket: WebSocket):
             auto_extension_enabled = bool(payload.get("auto_extension_enabled", False))
             loop_generation_enabled = bool(payload.get("loop_generation_enabled", False))
             single_clip_mode = bool(payload.get("single_clip_mode", False))
+            manual_continuation_mode = bool(payload.get("manual_continuation_mode", False))
 
             seed_prompt_memory = list(next_curated_prompts)
             curated_prompts = list(seed_prompt_memory)
@@ -544,10 +563,14 @@ async def websocket_endpoint(websocket: WebSocket):
             segment_idx = 0
             pending_seed_reset = False
             pending_seed_reset_reason = ""
-            pending_simple_submission = None
+            pending_simple_submission = (PromptSubmission(
+                prompt_id=str(payload.get("initial_rollout_prompt_id") or uuid.uuid4()),
+                raw_prompt=initial_rollout_prompt,
+            ) if manual_continuation_mode and initial_rollout_prompt else None)
             single_clip_waiting_for_request = False
             rollout_waiting_for_rewrite = False
-            generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
+            generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
+                                     and len(curated_prompts) == 0)
             initial_rollout_waiting_for_rewrite = generation_paused
             rewrite_restart_pending = False
             loop_iteration = 0
@@ -968,10 +991,11 @@ async def websocket_endpoint(websocket: WebSocket):
                     project_stream_started = True
                     await send_stream_start(pending_seed_reset_reason)
                     pending_seed_reset_reason = ""
-                    if pending_simple_submission is not None:
-                        submission = pending_simple_submission
-                        pending_simple_submission = None
-                        await promote_submission_to_ready(submission)
+
+                if pending_simple_submission is not None:
+                    submission = pending_simple_submission
+                    pending_simple_submission = None
+                    await promote_submission_to_ready(submission)
 
                 if generation_paused:
                     await asyncio.sleep(0.05)
@@ -981,8 +1005,8 @@ async def websocket_endpoint(websocket: WebSocket):
                     await asyncio.sleep(0.05)
                     continue
 
-                if (not single_clip_mode and not rollout_waiting_for_rewrite and GENERATION_SEGMENT_CAP > 0
-                        and segment_idx >= GENERATION_SEGMENT_CAP):
+                if (not single_clip_mode and not manual_continuation_mode and not rollout_waiting_for_rewrite
+                        and GENERATION_SEGMENT_CAP > 0 and segment_idx >= GENERATION_SEGMENT_CAP):
                     rollout_waiting_for_rewrite = True
                     loop_generation_enabled = False
                     project_stream_started = False
@@ -1217,7 +1241,9 @@ def cli() -> None:
     print(f"Starting mock server with {LATENCY_MS}ms latency on port {args.port}")
 
     _install_heartbeat_log_filter()
-    uvicorn.run(app, host="0.0.0.0", port=args.port)
+    # A 15MB init image (session_init_image.MAX_SESSION_INIT_IMAGE_BYTES) is ~20MB
+    # as a base64 ws message, above uvicorn's default 16MiB frame cap.
+    uvicorn.run(app, host="0.0.0.0", port=args.port, ws_max_size=32 * 1024 * 1024)
 
 
 if __name__ == "__main__":

apps/dreamverse/dreamverse/prompt_enhancer.py

@@ -313,6 +313,33 @@ def _extract_content_or_empty(response_json: dict[str, Any]) -> str:
         return ""
 
 
+def _find_balanced_object_end(text: str, start: int) -> int:
+    """Return the index just past the brace-balanced span opening at
+    ``text[start] == '{'``, honoring JSON string literals and escapes, or -1
+    if the braces never balance (i.e. the object was truncated)."""
+    depth = 0
+    in_string = False
+    escaped = False
+    for i in range(start, len(text)):
+        ch = text[i]
+        if in_string:
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == '"':
+                in_string = False
+        elif ch == '"':
+            in_string = True
+        elif ch == "{":
+            depth += 1
+        elif ch == "}":
+            depth -= 1
+            if depth == 0:
+                return i + 1
+    return -1
+
+
 def _parse_json_response(content: str) -> dict[str, Any]:
     text = content.strip()
     if not text:
@@ -330,6 +357,7 @@ def _parse_json_response(content: str) -> dict[str, Any]:
         r"```(?:json)?\s*([\s\S]*?)```",
         flags=re.IGNORECASE,
     )
+    last_fenced: dict[str, Any] | None = None
     for match in fence_pattern.finditer(text):
         block = match.group(1).strip()
         if not block:
@@ -337,21 +365,37 @@ def _parse_json_response(content: str) -> dict[str, Any]:
         try:
             parsed = json.loads(block)
             if isinstance(parsed, dict):
-                return parsed
+                last_fenced = parsed
         except json.JSONDecodeError:
             continue
+    if last_fenced is not None:
+        return last_fenced
 
-    # Fall back to scanning for the first decodable JSON object in free-form text.
+    # Scan for all decodable JSON objects and return the last — chain-of-thought
+    # models emit draft JSON mid-reasoning; the final answer is always last.
     decoder = json.JSONDecoder()
-    for idx, char in enumerate(text):
-        if char != "{":
-            continue
+    last_parsed: dict[str, Any] | None = None
+    pos = 0
+    while (idx := text.find("{", pos)) != -1:
         try:
-            parsed, _ = decoder.raw_decode(text[idx:])
+            parsed, consumed = decoder.raw_decode(text[idx:])
         except json.JSONDecodeError:
+            # Skip the whole failed object rather than rescanning inside it:
+            # fragments nested in a malformed or truncated (finish_reason=
+            # length) object must not override an earlier complete object.
+            span_end = _find_balanced_object_end(text, idx)
+            if span_end == -1:
+                break
+            pos = span_end
             continue
+        # Skip past the consumed span so nested braces inside a decoded
+        # object are not re-parsed as standalone objects.
+        pos = idx + consumed
         if isinstance(parsed, dict):
-            return parsed
+            last_parsed = parsed
+
+    if last_parsed is not None:
+        return last_parsed
 
     raise ValueError("No JSON object found in assistant response.")
 
@@ -381,7 +425,7 @@ def _format_locked_segments(locked_segments: list[str]) -> str:
 
 class PromptEnhancer:
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.provider = PROMPT_PROVIDER
         self.provider_label = _resolve_provider_label(PROMPT_PROVIDER)
         self.api_key = PROMPT_API_KEY
@@ -1417,15 +1461,12 @@ class PromptEnhancer:
             locked_text = _format_locked_segments(locked_segments_clean)
             request_system_prompt = self.enhance_system_prompt
             user_payload = {
-                "request": (
-                    "<locked_segments>\n"
-                    f"{locked_text}\n"
-                    "</locked_segments>\n\n"
-                    f"<conditioning_prompt>{cleaned}</conditioning_prompt>\n\n"
-                    f"Write exactly one new segment ({next_segment_key}) "
-                    "continuing from the locked segments. "
-                    'Respond with valid JSON only as {"next_prompt": "..."}.'  # noqa: E501
-                ),
+                "request": ("<locked_segments>\n"
+                            f"{locked_text}\n"
+                            "</locked_segments>\n\n"
+                            f"<conditioning_prompt>{cleaned}</conditioning_prompt>\n\n"
+                            f"Write exactly one new segment ({next_segment_key}) "
+                            "continuing from the locked segments."),
             }
 
         t0 = time.perf_counter()
@@ -1457,6 +1498,7 @@ class PromptEnhancer:
                             body=request_body,
                             timeout_seconds=timeout_seconds,
                         )
+                        _enhance_print("INFO", f"raw_response: {response_content}")
                         if is_single_clip_mode:
                             prompt = self._extract_single_clip_prompt(response_content)
                         else:
@@ -1539,7 +1581,7 @@ class PromptEnhancer:
                 f"Write exactly one new segment ({next_segment_key}) "
                 "that continues linearly from the locked segments. "
                 "Infer the next narrative beat from this history. "
-                'Respond with valid JSON only as {"next_prompt": "..."}.'  # noqa: E501
+                'Respond with valid JSON only: {"next_prompt": "<your segment description here>"}.'  # noqa: E501
             ),
         }
 

apps/dreamverse/dreamverse/session/controller.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 # mypy: ignore-errors
 
 import asyncio
+import os
 import time
 import uuid
 from typing import TYPE_CHECKING
@@ -52,6 +53,14 @@ if TYPE_CHECKING:
     from dreamverse.prompt_enhancer import PromptEnhancer
     from dreamverse.prompt_safety import PromptSafetyFilter
 
+# Optional append-only log of every generated segment prompt; unset disables it.
+SEGMENT_PROMPT_LOG_PATH = os.environ.get("DREAMVERSE_SEGMENT_PROMPT_LOG", "")
+
+
+def _append_segment_prompt_log(path: str, text: str) -> None:
+    with open(path, "a") as f:
+        f.write(text)
+
 
 class SessionController:
     """Runs one WebSocket session from accept() through disconnect."""
@@ -198,6 +207,7 @@ class SessionController:
             auto_extension_enabled = bool(init_data.get("auto_extension_enabled", False))
             loop_generation_enabled = bool(init_data.get("loop_generation_enabled", False))
             single_clip_mode = bool(init_data.get("single_clip_mode", False))
+            manual_continuation_mode = bool(init_data.get("manual_continuation_mode", False))
             rewrite_model = self.prompt_enhancer.resolve_rewrite_model(init_data.get("rewrite_model"))
             rewrite_system_prompt_override = str(init_data.get("rewrite_window_system_prompt") or "").strip()
             rewrite_user_system_prompt_override = str(init_data.get("rewrite_user_system_prompt") or "").strip()
@@ -282,6 +292,9 @@ class SessionController:
             # Session queues and mutable state.
             raw_prompt_queue: asyncio.Queue[PromptSubmission] = asyncio.Queue()
             ready_prompt_queue: asyncio.Queue[ReadyPrompt] = asyncio.Queue()
+            # Submissions dequeued by prompt_worker_loop but not yet resolved; while
+            # non-zero the prompt sources are busy, not drained.
+            prompt_enhancement_inflight = 0
 
             curated_idx = 0
             segment_idx = 0
@@ -291,13 +304,20 @@ class SessionController:
             generation_cap_blocked = False
             auto_extension_blocked_segment_idx: int | None = None
             prompt_sources_drained_logged = False
-            generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
+            generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
+                                     and len(curated_prompts) == 0)
             pending_seed_reset = False
             pending_seed_reset_reason = ""
             pending_reset_conditioning = False
             loop_iteration = 0 if generation_paused else 1
             force_curated_restart_segment = False
-            pending_simple_prompt_submission: PromptSubmission | None = None
+            # The frontend records the opening scene under this id; reuse it so
+            # prompt lifecycle events for the opening prompt reach that record.
+            pending_simple_prompt_submission: PromptSubmission | None = (PromptSubmission(
+                prompt_id=str(init_data.get("initial_rollout_prompt_id") or uuid.uuid4()),
+                raw_prompt=initial_rollout_prompt,
+                created_at_s=time.time(),
+            ) if manual_continuation_mode and initial_rollout_prompt else None)
             single_clip_waiting_for_request = False
             rollout_waiting_for_rewrite = False
             initial_rollout_waiting_for_rewrite = generation_paused
@@ -306,6 +326,7 @@ class SessionController:
             project_active = True
             project_stream_started = False
             pending_project_end = False
+            segment_prompt_log_warned = False
 
             def replace_session_init_image(initial_image_payload: object) -> None:
                 nonlocal session_init_image
@@ -424,6 +445,7 @@ class SessionController:
                 nonlocal auto_extension_enabled
                 nonlocal loop_generation_enabled
                 nonlocal single_clip_mode
+                nonlocal manual_continuation_mode
                 nonlocal generation_paused
                 nonlocal curated_prompts
                 nonlocal seed_prompt_memory
@@ -458,6 +480,7 @@ class SessionController:
                 next_auto_extension_enabled = bool(payload.get("auto_extension_enabled", False))
                 next_loop_generation_enabled = bool(payload.get("loop_generation_enabled", False))
                 next_single_clip_mode = bool(payload.get("single_clip_mode", False))
+                next_manual_continuation_mode = bool(payload.get("manual_continuation_mode", False))
 
                 next_preset_id = str(payload.get("preset_id") or "").strip()
                 if next_preset_id:
@@ -510,6 +533,7 @@ class SessionController:
                 auto_extension_enabled = next_auto_extension_enabled
                 loop_generation_enabled = next_loop_generation_enabled
                 single_clip_mode = next_single_clip_mode
+                manual_continuation_mode = next_manual_continuation_mode
                 rewrite_model = next_rewrite_model
                 rewrite_system_prompt_override = (next_rewrite_system_prompt_override)
                 rewrite_user_system_prompt_override = (next_rewrite_user_system_prompt_override)
@@ -530,10 +554,15 @@ class SessionController:
                 generation_cap_blocked = False
                 auto_extension_blocked_segment_idx = None
                 prompt_sources_drained_logged = False
-                pending_simple_prompt_submission = None
+                pending_simple_prompt_submission = (PromptSubmission(
+                    prompt_id=str(payload.get("initial_rollout_prompt_id") or uuid.uuid4()),
+                    raw_prompt=initial_rollout_prompt,
+                    created_at_s=time.time(),
+                ) if manual_continuation_mode and initial_rollout_prompt else None)
                 single_clip_waiting_for_request = False
                 rollout_waiting_for_rewrite = False
-                generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
+                generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt
+                                         and not single_clip_mode and len(curated_prompts) == 0)
                 initial_rollout_waiting_for_rewrite = generation_paused
                 rewrite_restart_pending = False
                 loop_iteration = 0
@@ -942,6 +971,7 @@ class SessionController:
                             _resolve_generation_segment_cap(
                                 single_clip_mode=single_clip_mode,
                                 cap=GENERATION_SEGMENT_CAP,
+                                manual_continuation_mode=manual_continuation_mode,
                             ),
                         })
                         continue
@@ -1005,11 +1035,9 @@ class SessionController:
                         continue
 
             async def prompt_worker_loop():
-                while not stop_event.is_set():
-                    try:
-                        submission = await asyncio.wait_for(raw_prompt_queue.get(), timeout=0.1)
-                    except asyncio.TimeoutError:
-                        continue
+                nonlocal prompt_enhancement_inflight
+
+                async def process_submission(submission: PromptSubmission) -> None:
                     _main_print("INFO", f"Received user prompt for enhancement: {submission.raw_prompt}")
                     prompt_id = submission.prompt_id
                     raw_prompt = submission.raw_prompt
@@ -1032,8 +1060,9 @@ class SessionController:
                         await ws_send_json({
                             "type": "error",
                             "message": blocked_raw_prompt_error,
+                            "prompt_id": prompt_id,
                         })
-                        continue
+                        return
                     await log_event(
                         "enhance_request",
                         {
@@ -1093,8 +1122,9 @@ class SessionController:
                             await ws_send_json({
                                 "type": "error",
                                 "message": blocked_final_prompt_error,
+                                "prompt_id": prompt_id,
                             })
-                            continue
+                            return
                         if result.fallback_used or not final_prompt:
                             source = "user_enhancement_failed"
                             _main_print(
@@ -1113,7 +1143,7 @@ class SessionController:
                             })
                             # Enhancement is strict JSON-only; do not enqueue raw
                             # prompt when enhancement fails.
-                            continue
+                            return
                         else:
                             source = "user_enhanced"
                             await ws_send_json({
@@ -1130,6 +1160,7 @@ class SessionController:
                                 source=source,
                                 fallback_used=result.fallback_used,
                                 loop_iteration=loop_iteration,
+                                raw_prompt=raw_prompt,
                             ))
                     else:
                         await ready_prompt_queue.put(
@@ -1139,6 +1170,7 @@ class SessionController:
                                 source="user_raw",
                                 fallback_used=False,
                                 loop_iteration=loop_iteration,
+                                raw_prompt=raw_prompt,
                             ))
                         await ws_send_json({
                             "type": "prompt_ready",
@@ -1148,6 +1180,21 @@ class SessionController:
                             "latency_ms": 0.0,
                         })
 
+                while not stop_event.is_set():
+                    try:
+                        submission = raw_prompt_queue.get_nowait()
+                    except asyncio.QueueEmpty:
+                        await asyncio.sleep(0.1)
+                        continue
+                    # Dequeue and increment without an await in between so the
+                    # generation loop never sees an empty queue with zero in flight
+                    # while this submission is still being enhanced.
+                    prompt_enhancement_inflight += 1
+                    try:
+                        await process_submission(submission)
+                    finally:
+                        prompt_enhancement_inflight -= 1
+
             def queue_snapshot() -> dict[str, object]:
                 return {
                     "user_ready": ready_prompt_queue.qsize(),
@@ -1300,6 +1347,7 @@ class SessionController:
                     _resolve_generation_segment_cap(
                         single_clip_mode=single_clip_mode,
                         cap=GENERATION_SEGMENT_CAP,
+                        manual_continuation_mode=manual_continuation_mode,
                     ),
                 })
                 await ws_send_json({
@@ -1359,6 +1407,7 @@ class SessionController:
                         _resolve_generation_segment_cap(
                             single_clip_mode=single_clip_mode,
                             cap=GENERATION_SEGMENT_CAP,
+                            manual_continuation_mode=manual_continuation_mode,
                         ),
                     })
                     if nonlocal_reason == "loop_restart":
@@ -1388,8 +1437,9 @@ class SessionController:
                     await raw_prompt_queue.put(pending_simple_prompt_submission)
                     pending_simple_prompt_submission = None
 
-                if (not single_clip_mode and not generation_cap_blocked and not rollout_waiting_for_rewrite
-                        and GENERATION_SEGMENT_CAP > 0 and generated_segment_count >= GENERATION_SEGMENT_CAP):
+                if (not single_clip_mode and not manual_continuation_mode and not generation_cap_blocked
+                        and not rollout_waiting_for_rewrite and GENERATION_SEGMENT_CAP > 0
+                        and generated_segment_count >= GENERATION_SEGMENT_CAP):
                     loop_generation_enabled = False
                     rollout_waiting_for_rewrite = True
                     _main_print(
@@ -1542,7 +1592,10 @@ class SessionController:
                     if single_clip_mode:
                         await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0)
                         continue
-                    if not prompt_sources_drained_logged:
+                    # A raw submission still queued or being enhanced will produce a
+                    # ready prompt shortly; that is not a drained/blocked state.
+                    enhancement_pending = (raw_prompt_queue.qsize() > 0 or prompt_enhancement_inflight > 0)
+                    if not prompt_sources_drained_logged and not enhancement_pending:
                         snapshot = queue_snapshot()
                         _main_print(
                             "WARN",
@@ -1581,6 +1634,24 @@ class SessionController:
                 total_segments_hint = max(segment_idx, len(curated_prompts))
                 prompt = selected.prompt
                 locked_segment_prompts.append(prompt)
+                if SEGMENT_PROMPT_LOG_PATH:
+                    _ts = time.strftime("%Y-%m-%d %H:%M:%S")
+                    _lines = [
+                        f"\n=== Segment {segment_idx} [{_ts}] source={selected.source} client={client_id[:8]} ===",
+                    ]
+                    if selected.raw_prompt and selected.raw_prompt != prompt:
+                        _lines.append(f"User:     {selected.raw_prompt}")
+                    _lines.append(f"Rewritten: {prompt}")
+                    try:
+                        await asyncio.to_thread(_append_segment_prompt_log, SEGMENT_PROMPT_LOG_PATH,
+                                                "\n".join(_lines) + "\n")
+                    except Exception as exc:
+                        if not segment_prompt_log_warned:
+                            segment_prompt_log_warned = True
+                            _main_print(
+                                "WARN",
+                                f"Failed to write segment prompt log {SEGMENT_PROMPT_LOG_PATH}: {exc}",
+                            )
                 if (auto_extension_blocked_segment_idx is not None
                         and auto_extension_blocked_segment_idx <= segment_idx):
                     auto_extension_blocked_segment_idx = None

apps/dreamverse/dreamverse/session/messages.py

@@ -19,3 +19,4 @@ class ReadyPrompt:
     fallback_used: bool = False
     seed_prompt_index: int | None = None
     loop_iteration: int | None = None
+    raw_prompt: str | None = None

apps/dreamverse/dreamverse/tests/test_config.py

@@ -150,12 +150,22 @@ def test_config_enables_prompt_safety_when_requested(monkeypatch):
 
 def test_config_uses_five_minute_session_timeout(monkeypatch):
     _set_required_prompt_keys(monkeypatch)
+    monkeypatch.delenv("DREAMVERSE_SESSION_TIMEOUT_SECONDS", raising=False)
 
     module = _load_config_module()
 
     assert module.SESSION_TIMEOUT_SECONDS == 300
 
 
+def test_config_session_timeout_env_override(monkeypatch):
+    _set_required_prompt_keys(monkeypatch)
+    monkeypatch.setenv("DREAMVERSE_SESSION_TIMEOUT_SECONDS", "1800")
+
+    module = _load_config_module()
+
+    assert module.SESSION_TIMEOUT_SECONDS == 1800
+
+
 def test_config_rejects_invalid_prompt_provider(monkeypatch):
     monkeypatch.setenv("FASTVIDEO_PROMPT_PROVIDER", "unsupported")
     _set_required_prompt_keys(monkeypatch)

apps/dreamverse/dreamverse/tests/test_entrypoints.py

@@ -75,12 +75,13 @@ def _run_cli(module, monkeypatch, argv: list[str]) -> list[dict[str, object]]:
     calls: list[dict[str, object]] = []
     uvicorn_stub = types.ModuleType("uvicorn")
 
-    def run(app, host: str, port: int) -> None:
+    def run(app, host: str, port: int, **kwargs) -> None:
         calls.append(
             {
                 "app": app,
                 "host": host,
                 "port": port,
+                **kwargs,
             }
         )
 
@@ -104,6 +105,7 @@ def test_server_cli_defaults_to_local_web_port(monkeypatch):
             "app": server_main.app,
             "host": "0.0.0.0",
             "port": 8009,
+            "ws_max_size": 32 * 1024 * 1024,
         }
     ]
 
@@ -121,6 +123,7 @@ def test_server_cli_allows_explicit_host_and_port(monkeypatch):
             "app": server_main.app,
             "host": "127.0.0.1",
             "port": 8123,
+            "ws_max_size": 32 * 1024 * 1024,
         }
     ]
 
@@ -147,6 +150,7 @@ def test_mock_server_cli_defaults_to_local_web_port(monkeypatch):
             "app": mock_server.app,
             "host": "0.0.0.0",
             "port": 8009,
+            "ws_max_size": 32 * 1024 * 1024,
         }
     ]
 
@@ -166,6 +170,7 @@ def test_mock_server_cli_updates_latency(monkeypatch):
                 "app": mock_server.app,
                 "host": "0.0.0.0",
                 "port": 8111,
+                "ws_max_size": 32 * 1024 * 1024,
             }
         ]
         assert mock_server.LATENCY_MS == 321

apps/dreamverse/dreamverse/tests/test_mock_server.py

@@ -290,6 +290,152 @@ def test_mock_server_supports_initial_custom_rollout_prompt():
         mock_server.LATENCY_MS = old_latency_ms
 
 
+def test_mock_server_manual_mode_streams_initial_prompt_without_rewrite_or_cap():
+    old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
+    old_latency_ms = mock_server.LATENCY_MS
+    old_generation_segment_cap = mock_server.GENERATION_SEGMENT_CAP
+    try:
+        mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
+        mock_server.LATENCY_MS = 1
+        mock_server.GENERATION_SEGMENT_CAP = 1
+
+        ws = _FakeWebSocket(
+            [
+                (
+                    0.0,
+                    {
+                        "type": "session_init_v2",
+                        "preset_id": "custom_editable",
+                        "preset_label": "Custom rollout",
+                        "curated_prompts": [],
+                        "initial_rollout_prompt": "A drone skims a neon canyon",
+                        "initial_rollout_prompt_id": "steer-prompt-1",
+                        "manual_continuation_mode": True,
+                        "enhancement_enabled": True,
+                        "auto_extension_enabled": False,
+                        "loop_generation_enabled": False,
+                    },
+                ),
+                (
+                    0.08,
+                    {
+                        "type": "append_prompt",
+                        "prompt": "The drone dives toward the river",
+                        "prompt_id": "steer-prompt-2",
+                    },
+                ),
+                (0.30, {"type": "leave"}),
+            ]
+        )
+
+        asyncio.run(mock_server.websocket_endpoint(ws))
+
+        message_types = [payload["type"] for payload in ws.sent_json]
+        assert "rewrite_seed_prompts_started" not in message_types
+        assert "rewrite_seed_prompts_complete" not in message_types
+        assert "ltx2_stream_start" in message_types
+        # cap=1 must not stop a manual-mode session after the first segment
+        assert "ltx2_stream_complete" not in message_types
+
+        prompt_ready_events = [
+            payload for payload in ws.sent_json if payload["type"] == "prompt_ready"
+        ]
+        assert [payload["prompt_id"] for payload in prompt_ready_events] == [
+            "steer-prompt-1",
+            "steer-prompt-2",
+        ]
+
+        segment_start_events = [
+            payload
+            for payload in ws.sent_json
+            if payload["type"] == "ltx2_segment_start"
+        ]
+        assert [payload["prompt"] for payload in segment_start_events] == [
+            "A drone skims a neon canyon",
+            "The drone dives toward the river",
+        ]
+        segment_source_events = [
+            payload
+            for payload in ws.sent_json
+            if payload["type"] == "segment_prompt_source"
+        ]
+        assert [payload["prompt_id"] for payload in segment_source_events] == [
+            "steer-prompt-1",
+            "steer-prompt-2",
+        ]
+    finally:
+        mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
+        mock_server.LATENCY_MS = old_latency_ms
+        mock_server.GENERATION_SEGMENT_CAP = old_generation_segment_cap
+
+
+def test_mock_server_project_init_manual_mode_streams_initial_prompt():
+    old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
+    old_latency_ms = mock_server.LATENCY_MS
+    try:
+        mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
+        mock_server.LATENCY_MS = 1
+
+        ws = _FakeWebSocket(
+            [
+                (
+                    0.0,
+                    {
+                        "type": "session_init_v2",
+                        "preset_id": "test_preset",
+                        "curated_prompts": ["segment one"],
+                        "enhancement_enabled": True,
+                        "auto_extension_enabled": False,
+                        "loop_generation_enabled": False,
+                    },
+                ),
+                (0.05, {"type": "end_project_keep_session"}),
+                (
+                    0.15,
+                    {
+                        "type": "project_init_v1",
+                        "preset_id": "custom_editable",
+                        "preset_label": "Custom rollout",
+                        "curated_prompts": [],
+                        "initial_rollout_prompt": "A drone skims a neon canyon",
+                        "initial_rollout_prompt_id": "steer-prompt-1",
+                        "manual_continuation_mode": True,
+                        "enhancement_enabled": True,
+                        "auto_extension_enabled": False,
+                        "loop_generation_enabled": False,
+                    },
+                ),
+                (0.45, {"type": "leave"}),
+            ]
+        )
+
+        asyncio.run(mock_server.websocket_endpoint(ws))
+
+        message_types = [payload["type"] for payload in ws.sent_json]
+        assert "project_idle" in message_types
+        project_idle_index = message_types.index("project_idle")
+        # manual-mode restart must not run the rewrite rollout
+        assert "rewrite_seed_prompts_started" not in message_types[project_idle_index:]
+        assert "ltx2_stream_start" in message_types[project_idle_index:]
+
+        prompt_ready_events = [
+            payload for payload in ws.sent_json if payload["type"] == "prompt_ready"
+        ]
+        assert [payload["prompt_id"] for payload in prompt_ready_events] == [
+            "steer-prompt-1",
+        ]
+
+        segment_start_events = [
+            payload
+            for payload in ws.sent_json
+            if payload["type"] == "ltx2_segment_start"
+        ]
+        assert segment_start_events[-1]["prompt"] == "A drone skims a neon canyon"
+    finally:
+        mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
+        mock_server.LATENCY_MS = old_latency_ms
+
+
 def test_mock_server_can_start_new_project_without_reconnecting():
     old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
     old_latency_ms = mock_server.LATENCY_MS

apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py

@@ -6,6 +6,7 @@ import os
 import re
 import time
 
+import pytest
 
 os.environ.setdefault("CEREBRAS_API_KEY", "dummy")
 os.environ.setdefault("GROQ_API_KEY", "dummy")
@@ -205,6 +206,59 @@ def test_parse_json_response_extracts_first_embedded_object():
     assert parsed == {"segment_prompts": ["A", "B"]}
 
 
+def test_parse_json_response_returns_outer_object_not_nested_value():
+    parsed = _parse_json_response(
+        "Final answer: {\"next_prompt\": \"a scene\", \"style\": {\"mood\": \"noir\"}} done."
+    )
+    assert parsed == {"next_prompt": "a scene", "style": {"mood": "noir"}}
+
+
+def test_parse_json_response_returns_last_of_multiple_objects():
+    parsed = _parse_json_response(
+        "Draft: {\"next_prompt\": \"draft\"}\nRefined: {\"next_prompt\": \"final\"}"
+    )
+    assert parsed == {"next_prompt": "final"}
+
+
+def test_parse_json_response_ignores_fragments_of_truncated_trailing_object():
+    # finish_reason=length cut the refined object short; the complete draft must
+    # win over a nested fragment of the truncated object.
+    parsed = _parse_json_response(
+        '{"next_prompt": "draft"} refined: {"next_prompt": "final", "style": {"mood": "noir"}'
+    )
+    assert parsed == {"next_prompt": "draft"}
+
+
+def test_parse_json_response_ignores_fragments_of_mid_string_truncated_object():
+    # Unterminated-string truncation reports the error at the opening quote,
+    # not end-of-text; nested fragments still must not win over the draft.
+    parsed = _parse_json_response(
+        '{"next_prompt": "draft"} refined: {"style": {"mood": "noir"}, "next_prompt": "cut off'
+    )
+    assert parsed == {"next_prompt": "draft"}
+
+
+def test_parse_json_response_ignores_fragments_of_malformed_object_with_trailing_prose():
+    parsed = _parse_json_response(
+        '{"next_prompt": "draft"} {"final": {"mood": "noir"}, "x": 1 and then some prose'
+    )
+    assert parsed == {"next_prompt": "draft"}
+
+
+def test_parse_json_response_raises_when_only_object_is_truncated():
+    with pytest.raises(ValueError):
+        _parse_json_response('{"style": {"mood": "noir"}, "next_prompt": "cut off')
+
+
+def test_parse_json_response_returns_outer_rollout_dict():
+    parsed = _parse_json_response(
+        "{\"rollout\": {\"segment_prompts\": [{\"prompt\": \"a\"}, {\"prompt\": \"b\"}]}}"
+    )
+    assert parsed == {
+        "rollout": {"segment_prompts": [{"prompt": "a"}, {"prompt": "b"}]}
+    }
+
+
 def test_load_prompt_required_falls_back_to_default_path(tmp_path):
     fallback_path = tmp_path / "next_segment_system_prompt.md"
     fallback_path.write_text("fallback prompt\n", encoding="utf-8")

apps/dreamverse/dreamverse/utils.py

@@ -15,5 +15,6 @@ def _utc_now_iso() -> str:
 PROMPT_EXTENSION_FAILURE_USER_MESSAGE = ("Prompt extension failed for this request.")
 
 
-def _resolve_generation_segment_cap(*, single_clip_mode: bool, cap: int) -> int:
-    return 0 if single_clip_mode else cap
+def _resolve_generation_segment_cap(*, single_clip_mode: bool, cap: int, manual_continuation_mode: bool = False) -> int:
+    # Steering (manual continuation) lets the user keep going indefinitely, like single-clip mode.
+    return 0 if (single_clip_mode or manual_continuation_mode) else cap

apps/dreamverse/dreamverse/video_generation.py

@@ -489,7 +489,7 @@ class VideoGenerationWorker:
             num_inference_steps=NUM_INFERENCE_STEPS,
             guidance_scale=1.0,
             seed=10,
-            ltx2_image_crf=0.0,
+            ltx2_image_crf=(33.0 if image_path and segment_idx == 1 else 0.0),
             image_path=image_path if segment_idx == 1 else None,
             return_continuation_state=False,
         )

apps/dreamverse/pyproject.toml

@@ -29,6 +29,9 @@ test = [
 dreamverse-server = "dreamverse.server_entry:cli"
 dreamverse-mock-server = "dreamverse.mock_server:cli"
 
+[tool.setuptools.packages.find]
+include = ["dreamverse*"]
+
 [tool.uv]
 package = false
 

apps/dreamverse/scripts/launch-dreamverse.sh

@@ -0,0 +1,67 @@
+#!/usr/bin/env bash
+# launch-dreamverse.sh — launch dreamverse-server on a compute node.
+#
+# Usage (from repo root):
+#   bash apps/dreamverse/scripts/launch-dreamverse.sh                          # GPUs 0-3, SP_SIZE=4
+#   CUDA_VISIBLE_DEVICES=0 bash apps/dreamverse/scripts/launch-dreamverse.sh   # single GPU, SP_SIZE=1
+
+set -euo pipefail
+
+CONDA_PREFIX="$HOME/miniconda3/envs/dreamverse"
+CUDA_RT_DIR="$CONDA_PREFIX/lib/python3.11/site-packages/nvidia/cuda_runtime/lib"
+GXX="$CONDA_PREFIX/bin/aarch64-conda-linux-gnu-g++"
+
+
+export CUDA_HOME="$CONDA_PREFIX"
+export FASTVIDEO_ENABLE_STARTUP_WARMUP=true
+export FASTVIDEO_ENABLE_PROMPT_SAFETY=false
+export DREAMVERSE_MAX_AUTOTUNE=true
+export LTX2_USE_DISTILLED_SIGMAS=0
+export LTX2_VIDEO_CONDITIONING_NUM_FRAMES=1
+export AUDIO_CONDITIONING_NUM_FRAMES=41
+export DREAMVERSE_SESSION_TIMEOUT_SECONDS="${DREAMVERSE_SESSION_TIMEOUT_SECONDS:-1800}"
+# GB200 max-autotune warmup compiles can run for hours; keep the watchdog generous here
+export FASTVIDEO_STARTUP_WARMUP_TIMEOUT_SECONDS="${FASTVIDEO_STARTUP_WARMUP_TIMEOUT_SECONDS:-24000}"
+export CEREBRAS_API_KEY="${CEREBRAS_API_KEY:-}"  # set this in your env or ~/.env
+export FASTVIDEO_PROMPT_CEREBRAS_MODEL="gpt-oss-120b"
+export TORCHINDUCTOR_CACHE_DIR="$HOME/.cache/torchinductor"
+export TRITON_CACHE_DIR="$HOME/.triton/cache"
+export TORCH_CUDA_ARCH_LIST="10.0a"
+
+# Compiler env (needed for flashinfer JIT compilation at server startup)
+export CXX="$CONDA_PREFIX/compiler_compat/g++"
+export CC="$CONDA_PREFIX/compiler_compat/gcc"
+export CUDAHOSTCXX="$GXX"
+export NVCC_PREPEND_FLAGS="-ccbin $GXX -allow-unsupported-compiler"
+
+# Link against libcudart.so.12 at JIT compile time; stubs for libcuda.so
+# cuda-compat has libcudart.so -> libcudart.so.12 (linker needs unversioned name)
+export LIBRARY_PATH="$CONDA_PREFIX/lib/cuda-compat:$CONDA_PREFIX/lib/stubs"
+
+# Only libcudart.so.12 at runtime — prevents cuDNN from seeing .so.13
+export LD_LIBRARY_PATH="$CUDA_RT_DIR"
+
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
+export FASTVIDEO_GPU_COUNT="${FASTVIDEO_GPU_COUNT:-all}"
+# Default SP size to the usable GPU count so single-GPU invocations work.
+# A numeric FASTVIDEO_GPU_COUNT caps the pool below the visible count, and an
+# SP size above the pool size fails GPUPool startup with "Not enough GPUs".
+IFS=',' read -ra _VISIBLE_GPUS <<< "$CUDA_VISIBLE_DEVICES"
+# Count only non-empty tokens, matching gpu_pool.get_available_gpus (e.g. ",0,1" is 2 GPUs).
+_USABLE_GPU_COUNT=0
+for _gpu in "${_VISIBLE_GPUS[@]}"; do
+    [[ -n "${_gpu//[[:space:]]/}" ]] && _USABLE_GPU_COUNT=$((_USABLE_GPU_COUNT + 1))
+done
+if [[ "$FASTVIDEO_GPU_COUNT" =~ ^[0-9]+$ ]] && (( FASTVIDEO_GPU_COUNT < _USABLE_GPU_COUNT )); then
+    _USABLE_GPU_COUNT="$FASTVIDEO_GPU_COUNT"
+fi
+export DREAMVERSE_SP_SIZE="${DREAMVERSE_SP_SIZE:-$_USABLE_GPU_COUNT}"
+PORT="${DREAMVERSE_PORT:-8009}"
+
+FFMPEG_ENV="$(dirname "$0")/ffmpeg-env.sh"
+# shellcheck source=ffmpeg-env.sh
+[[ -f "$FFMPEG_ENV" ]] && source "$FFMPEG_ENV"
+
+echo "==> Launching dreamverse-server on GPU $CUDA_VISIBLE_DEVICES port $PORT"
+CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" \
+    "$CONDA_PREFIX/bin/dreamverse-server" --host 0.0.0.0 --port "$PORT"

apps/dreamverse/scripts/launch-frontend.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# launch-frontend.sh — start the Dreamverse Next.js dev server and ngrok tunnel.
+#
+# Usage (from repo root):
+#   bash apps/dreamverse/scripts/launch-frontend.sh
+#
+# Override backend or ngrok URL via env:
+#   BACKEND_HOST=1.2.3.4 bash apps/dreamverse/scripts/launch-frontend.sh
+
+set -euo pipefail
+
+CONDA_PREFIX="$HOME/miniconda3/envs/dreamverse"
+BACKEND_HOST="${BACKEND_HOST:-10.244.18.228}"
+BACKEND_PORT="${BACKEND_PORT:-8009}"
+NGROK_URL="${NGROK_URL:-ltx23.ngrok.app}"
+WEB_DIR="$(git rev-parse --show-toplevel)/apps/dreamverse/web"
+
+cleanup() {
+    echo "==> Shutting down..."
+    kill "$FRONTEND_PID" 2>/dev/null || true
+}
+trap cleanup EXIT
+
+echo "==> Starting frontend (backend: $BACKEND_HOST:$BACKEND_PORT)"
+BACKEND_HOST="$BACKEND_HOST" BACKEND_PORT="$BACKEND_PORT" \
+    npm run --prefix "$WEB_DIR" dev &
+FRONTEND_PID=$!
+
+echo "==> Starting ngrok tunnel -> $NGROK_URL"
+"$CONDA_PREFIX/bin/ngrok" http --url="$NGROK_URL" 5299

apps/dreamverse/scripts/setup-dreamverse-env.sh

@@ -0,0 +1,105 @@
+#!/usr/bin/env bash
+# setup-dreamverse-env.sh — create and configure the dreamverse conda env
+# from scratch on this aarch64 NFS Slurm cluster.
+#
+# Run from the login node (from the repo root):
+#   bash apps/dreamverse/scripts/setup-dreamverse-env.sh
+#
+# After this script completes, use launch-dreamverse.sh on a compute node.
+
+set -euo pipefail
+
+REPO_ROOT="$(git rev-parse --show-toplevel)"
+ENV_NAME="dreamverse"
+LOCAL_DIR="/mnt/local/hal-kevin"       # cache/pkgs — keep on local disk
+CONDA_PREFIX="$HOME/miniconda3/envs/$ENV_NAME"  # env — on shared NFS so it survives node changes
+
+echo "==> Removing existing env if present"
+conda env remove -p "$CONDA_PREFIX" -y 2>/dev/null || true
+rm -rf "$CONDA_PREFIX" 2>/dev/null || true
+
+echo "==> Creating conda env at $CONDA_PREFIX"
+CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda create -p "$CONDA_PREFIX" python=3.11 -y
+
+GXX="$CONDA_PREFIX/bin/aarch64-conda-linux-gnu-g++"
+GCC="$CONDA_PREFIX/bin/aarch64-conda-linux-gnu-gcc"
+
+echo "==> Installing compiler"
+CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda install -p "$CONDA_PREFIX" gxx_linux-aarch64 -y
+
+echo "==> Installing CUDA toolkit (nvcc + headers)"
+CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda install -p "$CONDA_PREFIX" -c nvidia cuda-toolkit -y
+
+echo "==> Hiding conflicting libcudart.so.13 immediately"
+mkdir -p "$CONDA_PREFIX/lib/hidden"
+mv "$CONDA_PREFIX"/lib/libcudart.so* "$CONDA_PREFIX/lib/hidden/" 2>/dev/null || true
+
+echo "==> Fixing compiler_compat symlinks"
+mkdir -p "$CONDA_PREFIX/compiler_compat"
+ln -sf "$GXX" "$CONDA_PREFIX/compiler_compat/g++"
+ln -sf "$GCC" "$CONDA_PREFIX/compiler_compat/gcc"
+
+echo "==> Symlinking CUDA headers to standard location"
+for f in "$CONDA_PREFIX/targets/sbsa-linux/include/"*; do
+    ln -sf "$f" "$CONDA_PREFIX/include/$(basename "$f")" 2>/dev/null || true
+done
+
+echo "==> Installing ffmpeg (native build with x264 + NVENC)"
+CUDA_PREFIX="$CONDA_PREFIX" bash "$REPO_ROOT/apps/dreamverse/scripts/install_native_ffmpeg.sh"
+
+echo "==> Installing pip and uv"
+CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda install -p "$CONDA_PREFIX" pip -y
+"$CONDA_PREFIX/bin/pip" install uv
+
+echo "==> Setting compiler env vars"
+export UV_CACHE_DIR="$LOCAL_DIR/cache"
+export UV_LINK_MODE=copy
+export CXX="$CONDA_PREFIX/compiler_compat/g++"
+export CC="$CONDA_PREFIX/compiler_compat/gcc"
+export CUDAHOSTCXX="$GXX"
+export NVCC_PREPEND_FLAGS="-ccbin $GXX -allow-unsupported-compiler"
+export CUDA_HOME="$CONDA_PREFIX"
+# Only build for GB200 (sm_100a); CUDA 13 dropped support for older archs
+export TORCH_CUDA_ARCH_LIST="10.0a"
+
+echo "==> Installing torch with CUDA 12.8"
+UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install torch==2.11.0 torchvision \
+    --index-url https://download.pytorch.org/whl/cu128
+
+echo "==> Hiding any newly introduced libcudart.so.13"
+mv "$CONDA_PREFIX"/lib/libcudart.so* "$CONDA_PREFIX/lib/hidden/" 2>/dev/null || true
+
+# Set paths now that torch (and its nvidia packages) are installed
+CUDA_RT_DIR="$CONDA_PREFIX/lib/python3.11/site-packages/nvidia/cuda_runtime/lib"
+CUDA_RT_SO="$(ls "$CUDA_RT_DIR"/libcudart.so.* 2>/dev/null | head -1)"
+
+# The pip nvidia package only has libcudart.so.12 (versioned), not libcudart.so.
+# The linker needs the unversioned name to satisfy -lcudart. Create a compat dir.
+mkdir -p "$CONDA_PREFIX/lib/cuda-compat"
+ln -sf "$CUDA_RT_SO" "$CONDA_PREFIX/lib/cuda-compat/libcudart.so"
+
+export LIBRARY_PATH="$CONDA_PREFIX/lib/cuda-compat:$CONDA_PREFIX/lib/stubs"
+export CMAKE_ARGS="-DCUDA_CUDART_LIBRARY=$CUDA_RT_SO -DCUDA_INCLUDE_DIRS=$CONDA_PREFIX/targets/sbsa-linux/include"
+
+echo "==> Installing build tools"
+"$CONDA_PREFIX/bin/pip" install scikit-build-core cmake ninja
+
+echo "==> Initializing git submodules"
+cd "$REPO_ROOT"
+git submodule update --init fastvideo-kernel/include/cutlass fastvideo-kernel/include/tk
+
+echo "==> Building fastvideo-kernel from local source"
+UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install \
+    -e "./fastvideo-kernel" --no-build-isolation
+
+echo "==> Installing fastvideo + dreamverse extras"
+UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install \
+    -e ".[dreamverse]" --no-build-isolation
+
+echo "==> Installing flashinfer-python (pinned, must be last)"
+UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install \
+    https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.11.post3/flashinfer_python-0.6.11.post3-py3-none-any.whl
+
+echo ""
+echo "Done. On a compute node run (GPUs 0-3 by default; set CUDA_VISIBLE_DEVICES to restrict):"
+echo "  bash apps/dreamverse/scripts/launch-dreamverse.sh"

apps/dreamverse/serve_configs/streaming_demo.yaml

@@ -110,7 +110,7 @@ default_request:
     fps: 24                  # internal: gpu_pool.py:85 TARGET_FPS
 
 streaming:
-  # internal: config.py:33 SESSION_TIMEOUT_SECONDS = 300
+  # internal: config.py SESSION_TIMEOUT_SECONDS (env DREAMVERSE_SESSION_TIMEOUT_SECONDS, default 300)
   session_timeout_seconds: 300
   # internal: config.py:282-284 GENERATION_SEGMENT_CAP default 6
   generation_segment_cap: 6

apps/dreamverse/web/src/app/page.tsx

@@ -9,7 +9,7 @@ import SessionTimeoutModal from "@/components/SessionTimeoutModal";
 import Sidebar from "@/components/Sidebar";
 import Header from "@/components/Header";
 import VideoPlayer from "@/components/VideoPlayer";
-import Workspace from "@/components/Workspace";
+import Workspace, { SceneHistoryList } from "@/components/Workspace";
 import { saveProject, saveProjectMetadata, listProjects, loadProjectClips, deleteProject, pruneOldProjects, type StoredProject, type StoredClip } from "@/lib/projectStorage";
 import { isInfrastructureError } from "@/lib/ws/reducer";
 import { useStore } from "@/hooks/useStore";
@@ -31,7 +31,7 @@ import { applyNormalizedSocketEvent } from "@/lib/ws/reducer";
 import { createPromptWindowStore } from "@/stores/promptWindow";
 import { createRewriteStore } from "@/stores/rewrite";
 import { createSessionStore } from "@/stores/session";
-import { createStreamStore } from "@/stores/stream";
+import { createStreamStore, USER_PROMPT_SOURCES } from "@/stores/stream";
 import { createUiStore } from "@/stores/ui";
 import { Button } from "@/components/ui/button";
 
@@ -80,7 +80,7 @@ function yieldToEventLoop(): Promise<void> {
 
 const HERO_WAVE_LIGHT = ["#2A4A98", "#4878E5", "#6FA0F2", "#B0BCC8", "#E8D99E", "#D8C844", "#C2A620"];
 const HERO_WAVE_DARK = ["#143468", "#1E58B8", "#3892F0", "#80B8E8", "#B8D0EA", "#E2D498", "#DABB50"];
-const HERO_TEXT = "Direct scenes in seconds";
+const HERO_TEXT = "Direct scenes in seconds with";
 
 function HeroTagline() {
 	const ref = useRef<HTMLHeadingElement>(null);
@@ -166,6 +166,8 @@ function HeroTagline() {
 					</span>
 				</Fragment>
 			))}
+			<span data-char className="transition-[color,filter] duration-150">{" "}</span>
+			<img src="/logo.svg" alt="FastVideo" className="inline-block h-[1.1em] w-auto align-middle" />
 		</h1>
 	);
 }
@@ -237,6 +239,7 @@ export default function Page() {
 		enhancementEnabled,
 		promptExtensionError,
 		autoExtensionEnabled,
+		manualContinuationMode,
 		autoExtensionTimeoutHint,
 		loopGenerationEnabled,
 		generationPaused,
@@ -249,6 +252,8 @@ export default function Page() {
 		livePromptRewriteMode,
 		sessionExpired,
 		projectResetPending,
+		waitingForSegmentPrompt,
+		generatingNextScene,
 	} = sessionState;
 
 	const {
@@ -323,6 +328,11 @@ export default function Page() {
 	const [ttffValueMs, setTtffValueMs] = useState<number | null>(null);
 	const ttffIntervalRef = useRef<ReturnType<typeof setInterval> | null>(null);
 	const pendingInitialPromptRef = useRef("");
+	// Prompt id the opening scene is recorded under; sent as initial_rollout_prompt_id
+	// so the backend's pre-seeded opening PromptSubmission emits status updates
+	// (prompt_enhancing/prompt_ready/prompt_fallback_used) against the same id.
+	const pendingInitialPromptIdRef = useRef("");
+	const [initialImageDataUrl, setInitialImageDataUrl] = useState("");
 	const lastArchivedReplayKeyRef = useRef("");
 	const [sidebarOpen, setSidebarOpen] = useState(false);
 	const [currentThumbnail, setCurrentThumbnail] = useState<string | null>(null);
@@ -422,6 +432,11 @@ export default function Page() {
 			if (String(e?.source || "") === "user_rewrite" && typeof e?.text === "string" && e.text.trim()) {
 				return e.text.trim();
 			}
+			// Steering opening: the backend overwrites text/source with the enhanced
+			// prompt once ready, so fall back to the stable rawText record.
+			if (e?.steeringUserPrompt && typeof e?.rawText === "string" && e.rawText.trim()) {
+				return e.rawText.trim();
+			}
 		}
 		return "Untitled project";
 	}, [selectedPreset, promptEvents]);
@@ -432,11 +447,39 @@ export default function Page() {
 	const canDownloadVideo = useMemo(() => {
 		const currentActiveClip = activeClip as Record<string, any> | null;
 		if (currentActiveClip?.blob instanceof Blob) return true;
-		return (completedClips as Record<string, any>[]).some((clip) => clip?.blob instanceof Blob);
-	}, [activeClip, completedClips]);
+		if ((completedClips as Record<string, any>[]).some((clip) => clip?.blob instanceof Blob)) return true;
+		// Steering only: once playback has started the live AV pipeline holds playable segments, so
+		// the user can download the in-progress video at any time (handleDownloadVideo remuxes live
+		// segments). Auto mode keeps its original blob-gated behavior.
+		return Boolean(manualContinuationMode) && Boolean(avPlaybackStarted);
+	}, [activeClip, completedClips, avPlaybackStarted, manualContinuationMode]);
+
+	// Steering mode scene list (oldest first). Primary source is the user's own words, captured
+	// stably at submit time as `rawText` (the backend later overwrites text/source with the
+	// enhanced prompt, so we never read those). A segment with no user prompt — e.g. a preset's
+	// opening scene — falls back to promptHistory (the actual prompt that drove that segment).
+	const steeringScenes = useMemo(() => {
+		if (!manualContinuationMode) return [] as Record<string, any>[];
+		const userScenes = (promptEvents as Record<string, any>[])
+			.filter((e) => e?.steeringUserPrompt && !e?.steeringFailed && typeof e?.rawText === "string" && e.rawText.trim())
+			.slice()
+			.reverse() // oldest -> newest
+			.map((e) => ({ id: e.promptId, prompt: e.rawText as string }));
+		const scenes: Record<string, any>[] = [];
+		// Preset opening segments: curated seeds with no user prompt of their own.
+		const curatedHists = (promptHistory as Record<string, any>[])
+			.slice()
+			.reverse() // oldest first
+			.filter((h) => !USER_PROMPT_SOURCES.has(String(h?.source || "")) && typeof h?.prompt === "string" && (h.prompt as string).trim());
+		scenes.push(...curatedHists.map((h) => ({ id: h.id || "scene_open", prompt: h.prompt })));
+		scenes.push(...userScenes);
+		return scenes;
+	}, [manualContinuationMode, promptEvents, promptHistory]);
 
 	const hasEdits = useMemo(
-		() => Boolean(sessionStarted) && (promptEvents as Record<string, any>[]).some((e) => typeof e?.text === "string" && e.text.trim() && String(e?.source || "").trim() === "user_rewrite"),
+		() => Boolean(sessionStarted) && (
+			(promptEvents as Record<string, any>[]).some((e) => typeof e?.text === "string" && e.text.trim() && String(e?.source || "").trim() === "user_rewrite")
+		),
 		[sessionStarted, promptEvents],
 	);
 
@@ -1421,7 +1464,8 @@ export default function Page() {
 		if (!prompt) return;
 		lastSubmitTimeRef.current = now;
 		const rCAR = !uiStore.get().devtoolsMode && !uiStore.get().demoMode;
-		if (rCAR || (sessionStore.get().livePromptRewriteMode && !uiStore.get().demoMode)) {
+		const inManualMode = sessionStore.get().manualContinuationMode;
+		if (!inManualMode && (rCAR || (sessionStore.get().livePromptRewriteMode && !uiStore.get().demoMode))) {
 			if (rewriteStore.get().rewritingSeedPrompts) return;
 			const rewriteSourcePromptWindowPrompts = getActivePromptWindowPrompts();
 			const nextPendingClip = {
@@ -1462,6 +1506,11 @@ export default function Page() {
 			status: "submitted",
 			source: "user_raw",
 			text: prompt,
+			// Stable record of the user's own words for the steering scene list. The backend
+			// later overwrites `text`/`source` with the enhanced prompt via prompt/ready, but
+			// these two fields are never touched by trackPromptEvent.
+			steeringUserPrompt: true,
+			rawText: prompt,
 		});
 		ws.send(
 			JSON.stringify({
@@ -1475,7 +1524,14 @@ export default function Page() {
 			activeClipId: shouldUseArchivedPlaybackFallback() ? streamStore.get().activeClipId : "",
 			activePlaybackStartTime: shouldUseArchivedPlaybackFallback() ? streamStore.get().activePlaybackStartTime : 0,
 		});
-		sessionStore.patch({ livePromptDraft: "" });
+		sessionStore.patch({
+			livePromptDraft: "",
+			waitingForSegmentPrompt: false,
+			sessionNotice: "",
+			// Light the "Generating next scene" overlay immediately on a real submit;
+			// stream/media_init (or a fallback/error) clears it.
+			...(inManualMode ? { generatingNextScene: true } : {}),
+		});
 	}
 
 	function setLivePromptRewriteMode(enabled: boolean) {
@@ -1550,6 +1606,14 @@ export default function Page() {
 		);
 	}
 
+	// Steering (manual continuation) vs the automatic 6-segment rollout — a pre-session
+	// preference honored when the session starts.
+	function handleManualContinuationToggle(event: any) {
+		sessionStore.patch({
+			manualContinuationMode: Boolean(event.currentTarget.checked),
+		});
+	}
+
 	function handleLoopGenerationToggle(event: any) {
 		sessionStore.patch({
 			loopGenerationEnabled: Boolean(event.currentTarget.checked),
@@ -1725,6 +1789,9 @@ export default function Page() {
 			sessionNotice: preserveSessionNotice ? sessionStore.get().sessionNotice : "",
 			sessionExpired: preserveSessionNotice ? sessionStore.get().sessionExpired : false,
 			projectResetPending: false,
+			manualContinuationMode: true,
+			waitingForSegmentPrompt: false,
+			generatingNextScene: false,
 		});
 		rewriteStore.resetSessionState();
 		streamStore.resetSessionState();
@@ -1737,6 +1804,7 @@ export default function Page() {
 	function resetToProjectLobbyState() {
 		setVideoMuted(true);
 		pendingInitialPromptRef.current = "";
+		pendingInitialPromptIdRef.current = "";
 		sessionStore.patch({
 			sessionStarted: false,
 			livePromptDraft: "",
@@ -1750,6 +1818,9 @@ export default function Page() {
 			sessionNotice: "",
 			sessionExpired: false,
 			projectResetPending: false,
+			manualContinuationMode: true,
+			waitingForSegmentPrompt: false,
+			generatingNextScene: false,
 		});
 		rewriteStore.resetSessionState();
 		streamStore.resetSessionState();
@@ -1760,7 +1831,14 @@ export default function Page() {
 	}
 
 	function buildProjectInitPayload(type: "session_init_v2" | "project_init_v1") {
-		const segmentPrompts = getSessionInitPrompts();
+		const manualMode = Boolean(sessionStore.get().manualContinuationMode);
+		let segmentPrompts = getSessionInitPrompts();
+		// Steering mode: seed the first 2 segments from the preset so there's no
+		// gap between segment 1 and 2; the user drives every subsequent segment.
+		// Force auto/loop off so the backend waits after the seeded prompts run out.
+		if (manualMode) {
+			segmentPrompts = segmentPrompts.slice(0, 2);
+		}
 		setSeedPrompts(segmentPrompts);
 		return {
 			type,
@@ -1768,11 +1846,17 @@ export default function Page() {
 			preset_label: getInitialPresetLabel(),
 			curated_prompts: segmentPrompts,
 			initial_rollout_prompt: normalizeInitialPrompt(pendingInitialPromptRef.current),
-			initial_image: null,
+			// Ties the backend's pre-seeded opening PromptSubmission to the prompt event
+			// recorded in beginProjectLocally so its status updates land on that record.
+			initial_rollout_prompt_id: pendingInitialPromptIdRef.current,
+			initial_image: initialImageDataUrl
+				? { data_url: initialImageDataUrl, mime_type: initialImageDataUrl.split(";")[0].split(":")[1] || "image/png", name: "upload.png" }
+				: null,
 			single_clip_mode: false,
 			enhancement_enabled: sessionStore.get().enhancementEnabled,
-			auto_extension_enabled: sessionStore.get().autoExtensionEnabled,
-			loop_generation_enabled: sessionStore.get().loopGenerationEnabled,
+			auto_extension_enabled: manualMode ? false : sessionStore.get().autoExtensionEnabled,
+			loop_generation_enabled: manualMode ? false : sessionStore.get().loopGenerationEnabled,
+			manual_continuation_mode: manualMode,
 		};
 	}
 
@@ -1951,7 +2035,11 @@ export default function Page() {
 		setCurrentThumbnail(null);
 		const initialPrompt = normalizeInitialPrompt(sessionStore.get().livePromptDraft as string);
 		pendingInitialPromptRef.current = initialPrompt;
+		setInitialImageDataUrl("");
 		const rCAR = !uiStore.get().devtoolsMode && !uiStore.get().demoMode;
+		// The "Steering mode" toggle is authoritative: checked = manual per-segment steering,
+		// unchecked = automatic 6-segment rollout (default).
+		const nextManualContinuationMode = Boolean(sessionStore.get().manualContinuationMode);
 		sessionStore.patch({
 			sessionNotice: "",
 			sessionExpired: false,
@@ -1966,6 +2054,9 @@ export default function Page() {
 			autoExtensionTimeoutHint: "",
 			generationPaused: false,
 			projectResetPending: false,
+			manualContinuationMode: nextManualContinuationMode,
+			waitingForSegmentPrompt: false,
+			generatingNextScene: false,
 		});
 		resetPlaybackState();
 		streamStore.patch({
@@ -1979,12 +2070,15 @@ export default function Page() {
 			selectedHistoryId: "",
 		});
 		rewriteStore.resetSessionState();
+		pendingInitialPromptIdRef.current = initialPrompt ? makePromptId() : "";
 		if (initialPrompt) {
 			addPromptEvent({
-				promptId: makePromptId(),
+				promptId: pendingInitialPromptIdRef.current,
 				status: "rewrite_requested",
 				source: "user_rewrite",
 				text: initialPrompt,
+				// In steering mode the typed opening is the user's Scene 1 — record it stably.
+				...(nextManualContinuationMode ? { steeringUserPrompt: true, rawText: initialPrompt } : {}),
 			});
 		}
 		setSeedPrompts(getSessionInitPrompts());
@@ -2500,6 +2594,7 @@ export default function Page() {
 				selectedPresetId={selectedPresetId as string}
 				enhancementEnabled={enhancementEnabled as boolean}
 				autoExtensionEnabled={autoExtensionEnabled as boolean}
+				manualContinuationEnabled={manualContinuationMode as boolean}
 				loopGenerationEnabled={loopGenerationEnabled as boolean}
 				canJoinSession={canJoinSession as boolean}
 				canSubmitContinuation={canSubmitContinuation}
@@ -2512,6 +2607,7 @@ export default function Page() {
 				onEnhancementToggle={handleEnhancementToggle}
 				onCuratedPromptLimitChange={handleCuratedPromptLimitChange}
 				onAutoExtensionToggle={handleAutoExtensionToggle}
+				onManualContinuationToggle={handleManualContinuationToggle}
 				onLoopToggle={handleLoopGenerationToggle}
 				onJoin={joinSession}
 				onLeave={leaveSession}
@@ -2641,6 +2737,7 @@ export default function Page() {
 			/>
 			<Header timeLeft={headerTimeLeft} formatTime={formatTime} onToggleSidebar={() => setSidebarOpen((prev) => !prev)} />
 
+			<div className={cn("flex flex-1 min-h-0 flex-col", sessionStarted && "pb-16 sm:pb-28")}>
 			<div className="relative flex flex-1 min-h-0 flex-col justify-center px-4 pb-2 sm:px-6 sm:pb-12">
 				{isViewingMode && (
 					<>
@@ -2724,6 +2821,8 @@ export default function Page() {
 									showLivePlayback={showLivePlayback}
 									defaultMuted={videoMuted}
 									canDownload={canDownloadVideo}
+									waitingForSegmentPrompt={waitingForSegmentPrompt as boolean}
+									generatingNextScene={generatingNextScene as boolean}
 									onPlaying={markFirstFrameRendered}
 									onDownload={handleDownloadVideo}
 								/>
@@ -2734,6 +2833,7 @@ export default function Page() {
 					<section className={cn("mx-auto w-full max-w-2xl", hasEdits && "flex-1 min-h-0 overflow-y-auto")}>
 						<Workspace
 							promptEvents={promptEvents as any[]}
+							manualMode={manualContinuationMode as boolean}
 							currentThumbnail={currentThumbnail}
 							originalLabel={pendingInitialPromptRef.current || (selectedPreset as Record<string, any>)?.label || ""}
 							sessionStarted={sessionStarted as boolean}
@@ -2780,11 +2880,17 @@ export default function Page() {
 							isGenerating={loadingAnimation as boolean}
 							storyPresets={storyPresets as any[]}
 							continuationDraft={livePromptDraft as string}
+							manualContinuationEnabled={manualContinuationMode as boolean}
+							onModeChange={(manual) => sessionStore.patch({ manualContinuationMode: manual })}
+							initialImageDataUrl={initialImageDataUrl}
+							onImageUpload={(dataUrl) => setInitialImageDataUrl(dataUrl)}
+							onImageClear={() => setInitialImageDataUrl("")}
 							canJoinSession={canStartSession}
 							canSubmitContinuation={canSubmitContinuation}
 							sessionExpired={sessionExpired as boolean}
 							sessionNotice={sessionNotice as string}
 							projectResetPending={projectResetPending as boolean}
+							waitingForSegmentPrompt={waitingForSegmentPrompt as boolean}
 							onPresetGenerate={handlePresetGenerate}
 							onContinuationInput={handleLivePromptInput}
 							onContinuationKeydown={handleLivePromptKeydown}
@@ -2798,6 +2904,12 @@ export default function Page() {
 					</motion.div>
 				</div>
 			</div>
+			{manualContinuationMode && (
+				<div className="px-4 sm:px-6">
+					<SceneHistoryList sceneHistory={steeringScenes as any[]} />
+				</div>
+			)}
+			</div>
 		</main>
 	);
 }

apps/dreamverse/web/src/components/ChatBar.tsx

@@ -2,13 +2,16 @@
 
 import React, { useRef, useState, useCallback, useEffect } from "react";
 import Image from "next/image";
-import { Film, ArrowUp, X, Loader2, ArrowLeft } from "lucide-react";
+import { Film, ArrowUp, X, Loader2, ArrowLeft, ImagePlus } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import LeaveSessionModal, { shouldShowLeaveWarning } from "@/components/LeaveSessionModal";
 import SpeechToTextButton from "@/components/SpeechToTextButton";
 import { cn } from "@/lib/utils";
 
 const PROMPT_MAX_LENGTH = 500;
+// Must match backend session_init_image.py: MAX_SESSION_INIT_IMAGE_BYTES / SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES.
+const IMAGE_MAX_BYTES = 15 * 1024 * 1024;
+const IMAGE_ALLOWED_TYPES = ["image/png", "image/jpeg", "image/webp"];
 
 interface Props {
 	sessionStarted?: boolean;
@@ -22,6 +25,12 @@ interface Props {
 	sessionNotice?: string;
 	projectResetPending?: boolean;
 	viewingReadOnly?: boolean;
+	waitingForSegmentPrompt?: boolean;
+	manualContinuationEnabled?: boolean;
+	onModeChange?: (manual: boolean) => void;
+	initialImageDataUrl?: string;
+	onImageUpload?: (dataUrl: string, mimeType: string, name: string) => void;
+	onImageClear?: () => void;
 	onPresetGenerate?: (presetId: string) => void;
 	onContinuationInput?: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
 	onContinuationKeydown?: (e: React.KeyboardEvent<HTMLTextAreaElement>) => void;
@@ -46,6 +55,12 @@ export default function ChatBar({
 	sessionNotice = "",
 	projectResetPending = false,
 	viewingReadOnly = false,
+	waitingForSegmentPrompt = false,
+	manualContinuationEnabled = false,
+	onModeChange = () => {},
+	initialImageDataUrl = "",
+	onImageUpload = () => {},
+	onImageClear = () => {},
 	onPresetGenerate = () => {},
 	onContinuationInput = () => {},
 	onContinuationKeydown = () => {},
@@ -59,15 +74,51 @@ export default function ChatBar({
 }: Props) {
 	const [sttBusy, setSttBusy] = useState(false);
 	const [leaveModalOpen, setLeaveModalOpen] = useState(false);
+	const [imageError, setImageError] = useState("");
+	const fileInputRef = useRef<HTMLInputElement>(null);
+
+	const processImageFile = useCallback((file: File) => {
+		if (!IMAGE_ALLOWED_TYPES.includes(file.type)) {
+			setImageError("Unsupported image type. Use a PNG, JPEG, or WebP image.");
+			return;
+		}
+		if (file.size > IMAGE_MAX_BYTES) {
+			setImageError("Image is too large. The maximum size is 15MB.");
+			return;
+		}
+		const reader = new FileReader();
+		reader.onload = (e) => {
+			const dataUrl = e.target?.result as string;
+			if (dataUrl) {
+				setImageError("");
+				onImageUpload(dataUrl, file.type, file.name);
+			}
+		};
+		reader.onerror = () => {
+			setImageError("Could not read the image file. Please try again.");
+		};
+		reader.readAsDataURL(file);
+	}, [onImageUpload]);
+
+	const handleImagePaste = useCallback((e: React.ClipboardEvent) => {
+		if (sessionStarted) return;
+		const items = Array.from(e.clipboardData?.items ?? []);
+		const imageItem = items.find((item) => item.type.startsWith("image/"));
+		if (!imageItem) return;
+		const file = imageItem.getAsFile();
+		if (file) processImageFile(file);
+	}, [sessionStarted, processImageFile]);
 	const showSpinner = isGenerating || rewritingSeedPrompts;
 	const isBusy = isGenerating || rewritingSeedPrompts || projectResetPending;
 	const messagePlaceholder = projectResetPending
 		? "Starting new project\u2026"
 		: isBusy
 			? "Generating video\u2026"
-			: !sessionStarted
-				? "What video are you imagining?"
-				: "What do you want to edit?";
+			: waitingForSegmentPrompt
+				? "Describe the next scene\u2026"
+				: !sessionStarted
+					? "What video are you imagining?"
+					: "What do you want to edit?";
 	const actionLabel = !sessionStarted ? "Generate" : "Rewrite rollout";
 
 	const inputRef = useRef<HTMLTextAreaElement>(null);
@@ -267,9 +318,9 @@ export default function ChatBar({
 						<Button onClick={onStartNewProject} size="sm" className="rounded-full px-5">
 							New Project
 						</Button>
-						<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
+						<a href="https://haoailab.com/blogs/dreamverse/" target="_blank" rel="noopener noreferrer">
 							<Button variant="outline" size="sm" className="rounded-full px-5">
-								Join Waitlist
+								Blog
 							</Button>
 						</a>
 					</div>
@@ -346,6 +397,31 @@ export default function ChatBar({
 				</div>
 			)}
 
+
+			{!sessionStarted && imageError && (
+				<div className="rounded-xl border border-rose-500/20 bg-rose-500/10 px-4 py-2.5 text-center text-xs text-rose-700 dark:text-rose-300">
+					{imageError}
+				</div>
+			)}
+
+			{!sessionStarted && initialImageDataUrl && (
+				<div className="flex items-center gap-2 rounded-2xl border border-input bg-card/65 px-3 py-2">
+					<img src={initialImageDataUrl} alt="Initial frame" className="h-12 w-12 rounded-lg object-cover" />
+					<span className="flex-1 truncate text-xs text-muted-foreground">Starting image set</span>
+					<button type="button" onClick={() => { setImageError(""); onImageClear(); }} className="text-muted-foreground hover:text-foreground transition-colors">
+						<X className="size-4" />
+					</button>
+				</div>
+			)}
+
+			<input
+				ref={fileInputRef}
+				type="file"
+				accept={IMAGE_ALLOWED_TYPES.join(",")}
+				className="hidden"
+				onChange={(e) => { const f = e.target.files?.[0]; if (f) processImageFile(f); e.target.value = ""; }}
+			/>
+
 			<div
 				className={cn(
 					"flex min-w-0 items-center gap-1.5 rounded-4xl border py-2.5 pl-5 pr-2.5 shadow-md backdrop-blur-sm transition-all duration-200",
@@ -359,6 +435,7 @@ export default function ChatBar({
 					value={continuationDraft}
 					onChange={onContinuationInput}
 					onKeyDown={handleKeyDown}
+					onPaste={handleImagePaste}
 					placeholder={sttBusy ? "Listening\u2026" : messagePlaceholder}
 					maxLength={PROMPT_MAX_LENGTH}
 					disabled={isBusy || sttBusy}
@@ -369,6 +446,19 @@ export default function ChatBar({
 					)}
 				/>
 				{onSpeechTranscript && <SpeechToTextButton disabled={isBusy} onTranscript={onSpeechTranscript} onInterimChange={onSpeechInterimChange} onBusyChange={setSttBusy} />}
+				{!sessionStarted && (
+					<Button
+						type="button"
+						variant="ghost"
+						size="icon-sm"
+						title="Add starting image"
+						onClick={() => fileInputRef.current?.click()}
+						disabled={isBusy}
+						className="shrink-0 rounded-full text-muted-foreground hover:text-foreground"
+					>
+						<ImagePlus className="size-4" />
+					</Button>
+				)}
 				{!sessionStarted ? (
 					<Button
 						aria-label={actionLabel}

apps/dreamverse/web/src/components/Header.tsx

@@ -8,7 +8,6 @@ import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import { ThemeToggle } from "@/components/ui/theme-toggle";
 
-const FASTVIDEO_REPO_URL = "https://haoailab.com/blogs/dreamverse/";
 const FASTVIDEO_BLOG_URL = "https://haoailab.com/blogs/dreamverse/";
 
 interface Props {
@@ -29,13 +28,13 @@ export default function Header({ timeLeft = null, formatTime = (seconds) => `${s
 							<SidePanelOpenFilled size={20} />
 						</Button>
 					)}
-					<a href={FASTVIDEO_REPO_URL} target="_blank" rel="noopener noreferrer" title="FastVideo on GitHub">
+					<a href="/" title="FastVideo home">
 						<Image src="/logo.svg" alt="FastVideo" width={32} height={32} className="h-8 w-auto sm:h-9 transition-opacity hover:opacity-70" />
 					</a>
 					<div className="hidden sm:flex items-center gap-3">
-						<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
+						<a href={FASTVIDEO_BLOG_URL} target="_blank" rel="noopener noreferrer">
 							<Button variant="outline" size="sm" className="gap-1.5 rounded-full px-3 text-xs">
-								Join Waitlist
+								Blog
 								<ExternalLink className="size-3 opacity-60" />
 							</Button>
 						</a>
@@ -53,9 +52,9 @@ export default function Header({ timeLeft = null, formatTime = (seconds) => `${s
 			</div>
 
 			<div className="flex sm:hidden items-center gap-2 px-4 pb-3">
-				<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
+				<a href={FASTVIDEO_BLOG_URL} target="_blank" rel="noopener noreferrer">
 					<Button variant="outline" size="sm" className="gap-1.5 rounded-full px-3 text-xs">
-						Join Waitlist
+						Blog
 						<ExternalLink className="size-3 opacity-60" />
 					</Button>
 				</a>

apps/dreamverse/web/src/components/VideoPlayer.tsx

@@ -3,7 +3,7 @@
 import React, { useState, useEffect, useRef, useCallback } from "react";
 import { cn } from "@/lib/utils";
 import { PlayFilledAlt } from "@carbon/icons-react";
-import { Download, Loader2, Share } from "lucide-react";
+import { Check, ChevronDown, Download, Loader2, Share } from "lucide-react";
 import { Button } from "@/components/ui/button";
 interface VideoPlayerProps {
 	videoRef?: React.RefCallback<HTMLVideoElement>;
@@ -21,6 +21,8 @@ interface VideoPlayerProps {
 	showLivePlayback?: boolean;
 	defaultMuted?: boolean;
 	rewritePending?: boolean;
+	waitingForSegmentPrompt?: boolean;
+	generatingNextScene?: boolean;
 	onPlaying?: () => void;
 	onDownload?: () => void;
 }
@@ -57,6 +59,8 @@ export default function VideoPlayer({
 	showLivePlayback = true,
 	defaultMuted = true,
 	rewritePending = false,
+	waitingForSegmentPrompt = false,
+	generatingNextScene = false,
 	onPlaying = () => {},
 	onDownload,
 }: VideoPlayerProps) {
@@ -88,6 +92,137 @@ export default function VideoPlayer({
 		setCanShare(typeof navigator.canShare === "function" && window.matchMedia("(pointer: coarse)").matches);
 	}, []);
 
+	// Steering mode: the backend may signal "waiting for next prompt" while the current
+	// segment is still PLAYING (it generates ahead). Only surface the "Segment complete"
+	// overlay once the playhead actually reaches the end of the buffered segment.
+	const [playbackReachedEnd, setPlaybackReachedEnd] = useState(false);
+	// Steering mode: when the user submits the next scene, the segment is generated
+	// (a few seconds of latency) before frames stream. Show a "Generating next scene…"
+	// indicator across that gap so the frozen frame isn't silent. Driven by the explicit
+	// generatingNextScene state (set on scene submit / prompt selection), never inferred
+	// from waitingForSegmentPrompt edges.
+	const [generatingNext, setGeneratingNext] = useState(false);
+	// End of the buffered timeline captured the moment generation starts; the freshly
+	// generated segment extends the buffer past this, which is how we know it landed.
+	const genBoundaryRef = useRef(0);
+
+	// Steering mode: the backend may signal "waiting for next prompt" while the current
+	// segment is still PLAYING (it generates ahead). Track whether the playhead has reached
+	// the end of the buffered segment so end-overlays only show there. Keep tracking through
+	// the generating phase too, so scrubbing back and replaying to the end re-shows them.
+	useEffect(() => {
+		if (!waitingForSegmentPrompt && !generatingNext) {
+			setPlaybackReachedEnd(false);
+			return;
+		}
+		const el = liveVideoEl.current;
+		if (!el) return;
+		const check = () => {
+			try {
+				const buffered = el.buffered;
+				if (buffered.length === 0) return;
+				const end = buffered.end(buffered.length - 1);
+				// Track proximity both ways: scrubbing back off the end hides the overlay,
+				// playing forward to the end re-shows it.
+				setPlaybackReachedEnd(el.ended || end - el.currentTime <= 0.2);
+			} catch {
+				/* buffered access can throw mid-append */
+			}
+		};
+		check();
+		el.addEventListener("timeupdate", check);
+		el.addEventListener("ended", check);
+		el.addEventListener("waiting", check);
+		el.addEventListener("stalled", check);
+		el.addEventListener("pause", check);
+		el.addEventListener("seeking", check);
+		el.addEventListener("seeked", check);
+		el.addEventListener("playing", check);
+		el.addEventListener("progress", check);
+		return () => {
+			el.removeEventListener("timeupdate", check);
+			el.removeEventListener("ended", check);
+			el.removeEventListener("waiting", check);
+			el.removeEventListener("stalled", check);
+			el.removeEventListener("pause", check);
+			el.removeEventListener("seeking", check);
+			el.removeEventListener("seeked", check);
+			el.removeEventListener("playing", check);
+			el.removeEventListener("progress", check);
+		};
+	}, [waitingForSegmentPrompt, generatingNext]);
+
+	useEffect(() => {
+		if (waitingForSegmentPrompt || !sessionStarted) {
+			// Back to waiting (or session over): nothing is generating.
+			setGeneratingNext(false);
+			return;
+		}
+		if (!generatingNextScene) return;
+		// Snapshot the current end of the buffered timeline; the generated segment will
+		// extend the buffer past this boundary.
+		const el = liveVideoEl.current;
+		let boundary = el?.currentTime ?? 0;
+		try {
+			const b = el?.buffered;
+			if (b && b.length) boundary = Math.max(boundary, b.end(b.length - 1));
+		} catch {
+			/* buffered access can throw mid-append */
+		}
+		genBoundaryRef.current = boundary;
+		setGeneratingNext(true);
+	}, [generatingNextScene, waitingForSegmentPrompt, sessionStarted]);
+
+	useEffect(() => {
+		if (!generatingNext) return;
+		if (!sessionStarted) {
+			setGeneratingNext(false);
+			return;
+		}
+		const el = liveVideoEl.current;
+		if (!el) return;
+		// Clear the instant the freshly generated segment lands: the buffer grows past the
+		// boundary captured at generation start (or the playhead advances into the new
+		// frames). Deliberately NOT a bare "playing" handler — scrubbing back and replaying
+		// the EXISTING segment must keep "Generating" up until the new frames actually arrive.
+		const check = () => {
+			try {
+				const b = el.buffered;
+				const end = b.length ? b.end(b.length - 1) : 0;
+				if (end > genBoundaryRef.current + 0.25 || el.currentTime > genBoundaryRef.current + 0.1) {
+					setGeneratingNext(false);
+				}
+			} catch {
+				/* buffered access can throw mid-append */
+			}
+		};
+		check();
+		el.addEventListener("progress", check);
+		el.addEventListener("timeupdate", check);
+		el.addEventListener("durationchange", check);
+		return () => {
+			el.removeEventListener("progress", check);
+			el.removeEventListener("timeupdate", check);
+			el.removeEventListener("durationchange", check);
+		};
+	}, [generatingNext, sessionStarted]);
+
+	// Drive a ~10.5s progress bar during generation so the wait has a visible ETA.
+	const GEN_DURATION_MS = 10500;
+	const [genProgress, setGenProgress] = useState(0);
+	useEffect(() => {
+		if (!generatingNext) {
+			setGenProgress(0);
+			return;
+		}
+		const start = performance.now();
+		setGenProgress(0);
+		const id = setInterval(() => {
+			setGenProgress(Math.min((performance.now() - start) / GEN_DURATION_MS, 1));
+		}, 50);
+		return () => clearInterval(id);
+	}, [generatingNext]);
+
 	return (
 		<div className="mx-auto w-full max-w-3xl mb-2 sm:mb-6">
 			<div className="rounded-2xl border border-border bg-card/50 p-2 shadow-lg backdrop-blur-md">
@@ -104,7 +239,7 @@ export default function VideoPlayer({
 								<PlayFilledAlt className="size-10 text-white/25" />
 								<p className="text-sm text-white/50">Your video will appear here</p>
 							</div>
-						) : !avPlaybackStarted && !mediaAppendError && !inQueue && loadingAnimation ? (
+						) : !avPlaybackStarted && !mediaAppendError && !inQueue && !waitingForSegmentPrompt && loadingAnimation ? (
 							<div className="absolute inset-0 flex flex-col items-center justify-center gap-4 bg-slate-900/60 p-4 backdrop-blur-[2px]">
 								<div className="pointer-events-none absolute inset-0 overflow-hidden">
 									<div className="absolute inset-0 -translate-x-full animate-[shimmer_3s_ease-in-out_infinite] bg-gradient-to-r from-transparent via-white/[0.04] to-transparent" />
@@ -114,6 +249,35 @@ export default function VideoPlayer({
 							</div>
 						) : null}
 
+						{/* Steering mode: this segment finished — wait gracefully for the user's next scene
+							instead of spinning. The last frame stays visible behind a soft bottom gradient. */}
+						{sessionStarted && waitingForSegmentPrompt && playbackReachedEnd && !mediaAppendError && !inQueue && (
+							<div className="pointer-events-none absolute inset-0 flex flex-col items-center justify-end gap-2 bg-gradient-to-t from-slate-950/85 via-slate-950/15 to-transparent p-5 pb-6 text-center">
+								<div className="flex size-9 items-center justify-center rounded-full border border-white/25 bg-white/10 shadow-lg backdrop-blur-md">
+									<Check className="size-4 text-white/90" />
+								</div>
+								<div className="space-y-0.5">
+									<p className="text-sm font-medium text-white/95">Segment complete</p>
+									<p className="text-xs text-white/65">Describe the next scene below to keep going</p>
+								</div>
+								<ChevronDown className="size-4 animate-bounce text-white/45" />
+							</div>
+						)}
+
+						{/* Steering mode: generating the next segment — show a ~10.5s progress bar so the wait has an ETA.
+							Gated on playbackReachedEnd like "Segment complete": scrubbing back hides it, playing to the end re-shows it. */}
+						{sessionStarted && generatingNext && playbackReachedEnd && !mediaAppendError && !inQueue && (
+							<div className="pointer-events-none absolute inset-0 flex flex-col items-center justify-end gap-3 bg-gradient-to-t from-slate-950/85 via-slate-950/15 to-transparent p-5 pb-7 text-center">
+								<p className="text-sm font-medium text-white/95">Generating next scene&hellip;</p>
+								<div className="h-1.5 w-48 overflow-hidden rounded-full bg-white/15 shadow-sm">
+									<div
+										className="h-full rounded-full bg-white/85 transition-[width] duration-100 ease-linear"
+										style={{ width: `${Math.round(genProgress * 100)}%` }}
+									/>
+								</div>
+							</div>
+						)}
+
 						{rewritePending && avPlaybackStarted && (
 							<div className="absolute inset-x-0 bottom-0 z-10 flex items-center justify-center gap-2 bg-gradient-to-t from-black/60 to-transparent px-4 pb-12 pt-8 pointer-events-none">
 								<Loader2 className="size-4 animate-spin text-white/90" />

apps/dreamverse/web/src/components/Workspace.tsx

@@ -3,11 +3,90 @@ import React, { useRef, useMemo, useEffect, useCallback, useState } from "react"
 import { motion, useAnimationControls } from "framer-motion";
 import { Badge } from "@/components/ui/badge";
 import { cn } from "@/lib/utils";
-import { Check, Lightbulb, Pencil } from "lucide-react";
+import { Check, Clapperboard, Lightbulb, Pencil } from "lucide-react";
 
 export const WORKSPACE_ORIGINAL_SELECTION_KEY = "original";
 export const WORKSPACE_CURRENT_SELECTION_KEY = "current";
 
+export function SceneHistoryList({ sceneHistory = [] }: { sceneHistory?: Record<string, any>[] }) {
+	const bottomSentinelRef = useRef<HTMLDivElement>(null);
+	const topSentinelRef = useRef<HTMLDivElement>(null);
+	const [showTopFade, setShowTopFade] = useState(false);
+
+	const scenes = useMemo(
+		() => (sceneHistory || []).filter((s) => normalizeText(s?.prompt)),
+		[sceneHistory],
+	);
+
+	const scrollToBottom = useCallback(() => {
+		setTimeout(() => {
+			bottomSentinelRef.current?.scrollIntoView({ block: "end", behavior: "smooth" });
+		}, 60);
+	}, []);
+
+	useEffect(() => {
+		if (scenes.length > 0) scrollToBottom();
+	}, [scenes.length, scrollToBottom]);
+
+	useEffect(() => {
+		const sentinel = bottomSentinelRef.current;
+		if (!sentinel || typeof ResizeObserver === "undefined") return;
+		let container: HTMLElement | null = sentinel.parentElement;
+		while (container) {
+			const oy = getComputedStyle(container).overflowY;
+			if (oy === "auto" || oy === "scroll") break;
+			container = container.parentElement;
+		}
+		if (!container) return;
+		const ro = new ResizeObserver(() => {
+			const nearBottom = container!.scrollHeight - container!.scrollTop - container!.clientHeight < 96;
+			if (nearBottom) scrollToBottom();
+		});
+		ro.observe(container);
+		return () => ro.disconnect();
+	}, [scenes.length, scrollToBottom]);
+
+	useEffect(() => {
+		const el = topSentinelRef.current;
+		if (!el) return;
+		const observer = new IntersectionObserver(([entry]) => setShowTopFade(!entry.isIntersecting), { threshold: 0.1 });
+		observer.observe(el);
+		return () => observer.disconnect();
+	}, [scenes.length]);
+
+	if (scenes.length === 0) return null;
+
+	return (
+		<section className="relative z-10 flex flex-col mx-auto w-full max-w-2xl max-h-32 overflow-y-auto">
+			<div
+				className={cn(
+					"pointer-events-none sticky top-0 z-20 -mb-12 h-12 bg-linear-to-b from-background to-transparent transition-opacity duration-200",
+					showTopFade ? "opacity-100" : "opacity-0",
+				)}
+				aria-hidden="true"
+			/>
+			<div ref={topSentinelRef} className="h-0 w-0" aria-hidden="true" />
+			<div className="flex flex-col gap-2 pt-12 pb-4">
+				{scenes.map((scene, index) => (
+					<div
+						key={scene.id || index}
+						className="flex items-start gap-3 rounded-xl p-3 transition-colors duration-200 hover:bg-slate-200/50 hover:dark:bg-slate-800/30"
+					>
+						<div className="flex min-w-0 flex-1 flex-col gap-2">
+							<Badge variant="secondary" className="horizontal gap-2 items-center w-fit">
+								<Clapperboard className="size-3 opacity-70" />
+								{`Scene ${index + 1}`}
+							</Badge>
+							<p className="line-clamp-2 text-sm leading-5 text-muted-foreground">{scene.prompt}</p>
+						</div>
+					</div>
+				))}
+			</div>
+			<div ref={bottomSentinelRef} className="h-0 w-0" aria-hidden="true" />
+		</section>
+	);
+}
+
 interface WorkspaceProps {
 	promptEvents?: Record<string, any>[];
 	currentThumbnail?: string | null;
@@ -19,6 +98,7 @@ interface WorkspaceProps {
 	selectedClipId?: string;
 	selectedEntryKey?: string;
 	originalClipId?: string;
+	manualMode?: boolean;
 }
 
 function normalizeText(value: any): string {
@@ -218,7 +298,7 @@ function ChromaGradient({ sessionStarted = false }: { sessionStarted?: boolean }
 	);
 }
 
-export default function Workspace({ promptEvents = [], currentThumbnail = null, originalLabel = "", sessionStarted = false, onSelectOriginal, onSelectEvent, onSelectCurrent, selectedClipId, selectedEntryKey: selectedEntryKeyProp, originalClipId = "" }: WorkspaceProps) {
+export default function Workspace({ promptEvents = [], currentThumbnail = null, originalLabel = "", sessionStarted = false, onSelectOriginal, onSelectEvent, onSelectCurrent, selectedClipId, selectedEntryKey: selectedEntryKeyProp, originalClipId = "", manualMode = false }: WorkspaceProps) {
 	const bottomSentinelRef = useRef<HTMLDivElement>(null);
 	const topSentinelRef = useRef<HTMLDivElement>(null);
 	const [showTopFade, setShowTopFade] = useState(false);
@@ -259,6 +339,14 @@ export default function Workspace({ promptEvents = [], currentThumbnail = null,
 		return () => observer.disconnect();
 	}, [conversationEvents.length]);
 
+	if (manualMode) {
+		return (
+			<div className="mt-auto flex flex-col">
+				<ChromaGradient sessionStarted={sessionStarted} />
+			</div>
+		);
+	}
+
 	return (
 		<div className="mt-auto flex flex-col">
 			<ChromaGradient sessionStarted={sessionStarted} />

apps/dreamverse/web/src/components/devtools/DevtoolsComposer.tsx

@@ -34,6 +34,7 @@ interface DevtoolsComposerProps {
   demoMode?: boolean;
   enhancementEnabled?: boolean;
   autoExtensionEnabled?: boolean;
+  manualContinuationEnabled?: boolean;
   loopGenerationEnabled?: boolean;
   curatedPromptLimit?: number;
   maxCuratedPromptCount?: number;
@@ -49,6 +50,7 @@ interface DevtoolsComposerProps {
   onEnhancementToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onCuratedPromptLimitChange?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onAutoExtensionToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
+  onManualContinuationToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onLoopToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onLivePromptModeToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onSpeechTranscript?: (text: string) => void;
@@ -71,6 +73,7 @@ export default function DevtoolsComposer({
   demoMode = false,
   enhancementEnabled = true,
   autoExtensionEnabled = false,
+  manualContinuationEnabled = false,
   loopGenerationEnabled = false,
   curatedPromptLimit = 0,
   maxCuratedPromptCount = 0,
@@ -86,6 +89,7 @@ export default function DevtoolsComposer({
   onEnhancementToggle = () => {},
   onCuratedPromptLimitChange = () => {},
   onAutoExtensionToggle = () => {},
+  onManualContinuationToggle = () => {},
   onLoopToggle = () => {},
   onLivePromptModeToggle = () => {},
   onSpeechTranscript,
@@ -328,6 +332,28 @@ export default function DevtoolsComposer({
                   </div>
                 </div>
 
+                <div className="flex items-start gap-3">
+                  <Checkbox
+                    id="devtools-steering-mode"
+                    checked={manualContinuationEnabled}
+                    onCheckedChange={(checked) =>
+                      onManualContinuationToggle({
+                        target: { checked: Boolean(checked) },
+                        currentTarget: { checked: Boolean(checked) },
+                      } as React.ChangeEvent<HTMLInputElement>)
+                    }
+                  />
+                  <div className="space-y-1">
+                    <Label htmlFor="devtools-steering-mode">
+                      Steering mode
+                    </Label>
+                    <p className="text-sm text-muted-foreground">
+                      Drive each segment manually — type the next scene to
+                      continue (vs the automatic 6-segment rollout).
+                    </p>
+                  </div>
+                </div>
+
                 <div className="flex items-start gap-3">
                   <Checkbox
                     id="devtools-loop-generation"

apps/dreamverse/web/src/components/devtools/DevtoolsShell.tsx

@@ -21,6 +21,7 @@ interface DevtoolsShellProps {
   selectedPresetId?: string;
   enhancementEnabled?: boolean;
   autoExtensionEnabled?: boolean;
+  manualContinuationEnabled?: boolean;
   loopGenerationEnabled?: boolean;
   canJoinSession?: boolean;
   canSubmitContinuation?: boolean;
@@ -34,6 +35,7 @@ interface DevtoolsShellProps {
   onEnhancementToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onCuratedPromptLimitChange?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onAutoExtensionToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
+  onManualContinuationToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onLoopToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
   onJoin?: () => void;
   onLeave?: () => void;
@@ -128,6 +130,7 @@ export default function DevtoolsShell({
   selectedPresetId = '',
   enhancementEnabled = true,
   autoExtensionEnabled = false,
+  manualContinuationEnabled = false,
   loopGenerationEnabled = false,
   canJoinSession = false,
   canSubmitContinuation = false,
@@ -141,6 +144,7 @@ export default function DevtoolsShell({
   onEnhancementToggle = () => {},
   onCuratedPromptLimitChange = () => {},
   onAutoExtensionToggle = () => {},
+  onManualContinuationToggle = () => {},
   onLoopToggle = () => {},
   onJoin = () => {},
   onLeave = () => {},
@@ -284,6 +288,7 @@ export default function DevtoolsShell({
           demoMode={demoMode}
           enhancementEnabled={enhancementEnabled}
           autoExtensionEnabled={autoExtensionEnabled}
+          manualContinuationEnabled={manualContinuationEnabled}
           loopGenerationEnabled={loopGenerationEnabled}
           curatedPromptLimit={curatedPromptLimit}
           maxCuratedPromptCount={maxCuratedPromptCount}
@@ -299,6 +304,7 @@ export default function DevtoolsShell({
           onEnhancementToggle={onEnhancementToggle}
           onCuratedPromptLimitChange={onCuratedPromptLimitChange}
           onAutoExtensionToggle={onAutoExtensionToggle}
+          onManualContinuationToggle={onManualContinuationToggle}
           onLoopToggle={onLoopToggle}
           onLivePromptModeToggle={onLivePromptModeToggle}
           onSpeechTranscript={onSpeechTranscript}

apps/dreamverse/web/src/lib/promptEvents.test.ts

@@ -57,4 +57,32 @@ describe('prependPromptEvent', () => {
     expect(next[0].promptId).toBe('new');
     expect(next.some((item: any) => item.promptId === 'p-23')).toBe(false);
   });
+
+  it('never drops steering scene events when capping', () => {
+    // 30 scenes interleaved with 30 other events — well past the cap.
+    let events: Record<string, any>[] = [];
+    for (let i = 0; i < 30; i += 1) {
+      events = prependPromptEvent(events, {
+        promptId: `scene-${i}`,
+        status: 'submitted',
+        steeringUserPrompt: true,
+        rawText: `scene ${i}`,
+      });
+      events = prependPromptEvent(events, {
+        promptId: `other-${i}`,
+        status: 'submitted',
+      });
+    }
+
+    const scenes = events.filter((e) => e.steeringUserPrompt);
+    expect(scenes).toHaveLength(30);
+    // Oldest-first scene order (and therefore numbering) is stable and complete.
+    expect(scenes[scenes.length - 1].promptId).toBe('scene-0');
+    expect(scenes[0].promptId).toBe('scene-29');
+    // Non-scene events are still capped, oldest dropped first.
+    const others = events.filter((e) => !e.steeringUserPrompt);
+    expect(others.length).toBeLessThanOrEqual(24);
+    expect(others.some((e) => e.promptId === 'other-0')).toBe(false);
+    expect(others[0].promptId).toBe('other-29');
+  });
 });

apps/dreamverse/web/src/lib/promptEvents.ts

@@ -16,5 +16,19 @@ export function prependPromptEvent(
   events: Record<string, any>[],
   event: Record<string, any>,
 ): Record<string, any>[] {
-  return [event, ...events].slice(0, MAX_PROMPT_EVENTS);
+  const next = [event, ...events];
+  if (next.length <= MAX_PROMPT_EVENTS) {
+    return next;
+  }
+  // Steering scene events (steeringUserPrompt) are exempt from the cap: the
+  // scene list is derived from them and must stay complete and stably numbered
+  // for long sessions. Only the oldest non-scene events are dropped.
+  let nonSceneKept = 0;
+  return next.filter((e) => {
+    if (e?.steeringUserPrompt) {
+      return true;
+    }
+    nonSceneKept += 1;
+    return nonSceneKept <= MAX_PROMPT_EVENTS;
+  });
 }

apps/dreamverse/web/src/lib/ws/reducer.test.ts

@@ -1,6 +1,11 @@
 import { describe, expect, it } from 'vitest';
 
-import { resolveSessionErrorMessage } from './reducer';
+import { applyNormalizedSocketEvent, resolveSessionErrorMessage } from './reducer';
+import { createSessionStore } from '../../stores/session';
+import { createRewriteStore } from '../../stores/rewrite';
+import { createStreamStore } from '../../stores/stream';
+import { createUiStore } from '../../stores/ui';
+import { createPromptWindowStore } from '../../stores/promptWindow';
 
 describe('resolveSessionErrorMessage', () => {
   it('returns a dedicated message for IP session limit errors', () => {
@@ -19,3 +24,162 @@ describe('resolveSessionErrorMessage', () => {
     })).toBe('Backend replica unavailable. Rejoin session.');
   });
 });
+
+function buildContext(overrides: Record<string, unknown> = {}) {
+  const sessionStore = createSessionStore();
+  const rewriteStore = createRewriteStore();
+  const streamStore = createStreamStore();
+  const uiStore = createUiStore();
+  const promptWindowStore = createPromptWindowStore();
+  const avPipeline = {
+    reset: () => {},
+    setStreamCompleted: () => {},
+    noteSegmentInit: () => {},
+    noteSegmentComplete: () => {},
+    maybeStartPlayback: () => {},
+    ensurePipeline: async () => {},
+  };
+  return {
+    sessionStore,
+    promptWindowStore,
+    rewriteStore,
+    streamStore,
+    uiStore,
+    avPipeline,
+    tick: async () => {},
+    defaultAvMime: 'video/mp4',
+    fixedRewriteModel: 'model',
+    parseLatencyMs: () => null,
+    formatPromptWindowEventText: () => '',
+    makePromptId: () => 'generated-id',
+    buildClipLabel: () => 'clip',
+    startSessionCountdown: () => {},
+    clearCountdownInterval: () => {},
+    resetTtffTimer: () => {},
+    startTtffTimer: () => {},
+    preserveArchivedPlaybackSelection: false,
+    finalizeStreamCompletion: async () => {},
+    ...overrides,
+  };
+}
+
+describe('steering generatingNextScene flow', () => {
+  it('sets generatingNextScene on prompt/sources_resumed in manual mode', async () => {
+    const context = buildContext();
+    await applyNormalizedSocketEvent(
+      { type: 'prompt/sources_resumed', payload: { segment_idx: 2 } },
+      context,
+    );
+    expect(context.sessionStore.get().generatingNextScene).toBe(true);
+    expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(false);
+  });
+
+  it('does NOT set generatingNextScene on session/auto_extension_updated', async () => {
+    const context = buildContext();
+    await applyNormalizedSocketEvent(
+      { type: 'session/auto_extension_updated', payload: { enabled: true } },
+      context,
+    );
+    expect(context.sessionStore.get().generatingNextScene).toBe(false);
+    expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(false);
+  });
+
+  it('clears generatingNextScene when segment media arrives', async () => {
+    const context = buildContext();
+    context.sessionStore.patch({ generatingNextScene: true });
+    await applyNormalizedSocketEvent(
+      {
+        type: 'stream/media_init',
+        payload: { segment_idx: 2, stream_id: 's', mime: 'video/mp4' },
+      },
+      context,
+    );
+    expect(context.sessionStore.get().generatingNextScene).toBe(false);
+  });
+
+  it('clears generatingNextScene and returns to waiting on prompt/sources_blocked', async () => {
+    const context = buildContext();
+    context.sessionStore.patch({ generatingNextScene: true });
+    await applyNormalizedSocketEvent(
+      { type: 'prompt/sources_blocked', payload: { segment_idx: 3 } },
+      context,
+    );
+    expect(context.sessionStore.get().generatingNextScene).toBe(false);
+    expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(true);
+  });
+
+  it('clears generatingNextScene when the opening prompt falls back', async () => {
+    const context = buildContext();
+    context.sessionStore.patch({ generatingNextScene: true });
+    await applyNormalizedSocketEvent(
+      {
+        type: 'prompt/fallback_used',
+        payload: { prompt_id: 'p1', prompt: '', source: 'user_enhancement_failed' },
+      },
+      context,
+    );
+    expect(context.sessionStore.get().generatingNextScene).toBe(false);
+    expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(true);
+  });
+});
+
+describe('opening prompt id tracking', () => {
+  it('routes prompt lifecycle updates to the frontend-recorded opening event', async () => {
+    // The frontend records the opening scene under its own prompt id and sends it
+    // as initial_rollout_prompt_id; the backend echoes it in status updates.
+    const context = buildContext();
+    context.rewriteStore.addPromptEvent({
+      promptId: 'opening-id',
+      status: 'rewrite_requested',
+      source: 'user_rewrite',
+      text: 'a castle at dawn',
+      steeringUserPrompt: true,
+      rawText: 'a castle at dawn',
+    });
+
+    await applyNormalizedSocketEvent(
+      { type: 'prompt/enhancing', payload: { prompt_id: 'opening-id' } },
+      context,
+    );
+    let opening = (context.rewriteStore.get().promptEvents as Record<string, any>[])
+      .find((e) => e.promptId === 'opening-id');
+    expect(opening?.status).toBe('enhancing');
+
+    await applyNormalizedSocketEvent(
+      {
+        type: 'prompt/fallback_used',
+        payload: { prompt_id: 'opening-id', prompt: '', source: 'user_enhancement_failed' },
+      },
+      context,
+    );
+    opening = (context.rewriteStore.get().promptEvents as Record<string, any>[])
+      .find((e) => e.promptId === 'opening-id');
+    expect(opening?.status).toBe('ready_fallback');
+    // A failed opening is dropped from the steering scene list instead of
+    // lingering as a ghost "Scene 1".
+    expect(opening?.steeringFailed).toBe(true);
+  });
+
+  it('marks a prompt-scoped session/error (e.g. safety block) as steeringFailed', async () => {
+    const context = buildContext();
+    context.rewriteStore.addPromptEvent({
+      promptId: 'blocked-id',
+      status: 'queued',
+      source: 'user_raw',
+      text: 'a blocked prompt',
+      steeringUserPrompt: true,
+      rawText: 'a blocked prompt',
+    });
+
+    await applyNormalizedSocketEvent(
+      {
+        type: 'session/error',
+        payload: { message: 'Prompt blocked by safety filter.', prompt_id: 'blocked-id' },
+      },
+      context,
+    );
+    const blocked = (context.rewriteStore.get().promptEvents as Record<string, any>[])
+      .find((e) => e.promptId === 'blocked-id');
+    expect(blocked?.steeringFailed).toBe(true);
+  });
+});

apps/dreamverse/web/src/lib/ws/reducer.ts

@@ -99,7 +99,19 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
 				status: "ready_fallback",
 				source: payload.source || "user_raw",
 				text: payload.prompt,
+				// Steering: this prompt produced no segment — drop it from the scene list.
+				steeringFailed: true,
 			});
+			// Steering recovery: enhancement failed so the backend enqueued nothing AND won't
+			// re-emit prompt_sources_blocked (its drained flag is still set). Put the user back to
+			// "describe the next scene" ourselves so the generating overlay clears and they can retry.
+			if (!uiStore.get().simpleMode && sessionStore.get().manualContinuationMode) {
+				sessionStore.patch({
+					waitingForSegmentPrompt: true,
+					generatingNextScene: false,
+					sessionNotice: "Couldn't continue from that prompt — try rephrasing the next scene.",
+				});
+			}
 			console.warn("[PromptEnhanceFallback] Prompt extension failed for this request.");
 			return;
 
@@ -243,19 +255,32 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
 		}
 
 		case "prompt/sources_blocked":
-			sessionStore.patch({
-				autoExtensionTimeoutHint: uiStore.get().simpleMode ? "" : "blocked on user input, increase prompt count for smoother experience",
-			});
+			if (sessionStore.get().manualContinuationMode) {
+				sessionStore.patch({ waitingForSegmentPrompt: true, generatingNextScene: false, autoExtensionTimeoutHint: "" });
+			} else {
+				sessionStore.patch({
+					autoExtensionTimeoutHint: uiStore.get().simpleMode ? "" : "blocked on user input, increase prompt count for smoother experience",
+				});
+			}
 			return;
 
 		case "prompt/sources_resumed":
+			sessionStore.patch({
+				autoExtensionTimeoutHint: "",
+				waitingForSegmentPrompt: false,
+				// A real prompt was just selected for the next segment; media arriving
+				// (stream/media_init) clears this again.
+				...(sessionStore.get().manualContinuationMode ? { generatingNextScene: true } : {}),
+			});
+			return;
+
 		case "session/auto_extension_updated":
-			sessionStore.patch({ autoExtensionTimeoutHint: "" });
-			if (event.type === "session/auto_extension_updated") {
-				console.log("[AutoExtensionUpdated]", {
-					enabled: sessionStore.get().autoExtensionEnabled,
-				});
-			}
+			// Deliberately does NOT touch generatingNextScene: toggling auto extension
+			// starts no generation.
+			sessionStore.patch({ autoExtensionTimeoutHint: "", waitingForSegmentPrompt: false });
+			console.log("[AutoExtensionUpdated]", {
+				enabled: sessionStore.get().autoExtensionEnabled,
+			});
 			return;
 
 		case "segment/step_complete":
@@ -277,6 +302,7 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
 				projectResetPending: false,
 				sessionExpired: true,
 				sessionNotice: "",
+				generatingNextScene: false,
 			});
 			console.log("Session timed out");
 			clearCountdownInterval();
@@ -354,6 +380,8 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
 			return;
 
 		case "stream/media_init":
+			// Segment media is arriving — the "Generating next scene" phase is over.
+			sessionStore.patch({ generatingNextScene: false });
 			streamStore.patch({
 				mediaAppendError: null,
 				loadingAnimation: streamStore.get().avPlaybackStarted ? streamStore.get().loadingAnimation : true,
@@ -471,17 +499,30 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
 			sessionStore.patch({
 				generationCapReached: false,
 				sessionNotice: "",
+				generatingNextScene: false,
 			});
 			await finalizeStreamCompletion();
 			return;
 
 		case "session/error": {
 			const errorMessage = resolveSessionErrorMessage(payload);
+			if (payload?.prompt_id) {
+				// Prompt-scoped error (e.g. safety-blocked): the prompt produced no
+				// segment, so drop it from the steering scene list.
+				rewriteStore.trackPromptEvent(payload.prompt_id, {
+					steeringFailed: true,
+				});
+			}
 			sessionStore.patch({
 				generationCapReached: false,
 				preservePlaybackOnClose: false,
 				promptExtensionError: "",
 				sessionNotice: errorMessage,
+				// Steering: a blocked/failed prompt produced no segment and the backend won't re-emit
+				// prompt_sources_blocked, so recover the "describe the next scene" state ourselves.
+				...(!uiStore.get().simpleMode && sessionStore.get().manualContinuationMode
+					? { waitingForSegmentPrompt: true, generatingNextScene: false }
+					: {}),
 			});
 			rewriteStore.patch({
 				rewritingSeedPrompts: false,

apps/dreamverse/web/src/stores/session.ts

@@ -24,6 +24,9 @@ export interface SessionState {
   livePromptRewriteMode: boolean;
   sessionExpired: boolean;
   projectResetPending: boolean;
+  manualContinuationMode: boolean;
+  waitingForSegmentPrompt: boolean;
+  generatingNextScene: boolean;
 }
 
 const DEFAULT_SESSION_STATE: SessionState = {
@@ -49,6 +52,12 @@ const DEFAULT_SESSION_STATE: SessionState = {
   livePromptRewriteMode: false,
   sessionExpired: false,
   projectResetPending: false,
+  // Steering-only product: every session drives scenes manually. There is no
+  // auto-rollout mode and no UI selector, so this stays true throughout.
+  manualContinuationMode: true,
+  waitingForSegmentPrompt: false,
+  // True from scene submit / prompt selection until the segment's media arrives.
+  generatingNextScene: false,
 };
 
 export type SessionStore = ManagedStore<SessionState> & {

apps/dreamverse/web/src/stores/stream.ts

@@ -1,5 +1,18 @@
 import { createManagedStore, type ManagedStore } from "./createManagedStore";
 
+// Prompt-history sources driven by the user's own submissions. Curated (non-user)
+// entries — e.g. a preset's opening scene — feed the steering scene list and are
+// exempt from the history cap so Scene 1 survives long sessions.
+export const USER_PROMPT_SOURCES = new Set([
+	"user_raw",
+	"user",
+	"user_enhanced",
+	"user_rewrite",
+	"user_enhancement_failed",
+]);
+
+const PROMPT_HISTORY_CAP = 120;
+
 export interface StreamState {
 	playingSeedPromptIndex: number | null;
 	generatingSeedPromptIndex: number | null;
@@ -159,10 +172,14 @@ export function createStreamStore(initialState: Partial<StreamState> = {}): Stre
 					loopIteration: typeof loopIteration === "number" ? loopIteration : null,
 				};
 
+				const nextHistory = [entry, ...state.promptHistory];
 				return {
 					...state,
 					promptHistoryCounter: nextCounter,
-					promptHistory: [entry, ...state.promptHistory].slice(0, 120),
+					promptHistory: nextHistory.length > PROMPT_HISTORY_CAP
+						? nextHistory.filter((item, index) =>
+							index < PROMPT_HISTORY_CAP || !USER_PROMPT_SOURCES.has(String(item.source || "")))
+						: nextHistory,
 					selectedHistoryId: state.selectedHistoryId || (entry.id as string),
 				};
 			});

apps/performance_dashboard/README.md

@@ -104,6 +104,11 @@ and does not override stored status.
 - `GET /api/performance/trends?days=90&run_source=scheduled_main`
 - `GET /api/performance/records?days=90&run_source=local`
 
-The current v1 grouping key is `(model_id, gpu_type)`. Baselines are computed
-from the latest five previous successful records in each group for dashboard
-context. CI gating uses only records marked `baseline_eligible=true`.
+V2 records use the same comparison cohort as CI: `workload_id`, `variant_id`,
+`benchmark_version`, `recipe_fingerprint`, `hardware_profile_id`, and
+`software_profile_id`. `model_id` and `gpu_type` remain display/filter
+metadata, so renaming either does not split history. Legacy records still group
+by `(model_id, gpu_type)`. Dashboard baselines use the latest five previous
+successful, baseline-eligible records in each group. Summary and trend filters
+match the latest display metadata after grouping, while the raw records endpoint
+continues to filter individual records.

docs/contributing/performance_benchmarks.md

@@ -12,7 +12,8 @@ It serves three audiences:
 * **Maintainers** — surfaces regressions in a Markdown summary on every
   performance build and a long-form Plotly dashboard.
 * **Local developers** — lets you run the same benchmark on your own machine,
-  then compare against the historical baseline for the same model and GPU.
+  then compare against the historical baseline for the same comparable
+  identity.
 
 ## Quick start (local)
 
@@ -81,9 +82,9 @@ fastvideo/performance/
 The HF dataset (`FastVideo/performance-tracking` by default) holds one
 normalized JSON per run. For v2 records, the rolling baseline is the median of
 the last 5 successful, baseline-eligible records in the same comparison cohort:
-`model_id`, `gpu_type`, `workload_id`, `variant_id`, `benchmark_version`,
-`recipe_fingerprint`, `hardware_profile_id`, and `software_profile_id`. PR and
-local records are visible in the dashboard but are not baseline eligible.
+`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
+`hardware_profile_id`, and `software_profile_id`. PR and local records are
+visible in the dashboard but are not baseline eligible.
 
 ## Planned Coverage
 
@@ -163,12 +164,11 @@ headroom and almost never need touching.
 
 `compare_baseline.py` loads the last 5 successful, baseline-eligible records
 for the same comparison cohort from the HF dataset, computes the median for
-each available metric, and evaluates the current run with the metric's
-rolling regression policy. For v2 records, that cohort is `model_id`,
-`gpu_type`, `workload_id`, `variant_id`, `benchmark_version`,
-`recipe_fingerprint`, `hardware_profile_id`, and `software_profile_id`. For
-latency, memory, and component times, higher values are regressions. For
-throughput, lower values are regressions.
+each available metric, and evaluates the current run with the metric's rolling
+regression policy. For v2 records, that cohort is `workload_id`, `variant_id`,
+`benchmark_version`, `recipe_fingerprint`, `hardware_profile_id`, and
+`software_profile_id`. For latency, memory, and component times, higher values
+are regressions. For throughput, lower values are regressions.
 
 A metric exceeds its rolling threshold when both of these are true:
 
@@ -188,10 +188,33 @@ slowly add up. Only scheduled-main successful records are baseline eligible.
 Local and pull-request runs can upload dashboard-visible records, but they do
 not update future gating baselines.
 
+Comparator summaries and normalized artifacts include an explicit
+`comparison_status`:
+
+| Status | Meaning | CI behavior |
+|---|---|---|
+| `PASS` | Comparable baseline exists and no gated metric regressed. Legacy records with no baseline also keep the historical initialization behavior. | Passes |
+| `REGRESSION` | The record exceeds one of its static thresholds or at least one gated metric regressed against a comparable baseline. | Fails |
+| `CALIBRATION_NEEDED` | A v2 record has no exact comparable baseline. | Passes, may upload when `PERF_UPLOAD_POLICY=pass`, never seeds a baseline |
+| `RECIPE_MISMATCH` | The same workload, variant, and benchmark version has trusted successful records under another recipe fingerprint, including records from other hardware or software profiles. | Fails |
+| `INFRA_ERROR` | The comparator cannot safely classify the record, such as a v2 record missing required identity fields. | Fails |
+
+`QUALITY_BLOCKED` is reserved for a future variant-promotion workflow and is
+not emitted by normal rolling-baseline comparison.
+
+For `RECIPE_MISMATCH`, trusted records are scheduled-main records or records
+already marked `baseline_eligible`. PR and local calibration uploads remain
+visible but do not authorize future CI-gating recipe mismatch failures.
+
 When the baseline shifts for a legitimate reason (torch upgrade, kernel
 change, etc.) and CI starts failing, use the
 [`reseed-performance-baseline`](https://github.com/hao-ai-lab/FastVideo/blob/main/.agents/skills/reseed-performance-baseline/SKILL.md)
-agent skill to advance the rolling median.
+agent skill to advance the rolling median. To approve the first baseline for
+a new v2 exact identity, follow that skill with a reviewed scheduled-main
+full-suite `CALIBRATION_NEEDED` normalized artifact. Its prepare step uses
+`fastvideo/tests/performance/seed_baseline.py`; after a separate human gate,
+the skill rechecks the current HF revision and conditionally uploads the whole
+seed batch in one commit.
 
 ## Schemas
 
@@ -207,14 +230,14 @@ configs and remain loadable. New or migrated configs should use
   "config_schema_version": 2,
   "workload_id": "wan-t2v",
   "variant_id": "1.3b-sp2",
-  "benchmark_version": 2
+  "benchmark_version": 3
 }
 ```
 
-`benchmark_id` is still required in this phase because raw artifact names,
-generated-video directories, normalized record paths, and the current rolling
-baseline comparator still depend on it. The v2 identity fields are config
-metadata that make the measured workload explicit:
+`benchmark_id` is still required because raw artifact names, generated-video
+directories, normalized record paths, and legacy storage directories depend on
+it. The v2 comparator does not use it as part of the comparison cohort. The v2
+identity fields make the measured workload explicit:
 
 | Field | Purpose |
 |---|---|
@@ -225,20 +248,26 @@ metadata that make the measured workload explicit:
 If a config declares `config_schema_version: 2`, loading fails clearly when any
 required v2 identity field is missing. If v2 identity or metadata fields are
 added without `config_schema_version: 2`, loading also fails so partial
-migrations do not silently run as v1 configs. Optional v2 metadata fields
-reserved for follow-up work, such as `metric_threshold_policy` and
-`quality_metadata`, must be JSON objects when present. (`recipe` is emitted
-by the harness and is not config-declarable.)
+migrations do not silently run as v1 configs. Optional v2 `quality_metadata`
+and the v1/v2 `regression_thresholds` policy must be JSON objects when present.
+(`recipe` is emitted by the harness and is not config-declarable.)
 
-Recipe fingerprinting, hardware/software profile IDs, exact-identity
-comparison, and dashboard cohort grouping land with this change: v2 records
-compare only within their identity cohort, and a record that opens a NEW
-cohort is marked `baseline_status: "initialized_new_cohort"` (regression
-gating starts once that cohort accumulates history). Legacy v1 configs still
+V2 records compare only within their exact identity cohort. A record that opens
+a new cohort is marked `baseline_status: "initialized_new_cohort"` and
+`comparison_status: "CALIBRATION_NEEDED"`; it remains ineligible until a
+reviewed scheduled-main artifact is seeded explicitly. Legacy v1 configs still
 run and are normalized for reporting, but their records skip rolling-baseline
 comparison entirely (`baseline_status: "skipped_missing_identity"`, never
-baseline eligible); only static thresholds gate them. Metric-specific
-threshold policies and promoted baselines remain separate follow-ups.
+baseline eligible); only static thresholds gate them. Metric-specific threshold
+policies are active. `QUALITY_BLOCKED` remains reserved for future variant
+promotion policy.
+
+The shipped Wan benchmark uses `benchmark_version: 3` because recipe schema 2
+changed the recipe fingerprint by removing the legacy `benchmark_id` display
+name. This intentionally opens a new comparison cohort: after deployment, a
+reviewed scheduled-main full-suite `CALIBRATION_NEEDED` artifact must be seeded
+once before rolling regression gating resumes for that exact identity. Static
+thresholds remain active during calibration.
 
 ### Raw record (`results/perf_*.json`)
 
@@ -250,7 +279,7 @@ Written by `test_inference_performance.py`. One file per benchmark run.
   "result_schema_version": 2,
   "workload_id": "wan-t2v",
   "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "benchmark_version": 3,
   "model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
   "device": "NVIDIA L40S",
   "num_gpus": 2,
@@ -289,12 +318,12 @@ Written by `test_inference_performance.py`. One file per benchmark run.
   "dit_time_s": 8.437,
   "vae_decode_time_s": 3.208,
   "recipe": {
-    "recipe_schema_version": 1,
+    "recipe_schema_version": 2,
     "benchmark": {
       "benchmark_id": "wan-t2v-1.3b-2gpu",
       "workload_id": "wan-t2v",
       "variant_id": "1.3b-sp2",
-      "benchmark_version": 2
+      "benchmark_version": 3
     },
     "model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
     "init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
@@ -314,6 +343,9 @@ Written by `test_inference_performance.py`. One file per benchmark run.
     "python": "3.12",
     "pytorch": "2.12",
     "cuda": "13.0",
+    "attention_backend": "FLASH_ATTN",
+    "flash_attention_4_enabled": true,
+    "container_image_version": "py3.12-cuda13.0.0",
     "packages": {
       "fastvideo_kernel": "0.3.2",
       "flashinfer": "0.2.11",
@@ -322,7 +354,12 @@ Written by `test_inference_performance.py`. One file per benchmark run.
     }
   },
   "software_profile_id": "sw-<sha256-prefix>",
-  "environment_metadata": { "env": { "IMAGE_VERSION": "py3.12-cuda13.0.0" } },
+  "environment_metadata": {
+    "env": {
+      "IMAGE_VERSION": "py3.12-cuda13.0.0",
+      "FASTVIDEO_CONTAINER_IMAGE_REF": "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0@sha256:<digest>"
+    }
+  },
   "environment_fingerprint": "env-<sha256-prefix>"
 }
 ```
@@ -338,7 +375,7 @@ result, used as the rolling-baseline source of truth.
   "result_schema_version": 2,
   "workload_id": "wan-t2v",
   "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "benchmark_version": 3,
   "timestamp": "2026-05-08T22:00:00+00:00",
   "commit_sha": "<full sha>",
   "gpu_type": "NVIDIA L40S",
@@ -367,6 +404,10 @@ result, used as the rolling-baseline source of truth.
   "build_id": "<buildkite-build-id>",
   "job_id": "<buildkite-job-id>",
   "quality_metadata": { "quality_status": "canonical" },
+  "baseline_status": "compared",
+  "comparison_status": "PASS",
+  "comparison_status_reason": "Comparable baseline found with no gated regressions",
+  "baseline_eligible": false,
   "success": true
 }
 ```
@@ -384,9 +425,12 @@ Current `perf_*.json` artifacts that lack the v2 comparison identity are
 normalized for reporting but skip rolling-baseline comparison and are not marked
 baseline eligible.
 
-New records compare only against the same `model_id`, `gpu_type`,
-`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
-`hardware_profile_id`, and `software_profile_id` cohort.
+New v2 records compare only against the same `workload_id`, `variant_id`,
+`benchmark_version`, `recipe_fingerprint`, `hardware_profile_id`, and
+`software_profile_id` cohort, independent of the legacy `model_id` directory
+and `gpu_type` display string. Historical v1 records remain readable for
+reporting, but current legacy artifacts do not perform a `(model_id, gpu_type)`
+rolling comparison or seed new rolling baselines.
 `environment_metadata` and `environment_fingerprint` are audit data and are not
 part of the comparison key.
 The recipe prompt digests describe the prompts actually measured by the
@@ -406,11 +450,16 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
 | `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `fastvideo/performance/hf_store.py` | Required for upload or private dataset reads. |
 | `PERF_RUN_SOURCE` | inferred | `compare_baseline.py`, `test_inference_performance.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
 | `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
-| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
+| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Performance pytest exit code. Measured static-threshold failures are attributed per record; otherwise a nonzero code reports an infrastructure error. |
 | `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
 | `BUILDKITE_BRANCH`, `BUILDKITE_COMMIT`, `BUILDKITE_PULL_REQUEST` | unset | `compare_baseline.py`, `test_inference_performance.py` | CI metadata stamped into records. |
 | `DASHBOARD_DAYS` | `30` | `dashboard.py` | Lookback window for the Plotly trend pages. |
 | `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `fastvideo/performance/hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |
+| `FASTVIDEO_ATTENTION_BACKEND` | `auto` | `test_inference_performance.py` | Requested attention backend included in `software_profile_id`. |
+| `FASTVIDEO_FA4` | `0` | `test_inference_performance.py` | FlashAttention-4 toggle included in `software_profile_id`. |
+| `FASTVIDEO_PERFORMANCE_PROFILE_VERSION` | unset | `test_inference_performance.py` | Optional explicit software cohort/profile version included in `software_profile_id`. |
+| `IMAGE_VERSION` | unset | `test_inference_performance.py` | CI container image/profile version included in `software_profile_id` when available. |
+| `FASTVIDEO_CONTAINER_IMAGE_REF` | unset | `pr_test.py`, `launch_l40s_job.py`, `test_inference_performance.py` | Resolved CI container image ref or digest recorded in `environment_metadata` for audit without changing `software_profile_id`. |
 | `FASTVIDEO_STAGE_LOGGING` | set by the pytest test | `test_inference_performance.py` | Enables pipeline stage timing capture for component metrics during benchmark runs. |
 
 ## CI integration
@@ -425,9 +474,11 @@ Each performance build runs pytest first. PR and direct runs only continue to
 `compare_baseline.py` when that fixed-threshold phase passes; if pytest fails,
 Markdown summaries and normalized JSON artifacts are not emitted. Scheduled
 main runs set `PERF_UPLOAD_POLICY=always`, so they still run
-`compare_baseline.py` (with `PERF_PYTEST_RC` set) after a fixed-threshold
-failure. Those failed scheduled main runs emit summaries and normalized
-records, upload records with `success=false`, and are excluded from future
+`compare_baseline.py` (with `PERF_PYTEST_RC` set) after pytest fails. Each raw
+record is checked against its own static thresholds: a measured breach reports
+`REGRESSION`, while unaffected records retain their rolling-baseline status. A
+nonzero pytest exit with no attributable static-threshold breach reports
+`INFRA_ERROR`. Failed records have `success=false` and are excluded from future
 rolling baselines. The dashboard still runs best-effort for observability.
 When the rolling-baseline phase runs, it emits:
 
@@ -487,10 +538,13 @@ When the rolling-baseline phase runs, it emits:
 2. The pytest test auto-discovers all configs — no test code needed. CI
    picks it up on the next `/test performance` run.
 
-3. The first persisted main-branch run with no HF history initializes the
-   baseline (passes automatically). Subsequent runs compare against it. Local
-   and pull-request runs with no HF history also pass, but they do not seed the
-   shared baseline.
+3. Legacy benchmarks are gated only by their static thresholds. Their current
+   records skip rolling-baseline comparison and are never baseline eligible.
+   V2 benchmarks with no exact comparable baseline report
+   `CALIBRATION_NEEDED`; the record remains visible but does not become
+   baseline eligible until a comparable scheduled-main full-suite run is
+   reviewed and seeded through the prepare, review, and conditional-upload
+   steps in the `reseed-performance-baseline` skill.
 
 4. If the benchmark targets a GPU not currently in `thresholds`, either add
    that GPU as a key or rely on the `default` block. Note that `default` is
@@ -510,8 +564,12 @@ When the rolling-baseline phase runs, it emits:
 
 ## Troubleshooting
 
-**"No baseline for ... Initializing"** — first run for this comparison cohort.
-Run will pass and (if persisting) seed the first record.
+**`CALIBRATION_NEEDED` / "No baseline found for exact comparable identity"** —
+the v2 run passes, but its normalized record remains
+`baseline_eligible=false`. Review a successful scheduled-main full-suite
+normalized artifact, then follow the `reseed-performance-baseline` skill. The
+utility only prepares a digest-protected manifest; the separately confirmed
+upload rechecks remote state and commits the batch atomically.
 
 **Persistent failure right after a torch / kernel / image upgrade** —
 genuine regression *or* baseline drift. Compare the failing normalized record

docs/design/inference_schema_parity_inventory.yaml

@@ -458,6 +458,8 @@ surfaces:
       guidance_scale: request.sampling.guidance_scale
       guidance_scale_2: request.sampling.guidance_scale_2
       guidance_rescale: request.sampling.guidance_rescale
+      use_embedded_guidance: request.sampling.use_embedded_guidance
+      true_cfg_scale: request.sampling.true_cfg_scale
       boundary_ratio: request.sampling.boundary_ratio
       sigmas: request.sampling.sigmas
       enable_teacache: request.runtime.enable_teacache

docs/distillation/anyflow.md

@@ -0,0 +1,116 @@
+# 🌊 AnyFlow Any-Step Video Distillation
+
+**AnyFlow** ([paper](https://arxiv.org/abs/2605.13724), [project page](https://nvlabs.github.io/AnyFlow/), [official code](https://github.com/NVlabs/AnyFlow), [model weights](https://huggingface.co/collections/nvidia/anyflow)) is an any-step video diffusion framework built on flow maps. A single distilled checkpoint can be evaluated at NFE ∈ {1, 2, 4, 8, 16, 32} without retraining, and quality scales **monotonically** with steps — unlike consistency-based distillation, which often degrades as NFE grows.
+
+The student network ``u_θ(x_t, t, r)`` predicts the *average velocity* from time ``t`` back to time ``r``, so one Euler step is
+
+```
+x_r = x_t - ((t - r) / N) · u_θ(x_t, t, r)
+```
+
+for any ``t > r``.
+
+## 📊 Model Overview
+
+NVIDIA publishes four checkpoints under [`nvidia/anyflow`](https://huggingface.co/collections/nvidia/anyflow):
+
+- `nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers` — bidirectional T2V, Wan2.1 1.3B base
+- `nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers` — bidirectional T2V, Wan2.1 14B base
+- `nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers` — frame-autoregressive variant, 1.3B
+- `nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers` — frame-autoregressive variant, 14B
+
+FastVideo currently supports the bidirectional T2V variants for training; the FAR variants can be loaded for inference through the diffusers integration.
+
+## ⚙️ Inference
+
+For inference, load the published checkpoint directly through diffusers; FastVideo's training-side ``WanModel`` config maps the HF AnyFlow ``delta_embedder`` weights onto its internal layout via ``param_names_mapping`` so the same checkpoint can be used as the ``init_from`` for the on-policy YAML below.
+
+## 🧠 Algorithm
+
+Training runs in two stages. Both use the dual-timestep Wan backbone — enabled by ``pipeline.dit_config.r_embedder: true`` in the YAML, which allocates a sibling ``condition_embedder.delta_embedder`` and fuses its embedding with the standard timestep embedding via either an additive or a gated mixer.
+
+### Stage 1 — Pretrain (flow-map central-difference)
+
+Method: ``AnyFlowPretrainMethod`` (``fastvideo/train/methods/distribution_matching/anyflow_pretrain.py``)
+
+For each batch, sample ``(t, r) ∈ [0, 1]`` as ``(max, min)`` of two uniform draws, then:
+
+- a ``diffusion_ratio`` fraction (default 0.5) gets ``r = t`` — recovers plain flow matching;
+- a ``consistency_ratio`` fraction (default 0.25) gets ``r = 0`` — forces consistency to clean data;
+- the remainder is free.
+
+The student forward at ``(t, r)`` is trained against the central-difference target
+
+```
+target = (eps - x_0) - (t - r) · dF/dt
+```
+
+where ``dF/dt`` is estimated from the student's own forward at ``(t ± δ, r)`` with the sample also moved along the flow trajectory by ``v_pred · (δ / N)``. Per-timestep weighting uses ``beta08`` (``w(t) = t · sqrt(1 - t)``, renormalized). A stop-gradient scale-balance keeps the non-diffusion branches' loss magnitude aligned with the diffusion branch.
+
+### Stage 2 — On-policy DMD
+
+Method: ``AnyFlowMethod`` (``fastvideo/train/methods/distribution_matching/anyflow.py``)
+
+Inherits ``DMD2Method``. The student is rolled out for ``student_sample_steps`` Euler-flow steps from pure noise; one randomly-chosen step is gradient-enabled (broadcast from rank 0 so every worker agrees), the rest run under ``torch.no_grad``. With ``use_mean_velocity: true`` (default) the rollout uses ``r = t_next`` at each step, matching AnyFlow's ``WanAnyFlowPipeline.training_rollout``.
+
+The inherited ``_dmd_loss`` (VSD with fake-score critic) consumes the rollout output and the teacher's CFG prediction. The optional pinned ``t_list_override`` lets configs reproduce the paper's hand-tuned 4-step schedule ``[999, 937, 833, 624, 0]``.
+
+## 🚀 Training Scripts
+
+### Stage 1 — pretrain
+
+```bash
+bash examples/train/run.sh \
+    examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml
+```
+
+**Key configuration** (in ``examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml``):
+
+- Global batch size: 32 (8 GPUs × 4 per-GPU)
+- Learning rate: 5e-5
+- Flow shift: 5.0
+- ``diffusion_ratio`` / ``consistency_ratio``: 0.5 / 0.25
+- ``epsilon`` (finite-difference step): 5 (absolute train-timestep units)
+- ``weight_type``: ``beta08``
+- ``fuse_guidance_scale``: 3.0
+- Training steps: 6000
+
+### Stage 2 — on-policy
+
+```bash
+bash examples/train/run.sh \
+    examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml \
+    --models.student.init_from outputs/wan2.1_anyflow_pretrain/checkpoint-final
+```
+
+(Or point ``models.student.init_from`` directly at ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers`` to bootstrap from the paper weights and skip Stage 1.)
+
+**Key configuration**:
+
+- Global batch size: 8 (8 GPUs × 1 per-GPU)
+- Learning rate: 2e-6
+- Flow shift: 5.0
+- ``student_sample_steps``: 4
+- ``t_list_override``: ``[999, 937, 833, 624, 0]``
+- ``use_mean_velocity``: ``true`` (i.e. ``r = t_next`` during rollout)
+- ``real_score_guidance_scale``: 3.0
+- ``generator_update_interval``: 5 (DMD2 alternation)
+- Training steps: 4000
+
+## 🔌 Loading published AnyFlow checkpoints
+
+The HF AnyFlow checkpoints expose ``condition_embedder.delta_embedder.*`` weights that FastVideo internally maps onto its ``condition_embedder.delta_embedder.mlp.*`` layout. This rename happens automatically through the regex in ``WanVideoArchConfig.param_names_mapping`` — no separate adapter is needed. The same regex is a no-op on plain Wan checkpoints (which don't contain any ``delta_embedder`` keys).
+
+Set the YAML's ``pipeline.dit_config.r_embedder: true`` to allocate the ``delta_embedder`` module on the FastVideo side; when initializing from a plain Wan checkpoint the delta weights are deep-copied from ``time_embedder`` (matching AnyFlow's ``setup_flowmap_model()`` behavior).
+
+## 🧭 Note on ``fuse_guidance_scale``
+
+Stage 1 optionally fuses classifier-free guidance into the training target so the resulting checkpoint can be sampled at ``guidance_scale=1.0`` (no extra forward pass at inference time). The transformation is
+
+```
+noise_pred ← (noise_pred - (1 - g) · noise_pred_uncond) / g
+```
+
+with ``g = fuse_guidance_scale``. The negative prompt embedding comes from ``WanModel``'s ``ensure_negative_conditioning()`` — i.e. the dataset's configured ``sampling_param.negative_prompt``. Setting ``fuse_guidance_scale: 1.0`` skips the extra unconditional forward entirely.
+
+The on-policy stage's ``real_score_guidance_scale`` (inherited from DMD2) follows the same parameterization conventions documented in [``dmd.md``](dmd.md#-note-on-real_score_guidance_scale).

examples/inference/basic/basic_flux_dev.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+import re
+
+DEFAULT_PROMPTS = [
+    "a photo of a cat",
+    (
+        "a cinematic photo of a red panda wearing a tiny backpack, standing on a "
+        "rainy neon-lit street at night, shallow depth of field, sharp focus, "
+        "35mm, bokeh"
+    ),
+]
+
+
+def _safe_filename(text: str, max_len: int = 100) -> str:
+    """Make a stable, filesystem-friendly filename base."""
+    s = text[:max_len].strip()
+    s = s.replace(os.sep, "_")
+    if os.altsep:
+        s = s.replace(os.altsep, "_")
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
+    s = s.strip(" .")
+    return s or "prompt"
+
+
+def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
+    """Delete prior outputs so reruns do not get _1, _2 suffixes."""
+    if not os.path.isdir(out_dir):
+        return
+
+    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
+    for fn in os.listdir(out_dir):
+        if pattern.match(fn):
+            with contextlib.suppress(FileNotFoundError):
+                os.remove(os.path.join(out_dir, fn))
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.",
+    )
+    p.add_argument(
+        "--model-path",
+        default="official_weights/FLUX.1-dev",
+        help="Local Diffusers checkpoint dir or HF repo id.",
+    )
+    p.add_argument(
+        "--out-dir",
+        "--outdir",
+        default="outputs/flux_dev/samples",
+        help="Directory for saved PNG outputs.",
+    )
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=None,
+        help="Prompt. Repeat for multiple images.",
+    )
+    p.add_argument(
+        "--backend",
+        default=None,
+        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA).",
+    )
+    p.add_argument("--seed", type=int, default=42, help="Base seed; each prompt uses seed + index.")
+    p.add_argument("--height", type=int, default=1024, help="Output height.")
+    p.add_argument("--width", type=int, default=1024, help="Output width.")
+    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
+    p.add_argument("--guidance", type=float, default=3.5, help="Guidance scale.")
+    p.add_argument("--num-gpus", type=int, default=1, help="GPU count.")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    prompts: list[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
+
+    if args.backend:
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
+
+    from fastvideo import VideoGenerator
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    init_kwargs = {
+        "num_gpus": args.num_gpus,
+        "workload_type": "t2i",
+        "sp_size": 1,
+        "tp_size": 1,
+        "dit_cpu_offload": False,
+        "dit_layerwise_offload": False,
+        "text_encoder_cpu_offload": False,
+        "vae_cpu_offload": False,
+        "image_encoder_cpu_offload": False,
+        "pin_cpu_memory": False,
+        "use_fsdp_inference": False,
+    }
+
+    generator = VideoGenerator.from_pretrained(
+        model_path=args.model_path,
+        **init_kwargs,
+    )
+    try:
+        for i, prompt in enumerate(prompts):
+            seed = args.seed + i
+            filename_base = (
+                f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
+            )
+            _remove_existing_outputs(args.out_dir, filename_base)
+            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
+            print(f"[flux] prompt_idx={i} seed={seed} output_path={output_path}")
+
+            generation_kwargs = {
+                "output_path": output_path,
+                "height": args.height,
+                "width": args.width,
+                "num_frames": 1,
+                "fps": 1,
+                "num_inference_steps": args.steps,
+                "guidance_scale": args.guidance,
+                "use_embedded_guidance": True,
+                "true_cfg_scale": 1.0,
+                "seed": seed,
+                "save_video": True,
+            }
+
+            generator.generate_video(prompt, **generation_kwargs)
+
+        print(f"[flux] done. outputs written to: {args.out_dir}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_glm_image.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run GLM-Image text-to-image generation through FastVideo.
+
+User story:
+    "I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
+    text-to-image generation command, saved as a PNG."
+"""
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
+    parser.add_argument(
+        "--model-path",
+        default="zai-org/GLM-Image",
+        help="HF id or local diffusers-format GLM-Image weights directory.",
+    )
+    parser.add_argument(
+        "--output",
+        default="image_output/landscape.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default=("A beautiful landscape photography with rolling hills, "
+                 "a winding river, and a vibrant sunset in the background. "
+                 "Warm golden light, photorealistic style."),
+        help="Text prompt.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.5)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
+    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
+
+    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
+    # pipeline class come from the model's registered defaults — don't override.
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        trust_remote_code=True,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=str(output.parent),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+        result = generator.generate(request)
+        if isinstance(result, list):
+            result = result[0]
+
+        frames = result.frames
+        if frames is not None and len(frames):
+            Image.fromarray(frames[0]).save(output)
+            print(f"Saved image to {output}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/edit_glm_image.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run GLM-Image image-to-image (edit) generation through FastVideo.
+
+User story:
+    "I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
+    want a minimal edit command (text + image -> edited image), saved as a PNG."
+
+GLM-Image is a single unified pipeline: passing a condition image switches it
+from text-to-image to the edit path (the condition enters the DiT via a KV-cache
+write pass), so the generator config is identical to `basic_glm_image.py` — the
+`inputs.pil_image` on the request is what selects the edit mode.
+"""
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
+    parser.add_argument(
+        "--model-path",
+        default="zai-org/GLM-Image",
+        help="HF id or local diffusers-format GLM-Image weights directory.",
+    )
+    parser.add_argument(
+        "--image",
+        default="assets/images/couple.jpg",
+        help="Condition image to edit.",
+    )
+    parser.add_argument(
+        "--output",
+        default="image_output/edited.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="Change the background to a snowy mountain landscape at golden hour.",
+        help="Edit instruction.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.5)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    condition = Image.open(args.image).convert("RGB")
+    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
+    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
+
+    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
+    # pipeline class come from the model's registered defaults — don't override.
+    # The pipeline is registered as t2i; passing inputs.pil_image below switches
+    # it to the edit path.
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        trust_remote_code=True,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            inputs=InputConfig(pil_image=condition),
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=str(output.parent),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+        result = generator.generate(request)
+        if isinstance(result, list):
+            result = result[0]
+
+        frames = result.frames
+        if frames is not None and len(frames):
+            Image.fromarray(frames[0]).save(output)
+            print(f"Saved image to {output}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml

@@ -0,0 +1,105 @@
+# AnyFlow on-policy DMD — Wan 2.1 T2V 1.3B.
+#
+# Stage 2 of the AnyFlow two-stage recipe. Continues from the pretrain
+# checkpoint; refines the student via DMD2 with a multi-step Euler-flow
+# rollout from pure noise. Teacher provides the real score, critic
+# learns the fake score; both inherited from DMD2Method.
+#
+# Replace <PATH_TO_PRETRAIN_CKPT> with the output of the pretrain stage,
+# or with the NVIDIA-released checkpoint
+# nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers to bootstrap directly from
+# the paper weights (the delta_embedder rename is handled by the
+# param_names_mapping in WanVideoArchConfig).
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: <PATH_TO_PRETRAIN_CKPT>
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.0
+  dmd_denoising_steps: [999, 937, 833, 624]
+  warp_denoising_step: false
+
+  # AnyFlow rollout knobs.
+  student_sample_steps: 4
+  use_mean_velocity: true
+  t_list_override: [999.0, 937.0, 833.0, 624.0, 0.0]
+  dmd_score_r_value: 0.0  # DMD scoring conditioning is at r=0 (consistency target).
+
+  # Critic optimizer (DMD2 inherited).
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+  attn_kind: vsa
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 21
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_anyflow_onpolicy
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: anyflow-wan
+    run_name: wan2.1_t2v_anyflow_onpolicy
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5.0
+  dit_config:
+    r_embedder: true
+    r_embedder_fusion: gated
+    r_embedder_gate_value: 0.25
+    r_embedder_deltatime_type: r

examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml

@@ -0,0 +1,83 @@
+# AnyFlow pretrain (flow-map central-difference) — Wan 2.1 T2V 1.3B.
+#
+# Stage 1 of the AnyFlow two-stage recipe. Trains the dual-timestep
+# u_θ(x_t, t, r) on the central-difference target so the same checkpoint
+# can be sampled at arbitrary NFE in the on-policy stage.
+#
+# Initialize from base Wan 2.1 T2V 1.3B. No teacher or critic at this
+# stage; AnyFlowPretrainMethod owns a single student + one optimizer.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod
+  diffusion_ratio: 0.5
+  consistency_ratio: 0.25
+  epsilon: 5            # finite-difference step in absolute train-timestep units
+  weight_type: beta08   # per-timestep loss weight = t * sqrt(1 - t), renormalized
+  fuse_guidance_scale: 3.0
+  # shift is taken from pipeline.flow_shift below.
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 4
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 21
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 6000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_anyflow_pretrain
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: anyflow-wan
+    run_name: wan2.1_t2v_anyflow_pretrain
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5.0
+  dit_config:
+    # Enable AnyFlow dual-timestep conditioning. The student loads from
+    # base Wan 2.1 — its checkpoint has no delta_embedder weights, so they
+    # get initialized identically to time_embedder via deep-copy in
+    # WanTimeTextImageEmbedding.__init__.
+    r_embedder: true
+    r_embedder_fusion: gated
+    r_embedder_gate_value: 0.25
+    r_embedder_deltatime_type: r

fastvideo/api/flux.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from fastvideo.api.sampling_param import SamplingParam
+
+
+@dataclass
+class FluxSamplingParam(SamplingParam):
+
+    prompt: str | None = "a photo of a cat"
+    negative_prompt: str = ""
+
+    num_videos_per_prompt: int = 1
+    seed: int = 0
+
+    num_frames: int = 1
+    height: int = 1024
+    width: int = 1024
+    fps: int = 1
+
+    num_inference_steps: int = 28
+    guidance_scale: float = 3.5
+    use_embedded_guidance: bool = True
+    true_cfg_scale: float = 1.0

fastvideo/api/sampling_param.py

@@ -90,6 +90,10 @@ class SamplingParam:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
+    # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
+    use_embedded_guidance: bool = False
+    # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
+    true_cfg_scale: float = 1.0
     guidance_rescale: float = 0.0
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
@@ -325,6 +329,18 @@ class SamplingParam:
             default=SamplingParam.guidance_rescale,
             help="Guidance rescale factor",
         )
+        parser.add_argument(
+            "--use-embedded-guidance",
+            action="store_true",
+            default=SamplingParam.use_embedded_guidance,
+            help="Use embedded guidance scale (FLUX-style) instead of classic CFG",
+        )
+        parser.add_argument(
+            "--true-cfg-scale",
+            type=float,
+            default=SamplingParam.true_cfg_scale,
+            help="True CFG scale for FLUX when > 1 (requires negative prompt encoding)",
+        )
         parser.add_argument(
             "--boundary-ratio",
             type=float,

fastvideo/api/schema.py

@@ -150,6 +150,7 @@ class SamplingConfig:
     guidance_scale_2: float | None = None
     guidance_rescale: float = 0.0
     true_cfg_scale: float | None = None
+    use_embedded_guidance: bool | None = None
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
 

fastvideo/attention/backends/flash_attn.py

@@ -1,13 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import importlib.util
 import os
-
 import torch
 import torch.nn.functional as F
 from dataclasses import dataclass
 
-from fastvideo import envs
+from fastvideo.attention.utils.flash_attn_default import (
+    fa_version,
+    flash_attn_func_compilable,
+)
+
 from fastvideo.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
@@ -17,119 +19,6 @@ from fastvideo.attention.backends.abstract import (
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
-
-# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1, mirroring the
-# kernel package's FASTVIDEO_VSA_CUTEDSL: its CuTeDSL kernels JIT-compile per
-# shape family and can fail at runtime on some arch/shape combinations, so it
-# is never auto-selected just because it is installed. Below sm90 a capability
-# gate in flash_attn_cute routes to FA2 the calls FA4 cannot serve there:
-# grad-enabled (its backward asserts sm90+) and GQA (pack_gqa fails CuTeDSL
-# JIT, observed on sm_89).
-if envs.FASTVIDEO_FA4:
-    try:
-        from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
-    except ImportError as e:
-        raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
-                           "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
-                           "or unset FASTVIDEO_FA4.") from e
-    fa_version = "4"
-else:
-    try:
-        from flash_attn_interface import flash_attn_func as flash_attn_3_func
-
-        # flash_attn 3 no longer have a different API, see following commit:
-        # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
-        flash_attn_func = flash_attn_3_func
-        fa_version = "3"
-    except ImportError:
-        from flash_attn import flash_attn_func as flash_attn_2_func
-        flash_attn_func = flash_attn_2_func
-        fa_version = "2"
-    try:
-        if importlib.util.find_spec("flash_attn.cute") is not None:
-            logger.info("flash_attn.cute (FA4) is installed but not enabled; "
-                        "set FASTVIDEO_FA4=1 to use it for inference.")
-    except ImportError:
-        pass
-
-# torch.compile traceability: the FA4/cute path (fa_version=="4") is
-# already a registered torch.library custom op, so dynamo treats it as a
-# graph node. The external FA2/FA3 `flash_attn_func` is NOT — dynamo
-# breaks the graph at the call site (observed: wanvideo.py self-attn,
-# once per layer every step), which fragments the compiled region and
-# blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom
-# op (mirrors the FP4 `flash_attn_cute` template) so it becomes an
-# opaque-but-traceable node. The kernel still runs eager inside the op
-# (correct — flash-attn must run eager); only dynamo's treatment of the
-# boundary changes, so numerics are unchanged (SSIM-gate to confirm).
-if fa_version in ("2", "3"):
-    _fa_default = flash_attn_func
-
-    # Scope: this op covers exactly the q/k/v + softmax_scale + causal
-    # call shape used by FlashAttentionImpl.forward's default branch
-    # (see `flash_attn_func_compilable(...)` call site below). The
-    # masked/no-pad and varlen / cross-attn paths use different
-    # entry points (`flash_attn_no_pad`, `flash_attn_varlen_*`) which
-    # are intentionally out of scope for this PR — wrapping them is a
-    # natural follow-up. The wrapper's signature is the contract: any
-    # extra kwarg (dropout_p, window_size, alibi_slopes, deterministic,
-    # return_attn_probs, ...) raises TypeError at the call site, so
-    # silent loss of kwargs is not a failure mode.
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_default_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_default_forward(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> torch.Tensor:
-        return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
-
-    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
-    def _flash_attn_default_forward_fake(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> torch.Tensor:
-        del softmax_scale, causal
-        # FA2/FA3 default path: [batch, seqlen_q, nheads, head_dim_v],
-        # same dtype/device as q (head dim taken from v).
-        return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
-
-    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
-        # Autograd carve-out. The custom op above registers a forward + fake
-        # kernel but NO backward (register_autograd), so it is opaque to
-        # autograd. Inference runs under no_grad / inference_mode and routes
-        # through the traceable custom op — that is the torch.compile win, and
-        # the only path this PR claims. Training backprops through attention,
-        # so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
-        # (itself an autograd.Function, so backward is correct) at the cost of a
-        # dynamo graph break on the training path — i.e. pre-PR behavior, no
-        # regression. Full autograd parity for the custom op (mirroring the FP4
-        # cute template) is a tracked follow-up.
-        if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
-            return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
-        return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
-elif fa_version == "4":
-    # FA4 path: `flash_attn_func` (from `flash_attn_cute`) goes through a
-    # registered torch.library custom op (with an FA4 backward on sm90+;
-    # grad-enabled and GQA calls below sm90 route to FA2), so a passthrough
-    # is enough — no extra registration needed.
-    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
-        return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
-else:
-    # Defensive: the probe above only ever sets fa_version to "2", "3",
-    # or "4"; an unexpected value means an import/probe regression and
-    # we want a loud error at import, not a silent NameError later.
-    raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
-                       f"'2', '3', or '4' from the import probe above.")
-
 logger.info("Using FlashAttention-%s backend", fa_version)
 
 # FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
@@ -309,9 +198,17 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: FlashAttnMetadata,
     ):
         if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
+            # Route through the *_compilable wrappers so dynamo sees one
+            # traceable node for each masked entry point (the unpad/pad
+            # bookkeeping runs eager inside the custom op). On FA2 these
+            # wrappers go through ops with full register_autograd, so
+            # training also backprops through the op (no graph break on
+            # the training path); on FA3/FA4 they carve out to the
+            # autograd.Function for grad-enabled calls — see
+            # fastvideo/attention/utils/flash_attn_no_pad.py.
             from fastvideo.attention.utils.flash_attn_no_pad import (
-                flash_attn_no_pad,
-                flash_attn_varlen_qk_no_pad,
+                flash_attn_no_pad_compilable as flash_attn_no_pad,
+                flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
             )
 
             attn_mask = attn_metadata.attn_mask

fastvideo/attention/backends/sdpa.py

@@ -15,8 +15,14 @@ class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [32, 64, 96, 128, 160, 192, 224, 256]
+    def get_supported_head_sizes() -> list[int] | None:
+        # torch.nn.functional.scaled_dot_product_attention is head-size
+        # agnostic: the math backend handles any size and the fused kernels
+        # fall back internally when they cannot. The previous list
+        # ([32, 64, ..., 256]) was copied from FlashAttentionBackend in the
+        # v1 refactor (#270) and wrongly excluded sizes such as 80 (e.g.
+        # CLIP vision encoders). None means "no restriction".
+        return None
 
     @staticmethod
     def get_name() -> str:

fastvideo/attention/utils/flash_attn_default.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+"""torch.compile-traceable wrapper for the FA2/FA3/FA4 default attention path.
+
+The FA4/cute path (`fa_version == "4"`) is already a registered
+`torch.library.custom_op` in `fastvideo.attention.utils.flash_attn_cute`, so
+dynamo treats it as a graph node. The external FA2/FA3 ``flash_attn_func`` is
+NOT — dynamo breaks the graph at the call site (observed: wanvideo.py
+self-attn, once per layer every step), which fragments the compiled region
+and blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom op
+(mirrors the FP4 `flash_attn_cute` template) so it becomes an
+opaque-but-traceable node. The kernel still runs eager inside the op
+(correct — flash-attn must run eager); only dynamo's treatment of the
+boundary changes, so numerics are unchanged (SSIM-gated).
+
+Autograd: FA2 has full ``register_autograd`` parity — the custom op's
+backward calls flash_attn's ``_flash_attn_backward`` directly, so training
+backprops *through* the op (no graph break on the training path either).
+FA3 currently keeps the no-backward + carve-out pattern from PR #1373
+because FA3's private backward signature wants validation on a real Hopper
+box (gated on Kuan-Hao's Modal FA3 setup PR). Once that lands the FA3 path
+can mirror FA2.
+
+Lives in `attention/utils/` (sibling of `flash_attn_cute.py` and
+`flash_attn_no_pad.py`) so it can be imported by any backend that wants the
+traceable FA default call without pulling in backend dispatch logic. The
+backend (`attention/backends/flash_attn.py`) just imports
+`flash_attn_func_compilable` and `fa_version` from here.
+"""
+
+import importlib.util
+
+import torch
+
+from fastvideo import envs
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Pick the same backend the rest of FastVideo picked for `flash_attn_func`
+# (FA4/cute → FA3 → FA2). Mirror the precedence used in
+# `attention/utils/flash_attn_no_pad.py` so the two probes always agree.
+#
+# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1: its CuTeDSL
+# kernels JIT-compile per shape family and can fail at runtime on some
+# arch/shape combinations, so it is never auto-selected just because it is
+# installed.
+if envs.FASTVIDEO_FA4:
+    try:
+        from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
+    except ImportError as e:
+        raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
+                           "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
+                           "or unset FASTVIDEO_FA4.") from e
+    fa_version = "4"
+else:
+    try:
+        from flash_attn_interface import flash_attn_func as flash_attn_3_func
+
+        # flash_attn 3 no longer has a different API, see following commit:
+        # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
+        flash_attn_func = flash_attn_3_func
+        fa_version = "3"
+    except ImportError:
+        from flash_attn import flash_attn_func as flash_attn_2_func
+        flash_attn_func = flash_attn_2_func
+        fa_version = "2"
+    try:
+        if importlib.util.find_spec("flash_attn.cute") is not None:
+            logger.info("flash_attn.cute (FA4) is installed but not enabled; "
+                        "set FASTVIDEO_FA4=1 to use it for inference.")
+    except ImportError:
+        pass
+
+if fa_version == "2":
+    # Scope: this op covers exactly the q/k/v + softmax_scale + causal call
+    # shape used by FlashAttentionImpl.forward's default branch (see
+    # `flash_attn_func_compilable(...)` call site in
+    # `attention/backends/flash_attn.py`). The masked/no-pad and varlen /
+    # cross-attn paths use different entry points
+    # (`flash_attn_no_pad`, `flash_attn_varlen_*`) which live in
+    # `attention/utils/flash_attn_no_pad.py`. The wrapper's signature is the
+    # contract: any extra kwarg (dropout_p, window_size, alibi_slopes,
+    # deterministic, return_attn_probs, ...) raises TypeError at the call
+    # site, so silent loss of kwargs is not a failure mode.
+    from flash_attn.flash_attn_interface import _flash_attn_backward as _fa2_backward
+    _fa_default = flash_attn_func
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_default_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_default_forward(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # `return_attn_probs=True` asks FA2 to also return softmax_lse +
+        # S_dmask. We need softmax_lse to feed the backward; S_dmask is the
+        # dropout mask (always None here since dropout_p is fixed at 0).
+        out, softmax_lse, _ = _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal, return_attn_probs=True)
+        return out, softmax_lse
+
+    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
+    def _flash_attn_default_forward_fake(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del softmax_scale, causal
+        # FA2 default path: out = [batch, seqlen_q, nheads, head_dim_v],
+        # softmax_lse = [batch, nheads, seqlen_q], fp32 regardless of q dtype.
+        b, sq, hq = q.shape[0], q.shape[1], q.shape[2]
+        out = q.new_empty(b, sq, hq, v.shape[-1])
+        lse = q.new_empty(b, hq, sq, dtype=torch.float32)
+        return out, lse
+
+    def _flash_attn_default_setup_context(ctx, inputs, output):
+        q, k, v, softmax_scale, causal = inputs
+        out, lse = output
+        ctx.save_for_backward(q, k, v, out, lse)
+        # `lse` is an auxiliary output we save to feed FA2's backward; nobody
+        # should differentiate through it. Mark it non-differentiable so
+        # autograd errors loudly if a caller wires it into a loss, rather
+        # than silently producing zero/None grads through the `del grad_lse`
+        # in our backward.
+        ctx.mark_non_differentiable(lse)
+        # FA2's *forward* substitutes `1 / sqrt(head_dim)` for `softmax_scale=None`
+        # internally; FA2's *backward* (`_flash_attn_backward`) demands a concrete
+        # float in its C++ schema and rejects None at the binding boundary. Resolve
+        # the default here so the value saved on ctx (and passed to backward) is
+        # always a real float — matches what FA2's own autograd.Function does.
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1]**-0.5
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+
+    def _flash_attn_default_backward(ctx, grad_out, grad_lse):
+        # We only differentiate `out`; softmax_lse is saved-for-backward, not
+        # a real differentiable output. (Mirrors the FP4 cute template.)
+        del grad_lse
+        q, k, v, out, lse = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        # FA2's `_flash_attn_backward` writes into dq/dk/dv in place. The
+        # extra kwargs (window_size_*, softcap, alibi_slopes, deterministic,
+        # rng_state) are pinned to the same defaults the forward wrapper
+        # uses — flash-attn==2.8.1 (the version FastVideo pins) requires
+        # all of them explicitly. `rng_state=None` is correct for our
+        # `dropout_p=0` configuration.
+        _fa2_backward(
+            grad_out,
+            q,
+            k,
+            v,
+            out,
+            lse,
+            dq,
+            dk,
+            dv,
+            dropout_p=0.0,
+            softmax_scale=ctx.softmax_scale,
+            causal=ctx.causal,
+            window_size_left=-1,
+            window_size_right=-1,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=False,
+            rng_state=None,
+        )
+        return dq, dk, dv, None, None
+
+    torch.library.register_autograd(
+        "fastvideo::_flash_attn_default_forward",
+        _flash_attn_default_backward,
+        setup_context=_flash_attn_default_setup_context,
+    )
+
+    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
+        # Backward is registered: autograd flows through the op (training
+        # path is also traceable; no carve-out needed). Public API matches
+        # `flash_attn_func` — returns just `out`; we drop the saved-for-
+        # backward `lse` here so callers see the original single-tensor
+        # contract.
+        out, _ = torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
+        return out
+elif fa_version == "3":
+    # FA3 path: same forward+fake custom op as the original PR #1373, with
+    # the autograd carve-out kept. The full backward (mirroring the FA2 leg
+    # above) wants a Hopper box for grad-check validation, which we don't
+    # have until Kuan-Hao's Modal FA3 setup PR lands. Until then this keeps
+    # inference traceable + training correct (via the original
+    # autograd.Function path + a pre-PR-style graph break on training).
+    _fa_default = flash_attn_func
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_default_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_default_forward(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> torch.Tensor:
+        return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
+
+    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
+    def _flash_attn_default_forward_fake(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> torch.Tensor:
+        del softmax_scale, causal
+        return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
+
+    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
+        # Autograd carve-out. The custom op above registers a forward + fake
+        # kernel but NO backward (register_autograd), so it is opaque to
+        # autograd. Inference runs under no_grad / inference_mode and routes
+        # through the traceable custom op — that is the torch.compile win, and
+        # the only path this PR claims. Training backprops through attention,
+        # so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
+        # (itself an autograd.Function, so backward is correct) at the cost of a
+        # dynamo graph break on the training path — i.e. pre-PR behavior, no
+        # regression. Full autograd parity for the custom op (mirroring the FP4
+        # cute template) is a tracked follow-up.
+        if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
+            return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
+        return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
+elif fa_version == "4":
+    # FA4 path: `flash_attn_func` is already a torch.library custom op
+    # (registered in `fastvideo.attention.utils.flash_attn_cute`), so a
+    # passthrough is enough — no extra registration needed.
+    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
+        return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
+else:
+    # Defensive: the probe above only ever sets fa_version to "2", "3",
+    # or "4"; an unexpected value means an import/probe regression and
+    # we want a loud error at import, not a silent NameError later.
+    raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
+                       f"'2', '3', or '4' from the import probe above.")

fastvideo/attention/utils/flash_attn_no_pad.py

@@ -24,7 +24,7 @@ from flash_attn.bert_padding import pad_input, unpad_input
 from fastvideo import envs
 
 
-def _resolve_flash_attn_varlen_func() -> Any:
+def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
     if envs.FASTVIDEO_FA4:
         # FA4 cute is explicit opt-in (see fastvideo/attention/backends/
         # flash_attn.py); with FASTVIDEO_FA4=1 an unimportable FA4 build must
@@ -39,20 +39,28 @@ def _resolve_flash_attn_varlen_func() -> Any:
                                "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
                                "or unset FASTVIDEO_FA4.") from e
 
-        return flash_attn_varlen_func_cute
+        return flash_attn_varlen_func_cute, "4"
     try:
         from flash_attn_interface import (
             flash_attn_varlen_func as flash_attn_varlen_func_interface, )
 
-        return flash_attn_varlen_func_interface
+        return flash_attn_varlen_func_interface, "3"
     except ImportError:
         from flash_attn import (
             flash_attn_varlen_func as flash_attn_varlen_func_flash, )
 
-        return flash_attn_varlen_func_flash
+        return flash_attn_varlen_func_flash, "2"
 
 
-flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
+flash_attn_varlen_func_impl, _FA_VARLEN_VERSION = _resolve_flash_attn_varlen_func()
+
+# FA2-only: the private varlen backward we register against the custom ops
+# below. FA3 / FA4 have different private signatures and validation paths
+# (Hopper / Blackwell boxes) — those legs keep the autograd carve-out
+# pattern from PR #1373 until their setup PRs land.
+if _FA_VARLEN_VERSION == "2":
+    from flash_attn.flash_attn_interface import (
+        _flash_attn_varlen_backward as _fa2_varlen_backward, )
 
 
 def flash_attn_no_pad(
@@ -191,3 +199,473 @@ def flash_attn_varlen_qk_no_pad(
         h=nheads,
     )
     return output
+
+
+# ---------------------------------------------------------------------------
+# torch.compile traceability + register_autograd parity for the masked /
+# varlen attention paths.
+#
+# Wraps the two entry points `FlashAttentionImpl.forward` calls
+# (`flash_attn_no_pad`, `flash_attn_varlen_qk_no_pad`) as
+# `torch.library.custom_op`s so dynamo sees one traceable node — the
+# internal unpad / pad bookkeeping (data-dependent `nnz` shapes) runs
+# eager inside the op, and the op's outputs are the statically-shaped
+# padded tensors. This mirrors the FA2 default-path wrapper in
+# `fastvideo/attention/backends/flash_attn.py`.
+#
+# Autograd: on FA2 we register a real backward (`register_autograd`)
+# that calls FA2's `_flash_attn_varlen_backward` on the unpadded form
+# — re-unpadding the saved padded tensors using the saved mask. The
+# `softmax_lse` from the varlen forward is naturally unpadded
+# (`[nheads, total_q]`); we pad it to `[batch, nheads, seqlen]` on
+# the way out (statically shaped) and re-unpad in backward. So
+# training backprops *through* the op (no graph break on the training
+# path either).
+#
+# FA3 / FA4 keep the autograd carve-out pattern from PR #1373: the
+# custom op has forward + fake only, and `*_compilable` falls back to
+# the original autograd.Function for grad-enabled calls. Those legs
+# are gated on Hopper-class / Blackwell-class boxes for backward
+# validation and ship as separate follow-ups.
+
+if _FA_VARLEN_VERSION == "2":
+    # ---------- masked self-attention: flash_attn_no_pad (FA2) ----------
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_no_pad_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_no_pad_forward(
+        qkv: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, s, _three, h, d = qkv.shape
+        x = rearrange(qkv, "b s three h d -> b s (three h d)")
+        x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
+        x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
+        out_unpad, lse_unpad, _ = flash_attn_varlen_qkvpacked_func(x_unpad,
+                                                                   cu_seqlens,
+                                                                   max_s,
+                                                                   dropout_p,
+                                                                   softmax_scale=softmax_scale,
+                                                                   causal=causal,
+                                                                   deterministic=deterministic,
+                                                                   return_attn_probs=True)
+        # Pad out: [nnz, h, d] -> [b, s, h, d]
+        out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), indices, b, s),
+                               "b s (h d) -> b s h d",
+                               h=h)
+        # Pad lse: FA2 varlen returns [nheads, total_q]. Transpose to [total_q,
+        # nheads], pad to [b, s, nheads], permute to [b, nheads, s] — statically
+        # shaped so register_fake matches.
+        lse_padded = pad_input(lse_unpad.t().contiguous(), indices, b, s).permute(0, 2, 1).contiguous()
+        return out_padded, lse_padded
+
+    @torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
+    def _flash_attn_no_pad_forward_fake(
+        qkv: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
+        b, s, _three, h, d = qkv.shape
+        out = qkv.new_empty(b, s, h, d)
+        lse = qkv.new_empty(b, h, s, dtype=torch.float32)
+        return out, lse
+
+    def _flash_attn_no_pad_setup_context(ctx, inputs, output):
+        qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic = inputs
+        out, lse = output
+        ctx.save_for_backward(qkv, out, lse, key_padding_mask)
+        # Auxiliary output, not differentiable — see default-path note.
+        ctx.mark_non_differentiable(lse)
+        # FA2's varlen backward requires a concrete float for softmax_scale.
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1]**-0.5  # head_dim from qkv's last dim
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+        ctx.dropout_p = dropout_p
+        ctx.deterministic = deterministic
+
+    def _flash_attn_no_pad_backward(ctx, grad_out, grad_lse):
+        # lse is saved-for-backward, not differentiated.
+        del grad_lse
+        qkv, out_padded, lse_padded, key_padding_mask = ctx.saved_tensors
+        b, s, _three, h, d = qkv.shape
+
+        # One `unpad_input` call (on qkv) gives us indices + cu_seqlens + max_s;
+        # reuse those for out / dout / lse below via direct indexing instead
+        # of redundant `unpad_input` calls (each of which would re-run
+        # `nonzero` + `cumsum` + a `.max().item()` GPU→CPU sync).
+        x = rearrange(qkv, "b s three h d -> b s (three h d)")
+        x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
+        x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
+        q_unpad, k_unpad, v_unpad = (t.contiguous() for t in x_unpad.unbind(dim=1))
+
+        # Direct-index variants reuse `indices` (computed above).
+        out_unpad = out_padded.flatten(0, 1)[indices].view(-1, h, d).contiguous()
+        dout_unpad = grad_out.flatten(0, 1)[indices].view(-1, h, d).contiguous()
+        # lse_padded [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
+        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[indices].t().contiguous()
+
+        dq_unpad = torch.empty_like(q_unpad)
+        dk_unpad = torch.empty_like(k_unpad)
+        dv_unpad = torch.empty_like(v_unpad)
+        _fa2_varlen_backward(
+            dout_unpad,
+            q_unpad,
+            k_unpad,
+            v_unpad,
+            out_unpad,
+            lse_unpad,
+            dq_unpad,
+            dk_unpad,
+            dv_unpad,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_s,
+            max_seqlen_k=max_s,
+            dropout_p=ctx.dropout_p,
+            softmax_scale=ctx.softmax_scale,
+            causal=ctx.causal,
+            window_size_left=-1,
+            window_size_right=-1,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=ctx.deterministic,
+            rng_state=None,
+        )
+
+        # Re-pad each grad and stack into dqkv.
+        def _repad(dt_unpad: torch.Tensor) -> torch.Tensor:
+            padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, b, s)
+            return rearrange(padded, "b s (h d) -> b s h d", h=h)
+
+        dqkv = torch.stack([_repad(dq_unpad), _repad(dk_unpad), _repad(dv_unpad)], dim=2)
+        # 6 inputs total: qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic.
+        return dqkv, None, None, None, None, None
+
+    torch.library.register_autograd(
+        "fastvideo::_flash_attn_no_pad_forward",
+        _flash_attn_no_pad_backward,
+        setup_context=_flash_attn_no_pad_setup_context,
+    )
+
+    # ---------- cross-attention: flash_attn_varlen_qk_no_pad (FA2) ----------
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_varlen_qk_no_pad_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_varlen_qk_no_pad_forward(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        query_padding_mask: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, sq, h, d = query.shape
+        q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
+                                                                        query_padding_mask)
+        k_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
+                                                                key_padding_mask)
+        v_unpad, _, _, _, _ = unpad_input(rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
+        q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
+        k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
+        v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
+        out_unpad, lse_unpad, _ = flash_attn_varlen_func_impl(q_unpad,
+                                                              k_unpad,
+                                                              v_unpad,
+                                                              cu_seqlens_q,
+                                                              cu_seqlens_k,
+                                                              max_seqlen_q,
+                                                              max_seqlen_k,
+                                                              dropout_p=dropout_p,
+                                                              softmax_scale=softmax_scale,
+                                                              causal=causal,
+                                                              deterministic=deterministic,
+                                                              return_attn_probs=True)
+        # Pad out: [nnz_q, h, d] -> [b, sq, h, d]
+        out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), q_indices, b, sq),
+                               "b s (h d) -> b s h d",
+                               h=h)
+        # Pad lse: [h, nnz_q] -> [b, h, sq]
+        lse_padded = pad_input(lse_unpad.t().contiguous(), q_indices, b, sq).permute(0, 2, 1).contiguous()
+        return out_padded, lse_padded
+
+    @torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
+    def _flash_attn_varlen_qk_no_pad_forward_fake(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        query_padding_mask: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del key, query_padding_mask, key_padding_mask
+        del causal, dropout_p, softmax_scale, deterministic
+        b, sq, h, _ = query.shape
+        # `out`'s head_dim comes from value (d_v), matching the real forward's
+        # out_padded ([b, sq, h, d_v]); it can differ from query's d_q.
+        out = query.new_empty(b, sq, h, value.shape[-1])
+        lse = query.new_empty(b, h, sq, dtype=torch.float32)
+        return out, lse
+
+    def _flash_attn_varlen_qk_no_pad_setup_context(ctx, inputs, output):
+        (query, key, value, query_padding_mask, key_padding_mask, causal, dropout_p, softmax_scale,
+         deterministic) = inputs
+        out, lse = output
+        ctx.save_for_backward(query, key, value, out, lse, query_padding_mask, key_padding_mask)
+        # Auxiliary output, not differentiable — see default-path note.
+        ctx.mark_non_differentiable(lse)
+        if softmax_scale is None:
+            softmax_scale = query.shape[-1]**-0.5
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+        ctx.dropout_p = dropout_p
+        ctx.deterministic = deterministic
+
+    def _flash_attn_varlen_qk_no_pad_backward(ctx, grad_out, grad_lse):
+        del grad_lse
+        (query, key, value, out_padded, lse_padded, query_padding_mask, key_padding_mask) = ctx.saved_tensors
+        b, sq, h, d = query.shape
+        sk = key.shape[1]
+
+        # One `unpad_input` call per distinct mask; reuse the returned
+        # indices via direct indexing for everything else that shares
+        # the same mask (v with k_mask; out/dout/lse with q_mask; the
+        # final repad of dk/dv also reuses k_indices). Avoids ~4
+        # redundant `unpad_input` calls + their GPU→CPU `.max().item()`
+        # syncs.
+        q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
+                                                                        query_padding_mask)
+        k_unpad, k_indices, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
+                                                                        key_padding_mask)
+        q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
+        k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
+        v_unpad = value.flatten(0, 1)[k_indices].view(-1, h, d).contiguous()
+
+        # out / dout / lse follow q's shape, so index with q_indices.
+        out_unpad = out_padded.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
+        dout_unpad = grad_out.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
+        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[q_indices].t().contiguous()
+
+        dq_unpad = torch.empty_like(q_unpad)
+        dk_unpad = torch.empty_like(k_unpad)
+        dv_unpad = torch.empty_like(v_unpad)
+        _fa2_varlen_backward(
+            dout_unpad,
+            q_unpad,
+            k_unpad,
+            v_unpad,
+            out_unpad,
+            lse_unpad,
+            dq_unpad,
+            dk_unpad,
+            dv_unpad,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            dropout_p=ctx.dropout_p,
+            softmax_scale=ctx.softmax_scale,
+            causal=ctx.causal,
+            window_size_left=-1,
+            window_size_right=-1,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=ctx.deterministic,
+            rng_state=None,
+        )
+
+        # k_indices is already available from the unpad_input above —
+        # no need to recompute it for the dk/dv repad.
+        def _repad(dt_unpad: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
+            padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, batch, seqlen)
+            return rearrange(padded, "b s (h d) -> b s h d", h=h)
+
+        dq_padded = _repad(dq_unpad, q_indices, b, sq)
+        dk_padded = _repad(dk_unpad, k_indices, b, sk)
+        dv_padded = _repad(dv_unpad, k_indices, b, sk)
+        # 9 inputs total: query, key, value, q_mask, k_mask, causal, dropout_p,
+        # softmax_scale, deterministic.
+        return dq_padded, dk_padded, dv_padded, None, None, None, None, None, None
+
+    torch.library.register_autograd(
+        "fastvideo::_flash_attn_varlen_qk_no_pad_forward",
+        _flash_attn_varlen_qk_no_pad_backward,
+        setup_context=_flash_attn_varlen_qk_no_pad_setup_context,
+    )
+
+    # ---------- public dispatchers (FA2: autograd flows through the op) -----
+
+
+    def flash_attn_no_pad_compilable(qkv,
+                                     key_padding_mask,
+                                     causal=False,
+                                     dropout_p=0.0,
+                                     softmax_scale=None,
+                                     deterministic=False):
+        """dynamo-traceable wrapper around ``flash_attn_no_pad`` (registered op,
+        full register_autograd on FA2 — both inference and training go through
+        the op, no graph break on either)."""
+        out, _ = torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
+                                                                deterministic)
+        return out
+
+    def flash_attn_varlen_qk_no_pad_compilable(query,
+                                               key,
+                                               value,
+                                               query_padding_mask,
+                                               key_padding_mask,
+                                               causal=False,
+                                               dropout_p=0.0,
+                                               softmax_scale=None,
+                                               deterministic=False):
+        """dynamo-traceable wrapper around ``flash_attn_varlen_qk_no_pad`` (registered
+        op, full register_autograd on FA2)."""
+        out, _ = torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
+                                                                          key_padding_mask, causal, dropout_p,
+                                                                          softmax_scale, deterministic)
+        return out
+
+else:
+    # ---------- FA3 / FA4: carve-out (forward+fake only, no real backward) ---
+    # Same pattern as the parked varlen-extension and the FA3 default leg in
+    # `fastvideo/attention/backends/flash_attn.py`. Real backward for these
+    # versions is a follow-up gated on Hopper / Blackwell box validation.
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_no_pad_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_no_pad_forward(
+        qkv: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        return flash_attn_no_pad(  # type: ignore[no-untyped-call]
+            qkv,
+            key_padding_mask,
+            causal=causal,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            deterministic=deterministic)
+
+    @torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
+    def _flash_attn_no_pad_forward_fake(
+        qkv: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
+        b, s, _three, h, d = qkv.shape
+        return qkv.new_empty(b, s, h, d)
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_varlen_qk_no_pad_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_varlen_qk_no_pad_forward(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        query_padding_mask: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        return flash_attn_varlen_qk_no_pad(  # type: ignore[no-untyped-call]
+            query,
+            key,
+            value,
+            query_padding_mask=query_padding_mask,
+            key_padding_mask=key_padding_mask,
+            causal=causal,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            deterministic=deterministic)
+
+    @torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
+    def _flash_attn_varlen_qk_no_pad_forward_fake(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        query_padding_mask: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+        causal: bool,
+        dropout_p: float,
+        softmax_scale: float | None,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        del key, query_padding_mask, key_padding_mask
+        del causal, dropout_p, softmax_scale, deterministic
+        b, sq, h, _ = query.shape
+        # `out`'s head_dim comes from value (d_v), matching the real forward's
+        # output ([b, sq, h, d_v]); it can differ from query's d_q.
+        return query.new_empty(b, sq, h, value.shape[-1])
+
+    def flash_attn_no_pad_compilable(qkv,
+                                     key_padding_mask,
+                                     causal=False,
+                                     dropout_p=0.0,
+                                     softmax_scale=None,
+                                     deterministic=False):
+        if torch.is_grad_enabled() and qkv.requires_grad:
+            return flash_attn_no_pad(qkv,
+                                     key_padding_mask,
+                                     causal=causal,
+                                     dropout_p=dropout_p,
+                                     softmax_scale=softmax_scale,
+                                     deterministic=deterministic)
+        return torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
+                                                              deterministic)
+
+    def flash_attn_varlen_qk_no_pad_compilable(query,
+                                               key,
+                                               value,
+                                               query_padding_mask,
+                                               key_padding_mask,
+                                               causal=False,
+                                               dropout_p=0.0,
+                                               softmax_scale=None,
+                                               deterministic=False):
+        if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
+            return flash_attn_varlen_qk_no_pad(query,
+                                               key,
+                                               value,
+                                               query_padding_mask=query_padding_mask,
+                                               key_padding_mask=key_padding_mask,
+                                               causal=causal,
+                                               dropout_p=dropout_p,
+                                               softmax_scale=softmax_scale,
+                                               deterministic=deterministic)
+        return torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
+                                                                        key_padding_mask, causal, dropout_p,
+                                                                        softmax_scale, deterministic)

fastvideo/configs/models/dits/__init__.py

@@ -1,7 +1,9 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
+from fastvideo.configs.models.dits.flux import FluxDiTConfig
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
@@ -15,6 +17,7 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
-    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
-    "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
+    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
+    "StableAudioConfig", "GlmImageDiTConfig"
 ]

fastvideo/configs/models/dits/flux.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class FluxTransformer2DArchConfig(DiTArchConfig):
+
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int | None = None
+    num_layers: int = 19
+    num_single_layers: int = 38
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096
+    pooled_projection_dim: int = 768
+    guidance_embeds: bool = True
+    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
+
+
+@dataclass
+class FluxDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=FluxTransformer2DArchConfig)
+    prefix: str = "flux"

fastvideo/configs/models/dits/glm_image.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class GlmImageDiTArchConfig(DiTArchConfig):
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    hidden_size: int = 4096
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    num_layers: int = 30
+
+    text_embed_dim: int = 1472
+    time_embed_dim: int = 512
+    condition_dim: int = 256
+
+    prior_vq_quantizer_codebook_size: int = 16384
+
+    patch_size: int = 2
+
+    max_height: int = 2048
+    max_width: int = 2048
+
+    qk_norm: str = "layer_norm"
+    eps: float = 1e-5
+
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["image_projector", "glyph_projector", "prior_token_embedding"])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^glyph_projector\.net\.0\.proj\.(.*)$": r"glyph_projector.fc_in.\1",
+            r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1",
+            r"^prior_projector\.net\.0\.proj\.(.*)$": r"prior_projector.fc_in.\1",
+            r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class GlmImageDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig)
+    prefix: str = "GlmImage"

fastvideo/configs/models/dits/wanvideo.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
+from typing import Literal
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -19,6 +20,11 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
             r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
             r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            # AnyFlow dual-timestep checkpoints expose delta_embedder weights with the
+            # same internal layout as time_embedder. The regex is harmless on plain
+            # Wan checkpoints (no delta_embedder keys to match).
+            r"^condition_embedder\.delta_embedder\.linear_1\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.delta_embedder\.linear_2\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_out.\1",
             r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
             r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
             r"condition_embedder.image_embedder.ff.fc_in.\1",
@@ -86,6 +92,14 @@ class WanVideoArchConfig(DiTArchConfig):
     # "relativistic" keeps long rollouts in-distribution; a no-op unless sink_size > 0 and local_attn_size > 0.
     rope_cache_policy: str = "absolute"
 
+    # AnyFlow dual-timestep conditioning. Defaults preserve bit-identity with
+    # the legacy single-timestep forward (no delta_embedder allocated, no
+    # extra computation on the embedder forward path).
+    r_embedder: bool = False
+    r_embedder_fusion: Literal["additive", "gated"] = "additive"
+    r_embedder_gate_value: float = 0.25
+    r_embedder_deltatime_type: Literal["r", "t-r"] = "r"
+
     def __post_init__(self):
         super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels

fastvideo/configs/models/vaes/__init__.py

@@ -2,6 +2,7 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
 from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
 from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
+from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
@@ -21,4 +22,5 @@ __all__ = [
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
     "Flux2VAEConfig",
+    "GlmImageVAEConfig",
 ]

fastvideo/configs/models/vaes/glm_image.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.autoencoder_kl import (AutoencoderKLArchConfig, AutoencoderKLVAEConfig)
+
+_GLM_IMAGE_LATENTS_MEAN: tuple[float, ...] = (
+    -0.2080078125,
+    1.875,
+    -0.470703125,
+    -1.265625,
+    -1.421875,
+    0.77734375,
+    -0.3671875,
+    -0.9453125,
+    0.318359375,
+    0.7734375,
+    -0.1884765625,
+    -0.022216796875,
+    -0.220703125,
+    -1.59375,
+    -0.81640625,
+    -0.255859375,
+)
+_GLM_IMAGE_LATENTS_STD: tuple[float, ...] = (
+    3.0625,
+    2.203125,
+    2.265625,
+    4.84375,
+    2.5,
+    3.9375,
+    2.203125,
+    3.03125,
+    2.1875,
+    2.046875,
+    2.71875,
+    2.390625,
+    2.390625,
+    2.453125,
+    2.25,
+    2.15625,
+)
+
+
+@dataclass
+class GlmImageVAEArchConfig(AutoencoderKLArchConfig):
+    act_fn: str = "silu"
+    block_out_channels: tuple[int, ...] = (128, 512, 1024, 1024)
+    down_block_types: tuple[str, ...] = (
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+    )
+    up_block_types: tuple[str, ...] = (
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+    )
+    force_upcast: bool = True
+    in_channels: int = 3
+    out_channels: int = 3
+    latent_channels: int = 16
+    latents_mean: tuple[float, ...] = _GLM_IMAGE_LATENTS_MEAN
+    latents_std: tuple[float, ...] = _GLM_IMAGE_LATENTS_STD
+    layers_per_block: int = 3
+    mid_block_add_attention: bool = False
+    norm_num_groups: int = 32
+    sample_size: int = 1024
+    scaling_factor: float = 0.18215
+    shift_factor: float | None = None
+    use_quant_conv: bool = False
+    use_post_quant_conv: bool = False
+
+    temporal_compression_ratio: int = 1
+    spatial_compression_ratio: int = 8
+
+
+@dataclass
+class GlmImageVAEConfig(AutoencoderKLVAEConfig):
+    arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig)
+
+    use_tiling: bool = True
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    tile_sample_min_height: int = 512
+    tile_sample_min_width: int = 512
+    tile_sample_stride_height: int = 384
+    tile_sample_stride_width: int = 384
+
+    load_encoder: bool = True
+    load_decoder: bool = True

fastvideo/configs/pipelines/flux.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import EncoderConfig
+from fastvideo.configs.models.dits.flux import FluxDiTConfig
+from fastvideo.configs.models.encoders import (
+    BaseEncoderOutput,
+    CLIPTextConfig,
+    T5LargeConfig,
+)
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+
+def _flux_clip_pooled_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    """CLIP branch for FLUX: Diffusers uses pooled prompt embeddings only."""
+    if outputs.pooler_output is None:
+        raise RuntimeError(
+            "FLUX CLIP conditioning requires pooler_output. Ensure the CLIP text encoder returns pooled features.")
+    return outputs.pooler_output
+
+
+def _flux_t5_sequence_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.last_hidden_state is None:
+        raise RuntimeError("FLUX T5 conditioning requires last_hidden_state.")
+    return outputs.last_hidden_state
+
+
+@dataclass
+class FluxPipelineConfig(PipelineConfig):
+    """Pipeline layout for Diffusers FLUX.1-dev (CLIP + T5 + packed DiT + FlowMatch)."""
+
+    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
+    transformer_arch: str = "FluxTransformer2DModel"
+    vae_arch: str = "AutoencoderKL"
+    text_encoder_archs: tuple[str, ...] = ("CLIPTextModel", "T5EncoderModel")
+    tokenizer_archs: tuple[str, ...] = ("CLIPTokenizer", "T5TokenizerFast")
+
+    dit_config: FluxDiTConfig = field(default_factory=FluxDiTConfig)
+    vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
+
+    embedded_cfg_scale: float = 3.5
+    flow_shift: float | None = None
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (CLIPTextConfig(), T5LargeConfig()))
+    preprocess_text_funcs: tuple[Callable[[str], str],
+                                 ...] = field(default_factory=lambda: (preprocess_text, preprocess_text))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
+        default_factory=lambda: (_flux_clip_pooled_postprocess, _flux_t5_sequence_postprocess))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", "bf16"))
+
+    def __post_init__(self) -> None:
+        te_cfgs = list(self.text_encoder_configs)
+        if len(te_cfgs) >= 1:
+            te_cfgs[0].tokenizer_kwargs.setdefault("padding", "max_length")
+            te_cfgs[0].tokenizer_kwargs.setdefault("max_length", 77)
+            te_cfgs[0].tokenizer_kwargs.setdefault("truncation", True)
+            te_cfgs[0].tokenizer_kwargs.setdefault("return_tensors", "pt")
+        if len(te_cfgs) >= 2:
+            cap = 512
+            te_cfgs[1].tokenizer_kwargs["max_length"] = min(int(te_cfgs[1].tokenizer_kwargs.get("max_length", cap)),
+                                                            cap)
+            te_cfgs[1].tokenizer_kwargs.setdefault("padding", "max_length")
+            te_cfgs[1].tokenizer_kwargs.setdefault("truncation", True)
+            te_cfgs[1].tokenizer_kwargs.setdefault("return_tensors", "pt")

fastvideo/configs/pipelines/glm_image.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+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.glm_image import GlmImageDiTConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
+from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+def glm_image_t5_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+
+    assert torch.isnan(hidden_state).sum() == 0, "T5 hidden states contain NaN"
+
+    max_len = 512
+    prompt_embeds = [u[:min(v, max_len)] for u, v in zip(hidden_state, seq_lens, strict=True)]
+    prompt_embeds_tensor: torch.Tensor = torch.stack(
+        [torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
+
+    return prompt_embeds_tensor
+
+
+@dataclass
+class GlmImageConfig(PipelineConfig):
+
+    dit_config: DiTConfig = field(default_factory=GlmImageDiTConfig)
+    dit_precision: str = "bf16"
+
+    vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig)
+    vae_precision: str = "fp32"
+    vae_tiling: bool = True
+    vae_sp: bool = False
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (glm_image_t5_postprocess, ))
+
+    flow_shift: float | None = 1.0
+    embedded_cfg_scale: float = 7.5

fastvideo/entrypoints/streaming/server.py

@@ -191,10 +191,13 @@ def run_server(serve_config: ServeConfig, *, generator: _GeneratorProto | None =
 
         generator = VideoGenerator.from_pretrained(config=serve_config.generator)
     app = build_app(serve_config, generator)
+    # A 15MB init image is ~20MB as a base64 ws message, above uvicorn's
+    # default 16MiB frame cap.
     uvicorn.run(
         app,
         host=serve_config.server.host,
         port=serve_config.server.port,
+        ws_max_size=32 * 1024 * 1024,
     )
 
 

fastvideo/layers/rotary_embedding.py

@@ -295,6 +295,7 @@ def get_1d_rotary_pos_embed(
     interpolation_factor: float = 1.0,
     dtype: torch.dtype = torch.float32,
     use_real: bool = True,
+    freqs_dtype: torch.dtype | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Precompute the frequency tensor for complex exponential (cis) with given dimensions.
@@ -319,12 +320,16 @@ def get_1d_rotary_pos_embed(
     if isinstance(pos, int):
         pos = torch.arange(pos).float()
 
+    # freqs_dtype is an alias for dtype (Diffusers-compatible calling convention).
+    if freqs_dtype is not None:
+        dtype = freqs_dtype
+
     # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
     # has some connection to NTK literature
     if theta_rescale_factor != 1.0:
         theta *= theta_rescale_factor**(dim / (dim - 2))
 
-    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(dtype) / dim))  # [D/2]
+    freqs = 1.0 / (theta**(torch.arange(0, dim, 2, device=pos.device)[:(dim // 2)].to(dtype) / dim))  # [D/2]
     freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
     freqs_cos = freqs.cos()  # [S, D/2]
     freqs_sin = freqs.sin()  # [S, D/2]
@@ -445,6 +450,21 @@ def get_nd_rotary_pos_embed(
     return cos, sin
 
 
+_ROTARY_POS_EMBED_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+# Bound the table cache so long-running servers / causal models (which vary
+# start_frame per frame) cannot grow it without limit; entries are large float64
+# tensors. Least-recently-used eviction keeps the active resolution(s) hot while
+# capping memory.
+_ROTARY_POS_EMBED_CACHE_MAXSIZE = 16
+
+
+def _hashable(value: Any) -> Any:
+    """Return a hashable view of a scalar or sequence for use in a cache key."""
+    if isinstance(value, list | tuple):
+        return tuple(value)
+    return value
+
+
 def get_rotary_pos_embed(
     rope_sizes,
     hidden_size,
@@ -495,6 +515,31 @@ def get_rotary_pos_embed(
         sp_rank = 0
         sp_world_size = 1
 
+    # Memoize on every output-affecting argument; the table is constant across
+    # denoising steps, so this avoids recomputing the float64 cos/sin tables.
+    cache_key = (
+        _hashable(rope_sizes),
+        tuple(rope_dim_list),
+        rope_theta,
+        _hashable(theta_rescale_factor),
+        _hashable(interpolation_factor),
+        shard_dim,
+        sp_rank,
+        sp_world_size,
+        dtype,
+        start_frame,
+        use_real,
+    )
+    cached = _ROTARY_POS_EMBED_CACHE.get(cache_key)
+    if cached is not None:
+        # Move to most-recently-used position so the active table is not evicted
+        # when several resolutions / buckets share the process (LRU recency).
+        # Pop with a default: a concurrent eviction between the get() above and
+        # here would otherwise raise KeyError on the hit path.
+        if _ROTARY_POS_EMBED_CACHE.pop(cache_key, None) is not None:
+            _ROTARY_POS_EMBED_CACHE[cache_key] = cached
+        return cached
+
     freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
         rope_dim_list,
         rope_sizes,
@@ -508,6 +553,14 @@ def get_rotary_pos_embed(
         start_frame=start_frame,
         use_real=use_real,
     )
+    # The returned tensors are shared cache entries: callers must never mutate
+    # them in place. Note .to(device) is an identity alias when the tensor is
+    # already on the target device (e.g. CPU runs), so it does NOT guarantee a
+    # copy — treat the tables as read-only and copy before any in-place op.
+    # Reached only on a miss, so evict the least-recently-used entry at capacity.
+    if len(_ROTARY_POS_EMBED_CACHE) >= _ROTARY_POS_EMBED_CACHE_MAXSIZE:
+        _ROTARY_POS_EMBED_CACHE.pop(next(iter(_ROTARY_POS_EMBED_CACHE)))
+    _ROTARY_POS_EMBED_CACHE[cache_key] = (freqs_cos, freqs_sin)
     return freqs_cos, freqs_sin
 
 

fastvideo/models/dits/flux.py

@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.layers.rotary_embedding import apply_rotary_emb, get_1d_rotary_pos_embed
+
+from fastvideo.attention import DistributedAttention
+from fastvideo.configs.models import DiTConfig
+from fastvideo.forward_context import get_forward_context, set_forward_context
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.sd3 import (
+    CombinedTimestepTextProjEmbeddings,
+    SD3AdaLayerNormContinuous,
+    SD3AdaLayerNormZero,
+    SD3FeedForward,
+    SD3TextProjection,
+    SD3TimestepEmbedding,
+)
+from fastvideo.platforms import AttentionBackendEnum
+
+
+@dataclass
+class FluxTransformer2DModelOutput:
+    sample: torch.Tensor
+
+
+class FluxPosEmbed(nn.Module):
+    """1D RoPE axes concatenated per Diffusers `FluxPosEmbed`."""
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        n_axes = ids.shape[-1]
+        cos_out: list[torch.Tensor] = []
+        sin_out: list[torch.Tensor] = []
+        pos = ids.float()
+        is_mps = ids.device.type == "mps"
+        is_npu = ids.device.type == "npu"
+        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                use_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None:
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256,
+            flip_sin_to_cos=True,
+            downscale_freq_shift=0,
+        )
+        self.timestep_embedder = SD3TimestepEmbedding(
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            act_fn="silu",
+        )
+        self.guidance_embedder = SD3TimestepEmbedding(
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            act_fn="silu",
+        )
+        self.text_embedder = SD3TextProjection(
+            pooled_projection_dim,
+            embedding_dim,
+            act_fn="silu",
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor,
+        pooled_projection: torch.Tensor,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
+        guidance_proj = self.time_proj(guidance)
+        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
+        time_guidance_emb = timesteps_emb + guidance_emb
+        pooled_projections = self.text_embedder(pooled_projection)
+        return time_guidance_emb + pooled_projections
+
+
+class FluxAdaLayerNormZeroSingle(nn.Module):
+    def __init__(self, embedding_dim: int, bias: bool = True) -> None:
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(embedding_dim, 3 * embedding_dim, bias=bias)
+        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        emb, _ = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa
+
+
+class FluxJointAttention(nn.Module):
+    """Joint attention: text tokens precede image tokens (Diffusers order)."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_added_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_added_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_q_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_k_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_v_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(0.0),
+            ]
+        )
+        self.to_add_out = ReplicatedLinear(self.inner_dim, dim, bias=True)
+
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = hidden_states.shape[0]
+        text_seq_len = encoder_hidden_states.shape[1]
+        img_seq_len = hidden_states.shape[1]
+
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+        q = q.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        k = k.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        v = v.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+
+        enc_q, _ = self.add_q_proj(encoder_hidden_states)
+        enc_k, _ = self.add_k_proj(encoder_hidden_states)
+        enc_v, _ = self.add_v_proj(encoder_hidden_states)
+        enc_q = enc_q.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_k = enc_k.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_v = enc_v.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_q = self.norm_added_q(enc_q)
+        enc_k = self.norm_added_k(enc_k)
+
+        q = torch.cat([enc_q, q], dim=1)
+        k = torch.cat([enc_k, k], dim=1)
+        v = torch.cat([enc_v, v], dim=1)
+
+        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
+        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
+
+        joint_out, _ = self.attn(q, k, v)
+        joint_out = joint_out.reshape(batch_size, text_seq_len + img_seq_len, self.inner_dim)
+
+        enc_out = joint_out[:, :text_seq_len]
+        img_out = joint_out[:, text_seq_len:]
+
+        img_out, _ = self.to_out[0](img_out)
+        img_out = self.to_out[1](img_out)
+        enc_out, _ = self.to_add_out(enc_out)
+        return img_out, enc_out
+
+
+class FluxSingleStreamAttention(nn.Module):
+    """Self-attention on concatenated text+image sequence (single blocks)."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+        q = q.view(batch_size, seq_len, self.heads, self.head_dim)
+        k = k.view(batch_size, seq_len, self.heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.heads, self.head_dim)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
+        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
+        out, _ = self.attn(q, k, v)
+        return out.reshape(batch_size, seq_len, self.inner_dim)
+
+
+class FluxTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.norm1 = SD3AdaLayerNormZero(dim)
+        self.norm1_context = SD3AdaLayerNormZero(dim)
+        self.attn = FluxJointAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            supported_attention_backends=supported_attention_backends,
+        )
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = SD3FeedForward(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = SD3FeedForward(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del joint_attention_kwargs
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output)
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class FluxSingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 4.0,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.norm = FluxAdaLayerNormZeroSingle(dim)
+        self.proj_mlp = ReplicatedLinear(dim, mlp_hidden_dim, bias=True)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = ReplicatedLinear(dim + mlp_hidden_dim, dim, bias=True)
+        self.attn = FluxSingleStreamAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del joint_attention_kwargs
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)[0])
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)[0]
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+        encoder_hidden_states = hidden_states[:, :text_seq_len]
+        hidden_states = hidden_states[:, text_seq_len:]
+        return encoder_hidden_states, hidden_states
+
+
+class FluxTransformer2DModel(BaseDiT):
+    """FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1."""
+
+    _fsdp_shard_conditions = [
+        lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks."))
+        and n.split(".")[-1].isdigit(),
+    ]
+    _compile_conditions = _fsdp_shard_conditions
+    # HF weight names already match this module layout (cf. SGLang regex maps).
+    param_names_mapping: dict[str, Any] = {}
+    reverse_param_names_mapping: dict[str, Any] = {}
+    lora_param_names_mapping: dict[str, Any] = {}
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        del kwargs
+        super().__init__(config=config, hf_config=hf_config)
+        self.fastvideo_config = config
+        self.hf_config = hf_config
+        arch = config.arch_config
+
+        out_ch = arch.out_channels
+        self.out_channels = out_ch if out_ch is not None else arch.in_channels
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.in_channels
+
+        axes_list = list(arch.axes_dims_rope)
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_list)
+        if arch.guidance_embeds:
+            self.time_text_embed = FluxCombinedTimestepGuidanceTextProjEmbeddings(
+                embedding_dim=self.inner_dim,
+                pooled_projection_dim=arch.pooled_projection_dim,
+            )
+        else:
+            self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+                embedding_dim=self.inner_dim,
+                pooled_projection_dim=arch.pooled_projection_dim,
+            )
+        self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, self.inner_dim)
+        self.x_embedder = ReplicatedLinear(arch.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                )
+                for _ in range(arch.num_layers)
+            ]
+        )
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                )
+                for _ in range(arch.num_single_layers)
+            ]
+        )
+
+        self.norm_out = SD3AdaLayerNormContinuous(
+            self.inner_dim,
+            self.inner_dim,
+            elementwise_affine=False,
+            eps=1e-6,
+            bias=True,
+            norm_type="layer_norm",
+        )
+        self.proj_out = ReplicatedLinear(
+            self.inner_dim,
+            arch.patch_size * arch.patch_size * self.out_channels,
+            bias=True,
+        )
+        self.gradient_checkpointing = False
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        pooled_projections: torch.Tensor | None = None,
+        timestep: torch.LongTensor | torch.Tensor | None = None,
+        img_ids: torch.Tensor | None = None,
+        txt_ids: torch.Tensor | None = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+        controlnet_block_samples: Any | None = None,
+        controlnet_single_block_samples: Any | None = None,
+        controlnet_blocks_repeat: bool = False,
+        **kwargs: Any,
+    ) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]:
+        del kwargs
+        if encoder_hidden_states is None:
+            raise ValueError("encoder_hidden_states must be provided")
+        if pooled_projections is None:
+            raise ValueError("pooled_projections must be provided")
+        if timestep is None:
+            raise ValueError("timestep must be provided")
+        if img_ids is None or txt_ids is None:
+            raise ValueError("img_ids and txt_ids must be provided")
+
+        arch = self.fastvideo_config.arch_config
+        if arch.guidance_embeds and guidance is None:
+            raise ValueError("guidance must be provided when guidance_embeds=True")
+
+        if timestep.dim() == 0:
+            timestep = timestep[None]
+        if timestep.dim() > 1:
+            timestep = timestep.reshape(-1)
+        if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
+            timestep = timestep.expand(hidden_states.shape[0])
+
+        try:
+            get_forward_context()
+            forward_context = nullcontext()
+        except AssertionError:
+            if timestep.numel() == 0:
+                ts0 = 0
+            elif torch.is_floating_point(timestep):
+                ts0 = int(round(timestep[0].item() * 1000))
+            else:
+                ts0 = int(timestep[0].item())
+            forward_context = set_forward_context(current_timestep=ts0, attn_metadata=None)
+
+        with forward_context:
+            hidden_states, _ = self.x_embedder(hidden_states)
+
+            ts = timestep.to(hidden_states.dtype) * 1000
+            g = None if guidance is None else guidance.to(hidden_states.dtype) * 1000
+
+            if arch.guidance_embeds:
+                assert g is not None
+                temb = self.time_text_embed(ts, g, pooled_projections)
+            else:
+                temb = self.time_text_embed(timestep=ts, pooled_projection=pooled_projections)
+
+            encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)
+
+            if txt_ids.ndim == 3:
+                txt_ids = txt_ids[0]
+            if img_ids.ndim == 3:
+                img_ids = img_ids[0]
+
+            ids = torch.cat((txt_ids, img_ids), dim=0)
+            image_rotary_emb = self.pos_embed(ids)
+
+            jkwargs = joint_attention_kwargs or {}
+
+            for idx, block in enumerate(self.transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=jkwargs,
+                )
+                if controlnet_block_samples:
+                    interval = len(self.transformer_blocks) / len(controlnet_block_samples)
+                    interval = int(math.ceil(interval))
+                    if controlnet_blocks_repeat:
+                        hidden_states = hidden_states + controlnet_block_samples[idx % len(controlnet_block_samples)]
+                    else:
+                        hidden_states = hidden_states + controlnet_block_samples[idx // interval]
+
+            for idx, block in enumerate(self.single_transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=jkwargs,
+                )
+                if controlnet_single_block_samples:
+                    interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                    interval = int(math.ceil(interval))
+                    hidden_states = hidden_states + controlnet_single_block_samples[idx // interval]
+
+            hidden_states = self.norm_out(hidden_states, temb)
+            output, _ = self.proj_out(hidden_states)
+
+            if not return_dict:
+                return (output,)
+            return FluxTransformer2DModelOutput(sample=output)
+
+
+EntryClass = FluxTransformer2DModel

fastvideo/models/dits/glm_image.py

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from fastvideo.layers.mlp import MLP
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class GlmImageLayerKVCache:
+
+    def __init__(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode: Optional[str] = None
+
+    def store(self, k: torch.Tensor, v: torch.Tensor):
+        # Append along seq (dim=1).
+        if self.k_cache is None:
+            self.k_cache = k
+            self.v_cache = v
+        else:
+            self.k_cache = torch.cat([self.k_cache, k], dim=1)
+            self.v_cache = torch.cat([self.v_cache, v], dim=1)
+
+    def get(self):
+        return self.k_cache, self.v_cache
+
+    def clear(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode = None
+
+
+class GlmImageKVCache:
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
+
+    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
+        return self.caches[layer_idx]
+
+    def set_mode(self, mode: Optional[str]):
+        if mode is not None and mode not in ["write", "read", "skip"]:
+            raise ValueError(
+                f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'"
+            )
+        for cache in self.caches:
+            cache.mode = mode
+
+    def clear(self):
+        for cache in self.caches:
+            cache.clear()
+
+
+# =============================================================================
+# Timestep and Text Projection
+# =============================================================================
+class GlmImageTimestepEmbedding(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+    ):
+        super().__init__()
+        if out_dim is None:
+            out_dim = time_embed_dim
+        self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True)
+        if act_fn == "silu":
+            self.act = nn.SiLU()
+        elif act_fn == "gelu":
+            self.act = nn.GELU(approximate="tanh")
+        else:
+            self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        sample, _ = self.linear_1(sample)
+        sample = self.act(sample)
+        sample, _ = self.linear_2(sample)
+        return sample
+
+
+class GlmImageTextProjection(nn.Module):
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_size: int,
+        out_features: int = None,
+        act_fn: str = "silu",
+    ):
+        super().__init__()
+        if out_features is None:
+            out_features = hidden_size
+        self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True)
+        if act_fn == "silu":
+            self.act_1 = nn.SiLU()
+        elif act_fn == "gelu_tanh":
+            self.act_1 = nn.GELU(approximate="tanh")
+        else:
+            self.act_1 = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True)
+
+    def forward(self, caption: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.linear_1(caption)
+        hidden_states = self.act_1(hidden_states)
+        hidden_states, _ = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        condition_dim: int,
+        pooled_projection_dim: int,
+        timesteps_dim: int = 256,
+    ):
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.condition_proj = Timesteps(
+            num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = GlmImageTimestepEmbedding(
+            in_channels=timesteps_dim, time_embed_dim=embedding_dim
+        )
+        self.condition_embedder = GlmImageTextProjection(
+            pooled_projection_dim, embedding_dim, act_fn="silu"
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        hidden_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+
+        crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(
+            crop_coords.size(0), -1
+        )
+        target_size_proj = self.condition_proj(target_size.flatten()).view(
+            target_size.size(0), -1
+        )
+
+        condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1)
+
+        timesteps_emb = self.timestep_embedder(
+            timesteps_proj.to(dtype=hidden_dtype)
+        )
+        condition_emb = self.condition_embedder(
+            condition_proj.to(dtype=hidden_dtype)
+        )
+
+        conditioning = timesteps_emb + condition_emb
+        return conditioning
+
+
+# =============================================================================
+# Image Projector
+# =============================================================================
+class GlmImageImageProjector(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int = 16,
+        hidden_size: int = 2560,
+        patch_size: int = 2,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, channel, height, width = hidden_states.shape
+        post_patch_height = height // self.patch_size
+        post_patch_width = width // self.patch_size
+
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            channel,
+            post_patch_height,
+            self.patch_size,
+            post_patch_width,
+            self.patch_size,
+        )
+        hidden_states = (
+            hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+        )
+        hidden_states = self.proj(hidden_states)
+
+        return hidden_states
+
+
+# =============================================================================
+# AdaLayerNorm
+# =============================================================================
+class GlmImageAdaLayerNormZero(nn.Module):
+
+    def __init__(self, embedding_dim: int, dim: int) -> None:
+        super().__init__()
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+    ) -> Tuple[torch.Tensor, ...]:
+        dtype = hidden_states.dtype
+        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
+        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(
+            dtype=dtype
+        )
+
+        emb, _ = self.linear(temb)
+        (
+            shift_msa,
+            c_shift_msa,
+            scale_msa,
+            c_scale_msa,
+            gate_msa,
+            c_gate_msa,
+            shift_mlp,
+            c_shift_mlp,
+            scale_mlp,
+            c_scale_mlp,
+            gate_mlp,
+            c_gate_mlp,
+        ) = emb.chunk(12, dim=1)
+
+        hidden_states = norm_hidden_states * (
+            1 + scale_msa.unsqueeze(1)
+        ) + shift_msa.unsqueeze(1)
+        encoder_hidden_states = norm_encoder_hidden_states * (
+            1 + c_scale_msa.unsqueeze(1)
+        ) + c_shift_msa.unsqueeze(1)
+
+        return (
+            hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        )
+
+
+# =============================================================================
+# Attention
+# =============================================================================
+class GlmImageAttention(nn.Module):
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int,
+        dim_head: int,
+        out_dim: int,
+        bias: bool = True,
+        qk_norm: str = "layer_norm",
+        elementwise_affine: bool = False,
+        eps: float = 1e-5,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.num_kv_heads = self.heads
+        self.dim_head = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+
+        self.to_out = nn.ModuleList(
+            [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)]
+        )
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "layer_norm":
+            self.norm_q = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+            self.norm_k = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+        else:
+            raise ValueError(f"unknown qk_norm: {qk_norm}")
+
+        self.attn = LocalAttention(
+            num_heads=self.heads,
+            head_size=dim_head,
+            num_kv_heads=self.heads,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        dtype = encoder_hidden_states.dtype
+
+        batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
+        batch_size, image_seq_length, embed_dim = hidden_states.shape
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # 1. QKV projections
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(hidden_states)
+        value, _ = self.to_v(hidden_states)
+
+        query = query.unflatten(2, (self.heads, -1))
+        key = key.unflatten(2, (self.heads, -1))
+        value = value.unflatten(2, (self.heads, -1))
+
+        # 2. QK normalization
+        if self.norm_q is not None:
+            query = self.norm_q(query).to(dtype=dtype)
+        if self.norm_k is not None:
+            key = self.norm_k(key).to(dtype=dtype)
+
+        # 3. Rotational positional embeddings applied to latent stream
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+
+            query[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+            )
+            key[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+            )
+
+        # 4. KV Cache handling
+        if kv_cache is not None:
+            if kv_cache.mode == "write":
+                kv_cache.store(key, value)
+            elif kv_cache.mode == "read":
+                # Prepend cached condition k/v along seq (dim=1).
+                k_cache, v_cache = kv_cache.get()
+                key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key
+                value = (
+                    torch.cat([v_cache, value], dim=1) if v_cache is not None else value
+                )
+            elif kv_cache.mode == "skip":
+                pass
+
+        hidden_states = self.attn(query, key, value)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # 6. Output projection
+        hidden_states, _ = self.to_out[0](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+# =============================================================================
+# Transformer Block
+# =============================================================================
+class GlmImageTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        dim: int = 2560,
+        num_attention_heads: int = 64,
+        attention_head_dim: int = 40,
+        time_embed_dim: int = 512,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        # 1. Attention
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+
+        self.attn1 = GlmImageAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            out_dim=dim,
+            bias=True,
+            qk_norm="layer_norm",
+            elementwise_affine=False,
+            eps=1e-5,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn1",
+        )
+
+        # 2. Feedforward with fused ScaleResidualLayerNorm
+        self.norm2 = ScaleResidualLayerNormScaleShift(
+            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
+        )
+        self.norm2_context = ScaleResidualLayerNormScaleShift(
+            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
+        )
+        self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Timestep conditioning
+        (
+            norm_hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            norm_encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
+
+        # 2. Attention
+        if attention_kwargs is None:
+            attention_kwargs = {}
+
+        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            kv_cache=kv_cache,
+            **attention_kwargs,
+        )
+
+        # 3. Feedforward (fused residual + norm + scale/shift)
+        norm_hidden_states, hidden_states = self.norm2(
+            hidden_states,
+            attn_hidden_states,
+            gate_msa.unsqueeze(1),
+            shift_mlp.unsqueeze(1),
+            scale_mlp.unsqueeze(1),
+        )
+        norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context(
+            encoder_hidden_states,
+            attn_encoder_hidden_states,
+            c_gate_msa.unsqueeze(1),
+            c_shift_mlp.unsqueeze(1),
+            c_scale_mlp.unsqueeze(1),
+        )
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output_context = self.ff(norm_encoder_hidden_states)
+        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = (
+            encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
+        )
+
+        return hidden_states, encoder_hidden_states
+
+
+# =============================================================================
+# Rotary Positional Embedding
+# =============================================================================
+class GlmImageRotaryPosEmbed(nn.Module):
+
+    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.theta = theta
+        self._cache_key: tuple | None = None
+        self._cache_value: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        _, _, raw_h, raw_w = hidden_states.shape
+        height = raw_h // self.patch_size
+        width = raw_w // self.patch_size
+        device = hidden_states.device
+
+        cache_key = (height, width, device.type,
+                     device.index if device.index is not None else -1)
+        if self._cache_key == cache_key and self._cache_value is not None:
+            return self._cache_value
+
+        dim_h, dim_w = self.dim // 2, self.dim // 2
+        h_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[
+                    : (dim_h // 2)
+                ]
+                / dim_h
+            )
+        )
+        w_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[
+                    : (dim_w // 2)
+                ]
+                / dim_w
+            )
+        )
+        h_seq = torch.arange(height, device=device)
+        w_seq = torch.arange(width, device=device)
+        freqs_h = torch.outer(h_seq, h_inv_freq).unsqueeze(1).expand(height, width, -1)
+        freqs_w = torch.outer(w_seq, w_inv_freq).unsqueeze(0).expand(height, width, -1)
+
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1).reshape(height * width, -1)
+        result = (freqs.cos(), freqs.sin())
+        self._cache_key = cache_key
+        self._cache_value = result
+        return result
+
+
+# =============================================================================
+# Final AdaLayerNorm
+# =============================================================================
+class GlmImageAdaLayerNormContinuous(nn.Module):
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+        norm_type: str = "layer_norm",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(
+            conditioning_embedding_dim, embedding_dim * 2, bias=bias
+        )
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor
+    ) -> torch.Tensor:
+        emb = self.linear(conditioning_embedding.to(x.dtype))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# =============================================================================
+# Main Model
+# =============================================================================
+class GlmImageTransformer2DModel(BaseDiT):
+
+    _fsdp_shard_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = GlmImageDiTConfig().arch_config._compile_conditions
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    param_names_mapping = GlmImageDiTConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+    lora_param_names_mapping = {}
+
+    def __init__(
+        self,
+        config: GlmImageDiTConfig,
+        hf_config: dict[str, Any],
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+
+        arch_config = config.arch_config
+
+        self.in_channels = arch_config.in_channels
+        self.out_channels = arch_config.out_channels
+        self.patch_size = arch_config.patch_size
+        self.num_layers = arch_config.num_layers
+        self.attention_head_dim = arch_config.attention_head_dim
+        self.num_attention_heads = arch_config.num_attention_heads
+        self.text_embed_dim = arch_config.text_embed_dim
+        self.time_embed_dim = arch_config.time_embed_dim
+
+        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
+        # Each of these are sincos embeddings of shape 2 * condition_dim
+        pooled_projection_dim = 2 * 2 * arch_config.condition_dim
+        inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
+
+        self.hidden_size = inner_dim
+        self.num_channels_latents = arch_config.out_channels
+
+        # 1. RoPE
+        self.rotary_emb = GlmImageRotaryPosEmbed(
+            arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0
+        )
+
+        # 2. Patch & Text-timestep embedding
+        self.image_projector = GlmImageImageProjector(
+            arch_config.in_channels, inner_dim, arch_config.patch_size
+        )
+        self.glyph_projector = MLP(
+            input_dim=arch_config.text_embed_dim,
+            mlp_hidden_dim=inner_dim,
+            output_dim=inner_dim,
+            act_type="gelu",
+        )
+        self.prior_token_embedding = nn.Embedding(
+            arch_config.prior_vq_quantizer_codebook_size, inner_dim
+        )
+        self.prior_projector = MLP(
+            input_dim=inner_dim,
+            mlp_hidden_dim=inner_dim,
+            output_dim=inner_dim,
+            act_type="silu",
+        )
+
+        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
+            embedding_dim=arch_config.time_embed_dim,
+            condition_dim=arch_config.condition_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            timesteps_dim=arch_config.time_embed_dim,
+        )
+
+        # 3. Transformer blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                GlmImageTransformerBlock(
+                    inner_dim,
+                    arch_config.num_attention_heads,
+                    arch_config.attention_head_dim,
+                    arch_config.time_embed_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                    prefix=f"transformer_blocks.{i}",
+                )
+                for i in range(arch_config.num_layers)
+            ]
+        )
+
+        # 4. Output projection
+        self.norm_out = GlmImageAdaLayerNormContinuous(
+            inner_dim, arch_config.time_embed_dim, elementwise_affine=False
+        )
+        self.proj_out = nn.Linear(
+            inner_dim,
+            arch_config.patch_size * arch_config.patch_size * arch_config.out_channels,
+            bias=True,
+        )
+
+        self.gradient_checkpointing = False
+
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        prior_token_id: torch.Tensor,
+        prior_token_drop: torch.Tensor,
+        timestep: torch.LongTensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_caches: Optional[GlmImageKVCache] = None,
+        kv_caches_mode: Optional[str] = None,
+        freqs_cis: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        guidance: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if kv_caches is not None:
+            kv_caches.set_mode(kv_caches_mode)
+
+        batch_size, num_channels, height, width = hidden_states.shape
+
+        if isinstance(encoder_hidden_states, list):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        # 1. RoPE
+        image_rotary_emb = freqs_cis
+        if image_rotary_emb is None:
+            image_rotary_emb = self.rotary_emb(hidden_states)
+
+        # 2. Patch & Timestep embeddings
+        p = self.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.image_projector(hidden_states)
+        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
+        prior_embedding = self.prior_token_embedding(prior_token_id)
+        # Zero dropped priors by multiply: boolean indexing + .any() syncs each step.
+        keep = (~prior_token_drop).to(device=prior_embedding.device, dtype=prior_embedding.dtype)
+        while keep.dim() < prior_embedding.dim():
+            keep = keep.unsqueeze(-1)
+        prior_embedding = prior_embedding * keep
+        prior_hidden_states = self.prior_projector(prior_embedding)
+        hidden_states = hidden_states + prior_hidden_states
+
+        temb = self.time_condition_embed(
+            timestep, target_size, crop_coords, hidden_states.dtype
+        )
+        temb = F.silu(temb)
+
+        # 3. Transformer blocks
+        for idx, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                image_rotary_emb,
+                attention_kwargs,
+                kv_cache=kv_caches[idx] if kv_caches is not None else None,
+            )
+
+        # 4. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 5. Unpatchify
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_height, post_patch_width, -1, p, p
+        )
+        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+
+        return output.float()
+
+
+EntryClass = GlmImageTransformer2DModel

fastvideo/models/dits/ltx2.py

@@ -72,8 +72,13 @@ def _rms_norm_dispatch(
     """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
     if _is_ltx2_refine_stage():
         return _quack_rmsnorm(x, weight=weight, eps=eps)
+    # Official LTX-2 runs the DiT natively in bf16 with no autocast, so its
+    # rms_norm outputs stay bf16 on every torch version. torch 2.12 added
+    # rms_norm to autocast's float32 cast policy, which under our
+    # autocast(bfloat16) forward returns float32 instead. Restore the input
+    # dtype to match the official numerics (and torch <= 2.11 behavior).
     return torch.nn.functional.rms_norm(
-        x, (x.shape[-1], ), weight=weight, eps=eps)
+        x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
 
 
 class StageAwareRMSNorm(nn.RMSNorm):

fastvideo/models/dits/wanvideo.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import copy
 import math
 from typing import Any
 
@@ -59,6 +60,11 @@ class WanTimeTextImageEmbedding(nn.Module):
         time_freq_dim: int,
         text_embed_dim: int,
         image_embed_dim: int | None = None,
+        *,
+        r_embedder: bool = False,
+        r_embedder_fusion: str = "additive",
+        r_embedder_gate_value: float = 0.25,
+        r_embedder_deltatime_type: str = "r",
     ):
         super().__init__()
 
@@ -77,14 +83,57 @@ class WanTimeTextImageEmbedding(nn.Module):
         if image_embed_dim is not None:
             self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
 
+        # AnyFlow dual-timestep support. When r_embedder is False the forward
+        # path bypasses delta_embedder entirely and the output is byte-identical
+        # to the legacy single-timestep implementation.
+        self._r_embedder_enabled = bool(r_embedder)
+        self._r_embedder_fusion = r_embedder_fusion
+        self._r_embedder_deltatime_type = r_embedder_deltatime_type
+        if self._r_embedder_enabled:
+            if r_embedder_fusion not in ("additive", "gated"):
+                raise ValueError(
+                    "r_embedder_fusion must be one of {additive, gated}, "
+                    f"got {r_embedder_fusion!r}")
+            if r_embedder_deltatime_type not in ("r", "t-r"):
+                raise ValueError(
+                    "r_embedder_deltatime_type must be one of {r, t-r}, "
+                    f"got {r_embedder_deltatime_type!r}")
+            # Deep-copy preserves identical initialization with time_embedder,
+            # matching AnyFlow reference setup_flowmap_model() behavior.
+            self.delta_embedder = copy.deepcopy(self.time_embedder)
+            # Non-persistent buffer — gate is a hyperparameter, not learned.
+            self.register_buffer(
+                "_r_embedder_gate",
+                torch.tensor(float(r_embedder_gate_value)),
+                persistent=False,
+            )
+        else:
+            self.delta_embedder = None
+
     def forward(
         self,
         timestep: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: torch.Tensor | None = None,
         timestep_seq_len: int | None = None,
+        r_timestep: torch.Tensor | None = None,
     ):
         temb = self.time_embedder(timestep, timestep_seq_len)
+
+        if self._r_embedder_enabled and r_timestep is not None:
+            assert self.delta_embedder is not None
+            if self._r_embedder_deltatime_type == "r":
+                delta_input = r_timestep
+            else:
+                delta_input = timestep - r_timestep
+            delta_emb = self.delta_embedder(delta_input, timestep_seq_len)
+            gate = self._r_embedder_gate
+            if self._r_embedder_fusion == "gated":
+                temb = (1.0 - gate) * temb + gate * delta_emb
+            else:
+                # Additive (additional channel, no convex blend).
+                temb = temb + gate * delta_emb
+
         timestep_proj = self.time_modulation(temb)
 
         if self.text_embedder is not None:
@@ -595,6 +644,10 @@ class WanTransformer3DModel(BaseDiT):
             time_freq_dim=config.freq_dim,
             text_embed_dim=config.text_dim,
             image_embed_dim=config.image_dim,
+            r_embedder=config.r_embedder,
+            r_embedder_fusion=config.r_embedder_fusion,
+            r_embedder_gate_value=config.r_embedder_gate_value,
+            r_embedder_deltatime_type=config.r_embedder_deltatime_type,
         )
 
         # 3. Transformer blocks
@@ -636,6 +689,7 @@ class WanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 guidance=None,
+                r_timestep: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
@@ -684,8 +738,17 @@ class WanTransformer3DModel(BaseDiT):
         else:
             ts_seq_len = None
 
+        # AnyFlow dual-timestep — match timestep's flattening so embedder
+        # sees aligned shapes.
+        if r_timestep is not None and r_timestep.dim() == 2:
+            r_timestep = r_timestep.flatten()
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
+            timestep,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+            timestep_seq_len=ts_seq_len,
+            r_timestep=r_timestep)
         if ts_seq_len is not None:
             # batch_size, seq_len, 6, inner_dim
             timestep_proj = timestep_proj.unflatten(2, (6, -1))

fastvideo/models/encoders/glm_image_ar_loader.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sole HF-import boundary for GLM-Image's AR encoder (lazy-wrapper exception E001)."""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class GlmImageARLoader(nn.Module):
+
+    def __init__(self, model_path: str, processor_path: str | None = None,
+                 *, torch_dtype: torch.dtype = torch.bfloat16,
+                 trust_remote_code: bool = True) -> None:
+        super().__init__()
+        from transformers import (AutoProcessor,
+                                  GlmImageForConditionalGeneration)
+        logger.info("Loading GLM-Image AR encoder from %s", model_path)
+        self._model = GlmImageForConditionalGeneration.from_pretrained(
+            model_path,
+            torch_dtype=torch_dtype,
+            trust_remote_code=trust_remote_code,
+        )
+        if processor_path is not None:
+            logger.info("Loading GLM-Image processor from %s", processor_path)
+            self.processor = AutoProcessor.from_pretrained(
+                processor_path, trust_remote_code=trust_remote_code)
+        else:
+            self.processor = None
+
+    @torch.no_grad()
+    def generate(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        return self._model.generate(*args, **kwargs)
+
+    @torch.no_grad()
+    def get_image_features(self, pixel_values: torch.Tensor,
+                           image_grid_thw: torch.Tensor) -> Any:
+        return self._model.get_image_features(pixel_values, image_grid_thw)
+
+    @torch.no_grad()
+    def get_image_tokens(self, image_embeds: torch.Tensor,
+                         image_grid_thw: torch.Tensor) -> torch.Tensor:
+        return self._model.get_image_tokens(image_embeds, image_grid_thw)
+
+    @property
+    def config(self):  # type: ignore[no-untyped-def]
+        return self._model.config
+
+    @property
+    def generation_config(self):  # type: ignore[no-untyped-def]
+        return self._model.generation_config
+
+    def to(self, *args, **kwargs):  # type: ignore[override]
+        self._model = self._model.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+    def eval(self):  # type: ignore[override]
+        self._model = self._model.eval()
+        return super().eval()

fastvideo/models/encoders/qwen2_5_vl_custom.py

@@ -75,6 +75,25 @@ if not is_flash_attn_2_available():
 
 logger = logging.get_logger(__name__)
 
+
+def _resolve_torch_dtype(dtype, default: torch.dtype = torch.float32) -> torch.dtype:
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    if isinstance(dtype, str):
+        normalized = dtype.strip().removeprefix("torch.").lower()
+        return {
+            "bfloat16": torch.bfloat16,
+            "bf16": torch.bfloat16,
+            "float16": torch.float16,
+            "fp16": torch.float16,
+            "half": torch.float16,
+            "float32": torch.float32,
+            "fp32": torch.float32,
+            "float": torch.float32,
+        }.get(normalized, default)
+    return default
+
+
 class Qwen2_5_VLMLP(nn.Module):
     def __init__(self, config, bias: bool = False):
         super().__init__()
@@ -304,10 +323,13 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module):
     config_class = Qwen2_5_VLVisionConfig
     _no_split_modules = ["Qwen2_5_VLVisionBlock"]
 
-    def __init__(self, config) -> None:
+    def __init__(self, config, parent_torch_dtype=None) -> None:
         super().__init__()
 
-        self.dtype = torch.bfloat16 if config.torch_dtype == "bfloat16" else torch.float32
+        config_torch_dtype = getattr(config, "torch_dtype", None)
+        self.dtype = _resolve_torch_dtype(
+            config_torch_dtype if config_torch_dtype is not None else parent_torch_dtype
+        )
 
         self.spatial_merge_size = config.spatial_merge_size
         self.patch_size = config.patch_size
@@ -1458,7 +1480,10 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
         super().__init__()
         config = _flatten_text_config(config)
         self.config = config
-        self.visual = Qwen2_5_VisionTransformerPretrainedModel(config.vision_config)
+        self.visual = Qwen2_5_VisionTransformerPretrainedModel(
+            config.vision_config,
+            parent_torch_dtype=getattr(config, "torch_dtype", None),
+        )
 
         self.model = Qwen2_5_VLModel(config)
         self.vocab_size = config.vocab_size

fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py

@@ -42,7 +42,7 @@ def load_independent(files: list[str], device: str):
     """Before-PR behavior: every rank reads every tensor from disk to GPU."""
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f:
+            for name in f.keys():  # noqa: SIM118
                 param = f.get_tensor(name)
                 yield name, param
 
@@ -54,7 +54,7 @@ def load_broadcast(files: list[str], device: str, node_group,
     handles = []
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f:
+            for name in f.keys():  # noqa: SIM118
                 if local_rank == 0:
                     param = f.get_tensor(name)
                 else:

fastvideo/models/loader/component_loader.py

@@ -95,6 +95,8 @@ class ComponentLoader(ABC):
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
+            "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
+            "processor": (ProcessorLoader, "transformers"),
             "upsampler": (UpsamplerLoader, "diffusers"),
             "upsampler_2": (UpsamplerLoader, "diffusers"),
             # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
@@ -523,6 +525,39 @@ class ImageEncoderLoader(TextEncoderLoader):
         )
 
 
+class VisionLanguageEncoderLoader(ComponentLoader):
+    """Loader for vision-language autoregressive encoders."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        from fastvideo.distributed.parallel_state import get_local_torch_device
+        from fastvideo.models.encoders.glm_image_ar_loader import (
+            GlmImageARLoader)
+
+        logger.info("Loading vision-language encoder from %s", model_path)
+        target_device = get_local_torch_device()
+        loader = GlmImageARLoader(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=fastvideo_args.trust_remote_code,
+        ).to(target_device).eval()
+        return loader
+
+
+class ProcessorLoader(ComponentLoader):
+    """Loader for HF processors that pair with vision-language encoders."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        from transformers import AutoProcessor
+
+        logger.info("Loading processor from %s", model_path)
+        processor = AutoProcessor.from_pretrained(
+            model_path,
+            trust_remote_code=fastvideo_args.trust_remote_code,
+        )
+        logger.info("Loaded processor: %s", processor.__class__.__name__)
+        return processor
+
+
 class ImageProcessorLoader(ComponentLoader):
     """Loader for image processor."""
 

fastvideo/models/registry.py

@@ -61,6 +61,11 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
 }
 
+# Text-to-image DiT models (2D image generation)
+_TEXT_TO_IMAGE_DIT_MODELS = {
+    "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"),
+}
+
 _TEXT_ENCODER_MODELS = {
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection":
@@ -134,6 +139,7 @@ _UPSAMPLERS = {
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
+    **_TEXT_TO_IMAGE_DIT_MODELS,
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,

fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flow-map any-step Euler scheduler for AnyFlow.
+
+The model predicts the *average* velocity ``u_θ(x_t, t, r)`` from time
+``t`` back to time ``r``, so one Euler step is
+
+    x_r = x_t - ((t - r) / num_train_timesteps) * u_θ(x_t, t, r)
+
+regardless of how far apart ``t`` and ``r`` are. The scheduler also
+provides the AnyFlow training-time helpers ``apply_shift`` (flow-matching
+shift transform) and ``get_train_weight`` (per-timestep loss weight,
+including ``beta08``).
+
+Standalone — does not depend on diffusers' ConfigMixin/SchedulerMixin.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+
+from fastvideo.models.schedulers.base import BaseScheduler
+
+
+WeightType = Literal["uniform", "gaussian", "beta08"]
+
+
+class FlowMapEulerDiscreteScheduler(BaseScheduler):
+    """Minimal flow-map scheduler.
+
+    Parameters
+    ----------
+    num_train_timesteps:
+        Discretization granularity for training. ``t`` is expressed in
+        absolute units in ``[0, num_train_timesteps]``.
+    shift:
+        Flow-matching shift parameter (Wan video default: ``5.0``). Set
+        to ``1.0`` for an identity shift.
+    """
+
+    order: int = 1
+
+    def __init__(
+        self,
+        *,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+    ) -> None:
+        self.num_train_timesteps = int(num_train_timesteps)
+        self.shift = float(shift)
+        self.timesteps: torch.Tensor = torch.empty(0)
+        self.sigmas: torch.Tensor = torch.empty(0)
+        super().__init__()
+
+    # ------------------------------------------------------------------
+    # BaseScheduler abstract surface
+
+    def set_shift(self, shift: float) -> None:
+        self.shift = float(shift)
+
+    def scale_model_input(
+        self,
+        sample: torch.Tensor,
+        timestep: int | None = None,
+    ) -> torch.Tensor:
+        # Flow-matching has no per-step input scaling; pass through.
+        del timestep
+        return sample
+
+    # ------------------------------------------------------------------
+    # Public helpers used by the AnyFlow pretrain method.
+
+    def apply_shift(
+        self,
+        t: torch.Tensor,
+        *,
+        shift: float | None = None,
+    ) -> torch.Tensor:
+        """Apply the flow-matching shift: ``t' = s * t / (1 + (s - 1) * t)``.
+
+        Operates in the normalized ``[0, 1]`` domain — callers should pass
+        ``t / num_train_timesteps`` (or sample ``t`` directly from
+        ``[0, 1]``).
+        """
+        s = self.shift if shift is None else float(shift)
+        if s == 1.0:
+            return t
+        return s * t / (1.0 + (s - 1.0) * t)
+
+    def get_train_weight(
+        self,
+        t: torch.Tensor,
+        *,
+        weight_type: WeightType = "beta08",
+    ) -> torch.Tensor:
+        """Per-timestep training weight, renormalized so the total weight
+        mass equals ``num_train_timesteps`` (matching AnyFlow reference's
+        ``scheduling_flowmap_euler_discrete.py``).
+
+        ``beta08``: ``w(t) = t * sqrt(1 - t)`` (in normalized t-space).
+        """
+        # Auto-detect domain: if t was given in absolute units, normalize.
+        t_f = t.float()
+        max_val = t_f.max() if t_f.numel() > 0 else torch.tensor(0.0)
+        if max_val > 1.0 + 1e-6:
+            t_norm = t_f / self.num_train_timesteps
+        else:
+            t_norm = t_f
+        t_norm = t_norm.clamp(min=0.0, max=1.0)
+
+        if weight_type == "uniform":
+            w = torch.ones_like(t_norm)
+        elif weight_type == "gaussian":
+            w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2) ** 2)
+        elif weight_type == "beta08":
+            w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5)
+        else:
+            raise ValueError(f"Unknown weight_type: {weight_type!r}")
+
+        denom = w.sum().clamp_min(1e-8)
+        return w * (float(self.num_train_timesteps) / denom)
+
+    # ------------------------------------------------------------------
+
+    def set_timesteps(
+        self,
+        *,
+        num_inference_steps: int,
+        device: torch.device | str = "cpu",
+        custom_timesteps: list[float] | torch.Tensor | None = None,
+    ) -> None:
+        """Build a descending timestep schedule ending at 0.
+
+        With ``num_inference_steps=N`` the schedule has ``N + 1`` entries
+        ``[T_max, ..., 0]`` so a rollout consumes ``N`` Euler steps.
+
+        ``custom_timesteps`` overrides the linspace+shift schedule with
+        a pinned list (in absolute train-timestep units), useful for the
+        AnyFlow paper's hand-tuned ``[999, 937, 833, 624, 0]`` schedule.
+        """
+        if num_inference_steps <= 0:
+            raise ValueError(
+                "num_inference_steps must be positive, "
+                f"got {num_inference_steps}")
+        device = torch.device(device)
+
+        if custom_timesteps is not None:
+            ts = torch.as_tensor(
+                custom_timesteps, dtype=torch.float32, device=device)
+            if ts.ndim != 1:
+                raise ValueError(
+                    "custom_timesteps must be 1-D, got shape "
+                    f"{tuple(ts.shape)}")
+            if not torch.all(ts[:-1] >= ts[1:]):
+                raise ValueError(
+                    "custom_timesteps must be descending (largest first)")
+        else:
+            ts_norm = torch.linspace(
+                1.0, 0.0, num_inference_steps + 1, device=device)
+            ts_norm = self.apply_shift(ts_norm)
+            ts = ts_norm * self.num_train_timesteps
+
+        self.timesteps = ts
+        self.sigmas = ts / self.num_train_timesteps
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        *,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        r_timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """One Euler step from ``t`` to ``r``.
+
+        ``model_output`` is the average-velocity prediction
+        ``u_θ(x_t, t, r)``. Both ``timestep`` and ``r_timestep`` are in
+        absolute train-timestep units (``[0, num_train_timesteps]``).
+        """
+        t = timestep.to(sample.device, dtype=sample.dtype)
+        r = r_timestep.to(sample.device, dtype=sample.dtype)
+        dt_norm = (t - r) / float(self.num_train_timesteps)
+        # Broadcast dt over channel/spatial dims.
+        view: list[int] = [-1] + [1] * (sample.ndim - 1)
+        return sample - dt_norm.view(*view) * model_output
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``,
+        where ``σ = t / num_train_timesteps``.
+        """
+        sigma = (timestep.to(original_samples.device,
+                             dtype=original_samples.dtype)
+                 / float(self.num_train_timesteps))
+        view: list[int] = [-1] + [1] * (original_samples.ndim - 1)
+        sigma = sigma.view(*view)
+        return (1.0 - sigma) * original_samples + sigma * noise

fastvideo/performance/hf_store.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import pandas as pd
 from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.constants import ENDPOINT
 
 # ---------------------------------------------------------------------------
 # Configuration — read once at import time, shared across both consumers
@@ -35,7 +36,8 @@ SYNC_REUSE_TTL_SECONDS = int(os.environ.get("PERFORMANCE_TRACKING_SYNC_REUSE_TTL
 
 def sanitize(value: str) -> str:
     """Return a filesystem- and HF-path-safe version of *value*."""
-    return re.sub(r"[^A-Za-z0-9._-]", "_", value)
+    sanitized = re.sub(r"[^A-Za-z0-9._-]", "_", value)
+    return f"_{sanitized}" if not sanitized or sanitized.startswith(".") else sanitized
 
 
 def safe_float(value: Any) -> float | None:
@@ -60,6 +62,19 @@ def is_baseline_eligible_record(record: dict[str, Any]) -> bool:
     return "baseline_eligible" not in record and "run_source" not in record
 
 
+def _parse_record_timestamp(record: dict[str, Any]) -> datetime | None:
+    raw_ts = record.get("timestamp")
+    if not raw_ts:
+        return None
+    try:
+        ts = datetime.fromisoformat(str(raw_ts))
+    except ValueError:
+        return None
+    if ts.tzinfo is None:
+        ts = ts.replace(tzinfo=timezone.utc)
+    return ts
+
+
 def resolve_hf_token() -> str | None:
     """Return the first configured Hugging Face token env var."""
     for env_var in HF_TOKEN_ENV_VARS:
@@ -95,11 +110,23 @@ def _sync_marker_is_fresh(marker_path: str) -> bool:
     return age.total_seconds() <= SYNC_REUSE_TTL_SECONDS
 
 
+def _sync_marker_matches_request(marker_path: str, revision: str | None) -> bool:
+    try:
+        with open(marker_path, encoding="utf-8") as marker:
+            marker_data = json.load(marker)
+    except (OSError, json.JSONDecodeError):
+        return False
+    if marker_data.get("endpoint") != ENDPOINT or marker_data.get("repo_id") != HF_REPO_ID:
+        return False
+    return marker_data.get("revision") == revision
+
+
 def sync_from_hf(
     local_dir: str,
     *,
     strict: bool = False,
     reuse_existing: bool = False,
+    revision: str | None = None,
 ) -> str:
     """Download the HF dataset repo snapshot to *local_dir*.
 
@@ -116,14 +143,18 @@ def sync_from_hf(
     snapshot checks when compare and dashboard scripts run sequentially in the
     same CI job, without silently reusing stale data in persistent local or
     long-lived runner environments.
+
+    Pass ``revision`` to pin the snapshot to a previously read Hub commit. This
+    is used by conditional writers that must validate one exact remote state
+    before committing with that revision as their parent.
     """
     marker_path = _sync_marker_path(local_dir)
     if reuse_existing and os.path.exists(marker_path):
-        if _sync_marker_is_fresh(marker_path):
+        if _sync_marker_is_fresh(marker_path) and _sync_marker_matches_request(marker_path, revision):
             print(f"hf_store: reusing existing sync at {local_dir}")
             return local_dir
         os.remove(marker_path)
-        print(f"hf_store: existing sync at {local_dir} is stale; refreshing")
+        print(f"hf_store: existing sync at {local_dir} is stale or mismatched; refreshing")
 
     if not reuse_existing and os.path.exists(marker_path):
         os.remove(marker_path)
@@ -143,13 +174,17 @@ def sync_from_hf(
             local_dir=local_dir,
             token=resolve_hf_token(),
             allow_patterns="*.json",
+            revision=revision,
         )
         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)
+            json.dump(
+                {
+                    "endpoint": ENDPOINT,
+                    "repo_id": HF_REPO_ID,
+                    "synced_at": datetime.now(timezone.utc).isoformat(),
+                    "revision": revision,
+                }, marker)
     except Exception as exc:
         if strict:
             raise
@@ -231,7 +266,7 @@ def load_records(
     if days is not None:
         cutoff = datetime.now(timezone.utc) - timedelta(days=days)
 
-    records: list[dict[str, Any]] = []
+    records: list[tuple[datetime, str, dict[str, Any]]] = []
 
     for path in sorted(glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True)):
         try:
@@ -246,21 +281,17 @@ def load_records(
         if baseline_eligible_only and not is_baseline_eligible_record(data):
             continue
 
-        if cutoff is not None:
-            raw_ts = data.get("timestamp")
-            if raw_ts:
-                try:
-                    ts = datetime.fromisoformat(raw_ts)
-                    if ts.tzinfo is None:
-                        ts = ts.replace(tzinfo=timezone.utc)
-                    if ts < cutoff:
-                        continue
-                except ValueError:
-                    pass  # keep records with unparsable timestamps
+        ts = _parse_record_timestamp(data)
+        if cutoff is not None and ts is not None and ts < cutoff:
+            continue
 
-        records.append(data)
+        records.append((
+            ts or datetime.min.replace(tzinfo=timezone.utc),
+            path,
+            data,
+        ))
 
-    return records
+    return [data for _ts, _path, data in sorted(records)]
 
 
 def load_records_for_model(
@@ -332,6 +363,36 @@ def load_records_for_model(
     return records
 
 
+def load_records_for_identity(
+    local_dir: str,
+    identity_filters: dict[str, str],
+    *,
+    last_n: int | None = None,
+    successful_only: bool = True,
+    baseline_eligible_only: bool = False,
+) -> list[dict[str, Any]]:
+    """Return records matching comparable identity fields.
+
+    V2 performance comparison is intentionally independent from the legacy
+    ``model_id`` directory and ``gpu_type`` display string. The comparable
+    identity filters usually contain the full v2 comparison key, and may also
+    contain a subset when looking for same-cohort recipe mismatches.
+    """
+    records = load_records(
+        local_dir,
+        successful_only=successful_only,
+        baseline_eligible_only=baseline_eligible_only,
+    )
+
+    for key, expected in identity_filters.items():
+        records = [r for r in records if str(r.get(key)) == str(expected)]
+
+    if last_n is not None:
+        records = records[-last_n:]
+
+    return records
+
+
 # ---------------------------------------------------------------------------
 # DataFrame helpers (dashboard / analytics consumers)
 # ---------------------------------------------------------------------------

fastvideo/performance_dashboard/api.py

@@ -23,6 +23,10 @@ FRONTEND_DIST = os.path.abspath(
     os.path.join(os.path.dirname(__file__), "..", "..", "performance_dashboard", "frontend", "dist"))
 
 
+def _matches_display_filters(record: dict[str, Any], model_id: str | None, gpu_type: str | None) -> bool:
+    return (not model_id or record.get("model_id") == model_id) and (not gpu_type or record.get("gpu_type") == gpu_type)
+
+
 class PerformanceDataStore:
 
     def __init__(self, tracking_root: str | None = None) -> None:
@@ -147,11 +151,10 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         # query is kept only so the frontend can share filter state across
         # endpoints without affecting the summary semantics.
         loaded = data_store.load_records(days=None)
-        filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type)
-        rows = build_latest_summary(
-            filtered,
-            run_source=run_source,
-        )
+        rows = [
+            row for row in build_latest_summary(loaded, run_source=run_source)
+            if _matches_display_filters(row, model_id, gpu_type)
+        ]
         return {
             "rows": rows,
             "count": len(rows),
@@ -177,8 +180,10 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         run_source: str | None = None,
     ) -> dict[str, Any]:
         loaded = data_store.load_records(days=days)
-        filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type, run_source=run_source)
-        groups = build_trends(filtered)
+        source_filtered = filter_records(loaded, run_source=run_source)
+        groups = [
+            group for group in build_trends(source_filtered) if _matches_display_filters(group, model_id, gpu_type)
+        ]
         return {
             "groups": groups,
             "count": len(groups),

fastvideo/performance_dashboard/service.py

@@ -17,7 +17,7 @@ from fastvideo.performance.hf_store import is_baseline_eligible_record, safe_flo
 from fastvideo.performance.metric_policy import regression_delta, resolve_metric_policies
 
 Record = dict[str, Any]
-CohortKey = tuple[str, str, str, str, str, str, str, str]
+CohortKey = tuple[str, ...]
 COMPARISON_COHORT_KEYS = (
     "workload_id",
     "variant_id",
@@ -102,16 +102,18 @@ def record_comparison_metadata(record: Record) -> Record:
     return {key: _cohort_metadata_value(record.get(key)) for key in COMPARISON_COHORT_KEYS}
 
 
+def _record_has_complete_v2_identity(record: Record) -> bool:
+    return all(_cohort_metadata_value(record.get(key)) != "" for key in COMPARISON_COHORT_KEYS)
+
+
 def comparison_cohort_key(record: Record) -> CohortKey:
+    if _record_has_complete_v2_identity(record):
+        return ("v2", *(_cohort_key_value(record.get(key)) for key in COMPARISON_COHORT_KEYS))
     return (
+        "legacy",
         str(record.get("model_id") or "unknown"),
         str(record.get("gpu_type") or "unknown"),
-        _cohort_key_value(record.get("workload_id")),
-        _cohort_key_value(record.get("variant_id")),
-        _cohort_key_value(record.get("benchmark_version")),
-        _cohort_key_value(record.get("recipe_fingerprint")),
-        _cohort_key_value(record.get("hardware_profile_id")),
-        _cohort_key_value(record.get("software_profile_id")),
+        *(_cohort_key_value(record.get(key)) for key in COMPARISON_COHORT_KEYS),
     )
 
 
@@ -122,20 +124,18 @@ def group_by_comparison_cohort(records: list[Record]) -> dict[CohortKey, list[Re
     return {key: sorted(value, key=record_sort_key) for key, value in groups.items()}
 
 
-def comparison_sort_key(record: Record) -> CohortKey:
-    return comparison_cohort_key(record)
+def comparison_sort_key(record: Record) -> tuple[str, ...]:
+    return (
+        str(record.get("model_id") or "unknown"),
+        str(record.get("gpu_type") or "unknown"),
+        *(_cohort_key_value(record.get(key)) for key in COMPARISON_COHORT_KEYS),
+    )
 
 
 def latest_row_sort_key(row: Record) -> tuple[Any, ...]:
     return (row["status"] != "fail", *comparison_sort_key(row))
 
 
-def group_identity(key: CohortKey) -> tuple[str, str]:
-    model_id = key[0]
-    gpu_type = key[1]
-    return model_id, gpu_type
-
-
 def baseline_value(records: list[Record], metric_key: str) -> float | None:
     values = [safe_float(record.get(metric_key)) for record in records]
     values = [value for value in values if value is not None]
@@ -149,8 +149,7 @@ def build_latest_summary(records: list[Record],
                          baseline_window: int = 5,
                          run_source: str | None = None) -> list[Record]:
     rows: list[Record] = []
-    for key, group in group_by_comparison_cohort(records).items():
-        model_id, gpu_type = group_identity(key)
+    for group in group_by_comparison_cohort(records).values():
         latest_candidates = group
         if run_source:
             latest_candidates = [record for record in group if record_run_source(record) == run_source]
@@ -158,6 +157,8 @@ def build_latest_summary(records: list[Record],
             continue
 
         latest = latest_candidates[-1]
+        model_id = str(latest.get("model_id") or "unknown")
+        gpu_type = str(latest.get("gpu_type") or "unknown")
         latest_index = next(index for index, record in enumerate(group) if record is latest)
         baseline_pool = [
             record for record in group[:latest_index]
@@ -224,8 +225,10 @@ def build_latest_summary(records: list[Record],
 
 def build_trends(records: list[Record]) -> list[Record]:
     trends: list[Record] = []
-    for key, group in group_by_comparison_cohort(records).items():
-        model_id, gpu_type = group_identity(key)
+    for group in group_by_comparison_cohort(records).values():
+        latest = group[-1]
+        model_id = str(latest.get("model_id") or "unknown")
+        gpu_type = str(latest.get("gpu_type") or "unknown")
         points = []
         for record in group:
             metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
@@ -244,7 +247,7 @@ def build_trends(records: list[Record]) -> list[Record]:
         trends.append({
             "model_id": model_id,
             "gpu_type": gpu_type,
-            **record_comparison_metadata(group[-1]),
+            **record_comparison_metadata(latest),
             "points": points,
         })
     return sorted(trends, key=comparison_sort_key)

fastvideo/pipelines/basic/flux/__init__.py

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

fastvideo/pipelines/basic/flux/flux_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages.flux_stages import (
+    FluxConditioningStage,
+    FluxDecodingStage,
+    FluxDenoisingStage,
+    FluxInputValidationStage,
+    FluxLatentPreparationStage,
+    FluxTimestepPreparationStage,
+)
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+
+class FluxPipeline(ComposedPipelineBase):
+    """FLUX.1-dev T2I (Diffusers module layout, packed latents, embedded guidance)."""
+
+    _required_config_modules = [
+        "scheduler",
+        "transformer",
+        "vae",
+        "text_encoder",
+        "text_encoder_2",
+        "tokenizer",
+        "tokenizer_2",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=FluxInputValidationStage())
+
+        self.add_stage(
+            stage_name="text_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=[
+                    self.get_module("text_encoder"),
+                    self.get_module("text_encoder_2"),
+                ],
+                tokenizers=[
+                    self.get_module("tokenizer"),
+                    self.get_module("tokenizer_2"),
+                ],
+            ),
+        )
+
+        self.add_stage(stage_name="flux_conditioning_stage", stage=FluxConditioningStage())
+
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=FluxTimestepPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=FluxLatentPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=FluxDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=FluxDecodingStage(vae=self.get_module("vae")),
+        )
+
+
+EntryClass = FluxPipeline

fastvideo/pipelines/basic/glm_image/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-Image pipeline package."""
+
+from fastvideo.pipelines.basic.glm_image.glm_image_pipeline import (
+    GlmImagePipeline, )
+
+__all__ = ["GlmImagePipeline"]

fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.basic.glm_image.stages import (
+    GlmImageBeforeDenoisingStage,
+    GlmImageConditionEncodingStage,
+    GlmImageDecodingStage,
+    GlmImageDenoisingStage,
+)
+from fastvideo.pipelines.stages import InputValidationStage
+
+logger = init_logger(__name__)
+
+
+class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
+    pipeline_name = "GlmImagePipeline"
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+        "vision_language_encoder",
+        "processor",
+    ]
+
+    _optional_config_modules: list[str] = []
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=1.0)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage(),
+        )
+
+        self.add_stage(
+            stage_name="glm_image_before_denoising_stage",
+            stage=GlmImageBeforeDenoisingStage(
+                vae=self.get_module("vae"),
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+                processor=self.get_module("processor"),
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vision_language_encoder=self.get_module("vision_language_encoder"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="glm_image_condition_encoding_stage",
+            stage=GlmImageConditionEncodingStage(
+                vae=self.get_module("vae"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=GlmImageDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=GlmImageDecodingStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+
+EntryClass = GlmImagePipeline

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

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-Image pipeline stages."""
+
+from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (GlmImageBeforeDenoisingStage)
+from fastvideo.pipelines.basic.glm_image.stages.condition_encoding import (GlmImageConditionEncodingStage)
+from fastvideo.pipelines.basic.glm_image.stages.decoding import (GlmImageDecodingStage)
+from fastvideo.pipelines.basic.glm_image.stages.denoising import (GlmImageDenoisingStage)
+
+__all__ = [
+    "GlmImageBeforeDenoisingStage",
+    "GlmImageConditionEncodingStage",
+    "GlmImageDecodingStage",
+    "GlmImageDenoisingStage",
+]

fastvideo/pipelines/basic/glm_image/stages/before_denoising.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import re
+from math import sqrt
+
+import numpy as np
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+logger = init_logger(__name__)
+
+
+def calculate_shift(
+    image_seq_len: int,
+    base_seq_len: int = 256,
+    base_shift: float = 0.25,
+    max_shift: float = 0.75,
+) -> float:
+    return (image_seq_len / base_seq_len)**0.5 * max_shift + base_shift
+
+
+def get_glyph_texts(prompt: str | list[str]) -> list[str] | list[list[str]]:
+    if isinstance(prompt, str):
+        prompts: list[str] = [prompt]
+        is_batch = False
+    else:
+        prompts = prompt
+        is_batch = True
+    out: list[list[str]] = []
+    for p in prompts:
+        out.append(
+            re.findall(r"'([^']*)'", p) + re.findall(r"“([^“”]*)”", p) + re.findall(r'"([^"]*)"', p) +
+            re.findall(r"「([^「」]*)」", p))
+    return out if is_batch else out[0]
+
+
+def compute_glyph_embeds(
+    prompts: list[str],
+    tokenizer,
+    text_encoder,
+    device: torch.device,
+    dtype: torch.dtype,
+    max_sequence_length: int = 2048,
+) -> torch.Tensor:
+    all_glyph_texts = get_glyph_texts(prompts)
+    all_glyph_embeds = []
+    for glyph_texts in all_glyph_texts:
+        if len(glyph_texts) == 0:
+            glyph_texts = [""]
+        input_ids = tokenizer(
+            glyph_texts,
+            max_length=max_sequence_length,
+            truncation=True,
+        ).input_ids
+        input_ids = [[tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids for ids in input_ids]
+        max_length = max(len(ids) for ids in input_ids)
+        attention_mask = torch.tensor(
+            [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        input_ids_t = torch.tensor(
+            [ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        outputs = text_encoder(input_ids_t, attention_mask=attention_mask)
+        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
+        all_glyph_embeds.append(glyph_embeds)
+
+    max_seq_len = max(emb.size(1) for emb in all_glyph_embeds)
+    padded = []
+    for emb in all_glyph_embeds:
+        if emb.size(1) < max_seq_len:
+            pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype)
+            emb = torch.cat([pad, emb], dim=1)
+        padded.append(emb)
+    return torch.cat(padded, dim=0).to(device=device, dtype=dtype)
+
+
+def _grid_dims(height: int, width: int) -> tuple[int, int, int, int]:
+    th, tw = height // 32, width // 32
+    ratio = th / tw
+    pth = int(sqrt(ratio) * 16)
+    ptw = int(sqrt(1 / ratio) * 16)
+    return th, tw, pth, ptw
+
+
+def _upsample_d32_to_d16(tokens: torch.Tensor, th: int, tw: int) -> torch.Tensor:
+    tokens = tokens.view(1, 1, th, tw).float()
+    tokens = torch.nn.functional.interpolate(tokens, scale_factor=2, mode="nearest").long()
+    return tokens.view(1, -1)
+
+
+class GlmImageBeforeDenoisingStage(PipelineStage):
+
+    def __init__(self,
+                 vae,
+                 text_encoder,
+                 tokenizer,
+                 processor,
+                 transformer,
+                 scheduler,
+                 vision_language_encoder=None) -> None:
+        super().__init__()
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.processor = processor
+        self.transformer = transformer
+        self.scheduler = scheduler
+        if isinstance(vision_language_encoder, tuple):
+            self.vision_language_encoder, self.vl_processor = (vision_language_encoder[0], vision_language_encoder[1]
+                                                               or processor)
+        else:
+            self.vision_language_encoder = vision_language_encoder
+            self.vl_processor = processor
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        th, tw, pth, ptw = _grid_dims(batch.height, batch.width)
+
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+
+        # 1-3. AR token generation. I2I prepends the condition image and uses a
+        # single-scale target grid; T2I is multi-scale.
+        is_t2i = batch.pil_image is None
+        if self.vision_language_encoder is not None:
+            content = [{"type": "text", "text": batch.prompt}]
+            if not is_t2i:
+                content.insert(0, {"type": "image", "image": batch.pil_image})
+            messages = [{"role": "user", "content": content}]
+            inputs = self.vl_processor.apply_chat_template(messages,
+                                                           tokenize=True,
+                                                           target_h=batch.height,
+                                                           target_w=batch.width,
+                                                           return_dict=True,
+                                                           return_tensors="pt").to(device)
+
+            if is_t2i:
+                up_h, up_w = th, tw
+                large_start, large_count = pth * ptw, th * tw
+                max_new = large_count + (pth * ptw) + 1
+            else:
+                # Condition grid(s) first, target grid last.
+                _, t_h, t_w = inputs["image_grid_thw"][-1].tolist()
+                up_h, up_w = int(t_h), int(t_w)
+                large_start, large_count = 0, up_h * up_w
+                max_new = large_count + 1
+
+            outputs = self.vision_language_encoder.generate(**inputs, max_new_tokens=max_new, do_sample=True)
+            gen_tokens = outputs[0][inputs.input_ids.shape[-1]:]
+            if gen_tokens.shape[0] >= large_start + large_count:
+                large_tokens = gen_tokens[large_start:large_start + large_count]
+            else:
+                available = gen_tokens[large_start:]
+                large_tokens = torch.zeros(large_count, dtype=gen_tokens.dtype, device=gen_tokens.device)
+                if available.shape[0] > 0:
+                    large_tokens[:min(available.shape[0], large_count)] = available[:large_count]
+                logger.warning("AR generated %d tokens, expected %d. Padding with zeros.", gen_tokens.shape[0],
+                               large_start + large_count)
+            batch.prior_token_id = _upsample_d32_to_d16(large_tokens, up_h, up_w)
+            batch.prior_token_drop = torch.zeros(batch.prior_token_id.shape, dtype=torch.bool, device=device)
+
+            if not is_t2i:
+                self._compute_source_prior_tokens(batch, inputs)
+        else:
+            num_prior_tokens = 4 * th * tw
+            logger.warning("No vision_language_encoder provided; using random dropped priors.")
+            batch.prior_token_id = torch.randint(0, 16384, (1, num_prior_tokens), device=device)
+            batch.prior_token_drop = torch.ones(batch.prior_token_id.shape, dtype=torch.bool, device=device)
+
+        # 4. Glyph T5 encoding.
+        prompts = [batch.prompt] if isinstance(batch.prompt, str) else list(batch.prompt)
+        prompt_embeds = compute_glyph_embeds(prompts, self.tokenizer, self.text_encoder, device, dtype)
+
+        # 5. CFG-side negative encoding.
+        if batch.do_classifier_free_guidance:
+            neg_prompts = [batch.negative_prompt or ""] * len(prompts)
+            neg_embeds = compute_glyph_embeds(neg_prompts, self.tokenizer, self.text_encoder, device, dtype)
+            L_pos, L_neg = prompt_embeds.shape[1], neg_embeds.shape[1]
+            max_L = max(L_pos, L_neg)
+            if L_pos < max_L:
+                pad = torch.zeros(prompt_embeds.shape[0],
+                                  max_L - L_pos,
+                                  prompt_embeds.shape[2],
+                                  device=device,
+                                  dtype=dtype)
+                prompt_embeds = torch.cat([pad, prompt_embeds], dim=1)
+            if L_neg < max_L:
+                pad = torch.zeros(neg_embeds.shape[0], max_L - L_neg, neg_embeds.shape[2], device=device, dtype=dtype)
+                neg_embeds = torch.cat([pad, neg_embeds], dim=1)
+            # Row 0 conditional (positive), row 1 unconditional (negative).
+            prompt_embeds = torch.cat([prompt_embeds, neg_embeds], dim=0)
+            att_pos = torch.ones((1, max_L), device=device)
+            att_neg = torch.ones((1, max_L), device=device)
+            if L_pos < max_L:
+                att_pos[:, :max_L - L_pos] = 0
+            if L_neg < max_L:
+                att_neg[:, :max_L - L_neg] = 0
+            attention_mask = torch.cat([att_pos, att_neg], dim=0)
+        else:
+            attention_mask = torch.ones((1, prompt_embeds.shape[1]), device=device)
+
+        batch.prompt_embeds = [prompt_embeds]
+        batch.attention_mask = attention_mask
+
+        # 6. Latents + dynamic flow shift.
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+        batch.latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8), device=device, dtype=dtype)
+
+        # Integer-cast linspace timesteps with resolution-dependent shift applied to
+        # sigmas only; the DiT is conditioned on the unshifted integer timesteps.
+        ntt = self.scheduler.config.num_train_timesteps
+        patch_size = self.transformer.patch_size
+        image_seq_len = ((batch.height // 8) * (batch.width // 8)) // (patch_size**2)
+        sched_timesteps = np.linspace(ntt, 1.0, batch.num_inference_steps + 1)[:-1].astype(np.int64).astype(np.float32)
+        sched_sigmas = sched_timesteps / ntt
+        self.scheduler.set_shift(calculate_shift(image_seq_len))
+        self.scheduler.set_timesteps(batch.num_inference_steps,
+                                     device=device,
+                                     sigmas=sched_sigmas.tolist(),
+                                     timesteps=sched_timesteps.tolist())
+        batch.timesteps = self.scheduler.timesteps
+        return batch
+
+    @torch.no_grad()
+    def _compute_source_prior_tokens(self, batch: ForwardBatch, inputs) -> None:
+        image_grid_thw = inputs["image_grid_thw"]
+        num_condition_images = image_grid_thw.shape[0] - 1
+        source_grids = image_grid_thw[:num_condition_images]
+        image_features = self.vision_language_encoder.get_image_features(inputs["pixel_values"], source_grids)
+        image_feature_parts = getattr(image_features, "pooler_output", image_features)
+        embed = torch.cat(image_feature_parts, dim=0)
+        src_ids_d32 = self.vision_language_encoder.get_image_tokens(embed, source_grids)
+        split_sizes = source_grids.prod(dim=-1).tolist()
+        upsampled = [
+            _upsample_d32_to_d16(ids, int(grid[1]), int(grid[2])).squeeze(0)
+            for ids, grid in zip(torch.split(src_ids_d32, split_sizes), source_grids, strict=False)
+        ]
+        src_grids_up = source_grids.clone()
+        src_grids_up[:, 1] *= 2
+        src_grids_up[:, 2] *= 2
+        batch.extra["glm_prior_token_image_ids"] = torch.cat(upsampled, dim=0)
+        batch.extra["glm_source_image_grid_thw"] = src_grids_up
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_not_empty)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prior_token_id", batch.prior_token_id, V.is_tensor)
+        return result

fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.image_processor import ImageProcessor
+from fastvideo.models.dits.glm_image import GlmImageKVCache
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+
+_CONDITION_MULTIPLE_OF = 16  # vae_scale_factor (8) * DiT patch_size (2)
+
+
+class GlmImageConditionEncodingStage(PipelineStage):
+
+    def __init__(self, vae, transformer) -> None:
+        super().__init__()
+        self.vae = vae
+        self.transformer = transformer
+        self.image_processor = ImageProcessor(vae_scale_factor=_CONDITION_MULTIPLE_OF)
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.pil_image is None:
+            return batch
+
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        self.vae.to(device)
+
+        prior_ids = batch.extra["glm_prior_token_image_ids"].to(device)
+        if prior_ids.dim() == 1:
+            prior_ids = prior_ids.unsqueeze(0)
+
+        # Latent patch count must match the source prior tokens; mismatch is fatal.
+        src_grid = batch.extra["glm_source_image_grid_thw"][0]
+        cond_h = int(src_grid[1]) * _CONDITION_MULTIPLE_OF
+        cond_w = int(src_grid[2]) * _CONDITION_MULTIPLE_OF
+        cond_img = self.image_processor.preprocess(batch.pil_image, cond_h, cond_w).to(device=device,
+                                                                                       dtype=torch.float32)
+        latent = self.vae.encode(cond_img).latent_dist.mode()
+        # NOTE: at runtime self.vae.config is a diffusers FrozenDict with flat
+        # latents_mean/latents_std fields. Access the flat fields directly.
+        cfg = self.vae.config
+        mean = torch.tensor(cfg.latents_mean, device=device, dtype=torch.float32).view(1, -1, 1, 1)
+        std = torch.tensor(cfg.latents_std, device=device, dtype=torch.float32).view(1, -1, 1, 1)
+        latent = ((latent - mean) / std).to(dtype)
+
+        kv_caches = GlmImageKVCache(num_layers=self.transformer.num_layers)
+        empty_text = batch.prompt_embeds[0][:1, :0, :].to(device=device, dtype=dtype)
+        with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=batch):
+            self.transformer(
+                hidden_states=latent,
+                encoder_hidden_states=empty_text,
+                prior_token_id=prior_ids,
+                prior_token_drop=torch.zeros((prior_ids.shape[0], ), dtype=torch.bool, device=device),
+                timestep=torch.zeros((1, ), device=device),
+                target_size=torch.tensor([tuple(cond_img.shape[-2:])], device=device, dtype=torch.long),
+                crop_coords=torch.zeros((1, 2), device=device, dtype=torch.long),
+                kv_caches=kv_caches,
+                kv_caches_mode="write",
+            )
+        batch.extra["glm_kv_caches"] = kv_caches
+        return batch

fastvideo/pipelines/basic/glm_image/stages/decoding.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.stages.decoding import DecodingStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+
+class GlmImageDecodingStage(DecodingStage):
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast = (vae_dtype != torch.float32 and not fastvideo_args.disable_autocast)
+
+        latents = self._denormalize_latents(latents)
+        if latents.dim() == 5:
+            latents = latents.squeeze(2)
+
+        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            if not vae_autocast:
+                latents = latents.to(vae_dtype)
+            decoded = self.vae.decode(latents)
+
+        image = decoded.sample if hasattr(decoded, "sample") else decoded
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image.unsqueeze(2)

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

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CFG convention (both denoise paths): row 0 conditional (positive), row 1 unconditional."""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.attention.backends.sdpa import SDPAMetadata
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class GlmImageDenoisingStage(DenoisingStage):
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
+        latents = getattr(batch, "latent", getattr(batch, "latents", None))
+        result.add_check("latents", latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        return result
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        guidance_scale = batch.guidance_scale
+        do_cfg = guidance_scale > 1.0
+
+        latents = getattr(batch, "latent", getattr(batch, "latents", None))
+        if latents is None:
+            raise ValueError("No latents found in batch.")
+        if latents.dim() == 5:
+            latents = latents.squeeze(2)
+
+        prompt_embeds = batch.prompt_embeds[0]
+        text_attention_mask = getattr(batch, "attention_mask", None)
+        timesteps = batch.timesteps
+
+        patch_size = self.transformer.patch_size
+        _, _, h, w = latents.shape
+        image_seq_length = (h // patch_size) * (w // patch_size)
+        text_seq_length = prompt_embeds.shape[1] if prompt_embeds.dim() >= 2 else 0
+
+        first_block = self.transformer.transformer_blocks[0]
+        backend = getattr(first_block.attn1.attn, "backend", None)
+        sdpa = backend == AttentionBackendEnum.TORCH_SDPA and text_attention_mask is not None
+
+        kv_caches = batch.extra.get("glm_kv_caches")
+        if kv_caches is None:
+            self._denoise_t2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                              text_seq_length, image_seq_length, sdpa, device, dtype)
+        else:
+            self._denoise_i2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                              text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype)
+        return batch
+
+    def _denoise_t2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                     text_seq_length, image_seq_length, sdpa, device, dtype) -> None:
+        num_inference_steps = batch.num_inference_steps
+        bs = 2 if do_cfg else 1
+        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long).repeat(bs, 1)
+        crop_coords = torch.zeros((bs, 2), device=device, dtype=torch.long)
+
+        prior_token_id = batch.prior_token_id
+        if do_cfg and prior_token_id.shape[0] == 1:
+            prior_token_id = prior_token_id.repeat(2, 1)
+        if do_cfg:
+            prior_token_drop = torch.tensor([False, True], device=device)
+        else:
+            prior_token_drop = getattr(batch, "prior_token_drop", torch.tensor([False], device=device))
+
+        attention_mask_kv = None
+        if sdpa:
+            if (text_attention_mask.shape[0] == 1 and bs > 1):
+                text_attention_mask = text_attention_mask.repeat(bs, 1)
+            mix_attn_mask = torch.ones((bs, text_seq_length + image_seq_length), device=device, dtype=torch.float32)
+            mix_attn_mask[:, :text_seq_length] = (text_attention_mask.float().to(device))
+            attention_mask_kv = (mix_attn_mask > 0).unsqueeze(1).unsqueeze(2)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents] * 2) if do_cfg else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t).to(dtype)
+                t_expand = t.expand(latent_model_input.shape[0]) - 1
+
+                attn_metadata = (SDPAMetadata(current_timestep=i, attn_mask=attention_mask_kv)
+                                 if attention_mask_kv is not None else None)
+
+                with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                          attn_metadata=attn_metadata,
+                                                          forward_batch=batch):
+                    noise_pred = self.transformer(
+                        latent_model_input,
+                        prompt_embeds,
+                        prior_token_id,
+                        prior_token_drop,
+                        t_expand,
+                        target_size,
+                        crop_coords,
+                    )
+
+                    if do_cfg:
+                        noise_pred_cond, noise_pred_uncond = noise_pred.float().chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                        guidance_rescale = getattr(batch, "guidance_rescale", 0.0)
+                        if guidance_rescale > 0.0:
+                            dims = list(range(1, noise_pred_cond.ndim))
+                            std_text = noise_pred_cond.std(dim=dims, keepdim=True)
+                            std_cfg = noise_pred.std(dim=dims, keepdim=True)
+                            rescaled = noise_pred * (std_text / std_cfg)
+                            noise_pred = (guidance_rescale * rescaled + (1 - guidance_rescale) * noise_pred)
+
+                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        batch.latents = latents.unsqueeze(2)
+
+    def _denoise_i2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                     text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype) -> None:
+        """Two separate transformer calls (cond reads the cache, uncond skips it):
+        the cache mode is one global flag with batch-1 k/v, so a 2-row CFG call
+        cannot express both."""
+        num_inference_steps = batch.num_inference_steps
+        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long)
+        crop_coords = torch.zeros((1, 2), device=device, dtype=torch.long)
+        prior_token_id = batch.prior_token_id[:1]
+        drop_keep = torch.zeros((1, ), dtype=torch.bool, device=device)
+        drop_all = torch.ones((1, ), dtype=torch.bool, device=device)
+        cache_len = kv_caches[0].k_cache.shape[1] if kv_caches[0].k_cache is not None else 0
+
+        def _mask(row: int, with_cache: bool):
+            if not sdpa:
+                return None
+            prefix = cache_len if with_cache else 0
+            m = torch.ones((1, prefix + text_seq_length + image_seq_length), device=device, dtype=torch.float32)
+            m[:, prefix:prefix + text_seq_length] = text_attention_mask[row:row + 1].float().to(device)
+            return (m > 0).unsqueeze(1).unsqueeze(2)
+
+        cond_mask = _mask(0, with_cache=True)
+        uncond_mask = _mask(1, with_cache=False) if do_cfg else None
+        pos_embeds = prompt_embeds[:1]
+        neg_embeds = prompt_embeds[1:2] if do_cfg else None
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = self.scheduler.scale_model_input(latents, t).to(dtype)
+                t_expand = t.expand(1) - 1
+
+                cond_meta = (SDPAMetadata(current_timestep=i, attn_mask=cond_mask) if cond_mask is not None else None)
+                with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                          attn_metadata=cond_meta,
+                                                          forward_batch=batch):
+                    noise_pred = self.transformer(latent_model_input,
+                                                  pos_embeds,
+                                                  prior_token_id,
+                                                  drop_keep,
+                                                  t_expand,
+                                                  target_size,
+                                                  crop_coords,
+                                                  kv_caches=kv_caches,
+                                                  kv_caches_mode="read")
+
+                if do_cfg:
+                    uncond_meta = (SDPAMetadata(current_timestep=i, attn_mask=uncond_mask)
+                                   if uncond_mask is not None else None)
+                    with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                              attn_metadata=uncond_meta,
+                                                              forward_batch=batch):
+                        noise_pred_uncond = self.transformer(latent_model_input,
+                                                             neg_embeds,
+                                                             prior_token_id,
+                                                             drop_all,
+                                                             t_expand,
+                                                             target_size,
+                                                             crop_coords,
+                                                             kv_caches=kv_caches,
+                                                             kv_caches_mode="skip")
+                    noise_pred = noise_pred_uncond.float() + guidance_scale * (noise_pred.float() -
+                                                                               noise_pred_uncond.float())
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        kv_caches.clear()
+        batch.latents = latents.unsqueeze(2)

fastvideo/pipelines/pipeline_batch_info.py

@@ -109,6 +109,10 @@ class ForwardBatch:
     max_sequence_length: int | None = None
     prompt_template: dict[str, Any] | None = None
     do_classifier_free_guidance: bool = False
+    # When True, ``guidance_scale`` is passed into models that use embedded guidance (e.g. FLUX)
+    # and must not imply classic dual-forward CFG. Use ``true_cfg_scale > 1`` for true CFG.
+    use_embedded_guidance: bool = False
+    true_cfg_scale: float = 1.0
 
     # Batch info
     batch_size: int | None = None
@@ -252,9 +256,12 @@ class ForwardBatch:
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 
-        # Enable CFG for standard guidance_scale and LTX-2 text CFG scales.
+        # LTX-2 text CFG scales; FLUX uses ``use_embedded_guidance`` so ``guidance_scale > 1`` alone
+        # does not enable classifier-free guidance.
         ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0)
-        if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
+        if self.use_embedded_guidance:
+            self.do_classifier_free_guidance = (self.true_cfg_scale > 1.0) or ltx2_text_cfg_enabled
+        elif self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
             self.do_classifier_free_guidance = True
         if self.negative_prompt_embeds is None:
             self.negative_prompt_embeds = []

fastvideo/pipelines/stages/flux_stages.py

@@ -0,0 +1,424 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+def _pack_latents(
+    latents: torch.Tensor,
+    batch_size: int,
+    num_channels_latents: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Diffusers ``_pack_latents`` for FLUX (2×2 spatial pack in latent space)."""
+    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    return latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+
+def _unpack_latents(
+    latents: torch.Tensor,
+    batch_size: int,
+    num_channels_latents: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Inverse of ``_pack_latents``."""
+    latents = latents.reshape(batch_size, height // 2, width // 2, num_channels_latents, 2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+    return latents.reshape(batch_size, num_channels_latents, height, width)
+
+
+def _prepare_latent_image_ids(
+    patch_height: int,
+    patch_width: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.long,
+) -> torch.Tensor:
+    """Match Diffusers ``FluxPipeline._prepare_latent_image_ids`` (no batch dim)."""
+    latent_image_ids = torch.zeros(patch_height, patch_width, 3, device=device, dtype=torch.float32)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(patch_height, device=device)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(patch_width, device=device)[None, :]
+    h, w, c = latent_image_ids.shape
+    return latent_image_ids.reshape(h * w, c).to(dtype=dtype)
+
+
+class FluxInputValidationStage(InputValidationStage):
+    """Require height/width divisible by 16 (VAE scale × 2 for FLUX packing)."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if (batch.height is not None and batch.width is not None and (batch.height % 16 != 0 or batch.width % 16 != 0)):
+            raise ValueError("FLUX expects height and width divisible by 16 "
+                             f"(VAE latent grid × 2× packing); got {batch.height}×{batch.width}.")
+        return super().forward(batch, fastvideo_args)
+
+
+class FluxConditioningStage(PipelineStage):
+    """Build CLIP pooled + T5 sequence + ``text_ids`` (and optional negative for true CFG)."""
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if len(batch.prompt_embeds) < 2:
+            raise ValueError("FluxConditioningStage expects 2 prompt_embeds (CLIP pooled, T5 sequence), "
+                             f"got {len(batch.prompt_embeds)}")
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+
+        pooled = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
+        enc = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
+        seq_len = enc.shape[1]
+        text_ids = torch.zeros(seq_len, 3, device=device, dtype=torch.long)
+
+        batch.extra["flux_pooled_projections"] = pooled
+        batch.extra["flux_encoder_hidden_states"] = enc
+        batch.extra["flux_text_ids"] = text_ids
+
+        if batch.do_classifier_free_guidance:
+            if not batch.negative_prompt_embeds or len(batch.negative_prompt_embeds) < 2:
+                raise ValueError("True CFG requires two negative_prompt_embeds (CLIP, T5).")
+            neg_pooled = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
+            neg_enc = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
+            batch.extra["flux_negative_pooled_projections"] = neg_pooled
+            batch.extra["flux_negative_encoder_hidden_states"] = neg_enc
+
+        return batch
+
+
+class FluxTimestepPreparationStage(TimestepPreparationStage):
+    """Flow Match with resolution-dependent ``mu`` from packed image sequence length."""
+
+    @staticmethod
+    def _calculate_mu(
+        image_seq_len: int,
+        base_seq_len: int = 256,
+        max_seq_len: int = 4096,
+        base_shift: float = 0.5,
+        max_shift: float = 1.15,
+    ) -> float:
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        return float(image_seq_len) * m + b
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        sig = inspect.signature(self.scheduler.set_timesteps)
+        if "mu" not in sig.parameters:
+            logger.warning(
+                "FLUX timestep prep: scheduler %s.set_timesteps does not accept 'mu'; falling back to the base "
+                "timestep schedule. FLUX expects a FlowMatchEulerDiscreteScheduler with resolution-dependent "
+                "dynamic shifting — output quality may degrade.",
+                type(self.scheduler).__name__)
+            return super().forward(batch, fastvideo_args)
+
+        cfg = getattr(self.scheduler, "config", None)
+        use_dynamic = bool(getattr(cfg, "use_dynamic_shifting", False)) if cfg is not None else False
+        if not use_dynamic:
+            logger.warning(
+                "FLUX timestep prep: scheduler has use_dynamic_shifting=False; falling back to the base timestep "
+                "schedule and skipping the resolution-dependent 'mu' shift. FLUX requires dynamic shifting for "
+                "correct timesteps — output quality may degrade.")
+            return super().forward(batch, fastvideo_args)
+
+        if batch.height is None or batch.width is None:
+            raise ValueError("height/width must be set before FluxTimestepPreparationStage")
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
+        h_lat = batch.height // spatial_ratio
+        w_lat = batch.width // spatial_ratio
+        if h_lat % 2 != 0 or w_lat % 2 != 0:
+            raise ValueError(
+                f"Latent spatial dims must be even for FLUX packing; got {h_lat}×{w_lat} from {batch.height}×{batch.width}."
+            )
+        image_seq_len = (h_lat // 2) * (w_lat // 2)
+
+        base_seq_len = int(getattr(cfg, "base_image_seq_len", 256))
+        max_seq_len = int(getattr(cfg, "max_image_seq_len", 4096))
+        base_shift = float(getattr(cfg, "base_shift", 0.5))
+        max_shift = float(getattr(cfg, "max_shift", 1.15))
+
+        device = get_local_torch_device()
+        mu = self._calculate_mu(
+            image_seq_len=image_seq_len,
+            base_seq_len=base_seq_len,
+            max_seq_len=max_seq_len,
+            base_shift=base_shift,
+            max_shift=max_shift,
+        )
+        self.scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
+        batch.timesteps = self.scheduler.timesteps
+        return batch
+
+
+class FluxLatentPreparationStage(PipelineStage):
+
+    def __init__(self, scheduler) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.height is None or batch.width is None:
+            raise ValueError("height/width required for FluxLatentPreparationStage")
+
+        if isinstance(batch.prompt, list):
+            batch_size = len(batch.prompt)
+        elif batch.prompt is not None:
+            batch_size = 1
+        else:
+            if not batch.prompt_embeds:
+                raise ValueError("prompt or prompt_embeds must be provided")
+            batch_size = batch.prompt_embeds[0].shape[0]
+
+        batch_size *= batch.num_videos_per_prompt
+
+        if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
+            raise ValueError(f"generator list length {len(batch.generator)} does not match batch_size {batch_size}")
+
+        arch = fastvideo_args.pipeline_config.dit_config.arch_config
+        in_channels = int(getattr(arch, "in_channels", 64))
+        num_channels_latents = in_channels // 4
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
+
+        h_lat = batch.height // spatial_ratio
+        w_lat = batch.width // spatial_ratio
+
+        dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        device = get_local_torch_device()
+
+        shape = (batch_size, num_channels_latents, h_lat, w_lat)
+        latents = batch.latents
+        if latents is None:
+            latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+            if latents.shape != shape:
+                raise ValueError(f"Expected latents shape {shape}, got {tuple(latents.shape)}")
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+
+        packed = _pack_latents(latents, batch_size, num_channels_latents, h_lat, w_lat)
+
+        patch_h, patch_w = h_lat // 2, w_lat // 2
+        img_ids = _prepare_latent_image_ids(patch_h, patch_w, device, dtype=torch.long)
+
+        batch.latents = packed
+        batch.raw_latent_shape = shape
+        batch.extra["flux_h_lat"] = h_lat
+        batch.extra["flux_w_lat"] = w_lat
+        batch.extra["flux_num_channels_latents"] = num_channels_latents
+        batch.extra["flux_latent_image_ids"] = img_ids
+
+        return batch
+
+
+class FluxDenoisingStage(PipelineStage):
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    @staticmethod
+    def _step_kwargs(scheduler_step, batch: ForwardBatch) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {}
+        sig = inspect.signature(scheduler_step)
+        if "generator" in sig.parameters:
+            gen = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
+            kwargs["generator"] = gen
+        return kwargs
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.timesteps is None:
+            raise ValueError("timesteps must be set before FluxDenoisingStage")
+        if batch.latents is None:
+            raise ValueError("latents must be set before FluxDenoisingStage")
+
+        packed = batch.latents
+        timesteps = batch.timesteps
+
+        pooled = batch.extra["flux_pooled_projections"]
+        enc = batch.extra["flux_encoder_hidden_states"]
+        txt_ids = batch.extra["flux_text_ids"]
+        img_ids = batch.extra["flux_latent_image_ids"]
+
+        neg_pooled = batch.extra.get("flux_negative_pooled_projections")
+        neg_enc = batch.extra.get("flux_negative_encoder_hidden_states")
+
+        true_cfg_scale = float(batch.true_cfg_scale)
+        use_true_cfg = batch.do_classifier_free_guidance and true_cfg_scale > 1.0
+
+        # Prefer the loaded transformer's arch (HF ``guidance_embeds``), not static pipeline defaults.
+        tr_arch = self.transformer.fastvideo_config.arch_config
+        guidance_embeds = bool(getattr(tr_arch, "guidance_embeds", False))
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        bs = packed.shape[0]
+        if guidance_embeds:
+            guidance = torch.full((bs, ), float(batch.guidance_scale), device=device, dtype=torch.float32)
+        else:
+            guidance = None
+
+        step_extras = self._step_kwargs(self.scheduler.step, batch)
+
+        for t in timesteps:
+            t_scalar = t
+            if not isinstance(t_scalar, torch.Tensor):
+                t_scalar = torch.tensor([t_scalar], device=device, dtype=torch.float32)
+            t_scalar = t_scalar.to(device=device, dtype=torch.float32)
+
+            timestep_model = t_scalar.expand(bs).float() / 1000.0
+            timestep_model = timestep_model.to(dtype=target_dtype)
+
+            ts_ctx = int(t_scalar.reshape(-1)[0].item())
+            with (
+                    torch.autocast(
+                        device_type="cuda",
+                        dtype=target_dtype,
+                        enabled=autocast_enabled and device.type == "cuda",
+                    ),
+                    set_forward_context(
+                        current_timestep=ts_ctx,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                    ),
+            ):
+                if use_true_cfg:
+                    assert neg_enc is not None and neg_pooled is not None
+                    n_neg = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=neg_enc,
+                        pooled_projections=neg_pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+                    n_pos = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=enc,
+                        pooled_projections=pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = n_neg + true_cfg_scale * (n_pos - n_neg)
+                else:
+                    noise_pred = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=enc,
+                        pooled_projections=pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+
+            packed = self.scheduler.step(
+                noise_pred,
+                t_scalar,
+                packed,
+                return_dict=False,
+                **step_extras,
+            )[0]
+
+        batch.latents = packed
+        return batch
+
+
+class FluxDecodingStage(PipelineStage):
+    """Unpack latents, apply VAE scaling/shift, decode to pixels (5D output ``B×3×1×H×W``)."""
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    @staticmethod
+    def _denormalize_latents(latents: torch.Tensor, vae: Any) -> torch.Tensor:
+        cfg = getattr(vae, "config", None)
+        sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None
+        sh = getattr(cfg, "shift_factor", None) if cfg is not None else None
+        if sf is None and hasattr(vae, "scaling_factor"):
+            sf = vae.scaling_factor
+        if sh is None and hasattr(vae, "shift_factor"):
+            sh = vae.shift_factor
+        if sf is not None:
+            latents = latents / (sf.to(latents.device, latents.dtype) if isinstance(sf, torch.Tensor) else sf)
+            if sh is not None:
+                latents = latents + (sh.to(latents.device, latents.dtype) if isinstance(sh, torch.Tensor) else sh)
+        return latents
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        packed = batch.latents
+        if packed is None:
+            raise ValueError("latents must be set before FluxDecodingStage")
+
+        h_lat = int(batch.extra["flux_h_lat"])
+        w_lat = int(batch.extra["flux_w_lat"])
+        num_ch = int(batch.extra["flux_num_channels_latents"])
+        raw_shape = batch.raw_latent_shape
+        if raw_shape is None:
+            raise ValueError("raw_latent_shape missing; FluxLatentPreparationStage must run first.")
+        batch_size = int(raw_shape[0])
+
+        infer_device = get_local_torch_device()
+        packed = packed.to(infer_device)
+
+        latents_4d = _unpack_latents(packed, batch_size, num_ch, h_lat, w_lat)
+        latents_4d = self._denormalize_latents(latents_4d, self.vae)
+
+        vae_device = next(self.vae.parameters()).device
+        latents_4d = latents_4d.to(device=vae_device)
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+        use_cuda_autocast = autocast_enabled and vae_device.type == "cuda"
+
+        with torch.autocast(
+                device_type="cuda",
+                dtype=vae_dtype,
+                enabled=use_cuda_autocast,
+        ):
+            if not autocast_enabled:
+                latents_4d = latents_4d.to(dtype=vae_dtype)
+            dec = self.vae.decode(latents_4d)
+            image = dec.sample if hasattr(dec, "sample") else dec[0]
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        batch.output = image.unsqueeze(2).detach().float().cpu()
+        return batch

fastvideo/registry.py

@@ -58,11 +58,14 @@ from fastvideo.configs.pipelines.wan import (
     WanT2V480PConfig,
     WanT2V720PConfig,
 )
+from fastvideo.configs.pipelines.glm_image import GlmImageConfig
+from fastvideo.configs.pipelines.flux import FluxPipelineConfig
 from fastvideo.configs.pipelines.sd35 import SD35Config
 from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
 from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
+from fastvideo.api.flux import FluxSamplingParam
 
 from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.logger import init_logger
@@ -1074,6 +1077,33 @@ def _register_configs() -> None:
         default_preset="sd35_medium",
     )
 
+    # GLM-Image
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=GlmImageConfig,
+        hf_model_paths=[
+            "zai-org/GLM-Image",
+        ],
+        model_detectors=[lambda path: "glmimage" in path.lower() or "glm-image" in path.lower()],
+        workload_types=(WorkloadType.T2I, ),
+        model_family="glm_image",
+    )
+
+    # FLUX.1-dev (Diffusers)
+    register_configs(
+        sampling_param_cls=FluxSamplingParam,
+        pipeline_config_cls=FluxPipelineConfig,
+        workload_types=(WorkloadType.T2I, ),
+        hf_model_paths=[
+            "black-forest-labs/FLUX.1-dev",
+        ],
+        model_detectors=[
+            lambda path: "fluxpipeline" in path,
+            lambda path: "flux.1-dev" in path or "flux_1_dev" in path,
+            lambda path: "/flux/" in path or path.endswith("/flux"),
+        ],
+    )
+
 
 # --- Part 3: Main Resolver ---
 

fastvideo/tests/api/test_parser.py

@@ -183,6 +183,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "guidance_scale_2": None,
                 "guidance_rescale": 0.0,
                 "true_cfg_scale": None,
+                "use_embedded_guidance": None,
                 "boundary_ratio": None,
                 "sigmas": None,
             },

fastvideo/tests/attention/test_flash_attn_default_custom_op.py

@@ -31,23 +31,29 @@ def fa_default_impls():
     if not torch.cuda.is_available():
         pytest.skip("CUDA is required for the FA2/FA3 default custom-op tests")
 
+    # The registration + dispatcher live in `attention/utils/`, alongside the
+    # FP4 cute template and the masked/varlen wrappers. The backend
+    # (`attention/backends/flash_attn.py`) just imports `flash_attn_func_compilable`
+    # from there, so test references the utils module directly.
     try:
-        from fastvideo.attention.backends import flash_attn as fa_backend
+        from fastvideo.attention.utils import flash_attn_default as fa_module
     except ImportError as exc:
-        pytest.skip(f"flash_attn backend not importable: {exc}")
+        pytest.skip(f"flash_attn_default not importable: {exc}")
 
-    if fa_backend.fa_version not in ("2", "3"):
+    if fa_module.fa_version not in ("2", "3"):
         pytest.skip(
             f"FA2/FA3 default custom op only exists for fa_version in (2, 3); "
-            f"got {fa_backend.fa_version!r}"
+            f"got {fa_module.fa_version!r}"
         )
 
-    # compilable dispatcher, the original FA wrapper it falls back to, and
-    # the raw custom op for opcheck.
+    # compilable dispatcher, the original FA wrapper it falls back to, the
+    # raw custom op for opcheck, and the fa_version (FA2 has full register_
+    # autograd; FA3 keeps the carve-out so some tests gate on this).
     return (
-        fa_backend.flash_attn_func_compilable,
-        fa_backend._fa_default,
+        fa_module.flash_attn_func_compilable,
+        fa_module._fa_default,
         torch.ops.fastvideo._flash_attn_default_forward,
+        fa_module.fa_version,
     )
 
 
@@ -71,7 +77,7 @@ def test_default_compilable_inference_matches_original(fa_default_impls, dtype,
     """No-grad path routes through the custom op and is numerically identical."""
     if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
         pytest.skip("bfloat16 is not supported on this GPU")
-    compilable, original, _ = fa_default_impls
+    compilable, original, _, _ = fa_default_impls
 
     torch.manual_seed(0)
     q, k, v = _qkv(dtype, requires_grad=False)
@@ -91,7 +97,7 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
     """
     if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
         pytest.skip("bfloat16 is not supported on this GPU")
-    compilable, original, _ = fa_default_impls
+    compilable, original, _, _ = fa_default_impls
 
     torch.manual_seed(0)
     q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
@@ -115,7 +121,65 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
 @pytest.mark.parametrize("causal", [False, True])
 def test_default_forward_opcheck(fa_default_impls, causal):
     """Schema / fake-kernel consistency for the custom op (forward only)."""
-    _, _, op = fa_default_impls
+    _, _, op, _ = fa_default_impls
     torch.manual_seed(0)
     q, k, v = _qkv(torch.float16, requires_grad=False)
     torch.library.opcheck(op, (q, k, v, None, causal))
+
+
+# --------------------------------------------------------------------------- #
+# FA2-only: backward is registered on the custom op itself. Exercises the     #
+# `register_autograd` wiring directly (not via the dispatcher's carve-out).   #
+# Skipped on FA3 until Kuan-Hao's Modal FA3 setup PR lands and we mirror the  #
+# pattern there.                                                              #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("causal", [False, True])
+def test_default_op_backward_through_registered_autograd(fa_default_impls, dtype, causal):
+    """FA2: gradients flow through ``torch.ops.fastvideo._flash_attn_default_forward``
+    itself (no dispatcher carve-out involved) and match the original
+    ``flash_attn_func``'s gradients."""
+    if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported on this GPU")
+    _, original, op, fa_version = fa_default_impls
+    if fa_version != "2":
+        pytest.skip(f"register_autograd is only wired for FA2 right now; got {fa_version!r}")
+
+    torch.manual_seed(0)
+    q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
+    q_test, k_test, v_test = _clone(q_ref, k_ref, v_ref)
+
+    # Reference grads via the original autograd.Function.
+    out_ref = original(q_ref, k_ref, v_ref, softmax_scale=None, causal=causal)
+    # Custom-op grads via the registered backward — unpack (out, lse), discard lse.
+    out_test, _ = op(q_test, k_test, v_test, None, causal)
+
+    torch.testing.assert_close(out_test, out_ref,
+                               atol=0 if dtype == torch.float16 else 1e-3,
+                               rtol=0 if dtype == torch.float16 else 1e-3)
+
+    dout = torch.randn_like(out_ref)
+    dq_ref, dk_ref, dv_ref = torch.autograd.grad(
+        (out_ref * dout).sum(), (q_ref, k_ref, v_ref))
+    dq_test, dk_test, dv_test = torch.autograd.grad(
+        (out_test * dout).sum(), (q_test, k_test, v_test))
+
+    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
+    torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("causal", [False, True])
+def test_default_op_opcheck_with_grad_inputs(fa_default_impls, causal):
+    """FA2: full ``opcheck`` including ``test_autograd_registration`` —
+    catches a missing/inconsistent backward at unit-test time, which was
+    exactly the gap that #1373's first revision shipped."""
+    _, _, op, fa_version = fa_default_impls
+    if fa_version != "2":
+        pytest.skip(f"autograd registration only wired for FA2; got {fa_version!r}")
+    torch.manual_seed(0)
+    q, k, v = _qkv(torch.float16, requires_grad=True)
+    torch.library.opcheck(op, (q, k, v, None, causal))

fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression guard for the masked/varlen custom ops.
+
+`flash_attn_no_pad_compilable` / `flash_attn_varlen_qk_no_pad_compilable` wrap
+the whole masked-attention functions in `torch.library.custom_op`s so dynamo
+sees one traceable node (the internal unpad/pad bookkeeping runs eager inside).
+
+On FA2 these ops register a real backward — `softmax_lse` is padded back to a
+statically-shaped `[batch, nheads, seqlen]` form on the way out and re-unpadded
+in backward, which calls FA2's `_flash_attn_varlen_backward`. So training also
+backprops through the op (no graph break on the training path either).
+
+On FA3/FA4 these ops are forward+fake only and the `*_compilable` dispatchers
+carve out to the autograd.Function for grad-enabled calls (PR #1373 pattern).
+Tests gating on FA2 are skipped on FA3/FA4.
+
+These tests pin:
+  - inference (no grad): output through the custom op is bit-identical to the
+    original function;
+  - training (requires_grad): gradients through the registered op match the
+    original autograd.Function;
+  - schema/fake-kernel consistency via torch.library.opcheck, both with and
+    without grad-requiring inputs (the latter exercises
+    `test_autograd_registration` — would have caught a missing/broken backward
+    at unit-test time).
+
+GPU assumptions: requires CUDA and the FA2 `flash_attn` varlen package.
+Skips on CPU and when flash_attn is unavailable.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+
+@pytest.fixture(scope="module")
+def no_pad_impls():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the masked/varlen custom-op tests")
+    try:
+        from fastvideo.attention.utils import flash_attn_no_pad as mod
+    except ImportError as exc:
+        pytest.skip(f"flash_attn_no_pad not importable (flash_attn missing?): {exc}")
+    return mod
+
+
+def _dtype_skip(dtype):
+    if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported on this GPU")
+
+
+def _fa2_only(mod):
+    if mod._FA_VARLEN_VERSION != "2":
+        pytest.skip(
+            f"register_autograd is only wired for FA2; got "
+            f"_FA_VARLEN_VERSION={mod._FA_VARLEN_VERSION!r}"
+        )
+
+
+def _padding_mask(batch, seqlen, valid_lens, device):
+    mask = torch.zeros(batch, seqlen, dtype=torch.bool, device=device)
+    for i, n in enumerate(valid_lens):
+        mask[i, :n] = True
+    return mask
+
+
+# --------------------------------------------------------------------------- #
+# flash_attn_no_pad (masked self-attention)                                    #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_no_pad_inference_matches_original(no_pad_impls, dtype):
+    """No-grad path through the custom op is bit-identical to the original."""
+    _dtype_skip(dtype)
+    mod = no_pad_impls
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    b, s, h, d = 2, 64, 4, 64
+    qkv = torch.randn(b, s, 3, h, d, device=device, dtype=dtype)
+    mask = _padding_mask(b, s, [64, 48], device)
+
+    with torch.inference_mode():
+        out_ref = mod.flash_attn_no_pad(qkv, mask, causal=False, dropout_p=0.0)
+        out_test = mod.flash_attn_no_pad_compilable(qkv, mask, causal=False, dropout_p=0.0)
+    torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_no_pad_training_backward_through_registered_autograd(no_pad_impls, dtype):
+    """FA2: grads flow through the registered op and match the original."""
+    _dtype_skip(dtype)
+    mod = no_pad_impls
+    _fa2_only(mod)
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    b, s, h, d = 2, 64, 4, 64
+    mask = _padding_mask(b, s, [64, 48], device)
+
+    qkv_ref = torch.randn(b, s, 3, h, d, device=device, dtype=dtype, requires_grad=True)
+    qkv_test = qkv_ref.detach().clone().requires_grad_(True)
+
+    out_ref = mod.flash_attn_no_pad(qkv_ref, mask, causal=False, dropout_p=0.0)
+    # Go through the compilable wrapper (which on FA2 unconditionally routes
+    # to the registered op — no carve-out).
+    out_test = mod.flash_attn_no_pad_compilable(qkv_test, mask, causal=False, dropout_p=0.0)
+
+    dout = torch.randn_like(out_ref)
+    (dqkv_ref,) = torch.autograd.grad((out_ref * dout).sum(), (qkv_ref,))
+    (dqkv_test,) = torch.autograd.grad((out_test * dout).sum(), (qkv_test,))
+
+    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
+    torch.testing.assert_close(dqkv_test, dqkv_ref, atol=atol, rtol=rtol)
+
+
+def test_no_pad_forward_opcheck(no_pad_impls):
+    """Schema/fake-kernel consistency for the custom op (forward only)."""
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    b, s, h, d = 2, 64, 4, 64
+    qkv = torch.randn(b, s, 3, h, d, device=device, dtype=torch.float16)
+    mask = _padding_mask(b, s, [64, 48], device)
+    torch.library.opcheck(
+        torch.ops.fastvideo._flash_attn_no_pad_forward,
+        (qkv, mask, False, 0.0, None, False),
+    )
+
+
+def test_no_pad_opcheck_with_grad_inputs(no_pad_impls):
+    """FA2: full opcheck including ``test_autograd_registration`` — catches a
+    missing/inconsistent backward at unit-test time."""
+    mod = no_pad_impls
+    _fa2_only(mod)
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    b, s, h, d = 2, 64, 4, 64
+    qkv = torch.randn(b, s, 3, h, d, device=device, dtype=torch.float16, requires_grad=True)
+    mask = _padding_mask(b, s, [64, 48], device)
+    torch.library.opcheck(
+        torch.ops.fastvideo._flash_attn_no_pad_forward,
+        (qkv, mask, False, 0.0, None, False),
+    )
+
+
+# --------------------------------------------------------------------------- #
+# flash_attn_varlen_qk_no_pad (cross-attn / unequal q-k seqlen)                #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_varlen_qk_inference_matches_original(no_pad_impls, dtype):
+    _dtype_skip(dtype)
+    mod = no_pad_impls
+    # The varlen-qk custom op's real forward goes through FA's varlen func; off
+    # FA2 that path is the pre-existing FA3 `dropout_p` carve-out (out of scope
+    # here), so skip cleanly like the autograd tests until that's fixed.
+    _fa2_only(mod)
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    b, sq, sk, h, d = 2, 48, 64, 4, 64
+    q = torch.randn(b, sq, h, d, device=device, dtype=dtype)
+    k = torch.randn(b, sk, h, d, device=device, dtype=dtype)
+    v = torch.randn(b, sk, h, d, device=device, dtype=dtype)
+    qmask = _padding_mask(b, sq, [48, 40], device)
+    kmask = _padding_mask(b, sk, [64, 56], device)
+
+    with torch.inference_mode():
+        out_ref = mod.flash_attn_varlen_qk_no_pad(q, k, v, qmask, kmask, causal=False, dropout_p=0.0)
+        out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q, k, v, qmask, kmask, causal=False, dropout_p=0.0)
+    torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_varlen_qk_training_backward_through_registered_autograd(no_pad_impls, dtype):
+    """FA2: grads through q, k, v all flow via the registered op and match."""
+    _dtype_skip(dtype)
+    mod = no_pad_impls
+    _fa2_only(mod)
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    b, sq, sk, h, d = 2, 48, 64, 4, 64
+    qmask = _padding_mask(b, sq, [48, 40], device)
+    kmask = _padding_mask(b, sk, [64, 56], device)
+
+    q_ref = torch.randn(b, sq, h, d, device=device, dtype=dtype, requires_grad=True)
+    k_ref = torch.randn(b, sk, h, d, device=device, dtype=dtype, requires_grad=True)
+    v_ref = torch.randn(b, sk, h, d, device=device, dtype=dtype, requires_grad=True)
+    q_test = q_ref.detach().clone().requires_grad_(True)
+    k_test = k_ref.detach().clone().requires_grad_(True)
+    v_test = v_ref.detach().clone().requires_grad_(True)
+
+    out_ref = mod.flash_attn_varlen_qk_no_pad(q_ref, k_ref, v_ref, qmask, kmask, causal=False, dropout_p=0.0)
+    out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q_test, k_test, v_test, qmask, kmask,
+                                                          causal=False, dropout_p=0.0)
+
+    dout = torch.randn_like(out_ref)
+    dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref))
+    dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test))
+
+    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
+    torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
+
+
+def test_varlen_qk_forward_opcheck(no_pad_impls):
+    mod = no_pad_impls
+    _fa2_only(mod)
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    b, sq, sk, h, d = 2, 48, 64, 4, 64
+    q = torch.randn(b, sq, h, d, device=device, dtype=torch.float16)
+    k = torch.randn(b, sk, h, d, device=device, dtype=torch.float16)
+    v = torch.randn(b, sk, h, d, device=device, dtype=torch.float16)
+    qmask = _padding_mask(b, sq, [48, 40], device)
+    kmask = _padding_mask(b, sk, [64, 56], device)
+    torch.library.opcheck(
+        torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
+        (q, k, v, qmask, kmask, False, 0.0, None, False),
+    )
+
+
+def test_varlen_qk_opcheck_with_grad_inputs(no_pad_impls):
+    """FA2: full opcheck with requires_grad inputs (autograd-registration check)."""
+    mod = no_pad_impls
+    _fa2_only(mod)
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    b, sq, sk, h, d = 2, 48, 64, 4, 64
+    q = torch.randn(b, sq, h, d, device=device, dtype=torch.float16, requires_grad=True)
+    k = torch.randn(b, sk, h, d, device=device, dtype=torch.float16, requires_grad=True)
+    v = torch.randn(b, sk, h, d, device=device, dtype=torch.float16, requires_grad=True)
+    qmask = _padding_mask(b, sq, [48, 40], device)
+    kmask = _padding_mask(b, sk, [64, 56], device)
+    torch.library.opcheck(
+        torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
+        (q, k, v, qmask, kmask, False, 0.0, None, False),
+    )

fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py

@@ -50,9 +50,10 @@ def test_resolver_skips_cute_without_opt_in(monkeypatch) -> None:
     monkeypatch.setattr(builtins, "__import__", spying_import)
 
     mod = _reload_resolver_module()
-    resolved = mod._resolve_flash_attn_varlen_func()
+    resolved, version = mod._resolve_flash_attn_varlen_func()
     assert resolved is not None
     assert resolved.__name__ == "flash_attn_varlen_func"
+    assert version in ("2", "3")
     assert "fastvideo.attention.utils.flash_attn_cute" not in attempted
 
 
@@ -92,9 +93,10 @@ def test_resolver_returns_flash_attn_when_interface_unavailable(monkeypatch) ->
     monkeypatch.setattr(builtins, "__import__", patched_import)
 
     mod = _reload_resolver_module()
-    resolved = mod._resolve_flash_attn_varlen_func()
+    resolved, version = mod._resolve_flash_attn_varlen_func()
     assert resolved is not None
     assert resolved.__module__.startswith("flash_attn")
+    assert version == "2"
 
 
 def test_module_level_impl_is_resolver_output() -> None:

fastvideo/tests/attention/test_sdpa_head_sizes.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU tests: SDPA must not restrict attention head sizes.
+
+``torch.nn.functional.scaled_dot_product_attention`` is head-size
+agnostic, but ``SDPABackend.get_supported_head_sizes()`` used to declare
+a list copied from FlashAttentionBackend that omitted e.g. 80 (the head
+size of CLIP-L vision encoders such as matrixgame2's). That made
+capability validation flag SDPA as incompatible with head-80 layers even
+though it handles them fine.
+
+CPU-only; no GPU or flash-attn dependency.
+"""
+
+import torch
+
+from fastvideo.attention.backends.sdpa import SDPABackend, SDPAImpl, SDPAMetadata
+from fastvideo.platforms.cuda import CudaPlatformBase
+from fastvideo.platforms.interface import AttentionBackendEnum
+
+
+def test_sdpa_head_sizes_unrestricted() -> None:
+    """SDPA declares no head-size restriction (None), so 80 is allowed."""
+    assert SDPABackend.get_supported_head_sizes() is None
+    # Once the capability-validation API (#1494) lands, head-80 must pass.
+    validate = getattr(SDPABackend, "validate_compatibility", None)
+    if validate is not None:
+        assert validate(head_size=80, dtype=torch.bfloat16) is None
+
+
+def test_explicit_torch_sdpa_pin_head_80_resolves() -> None:
+    """An explicit TORCH_SDPA pin on a head-80 layer resolves to SDPABackend."""
+    cls_str = CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.TORCH_SDPA, head_size=80, dtype=torch.bfloat16)
+    assert cls_str == "fastvideo.attention.backends.sdpa.SDPABackend"
+
+
+def test_sdpa_impl_computes_head_80_attention() -> None:
+    """SDPAImpl produces correct output for head_size=80 (not in the old list)."""
+    bs, seq, heads, head_dim = 2, 8, 2, 80
+    softmax_scale = head_dim**-0.5
+    impl = SDPAImpl(num_heads=heads, head_size=head_dim, causal=False, softmax_scale=softmax_scale)
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(bs, seq, heads, head_dim) for _ in range(3))
+
+    out = impl.forward(q, k, v, SDPAMetadata(current_timestep=0))
+
+    assert out.shape == (bs, seq, heads, head_dim)
+    qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v))
+    ref = torch.softmax(qt @ kt.transpose(-2, -1) * softmax_scale, dim=-1) @ vt
+    torch.testing.assert_close(out, ref.transpose(1, 2), rtol=1e-4, atol=1e-4)

fastvideo/tests/contract/test_modal_fa4_policy.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Guard Modal runtime policies that keep CI lanes on their intended backend.
+
+Pure text/AST analysis: no fastvideo imports, no torch, no Modal client.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+MODAL_ROOT = REPO_ROOT / "fastvideo" / "tests" / "modal"
+PR_TEST = MODAL_ROOT / "pr_test.py"
+LAUNCH_L40S_JOB = MODAL_ROOT / "launch_l40s_job.py"
+SSIM_TEST = MODAL_ROOT / "ssim_test.py"
+
+
+def _function_strings(path: Path, function_name: str) -> str:
+    source = path.read_text(encoding="utf-8")
+    tree = ast.parse(source)
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            return "\n".join(
+                child.value
+                for child in ast.walk(node)
+                if isinstance(child, ast.Constant)
+                and isinstance(child.value, str)
+            )
+    raise AssertionError(f"{function_name} not found in {path}")
+
+
+def test_generic_l40s_launcher_defaults_fa4_off():
+    source = LAUNCH_L40S_JOB.read_text(encoding="utf-8")
+    assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0")' in source
+
+
+def test_ssim_launcher_keeps_fa4_enabled_by_default():
+    source = SSIM_TEST.read_text(encoding="utf-8")
+    assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1")' in source
+
+
+def test_performance_identity_env_reaches_modal_runtime():
+    pr_source = PR_TEST.read_text(encoding="utf-8")
+    launch_source = LAUNCH_L40S_JOB.read_text(encoding="utf-8")
+    runtime_secret = pr_source.split("ci_env_secret =", 1)[1].split("hf_secret =", 1)[0]
+
+    for key in ("FASTVIDEO_ATTENTION_BACKEND", "FASTVIDEO_PERFORMANCE_PROFILE_VERSION"):
+        assert key in runtime_secret
+    assert "FASTVIDEO_PERFORMANCE_PROFILE_VERSION" in launch_source.split(".env({", 1)[1]
+
+
+def test_performance_lane_classifies_pull_requests_before_main():
+    function_strings = _function_strings(PR_TEST, "run_performance_tests")
+    assert function_strings.index(
+        "BUILDKITE_PULL_REQUEST") < function_strings.index("BUILDKITE_BRANCH")
+
+
+def test_pr_model_load_and_training_lanes_disable_fa4():
+    lanes = {
+        "run_transformer_tests": "pytest ./fastvideo/tests/transformers -vs",
+        "run_training_tests": "pytest ./fastvideo/tests/training/Vanilla -srP",
+        "run_training_lora_tests": "pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP",
+        "run_training_tests_VSA": "pytest ./fastvideo/tests/training/VSA -srP",
+        "run_distill_dmd_tests": "pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs",
+        "run_self_forcing_tests": "pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs",
+        "run_train_framework_tests": "pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs",
+        "seed_grad_norm_references": "pytest ./fastvideo/tests/train/methods -vs -rs",
+    }
+
+    for function_name, pytest_command in lanes.items():
+        function_strings = _function_strings(PR_TEST, function_name)
+        assert "FASTVIDEO_FA4=0" in function_strings
+        assert pytest_command in function_strings

fastvideo/tests/encoders/test_qwen2_5_vl_vision_dtype.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+
+from fastvideo.models.encoders.qwen2_5_vl_custom import (
+    Qwen2_5_VisionTransformerPretrainedModel,
+    Qwen2_5_VLForConditionalGenerationSimple,
+)
+
+
+def _vision_config(torch_dtype=None):
+    return SimpleNamespace(
+        torch_dtype=torch_dtype,
+        spatial_merge_size=2,
+        patch_size=14,
+        fullatt_block_indexes=[],
+        window_size=112,
+        temporal_patch_size=2,
+        in_channels=3,
+        hidden_size=8,
+        num_heads=2,
+        depth=0,
+        _attn_implementation="sdpa",
+        out_hidden_size=8,
+    )
+
+
+def _full_config(torch_dtype):
+    return SimpleNamespace(
+        vision_config=_vision_config(),
+        torch_dtype=torch_dtype,
+        vocab_size=16,
+        hidden_size=8,
+        num_hidden_layers=0,
+        num_attention_heads=2,
+        pad_token_id=0,
+        _attn_implementation="sdpa",
+        rms_norm_eps=1e-6,
+        max_position_embeddings=32,
+        rope_theta=1_000_000.0,
+        rope_scaling=None,
+    )
+
+
+def test_vision_dtype_uses_parent_bfloat16_string_when_vision_dtype_missing():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype="bfloat16",
+    )
+
+    assert model.dtype == torch.bfloat16
+
+
+def test_vision_dtype_accepts_parent_torch_dtype_object():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype=torch.bfloat16,
+    )
+
+    assert model.dtype == torch.bfloat16
+
+
+def test_vision_dtype_strips_parent_dtype_string_whitespace():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype=" torch.bfloat16 ",
+    )
+
+    assert model.dtype == torch.bfloat16
+
+
+def test_vision_dtype_prefers_explicit_vision_dtype_over_parent_dtype():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(torch_dtype="float16"),
+        parent_torch_dtype="bfloat16",
+    )
+
+    assert model.dtype == torch.float16
+
+
+def test_vision_dtype_falls_back_to_float32_for_missing_or_unknown_dtype():
+    missing = Qwen2_5_VisionTransformerPretrainedModel(_vision_config())
+    unknown = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype="not-a-real-dtype",
+    )
+
+    assert missing.dtype == torch.float32
+    assert unknown.dtype == torch.float32
+
+
+def test_conditional_generation_passes_parent_dtype_to_visual_tower():
+    model = Qwen2_5_VLForConditionalGenerationSimple(_full_config("torch.bfloat16"))
+
+    assert model.visual.dtype == torch.bfloat16

fastvideo/tests/modal/launch_l40s_job.py

@@ -95,11 +95,17 @@ image = (
         "PATH": "/root/.cargo/bin:$PATH",
         "HF_HOME": "/root/data/.cache",
         "TOKENIZERS_PARALLELISM": "false",
+        "IMAGE_VERSION": IMAGE_VERSION,
+        "FASTVIDEO_CONTAINER_IMAGE_REF": IMAGE_REF,
         **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
         "FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
-        # FA4 is opt-in (FASTVIDEO_FA4); keep CI parity with the seeded
-        # references. Caller override wins.
-        "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
+        **({
+            "FASTVIDEO_PERFORMANCE_PROFILE_VERSION": os.environ["FASTVIDEO_PERFORMANCE_PROFILE_VERSION"]
+        } if os.environ.get("FASTVIDEO_PERFORMANCE_PROFILE_VERSION") else {}),
+        # FA4 is opt-in (FASTVIDEO_FA4). Generic ad hoc jobs should follow the
+        # product default unless a caller opts in through the local env or
+        # --env-vars.
+        "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0"),
     })
 )
 

fastvideo/tests/modal/pr_test.py

@@ -65,16 +65,28 @@ ci_env_secret = modal.Secret.from_dict({
     "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
     "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
     "BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
+    "BUILDKITE_SOURCE": os.environ.get("BUILDKITE_SOURCE", ""),
     "BUILDKITE_BUILD_URL": os.environ.get("BUILDKITE_BUILD_URL", ""),
     "BUILDKITE_BUILD_ID": os.environ.get("BUILDKITE_BUILD_ID", ""),
     "BUILDKITE_JOB_ID": os.environ.get("BUILDKITE_JOB_ID", ""),
     "TEST_SCOPE": os.environ.get("TEST_SCOPE", ""),
     "IMAGE_VERSION": image_version,
+    "FASTVIDEO_CONTAINER_IMAGE_REF": image_ref,
+    **{
+        key: os.environ[key]
+        for key in (
+            "FASTVIDEO_ATTENTION_BACKEND",
+            "FASTVIDEO_PERFORMANCE_PROFILE_VERSION",
+        )
+        if os.environ.get(key)
+    },
     **({
         "UV_TORCH_BACKEND": uv_torch_backend_override
     } if uv_torch_backend_override else {}),
-    # FA4 is opt-in (FASTVIDEO_FA4); CI lanes keep it enabled to match the
-    # SSIM/perf baselines. Caller override wins.
+    # FA4 is opt-in (FASTVIDEO_FA4). Keep the default enabled for
+    # inference/perf parity; model-load and training lanes that do not exercise
+    # FA4 explicitly set FASTVIDEO_FA4=0 in their command strings below.
+    # Caller override wins.
     "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
 })
 
@@ -184,7 +196,8 @@ def run_vae_tests():
               volumes={"/root/data": model_vol})
 def run_transformer_tests():
     run_test(
-        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs"
+        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
+        "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/transformers -vs"
     )
 
 
@@ -197,7 +210,8 @@ def run_transformer_tests():
               volumes={"/root/data": model_vol})
 def run_training_tests():
     run_test(
-        "export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP"
+        "export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
+        "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/Vanilla -srP"
     )
 
 
@@ -210,7 +224,8 @@ def run_training_tests():
               volumes={"/root/data": model_vol})
 def run_training_lora_tests():
     run_test(
-        "export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP"
+        "export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
+        "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP"
     )
 
 
@@ -220,7 +235,7 @@ def run_training_lora_tests():
               secrets=[wandb_secret, ci_env_secret])
 def run_training_tests_VSA():
     run_test(
-        "wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/VSA -srP"
+        "wandb login $WANDB_API_KEY && FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/VSA -srP"
     )
 
 
@@ -254,7 +269,7 @@ def run_inference_lora_tests():
 @app.function(gpu="L40S:2", image=image, timeout=900, secrets=[ci_env_secret])
 def run_distill_dmd_tests():
     run_test(
-        "pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
+        "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
 
 
 @app.function(gpu="L40S:2",
@@ -263,7 +278,8 @@ def run_distill_dmd_tests():
               secrets=[wandb_secret, ci_env_secret])
 def run_self_forcing_tests():
     run_test(
-        "wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
+        "wandb login $WANDB_API_KEY && "
+        "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
     )
 
 
@@ -323,7 +339,8 @@ def run_dreamverse_app_tests():
               volumes={"/root/data": model_vol})
 def run_train_framework_tests():
     run_test(
-        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs"
+        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
+        "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs"
     )
 
 
@@ -349,7 +366,8 @@ def seed_grad_norm_references():
     the local command and the ``_DEVICE_MAPPINGS`` table.
     """
     run_test(
-        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs"
+        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
+        "FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs"
     )
 
 
@@ -400,12 +418,12 @@ def run_performance_tests():
         "export HF_HOME='/root/data/.cache' && "
         "export PERFORMANCE_TRACKING_ROOT='/tmp/perf-tracking' && "
         "hf auth login --token $HF_API_KEY && "
-        "if [ \"${BUILDKITE_BRANCH:-}\" = 'main' ] && [ \"${TEST_SCOPE:-}\" = 'full' ]; then "
-        "export PERF_RUN_SOURCE='scheduled_main'; "
-        "export PERF_UPLOAD_POLICY='always'; "
-        "elif [ -n \"${BUILDKITE_PULL_REQUEST:-}\" ] && [ \"${BUILDKITE_PULL_REQUEST:-false}\" != 'false' ]; then "
+        "if [ -n \"${BUILDKITE_PULL_REQUEST:-}\" ] && [ \"${BUILDKITE_PULL_REQUEST:-false}\" != 'false' ]; then "
         "export PERF_RUN_SOURCE='pr'; "
         "export PERF_UPLOAD_POLICY='pass'; "
+        "elif [ \"${BUILDKITE_BRANCH:-}\" = 'main' ] && ( [ \"${BUILDKITE_SOURCE:-}\" = 'schedule' ] || [ \"${TEST_SCOPE:-}\" = 'full' ] ); then "
+        "export PERF_RUN_SOURCE='scheduled_main'; "
+        "export PERF_UPLOAD_POLICY='always'; "
         "elif [ \"${TEST_SCOPE:-}\" = 'direct' ]; then "
         "export PERF_RUN_SOURCE='unknown'; "
         "export PERF_UPLOAD_POLICY='pass'; "

fastvideo/tests/ops/test_rotary_pos_embed_cache.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the get_rotary_pos_embed memoization cache."""
+
+import pytest
+import torch
+
+from fastvideo.layers.rotary_embedding import (
+    _ROTARY_POS_EMBED_CACHE,
+    _ROTARY_POS_EMBED_CACHE_MAXSIZE,
+    get_rotary_pos_embed,
+)
+
+
+def _rope_dim_list(hidden_size: int, heads_num: int) -> list[int]:
+    """Return the default 3-axis rope_dim_list used by the video DiTs."""
+    d = hidden_size // heads_num
+    return [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+
+
+def _call(
+    rope_sizes=(21, 30, 52),
+    hidden_size=1536,
+    heads_num=12,
+    rope_dim_list="default",
+    rope_theta=10000.0,
+    dtype=torch.float64,
+    start_frame=0,
+    use_real=True,
+    **kwargs,
+):
+    """Thin wrapper around get_rotary_pos_embed with DiT-like defaults."""
+    if rope_dim_list == "default":
+        rope_dim_list = _rope_dim_list(hidden_size, heads_num)
+    return get_rotary_pos_embed(
+        rope_sizes,
+        hidden_size,
+        heads_num,
+        rope_dim_list,
+        rope_theta,
+        dtype=dtype,
+        start_frame=start_frame,
+        use_real=use_real,
+        **kwargs,
+    )
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    """Isolate every test by clearing the module-level cache around it."""
+    _ROTARY_POS_EMBED_CACHE.clear()
+    yield
+    _ROTARY_POS_EMBED_CACHE.clear()
+
+
+def test_repeated_call_hits_cache():
+    """A second identical call returns the exact same tensor objects."""
+    cos1, sin1 = _call()
+    cos2, sin2 = _call()
+    assert cos1 is cos2 and sin1 is sin2
+    assert len(_ROTARY_POS_EMBED_CACHE) == 1
+
+
+def test_many_identical_calls_keep_single_entry():
+    """Many identical calls never grow the cache beyond one entry."""
+    for _ in range(10):
+        _call()
+    assert len(_ROTARY_POS_EMBED_CACHE) == 1
+
+
+@pytest.mark.parametrize(
+    "rope_sizes,hidden_size,heads_num,dtype,use_real",
+    [
+        ((21, 30, 52), 1536, 12, torch.float64, True),
+        ((21, 45, 80), 5120, 40, torch.float64, True),
+        ((1, 16, 16), 1536, 12, torch.float32, True),
+        ((4, 8, 8), 1536, 12, torch.float64, False),
+    ],
+)
+def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num,
+                                        dtype, use_real):
+    """Cached tensors are bitwise-equal to a fresh uncached recompute."""
+    cos_cached, sin_cached = _call(rope_sizes=rope_sizes,
+                                   hidden_size=hidden_size,
+                                   heads_num=heads_num,
+                                   dtype=dtype,
+                                   use_real=use_real)
+    _ROTARY_POS_EMBED_CACHE.clear()
+    cos_fresh, sin_fresh = _call(rope_sizes=rope_sizes,
+                                 hidden_size=hidden_size,
+                                 heads_num=heads_num,
+                                 dtype=dtype,
+                                 use_real=use_real)
+    assert torch.equal(cos_cached, cos_fresh)
+    assert torch.equal(sin_cached, sin_fresh)
+
+
+@pytest.mark.parametrize(
+    "kwargs_a,kwargs_b",
+    [
+        ({"rope_sizes": (21, 30, 52)}, {"rope_sizes": (21, 45, 80)}),
+        ({"dtype": torch.float64}, {"dtype": torch.float32}),
+        ({"use_real": True}, {"use_real": False}),
+        ({"start_frame": 0}, {"start_frame": 3}),
+        ({"rope_theta": 10000.0}, {"rope_theta": 5000.0}),
+        ({"shard_dim": 0}, {"shard_dim": 1}),
+    ],
+)
+def test_distinct_args_create_distinct_entries(kwargs_a, kwargs_b):
+    """Any output-affecting argument difference yields a separate cache entry."""
+    _call(**kwargs_a)
+    _call(**kwargs_b)
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_none_rope_dim_list_shares_key_with_equivalent_list():
+    """None rope_dim_list and its derived explicit list map to one entry."""
+    # head_dim must be divisible by 3 for the None branch to stay valid.
+    hidden_size, heads_num = 1536, 16  # head_dim == 96 -> [32, 32, 32]
+    _call(rope_dim_list=None, hidden_size=hidden_size, heads_num=heads_num)
+    before = len(_ROTARY_POS_EMBED_CACHE)
+    _call(rope_dim_list=[32, 32, 32], hidden_size=hidden_size, heads_num=heads_num)
+    assert len(_ROTARY_POS_EMBED_CACHE) == before == 1
+
+
+def test_use_real_controls_last_dim():
+    """use_real=True spans full head_dim; use_real=False spans half."""
+    cos_full, _ = _call(use_real=True)
+    cos_half, _ = _call(use_real=False)
+    assert cos_full.shape[-1] == 128
+    assert cos_half.shape[-1] == 64
+
+
+@pytest.mark.parametrize("rope_sizes", [(1, 1, 1), (1, 30, 52), (21, 1, 1)])
+def test_degenerate_grid_shapes(rope_sizes):
+    """Degenerate single-element axes still produce a correctly sized table."""
+    cos, sin = _call(rope_sizes=rope_sizes)
+    expected = rope_sizes[0] * rope_sizes[1] * rope_sizes[2]
+    assert cos.shape[0] == expected
+    assert sin.shape[0] == expected
+
+
+def test_scalar_and_list_factors_are_hashable_and_distinct():
+    """List-valued rescale factors are hashable and keyed apart from scalars."""
+    _call(theta_rescale_factor=1.0)
+    _call(theta_rescale_factor=[1.0, 1.0, 1.0])
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_caller_device_copy_does_not_corrupt_cache():
+    """The .to()/.float() copy callers perform must not mutate cached tensors."""
+    cos, _ = _call()
+    snapshot = cos.clone()
+    _ = cos.to("cpu").float()
+    cos_again, _ = _call()
+    assert torch.equal(cos_again, snapshot)
+
+
+def test_start_frame_offsets_values():
+    """A non-zero start_frame shifts the temporal positions, changing output."""
+    cos0, _ = _call(start_frame=0)
+    cos3, _ = _call(start_frame=3)
+    assert not torch.equal(cos0, cos3)
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_cache_is_bounded_and_evicts_oldest():
+    """The cache caps at the max size and evicts the oldest entry first."""
+    # Tiny grids keep this lightweight; each start_frame is a distinct key.
+    overshoot = _ROTARY_POS_EMBED_CACHE_MAXSIZE + 4
+    for frame in range(overshoot):
+        _call(rope_sizes=(2, 2, 2), start_frame=frame)
+        assert len(_ROTARY_POS_EMBED_CACHE) <= _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    # The earliest-inserted frames must have been evicted; the latest survive.
+    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}  # start_frame slot
+    assert overshoot - 1 in surviving
+    assert 0 not in surviving
+
+
+def test_cache_hit_refreshes_recency():
+    """Re-accessing an entry protects it from eviction over an untouched one."""
+    _call(rope_sizes=(2, 2, 2), start_frame=0)  # entry we will keep hot
+    for frame in range(1, _ROTARY_POS_EMBED_CACHE_MAXSIZE):
+        _call(rope_sizes=(2, 2, 2), start_frame=frame)
+    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    _call(rope_sizes=(2, 2, 2), start_frame=0)   # hit -> frame 0 becomes most recent
+    _call(rope_sizes=(2, 2, 2), start_frame=99)  # miss -> evicts now-oldest (frame 1)
+    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}
+    assert 0 in surviving
+    assert 1 not in surviving

fastvideo/tests/performance/compare_baseline.py

@@ -23,7 +23,7 @@ from typing import Any
 
 try:
     from fastvideo.performance.hf_store import (
-        load_records_for_model,
+        load_records_for_identity,
         safe_float,
         sanitize,
         sync_from_hf,
@@ -40,7 +40,7 @@ except ImportError:
     if repo_root not in sys.path:
         sys.path.insert(0, repo_root)
     from fastvideo.performance.hf_store import (
-        load_records_for_model,
+        load_records_for_identity,
         safe_float,
         sanitize,
         sync_from_hf,
@@ -89,6 +89,20 @@ COMPARISON_IDENTITY_KEYS = (
     "hardware_profile_id",
     "software_profile_id",
 )
+STATUS_PASS = "PASS"
+STATUS_REGRESSION = "REGRESSION"
+STATUS_CALIBRATION_NEEDED = "CALIBRATION_NEEDED"
+STATUS_RECIPE_MISMATCH = "RECIPE_MISMATCH"
+STATUS_INFRA_ERROR = "INFRA_ERROR"
+# Reserved for the promoted-baseline workflow; never emitted here.
+STATUS_QUALITY_BLOCKED = "QUALITY_BLOCKED"
+STATIC_THRESHOLD_FIELDS = (
+    ("avg_generation_time_s", "max_generation_time_s"),
+    ("max_peak_memory_mb", "max_peak_memory_mb"),
+    ("text_encoder_time_s", "max_text_encoder_time_s"),
+    ("dit_time_s", "max_dit_time_s"),
+    ("vae_decode_time_s", "max_vae_decode_time_s"),
+)
 
 
 def _should_persist_tracking() -> bool:
@@ -122,7 +136,13 @@ def _detect_run_source() -> str:
     return "unknown"
 
 
-def _is_baseline_eligible(run_source: str, success: bool) -> bool:
+def _is_baseline_eligible(
+    run_source: str,
+    success: bool,
+    comparison_status: str | None = None,
+) -> bool:
+    if comparison_status is not None and comparison_status != STATUS_PASS:
+        return False
     return run_source == "scheduled_main" and success
 
 
@@ -135,7 +155,7 @@ def _upload_allowed(record: dict[str, Any]) -> bool:
     return False
 
 
-def _result_failed_static_thresholds() -> bool:
+def _performance_pytest_failed() -> bool:
     value = os.environ.get("PERF_PYTEST_RC", "")
     if not value:
         return False
@@ -193,18 +213,52 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
     }
 
 
+def _record_uses_v2_identity(record: dict[str, Any]) -> bool:
+    schema_version = record.get("result_schema_version")
+    return str(schema_version) == "2" or any(key in record for key in COMPARISON_IDENTITY_KEYS)
+
+
+def _recipe_cohort_filters(record: dict[str, Any]) -> dict[str, str]:
+    identity = _comparison_identity_filters(record)
+    return {
+        key: identity[key]
+        for key in ("workload_id", "variant_id", "benchmark_version")
+    }
+
+
 def _comparison_identity_filters_or_none(record: dict[str, Any]) -> dict[str, str] | None:
     try:
         return _comparison_identity_filters(record)
     except ValueError as exc:
-        print("Skipping rolling baseline comparison for "
-              f"{record.get('model_id', 'unknown')} on "
-              f"{record.get('gpu_type', 'unknown')}: {exc}. "
-              "The normalized artifact will still be written, but this record "
-              "will not be baseline eligible.")
+        if _record_uses_v2_identity(record):
+            print(f"Invalid v2 comparison identity for {record.get('model_id', 'unknown')}: {exc}")
+        else:
+            print("Skipping rolling baseline comparison for "
+                  f"{record.get('model_id', 'unknown')} on "
+                  f"{record.get('gpu_type', 'unknown')}: {exc}. "
+                  "The normalized artifact will still be written, but this record "
+                  "will not be baseline eligible.")
         return None
 
 
+def _recipe_mismatch_records(
+    record: dict[str, Any],
+    cohort_records: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    current_recipe = str(record.get("recipe_fingerprint"))
+    return [
+        item for item in cohort_records
+        if _record_can_authorize_recipe_mismatch(item)
+        and item.get("recipe_fingerprint") is not None
+        and str(item.get("recipe_fingerprint")).strip()
+        and str(item.get("recipe_fingerprint")) != current_recipe
+    ]
+
+
+def _record_can_authorize_recipe_mismatch(record: dict[str, Any]) -> bool:
+    return record.get("run_source") == "scheduled_main" or record.get("baseline_eligible") is True
+
+
 def _format_identity_filters(filters: dict[str, str]) -> str:
     if not filters:
         return ""
@@ -343,12 +397,99 @@ def _check_regressions(
     return failures
 
 
+def _fixed_threshold_failures(raw_result: dict[str, Any]) -> list[str]:
+    thresholds = raw_result.get("thresholds")
+    if not isinstance(thresholds, dict):
+        return []
+
+    failures = []
+    for result_key, threshold_key in STATIC_THRESHOLD_FIELDS:
+        current = safe_float(raw_result.get(result_key))
+        threshold = safe_float(thresholds.get(threshold_key))
+        if current is None or threshold is None or current <= threshold:
+            continue
+        failures.append(
+            f"{raw_result.get('benchmark_id', 'unknown')} {result_key} exceeded fixed threshold "
+            f"(current={current:.3f}, threshold={threshold:.3f})"
+        )
+    return failures
+
+
+def _pytest_infra_failure(record: dict[str, Any]) -> str:
+    return (f"{record['model_id']} performance pytest failed without an "
+            "attributable static-threshold regression "
+            f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
+
+
+def _recipe_mismatch_failure(
+    record: dict[str, Any],
+    recipe_mismatch_records: list[dict[str, Any]],
+) -> str:
+    seen = sorted({
+        str(item.get("recipe_fingerprint"))
+        for item in recipe_mismatch_records
+        if item.get("recipe_fingerprint") is not None
+    })
+    suffix = f"; existing recipe_fingerprint values: {', '.join(seen)}" if seen else ""
+    return (f"{record['model_id']} recipe_fingerprint={record.get('recipe_fingerprint')} "
+            "does not match baseline records for the same workload, variant, "
+            "and benchmark version across hardware and software profiles"
+            f"{suffix}")
+
+
+def _evaluate_record_comparison(
+    record: dict[str, Any],
+    baseline_records: list[dict[str, Any]],
+    recipe_mismatch_records: list[dict[str, Any]],
+    metric_policies: tuple[MetricPolicy, ...],
+    static_threshold_failures: list[str],
+    unattributed_pytest_failure: bool,
+) -> tuple[list[str], str, str]:
+    identity_filters = None
+    if _record_uses_v2_identity(record):
+        try:
+            identity_filters = _comparison_identity_filters(record)
+        except ValueError as exc:
+            failure = f"{record['model_id']} cannot compare v2 record: {exc}"
+            return [failure], STATUS_INFRA_ERROR, failure
+        if not baseline_records and recipe_mismatch_records:
+            failure = _recipe_mismatch_failure(record, recipe_mismatch_records)
+            return [failure], STATUS_RECIPE_MISMATCH, failure
+
+    if static_threshold_failures:
+        return static_threshold_failures, STATUS_REGRESSION, "; ".join(static_threshold_failures)
+
+    if unattributed_pytest_failure:
+        failure = _pytest_infra_failure(record)
+        return [failure], STATUS_INFRA_ERROR, failure
+
+    if not baseline_records:
+        if identity_filters is not None:
+            return [], STATUS_CALIBRATION_NEEDED, (
+                "No baseline found for exact comparable identity"
+                f"{_format_identity_filters(identity_filters)}"
+            )
+
+        return [], STATUS_PASS, f"No legacy baseline for {record['model_id']} on {record['gpu_type']}"
+
+    failures = _check_regressions(record, baseline_records, metric_policies)
+    if failures:
+        return failures, STATUS_REGRESSION, "; ".join(failures)
+
+    return [], STATUS_PASS, "Comparable baseline found with no gated regressions"
+
+
 def _compact_value(value: float | None, precision: int = 3) -> str:
     if value is None:
         return "n/a"
     return f"{value:.{precision}f}"
 
 
+def _markdown_cell(value: Any) -> str:
+    text = str(value) if value is not None else ""
+    return text.replace("\n", " ").replace("|", "/")
+
+
 def _build_summary_row(
     record: dict[str, Any],
     baseline_records: list[dict[str, Any]],
@@ -400,6 +541,8 @@ def _build_summary_row(
         "threshold_exceeded_metrics": threshold_exceeded_metrics,
         "failing_metrics": failing_metrics,
         "failed": has_failed,
+        "status": record.get("comparison_status", STATUS_REGRESSION if has_failed else STATUS_PASS),
+        "status_reason": record.get("comparison_status_reason", ""),
     }
 
 
@@ -417,8 +560,8 @@ def _build_markdown_summary(
          "Throughput (curr/base) | Memory (curr/base) | "
          "Text Enc (curr/base) | DiT (curr/base) | "
          "VAE Decode (curr/base) | Worst Regression | Exceeded Metrics | "
-         "Failing Metrics | Status |"),
-        "|---|---|---:|---|---|---|---|---|---|---:|---|---|---|",
+         "Failing Metrics | Status | Status Detail |"),
+        "|---|---|---:|---|---|---|---|---|---|---:|---|---|---|---|",
     ]
 
     for row in summary_rows:
@@ -435,12 +578,14 @@ def _build_markdown_summary(
             else "none"
         )
         failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none"
-        status = "FAIL" if row["failed"] else "PASS"
+        status = row["status"]
+        status_detail = row["status_reason"]
 
         lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
                      f"{row['baseline_n']} | "
                      f"{' | '.join(metric_cells)} | "
-                     f"{worst_reg} | {exceeded_metrics} | {failing_metrics} | {status} |")
+                     f"{worst_reg} | {exceeded_metrics} | {failing_metrics} | "
+                     f"{_markdown_cell(status)} | {_markdown_cell(status_detail)} |")
 
     return "\n".join(lines) + "\n"
 
@@ -470,7 +615,7 @@ def _emit_markdown_summary(markdown: str, commit_sha: str) -> None:
 def main() -> int:
     persist_tracking = _should_persist_tracking()
     upload_policy = _normalized_upload_policy()
-    static_threshold_failed = _result_failed_static_thresholds()
+    performance_pytest_failed = _performance_pytest_failed()
 
     # Strict on upload-enabled runs: silent sync failure would make comparison
     # and upload state ambiguous.
@@ -481,6 +626,9 @@ def main() -> int:
         print(f"No performance result files found in {RESULTS_DIR}")
         return 0
 
+    static_threshold_failures = [_fixed_threshold_failures(raw) for raw in current_results]
+    unattributed_pytest_failure = performance_pytest_failed and not any(static_threshold_failures)
+
     all_failures: list[str] = []
     summary_rows: list[dict[str, Any]] = []
 
@@ -489,61 +637,75 @@ def main() -> int:
     else:
         print("Tracking persistence disabled: PERF_UPLOAD_POLICY=never")
 
-    if static_threshold_failed:
-        print(f"Static-threshold phase failed: PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
+    if performance_pytest_failed:
+        if unattributed_pytest_failure:
+            print("Performance pytest failed without a measured static-threshold regression: "
+                  f"PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
+        else:
+            print("Performance pytest failure attributed to per-record static thresholds: "
+                  f"PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
 
-    for raw in current_results:
+    for raw, fixed_threshold_failures in zip(current_results, static_threshold_failures, strict=True):
         record = _normalize_record(raw)
         metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
         identity_filters = _comparison_identity_filters_or_none(record)
 
         baseline_records: list[dict[str, Any]] = []
-        if identity_filters is None:
-            # Records without the full v2 identity block skip rolling-baseline
-            # comparison entirely: only the static thresholds gate them and
-            # they never become baseline eligible.
-            record["baseline_status"] = "skipped_missing_identity"
-            failures: list[str] = []
-        else:
-            baseline_records = load_records_for_model(
-                TRACKING_ROOT,
-                record["model_id"],
-                record["gpu_type"],
-                **identity_filters,
-                last_n=5,
-                successful_only=True,
-                baseline_eligible_only=True,
-            )
-
-            if not baseline_records:
-                # A brand-new cohort has NO regression gating until history
-                # accumulates — make that loud and machine-readable instead of
-                # an indistinguishable pass, so a cohort shift (intended or
-                # accidental, e.g. an identity-field change) never silently
-                # blinds the comparison.
-                record["baseline_status"] = "initialized_new_cohort"
-                print("=" * 72)
-                print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
-                      f"{record['model_id']} on {record['gpu_type']}"
-                      f"{_format_identity_filters(identity_filters)}")
-                print("Regression gating is INACTIVE for this cohort until "
-                      "baseline history accumulates. If this cohort shift is "
-                      "unexpected, check the identity fields above.")
-                print("=" * 72)
-                failures = []
+        recipe_mismatch_records: list[dict[str, Any]] = []
+        try:
+            if identity_filters is None:
+                record["baseline_status"] = (
+                    "invalid_identity" if _record_uses_v2_identity(record) else "skipped_missing_identity"
+                )
             else:
-                record["baseline_status"] = "compared"
-                failures = _check_regressions(record, baseline_records, metric_policies)
-        if static_threshold_failed:
-            failures.append(f"{record['model_id']} fixed-threshold phase failed "
-                            f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
+                baseline_records = load_records_for_identity(
+                    TRACKING_ROOT,
+                    identity_filters,
+                    last_n=5,
+                    successful_only=True,
+                    baseline_eligible_only=True,
+                )
+                if baseline_records:
+                    record["baseline_status"] = "compared"
+                else:
+                    recipe_cohort_records = load_records_for_identity(
+                        TRACKING_ROOT,
+                        _recipe_cohort_filters(record),
+                        successful_only=True,
+                    )
+                    recipe_mismatch_records = _recipe_mismatch_records(record, recipe_cohort_records)
+                    record["baseline_status"] = (
+                        "recipe_mismatch" if recipe_mismatch_records else "initialized_new_cohort"
+                    )
+
+            failures, comparison_status, comparison_status_reason = _evaluate_record_comparison(
+                record,
+                baseline_records,
+                recipe_mismatch_records,
+                metric_policies,
+                fixed_threshold_failures,
+                unattributed_pytest_failure,
+            )
+        except Exception as exc:
+            comparison_status = STATUS_INFRA_ERROR
+            comparison_status_reason = f"{record['model_id']} baseline comparison hit an infra error: {exc}"
+            failures = [comparison_status_reason]
+            baseline_records = []
+            record["baseline_status"] = "infra_error"
+
+        record["comparison_status"] = comparison_status
+        record["comparison_status_reason"] = comparison_status_reason
 
         record["success"] = not failures
         record["baseline_eligible"] = (
-            identity_filters is not None and _is_baseline_eligible(record["run_source"], record["success"])
+            identity_filters is not None
+            and _is_baseline_eligible(record["run_source"], record["success"], comparison_status)
         )
         all_failures.extend(failures)
 
+        print(f"{record['model_id']} comparison status: "
+              f"{comparison_status} - {comparison_status_reason}")
+
         _write_normalized_artifact(record)
 
         if _upload_allowed(record):

fastvideo/tests/performance/dashboard.py

@@ -37,7 +37,7 @@ COMPARISON_COHORT_KEYS = (
     "hardware_profile_id",
     "software_profile_id",
 )
-GROUP_KEYS = ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)
+GROUP_KEYS = ("_cohort_schema", "_cohort_model_id", "_cohort_gpu_type", *COMPARISON_COHORT_KEYS)
 
 
 def _cohort_value(value: object) -> str:
@@ -62,35 +62,43 @@ def _short_value(value: object) -> str:
     return text[:12]
 
 
-def _cohort_title(group_key: tuple[object, ...]) -> str:
-    workload = _display_value(group_key[2])
-    variant = _display_value(group_key[3])
-    version = _display_value(group_key[4])
+def _cohort_title(record: object) -> str:
+    workload = _display_value(record.get("workload_id"))
+    variant = _display_value(record.get("variant_id"))
+    version = _display_value(record.get("benchmark_version"))
     version_label = version if version == "legacy" else f"v{version}"
     return f"{workload} / {variant} / {version_label}"
 
 
-def _cohort_detail(group_key: tuple[object, ...]) -> str:
+def _cohort_detail(record: object) -> str:
     return " | ".join((
-        f"recipe {_short_value(group_key[5])}",
-        _short_value(group_key[6]),
-        _short_value(group_key[7]),
+        f"recipe {_short_value(record.get('recipe_fingerprint'))}",
+        _short_value(record.get("hardware_profile_id")),
+        _short_value(record.get("software_profile_id")),
     ))
 
 
-def _group_record(group_key: tuple[object, ...]) -> dict[str, str]:
+def _group_record(record: object) -> dict[str, str]:
     return {
-        key: _cohort_value(value)
-        for key, value in zip(GROUP_KEYS, group_key)
+        key: _cohort_value(record.get(key))
+        for key in ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)
     }
 
 
+def _has_complete_v2_identity(record: object) -> bool:
+    return all(_cohort_value(record.get(key)) for key in COMPARISON_COHORT_KEYS)
+
+
 def _dashboard_frame(df: pd.DataFrame) -> pd.DataFrame:
     dashboard_df = df.copy()
-    for key in GROUP_KEYS:
+    for key in ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS):
         if key not in dashboard_df.columns:
             dashboard_df[key] = ""
         dashboard_df[key] = dashboard_df[key].map(_cohort_value)
+    uses_v2_identity = dashboard_df.apply(_has_complete_v2_identity, axis=1)
+    dashboard_df["_cohort_schema"] = uses_v2_identity.map({True: "v2", False: "legacy"})
+    dashboard_df["_cohort_model_id"] = dashboard_df["model_id"].where(~uses_v2_identity, "")
+    dashboard_df["_cohort_gpu_type"] = dashboard_df["gpu_type"].where(~uses_v2_identity, "")
     return dashboard_df
 
 # -----------------------------
@@ -107,12 +115,14 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
     figs = []
     skipped_metrics: list[dict[str, object]] = []
 
-    for group_key, g in group_data(df):
-        model_id, gpu_type = group_key[:2]
-        group_record = _group_record(group_key)
-        cohort_title = _cohort_title(group_key)
-        cohort_detail = _cohort_detail(group_key)
+    for _group_key, g in group_data(df):
         g = g.sort_values("timestamp")
+        display_record = g.iloc[-1]
+        model_id = display_record["model_id"]
+        gpu_type = display_record["gpu_type"]
+        group_record = _group_record(display_record)
+        cohort_title = _cohort_title(display_record)
+        cohort_detail = _cohort_detail(display_record)
 
         # One chart per metric so the y-axes aren't on wildly different scales
         for metric in METRICS:
@@ -146,7 +156,7 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
                 x="timestamp",
                 y=metric,
                 markers=True,
-                hover_data=["config_id", "commit_sha", *COMPARISON_COHORT_KEYS],
+                hover_data=["model_id", "gpu_type", "config_id", "commit_sha", *COMPARISON_COHORT_KEYS],
                 title=f"{model_id} | {gpu_type} | {cohort_title} | {cohort_detail} | {metric}",
                 labels={"timestamp": "Time", metric: metric},
             )

fastvideo/tests/performance/identity.py

@@ -17,7 +17,7 @@ from typing import Any
 
 import torch
 
-RECIPE_SCHEMA_VERSION = 1
+RECIPE_SCHEMA_VERSION = 2
 PROFILE_ID_LENGTH = 16
 _PATH_EXCLUDED_GENERATION_KEYS = {"output_path", "output_video_name"}
 _PROMPT_KEYS = {"prompt", "negative_prompt", "neg_prompt"}
@@ -32,6 +32,7 @@ _PROFILE_ENV_VARS = (
     "FASTVIDEO_ATTENTION_BACKEND",
     "IMAGE_VERSION",
     "UV_TORCH_BACKEND",
+    "FASTVIDEO_CONTAINER_IMAGE_REF",
 )
 _PACKAGE_DISTRIBUTIONS = {
     "fastvideo_kernel": ("fastvideo-kernel", "fastvideo_kernel"),
@@ -87,6 +88,12 @@ def recipe_fingerprint(recipe: Mapping[str, Any]) -> str:
         # Intentional backend changes are declared via requested_backend,
         # which stays in the hash.
         attention.pop("resolved_backend", None)
+    benchmark = pruned.get("benchmark")
+    if isinstance(benchmark, dict):
+        # benchmark_id names artifacts and storage directories, but is not one
+        # of the six v2 comparison fields. Keep it in the recipe for audit
+        # while allowing display-only renames to stay in the same cohort.
+        benchmark.pop("benchmark_id", None)
     return sha256_hexdigest(canonical_json(pruned))