Compare commits

..
243 changed files with 932 additions and 24174 deletions
@@ -1,23 +1,20 @@
---
name: reseed-performance-baseline
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.
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.
---
# Re-seed Performance Baseline
## Purpose
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`.
Replace or advance the rolling performance baseline for a single
`(model_id, gpu_type)` pair in the HF dataset
`FastVideo/performance-tracking`.
Performance comparison uses the median of up to the last 5 successful,
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`.
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`.
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
@@ -25,13 +22,11 @@ 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.
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`.
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.
These records are intentional operator-approved baseline resets, not ordinary
independent main-branch persistence. Mark them clearly with provenance fields
@@ -71,8 +66,8 @@ approval, then upload reviewed accepted baseline records.
| Parameter | Required | Description |
|-----------|----------|-------------|
| `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. |
| `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
| `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
| `source_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. |
@@ -83,14 +78,10 @@ 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`, `baseline_eligible=true` records
for the same legacy `(model_id, gpu_type)` target or the same v2 exact
comparable identity.
- Baseline window: last 5 `success=true` records for the same
`(model_id, gpu_type)`.
- Reseed count: dynamic. Upload exactly one accepted seed record per validated
source JSON.
@@ -124,24 +115,12 @@ with open(source_result, encoding="utf-8") as f:
record = json.load(f)
```
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.
Stop if any normalized record's `model_id` or `gpu_type` does not match the
requested `model_id` and `gpu_type`.
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,
@@ -215,7 +194,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)'
```
For legacy records, back up the sanitized model directory under `/tmp`:
Then back up only the sanitized model directory under `/tmp`:
```bash
SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
@@ -230,16 +209,6 @@ 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
@@ -262,9 +231,7 @@ first baseline seed. Continue, but report that baseline history was empty.
### 3. Compute old baseline and candidate shift
Load the last 5 successful baseline records for the target.
For legacy targets:
Load the last 5 successful records for the target:
```python
from fastvideo.performance.hf_store import load_records_for_model
@@ -275,28 +242,6 @@ 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,
)
```
@@ -312,8 +257,7 @@ 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 an existing last-5 median by itself, but
it is enough to establish the first v2 baseline for a new exact identity.
- 1 seed record usually does not move a last-5 median by itself.
- 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
@@ -323,10 +267,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 `<target description>`.
> About to RE-SEED performance baseline for `<model_id>` on `<gpu_type>`.
> This will upload `<N>` new `success=true` records to
> `FastVideo/performance-tracking/<sanitize(model_id)>/` or the source
> artifact's v2 model directory, one per accepted source JSON.
> `FastVideo/performance-tracking/<sanitize(model_id)>/`, one per accepted
> source JSON.
>
> Reason: `<intent_rationale>`
> Source results: `<source_results>`
@@ -344,63 +288,25 @@ 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.
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
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
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.
`(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.
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, fall back to this default
baseline field list:
If there are no previous HF records for the target model/GPU, fall back to this
default baseline field list:
- `model_id`
- `timestamp`
@@ -413,22 +319,6 @@ 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.
@@ -444,22 +334,6 @@ 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
@@ -482,8 +356,7 @@ Prefer uploading new accepted seed records so failed history remains visible.
Print:
- Backup directory path under `/tmp`.
- Prepared local record paths under `PERFORMANCE_RESEED_STAGING_ROOT`.
- Prepared upload-manifest path under the identity reservation.
- Prepared local record paths under `PERFORMANCE_TRACKING_ROOT`.
- HF paths that will receive the new records.
- Old rolling medians.
- Source batch medians, source batch spread, reseed count, and candidate
@@ -495,36 +368,22 @@ prepared records plus backup on disk.
### 7. Upload only the scoped records
For a first v2 calibration seed, use the manifest uploader after the user
replies exactly `upload`:
Use the shared storage helper so the path and repo type match CI:
```bash
python -c 'from fastvideo.tests.performance.seed_baseline import upload_prepared_seed_manifest; print(upload_prepared_seed_manifest("<prepared_manifest>"))'
```python
from fastvideo.performance.hf_store import upload_record
upload_record("<local_record_path>", record, strict=True)
```
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:
Run it once per prepared record. Each upload goes to:
```text
FastVideo/performance-tracking/<sanitize(model_id)>/<record_filename>.json
```
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.
Never bulk upload the whole tracking root. Never modify another model's
directory in the same operation.
### 8. Report outcome and offer cleanup
@@ -546,14 +405,9 @@ 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`: 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.
- `PERFORMANCE_TRACKING_ROOT`, usually `/tmp/perf-tracking`: local synced
mirror of `FastVideo/performance-tracking` plus the prepared local seed
records used for scoped upload.
- `/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.
@@ -563,18 +417,14 @@ local state. Explain what each directory is for:
Ask:
> Reseed succeeded. Do you want me to delete the local temp tracking mirror,
> 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.
> 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 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.
directories created for this reseed. Never remove unrelated `/tmp` contents.
## Failure modes and handling
@@ -586,34 +436,19 @@ 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 an existing last-5
median. This warning does not block a first v2 calibration seed for an exact
identity with no eligible baseline yet.
clear that one or two records may not immediately move the last-5 median.
- **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 prepared seed
records under `/tmp/performance_reseed_prepared`, 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 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.
@@ -624,9 +459,8 @@ 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 and record loading helpers.
- `fastvideo/tests/performance/seed_baseline.py` — first-seed preparation,
staging reservation, manifest validation, and conditional batch upload.
- `fastvideo/performance/hf_store.py` — HF sync, record loading,
`sanitize()`, and `upload_record()`.
- `fastvideo/tests/performance/test_inference_performance.py` — source result
JSON schema.
- `.buildkite/performance-benchmarks/tests/*.json` — fixed absolute benchmark
@@ -639,4 +473,3 @@ 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. |
@@ -3,7 +3,7 @@
"config_schema_version": 2,
"workload_id": "wan-t2v",
"variant_id": "1.3b-sp2",
"benchmark_version": 3,
"benchmark_version": 2,
"description": "Wan2.1 T2V 1.3B inference performance",
"model": {
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+1 -2
View File
@@ -436,7 +436,7 @@ steps:
- "pyproject.toml"
- "docker/Dockerfile"
config:
command: "timeout 25m .buildkite/scripts/pr_test.sh"
command: "timeout 15m .buildkite/scripts/pr_test.sh"
label: ":test_tube: LoRA Training Tests"
env:
- TEST_TYPE=training_lora
@@ -481,7 +481,6 @@ steps:
- "fastvideo/layers/**"
- "fastvideo/worker/**"
- "fastvideo/entrypoints/**"
- "fastvideo/performance/**"
- "fastvideo/tests/performance/**"
- ".buildkite/performance-benchmarks/**"
- "pyproject.toml"
+2 -24
View File
@@ -76,27 +76,10 @@ 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:-} 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"
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"
POST_RUN_HOOK=""
is_truthy() {
case "${1:-}" in
1|true|TRUE|yes|YES|on|ON) return 0 ;;
*) return 1 ;;
esac
}
ssim_bootstrap_args() {
local title="${PR_TITLE:-}"
local message="${BUILDKITE_MESSAGE:-}"
if is_truthy "${FASTVIDEO_SSIM_BOOTSTRAP_MODE:-}" \
|| [[ "$title" == *"[new-model]"* ]] \
|| [[ "$message" == *"[new-model]"* ]]; then
printf ' --bootstrap-mode'
fi
}
upload_performance_artifacts() {
SHORT_SHA=${BUILDKITE_COMMIT:0:7}
LOCAL_DIR="downloaded_reports"
@@ -189,12 +172,7 @@ case "$TEST_TYPE" in
;;
"ssim")
log "Running SSIM tests..."
SSIM_BOOTSTRAP_ARGS=$(ssim_bootstrap_args)
if [ -n "$SSIM_BOOTSTRAP_ARGS" ]; then
log "SSIM bootstrap mode enabled for new-model reference draft generation"
fi
MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run "
MODAL_COMMAND+="$MODAL_SSIM_TEST_FILE::run_ssim_tests$SSIM_BOOTSTRAP_ARGS"
MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_SSIM_TEST_FILE::run_ssim_tests"
;;
"training")
log "Running training tests..."
-133
View File
@@ -1,133 +0,0 @@
#!/usr/bin/env bash
# Gate the expensive Buildkite full suite on the cheap GitHub checks.
#
# Polls the workflow runs for the PR head commit and only exits 0 once the
# watched cheap workflows (pre-commit, docs build) have succeeded, so the
# 'ready' label cannot burn ~20 GPU lanes on a head that a cheap check has
# already doomed.
#
# Semantics:
# - watched run completed with a bad conclusion -> exit 1 (fail CLOSED:
# no full suite; the next push re-arms via the 'synchronize' trigger)
# - watched run cancelled -> still pending: the docs
# workflow's repo-global 'pages' concurrency group cancels runs superseded
# by unrelated pushes, so 'cancelled' is not a verdict on this PR
# - watched runs pending -> poll until done
# - docs run absent -> not applicable after a
# short grace period ('Deploy Documentation' is path-filtered on PRs)
# - pre-commit run absent -> keep polling: pre-commit
# is never path-filtered, so its absence is always anomalous
# - 'ready' label removed while waiting -> exit 1 (fail CLOSED:
# un-labeling is a deliberate maintainer action)
# - GitHub API unreachable or timeout -> exit 0 (fail OPEN,
# loud warning: never brick CI on a GitHub outage)
#
# Required env: PR_SHA (PR head commit), PR_NUMBER, GITHUB_REPOSITORY, GH_TOKEN.
set -euo pipefail
: "${PR_SHA:?PR_SHA (PR head commit) is required}"
: "${PR_NUMBER:?PR_NUMBER (pull request number) is required}"
: "${GITHUB_REPOSITORY:?GITHUB_REPOSITORY is required}"
# Workflow-level `name:` values that must be green before the full suite
# may start. "Deploy Documentation" is path-filtered on PRs, so its run may
# legitimately never exist; pre-commit always runs, so it must appear.
WATCHED_NAMES='["pre-commit", "Deploy Documentation"]'
WATCHED_REGEX='^(pre-commit|Deploy Documentation)$'
POLL_SECS="${POLL_SECS:-20}"
GRACE_SECS="${GRACE_SECS:-60}"
MAX_WAIT_SECS="${MAX_WAIT_SECS:-1500}"
# Bound each API call so a hung connection hits the 3-strike fail-open path
# instead of pinning the loop until the job timeout (which would fail closed
# on exactly the GitHub-outage case this script is meant to survive).
if command -v timeout >/dev/null 2>&1; then
gh_api() { timeout 30 gh api "$@"; }
else
gh_api() { gh api "$@"; } # macOS dev boxes; CI always has coreutils timeout
fi
# The workflow checked the label before starting the gate, but the wait can
# last ~25 min: re-check once before any exit 0 and fail closed if 'ready'
# was removed in the meantime. An API error here proceeds (the label was
# present when the gate started; never brick CI on an outage).
recheck_ready_label() {
local pr_json
if pr_json=$(gh_api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" 2>/dev/null); then
if ! jq -e '[.labels[]?.name] | index("ready")' <<<"$pr_json" >/dev/null 2>&1; then
echo "::error::PR #${PR_NUMBER} no longer has the 'ready' label —" \
"NOT triggering the Buildkite full suite. Re-add the label to re-arm."
exit 1
fi
else
echo "::warning::Could not re-check the 'ready' label on PR #${PR_NUMBER}; proceeding (it was present when the gate started)."
fi
}
start=$(date +%s)
api_fails=0
missing=""
while true; do
elapsed=$(( $(date +%s) - start ))
if runs_json=$(gh_api "repos/${GITHUB_REPOSITORY}/actions/runs?head_sha=${PR_SHA}&per_page=100" 2>/dev/null) \
&& state=$(jq --arg re "$WATCHED_REGEX" '
[.workflow_runs[]? | select(.name // "" | test($re))]
| group_by(.name) | map(max_by(.id))
| map({name, status, conclusion})' <<<"$runs_json" 2>/dev/null); then
api_fails=0
echo "t+${elapsed}s watched checks: $(jq -c . <<<"$state")"
failed=$(jq -r '[.[] | select(.status == "completed"
and (.conclusion | IN("success", "skipped", "neutral", "cancelled") | not))]
| map(.name) | join(", ")' <<<"$state")
if [ -n "$failed" ]; then
echo "::error::Cheap check(s) failed on ${PR_SHA}: ${failed}." \
"NOT triggering the Buildkite full suite. Push a fix (the 'ready'" \
"label re-arms on every push), or re-run the failed check and then" \
"re-run this workflow."
exit 1
fi
# 'cancelled' counts as pending: wait for a re-run to reach a real verdict
# (bounded by MAX_WAIT, then the fail-open below).
pending=$(jq '[.[] | select(.status != "completed" or .conclusion == "cancelled")] | length' <<<"$state")
missing=$(jq -r --argjson watched "$WATCHED_NAMES" '($watched - map(.name)) | join(", ")' <<<"$state")
if [ "$pending" -eq 0 ]; then
if [ -z "$missing" ]; then
recheck_ready_label
echo "All watched cheap checks are green — full suite may proceed."
exit 0
fi
case "$missing" in
*pre-commit*)
echo "pre-commit run not found for ${PR_SHA} yet; waiting (pre-commit is never path-filtered, so its absence is anomalous)."
;;
*)
if [ "$elapsed" -ge "$GRACE_SECS" ]; then
recheck_ready_label
echo "::warning::Watched run(s) never appeared for ${PR_SHA}: ${missing} (path-filtered, likely not applicable). Proceeding on the checks that did run."
exit 0
fi
echo "Waiting up to ${GRACE_SECS}s grace for path-filtered run(s) to appear: ${missing}."
;;
esac
fi
else
api_fails=$(( api_fails + 1 ))
echo "::warning::GitHub API error querying workflow runs for ${PR_SHA} (attempt ${api_fails}/3)."
if [ "$api_fails" -ge 3 ]; then
recheck_ready_label
echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the full suite WITHOUT the cheap-check gate."
exit 0
fi
fi
if [ "$elapsed" -ge "$MAX_WAIT_SECS" ]; then
recheck_ready_label
echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the full suite anyway."
exit 0
fi
sleep "$POLL_SECS"
done
-122
View File
@@ -1,122 +0,0 @@
#!/usr/bin/env bash
# Self-test for gate_full_suite.sh using a mocked `gh`. No network, runs on
# any dev box: bash .github/scripts/test_gate_full_suite.sh
set -u
here=$(cd "$(dirname "$0")" && pwd)
tmp=$(mktemp -d)
trap 'rm -rf "$tmp"' EXIT
# Mock gh. Asserts the exact endpoint (including head_sha) it is called
# with — an endpoint typo in the gate script fails the test rather than
# silently serving canned data. On the runs endpoint it serves
# $MOCK_DIR/response_<call#>.json, sticking on the highest existing file,
# and exits 1 if none exist (simulates a GitHub API outage). On the pulls
# endpoint it serves $MOCK_DIR/pr.json, defaulting to a 'ready'-labeled PR.
cat > "$tmp/gh" <<'EOF'
#!/usr/bin/env bash
if [ "${1:-}" != "api" ]; then
echo "unexpected gh invocation: $*" >> "$MOCK_DIR/endpoint_error"
exit 2
fi
case "${2:-}" in
"repos/o/r/actions/runs?head_sha=deadbeef&per_page=100")
n=$(( $(cat "$MOCK_DIR/count" 2>/dev/null || echo 0) + 1 ))
echo "$n" > "$MOCK_DIR/count"
while [ "$n" -gt 0 ]; do
if [ -f "$MOCK_DIR/response_$n.json" ]; then
cat "$MOCK_DIR/response_$n.json"
exit 0
fi
n=$(( n - 1 ))
done
echo "api outage" >&2
exit 1
;;
"repos/o/r/pulls/42")
if [ -f "$MOCK_DIR/pr.json" ]; then
cat "$MOCK_DIR/pr.json"
else
echo '{"labels": [{"name": "ready"}]}'
fi
;;
*)
echo "unexpected gh endpoint: $2" >> "$MOCK_DIR/endpoint_error"
exit 2
;;
esac
EOF
chmod +x "$tmp/gh"
PC_OK='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "success"}'
PC_BAD='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "failure"}'
PC_PENDING='{"name": "pre-commit", "id": 1, "status": "in_progress", "conclusion": null}'
DOCS_OK='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "success"}'
DOCS_BAD='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "failure"}'
DOCS_CANCELLED='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "cancelled"}'
OTHER='{"name": "Trigger Full Suite", "id": 3, "status": "in_progress", "conclusion": null}'
NULL_NAME='{"name": null, "id": 4, "status": "completed", "conclusion": "failure"}'
PC_OK_RERUN='{"name": "pre-commit", "id": 5, "status": "completed", "conclusion": "success"}'
fails=0
want_log="" # optional: expect() also greps out.log for this regex, then resets
pr_json="" # optional: served for the pulls (label re-check) endpoint, then resets
raw_body="" # optional: serve responses verbatim instead of wrapping in workflow_runs
expect() { # <name> <expected-exit> <response json>...
local name=$1 want=$2 dir i=1
shift 2
dir=$(mktemp -d "$tmp/test_XXXXXX")
for body in "$@"; do
if [ -n "$raw_body" ]; then
printf '%s' "$body" > "$dir/response_$i.json"
else
printf '{"workflow_runs": [%s]}' "$body" > "$dir/response_$i.json"
fi
i=$(( i + 1 ))
done
[ -n "$pr_json" ] && printf '%s' "$pr_json" > "$dir/pr.json"
( export PATH="$tmp:$PATH" MOCK_DIR="$dir" PR_SHA=deadbeef PR_NUMBER=42 \
GITHUB_REPOSITORY=o/r POLL_SECS=0 GRACE_SECS=1 MAX_WAIT_SECS=3
bash "$here/gate_full_suite.sh" > "$dir/out.log" 2>&1 )
local rc=$?
if [ "$rc" -ne "$want" ]; then
echo "FAIL: $name (exit $rc, want $want)"
cat "$dir/out.log"
fails=1
elif [ -f "$dir/endpoint_error" ]; then
echo "FAIL: $name (mock gh got an unexpected call)"
cat "$dir/endpoint_error"
fails=1
elif [ -n "$want_log" ] && ! grep -Eq "$want_log" "$dir/out.log"; then
echo "FAIL: $name (log does not match: $want_log)"
cat "$dir/out.log"
fails=1
else
echo "ok: $name"
fi
want_log="" pr_json="" raw_body=""
}
expect "both green -> proceed" 0 "$PC_OK, $DOCS_OK, $OTHER, $NULL_NAME"
expect "docs build failed -> blocked" 1 "$PC_OK, $DOCS_BAD"
expect "pre-commit failed -> blocked" 1 "$PC_BAD"
expect "pending then green -> proceed" 0 "$PC_PENDING" "$PC_OK, $DOCS_OK"
want_log="never appeared.*Deploy Documentation"
expect "docs run absent (path-filtered) -> proceed after grace" 0 "$PC_OK"
expect "API outage -> fail open" 0
want_log="FAILING OPEN"
expect "pending past MAX_WAIT -> fail open" 0 "$PC_PENDING"
want_log="FAILING OPEN"
expect "unrelated runs only -> no grace, fail open at MAX_WAIT" 0 "$OTHER"
expect "cancelled docs then green -> proceed" 0 \
"$PC_OK, $DOCS_CANCELLED" "$PC_OK, $DOCS_OK"
want_log="FAILING OPEN"
expect "cancelled docs forever -> fail open at MAX_WAIT" 0 "$PC_OK, $DOCS_CANCELLED"
want_log="FAILING OPEN"
expect "pre-commit absent -> no grace, fail open at MAX_WAIT" 0 "$DOCS_OK"
expect "duplicate run names -> latest wins" 0 "$PC_BAD, $PC_OK_RERUN, $DOCS_OK"
raw_body=1
expect "garbage response body -> fail open" 0 "this is not json"
pr_json='{"labels": [{"name": "other"}]}'
expect "ready label removed mid-gate -> blocked" 1 "$PC_OK, $DOCS_OK"
exit "$fails"
-3
View File
@@ -47,6 +47,3 @@ jobs:
- uses: pre-commit/action@v3.0.1
with:
extra_args: --all-files --hook-stage manual
# After pre-commit so a self-test failure cannot mask lint failures.
- name: Full-suite gate self-test
run: bash .github/scripts/test_gate_full_suite.sh
+10 -17
View File
@@ -52,7 +52,6 @@ jobs:
core.setOutput('pr_sha', pr.head.sha);
core.setOutput('pr_branch', pr.head.ref);
core.setOutput('pr_number', String(prNumber));
core.setOutput('pr_title', pr.title);
- name: Trigger Full Suite
if: steps.perm.outputs.has_write == 'true'
@@ -61,7 +60,6 @@ jobs:
PR_SHA: ${{ steps.label.outputs.pr_sha }}
PR_BRANCH: ${{ steps.label.outputs.pr_branch }}
PR_NUMBER: ${{ steps.label.outputs.pr_number }}
PR_TITLE: ${{ steps.label.outputs.pr_title }}
BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
run: |
@@ -73,7 +71,6 @@ jobs:
--arg commit "$PR_SHA" \
--arg branch "$PR_BRANCH" \
--arg message "Full Suite for PR #${PR_NUMBER} (via /merge)" \
--arg pr_title "$PR_TITLE" \
--argjson pr_id "$PR_NUMBER" \
'{
commit: $commit,
@@ -83,12 +80,11 @@ jobs:
pull_request_id: $pr_id,
pull_request_base_branch: "main",
env: {
TEST_SCOPE: "full",
FULL_SUITE: "true",
PR_NUMBER: ($pr_id | tostring),
PR_TITLE: $pr_title
}
}')"
TEST_SCOPE: "full",
FULL_SUITE: "true",
PR_NUMBER: ($pr_id | tostring)
}
}')"
parse-command:
if: >-
@@ -245,7 +241,6 @@ jobs:
TEST_SCOPE: ${{ needs.parse-command.outputs.test_scope }}
FULL_SUITE: ${{ needs.parse-command.outputs.full_suite }}
TEST_TYPE: ${{ needs.parse-command.outputs.test_type }}
PR_TITLE: ${{ github.event.issue.title }}
BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
run: |
@@ -262,7 +257,6 @@ jobs:
--arg full_suite "$FULL_SUITE" \
--arg test_type "$TEST_TYPE" \
--arg pr_number "$PR_NUMBER" \
--arg pr_title "$PR_TITLE" \
'{
commit: $commit,
branch: $branch,
@@ -272,9 +266,8 @@ jobs:
pull_request_base_branch: "main",
env: {
TEST_SCOPE: $test_scope,
FULL_SUITE: $full_suite,
TEST_TYPE: $test_type,
PR_NUMBER: $pr_number,
PR_TITLE: $pr_title
}
}')"
FULL_SUITE: $full_suite,
TEST_TYPE: $test_type,
PR_NUMBER: $pr_number
}
}')"
+1 -21
View File
@@ -7,7 +7,6 @@ on:
permissions:
contents: read
pull-requests: read
actions: read
concurrency:
group: full-suite-${{ github.event.pull_request.number }}
@@ -19,8 +18,6 @@ jobs:
(github.event.action == 'labeled' && github.event.label.name == 'ready')
|| github.event.action == 'synchronize'
runs-on: ubuntu-latest
# Gate below may wait for cheap checks (up to MAX_WAIT_SECS = 25 min).
timeout-minutes: 35
steps:
- name: Check ready label
id: check
@@ -52,20 +49,6 @@ jobs:
"https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel"
done
# Checks out the BASE branch (default for pull_request_target), so PR
# authors cannot tamper with the gate script.
- name: Checkout gate script
if: steps.check.outputs.has_ready == 'true'
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Wait for pre-commit and docs build
if: steps.check.outputs.has_ready == 'true'
env:
GH_TOKEN: ${{ github.token }}
PR_SHA: ${{ github.event.pull_request.head.sha }}
PR_NUMBER: ${{ github.event.pull_request.number }}
run: bash .github/scripts/gate_full_suite.sh
- name: Trigger Buildkite Full Suite
if: steps.check.outputs.has_ready == 'true'
env:
@@ -73,7 +56,6 @@ jobs:
PR_SHA: ${{ github.event.pull_request.head.sha }}
PR_BRANCH: ${{ github.event.pull_request.head.ref }}
PR_NUMBER: ${{ github.event.pull_request.number }}
PR_TITLE: ${{ github.event.pull_request.title }}
BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
run: |
@@ -85,7 +67,6 @@ jobs:
--arg commit "$PR_SHA" \
--arg branch "$PR_BRANCH" \
--arg message "Full Suite for PR #${PR_NUMBER}" \
--arg pr_title "$PR_TITLE" \
--argjson pr_id "$PR_NUMBER" \
'{
commit: $commit,
@@ -97,7 +78,6 @@ jobs:
env: {
TEST_SCOPE: "full",
FULL_SUITE: "true",
PR_NUMBER: ($pr_id | tostring),
PR_TITLE: $pr_title
PR_NUMBER: ($pr_id | tostring)
}
}')"
+2 -3
View File
@@ -72,7 +72,8 @@ docs/distillation/examples/
# Python pickle files
*.pkl
# Reference videos (negations must come after the catch-all on line below)
# Reference videos
!fastvideo/tests/ssim/reference_videos/**/*.mp4
# Static images
!docs/assets/images/**/*.png
@@ -126,8 +127,6 @@ 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
+5 -5
View File
@@ -67,7 +67,7 @@ MODEL_REGISTRY = {
},
}
DEFAULT_MODEL_ID = "fast-ltx23"
DEFAULT_MODEL_ID = "fast-ltx2"
ACTIVE_MODEL_ID = (os.getenv("DREAMVERSE_MODEL_ID", "").strip() or DEFAULT_MODEL_ID)
if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
@@ -76,11 +76,14 @@ 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 = 6
NUM_INFERENCE_STEPS = 5
JPEG_QUALITY = 100
BATCH_SIZE = 3
@@ -165,9 +168,6 @@ 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)
+1 -5
View File
@@ -1023,11 +1023,7 @@ 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:
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
visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
return _limit_gpu_ids(visible_gpu_ids)
# Auto-detect available GPUs
+1 -3
View File
@@ -206,9 +206,7 @@ def cli() -> None:
args = parser.parse_args()
_install_heartbeat_log_filter()
# 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)
uvicorn.run(app, host=args.host, port=args.port)
if __name__ == "__main__":
+21 -47
View File
@@ -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, SESSION_TIMEOUT_SECONDS
from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP
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,7 +334,6 @@ 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":
@@ -345,8 +344,7 @@ 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(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
and len(curated_prompts) == 0)
generation_paused = bool(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"))
@@ -375,10 +373,7 @@ async def websocket_endpoint(websocket: WebSocket):
prompt_sources_blocked = False
pending_seed_reset = False
pending_seed_reset_reason = ""
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)
pending_simple_submission: PromptSubmission | None = None
single_clip_waiting_for_request = False
rollout_waiting_for_rewrite = False
initial_rollout_waiting_for_rewrite = generation_paused
@@ -398,26 +393,14 @@ 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":
_resolve_generation_segment_cap(
single_clip_mode=single_clip_mode,
cap=GENERATION_SEGMENT_CAP,
manual_continuation_mode=manual_continuation_mode,
),
"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,
})
if seed_reason == "init":
await ws_send_json({
@@ -510,7 +493,6 @@ 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
@@ -555,7 +537,6 @@ 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)
@@ -563,14 +544,10 @@ async def websocket_endpoint(websocket: WebSocket):
segment_idx = 0
pending_seed_reset = False
pending_seed_reset_reason = ""
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)
pending_simple_submission = None
single_clip_waiting_for_request = False
rollout_waiting_for_rewrite = False
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
and len(curated_prompts) == 0)
generation_paused = bool(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
@@ -991,11 +968,10 @@ 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)
@@ -1005,8 +981,8 @@ async def websocket_endpoint(websocket: WebSocket):
await asyncio.sleep(0.05)
continue
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):
if (not single_clip_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
@@ -1241,9 +1217,7 @@ def cli() -> None:
print(f"Starting mock server with {LATENCY_MS}ms latency on port {args.port}")
_install_heartbeat_log_filter()
# 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)
uvicorn.run(app, host="0.0.0.0", port=args.port)
if __name__ == "__main__":
+18 -60
View File
@@ -313,33 +313,6 @@ 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:
@@ -357,7 +330,6 @@ 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:
@@ -365,37 +337,21 @@ def _parse_json_response(content: str) -> dict[str, Any]:
try:
parsed = json.loads(block)
if isinstance(parsed, dict):
last_fenced = parsed
return parsed
except json.JSONDecodeError:
continue
if last_fenced is not None:
return last_fenced
# 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.
# Fall back to scanning for the first decodable JSON object in free-form text.
decoder = json.JSONDecoder()
last_parsed: dict[str, Any] | None = None
pos = 0
while (idx := text.find("{", pos)) != -1:
for idx, char in enumerate(text):
if char != "{":
continue
try:
parsed, consumed = decoder.raw_decode(text[idx:])
parsed, _ = 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):
last_parsed = parsed
if last_parsed is not None:
return last_parsed
return parsed
raise ValueError("No JSON object found in assistant response.")
@@ -425,7 +381,7 @@ def _format_locked_segments(locked_segments: list[str]) -> str:
class PromptEnhancer:
def __init__(self) -> None:
def __init__(self):
self.provider = PROMPT_PROVIDER
self.provider_label = _resolve_provider_label(PROMPT_PROVIDER)
self.api_key = PROMPT_API_KEY
@@ -1461,12 +1417,15 @@ 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."),
"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
),
}
t0 = time.perf_counter()
@@ -1498,7 +1457,6 @@ 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:
@@ -1581,7 +1539,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: {"next_prompt": "<your segment description here>"}.' # noqa: E501
'Respond with valid JSON only as {"next_prompt": "..."}.' # noqa: E501
),
}
@@ -20,7 +20,6 @@ from __future__ import annotations
# mypy: ignore-errors
import asyncio
import os
import time
import uuid
from typing import TYPE_CHECKING
@@ -53,14 +52,6 @@ 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."""
@@ -207,7 +198,6 @@ 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()
@@ -292,9 +282,6 @@ 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
@@ -304,20 +291,13 @@ class SessionController:
generation_cap_blocked = False
auto_extension_blocked_segment_idx: int | None = None
prompt_sources_drained_logged = False
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
and len(curated_prompts) == 0)
generation_paused = bool(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
# 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)
pending_simple_prompt_submission: PromptSubmission | None = None
single_clip_waiting_for_request = False
rollout_waiting_for_rewrite = False
initial_rollout_waiting_for_rewrite = generation_paused
@@ -326,7 +306,6 @@ 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
@@ -445,7 +424,6 @@ 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
@@ -480,7 +458,6 @@ 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:
@@ -533,7 +510,6 @@ 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)
@@ -554,15 +530,10 @@ class SessionController:
generation_cap_blocked = False
auto_extension_blocked_segment_idx = None
prompt_sources_drained_logged = False
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)
pending_simple_prompt_submission = None
single_clip_waiting_for_request = False
rollout_waiting_for_rewrite = False
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt
and not single_clip_mode and len(curated_prompts) == 0)
generation_paused = bool(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
@@ -971,7 +942,6 @@ class SessionController:
_resolve_generation_segment_cap(
single_clip_mode=single_clip_mode,
cap=GENERATION_SEGMENT_CAP,
manual_continuation_mode=manual_continuation_mode,
),
})
continue
@@ -1035,9 +1005,11 @@ class SessionController:
continue
async def prompt_worker_loop():
nonlocal prompt_enhancement_inflight
async def process_submission(submission: PromptSubmission) -> None:
while not stop_event.is_set():
try:
submission = await asyncio.wait_for(raw_prompt_queue.get(), timeout=0.1)
except asyncio.TimeoutError:
continue
_main_print("INFO", f"Received user prompt for enhancement: {submission.raw_prompt}")
prompt_id = submission.prompt_id
raw_prompt = submission.raw_prompt
@@ -1060,9 +1032,8 @@ class SessionController:
await ws_send_json({
"type": "error",
"message": blocked_raw_prompt_error,
"prompt_id": prompt_id,
})
return
continue
await log_event(
"enhance_request",
{
@@ -1122,9 +1093,8 @@ class SessionController:
await ws_send_json({
"type": "error",
"message": blocked_final_prompt_error,
"prompt_id": prompt_id,
})
return
continue
if result.fallback_used or not final_prompt:
source = "user_enhancement_failed"
_main_print(
@@ -1143,7 +1113,7 @@ class SessionController:
})
# Enhancement is strict JSON-only; do not enqueue raw
# prompt when enhancement fails.
return
continue
else:
source = "user_enhanced"
await ws_send_json({
@@ -1160,7 +1130,6 @@ class SessionController:
source=source,
fallback_used=result.fallback_used,
loop_iteration=loop_iteration,
raw_prompt=raw_prompt,
))
else:
await ready_prompt_queue.put(
@@ -1170,7 +1139,6 @@ class SessionController:
source="user_raw",
fallback_used=False,
loop_iteration=loop_iteration,
raw_prompt=raw_prompt,
))
await ws_send_json({
"type": "prompt_ready",
@@ -1180,21 +1148,6 @@ 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(),
@@ -1347,7 +1300,6 @@ 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({
@@ -1407,7 +1359,6 @@ 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":
@@ -1437,9 +1388,8 @@ class SessionController:
await raw_prompt_queue.put(pending_simple_prompt_submission)
pending_simple_prompt_submission = None
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):
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):
loop_generation_enabled = False
rollout_waiting_for_rewrite = True
_main_print(
@@ -1592,10 +1542,7 @@ class SessionController:
if single_clip_mode:
await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0)
continue
# 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:
if not prompt_sources_drained_logged:
snapshot = queue_snapshot()
_main_print(
"WARN",
@@ -1634,24 +1581,6 @@ 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
@@ -19,4 +19,3 @@ class ReadyPrompt:
fallback_used: bool = False
seed_prompt_index: int | None = None
loop_iteration: int | None = None
raw_prompt: str | None = None
@@ -150,22 +150,12 @@ 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)
@@ -75,13 +75,12 @@ 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, **kwargs) -> None:
def run(app, host: str, port: int) -> None:
calls.append(
{
"app": app,
"host": host,
"port": port,
**kwargs,
}
)
@@ -105,7 +104,6 @@ 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,
}
]
@@ -123,7 +121,6 @@ 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,
}
]
@@ -150,7 +147,6 @@ 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,
}
]
@@ -170,7 +166,6 @@ 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
@@ -290,152 +290,6 @@ 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
@@ -6,7 +6,6 @@ import os
import re
import time
import pytest
os.environ.setdefault("CEREBRAS_API_KEY", "dummy")
os.environ.setdefault("GROQ_API_KEY", "dummy")
@@ -206,59 +205,6 @@ 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")
+2 -3
View File
@@ -15,6 +15,5 @@ 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, 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
def _resolve_generation_segment_cap(*, single_clip_mode: bool, cap: int) -> int:
return 0 if single_clip_mode else cap
@@ -489,7 +489,7 @@ class VideoGenerationWorker:
num_inference_steps=NUM_INFERENCE_STEPS,
guidance_scale=1.0,
seed=10,
ltx2_image_crf=(33.0 if image_path and segment_idx == 1 else 0.0),
ltx2_image_crf=0.0,
image_path=image_path if segment_idx == 1 else None,
return_continuation_state=False,
)
-3
View File
@@ -29,9 +29,6 @@ test = [
dreamverse-server = "dreamverse.server_entry:cli"
dreamverse-mock-server = "dreamverse.mock_server:cli"
[tool.setuptools.packages.find]
include = ["dreamverse*"]
[tool.uv]
package = false
@@ -1,67 +0,0 @@
#!/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"
@@ -1,30 +0,0 @@
#!/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
@@ -1,105 +0,0 @@
#!/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"
@@ -110,7 +110,7 @@ default_request:
fps: 24 # internal: gpu_pool.py:85 TARGET_FPS
streaming:
# internal: config.py SESSION_TIMEOUT_SECONDS (env DREAMVERSE_SESSION_TIMEOUT_SECONDS, default 300)
# internal: config.py:33 SESSION_TIMEOUT_SECONDS = 300
session_timeout_seconds: 300
# internal: config.py:282-284 GENERATION_SEGMENT_CAP default 6
generation_segment_cap: 6
+13 -125
View File
@@ -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, { SceneHistoryList } from "@/components/Workspace";
import Workspace 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, USER_PROMPT_SOURCES } from "@/stores/stream";
import { createStreamStore } 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 with";
const HERO_TEXT = "Direct scenes in seconds";
function HeroTagline() {
const ref = useRef<HTMLHeadingElement>(null);
@@ -166,8 +166,6 @@ 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>
);
}
@@ -239,7 +237,6 @@ export default function Page() {
enhancementEnabled,
promptExtensionError,
autoExtensionEnabled,
manualContinuationMode,
autoExtensionTimeoutHint,
loopGenerationEnabled,
generationPaused,
@@ -252,8 +249,6 @@ export default function Page() {
livePromptRewriteMode,
sessionExpired,
projectResetPending,
waitingForSegmentPrompt,
generatingNextScene,
} = sessionState;
const {
@@ -328,11 +323,6 @@ 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);
@@ -432,11 +422,6 @@ 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]);
@@ -447,39 +432,11 @@ export default function Page() {
const canDownloadVideo = useMemo(() => {
const currentActiveClip = activeClip as Record<string, any> | null;
if (currentActiveClip?.blob instanceof Blob) return true;
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]);
return (completedClips as Record<string, any>[]).some((clip) => clip?.blob instanceof Blob);
}, [activeClip, completedClips]);
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],
);
@@ -1464,8 +1421,7 @@ export default function Page() {
if (!prompt) return;
lastSubmitTimeRef.current = now;
const rCAR = !uiStore.get().devtoolsMode && !uiStore.get().demoMode;
const inManualMode = sessionStore.get().manualContinuationMode;
if (!inManualMode && (rCAR || (sessionStore.get().livePromptRewriteMode && !uiStore.get().demoMode))) {
if (rCAR || (sessionStore.get().livePromptRewriteMode && !uiStore.get().demoMode)) {
if (rewriteStore.get().rewritingSeedPrompts) return;
const rewriteSourcePromptWindowPrompts = getActivePromptWindowPrompts();
const nextPendingClip = {
@@ -1506,11 +1462,6 @@ 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({
@@ -1524,14 +1475,7 @@ export default function Page() {
activeClipId: shouldUseArchivedPlaybackFallback() ? streamStore.get().activeClipId : "",
activePlaybackStartTime: shouldUseArchivedPlaybackFallback() ? streamStore.get().activePlaybackStartTime : 0,
});
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 } : {}),
});
sessionStore.patch({ livePromptDraft: "" });
}
function setLivePromptRewriteMode(enabled: boolean) {
@@ -1606,14 +1550,6 @@ 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),
@@ -1789,9 +1725,6 @@ 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();
@@ -1804,7 +1737,6 @@ export default function Page() {
function resetToProjectLobbyState() {
setVideoMuted(true);
pendingInitialPromptRef.current = "";
pendingInitialPromptIdRef.current = "";
sessionStore.patch({
sessionStarted: false,
livePromptDraft: "",
@@ -1818,9 +1750,6 @@ export default function Page() {
sessionNotice: "",
sessionExpired: false,
projectResetPending: false,
manualContinuationMode: true,
waitingForSegmentPrompt: false,
generatingNextScene: false,
});
rewriteStore.resetSessionState();
streamStore.resetSessionState();
@@ -1831,14 +1760,7 @@ export default function Page() {
}
function buildProjectInitPayload(type: "session_init_v2" | "project_init_v1") {
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);
}
const segmentPrompts = getSessionInitPrompts();
setSeedPrompts(segmentPrompts);
return {
type,
@@ -1846,17 +1768,11 @@ export default function Page() {
preset_label: getInitialPresetLabel(),
curated_prompts: segmentPrompts,
initial_rollout_prompt: normalizeInitialPrompt(pendingInitialPromptRef.current),
// 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,
initial_image: null,
single_clip_mode: false,
enhancement_enabled: sessionStore.get().enhancementEnabled,
auto_extension_enabled: manualMode ? false : sessionStore.get().autoExtensionEnabled,
loop_generation_enabled: manualMode ? false : sessionStore.get().loopGenerationEnabled,
manual_continuation_mode: manualMode,
auto_extension_enabled: sessionStore.get().autoExtensionEnabled,
loop_generation_enabled: sessionStore.get().loopGenerationEnabled,
};
}
@@ -2035,11 +1951,7 @@ 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,
@@ -2054,9 +1966,6 @@ export default function Page() {
autoExtensionTimeoutHint: "",
generationPaused: false,
projectResetPending: false,
manualContinuationMode: nextManualContinuationMode,
waitingForSegmentPrompt: false,
generatingNextScene: false,
});
resetPlaybackState();
streamStore.patch({
@@ -2070,15 +1979,12 @@ export default function Page() {
selectedHistoryId: "",
});
rewriteStore.resetSessionState();
pendingInitialPromptIdRef.current = initialPrompt ? makePromptId() : "";
if (initialPrompt) {
addPromptEvent({
promptId: pendingInitialPromptIdRef.current,
promptId: makePromptId(),
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());
@@ -2594,7 +2500,6 @@ 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}
@@ -2607,7 +2512,6 @@ export default function Page() {
onEnhancementToggle={handleEnhancementToggle}
onCuratedPromptLimitChange={handleCuratedPromptLimitChange}
onAutoExtensionToggle={handleAutoExtensionToggle}
onManualContinuationToggle={handleManualContinuationToggle}
onLoopToggle={handleLoopGenerationToggle}
onJoin={joinSession}
onLeave={leaveSession}
@@ -2737,7 +2641,6 @@ 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 && (
<>
@@ -2821,8 +2724,6 @@ export default function Page() {
showLivePlayback={showLivePlayback}
defaultMuted={videoMuted}
canDownload={canDownloadVideo}
waitingForSegmentPrompt={waitingForSegmentPrompt as boolean}
generatingNextScene={generatingNextScene as boolean}
onPlaying={markFirstFrameRendered}
onDownload={handleDownloadVideo}
/>
@@ -2833,7 +2734,6 @@ 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}
@@ -2880,17 +2780,11 @@ 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}
@@ -2904,12 +2798,6 @@ export default function Page() {
</motion.div>
</div>
</div>
{manualContinuationMode && (
<div className="px-4 sm:px-6">
<SceneHistoryList sceneHistory={steeringScenes as any[]} />
</div>
)}
</div>
</main>
);
}
+6 -96
View File
@@ -2,16 +2,13 @@
import React, { useRef, useState, useCallback, useEffect } from "react";
import Image from "next/image";
import { Film, ArrowUp, X, Loader2, ArrowLeft, ImagePlus } from "lucide-react";
import { Film, ArrowUp, X, Loader2, ArrowLeft } 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;
@@ -25,12 +22,6 @@ 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;
@@ -55,12 +46,6 @@ export default function ChatBar({
sessionNotice = "",
projectResetPending = false,
viewingReadOnly = false,
waitingForSegmentPrompt = false,
manualContinuationEnabled = false,
onModeChange = () => {},
initialImageDataUrl = "",
onImageUpload = () => {},
onImageClear = () => {},
onPresetGenerate = () => {},
onContinuationInput = () => {},
onContinuationKeydown = () => {},
@@ -74,51 +59,15 @@ 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"
: waitingForSegmentPrompt
? "Describe the next scene\u2026"
: !sessionStarted
? "What video are you imagining?"
: "What do you want to edit?";
: !sessionStarted
? "What video are you imagining?"
: "What do you want to edit?";
const actionLabel = !sessionStarted ? "Generate" : "Rewrite rollout";
const inputRef = useRef<HTMLTextAreaElement>(null);
@@ -318,9 +267,9 @@ export default function ChatBar({
<Button onClick={onStartNewProject} size="sm" className="rounded-full px-5">
New Project
</Button>
<a href="https://haoailab.com/blogs/dreamverse/" target="_blank" rel="noopener noreferrer">
<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
<Button variant="outline" size="sm" className="rounded-full px-5">
Blog
Join Waitlist
</Button>
</a>
</div>
@@ -397,31 +346,6 @@ 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",
@@ -435,7 +359,6 @@ export default function ChatBar({
value={continuationDraft}
onChange={onContinuationInput}
onKeyDown={handleKeyDown}
onPaste={handleImagePaste}
placeholder={sttBusy ? "Listening\u2026" : messagePlaceholder}
maxLength={PROMPT_MAX_LENGTH}
disabled={isBusy || sttBusy}
@@ -446,19 +369,6 @@ 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}
@@ -8,6 +8,7 @@ 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 {
@@ -28,13 +29,13 @@ export default function Header({ timeLeft = null, formatTime = (seconds) => `${s
<SidePanelOpenFilled size={20} />
</Button>
)}
<a href="/" title="FastVideo home">
<a href={FASTVIDEO_REPO_URL} target="_blank" rel="noopener noreferrer" title="FastVideo on GitHub">
<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={FASTVIDEO_BLOG_URL} target="_blank" rel="noopener noreferrer">
<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
<Button variant="outline" size="sm" className="gap-1.5 rounded-full px-3 text-xs">
Blog
Join Waitlist
<ExternalLink className="size-3 opacity-60" />
</Button>
</a>
@@ -52,9 +53,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={FASTVIDEO_BLOG_URL} target="_blank" rel="noopener noreferrer">
<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
<Button variant="outline" size="sm" className="gap-1.5 rounded-full px-3 text-xs">
Blog
Join Waitlist
<ExternalLink className="size-3 opacity-60" />
</Button>
</a>
@@ -3,7 +3,7 @@
import React, { useState, useEffect, useRef, useCallback } from "react";
import { cn } from "@/lib/utils";
import { PlayFilledAlt } from "@carbon/icons-react";
import { Check, ChevronDown, Download, Loader2, Share } from "lucide-react";
import { Download, Loader2, Share } from "lucide-react";
import { Button } from "@/components/ui/button";
interface VideoPlayerProps {
videoRef?: React.RefCallback<HTMLVideoElement>;
@@ -21,8 +21,6 @@ interface VideoPlayerProps {
showLivePlayback?: boolean;
defaultMuted?: boolean;
rewritePending?: boolean;
waitingForSegmentPrompt?: boolean;
generatingNextScene?: boolean;
onPlaying?: () => void;
onDownload?: () => void;
}
@@ -59,8 +57,6 @@ export default function VideoPlayer({
showLivePlayback = true,
defaultMuted = true,
rewritePending = false,
waitingForSegmentPrompt = false,
generatingNextScene = false,
onPlaying = () => {},
onDownload,
}: VideoPlayerProps) {
@@ -92,137 +88,6 @@ 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">
@@ -239,7 +104,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 && !waitingForSegmentPrompt && loadingAnimation ? (
) : !avPlaybackStarted && !mediaAppendError && !inQueue && 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" />
@@ -249,35 +114,6 @@ 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" />
@@ -3,90 +3,11 @@ 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, Clapperboard, Lightbulb, Pencil } from "lucide-react";
import { Check, 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;
@@ -98,7 +19,6 @@ interface WorkspaceProps {
selectedClipId?: string;
selectedEntryKey?: string;
originalClipId?: string;
manualMode?: boolean;
}
function normalizeText(value: any): string {
@@ -298,7 +218,7 @@ function ChromaGradient({ sessionStarted = false }: { sessionStarted?: boolean }
);
}
export default function Workspace({ promptEvents = [], currentThumbnail = null, originalLabel = "", sessionStarted = false, onSelectOriginal, onSelectEvent, onSelectCurrent, selectedClipId, selectedEntryKey: selectedEntryKeyProp, originalClipId = "", manualMode = false }: WorkspaceProps) {
export default function Workspace({ promptEvents = [], currentThumbnail = null, originalLabel = "", sessionStarted = false, onSelectOriginal, onSelectEvent, onSelectCurrent, selectedClipId, selectedEntryKey: selectedEntryKeyProp, originalClipId = "" }: WorkspaceProps) {
const bottomSentinelRef = useRef<HTMLDivElement>(null);
const topSentinelRef = useRef<HTMLDivElement>(null);
const [showTopFade, setShowTopFade] = useState(false);
@@ -339,14 +259,6 @@ 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} />
@@ -34,7 +34,6 @@ interface DevtoolsComposerProps {
demoMode?: boolean;
enhancementEnabled?: boolean;
autoExtensionEnabled?: boolean;
manualContinuationEnabled?: boolean;
loopGenerationEnabled?: boolean;
curatedPromptLimit?: number;
maxCuratedPromptCount?: number;
@@ -50,7 +49,6 @@ 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;
@@ -73,7 +71,6 @@ export default function DevtoolsComposer({
demoMode = false,
enhancementEnabled = true,
autoExtensionEnabled = false,
manualContinuationEnabled = false,
loopGenerationEnabled = false,
curatedPromptLimit = 0,
maxCuratedPromptCount = 0,
@@ -89,7 +86,6 @@ export default function DevtoolsComposer({
onEnhancementToggle = () => {},
onCuratedPromptLimitChange = () => {},
onAutoExtensionToggle = () => {},
onManualContinuationToggle = () => {},
onLoopToggle = () => {},
onLivePromptModeToggle = () => {},
onSpeechTranscript,
@@ -332,28 +328,6 @@ 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"
@@ -21,7 +21,6 @@ interface DevtoolsShellProps {
selectedPresetId?: string;
enhancementEnabled?: boolean;
autoExtensionEnabled?: boolean;
manualContinuationEnabled?: boolean;
loopGenerationEnabled?: boolean;
canJoinSession?: boolean;
canSubmitContinuation?: boolean;
@@ -35,7 +34,6 @@ 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;
@@ -130,7 +128,6 @@ export default function DevtoolsShell({
selectedPresetId = '',
enhancementEnabled = true,
autoExtensionEnabled = false,
manualContinuationEnabled = false,
loopGenerationEnabled = false,
canJoinSession = false,
canSubmitContinuation = false,
@@ -144,7 +141,6 @@ export default function DevtoolsShell({
onEnhancementToggle = () => {},
onCuratedPromptLimitChange = () => {},
onAutoExtensionToggle = () => {},
onManualContinuationToggle = () => {},
onLoopToggle = () => {},
onJoin = () => {},
onLeave = () => {},
@@ -288,7 +284,6 @@ export default function DevtoolsShell({
demoMode={demoMode}
enhancementEnabled={enhancementEnabled}
autoExtensionEnabled={autoExtensionEnabled}
manualContinuationEnabled={manualContinuationEnabled}
loopGenerationEnabled={loopGenerationEnabled}
curatedPromptLimit={curatedPromptLimit}
maxCuratedPromptCount={maxCuratedPromptCount}
@@ -304,7 +299,6 @@ export default function DevtoolsShell({
onEnhancementToggle={onEnhancementToggle}
onCuratedPromptLimitChange={onCuratedPromptLimitChange}
onAutoExtensionToggle={onAutoExtensionToggle}
onManualContinuationToggle={onManualContinuationToggle}
onLoopToggle={onLoopToggle}
onLivePromptModeToggle={onLivePromptModeToggle}
onSpeechTranscript={onSpeechTranscript}
@@ -57,32 +57,4 @@ 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');
});
});
+1 -15
View File
@@ -16,19 +16,5 @@ export function prependPromptEvent(
events: Record<string, any>[],
event: Record<string, any>,
): Record<string, any>[] {
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;
});
return [event, ...events].slice(0, MAX_PROMPT_EVENTS);
}
+1 -165
View File
@@ -1,11 +1,6 @@
import { describe, expect, it } from 'vitest';
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';
import { resolveSessionErrorMessage } from './reducer';
describe('resolveSessionErrorMessage', () => {
it('returns a dedicated message for IP session limit errors', () => {
@@ -24,162 +19,3 @@ 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);
});
});
+9 -50
View File
@@ -99,19 +99,7 @@ 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;
@@ -255,32 +243,19 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
}
case "prompt/sources_blocked":
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",
});
}
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":
// Deliberately does NOT touch generatingNextScene: toggling auto extension
// starts no generation.
sessionStore.patch({ autoExtensionTimeoutHint: "", waitingForSegmentPrompt: false });
console.log("[AutoExtensionUpdated]", {
enabled: sessionStore.get().autoExtensionEnabled,
});
sessionStore.patch({ autoExtensionTimeoutHint: "" });
if (event.type === "session/auto_extension_updated") {
console.log("[AutoExtensionUpdated]", {
enabled: sessionStore.get().autoExtensionEnabled,
});
}
return;
case "segment/step_complete":
@@ -302,7 +277,6 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
projectResetPending: false,
sessionExpired: true,
sessionNotice: "",
generatingNextScene: false,
});
console.log("Session timed out");
clearCountdownInterval();
@@ -380,8 +354,6 @@ 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,
@@ -499,30 +471,17 @@ 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,
@@ -24,9 +24,6 @@ export interface SessionState {
livePromptRewriteMode: boolean;
sessionExpired: boolean;
projectResetPending: boolean;
manualContinuationMode: boolean;
waitingForSegmentPrompt: boolean;
generatingNextScene: boolean;
}
const DEFAULT_SESSION_STATE: SessionState = {
@@ -52,12 +49,6 @@ 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> & {
+1 -18
View File
@@ -1,18 +1,5 @@
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;
@@ -172,14 +159,10 @@ export function createStreamStore(initialState: Partial<StreamState> = {}): Stre
loopIteration: typeof loopIteration === "number" ? loopIteration : null,
};
const nextHistory = [entry, ...state.promptHistory];
return {
...state,
promptHistoryCounter: nextCounter,
promptHistory: nextHistory.length > PROMPT_HISTORY_CAP
? nextHistory.filter((item, index) =>
index < PROMPT_HISTORY_CAP || !USER_PROMPT_SOURCES.has(String(item.source || "")))
: nextHistory,
promptHistory: [entry, ...state.promptHistory].slice(0, 120),
selectedHistoryId: state.selectedHistoryId || (entry.id as string),
};
});
+3 -8
View File
@@ -104,11 +104,6 @@ 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`
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.
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`.
@@ -1,22 +0,0 @@
{
"alpha_yaw": 0.08734091699186919,
"alpha_pitch": 0.08169667696275307,
"alpha_turn": 5.724587470723463e-17,
"beta_fwd": 0.02842768078408099,
"beta_strafe": 0.022531015077067108,
"focal_length": 457.0,
"frame_shape": [
352,
640
],
"calibrated_from": [
"1_wasd_only",
"camera",
"camera4hold_alpha1",
"fully_random",
"wasdonly_alpha1",
"wasd4holdrandview_simple_1key1mouse1"
],
"residual_rms": 15.890399609478676,
"n_equations": 4125232
}
-7
View File
@@ -99,13 +99,6 @@ status.
Full Suite is also path-filtered. It validates broader behavior before Mergify
can merge a PR.
A `ready`-labeled PR does not hit Buildkite immediately:
`ci-trigger-full-suite.yml` first runs `.github/scripts/gate_full_suite.sh`,
which waits for the cheap Tier-1 checks (pre-commit, docs build) on the PR
head. A red cheap check blocks the suite (fail closed; the next push re-arms
it), while a GitHub outage or a >25 min wait lets it run anyway (fail open).
`/test full` bypasses the gate.
| Buildkite label | `TEST_TYPE` | Main watched paths |
|---|---|---|
| SSIM Tests | `ssim` | `fastvideo/**/*.py`, `pyproject.toml`, `docker/Dockerfile` |
+55 -143
View File
@@ -12,8 +12,7 @@ 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 comparable
identity.
then compare against the historical baseline for the same model and GPU.
## Quick start (local)
@@ -80,11 +79,10 @@ 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:
`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.
normalized JSON per `(model_id, gpu_type, run)` tuple. The rolling baseline is
the median of the last 5 successful, baseline-eligible records for that
model+GPU. PR and local records are visible in the dashboard but are not
baseline eligible.
## Planned Coverage
@@ -160,15 +158,13 @@ unrealistic memory growth, and optionally large component-specific slowdowns
even when the rolling baseline is empty. They are hand-set with generous
headroom and almost never need touching.
### Rolling baseline (per comparison cohort)
### Rolling baseline (per `(model_id, gpu_type)`)
`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 `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.
for the same `(model_id, gpu_type)` from the HF dataset, computes the median
for each available metric, and evaluates the current run with the metric's
rolling regression policy. 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,33 +184,10 @@ 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. 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.
agent skill to advance the rolling median.
## Schemas
@@ -230,14 +203,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": 3
"benchmark_version": 2
}
```
`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:
`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:
| Field | Purpose |
|---|---|
@@ -248,26 +221,18 @@ identity fields 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 `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.)
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.)
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 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.
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 keep
the `(model_id, gpu_type)` rolling-baseline comparison. Metric-specific
threshold policies and promoted baselines remain separate follow-ups.
### Raw record (`results/perf_*.json`)
@@ -276,10 +241,10 @@ Written by `test_inference_performance.py`. One file per benchmark run.
```jsonc
{
"benchmark_id": "wan-t2v-1.3b-2gpu",
"result_schema_version": 2,
"config_schema_version": 2,
"workload_id": "wan-t2v",
"variant_id": "1.3b-sp2",
"benchmark_version": 3,
"benchmark_version": 2,
"model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
"device": "NVIDIA L40S",
"num_gpus": 2,
@@ -305,25 +270,18 @@ Written by `test_inference_performance.py`. One file per benchmark run.
}
},
"commit": "<full sha>",
"run_source": "pr",
"branch": "feature/perf-change",
"pr_number": "1234",
"test_scope": "direct",
"build_url": "https://buildkite.example/build",
"build_id": "<buildkite-build-id>",
"job_id": "<buildkite-job-id>",
"timestamp": "2026-05-08T22:00:00+00:00",
"quality_metadata": { "quality_status": "canonical" },
"text_encoder_time_s": 2.141,
"dit_time_s": 8.437,
"vae_decode_time_s": 3.208,
"recipe": {
"recipe_schema_version": 2,
"recipe_schema_version": 1,
"benchmark": {
"benchmark_id": "wan-t2v-1.3b-2gpu",
"workload_id": "wan-t2v",
"variant_id": "1.3b-sp2",
"benchmark_version": 3
"benchmark_version": 2
},
"model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
"init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
@@ -343,9 +301,6 @@ 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",
@@ -354,12 +309,7 @@ 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",
"FASTVIDEO_CONTAINER_IMAGE_REF": "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0@sha256:<digest>"
}
},
"environment_metadata": { "env": { "IMAGE_VERSION": "py3.12-cuda13.0.0" } },
"environment_fingerprint": "env-<sha256-prefix>"
}
```
@@ -372,10 +322,9 @@ result, used as the rolling-baseline source of truth.
```jsonc
{
"model_id": "wan-t2v-1.3b-2gpu",
"result_schema_version": 2,
"workload_id": "wan-t2v",
"variant_id": "1.3b-sp2",
"benchmark_version": 3,
"benchmark_version": 2,
"timestamp": "2026-05-08T22:00:00+00:00",
"commit_sha": "<full sha>",
"gpu_type": "NVIDIA L40S",
@@ -396,41 +345,21 @@ result, used as the rolling-baseline source of truth.
"hardware_profile_id": "hw-<sha256-prefix>",
"software_profile_id": "sw-<sha256-prefix>",
"environment_fingerprint": "env-<sha256-prefix>",
"run_source": "pr",
"branch": "feature/perf-change",
"pr_number": "1234",
"test_scope": "direct",
"build_url": "https://buildkite.example/build",
"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
}
```
### Compatibility with legacy records
Older records in the HF dataset may not have `result_schema_version`,
component timing fields, or v2 identity/profile fields. Records without
`result_schema_version` are treated as v1. The comparator ignores missing or
`null` metrics when computing a median, and the dashboard lists skipped plots
for metric series that have no non-null values. Records missing both
`run_source` and `baseline_eligible` are treated as legacy successful
main/full-suite uploads and remain eligible for rolling baselines.
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.
Older records in the HF dataset may not have component timing fields. The
comparator ignores missing or `null` metrics when computing a median, and the
dashboard lists skipped plots for metric series that have no non-null values.
Records missing both `run_source` and `baseline_eligible` are treated as legacy
successful main/full-suite uploads and remain eligible for rolling baselines.
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.
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.
`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
@@ -448,18 +377,13 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
| `PERF_REPORTS_DIR` | `/root/data/perf_reports` | `compare_baseline.py`, `dashboard.py` | Where the Markdown summary and Plotly HTML get written for Buildkite to pick up. |
| `HF_REPO_ID` | `FastVideo/performance-tracking` | `fastvideo/performance/hf_store.py` | HF dataset repo holding rolling-baseline records. |
| `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_RUN_SOURCE` | inferred | `compare_baseline.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` | Performance pytest exit code. Measured static-threshold failures are attributed per record; otherwise a nonzero code reports an infrastructure error. |
| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
| `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
@@ -470,16 +394,11 @@ point is `fastvideo/tests/modal/pr_test.py:run_performance_tests` and the
Buildkite artifact upload is in
`.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.
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 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.
Each performance build runs pytest first. If that fixed-threshold phase fails,
PR/direct runs skip `compare_baseline.py` because they only upload passing
records. Scheduled-main runs still execute `compare_baseline.py` with
`PERF_PYTEST_RC` set so the failed canonical attempt is visible in normalized
JSON and dashboard history. The dashboard runs best-effort for observability.
When the rolling-baseline phase runs, it emits:
* **Markdown summary** — appended to `$GITHUB_STEP_SUMMARY` when that variable
@@ -487,7 +406,7 @@ When the rolling-baseline phase runs, it emits:
per-benchmark row with current vs. baseline values for latency, throughput,
memory, text encoder time, DiT time, and VAE decode time.
* **Plotly dashboard** — `dashboard_<sha>_<ts>.html` showing time-series for
each metric grouped by comparison cohort.
each metric grouped by `(model_id, gpu_type)`.
* **Normalized records** — `normalized_perf_*.json`, one per benchmark.
Useful as input to the
[`reseed-performance-baseline`](https://github.com/hao-ai-lab/FastVideo/blob/main/.agents/skills/reseed-performance-baseline/SKILL.md)
@@ -538,13 +457,10 @@ 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. 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.
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.
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
@@ -564,12 +480,8 @@ When the rolling-baseline phase runs, it emits:
## Troubleshooting
**`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.
**"No baseline for ... Initializing"** — first run for this `(model_id,
gpu_type)`. Run will pass and (if persisting) seed the first record.
**Persistent failure right after a torch / kernel / image upgrade** —
genuine regression *or* baseline drift. Compare the failing normalized record
-24
View File
@@ -180,30 +180,6 @@ python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
--device-folder L40S_reference_videos
```
### SSIM Bootstrap Mode
Normal SSIM runs are strict: if a reference video or latent is missing, the
test fails. For new-model PRs, CI can run SSIM in bootstrap mode so missing
references are uploaded as draft artifacts for review instead of immediately
blocking on a missing canonical reference.
Buildkite enables SSIM bootstrap mode when either condition is true:
- the PR title or Buildkite message contains `[new-model]`;
- `FASTVIDEO_SSIM_BOOTSTRAP_MODE=1` is set for the Buildkite job.
Bootstrap mode passes `--ssim-bootstrap-mode` to pytest. When a generated
artifact is available, the test uploads it under the `drafts/...` namespace in
the SSIM reference repo and marks that case as expected-failed. After reviewing
the draft, promote it into the canonical reference layout:
```bash
python fastvideo/tests/ssim/reference_videos_cli.py promote-draft \
--quality-tier default \
--device-folder L40S_reference_videos \
--model-id <model_id>
```
## CI Integration
FastVideo CI tests are orchestrated by Buildkite and run on Modal GPU
@@ -458,8 +458,6 @@ 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
-116
View File
@@ -1,116 +0,0 @@
# 🌊 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).
-46
View File
@@ -161,18 +161,6 @@ training:
decay_interval_steps: 0
```
`training.data.data_path` can also mix multiple preprocessed datasets by using a mapping from dataset path to repeat count:
```yaml
training:
data:
data_path:
data/zeldam2-clean: 1
data/multi3d_games: 2
```
The repeat count duplicates that dataset's parquet file list before shuffling/sampling, so the example above trains with roughly twice as much `multi3d_games` exposure as `zeldam2-clean`. Paths are just suggested locations; use any local path that contains a FastVideo preprocessed parquet dataset.
See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
including SwanLab installation and authentication.
@@ -338,40 +326,6 @@ Self-Forcing inherits all DMD2 parameters, plus:
| `enable_gradient_in_rollout` | `true` | Enable backprop through rollout |
| `start_gradient_frame` | `0` | Frame index where gradients begin |
### Streaming Long Tuning
`StreamingLongTuningMethod` extends Self-Forcing for LongLive-style rollouts. It
keeps a streaming state, generates overlapping chunks, and trains only the new
frames while preserving context from earlier chunks.
For the MatrixGame2/Zelda world-model example, self-forcing and long tuning are
separate runs: first train or load the 1k-step self-forcing checkpoint using
`examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml`,
then run
`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
from that checkpoint for the 3k-step streaming long-tuning stage.
```yaml
method:
_target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
streaming_chunk_size: 9
streaming_max_length: 39
streaming_fixed_overlap_latents: 3
streaming_reencode_overlap_anchor: true
streaming_anchor_inject_k: 1
streaming_require_full_blocks: true
multi_phased_distill_schedule:
- stage: streaming_long
start_step: 0
end_step: 3000
num_latent_t: 39
streaming_training: true
```
See
`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
for a complete MatrixGame2/Zelda configuration.
---
## Callbacks
-140
View File
@@ -1,140 +0,0 @@
# 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()
-107
View File
@@ -1,107 +0,0 @@
# 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()
@@ -1,37 +0,0 @@
from fastvideo import VideoGenerator
OUTPUT_PATH = "video_samples_kandinsky5_i2v"
IMAGE_PATH = "assets/girl.png"
def main():
generator = VideoGenerator.from_pretrained(
"kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers",
# "kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"
# "kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"
num_gpus=1,
use_fsdp_inference=False,
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
pin_cpu_memory=True,
# image_encoder_cpu_offload=False,
)
prompt = (
"A woman stands up and walks away"
)
_ = generator.generate_video(
prompt,
image_path=IMAGE_PATH,
output_path=OUTPUT_PATH,
save_video=True,
height=1024,
width=1024,
num_frames=121,
)
if __name__ == "__main__":
main()
@@ -1,37 +0,0 @@
from fastvideo import VideoGenerator
OUTPUT_PATH = "video_samples_kandinsky5_t2v"
def main():
generator = VideoGenerator.from_pretrained(
"kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
# "kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers"
# "kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers"
# "kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers"
num_gpus=1,
use_fsdp_inference=False,
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
pin_cpu_memory=True,
# image_encoder_cpu_offload=False,
)
prompt = (
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
"wide with interest. The playful yet serene atmosphere is complemented by soft "
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
)
_ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True,height=512, width=768, num_frames=121)
prompt2 = (
"A majestic lion strides across the golden savanna, its powerful frame "
"glistening under the warm afternoon sun. The tall grass ripples gently in "
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
"cinematic.")
_ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=512, width=768, num_frames=121)
if __name__ == "__main__":
main()
-120
View File
@@ -1,120 +0,0 @@
# 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()
-69
View File
@@ -1,69 +0,0 @@
# LTX-2.3 distilled inference configs
Ready-to-run `fastvideo generate` run configs for the LTX-2.3
distilled model (`FastVideo/LTX-2.3-Distilled-Diffusers`), covering both
workloads (t2v / i2v), both two-stage step schedules (`5+2`, `8+3` = denoise
+ refine), and four resolutions.
```bash
fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml
```
Each config is self-contained (no preset registry needed): the two-stage
refine is wired via `generator.pipeline.preset_overrides.refine`, and the
base sampling knobs live under `request.sampling`. The refine upsampler
auto-resolves from the model's `spatial_upscaler`.
## Configs
| workload | schedule | resolution (HxW) | file |
|---|---|---|---|
| t2v | 5+2 | 1280x832 | `t2v_5s2_1280x832.yaml` |
| t2v | 5+2 | 1024x1536 | `t2v_5s2_1024x1536.yaml` |
| t2v | 5+2 | 768x1280 | `t2v_5s2_768x1280.yaml` |
| t2v | 5+2 | 512x768 | `t2v_5s2_512x768.yaml` |
| t2v | 8+3 | 1280x832 | `t2v_8s3_1280x832.yaml` |
| t2v | 8+3 | 1024x1536 | `t2v_8s3_1024x1536.yaml` |
| t2v | 8+3 | 768x1280 | `t2v_8s3_768x1280.yaml` |
| t2v | 8+3 | 512x768 | `t2v_8s3_512x768.yaml` |
| i2v | 5+2 | 1280x832 | `i2v_5s2_1280x832.yaml` |
| i2v | 5+2 | 1024x1536 | `i2v_5s2_1024x1536.yaml` |
| i2v | 5+2 | 768x1280 | `i2v_5s2_768x1280.yaml` |
| i2v | 5+2 | 512x768 | `i2v_5s2_512x768.yaml` |
| i2v | 8+3 | 1280x832 | `i2v_8s3_1280x832.yaml` |
| i2v | 8+3 | 1024x1536 | `i2v_8s3_1024x1536.yaml` |
| i2v | 8+3 | 768x1280 | `i2v_8s3_768x1280.yaml` |
| i2v | 8+3 | 512x768 | `i2v_8s3_512x768.yaml` |
## Overriding without editing a file
Dotted overrides (prefixes `generator.` / `request.`) let you tweak any field:
```bash
# swap prompt
fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml \
--request.prompt "a red fox running through fresh snow"
# change output path / gpu count
fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_512x768.yaml \
--request.output.output_path outputs/preview.mp4 \
--generator.engine.num_gpus 4
```
## i2v
The `i2v_*` configs take a first-frame image via
`request.extensions.ltx2_images` (`[[path, frame_offset, weight]]`). Edit the
path in the file, or override it:
```bash
fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1280x832.yaml \
--request.extensions.ltx2_images '[["/data/portrait.jpg", 0, 1.0]]'
```
## Schedules
`5+2` is the fast preview schedule; `8+3` is the higher-quality distilled
recipe. Refine (`preset_overrides.refine.num_inference_steps`) only accepts 2
or 3 steps.
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 5+2 two-stage at 1024x1536.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_1024x1536.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 1024
width: 1536
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_5s2_1024x1536.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 5+2 two-stage at 1280x832.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_1280x832.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 1280
width: 832
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_5s2_1280x832.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 5+2 two-stage at 512x768.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_512x768.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 512
width: 768
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_5s2_512x768.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 5+2 two-stage at 768x1280.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_768x1280.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 768
width: 1280
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_5s2_768x1280.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 8+3 two-stage at 1024x1536.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1024x1536.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 1024
width: 1536
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_8s3_1024x1536.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 8+3 two-stage at 1280x832.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1280x832.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 1280
width: 832
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_8s3_1280x832.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 8+3 two-stage at 512x768.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_512x768.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 512
width: 768
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_8s3_512x768.mp4
save_video: true
@@ -1,38 +0,0 @@
# LTX-2.3 distilled i2v — 8+3 two-stage at 768x1280.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_768x1280.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: i2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
sampling:
height: 768
width: 1280
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
# i2v conditioning — replace the path with your own first-frame image.
extensions:
ltx2_images:
- ["/path/to/your/first_frame.jpg", 0, 1.0]
ltx2_image_crf: 0.0
output:
output_path: outputs/ltx2_3_i2v_8s3_768x1280.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 5+2 two-stage at 1024x1536.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_1024x1536.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 1024
width: 1536
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
output:
output_path: outputs/ltx2_3_t2v_5s2_1024x1536.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 5+2 two-stage at 1280x832.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_1280x832.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 1280
width: 832
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
output:
output_path: outputs/ltx2_3_t2v_5s2_1280x832.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 5+2 two-stage at 512x768.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_512x768.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 512
width: 768
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
output:
output_path: outputs/ltx2_3_t2v_5s2_512x768.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 5+2 two-stage at 768x1280.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_768x1280.yaml
#
# Stage 1 denoises for 5 steps at half resolution; the latents are then
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 2
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 768
width: 1280
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 5
output:
output_path: outputs/ltx2_3_t2v_5s2_768x1280.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 8+3 two-stage at 1024x1536.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1024x1536.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 1024
width: 1536
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
output:
output_path: outputs/ltx2_3_t2v_8s3_1024x1536.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 8+3 two-stage at 1280x832.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 1280
width: 832
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
output:
output_path: outputs/ltx2_3_t2v_8s3_1280x832.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 8+3 two-stage at 512x768.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_512x768.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 512
width: 768
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
output:
output_path: outputs/ltx2_3_t2v_8s3_512x768.mp4
save_video: true
@@ -1,33 +0,0 @@
# LTX-2.3 distilled t2v — 8+3 two-stage at 768x1280.
#
# Run:
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_768x1280.yaml
#
# Stage 1 denoises for 8 steps at half resolution; the latents are then
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
generator:
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
engine:
num_gpus: 1
pipeline:
workload_type: t2v
preset_overrides:
refine:
enabled: true
num_inference_steps: 3
guidance_scale: 1.0
add_noise: true
request:
prompt: >-
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
sampling:
height: 768
width: 1280
num_frames: 121
fps: 24
guidance_scale: 1.0
num_inference_steps: 8
output:
output_path: outputs/ltx2_3_t2v_8s3_768x1280.mp4
save_video: true
@@ -1,105 +0,0 @@
# 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
@@ -1,83 +0,0 @@
# 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
@@ -100,6 +100,9 @@ callbacks:
sampling_steps: [4]
sampling_timesteps: [1000, 750, 500, 250]
num_frames: 81
# Validation/inference uses standard CFG in both clean and Self-Forcing,
# so this directly matches Self-Forcing guidance_scale=3.0.
guidance_scale: 3.0
pipeline:
flow_shift: 5
+9 -96
View File
@@ -1,119 +1,32 @@
# World-Model: Matrix-Game 2.0 I2V
Training scenarios for the Matrix-Game 2.0 I2V world model on Solaris (Minecraft)
data and Zelda data, using the new YAML-driven trainer
(`fastvideo/train/entrypoint/train.py`).
## Solaris Configs
Three training scenarios for the Matrix-Game 2.0 I2V world model on the
new YAML-driven trainer (`fastvideo/train/entrypoint/train.py`).
| Config | Method | Student | Notes |
|---|---|---|---|
| `solaris/finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. |
| `solaris/dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. |
| `solaris/self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | Matrix-Game 2.0 DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. |
## Zelda Configs
| Config | Method | Student | Notes |
|---|---|---|---|
| `zelda/finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Zelda bidirectional I2V finetuning from `FastVideo/Matrix-Game-2.0-Base-Diffusers`. Uses 33-frame clips and Zelda validation with action overlays. |
| `zelda/dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Zelda causal Diffusion-Forcing SFT from `mignonjia/mg_bidirectional_zelda`. Uses the same Zelda data, resolution, optimizer, and validation defaults as the Zelda finetune config. |
| `zelda/self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | Zelda DMD/Self-Forcing distillation; student init = `mignonjia/mg_causal_zelda`, teacher = bidirectional (`mignonjia/mg_bidirectional_zelda`), critic = bidirectional. |
| `zelda/streaming_long_tuning_causal_i2v.yaml` | `StreamingLongTuningMethod` | `MatrixGame2CausalModel` | LongLive-style streaming long tuning from the 1k-step Zelda self-forcing checkpoint. |
Zelda world-model distillation is a two-run workflow: first run
`zelda/self_forcing_causal_i2v.yaml` to train or load the 1k-step
self-forcing checkpoint (`mignonjia/mg_sf_distilled_zelda_1k_steps`), then run
`zelda/streaming_long_tuning_causal_i2v.yaml` for the 3k-step streaming
long-tuning stage. The long-tuning YAML starts from that 1k-step checkpoint; it
does not run the short self-forcing stage inside the same config.
## Zelda Training Data
The Zelda training configs use `data/zeldam2-clean` as a suggested local path.
Download the dataset from Hugging Face before running those configs:
```bash
python scripts/huggingface/download_hf.py \
--repo_id mignonjia/zeldam2-clean \
--local_dir data/zeldam2-clean \
--repo_type dataset
```
You can store the dataset elsewhere; update `training.data.data_path` in the
YAML to point at that location.
## Multi3D Training Data
`zelda/finetune_i2v.yaml` includes an optional, commented-out Multi3D entry.
Enable it only when you want to mix Zelda with multi-game data from
`data/multi3d_games`. You can store this dataset anywhere; before enabling it,
update the matching commented `training.data.data_path` key in the YAML to the
correct location.
To mix datasets in a training YAML, set `training.data.data_path` to a
path-to-repeat-count mapping. For example, `zelda/finetune_i2v.yaml` can use
`data/zeldam2-clean: 1` and `# data/multi3d_games: 10`; uncommenting the
Multi3D entry repeats the multi-game parquet list ten times before training
samples are shuffled.
## World Model Validation Data
The Zelda validation configs expect a small public validation bundle under
`data/zelda_validation_data`.
Download it from Hugging Face before running the Zelda scenarios:
```bash
python scripts/huggingface/download_hf.py \
--repo_id mignonjia/zelda_validation_data \
--local_dir data/zelda_validation_data \
--repo_type dataset
```
The bundle contains `validation_zelda.json`, `images/`, and `actions/`.
The Zelda configs point
`callbacks.validation.dataset_file` at
`data/zelda_validation_data/validation_zelda.json`.
| `finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. |
| `dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. |
| `self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. |
## Usage
### Solaris
```bash
bash examples/train/run.sh \
examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml
examples/train/scenario/worldmodel/finetune_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml
examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml
```
### Zelda
```bash
# Finetuning / DFSFT
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml
# Distillation / long tuning
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml
examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml
```
Override any field on the command line:
```bash
bash examples/train/run.sh \
examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml \
examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml \
--training.distributed.num_gpus 8 \
--training.optimizer.learning_rate 1e-5
```
@@ -97,4 +97,4 @@ callbacks:
guidance_scale: 6.0
pipeline:
flow_shift: 5
flow_shift: 5
@@ -1,94 +0,0 @@
# Diffusion-Forcing SFT: Zelda world model I2V Causal
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_bidirectional_zelda
trainable: true
method:
_target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
chunk_size: 3
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path:
data/zeldam2-clean: 1
dataloader_num_workers: 1
train_batch_size: 1
training_cfg_rate: 0.0
seed: 42
num_latent_t: 9
num_height: 480
num_width: 832
num_frames: 33
optimizer:
learning_rate: 2.0e-5
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 60000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/matrixgame_finetune/checkpoints/zelda_causal_dfsft
training_state_checkpointing_steps: 5000
checkpoints_total_limit: 3
tracker:
entity: hapo-exp
project_name: mg_1.3b_zelda
run_name: zelda_causal_dfsft
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
grad_clip:
max_grad_norm: 1.0
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 200
sampling_steps: [40]
sampling_timesteps: [1000, 975, 950, 925, 900, 875, 850, 825, 800, 775,
750, 725, 700, 675, 650, 625, 600, 575, 550, 525,
500, 475, 450, 425, 400, 375, 350, 325, 300, 275,
250, 225, 200, 175, 150, 125, 100, 75, 50, 25]
num_frames: 33
overlay_actions: true
guidance_scale: 6.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
dit_config:
local_attn_size: 6
sink_size: 1
@@ -1,88 +0,0 @@
# Matrix-Game 2.0 Zelda + multi-game I2V finetune.
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: FastVideo/Matrix-Game-2.0-Base-Diffusers
trainable: true
method:
_target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path:
data/zeldam2-clean: 1
# data/multi3d_games: 10
dataloader_num_workers: 1
train_batch_size: 1
training_cfg_rate: 0.0 # unused for MatrixGame2 I2V; no text_embedding CFG dropout
seed: 42
num_latent_t: 9
num_height: 480
num_width: 832
num_frames: 33
optimizer:
learning_rate: 2.0e-5
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 60000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/matrixgame_finetune/checkpoints/zelda_with_mg_init
training_state_checkpointing_steps: 5000
checkpoints_total_limit: 3
tracker:
project_name: mg_1.3b_zelda
run_name: zelda_with_mg_init
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
grad_clip:
_target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
max_grad_norm: 1.0
validation:
_target_: fastvideo.train.callbacks.validation.ValidationCallback
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_i2v_pipeline.MatrixGame2I2VPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 200
sampling_steps: [40]
num_frames: 33
overlay_actions: true
guidance_scale: 6.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
@@ -1,120 +0,0 @@
# Self-forcing distillation: Zelda world model I2V Causal
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_causal_zelda
trainable: true
teacher:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: false
disable_custom_init_weights: true
critic:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: true
disable_custom_init_weights: true
method:
_target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
rollout_mode: simulate
generator_update_interval: 5
dmd_denoising_steps: [1000, 750, 500, 250]
warp_denoising_step: true
min_timestep_ratio: 0.02
max_timestep_ratio: 0.98
chunk_size: 3
student_sample_type: sde
same_step_across_blocks: true
last_step_only: false
context_noise: 0.0
enable_gradient_in_rollout: true
start_gradient_frame: 0
# Critic optimizer
fake_score_learning_rate: 3.0e-7
fake_score_betas: [0.9, 0.95]
fake_score_lr_scheduler: constant
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path: data/zeldam2-clean
dataloader_num_workers: 4
train_batch_size: 1
training_cfg_rate: 0.0
seed: 1001
num_latent_t: 9
num_height: 480
num_width: 832
num_frames: 33
optimizer:
learning_rate: 3.0e-6
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 1000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/zelda_causal_self_forcing
training_state_checkpointing_steps: 1000
checkpoints_total_limit: 1
tracker:
project_name: wangame_sf
run_name: mg2_self_forcing_9_latents
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
ema:
decay: 0.99
start_iter: 200
grad_clip:
max_grad_norm: 1.0
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 100
sampling_steps: [4]
sampling_timesteps: [1000, 750, 500, 250]
num_frames: 153
overlay_actions: true
keyboard_value_scale: 1.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
dit_config:
local_attn_size: 6
sink_size: 1
@@ -1,140 +0,0 @@
# MatrixGame2 I2V LongLive-style streaming distillation: Zelda world model I2V Causal
# Student init from self forcing after 1k steps
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_sf_distilled_zelda_1k_steps
trainable: true
# transformer_override_safetensor: outputs/matrixgame_dmd/checkpoints/mg_zelda_sf_m2/checkpoint-1000_weight_only/ema/generator_ema.safetensors
teacher:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: false
disable_custom_init_weights: true
critic:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: true
disable_custom_init_weights: true
method:
_target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
rollout_mode: simulate
generator_update_interval: 5
dmd_denoising_steps: [1000, 750, 500, 250]
warp_denoising_step: true
min_timestep_ratio: 0.02
max_timestep_ratio: 0.98
chunk_size: 3
student_sample_type: sde
same_step_across_blocks: true
last_step_only: false
context_noise: 0.0
enable_gradient_in_rollout: true
start_gradient_frame: 0
streaming_training: true
streaming_chunk_size: 9
streaming_max_length: 39
streaming_fixed_overlap_latents: 3
streaming_reencode_overlap_anchor: true
streaming_anchor_inject_k: 1
streaming_require_full_blocks: true
multi_phased_distill_schedule:
- stage: streaming_long
start_step: 0
end_step: 3000
num_latent_t: 39
streaming_training: true
streaming_chunk_size: 9
streaming_max_length: 39
streaming_fixed_overlap_latents: 3
# Critic optimizer
fake_score_learning_rate: 3.0e-7
fake_score_betas: [0.9, 0.95]
fake_score_lr_scheduler: constant
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path: data/zeldam2-clean
dataloader_num_workers: 4
train_batch_size: 1
training_cfg_rate: 0.0
seed: 1001
num_latent_t: 39
num_height: 480
num_width: 832
num_frames: 153
optimizer:
learning_rate: 3.0e-6
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 3000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/zelda_causal_long_tuning
training_state_checkpointing_steps: 1000
checkpoints_total_limit: 2
tracker:
project_name: wangame_sf
run_name: mg2_39only_streaming_long
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
ema:
decay: 0.99
start_iter: 200
grad_clip:
max_grad_norm: 1.0
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 100
sampling_steps: [4]
sampling_timesteps: [1000, 750, 500, 250]
num_frames: 153
overlay_actions: true
keyboard_value_scale: 1.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
dit_config:
local_attn_size: 6
sink_size: 1
-27
View File
@@ -1,27 +0,0 @@
# 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
-16
View File
@@ -90,10 +90,6 @@ 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
@@ -329,18 +325,6 @@ 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,
-1
View File
@@ -150,7 +150,6 @@ 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
+119 -20
View File
@@ -1,15 +1,13 @@
# 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.attention.utils.flash_attn_default import (
fa_version,
flash_attn_func_compilable,
)
from fastvideo import envs
from fastvideo.attention.backends.abstract import (
AttentionBackend,
AttentionImpl,
@@ -19,6 +17,119 @@ 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.
@@ -198,17 +309,9 @@ 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_compilable as flash_attn_no_pad,
flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
flash_attn_no_pad,
flash_attn_varlen_qk_no_pad,
)
attn_mask = attn_metadata.attn_mask
@@ -235,11 +338,7 @@ class FlashAttentionImpl(AttentionImpl):
)
qkv = torch.stack([query, key, value], dim=2)
key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, attn_mask.shape[-1]).to(device=query.device)
if key_padding_mask.shape[-1] > qkv.shape[1]:
raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
f"expected at most {qkv.shape[1]}, got {key_padding_mask.shape[-1]}")
attn_mask_padded = F.pad(key_padding_mask, (qkv.shape[1] - key_padding_mask.shape[-1], 0), value=True)
attn_mask_padded = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True)
output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
elif self.nvfp4_fa4:
output = self._forward_nvfp4(query, key, value)
-147
View File
@@ -1,147 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""NABLA block-sparse flex-attention backend (Kandinsky5 "nabla" checkpoints).
The block mask is data-dependent: nablaT_v2 mean-pools 64-token blocks of the
fractal-ordered sequence, thresholds the softmaxed block map, and ORs it with a
precomputed spatio-temporal-window (STA) mask carried on the attention
metadata. The mask spans the full sequence, so this backend does not support
sequence parallelism — use it via LocalAttention only.
"""
import math
from dataclasses import dataclass
from typing import Any
import torch
try:
from torch.nn.attention.flex_attention import BlockMask, flex_attention
flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
CAN_USE_FLEX_ATTN = True
except ImportError:
CAN_USE_FLEX_ATTN = False
from fastvideo.attention.backends.abstract import (
AttentionBackend,
AttentionImpl,
AttentionMetadata,
AttentionMetadataBuilder,
)
def nablaT_v2(
q: torch.Tensor,
k: torch.Tensor,
sta: torch.Tensor,
thr: float = 0.9,
) -> "BlockMask":
q = q.transpose(1, 2).contiguous()
k = k.transpose(1, 2).contiguous()
# Map estimation
B, h, S, D = q.shape
s1 = S // 64
qa = q.reshape(B, h, s1, 64, D).mean(-2)
ka = k.reshape(B, h, s1, 64, D).mean(-2).transpose(-2, -1)
map = qa @ ka
map = torch.softmax(map / math.sqrt(D), dim=-1)
# Map binarization
vals, inds = map.sort(-1)
cvals = vals.cumsum_(-1)
mask = (cvals >= 1 - thr).int()
mask = mask.gather(-1, inds.argsort(-1))
mask = torch.logical_or(mask, sta)
# BlockMask creation
kv_nb = mask.sum(-1).to(torch.int32)
kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32)
return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None)
class NablaAttentionBackend(AttentionBackend):
@staticmethod
def get_name() -> str:
return "NABLA_ATTN"
@staticmethod
def get_impl_cls() -> type["NablaAttentionImpl"]:
return NablaAttentionImpl
@staticmethod
def get_metadata_cls() -> type["NablaAttentionMetadata"]:
return NablaAttentionMetadata
@staticmethod
def get_builder_cls() -> type["NablaAttentionMetadataBuilder"]:
return NablaAttentionMetadataBuilder
@dataclass
class NablaAttentionMetadata(AttentionMetadata):
# Block-level STA window mask [1, 1, S/64, S/64], precomputed once per run.
sta_mask: torch.Tensor = None # type: ignore[assignment]
# Cumulative-probability threshold for block-map binarization.
P: float = 0.9
visual_shape: tuple[int, int, int] = (0, 0, 0)
class NablaAttentionMetadataBuilder(AttentionMetadataBuilder):
def __init__(self) -> None:
pass
def prepare(self) -> None:
pass
def build(
self,
current_timestep: int,
sta_mask: torch.Tensor,
P: float,
visual_shape: tuple[int, int, int],
**kwargs: Any,
) -> NablaAttentionMetadata:
return NablaAttentionMetadata(
current_timestep=current_timestep,
sta_mask=sta_mask,
P=P,
visual_shape=visual_shape,
)
class NablaAttentionImpl(AttentionImpl):
def __init__(
self,
num_heads: int,
head_size: int,
softmax_scale: float,
causal: bool = False,
num_kv_heads: int | None = None,
prefix: str = "",
**extra_impl_args,
) -> None:
if not CAN_USE_FLEX_ATTN:
raise RuntimeError("NABLA attention requires torch.nn.attention.flex_attention, "
"which is unavailable in this PyTorch build.")
if causal:
raise ValueError("NABLA attention does not support causal masking.")
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: NablaAttentionMetadata,
) -> torch.Tensor:
# q/k/v: [B, S, heads, head_dim], fractal-ordered by the model; S % 64 == 0.
block_mask = nablaT_v2(query, key, attn_metadata.sta_mask, thr=attn_metadata.P)
return flex_attention(
query=query.transpose(1, 2),
key=key.transpose(1, 2),
value=value.transpose(1, 2),
block_mask=block_mask,
).transpose(1, 2)
+3 -53
View File
@@ -2,7 +2,6 @@
import torch
from dataclasses import dataclass
from torch.nn import functional as F
from fastvideo.attention.backends.abstract import ( # FlashAttentionMetadata,
AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder)
from fastvideo.logger import init_logger
@@ -15,18 +14,12 @@ class SDPABackend(AttentionBackend):
accept_output_buffer: bool = True
@staticmethod
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
def get_supported_head_sizes() -> list[int]:
return [32, 64, 96, 128, 160, 192, 224, 256]
@staticmethod
def get_name() -> str:
return "TORCH_SDPA"
return "SDPA"
@staticmethod
def get_impl_cls() -> type["SDPAImpl"]:
@@ -56,51 +49,9 @@ class SDPAMetadataBuilder(AttentionMetadataBuilder):
current_timestep: int,
attn_mask: torch.Tensor,
) -> SDPAMetadata:
# Store the mask exactly as passed. The metadata is cross-backend:
# call sites (HYWorld, HunyuanVideo15) build SDPAMetadata while the
# layer's selector may pick FLASH_ATTN, and the shared convention for
# padding masks is the tokenizer-style 2D [batch, key_len]. Any
# reshaping for torch.sdpa happens inside the SDPA impl
# (_normalize_attn_mask_for_sdpa).
return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask)
def _normalize_attn_mask_for_sdpa(
attn_mask: torch.Tensor | None,
query: torch.Tensor,
key: torch.Tensor,
) -> torch.Tensor | None:
if attn_mask is None:
return None
attn_mask = attn_mask.to(device=query.device)
# F.scaled_dot_product_attention only accepts bool or float masks;
# tokenizers commonly produce int64 0/1 padding masks.
if attn_mask.dtype != torch.bool and not attn_mask.dtype.is_floating_point:
attn_mask = attn_mask != 0
key_len = key.shape[-2]
if attn_mask.shape[-1] > key_len:
raise ValueError("Invalid attention mask length for SDPA: "
f"expected at most {key_len}, got {attn_mask.shape[-1]}")
if attn_mask.shape[-1] < key_len:
# Front-pad as "attend": double-stream layouts (HYWorld) prepend
# non-text tokens the tokenizer mask does not cover.
valid_value = True if attn_mask.dtype == torch.bool else 0.0
attn_mask = F.pad(attn_mask, (key_len - attn_mask.shape[-1], 0), value=valid_value)
if attn_mask.dim() == 2:
# In-tree producers pass 2D [batch, key_len] padding masks; lift to a
# broadcastable [batch, 1, 1, key_len] here so torch.sdpa does not
# reinterpret 2D as its documented [query_len, key_len] broadcast.
return attn_mask[:, None, None, :]
if attn_mask.dim() == 3:
return attn_mask[:, None, :, :]
if attn_mask.dim() == 4:
return attn_mask
raise ValueError(f"Unsupported attention mask shape for SDPA: {attn_mask.shape}")
class SDPAImpl(AttentionImpl):
def __init__(
@@ -131,7 +82,6 @@ class SDPAImpl(AttentionImpl):
attn_mask = attn_metadata.attn_mask if (attn_metadata is not None
and hasattr(attn_metadata, "attn_mask")) else None
attn_mask = _normalize_attn_mask_for_sdpa(attn_mask, query, key)
attn_kwargs = {
"attn_mask": attn_mask,
"dropout_p": self.dropout,
+1 -5
View File
@@ -252,7 +252,6 @@ class LocalAttention(nn.Module):
causal: bool = False,
supported_attention_backends: tuple[AttentionBackendEnum, ...]
| None = None,
default_backend: AttentionBackendEnum | None = None,
**extra_impl_args) -> None:
super().__init__()
if softmax_scale is None:
@@ -263,10 +262,7 @@ class LocalAttention(nn.Module):
num_kv_heads = num_heads
dtype = get_compute_dtype()
attn_backend = get_attn_backend(head_size,
dtype,
supported_attention_backends=supported_attention_backends,
default_backend=default_backend)
attn_backend = get_attn_backend(head_size, dtype, supported_attention_backends=supported_attention_backends)
impl_cls = attn_backend.get_impl_cls()
self.attn_impl = impl_cls(num_heads=num_heads,
head_size=head_size,
+1 -9
View File
@@ -84,9 +84,8 @@ def get_attn_backend(
dtype: torch.dtype,
supported_attention_backends: tuple[AttentionBackendEnum, ...]
| None = None,
default_backend: AttentionBackendEnum | None = None,
) -> type[AttentionBackend]:
return _cached_get_attn_backend(head_size, dtype, supported_attention_backends, default_backend)
return _cached_get_attn_backend(head_size, dtype, supported_attention_backends)
@cache
@@ -95,7 +94,6 @@ def _cached_get_attn_backend(
dtype: torch.dtype,
supported_attention_backends: tuple[AttentionBackendEnum, ...]
| None = None,
default_backend: AttentionBackendEnum | None = None,
) -> type[AttentionBackend]:
# Check whether a particular choice of backend was
# previously forced.
@@ -114,12 +112,6 @@ def _cached_get_attn_backend(
if backend_by_env_var is not None:
selected_backend = backend_name_to_enum(backend_by_env_var)
# Layer-level default (e.g. a checkpoint that requires a specific sparse
# backend). Lower precedence than the global force and the env var, so
# users can still override it.
if selected_backend is None and default_backend is not None:
selected_backend = default_backend
# get device-specific attn_backend
from fastvideo.platforms import current_platform
@@ -1,251 +0,0 @@
# 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.")
+5 -483
View File
@@ -24,7 +24,7 @@ from flash_attn.bert_padding import pad_input, unpad_input
from fastvideo import envs
def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
def _resolve_flash_attn_varlen_func() -> Any:
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,28 +39,20 @@ def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
"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, "4"
return flash_attn_varlen_func_cute
try:
from flash_attn_interface import (
flash_attn_varlen_func as flash_attn_varlen_func_interface, )
return flash_attn_varlen_func_interface, "3"
return flash_attn_varlen_func_interface
except ImportError:
from flash_attn import (
flash_attn_varlen_func as flash_attn_varlen_func_flash, )
return flash_attn_varlen_func_flash, "2"
return flash_attn_varlen_func_flash
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, )
flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
def flash_attn_no_pad(
@@ -199,473 +191,3 @@ 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)
+2 -5
View File
@@ -1,9 +1,7 @@
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
@@ -17,7 +15,6 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
__all__ = [
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
"DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
"LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
"StableAudioConfig", "GlmImageDiTConfig"
"DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
"HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
]
-27
View File
@@ -1,27 +0,0 @@
# 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"
@@ -1,61 +0,0 @@
# 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"
+4 -14
View File
@@ -2,24 +2,14 @@
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.platforms import AttentionBackendEnum
def _is_kandinsky5_transformer_block(n: str, m) -> bool:
return ("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
@dataclass
class Kandinsky5ArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [_is_kandinsky5_transformer_block])
# NABLA block-sparse attention for attention_type="nabla" checkpoints, plus
# the dense backends every DiT supports.
_supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
AttentionBackendEnum.NABLA_ATTN,
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
)
_fsdp_shard_conditions: list = field(default_factory=lambda: [
lambda n, m:
("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
])
# Native FastVideo implementation uses the same parameter names as diffusers
# except FFN internals: Diffusers FFN uses `in_layer/out_layer`, while
-14
View File
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Literal
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
@@ -20,11 +19,6 @@ 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",
@@ -92,14 +86,6 @@ 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
+1 -6
View File
@@ -43,16 +43,11 @@ class TextEncoderArchConfig(EncoderArchConfig):
require_processor: bool = False
def __post_init__(self) -> None:
# update_model_arch re-runs __post_init__ after pipeline configs may
# have customized tokenizer_kwargs (e.g. kandinsky5/gen3c/longcat set
# "padding"); rebuilding the dict here would silently wipe those
# customizations, so only fill in defaults for keys not already set.
defaults = {
self.tokenizer_kwargs = {
"truncation": True,
"max_length": self.text_len,
"return_tensors": "pt",
}
self.tokenizer_kwargs = defaults | self.tokenizer_kwargs
@dataclass
@@ -2,7 +2,6 @@ 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
@@ -22,5 +21,4 @@ __all__ = [
"OobleckVAEArchConfig",
"OobleckVAEConfig",
"Flux2VAEConfig",
"GlmImageVAEConfig",
]
@@ -1,95 +0,0 @@
# 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

Some files were not shown because too many files have changed in this diff Show More