Compare commits
24
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
59cf9d6fa2 | ||
|
|
0c90c328c5 | ||
|
|
cfb54a3b2a | ||
|
|
deb51f1dcc | ||
|
|
79a5930340 | ||
|
|
3a0ce7c20d | ||
|
|
0b006a9e46 | ||
|
|
fb9acac514 | ||
|
|
2d78957219 | ||
|
|
63991d2017 | ||
|
|
b6eafbea50 | ||
|
|
7b84041642 | ||
|
|
48801c29c7 | ||
|
|
cb8be0d3f3 | ||
|
|
1bc7ff79e6 | ||
|
|
76fbe472ae | ||
|
|
eba74c43ef | ||
|
|
7744a74c13 | ||
|
|
d746b8b956 | ||
|
|
65f12605d8 | ||
|
|
1e1ac08cd0 | ||
|
|
a253856147 | ||
|
|
6e25d94ebc | ||
|
|
cae8fa18dc |
@@ -1,20 +1,23 @@
|
||||
---
|
||||
name: reseed-performance-baseline
|
||||
description: Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, or environment-caused benchmark shift using one or more reviewed normalized performance JSONs. Use when performance CI fails because metrics such as latency, throughput, component time, or peak memory changed for an accepted reason and the rolling median baseline in FastVideo/performance-tracking must be advanced from a consistent batch of reviewed source results. The workflow backs up existing history under /tmp, validates all source JSONs for the same (model_id, gpu_type), rejects internally inconsistent source batches, uploads one success=true reseed record per accepted source JSON, and offers to clean local temp state after a successful upload.
|
||||
description: Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, environment-caused benchmark shift, or reviewed v2 calibration using one or more reviewed normalized performance JSONs. Use when performance CI fails because metrics such as latency, throughput, component time, or peak memory changed for an accepted reason and the rolling median baseline in FastVideo/performance-tracking must be advanced, or when a new v2 exact comparable identity needs its first approved baseline. The workflow backs up existing history under /tmp, validates all source JSONs for the same legacy (model_id, gpu_type) target or the same v2 exact identity, rejects internally inconsistent source batches, uploads one success=true baseline record per accepted source JSON, and offers to clean local temp state after a successful upload.
|
||||
---
|
||||
|
||||
# Re-seed Performance Baseline
|
||||
|
||||
## Purpose
|
||||
|
||||
Replace or advance the rolling performance baseline for a single
|
||||
`(model_id, gpu_type)` pair in the HF dataset
|
||||
`FastVideo/performance-tracking`.
|
||||
Replace or advance the rolling performance baseline in the HF dataset
|
||||
`FastVideo/performance-tracking`. Legacy targets are scoped by
|
||||
`(model_id, gpu_type)`. V2 targets are scoped by exact comparable identity:
|
||||
`workload_id`, `variant_id`, `benchmark_version`, `hardware_profile_id`,
|
||||
`software_profile_id`, and `recipe_fingerprint`.
|
||||
|
||||
Performance comparison uses the median of up to the last 5 successful records
|
||||
for the same model and GPU. Failed records are useful audit history, but they
|
||||
do not move the future baseline because `compare_baseline.py` loads records
|
||||
with `successful_only=True`.
|
||||
Performance comparison uses the median of up to the last 5 successful,
|
||||
baseline-eligible records for the same target. Failed or calibration-only
|
||||
records are useful audit history, but they do not move the future baseline
|
||||
because `compare_baseline.py` loads records with `successful_only=True` and
|
||||
`baseline_eligible_only=True`.
|
||||
|
||||
This skill now reseeds from a reviewed batch of one or more source performance
|
||||
JSONs. It uploads one new `success=true` record per accepted source JSON; it
|
||||
@@ -22,11 +25,13 @@ does not blindly replicate one measurement into 3 or 5 records. The effective
|
||||
reseed size is therefore dynamic and equals the number of provided, validated,
|
||||
internally consistent source JSONs.
|
||||
|
||||
If the operator provides fewer than 3 records, call out that the last-5 rolling
|
||||
median may not move immediately. If the operator provides 3 consistent shifted
|
||||
records, the rolling median usually moves immediately. If the operator provides
|
||||
5 consistent shifted records, the last-5 window is effectively reset to the new
|
||||
runtime profile.
|
||||
For baseline shifts with existing history, if the operator provides fewer than
|
||||
3 records, call out that the last-5 rolling median may not move immediately. If
|
||||
the operator provides 3 consistent shifted records, the rolling median usually
|
||||
moves immediately. If the operator provides 5 consistent shifted records, the
|
||||
last-5 window is effectively reset to the new runtime profile. For the first
|
||||
approved v2 baseline of a new exact identity, one reviewed calibration seed is
|
||||
enough for the next comparable run to leave `CALIBRATION_NEEDED`.
|
||||
|
||||
These records are intentional operator-approved baseline resets, not ordinary
|
||||
independent main-branch persistence. Mark them clearly with provenance fields
|
||||
@@ -66,8 +71,8 @@ approval, then upload reviewed accepted baseline records.
|
||||
|
||||
| Parameter | Required | Description |
|
||||
|-----------|----------|-------------|
|
||||
| `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
|
||||
| `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
|
||||
| `model_id` | Legacy required; v2 inferred | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. For v2 records, use the `model_id` from each source artifact only as the upload directory; comparison is by exact identity. |
|
||||
| `gpu_type` | Legacy required; v2 inferred | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. V2 hardware matching uses `hardware_profile_id`; preserve `gpu_type` as display metadata. |
|
||||
| `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
|
||||
| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
|
||||
| `intent_rationale` | Yes | One-line explanation for why the baseline shift is legitimate. This is written into provenance and should be reused in the PR. |
|
||||
@@ -78,10 +83,14 @@ Hardcoded defaults:
|
||||
supported by the code, but use the default unless the user explicitly asks).
|
||||
- Local sync root: `/tmp/perf-tracking` (`PERFORMANCE_TRACKING_ROOT` override
|
||||
is supported).
|
||||
- Prepared-record staging root: `/tmp/performance_reseed_prepared`
|
||||
(`PERFORMANCE_RESEED_STAGING_ROOT` override is supported). Keep it separate
|
||||
and non-nested from the sync root.
|
||||
- Backup root: `/tmp/performance_reseed_backup`.
|
||||
- Download scratch root for source artifact URLs: `/tmp/performance_reseed_source`.
|
||||
- Baseline window: last 5 `success=true` records for the same
|
||||
`(model_id, gpu_type)`.
|
||||
- Baseline window: last 5 `success=true`, `baseline_eligible=true` records
|
||||
for the same legacy `(model_id, gpu_type)` target or the same v2 exact
|
||||
comparable identity.
|
||||
- Reseed count: dynamic. Upload exactly one accepted seed record per validated
|
||||
source JSON.
|
||||
|
||||
@@ -115,12 +124,24 @@ with open(source_result, encoding="utf-8") as f:
|
||||
record = json.load(f)
|
||||
```
|
||||
|
||||
Stop if any normalized record's `model_id` or `gpu_type` does not match the
|
||||
requested `model_id` and `gpu_type`.
|
||||
Classify the source batch before continuing:
|
||||
|
||||
- **Legacy source records** have no v2 exact identity fields. Stop if any
|
||||
normalized record's `model_id` or `gpu_type` does not match the requested
|
||||
`model_id` and `gpu_type`.
|
||||
- **V2 source records** have exact identity fields. Stop unless every source
|
||||
record has all six comparable identity fields and they are identical across
|
||||
the batch: `workload_id`, `variant_id`, `benchmark_version`,
|
||||
`hardware_profile_id`, `software_profile_id`, and `recipe_fingerprint`.
|
||||
Do not fall back to legacy `(model_id, gpu_type)` matching for v2 records.
|
||||
|
||||
The source records may have `success: false` when they came from failed
|
||||
rolling baseline comparisons. That is expected; only the reviewed reseed
|
||||
records become new `success: true` baseline records after explicit approval.
|
||||
For a first v2 baseline seed, the source records must instead be successful
|
||||
scheduled-main full-suite `CALIBRATION_NEEDED` normalized artifacts. Reject PR,
|
||||
local, direct-run, non-main-branch, or non-full-suite calibration artifacts as
|
||||
seed sources.
|
||||
|
||||
Sort validated source records by their original `timestamp` ascending before
|
||||
preparing the seed records. If a source timestamp is missing or unparsable,
|
||||
@@ -194,7 +215,7 @@ export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
|
||||
python -c 'from fastvideo.performance.hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
|
||||
```
|
||||
|
||||
Then back up only the sanitized model directory under `/tmp`:
|
||||
For legacy records, back up the sanitized model directory under `/tmp`:
|
||||
|
||||
```bash
|
||||
SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
|
||||
@@ -209,6 +230,16 @@ mkdir -p "$BACKUP_DIR"
|
||||
cp -R "${PERFORMANCE_TRACKING_ROOT}/${MODEL_SAFE}" "$BACKUP_DIR/" 2>/dev/null || true
|
||||
```
|
||||
|
||||
For v2 records, back up the full local tracking root after sync. Exact identity
|
||||
lookup scans across model directories, so a benchmark rename may have relevant
|
||||
history outside the current source artifact's `model_id` directory:
|
||||
|
||||
```bash
|
||||
BACKUP_DIR="/tmp/performance_reseed_backup/${TIMESTAMP}_${SHORT_COMMIT}_v2_exact_identity"
|
||||
mkdir -p "$BACKUP_DIR"
|
||||
cp -R "${PERFORMANCE_TRACKING_ROOT}" "$BACKUP_DIR/tracking-root"
|
||||
```
|
||||
|
||||
Write provenance next to the backup:
|
||||
|
||||
```bash
|
||||
@@ -231,7 +262,9 @@ first baseline seed. Continue, but report that baseline history was empty.
|
||||
|
||||
### 3. Compute old baseline and candidate shift
|
||||
|
||||
Load the last 5 successful records for the target:
|
||||
Load the last 5 successful baseline records for the target.
|
||||
|
||||
For legacy targets:
|
||||
|
||||
```python
|
||||
from fastvideo.performance.hf_store import load_records_for_model
|
||||
@@ -242,6 +275,28 @@ records = load_records_for_model(
|
||||
"<gpu_type>",
|
||||
last_n=5,
|
||||
successful_only=True,
|
||||
baseline_eligible_only=True,
|
||||
)
|
||||
```
|
||||
|
||||
For v2 exact-identity targets:
|
||||
|
||||
```python
|
||||
from fastvideo.performance.hf_store import load_records_for_identity
|
||||
|
||||
records = load_records_for_identity(
|
||||
"/tmp/perf-tracking",
|
||||
{
|
||||
"workload_id": "<workload_id>",
|
||||
"variant_id": "<variant_id>",
|
||||
"benchmark_version": "<benchmark_version>",
|
||||
"hardware_profile_id": "<hardware_profile_id>",
|
||||
"software_profile_id": "<software_profile_id>",
|
||||
"recipe_fingerprint": "<recipe_fingerprint>",
|
||||
},
|
||||
last_n=5,
|
||||
successful_only=True,
|
||||
baseline_eligible_only=True,
|
||||
)
|
||||
```
|
||||
|
||||
@@ -257,7 +312,8 @@ medians after appending the proposed seed records, and source batch spread for:
|
||||
|
||||
Also print how many successful old records exist. Make clear:
|
||||
|
||||
- 1 seed record usually does not move a last-5 median by itself.
|
||||
- 1 seed record usually does not move an existing last-5 median by itself, but
|
||||
it is enough to establish the first v2 baseline for a new exact identity.
|
||||
- 3 consistent seed records usually move the last-5 median immediately.
|
||||
- 5 consistent seed records effectively reset the last-5 window.
|
||||
- The records are intentional approved baseline resets and must be labeled
|
||||
@@ -267,10 +323,10 @@ Also print how many successful old records exist. Make clear:
|
||||
|
||||
Require an explicit confirmation phrase before preparing the upload:
|
||||
|
||||
> About to RE-SEED performance baseline for `<model_id>` on `<gpu_type>`.
|
||||
> About to RE-SEED performance baseline for `<target description>`.
|
||||
> This will upload `<N>` new `success=true` records to
|
||||
> `FastVideo/performance-tracking/<sanitize(model_id)>/`, one per accepted
|
||||
> source JSON.
|
||||
> `FastVideo/performance-tracking/<sanitize(model_id)>/` or the source
|
||||
> artifact's v2 model directory, one per accepted source JSON.
|
||||
>
|
||||
> Reason: `<intent_rationale>`
|
||||
> Source results: `<source_results>`
|
||||
@@ -288,25 +344,63 @@ Do not continue unless the user types exactly `confirm performance reseed`.
|
||||
|
||||
### 5. Create the accepted seed records
|
||||
|
||||
Create one seed record from each normalized source result. Do not copy the
|
||||
Create one seed record from each normalized source result.
|
||||
|
||||
For first v2 baseline seeds, use the scoped utility. It validates exact
|
||||
identity, requires successful scheduled-main full-suite `CALIBRATION_NEEDED`
|
||||
source artifacts, preserves the normalized v2 identity and metadata fields,
|
||||
and writes seed records with `success=true`, `baseline_eligible=true`, and
|
||||
`comparison_status=PASS`:
|
||||
|
||||
```bash
|
||||
python fastvideo/tests/performance/seed_baseline.py \
|
||||
--source-result <normalized_perf_1.json> \
|
||||
--source-result <normalized_perf_2.json> \
|
||||
--intent-rationale "<intent_rationale>" \
|
||||
--max-intra-batch-regression 0.05 \
|
||||
--tracking-root "${PERFORMANCE_TRACKING_ROOT}" \
|
||||
--staging-root "${PERFORMANCE_RESEED_STAGING_ROOT:-/tmp/performance_reseed_prepared}"
|
||||
```
|
||||
|
||||
The utility is prepare-only and intentionally has no upload option. Upload the
|
||||
scoped records only after the separate confirmation in step 6.
|
||||
|
||||
The utility validates against an isolated fresh HF snapshot and leaves
|
||||
`PERFORMANCE_TRACKING_ROOT` untouched; that argument only proves the staging
|
||||
root is separate from the operator's tracking mirror. Before writing, it stops
|
||||
if the exact identity already has a successful baseline-eligible record or if
|
||||
the workload/variant/version already trusts another recipe. It atomically
|
||||
reserves the exact identity and writes a digest-protected upload manifest bound
|
||||
to the current HF endpoint, repository id, and repository type. Keep the
|
||||
prepared records, manifest, source files, and reservation unchanged until the
|
||||
operation is uploaded or explicitly cleaned up.
|
||||
|
||||
If the prepared seed records look correct, upload only those scoped records in
|
||||
step 7. Do not rerun the utility with a different source list after approval.
|
||||
|
||||
For legacy reseeds or accepted v2 baseline shifts from regression artifacts,
|
||||
create one seed record from each normalized source result. Do not copy the
|
||||
source JSON wholesale.
|
||||
|
||||
Infer the baseline field allowlist from all existing HF records for the target
|
||||
`(model_id, gpu_type)` after syncing, including both `success=true` and
|
||||
`success=false` records. Use the union of non-provenance keys present in those
|
||||
target records, preserving only fields that also exist in the normalized
|
||||
source record or are explicitly set by the reseed workflow. Always include
|
||||
`model_id`, `timestamp`, and `success` because the upload path and baseline
|
||||
loader depend on them. Always set `timestamp` to a fresh reseed timestamp and
|
||||
`success` to `true`. Do not include unrelated source-only fields that are
|
||||
absent from existing HF records.
|
||||
after syncing, including both `success=true` and `success=false` records. For
|
||||
legacy targets the target is `(model_id, gpu_type)`. For v2 baseline-shift
|
||||
reseeds the target is the exact comparable identity. Use the union of
|
||||
non-provenance keys present in those target records, preserving only fields
|
||||
that also exist in the normalized source record or are explicitly set by the
|
||||
reseed workflow. Always include `model_id`, `timestamp`, `success`,
|
||||
`baseline_eligible`, and `comparison_status` because the upload path and
|
||||
baseline loader depend on them. Always set `timestamp` to a fresh reseed
|
||||
timestamp, `success` to `true`, `baseline_eligible` to `true`, and
|
||||
`comparison_status` to `PASS`. Do not include unrelated source-only fields
|
||||
that are absent from existing HF records.
|
||||
|
||||
Exclude existing provenance or operator metadata from the inferred baseline
|
||||
field allowlist. At minimum, exclude keys prefixed with `baseline_reseed` and
|
||||
any fields known to be local-only audit metadata.
|
||||
|
||||
If there are no previous HF records for the target model/GPU, fall back to this
|
||||
default baseline field list:
|
||||
If there are no previous HF records for the target, fall back to this default
|
||||
baseline field list:
|
||||
|
||||
- `model_id`
|
||||
- `timestamp`
|
||||
@@ -319,6 +413,22 @@ default baseline field list:
|
||||
- `dit_time_s`
|
||||
- `vae_decode_time_s`
|
||||
- `success`
|
||||
- `baseline_eligible`
|
||||
- `comparison_status`
|
||||
|
||||
For v2 baseline-shift reseeds with no previous HF records for the exact
|
||||
identity, also preserve:
|
||||
|
||||
- `workload_id`
|
||||
- `variant_id`
|
||||
- `benchmark_version`
|
||||
- `recipe_fingerprint`
|
||||
- `hardware_profile_id`
|
||||
- `software_profile_id`
|
||||
- `recipe`
|
||||
- `hardware_profile`
|
||||
- `software_profile`
|
||||
- `software_comparison_profile`
|
||||
|
||||
Do not upload extra fields from the source artifact.
|
||||
|
||||
@@ -334,6 +444,22 @@ Optional provenance fields are allowed and useful:
|
||||
- `baseline_reseed_operator`
|
||||
- `baseline_reseed_max_intra_batch_regression`
|
||||
|
||||
The v2 calibration seed utility writes analogous first-seed provenance:
|
||||
|
||||
- `baseline_seed: true`
|
||||
- `baseline_seed_reason`
|
||||
- `baseline_seed_source_result`
|
||||
- `baseline_seed_source_status`
|
||||
- `baseline_seed_source_timestamp`
|
||||
- `baseline_seed_source_success`
|
||||
- `baseline_seed_source_run_source`
|
||||
- `baseline_seed_source_branch`
|
||||
- `baseline_seed_source_test_scope`
|
||||
- `baseline_seed_source_pr_number`
|
||||
- `baseline_seed_batch_size`
|
||||
- `baseline_seed_batch_index`
|
||||
- `baseline_seed_operator`
|
||||
|
||||
Use a fresh reseed timestamp for each seed record, not the original source
|
||||
result timestamp. This is required because
|
||||
`load_records_for_model(..., last_n=5)` keeps the last records after loading
|
||||
@@ -356,7 +482,8 @@ Prefer uploading new accepted seed records so failed history remains visible.
|
||||
Print:
|
||||
|
||||
- Backup directory path under `/tmp`.
|
||||
- Prepared local record paths under `PERFORMANCE_TRACKING_ROOT`.
|
||||
- Prepared local record paths under `PERFORMANCE_RESEED_STAGING_ROOT`.
|
||||
- Prepared upload-manifest path under the identity reservation.
|
||||
- HF paths that will receive the new records.
|
||||
- Old rolling medians.
|
||||
- Source batch medians, source batch spread, reseed count, and candidate
|
||||
@@ -368,22 +495,36 @@ prepared records plus backup on disk.
|
||||
|
||||
### 7. Upload only the scoped records
|
||||
|
||||
Use the shared storage helper so the path and repo type match CI:
|
||||
For a first v2 calibration seed, use the manifest uploader after the user
|
||||
replies exactly `upload`:
|
||||
|
||||
```python
|
||||
from fastvideo.performance.hf_store import upload_record
|
||||
|
||||
upload_record("<local_record_path>", record, strict=True)
|
||||
```bash
|
||||
python -c 'from fastvideo.tests.performance.seed_baseline import upload_prepared_seed_manifest; print(upload_prepared_seed_manifest("<prepared_manifest>"))'
|
||||
```
|
||||
|
||||
Run it once per prepared record. Each upload goes to:
|
||||
The uploader verifies the source and prepared-record digests, pins and scans
|
||||
the current HF revision, rechecks exact-identity and recipe-cohort conflicts,
|
||||
and writes the entire batch in one commit whose `parent_commit` must still be
|
||||
current. A concurrent Hub update makes the commit fail. Do not retry
|
||||
automatically: preserve staging, refresh/review remote state, and request a new
|
||||
explicit `upload` after the conflict is understood. Each record goes to:
|
||||
|
||||
```text
|
||||
FastVideo/performance-tracking/<sanitize(model_id)>/<record_filename>.json
|
||||
```
|
||||
|
||||
Never bulk upload the whole tracking root. Never modify another model's
|
||||
directory in the same operation.
|
||||
Never call `upload_record()` once per first-seed record: that can partially
|
||||
land the batch and has no compare-and-swap guard.
|
||||
|
||||
For a legacy reseed or an accepted v2 baseline shift, the first-seed manifest
|
||||
validator does not apply because an eligible baseline already exists. Upload
|
||||
only the individually reviewed records prepared in step 5 with the shared
|
||||
`upload_record(local_path, record, strict=True)` helper. Stop on the first
|
||||
failure and report exactly which records reached HF; do not silently rerun or
|
||||
replicate the remainder.
|
||||
|
||||
Never bulk upload the tracking or staging root, and never modify another
|
||||
model's directory in the same operation.
|
||||
|
||||
### 8. Report outcome and offer cleanup
|
||||
|
||||
@@ -405,9 +546,14 @@ distinguish an accepted baseline shift from a hidden regression.
|
||||
After the upload is verified, ask whether the user wants to clear temporary
|
||||
local state. Explain what each directory is for:
|
||||
|
||||
- `PERFORMANCE_TRACKING_ROOT`, usually `/tmp/perf-tracking`: local synced
|
||||
mirror of `FastVideo/performance-tracking` plus the prepared local seed
|
||||
records used for scoped upload.
|
||||
- `PERFORMANCE_TRACKING_ROOT`, usually `/tmp/perf-tracking`: read-only local
|
||||
synced mirror used for operator review and reporting. First-v2 preparation
|
||||
independently proves remote state from a fresh temporary HF snapshot.
|
||||
- `PERFORMANCE_RESEED_STAGING_ROOT`, usually
|
||||
`/tmp/performance_reseed_prepared`: prepared local seed records used for the
|
||||
scoped upload, plus the identity reservation and digest manifest. Keeping
|
||||
this separate prevents aborted preparations from appearing in later
|
||||
baseline reads.
|
||||
- `/tmp/performance_reseed_backup/<...>`: local backup of the target model's
|
||||
pre-reseed HF history plus `PROVENANCE.txt`, kept so a bad reseed can be
|
||||
audited or corrected.
|
||||
@@ -417,14 +563,18 @@ local state. Explain what each directory is for:
|
||||
Ask:
|
||||
|
||||
> Reseed succeeded. Do you want me to delete the local temp tracking mirror,
|
||||
> source downloads, and reseed backup under `/tmp`? These files are local
|
||||
> safety/audit artifacts only; HF already has the uploaded records.
|
||||
> this reseed's prepared staging records, source downloads, and reseed backup
|
||||
> under `/tmp`? These files are local safety/audit artifacts only; HF already
|
||||
> has the uploaded records.
|
||||
>
|
||||
> Reply `cleanup reseed temp` to delete them, anything else to keep them.
|
||||
|
||||
Do not delete anything unless the user replies exactly
|
||||
`cleanup reseed temp`. If cleanup is requested, remove only the specific
|
||||
directories created for this reseed. Never remove unrelated `/tmp` contents.
|
||||
directories and prepared record paths created for this reseed. Do not remove
|
||||
the shared staging root when it contains other records. Remove this operation's
|
||||
identity reservation only with its prepared records and manifest, and never
|
||||
remove unrelated `/tmp` contents.
|
||||
|
||||
## Failure modes and handling
|
||||
|
||||
@@ -436,19 +586,34 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
|
||||
against the source batch median by more than `max_intra_batch_regression`.
|
||||
Ask for cleaner sources or a reviewed explanation before continuing.
|
||||
- **Too few source records to move the median.** Continue only after making
|
||||
clear that one or two records may not immediately move the last-5 median.
|
||||
clear that one or two records may not immediately move an existing last-5
|
||||
median. This warning does not block a first v2 calibration seed for an exact
|
||||
identity with no eligible baseline yet.
|
||||
- **The source results are noisy or suspicious.** Stop. Reseeding amplifies
|
||||
those measurements into the baseline, so they must be reviewed first.
|
||||
- **HF sync fails.** Stop for destructive reseeds. A stale or empty sync can
|
||||
make the old baseline look missing.
|
||||
- **The exact v2 identity already has an eligible baseline.** Stop. The
|
||||
`CALIBRATION_NEEDED` artifact is stale; use the reviewed baseline-shift path
|
||||
instead of the first-seed utility.
|
||||
- **The workload/variant/version trusts another recipe.** Stop. The source is
|
||||
stale relative to the current recipe cohort and must not bypass
|
||||
`RECIPE_MISMATCH` by creating a second trusted recipe.
|
||||
- **The staging root already has a prepared seed for the exact identity.**
|
||||
Stop and reuse, upload, or explicitly clean that preparation. Do not prepare
|
||||
another copy of the same measurement.
|
||||
- **The conditional Hub commit loses its parent race.** Stop without retrying.
|
||||
Keep the preparation, refresh and review the new remote state, then request
|
||||
a new explicit `upload` only if the seed is still valid.
|
||||
- **Candidate still violates fixed thresholds.** Report that this skill only
|
||||
handles the rolling HF baseline; update benchmark JSON thresholds in code
|
||||
review if maintainers accept the new absolute limit.
|
||||
- **The user aborts at either confirmation.** Leave the backup and prepared
|
||||
records on disk. Nothing should be uploaded.
|
||||
- **The user declines cleanup.** Keep `/tmp/perf-tracking`, the source
|
||||
download directory if any, and `/tmp/performance_reseed_backup/<...>` in
|
||||
place for audit/debugging.
|
||||
- **The user declines cleanup.** Keep `/tmp/perf-tracking`, the prepared seed
|
||||
records under `/tmp/performance_reseed_prepared`, the source download
|
||||
directory if any, and `/tmp/performance_reseed_backup/<...>` in place for
|
||||
audit/debugging.
|
||||
- **A bad seed was uploaded.** Use the backup and HF history to identify the
|
||||
uploaded file, then remove or supersede it with an explicitly reviewed
|
||||
corrective record. Do not silently rewrite unrelated history.
|
||||
@@ -459,8 +624,9 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
|
||||
intentional baseline replacement.
|
||||
- `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
|
||||
median comparison, and persistence rules.
|
||||
- `fastvideo/performance/hf_store.py` — HF sync, record loading,
|
||||
`sanitize()`, and `upload_record()`.
|
||||
- `fastvideo/performance/hf_store.py` — HF sync and record loading helpers.
|
||||
- `fastvideo/tests/performance/seed_baseline.py` — first-seed preparation,
|
||||
staging reservation, manifest validation, and conditional batch upload.
|
||||
- `fastvideo/tests/performance/test_inference_performance.py` — source result
|
||||
JSON schema.
|
||||
- `.buildkite/performance-benchmarks/tests/*.json` — fixed absolute benchmark
|
||||
@@ -473,3 +639,4 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
|
||||
| 2026-05-03 | Initial version. Sister workflow to `reseed-ssim-references`, scoped to one performance `(model_id, gpu_type)` baseline seed with backup, confirmation, provenance, and `success=true` upload. |
|
||||
| 2026-05-03 | Previous policy: replicate one approved shifted source result into 3 success records by default, or 5 only when explicitly requested. Add provenance marker for replicated-source reseeds. Superseded by the 2026-05-08 dynamic multi-source policy. |
|
||||
| 2026-05-08 | Replace fixed 3/5 replication with dynamic multi-source reseeding: upload one seed record per reviewed source JSON, validate intra-batch consistency, move backup/source scratch under `/tmp`, and ask whether to clean temp state after successful upload. |
|
||||
| 2026-07-13 | Keep first-v2-seed preparation outside the canonical mirror, reserve staging identities atomically, reject stale or replayed calibration seeds, and upload reviewed manifests with a single parent-guarded Hub commit. |
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"benchmark_version": 3,
|
||||
"description": "Wan2.1 T2V 1.3B inference performance",
|
||||
"model": {
|
||||
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
|
||||
@@ -76,7 +76,7 @@ EFFECTIVE_PR=${BUILDKITE_PULL_REQUEST:-false}
|
||||
if [ "$EFFECTIVE_PR" = "false" ] && [ -n "${PR_NUMBER:-}" ]; then
|
||||
EFFECTIVE_PR=$PR_NUMBER
|
||||
fi
|
||||
MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} BUILDKITE_BUILD_URL=${BUILDKITE_BUILD_URL:-} BUILDKITE_BUILD_ID=${BUILDKITE_BUILD_ID:-} BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID:-} IMAGE_VERSION=$IMAGE_VERSION"
|
||||
MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} BUILDKITE_SOURCE=${BUILDKITE_SOURCE:-} TEST_SCOPE=${TEST_SCOPE:-} BUILDKITE_BUILD_URL=${BUILDKITE_BUILD_URL:-} BUILDKITE_BUILD_ID=${BUILDKITE_BUILD_ID:-} BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID:-} IMAGE_VERSION=$IMAGE_VERSION"
|
||||
|
||||
POST_RUN_HOOK=""
|
||||
|
||||
|
||||
@@ -34,7 +34,6 @@ env
|
||||
*.log
|
||||
weights/
|
||||
logs/
|
||||
/Z-Image/
|
||||
official_weights/
|
||||
converted_weights/
|
||||
|
||||
|
||||
@@ -67,7 +67,7 @@ MODEL_REGISTRY = {
|
||||
},
|
||||
}
|
||||
|
||||
DEFAULT_MODEL_ID = "fast-ltx2"
|
||||
DEFAULT_MODEL_ID = "fast-ltx23"
|
||||
|
||||
ACTIVE_MODEL_ID = (os.getenv("DREAMVERSE_MODEL_ID", "").strip() or DEFAULT_MODEL_ID)
|
||||
if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
|
||||
@@ -76,14 +76,11 @@ if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
|
||||
# Active model configuration
|
||||
MODEL_CONFIG = MODEL_REGISTRY[ACTIVE_MODEL_ID]
|
||||
|
||||
# Generation limits
|
||||
SESSION_TIMEOUT_SECONDS = 300
|
||||
|
||||
# Frame settings
|
||||
NUM_FRAMES = 121
|
||||
FRAME_HEIGHT = 1088
|
||||
FRAME_WIDTH = 1920
|
||||
NUM_INFERENCE_STEPS = 5
|
||||
NUM_INFERENCE_STEPS = 6
|
||||
JPEG_QUALITY = 100
|
||||
BATCH_SIZE = 3
|
||||
|
||||
@@ -168,6 +165,9 @@ def _optional_env(*names: str) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
# Generation limits
|
||||
SESSION_TIMEOUT_SECONDS = _env_int("DREAMVERSE_SESSION_TIMEOUT_SECONDS", 300)
|
||||
|
||||
DEVTOOLS_ENABLED = _env_bool("FASTVIDEO_ENABLE_DEVTOOLS", False)
|
||||
PROMPT_SAFETY_ENABLED = _env_bool("FASTVIDEO_ENABLE_PROMPT_SAFETY", False)
|
||||
DREAMVERSE_MAX_AUTOTUNE = _env_bool("DREAMVERSE_MAX_AUTOTUNE", True)
|
||||
|
||||
@@ -1023,7 +1023,11 @@ def get_available_gpus() -> list[int]:
|
||||
"""Get list of available GPU IDs from environment or auto-detect."""
|
||||
cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
|
||||
if cuda_visible:
|
||||
visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
|
||||
try:
|
||||
visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
|
||||
except ValueError as exc:
|
||||
raise RuntimeError("CUDA_VISIBLE_DEVICES must be a comma-separated list of integer GPU "
|
||||
f"indices (got {cuda_visible!r}); GPU UUIDs are not supported.") from exc
|
||||
return _limit_gpu_ids(visible_gpu_ids)
|
||||
|
||||
# Auto-detect available GPUs
|
||||
|
||||
@@ -206,7 +206,9 @@ def cli() -> None:
|
||||
args = parser.parse_args()
|
||||
|
||||
_install_heartbeat_log_filter()
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
# A 15MB init image (session_init_image.MAX_SESSION_INIT_IMAGE_BYTES) is ~20MB
|
||||
# as a base64 ws message, above uvicorn's default 16MiB frame cap.
|
||||
uvicorn.run(app, host=args.host, port=args.port, ws_max_size=32 * 1024 * 1024)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -30,11 +30,11 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from dreamverse._deps import require_dreamverse_runtime_deps
|
||||
from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP
|
||||
from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP, SESSION_TIMEOUT_SECONDS
|
||||
from dreamverse.session_init_image import cleanup_session_init_image, persist_session_init_image
|
||||
from dreamverse.utils import _resolve_generation_segment_cap
|
||||
|
||||
LATENCY_MS = 200
|
||||
SESSION_TIMEOUT_SECONDS = 300
|
||||
MOCK_FRAME_WIDTH = 640
|
||||
MOCK_FRAME_HEIGHT = 352
|
||||
MOCK_FPS = 24
|
||||
@@ -334,6 +334,7 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
auto_extension_enabled = bool(init_data.get("auto_extension_enabled", False))
|
||||
loop_generation_enabled = bool(init_data.get("loop_generation_enabled", False))
|
||||
single_clip_mode = bool(init_data.get("single_clip_mode", False))
|
||||
manual_continuation_mode = bool(init_data.get("manual_continuation_mode", False))
|
||||
generation_paused = False
|
||||
|
||||
if init_type == "session_init_v2":
|
||||
@@ -344,7 +345,8 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
incoming_prompts = []
|
||||
|
||||
curated_prompts = [prompt.strip() for prompt in incoming_prompts if isinstance(prompt, str) and prompt.strip()]
|
||||
generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
|
||||
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
|
||||
and len(curated_prompts) == 0)
|
||||
|
||||
try:
|
||||
session_init_image = persist_session_init_image(init_data.get("initial_image"))
|
||||
@@ -373,7 +375,10 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
prompt_sources_blocked = False
|
||||
pending_seed_reset = False
|
||||
pending_seed_reset_reason = ""
|
||||
pending_simple_submission: PromptSubmission | None = None
|
||||
pending_simple_submission: PromptSubmission | None = (PromptSubmission(
|
||||
prompt_id=str(init_data.get("initial_rollout_prompt_id") or uuid.uuid4()),
|
||||
raw_prompt=initial_rollout_prompt,
|
||||
) if manual_continuation_mode and initial_rollout_prompt else None)
|
||||
single_clip_waiting_for_request = False
|
||||
rollout_waiting_for_rewrite = False
|
||||
initial_rollout_waiting_for_rewrite = generation_paused
|
||||
@@ -393,14 +398,26 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
|
||||
async def send_stream_start(seed_reason: str) -> None:
|
||||
await ws_send_json({
|
||||
"type": "ltx2_stream_start",
|
||||
"total_segments": len(curated_prompts),
|
||||
"preset_id": preset_id,
|
||||
"stream_mode": "av_fmp4",
|
||||
"live_mode": True,
|
||||
"loop_generation_enabled": loop_generation_enabled,
|
||||
"loop_iteration": loop_iteration,
|
||||
"generation_segment_cap": 0,
|
||||
"type":
|
||||
"ltx2_stream_start",
|
||||
"total_segments":
|
||||
len(curated_prompts),
|
||||
"preset_id":
|
||||
preset_id,
|
||||
"stream_mode":
|
||||
"av_fmp4",
|
||||
"live_mode":
|
||||
True,
|
||||
"loop_generation_enabled":
|
||||
loop_generation_enabled,
|
||||
"loop_iteration":
|
||||
loop_iteration,
|
||||
"generation_segment_cap":
|
||||
_resolve_generation_segment_cap(
|
||||
single_clip_mode=single_clip_mode,
|
||||
cap=GENERATION_SEGMENT_CAP,
|
||||
manual_continuation_mode=manual_continuation_mode,
|
||||
),
|
||||
})
|
||||
if seed_reason == "init":
|
||||
await ws_send_json({
|
||||
@@ -493,6 +510,7 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
nonlocal auto_extension_enabled
|
||||
nonlocal loop_generation_enabled
|
||||
nonlocal single_clip_mode
|
||||
nonlocal manual_continuation_mode
|
||||
nonlocal generation_paused
|
||||
nonlocal seed_prompt_memory
|
||||
nonlocal curated_prompts
|
||||
@@ -537,6 +555,7 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
auto_extension_enabled = bool(payload.get("auto_extension_enabled", False))
|
||||
loop_generation_enabled = bool(payload.get("loop_generation_enabled", False))
|
||||
single_clip_mode = bool(payload.get("single_clip_mode", False))
|
||||
manual_continuation_mode = bool(payload.get("manual_continuation_mode", False))
|
||||
|
||||
seed_prompt_memory = list(next_curated_prompts)
|
||||
curated_prompts = list(seed_prompt_memory)
|
||||
@@ -544,10 +563,14 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
segment_idx = 0
|
||||
pending_seed_reset = False
|
||||
pending_seed_reset_reason = ""
|
||||
pending_simple_submission = None
|
||||
pending_simple_submission = (PromptSubmission(
|
||||
prompt_id=str(payload.get("initial_rollout_prompt_id") or uuid.uuid4()),
|
||||
raw_prompt=initial_rollout_prompt,
|
||||
) if manual_continuation_mode and initial_rollout_prompt else None)
|
||||
single_clip_waiting_for_request = False
|
||||
rollout_waiting_for_rewrite = False
|
||||
generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
|
||||
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
|
||||
and len(curated_prompts) == 0)
|
||||
initial_rollout_waiting_for_rewrite = generation_paused
|
||||
rewrite_restart_pending = False
|
||||
loop_iteration = 0
|
||||
@@ -968,10 +991,11 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
project_stream_started = True
|
||||
await send_stream_start(pending_seed_reset_reason)
|
||||
pending_seed_reset_reason = ""
|
||||
if pending_simple_submission is not None:
|
||||
submission = pending_simple_submission
|
||||
pending_simple_submission = None
|
||||
await promote_submission_to_ready(submission)
|
||||
|
||||
if pending_simple_submission is not None:
|
||||
submission = pending_simple_submission
|
||||
pending_simple_submission = None
|
||||
await promote_submission_to_ready(submission)
|
||||
|
||||
if generation_paused:
|
||||
await asyncio.sleep(0.05)
|
||||
@@ -981,8 +1005,8 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
await asyncio.sleep(0.05)
|
||||
continue
|
||||
|
||||
if (not single_clip_mode and not rollout_waiting_for_rewrite and GENERATION_SEGMENT_CAP > 0
|
||||
and segment_idx >= GENERATION_SEGMENT_CAP):
|
||||
if (not single_clip_mode and not manual_continuation_mode and not rollout_waiting_for_rewrite
|
||||
and GENERATION_SEGMENT_CAP > 0 and segment_idx >= GENERATION_SEGMENT_CAP):
|
||||
rollout_waiting_for_rewrite = True
|
||||
loop_generation_enabled = False
|
||||
project_stream_started = False
|
||||
@@ -1217,7 +1241,9 @@ def cli() -> None:
|
||||
print(f"Starting mock server with {LATENCY_MS}ms latency on port {args.port}")
|
||||
|
||||
_install_heartbeat_log_filter()
|
||||
uvicorn.run(app, host="0.0.0.0", port=args.port)
|
||||
# A 15MB init image (session_init_image.MAX_SESSION_INIT_IMAGE_BYTES) is ~20MB
|
||||
# as a base64 ws message, above uvicorn's default 16MiB frame cap.
|
||||
uvicorn.run(app, host="0.0.0.0", port=args.port, ws_max_size=32 * 1024 * 1024)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -313,6 +313,33 @@ def _extract_content_or_empty(response_json: dict[str, Any]) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
def _find_balanced_object_end(text: str, start: int) -> int:
|
||||
"""Return the index just past the brace-balanced span opening at
|
||||
``text[start] == '{'``, honoring JSON string literals and escapes, or -1
|
||||
if the braces never balance (i.e. the object was truncated)."""
|
||||
depth = 0
|
||||
in_string = False
|
||||
escaped = False
|
||||
for i in range(start, len(text)):
|
||||
ch = text[i]
|
||||
if in_string:
|
||||
if escaped:
|
||||
escaped = False
|
||||
elif ch == "\\":
|
||||
escaped = True
|
||||
elif ch == '"':
|
||||
in_string = False
|
||||
elif ch == '"':
|
||||
in_string = True
|
||||
elif ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return i + 1
|
||||
return -1
|
||||
|
||||
|
||||
def _parse_json_response(content: str) -> dict[str, Any]:
|
||||
text = content.strip()
|
||||
if not text:
|
||||
@@ -330,6 +357,7 @@ def _parse_json_response(content: str) -> dict[str, Any]:
|
||||
r"```(?:json)?\s*([\s\S]*?)```",
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
last_fenced: dict[str, Any] | None = None
|
||||
for match in fence_pattern.finditer(text):
|
||||
block = match.group(1).strip()
|
||||
if not block:
|
||||
@@ -337,21 +365,37 @@ def _parse_json_response(content: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = json.loads(block)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
last_fenced = parsed
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if last_fenced is not None:
|
||||
return last_fenced
|
||||
|
||||
# Fall back to scanning for the first decodable JSON object in free-form text.
|
||||
# Scan for all decodable JSON objects and return the last — chain-of-thought
|
||||
# models emit draft JSON mid-reasoning; the final answer is always last.
|
||||
decoder = json.JSONDecoder()
|
||||
for idx, char in enumerate(text):
|
||||
if char != "{":
|
||||
continue
|
||||
last_parsed: dict[str, Any] | None = None
|
||||
pos = 0
|
||||
while (idx := text.find("{", pos)) != -1:
|
||||
try:
|
||||
parsed, _ = decoder.raw_decode(text[idx:])
|
||||
parsed, consumed = decoder.raw_decode(text[idx:])
|
||||
except json.JSONDecodeError:
|
||||
# Skip the whole failed object rather than rescanning inside it:
|
||||
# fragments nested in a malformed or truncated (finish_reason=
|
||||
# length) object must not override an earlier complete object.
|
||||
span_end = _find_balanced_object_end(text, idx)
|
||||
if span_end == -1:
|
||||
break
|
||||
pos = span_end
|
||||
continue
|
||||
# Skip past the consumed span so nested braces inside a decoded
|
||||
# object are not re-parsed as standalone objects.
|
||||
pos = idx + consumed
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
last_parsed = parsed
|
||||
|
||||
if last_parsed is not None:
|
||||
return last_parsed
|
||||
|
||||
raise ValueError("No JSON object found in assistant response.")
|
||||
|
||||
@@ -381,7 +425,7 @@ def _format_locked_segments(locked_segments: list[str]) -> str:
|
||||
|
||||
class PromptEnhancer:
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
self.provider = PROMPT_PROVIDER
|
||||
self.provider_label = _resolve_provider_label(PROMPT_PROVIDER)
|
||||
self.api_key = PROMPT_API_KEY
|
||||
@@ -1417,15 +1461,12 @@ class PromptEnhancer:
|
||||
locked_text = _format_locked_segments(locked_segments_clean)
|
||||
request_system_prompt = self.enhance_system_prompt
|
||||
user_payload = {
|
||||
"request": (
|
||||
"<locked_segments>\n"
|
||||
f"{locked_text}\n"
|
||||
"</locked_segments>\n\n"
|
||||
f"<conditioning_prompt>{cleaned}</conditioning_prompt>\n\n"
|
||||
f"Write exactly one new segment ({next_segment_key}) "
|
||||
"continuing from the locked segments. "
|
||||
'Respond with valid JSON only as {"next_prompt": "..."}.' # noqa: E501
|
||||
),
|
||||
"request": ("<locked_segments>\n"
|
||||
f"{locked_text}\n"
|
||||
"</locked_segments>\n\n"
|
||||
f"<conditioning_prompt>{cleaned}</conditioning_prompt>\n\n"
|
||||
f"Write exactly one new segment ({next_segment_key}) "
|
||||
"continuing from the locked segments."),
|
||||
}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
@@ -1457,6 +1498,7 @@ class PromptEnhancer:
|
||||
body=request_body,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
_enhance_print("INFO", f"raw_response: {response_content}")
|
||||
if is_single_clip_mode:
|
||||
prompt = self._extract_single_clip_prompt(response_content)
|
||||
else:
|
||||
@@ -1539,7 +1581,7 @@ class PromptEnhancer:
|
||||
f"Write exactly one new segment ({next_segment_key}) "
|
||||
"that continues linearly from the locked segments. "
|
||||
"Infer the next narrative beat from this history. "
|
||||
'Respond with valid JSON only as {"next_prompt": "..."}.' # noqa: E501
|
||||
'Respond with valid JSON only: {"next_prompt": "<your segment description here>"}.' # noqa: E501
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ from __future__ import annotations
|
||||
# mypy: ignore-errors
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
@@ -52,6 +53,14 @@ if TYPE_CHECKING:
|
||||
from dreamverse.prompt_enhancer import PromptEnhancer
|
||||
from dreamverse.prompt_safety import PromptSafetyFilter
|
||||
|
||||
# Optional append-only log of every generated segment prompt; unset disables it.
|
||||
SEGMENT_PROMPT_LOG_PATH = os.environ.get("DREAMVERSE_SEGMENT_PROMPT_LOG", "")
|
||||
|
||||
|
||||
def _append_segment_prompt_log(path: str, text: str) -> None:
|
||||
with open(path, "a") as f:
|
||||
f.write(text)
|
||||
|
||||
|
||||
class SessionController:
|
||||
"""Runs one WebSocket session from accept() through disconnect."""
|
||||
@@ -198,6 +207,7 @@ class SessionController:
|
||||
auto_extension_enabled = bool(init_data.get("auto_extension_enabled", False))
|
||||
loop_generation_enabled = bool(init_data.get("loop_generation_enabled", False))
|
||||
single_clip_mode = bool(init_data.get("single_clip_mode", False))
|
||||
manual_continuation_mode = bool(init_data.get("manual_continuation_mode", False))
|
||||
rewrite_model = self.prompt_enhancer.resolve_rewrite_model(init_data.get("rewrite_model"))
|
||||
rewrite_system_prompt_override = str(init_data.get("rewrite_window_system_prompt") or "").strip()
|
||||
rewrite_user_system_prompt_override = str(init_data.get("rewrite_user_system_prompt") or "").strip()
|
||||
@@ -282,6 +292,9 @@ class SessionController:
|
||||
# Session queues and mutable state.
|
||||
raw_prompt_queue: asyncio.Queue[PromptSubmission] = asyncio.Queue()
|
||||
ready_prompt_queue: asyncio.Queue[ReadyPrompt] = asyncio.Queue()
|
||||
# Submissions dequeued by prompt_worker_loop but not yet resolved; while
|
||||
# non-zero the prompt sources are busy, not drained.
|
||||
prompt_enhancement_inflight = 0
|
||||
|
||||
curated_idx = 0
|
||||
segment_idx = 0
|
||||
@@ -291,13 +304,20 @@ class SessionController:
|
||||
generation_cap_blocked = False
|
||||
auto_extension_blocked_segment_idx: int | None = None
|
||||
prompt_sources_drained_logged = False
|
||||
generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
|
||||
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt and not single_clip_mode
|
||||
and len(curated_prompts) == 0)
|
||||
pending_seed_reset = False
|
||||
pending_seed_reset_reason = ""
|
||||
pending_reset_conditioning = False
|
||||
loop_iteration = 0 if generation_paused else 1
|
||||
force_curated_restart_segment = False
|
||||
pending_simple_prompt_submission: PromptSubmission | None = None
|
||||
# The frontend records the opening scene under this id; reuse it so
|
||||
# prompt lifecycle events for the opening prompt reach that record.
|
||||
pending_simple_prompt_submission: PromptSubmission | None = (PromptSubmission(
|
||||
prompt_id=str(init_data.get("initial_rollout_prompt_id") or uuid.uuid4()),
|
||||
raw_prompt=initial_rollout_prompt,
|
||||
created_at_s=time.time(),
|
||||
) if manual_continuation_mode and initial_rollout_prompt else None)
|
||||
single_clip_waiting_for_request = False
|
||||
rollout_waiting_for_rewrite = False
|
||||
initial_rollout_waiting_for_rewrite = generation_paused
|
||||
@@ -306,6 +326,7 @@ class SessionController:
|
||||
project_active = True
|
||||
project_stream_started = False
|
||||
pending_project_end = False
|
||||
segment_prompt_log_warned = False
|
||||
|
||||
def replace_session_init_image(initial_image_payload: object) -> None:
|
||||
nonlocal session_init_image
|
||||
@@ -424,6 +445,7 @@ class SessionController:
|
||||
nonlocal auto_extension_enabled
|
||||
nonlocal loop_generation_enabled
|
||||
nonlocal single_clip_mode
|
||||
nonlocal manual_continuation_mode
|
||||
nonlocal generation_paused
|
||||
nonlocal curated_prompts
|
||||
nonlocal seed_prompt_memory
|
||||
@@ -458,6 +480,7 @@ class SessionController:
|
||||
next_auto_extension_enabled = bool(payload.get("auto_extension_enabled", False))
|
||||
next_loop_generation_enabled = bool(payload.get("loop_generation_enabled", False))
|
||||
next_single_clip_mode = bool(payload.get("single_clip_mode", False))
|
||||
next_manual_continuation_mode = bool(payload.get("manual_continuation_mode", False))
|
||||
|
||||
next_preset_id = str(payload.get("preset_id") or "").strip()
|
||||
if next_preset_id:
|
||||
@@ -510,6 +533,7 @@ class SessionController:
|
||||
auto_extension_enabled = next_auto_extension_enabled
|
||||
loop_generation_enabled = next_loop_generation_enabled
|
||||
single_clip_mode = next_single_clip_mode
|
||||
manual_continuation_mode = next_manual_continuation_mode
|
||||
rewrite_model = next_rewrite_model
|
||||
rewrite_system_prompt_override = (next_rewrite_system_prompt_override)
|
||||
rewrite_user_system_prompt_override = (next_rewrite_user_system_prompt_override)
|
||||
@@ -530,10 +554,15 @@ class SessionController:
|
||||
generation_cap_blocked = False
|
||||
auto_extension_blocked_segment_idx = None
|
||||
prompt_sources_drained_logged = False
|
||||
pending_simple_prompt_submission = None
|
||||
pending_simple_prompt_submission = (PromptSubmission(
|
||||
prompt_id=str(payload.get("initial_rollout_prompt_id") or uuid.uuid4()),
|
||||
raw_prompt=initial_rollout_prompt,
|
||||
created_at_s=time.time(),
|
||||
) if manual_continuation_mode and initial_rollout_prompt else None)
|
||||
single_clip_waiting_for_request = False
|
||||
rollout_waiting_for_rewrite = False
|
||||
generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0)
|
||||
generation_paused = bool(not manual_continuation_mode and initial_rollout_prompt
|
||||
and not single_clip_mode and len(curated_prompts) == 0)
|
||||
initial_rollout_waiting_for_rewrite = generation_paused
|
||||
rewrite_restart_pending = False
|
||||
loop_iteration = 0
|
||||
@@ -942,6 +971,7 @@ class SessionController:
|
||||
_resolve_generation_segment_cap(
|
||||
single_clip_mode=single_clip_mode,
|
||||
cap=GENERATION_SEGMENT_CAP,
|
||||
manual_continuation_mode=manual_continuation_mode,
|
||||
),
|
||||
})
|
||||
continue
|
||||
@@ -1005,11 +1035,9 @@ class SessionController:
|
||||
continue
|
||||
|
||||
async def prompt_worker_loop():
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
submission = await asyncio.wait_for(raw_prompt_queue.get(), timeout=0.1)
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
nonlocal prompt_enhancement_inflight
|
||||
|
||||
async def process_submission(submission: PromptSubmission) -> None:
|
||||
_main_print("INFO", f"Received user prompt for enhancement: {submission.raw_prompt}")
|
||||
prompt_id = submission.prompt_id
|
||||
raw_prompt = submission.raw_prompt
|
||||
@@ -1032,8 +1060,9 @@ class SessionController:
|
||||
await ws_send_json({
|
||||
"type": "error",
|
||||
"message": blocked_raw_prompt_error,
|
||||
"prompt_id": prompt_id,
|
||||
})
|
||||
continue
|
||||
return
|
||||
await log_event(
|
||||
"enhance_request",
|
||||
{
|
||||
@@ -1093,8 +1122,9 @@ class SessionController:
|
||||
await ws_send_json({
|
||||
"type": "error",
|
||||
"message": blocked_final_prompt_error,
|
||||
"prompt_id": prompt_id,
|
||||
})
|
||||
continue
|
||||
return
|
||||
if result.fallback_used or not final_prompt:
|
||||
source = "user_enhancement_failed"
|
||||
_main_print(
|
||||
@@ -1113,7 +1143,7 @@ class SessionController:
|
||||
})
|
||||
# Enhancement is strict JSON-only; do not enqueue raw
|
||||
# prompt when enhancement fails.
|
||||
continue
|
||||
return
|
||||
else:
|
||||
source = "user_enhanced"
|
||||
await ws_send_json({
|
||||
@@ -1130,6 +1160,7 @@ class SessionController:
|
||||
source=source,
|
||||
fallback_used=result.fallback_used,
|
||||
loop_iteration=loop_iteration,
|
||||
raw_prompt=raw_prompt,
|
||||
))
|
||||
else:
|
||||
await ready_prompt_queue.put(
|
||||
@@ -1139,6 +1170,7 @@ class SessionController:
|
||||
source="user_raw",
|
||||
fallback_used=False,
|
||||
loop_iteration=loop_iteration,
|
||||
raw_prompt=raw_prompt,
|
||||
))
|
||||
await ws_send_json({
|
||||
"type": "prompt_ready",
|
||||
@@ -1148,6 +1180,21 @@ class SessionController:
|
||||
"latency_ms": 0.0,
|
||||
})
|
||||
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
submission = raw_prompt_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
await asyncio.sleep(0.1)
|
||||
continue
|
||||
# Dequeue and increment without an await in between so the
|
||||
# generation loop never sees an empty queue with zero in flight
|
||||
# while this submission is still being enhanced.
|
||||
prompt_enhancement_inflight += 1
|
||||
try:
|
||||
await process_submission(submission)
|
||||
finally:
|
||||
prompt_enhancement_inflight -= 1
|
||||
|
||||
def queue_snapshot() -> dict[str, object]:
|
||||
return {
|
||||
"user_ready": ready_prompt_queue.qsize(),
|
||||
@@ -1300,6 +1347,7 @@ class SessionController:
|
||||
_resolve_generation_segment_cap(
|
||||
single_clip_mode=single_clip_mode,
|
||||
cap=GENERATION_SEGMENT_CAP,
|
||||
manual_continuation_mode=manual_continuation_mode,
|
||||
),
|
||||
})
|
||||
await ws_send_json({
|
||||
@@ -1359,6 +1407,7 @@ class SessionController:
|
||||
_resolve_generation_segment_cap(
|
||||
single_clip_mode=single_clip_mode,
|
||||
cap=GENERATION_SEGMENT_CAP,
|
||||
manual_continuation_mode=manual_continuation_mode,
|
||||
),
|
||||
})
|
||||
if nonlocal_reason == "loop_restart":
|
||||
@@ -1388,8 +1437,9 @@ class SessionController:
|
||||
await raw_prompt_queue.put(pending_simple_prompt_submission)
|
||||
pending_simple_prompt_submission = None
|
||||
|
||||
if (not single_clip_mode and not generation_cap_blocked and not rollout_waiting_for_rewrite
|
||||
and GENERATION_SEGMENT_CAP > 0 and generated_segment_count >= GENERATION_SEGMENT_CAP):
|
||||
if (not single_clip_mode and not manual_continuation_mode and not generation_cap_blocked
|
||||
and not rollout_waiting_for_rewrite and GENERATION_SEGMENT_CAP > 0
|
||||
and generated_segment_count >= GENERATION_SEGMENT_CAP):
|
||||
loop_generation_enabled = False
|
||||
rollout_waiting_for_rewrite = True
|
||||
_main_print(
|
||||
@@ -1542,7 +1592,10 @@ class SessionController:
|
||||
if single_clip_mode:
|
||||
await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0)
|
||||
continue
|
||||
if not prompt_sources_drained_logged:
|
||||
# A raw submission still queued or being enhanced will produce a
|
||||
# ready prompt shortly; that is not a drained/blocked state.
|
||||
enhancement_pending = (raw_prompt_queue.qsize() > 0 or prompt_enhancement_inflight > 0)
|
||||
if not prompt_sources_drained_logged and not enhancement_pending:
|
||||
snapshot = queue_snapshot()
|
||||
_main_print(
|
||||
"WARN",
|
||||
@@ -1581,6 +1634,24 @@ class SessionController:
|
||||
total_segments_hint = max(segment_idx, len(curated_prompts))
|
||||
prompt = selected.prompt
|
||||
locked_segment_prompts.append(prompt)
|
||||
if SEGMENT_PROMPT_LOG_PATH:
|
||||
_ts = time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
_lines = [
|
||||
f"\n=== Segment {segment_idx} [{_ts}] source={selected.source} client={client_id[:8]} ===",
|
||||
]
|
||||
if selected.raw_prompt and selected.raw_prompt != prompt:
|
||||
_lines.append(f"User: {selected.raw_prompt}")
|
||||
_lines.append(f"Rewritten: {prompt}")
|
||||
try:
|
||||
await asyncio.to_thread(_append_segment_prompt_log, SEGMENT_PROMPT_LOG_PATH,
|
||||
"\n".join(_lines) + "\n")
|
||||
except Exception as exc:
|
||||
if not segment_prompt_log_warned:
|
||||
segment_prompt_log_warned = True
|
||||
_main_print(
|
||||
"WARN",
|
||||
f"Failed to write segment prompt log {SEGMENT_PROMPT_LOG_PATH}: {exc}",
|
||||
)
|
||||
if (auto_extension_blocked_segment_idx is not None
|
||||
and auto_extension_blocked_segment_idx <= segment_idx):
|
||||
auto_extension_blocked_segment_idx = None
|
||||
|
||||
@@ -19,3 +19,4 @@ class ReadyPrompt:
|
||||
fallback_used: bool = False
|
||||
seed_prompt_index: int | None = None
|
||||
loop_iteration: int | None = None
|
||||
raw_prompt: str | None = None
|
||||
|
||||
@@ -150,12 +150,22 @@ def test_config_enables_prompt_safety_when_requested(monkeypatch):
|
||||
|
||||
def test_config_uses_five_minute_session_timeout(monkeypatch):
|
||||
_set_required_prompt_keys(monkeypatch)
|
||||
monkeypatch.delenv("DREAMVERSE_SESSION_TIMEOUT_SECONDS", raising=False)
|
||||
|
||||
module = _load_config_module()
|
||||
|
||||
assert module.SESSION_TIMEOUT_SECONDS == 300
|
||||
|
||||
|
||||
def test_config_session_timeout_env_override(monkeypatch):
|
||||
_set_required_prompt_keys(monkeypatch)
|
||||
monkeypatch.setenv("DREAMVERSE_SESSION_TIMEOUT_SECONDS", "1800")
|
||||
|
||||
module = _load_config_module()
|
||||
|
||||
assert module.SESSION_TIMEOUT_SECONDS == 1800
|
||||
|
||||
|
||||
def test_config_rejects_invalid_prompt_provider(monkeypatch):
|
||||
monkeypatch.setenv("FASTVIDEO_PROMPT_PROVIDER", "unsupported")
|
||||
_set_required_prompt_keys(monkeypatch)
|
||||
|
||||
@@ -75,12 +75,13 @@ def _run_cli(module, monkeypatch, argv: list[str]) -> list[dict[str, object]]:
|
||||
calls: list[dict[str, object]] = []
|
||||
uvicorn_stub = types.ModuleType("uvicorn")
|
||||
|
||||
def run(app, host: str, port: int) -> None:
|
||||
def run(app, host: str, port: int, **kwargs) -> None:
|
||||
calls.append(
|
||||
{
|
||||
"app": app,
|
||||
"host": host,
|
||||
"port": port,
|
||||
**kwargs,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -104,6 +105,7 @@ def test_server_cli_defaults_to_local_web_port(monkeypatch):
|
||||
"app": server_main.app,
|
||||
"host": "0.0.0.0",
|
||||
"port": 8009,
|
||||
"ws_max_size": 32 * 1024 * 1024,
|
||||
}
|
||||
]
|
||||
|
||||
@@ -121,6 +123,7 @@ def test_server_cli_allows_explicit_host_and_port(monkeypatch):
|
||||
"app": server_main.app,
|
||||
"host": "127.0.0.1",
|
||||
"port": 8123,
|
||||
"ws_max_size": 32 * 1024 * 1024,
|
||||
}
|
||||
]
|
||||
|
||||
@@ -147,6 +150,7 @@ def test_mock_server_cli_defaults_to_local_web_port(monkeypatch):
|
||||
"app": mock_server.app,
|
||||
"host": "0.0.0.0",
|
||||
"port": 8009,
|
||||
"ws_max_size": 32 * 1024 * 1024,
|
||||
}
|
||||
]
|
||||
|
||||
@@ -166,6 +170,7 @@ def test_mock_server_cli_updates_latency(monkeypatch):
|
||||
"app": mock_server.app,
|
||||
"host": "0.0.0.0",
|
||||
"port": 8111,
|
||||
"ws_max_size": 32 * 1024 * 1024,
|
||||
}
|
||||
]
|
||||
assert mock_server.LATENCY_MS == 321
|
||||
|
||||
@@ -290,6 +290,152 @@ def test_mock_server_supports_initial_custom_rollout_prompt():
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
|
||||
|
||||
def test_mock_server_manual_mode_streams_initial_prompt_without_rewrite_or_cap():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
old_generation_segment_cap = mock_server.GENERATION_SEGMENT_CAP
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 1
|
||||
mock_server.GENERATION_SEGMENT_CAP = 1
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "custom_editable",
|
||||
"preset_label": "Custom rollout",
|
||||
"curated_prompts": [],
|
||||
"initial_rollout_prompt": "A drone skims a neon canyon",
|
||||
"initial_rollout_prompt_id": "steer-prompt-1",
|
||||
"manual_continuation_mode": True,
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
0.08,
|
||||
{
|
||||
"type": "append_prompt",
|
||||
"prompt": "The drone dives toward the river",
|
||||
"prompt_id": "steer-prompt-2",
|
||||
},
|
||||
),
|
||||
(0.30, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
message_types = [payload["type"] for payload in ws.sent_json]
|
||||
assert "rewrite_seed_prompts_started" not in message_types
|
||||
assert "rewrite_seed_prompts_complete" not in message_types
|
||||
assert "ltx2_stream_start" in message_types
|
||||
# cap=1 must not stop a manual-mode session after the first segment
|
||||
assert "ltx2_stream_complete" not in message_types
|
||||
|
||||
prompt_ready_events = [
|
||||
payload for payload in ws.sent_json if payload["type"] == "prompt_ready"
|
||||
]
|
||||
assert [payload["prompt_id"] for payload in prompt_ready_events] == [
|
||||
"steer-prompt-1",
|
||||
"steer-prompt-2",
|
||||
]
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
assert [payload["prompt"] for payload in segment_start_events] == [
|
||||
"A drone skims a neon canyon",
|
||||
"The drone dives toward the river",
|
||||
]
|
||||
segment_source_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "segment_prompt_source"
|
||||
]
|
||||
assert [payload["prompt_id"] for payload in segment_source_events] == [
|
||||
"steer-prompt-1",
|
||||
"steer-prompt-2",
|
||||
]
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
mock_server.GENERATION_SEGMENT_CAP = old_generation_segment_cap
|
||||
|
||||
|
||||
def test_mock_server_project_init_manual_mode_streams_initial_prompt():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
try:
|
||||
mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes"
|
||||
mock_server.LATENCY_MS = 1
|
||||
|
||||
ws = _FakeWebSocket(
|
||||
[
|
||||
(
|
||||
0.0,
|
||||
{
|
||||
"type": "session_init_v2",
|
||||
"preset_id": "test_preset",
|
||||
"curated_prompts": ["segment one"],
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(0.05, {"type": "end_project_keep_session"}),
|
||||
(
|
||||
0.15,
|
||||
{
|
||||
"type": "project_init_v1",
|
||||
"preset_id": "custom_editable",
|
||||
"preset_label": "Custom rollout",
|
||||
"curated_prompts": [],
|
||||
"initial_rollout_prompt": "A drone skims a neon canyon",
|
||||
"initial_rollout_prompt_id": "steer-prompt-1",
|
||||
"manual_continuation_mode": True,
|
||||
"enhancement_enabled": True,
|
||||
"auto_extension_enabled": False,
|
||||
"loop_generation_enabled": False,
|
||||
},
|
||||
),
|
||||
(0.45, {"type": "leave"}),
|
||||
]
|
||||
)
|
||||
|
||||
asyncio.run(mock_server.websocket_endpoint(ws))
|
||||
|
||||
message_types = [payload["type"] for payload in ws.sent_json]
|
||||
assert "project_idle" in message_types
|
||||
project_idle_index = message_types.index("project_idle")
|
||||
# manual-mode restart must not run the rewrite rollout
|
||||
assert "rewrite_seed_prompts_started" not in message_types[project_idle_index:]
|
||||
assert "ltx2_stream_start" in message_types[project_idle_index:]
|
||||
|
||||
prompt_ready_events = [
|
||||
payload for payload in ws.sent_json if payload["type"] == "prompt_ready"
|
||||
]
|
||||
assert [payload["prompt_id"] for payload in prompt_ready_events] == [
|
||||
"steer-prompt-1",
|
||||
]
|
||||
|
||||
segment_start_events = [
|
||||
payload
|
||||
for payload in ws.sent_json
|
||||
if payload["type"] == "ltx2_segment_start"
|
||||
]
|
||||
assert segment_start_events[-1]["prompt"] == "A drone skims a neon canyon"
|
||||
finally:
|
||||
mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes
|
||||
mock_server.LATENCY_MS = old_latency_ms
|
||||
|
||||
|
||||
def test_mock_server_can_start_new_project_without_reconnecting():
|
||||
old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES
|
||||
old_latency_ms = mock_server.LATENCY_MS
|
||||
|
||||
@@ -6,6 +6,7 @@ import os
|
||||
import re
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("CEREBRAS_API_KEY", "dummy")
|
||||
os.environ.setdefault("GROQ_API_KEY", "dummy")
|
||||
@@ -205,6 +206,59 @@ def test_parse_json_response_extracts_first_embedded_object():
|
||||
assert parsed == {"segment_prompts": ["A", "B"]}
|
||||
|
||||
|
||||
def test_parse_json_response_returns_outer_object_not_nested_value():
|
||||
parsed = _parse_json_response(
|
||||
"Final answer: {\"next_prompt\": \"a scene\", \"style\": {\"mood\": \"noir\"}} done."
|
||||
)
|
||||
assert parsed == {"next_prompt": "a scene", "style": {"mood": "noir"}}
|
||||
|
||||
|
||||
def test_parse_json_response_returns_last_of_multiple_objects():
|
||||
parsed = _parse_json_response(
|
||||
"Draft: {\"next_prompt\": \"draft\"}\nRefined: {\"next_prompt\": \"final\"}"
|
||||
)
|
||||
assert parsed == {"next_prompt": "final"}
|
||||
|
||||
|
||||
def test_parse_json_response_ignores_fragments_of_truncated_trailing_object():
|
||||
# finish_reason=length cut the refined object short; the complete draft must
|
||||
# win over a nested fragment of the truncated object.
|
||||
parsed = _parse_json_response(
|
||||
'{"next_prompt": "draft"} refined: {"next_prompt": "final", "style": {"mood": "noir"}'
|
||||
)
|
||||
assert parsed == {"next_prompt": "draft"}
|
||||
|
||||
|
||||
def test_parse_json_response_ignores_fragments_of_mid_string_truncated_object():
|
||||
# Unterminated-string truncation reports the error at the opening quote,
|
||||
# not end-of-text; nested fragments still must not win over the draft.
|
||||
parsed = _parse_json_response(
|
||||
'{"next_prompt": "draft"} refined: {"style": {"mood": "noir"}, "next_prompt": "cut off'
|
||||
)
|
||||
assert parsed == {"next_prompt": "draft"}
|
||||
|
||||
|
||||
def test_parse_json_response_ignores_fragments_of_malformed_object_with_trailing_prose():
|
||||
parsed = _parse_json_response(
|
||||
'{"next_prompt": "draft"} {"final": {"mood": "noir"}, "x": 1 and then some prose'
|
||||
)
|
||||
assert parsed == {"next_prompt": "draft"}
|
||||
|
||||
|
||||
def test_parse_json_response_raises_when_only_object_is_truncated():
|
||||
with pytest.raises(ValueError):
|
||||
_parse_json_response('{"style": {"mood": "noir"}, "next_prompt": "cut off')
|
||||
|
||||
|
||||
def test_parse_json_response_returns_outer_rollout_dict():
|
||||
parsed = _parse_json_response(
|
||||
"{\"rollout\": {\"segment_prompts\": [{\"prompt\": \"a\"}, {\"prompt\": \"b\"}]}}"
|
||||
)
|
||||
assert parsed == {
|
||||
"rollout": {"segment_prompts": [{"prompt": "a"}, {"prompt": "b"}]}
|
||||
}
|
||||
|
||||
|
||||
def test_load_prompt_required_falls_back_to_default_path(tmp_path):
|
||||
fallback_path = tmp_path / "next_segment_system_prompt.md"
|
||||
fallback_path.write_text("fallback prompt\n", encoding="utf-8")
|
||||
|
||||
@@ -15,5 +15,6 @@ def _utc_now_iso() -> str:
|
||||
PROMPT_EXTENSION_FAILURE_USER_MESSAGE = ("Prompt extension failed for this request.")
|
||||
|
||||
|
||||
def _resolve_generation_segment_cap(*, single_clip_mode: bool, cap: int) -> int:
|
||||
return 0 if single_clip_mode else cap
|
||||
def _resolve_generation_segment_cap(*, single_clip_mode: bool, cap: int, manual_continuation_mode: bool = False) -> int:
|
||||
# Steering (manual continuation) lets the user keep going indefinitely, like single-clip mode.
|
||||
return 0 if (single_clip_mode or manual_continuation_mode) else cap
|
||||
|
||||
@@ -489,7 +489,7 @@ class VideoGenerationWorker:
|
||||
num_inference_steps=NUM_INFERENCE_STEPS,
|
||||
guidance_scale=1.0,
|
||||
seed=10,
|
||||
ltx2_image_crf=0.0,
|
||||
ltx2_image_crf=(33.0 if image_path and segment_idx == 1 else 0.0),
|
||||
image_path=image_path if segment_idx == 1 else None,
|
||||
return_continuation_state=False,
|
||||
)
|
||||
|
||||
@@ -29,6 +29,9 @@ test = [
|
||||
dreamverse-server = "dreamverse.server_entry:cli"
|
||||
dreamverse-mock-server = "dreamverse.mock_server:cli"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
include = ["dreamverse*"]
|
||||
|
||||
[tool.uv]
|
||||
package = false
|
||||
|
||||
|
||||
Executable
+67
@@ -0,0 +1,67 @@
|
||||
#!/usr/bin/env bash
|
||||
# launch-dreamverse.sh — launch dreamverse-server on a compute node.
|
||||
#
|
||||
# Usage (from repo root):
|
||||
# bash apps/dreamverse/scripts/launch-dreamverse.sh # GPUs 0-3, SP_SIZE=4
|
||||
# CUDA_VISIBLE_DEVICES=0 bash apps/dreamverse/scripts/launch-dreamverse.sh # single GPU, SP_SIZE=1
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
CONDA_PREFIX="$HOME/miniconda3/envs/dreamverse"
|
||||
CUDA_RT_DIR="$CONDA_PREFIX/lib/python3.11/site-packages/nvidia/cuda_runtime/lib"
|
||||
GXX="$CONDA_PREFIX/bin/aarch64-conda-linux-gnu-g++"
|
||||
|
||||
|
||||
export CUDA_HOME="$CONDA_PREFIX"
|
||||
export FASTVIDEO_ENABLE_STARTUP_WARMUP=true
|
||||
export FASTVIDEO_ENABLE_PROMPT_SAFETY=false
|
||||
export DREAMVERSE_MAX_AUTOTUNE=true
|
||||
export LTX2_USE_DISTILLED_SIGMAS=0
|
||||
export LTX2_VIDEO_CONDITIONING_NUM_FRAMES=1
|
||||
export AUDIO_CONDITIONING_NUM_FRAMES=41
|
||||
export DREAMVERSE_SESSION_TIMEOUT_SECONDS="${DREAMVERSE_SESSION_TIMEOUT_SECONDS:-1800}"
|
||||
# GB200 max-autotune warmup compiles can run for hours; keep the watchdog generous here
|
||||
export FASTVIDEO_STARTUP_WARMUP_TIMEOUT_SECONDS="${FASTVIDEO_STARTUP_WARMUP_TIMEOUT_SECONDS:-24000}"
|
||||
export CEREBRAS_API_KEY="${CEREBRAS_API_KEY:-}" # set this in your env or ~/.env
|
||||
export FASTVIDEO_PROMPT_CEREBRAS_MODEL="gpt-oss-120b"
|
||||
export TORCHINDUCTOR_CACHE_DIR="$HOME/.cache/torchinductor"
|
||||
export TRITON_CACHE_DIR="$HOME/.triton/cache"
|
||||
export TORCH_CUDA_ARCH_LIST="10.0a"
|
||||
|
||||
# Compiler env (needed for flashinfer JIT compilation at server startup)
|
||||
export CXX="$CONDA_PREFIX/compiler_compat/g++"
|
||||
export CC="$CONDA_PREFIX/compiler_compat/gcc"
|
||||
export CUDAHOSTCXX="$GXX"
|
||||
export NVCC_PREPEND_FLAGS="-ccbin $GXX -allow-unsupported-compiler"
|
||||
|
||||
# Link against libcudart.so.12 at JIT compile time; stubs for libcuda.so
|
||||
# cuda-compat has libcudart.so -> libcudart.so.12 (linker needs unversioned name)
|
||||
export LIBRARY_PATH="$CONDA_PREFIX/lib/cuda-compat:$CONDA_PREFIX/lib/stubs"
|
||||
|
||||
# Only libcudart.so.12 at runtime — prevents cuDNN from seeing .so.13
|
||||
export LD_LIBRARY_PATH="$CUDA_RT_DIR"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
|
||||
export FASTVIDEO_GPU_COUNT="${FASTVIDEO_GPU_COUNT:-all}"
|
||||
# Default SP size to the usable GPU count so single-GPU invocations work.
|
||||
# A numeric FASTVIDEO_GPU_COUNT caps the pool below the visible count, and an
|
||||
# SP size above the pool size fails GPUPool startup with "Not enough GPUs".
|
||||
IFS=',' read -ra _VISIBLE_GPUS <<< "$CUDA_VISIBLE_DEVICES"
|
||||
# Count only non-empty tokens, matching gpu_pool.get_available_gpus (e.g. ",0,1" is 2 GPUs).
|
||||
_USABLE_GPU_COUNT=0
|
||||
for _gpu in "${_VISIBLE_GPUS[@]}"; do
|
||||
[[ -n "${_gpu//[[:space:]]/}" ]] && _USABLE_GPU_COUNT=$((_USABLE_GPU_COUNT + 1))
|
||||
done
|
||||
if [[ "$FASTVIDEO_GPU_COUNT" =~ ^[0-9]+$ ]] && (( FASTVIDEO_GPU_COUNT < _USABLE_GPU_COUNT )); then
|
||||
_USABLE_GPU_COUNT="$FASTVIDEO_GPU_COUNT"
|
||||
fi
|
||||
export DREAMVERSE_SP_SIZE="${DREAMVERSE_SP_SIZE:-$_USABLE_GPU_COUNT}"
|
||||
PORT="${DREAMVERSE_PORT:-8009}"
|
||||
|
||||
FFMPEG_ENV="$(dirname "$0")/ffmpeg-env.sh"
|
||||
# shellcheck source=ffmpeg-env.sh
|
||||
[[ -f "$FFMPEG_ENV" ]] && source "$FFMPEG_ENV"
|
||||
|
||||
echo "==> Launching dreamverse-server on GPU $CUDA_VISIBLE_DEVICES port $PORT"
|
||||
CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" \
|
||||
"$CONDA_PREFIX/bin/dreamverse-server" --host 0.0.0.0 --port "$PORT"
|
||||
Executable
+30
@@ -0,0 +1,30 @@
|
||||
#!/usr/bin/env bash
|
||||
# launch-frontend.sh — start the Dreamverse Next.js dev server and ngrok tunnel.
|
||||
#
|
||||
# Usage (from repo root):
|
||||
# bash apps/dreamverse/scripts/launch-frontend.sh
|
||||
#
|
||||
# Override backend or ngrok URL via env:
|
||||
# BACKEND_HOST=1.2.3.4 bash apps/dreamverse/scripts/launch-frontend.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
CONDA_PREFIX="$HOME/miniconda3/envs/dreamverse"
|
||||
BACKEND_HOST="${BACKEND_HOST:-10.244.18.228}"
|
||||
BACKEND_PORT="${BACKEND_PORT:-8009}"
|
||||
NGROK_URL="${NGROK_URL:-ltx23.ngrok.app}"
|
||||
WEB_DIR="$(git rev-parse --show-toplevel)/apps/dreamverse/web"
|
||||
|
||||
cleanup() {
|
||||
echo "==> Shutting down..."
|
||||
kill "$FRONTEND_PID" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
echo "==> Starting frontend (backend: $BACKEND_HOST:$BACKEND_PORT)"
|
||||
BACKEND_HOST="$BACKEND_HOST" BACKEND_PORT="$BACKEND_PORT" \
|
||||
npm run --prefix "$WEB_DIR" dev &
|
||||
FRONTEND_PID=$!
|
||||
|
||||
echo "==> Starting ngrok tunnel -> $NGROK_URL"
|
||||
"$CONDA_PREFIX/bin/ngrok" http --url="$NGROK_URL" 5299
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
#!/usr/bin/env bash
|
||||
# setup-dreamverse-env.sh — create and configure the dreamverse conda env
|
||||
# from scratch on this aarch64 NFS Slurm cluster.
|
||||
#
|
||||
# Run from the login node (from the repo root):
|
||||
# bash apps/dreamverse/scripts/setup-dreamverse-env.sh
|
||||
#
|
||||
# After this script completes, use launch-dreamverse.sh on a compute node.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(git rev-parse --show-toplevel)"
|
||||
ENV_NAME="dreamverse"
|
||||
LOCAL_DIR="/mnt/local/hal-kevin" # cache/pkgs — keep on local disk
|
||||
CONDA_PREFIX="$HOME/miniconda3/envs/$ENV_NAME" # env — on shared NFS so it survives node changes
|
||||
|
||||
echo "==> Removing existing env if present"
|
||||
conda env remove -p "$CONDA_PREFIX" -y 2>/dev/null || true
|
||||
rm -rf "$CONDA_PREFIX" 2>/dev/null || true
|
||||
|
||||
echo "==> Creating conda env at $CONDA_PREFIX"
|
||||
CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda create -p "$CONDA_PREFIX" python=3.11 -y
|
||||
|
||||
GXX="$CONDA_PREFIX/bin/aarch64-conda-linux-gnu-g++"
|
||||
GCC="$CONDA_PREFIX/bin/aarch64-conda-linux-gnu-gcc"
|
||||
|
||||
echo "==> Installing compiler"
|
||||
CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda install -p "$CONDA_PREFIX" gxx_linux-aarch64 -y
|
||||
|
||||
echo "==> Installing CUDA toolkit (nvcc + headers)"
|
||||
CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda install -p "$CONDA_PREFIX" -c nvidia cuda-toolkit -y
|
||||
|
||||
echo "==> Hiding conflicting libcudart.so.13 immediately"
|
||||
mkdir -p "$CONDA_PREFIX/lib/hidden"
|
||||
mv "$CONDA_PREFIX"/lib/libcudart.so* "$CONDA_PREFIX/lib/hidden/" 2>/dev/null || true
|
||||
|
||||
echo "==> Fixing compiler_compat symlinks"
|
||||
mkdir -p "$CONDA_PREFIX/compiler_compat"
|
||||
ln -sf "$GXX" "$CONDA_PREFIX/compiler_compat/g++"
|
||||
ln -sf "$GCC" "$CONDA_PREFIX/compiler_compat/gcc"
|
||||
|
||||
echo "==> Symlinking CUDA headers to standard location"
|
||||
for f in "$CONDA_PREFIX/targets/sbsa-linux/include/"*; do
|
||||
ln -sf "$f" "$CONDA_PREFIX/include/$(basename "$f")" 2>/dev/null || true
|
||||
done
|
||||
|
||||
echo "==> Installing ffmpeg (native build with x264 + NVENC)"
|
||||
CUDA_PREFIX="$CONDA_PREFIX" bash "$REPO_ROOT/apps/dreamverse/scripts/install_native_ffmpeg.sh"
|
||||
|
||||
echo "==> Installing pip and uv"
|
||||
CONDA_PKGS_DIRS="$LOCAL_DIR/conda/pkgs" conda install -p "$CONDA_PREFIX" pip -y
|
||||
"$CONDA_PREFIX/bin/pip" install uv
|
||||
|
||||
echo "==> Setting compiler env vars"
|
||||
export UV_CACHE_DIR="$LOCAL_DIR/cache"
|
||||
export UV_LINK_MODE=copy
|
||||
export CXX="$CONDA_PREFIX/compiler_compat/g++"
|
||||
export CC="$CONDA_PREFIX/compiler_compat/gcc"
|
||||
export CUDAHOSTCXX="$GXX"
|
||||
export NVCC_PREPEND_FLAGS="-ccbin $GXX -allow-unsupported-compiler"
|
||||
export CUDA_HOME="$CONDA_PREFIX"
|
||||
# Only build for GB200 (sm_100a); CUDA 13 dropped support for older archs
|
||||
export TORCH_CUDA_ARCH_LIST="10.0a"
|
||||
|
||||
echo "==> Installing torch with CUDA 12.8"
|
||||
UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install torch==2.11.0 torchvision \
|
||||
--index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
echo "==> Hiding any newly introduced libcudart.so.13"
|
||||
mv "$CONDA_PREFIX"/lib/libcudart.so* "$CONDA_PREFIX/lib/hidden/" 2>/dev/null || true
|
||||
|
||||
# Set paths now that torch (and its nvidia packages) are installed
|
||||
CUDA_RT_DIR="$CONDA_PREFIX/lib/python3.11/site-packages/nvidia/cuda_runtime/lib"
|
||||
CUDA_RT_SO="$(ls "$CUDA_RT_DIR"/libcudart.so.* 2>/dev/null | head -1)"
|
||||
|
||||
# The pip nvidia package only has libcudart.so.12 (versioned), not libcudart.so.
|
||||
# The linker needs the unversioned name to satisfy -lcudart. Create a compat dir.
|
||||
mkdir -p "$CONDA_PREFIX/lib/cuda-compat"
|
||||
ln -sf "$CUDA_RT_SO" "$CONDA_PREFIX/lib/cuda-compat/libcudart.so"
|
||||
|
||||
export LIBRARY_PATH="$CONDA_PREFIX/lib/cuda-compat:$CONDA_PREFIX/lib/stubs"
|
||||
export CMAKE_ARGS="-DCUDA_CUDART_LIBRARY=$CUDA_RT_SO -DCUDA_INCLUDE_DIRS=$CONDA_PREFIX/targets/sbsa-linux/include"
|
||||
|
||||
echo "==> Installing build tools"
|
||||
"$CONDA_PREFIX/bin/pip" install scikit-build-core cmake ninja
|
||||
|
||||
echo "==> Initializing git submodules"
|
||||
cd "$REPO_ROOT"
|
||||
git submodule update --init fastvideo-kernel/include/cutlass fastvideo-kernel/include/tk
|
||||
|
||||
echo "==> Building fastvideo-kernel from local source"
|
||||
UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install \
|
||||
-e "./fastvideo-kernel" --no-build-isolation
|
||||
|
||||
echo "==> Installing fastvideo + dreamverse extras"
|
||||
UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install \
|
||||
-e ".[dreamverse]" --no-build-isolation
|
||||
|
||||
echo "==> Installing flashinfer-python (pinned, must be last)"
|
||||
UV_CACHE_DIR="$LOCAL_DIR/cache" "$CONDA_PREFIX/bin/uv" pip install \
|
||||
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.11.post3/flashinfer_python-0.6.11.post3-py3-none-any.whl
|
||||
|
||||
echo ""
|
||||
echo "Done. On a compute node run (GPUs 0-3 by default; set CUDA_VISIBLE_DEVICES to restrict):"
|
||||
echo " bash apps/dreamverse/scripts/launch-dreamverse.sh"
|
||||
@@ -110,7 +110,7 @@ default_request:
|
||||
fps: 24 # internal: gpu_pool.py:85 TARGET_FPS
|
||||
|
||||
streaming:
|
||||
# internal: config.py:33 SESSION_TIMEOUT_SECONDS = 300
|
||||
# internal: config.py SESSION_TIMEOUT_SECONDS (env DREAMVERSE_SESSION_TIMEOUT_SECONDS, default 300)
|
||||
session_timeout_seconds: 300
|
||||
# internal: config.py:282-284 GENERATION_SEGMENT_CAP default 6
|
||||
generation_segment_cap: 6
|
||||
|
||||
@@ -9,7 +9,7 @@ import SessionTimeoutModal from "@/components/SessionTimeoutModal";
|
||||
import Sidebar from "@/components/Sidebar";
|
||||
import Header from "@/components/Header";
|
||||
import VideoPlayer from "@/components/VideoPlayer";
|
||||
import Workspace from "@/components/Workspace";
|
||||
import Workspace, { SceneHistoryList } from "@/components/Workspace";
|
||||
import { saveProject, saveProjectMetadata, listProjects, loadProjectClips, deleteProject, pruneOldProjects, type StoredProject, type StoredClip } from "@/lib/projectStorage";
|
||||
import { isInfrastructureError } from "@/lib/ws/reducer";
|
||||
import { useStore } from "@/hooks/useStore";
|
||||
@@ -31,7 +31,7 @@ import { applyNormalizedSocketEvent } from "@/lib/ws/reducer";
|
||||
import { createPromptWindowStore } from "@/stores/promptWindow";
|
||||
import { createRewriteStore } from "@/stores/rewrite";
|
||||
import { createSessionStore } from "@/stores/session";
|
||||
import { createStreamStore } from "@/stores/stream";
|
||||
import { createStreamStore, USER_PROMPT_SOURCES } from "@/stores/stream";
|
||||
import { createUiStore } from "@/stores/ui";
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
||||
@@ -80,7 +80,7 @@ function yieldToEventLoop(): Promise<void> {
|
||||
|
||||
const HERO_WAVE_LIGHT = ["#2A4A98", "#4878E5", "#6FA0F2", "#B0BCC8", "#E8D99E", "#D8C844", "#C2A620"];
|
||||
const HERO_WAVE_DARK = ["#143468", "#1E58B8", "#3892F0", "#80B8E8", "#B8D0EA", "#E2D498", "#DABB50"];
|
||||
const HERO_TEXT = "Direct scenes in seconds";
|
||||
const HERO_TEXT = "Direct scenes in seconds with";
|
||||
|
||||
function HeroTagline() {
|
||||
const ref = useRef<HTMLHeadingElement>(null);
|
||||
@@ -166,6 +166,8 @@ function HeroTagline() {
|
||||
</span>
|
||||
</Fragment>
|
||||
))}
|
||||
<span data-char className="transition-[color,filter] duration-150">{" "}</span>
|
||||
<img src="/logo.svg" alt="FastVideo" className="inline-block h-[1.1em] w-auto align-middle" />
|
||||
</h1>
|
||||
);
|
||||
}
|
||||
@@ -237,6 +239,7 @@ export default function Page() {
|
||||
enhancementEnabled,
|
||||
promptExtensionError,
|
||||
autoExtensionEnabled,
|
||||
manualContinuationMode,
|
||||
autoExtensionTimeoutHint,
|
||||
loopGenerationEnabled,
|
||||
generationPaused,
|
||||
@@ -249,6 +252,8 @@ export default function Page() {
|
||||
livePromptRewriteMode,
|
||||
sessionExpired,
|
||||
projectResetPending,
|
||||
waitingForSegmentPrompt,
|
||||
generatingNextScene,
|
||||
} = sessionState;
|
||||
|
||||
const {
|
||||
@@ -323,6 +328,11 @@ export default function Page() {
|
||||
const [ttffValueMs, setTtffValueMs] = useState<number | null>(null);
|
||||
const ttffIntervalRef = useRef<ReturnType<typeof setInterval> | null>(null);
|
||||
const pendingInitialPromptRef = useRef("");
|
||||
// Prompt id the opening scene is recorded under; sent as initial_rollout_prompt_id
|
||||
// so the backend's pre-seeded opening PromptSubmission emits status updates
|
||||
// (prompt_enhancing/prompt_ready/prompt_fallback_used) against the same id.
|
||||
const pendingInitialPromptIdRef = useRef("");
|
||||
const [initialImageDataUrl, setInitialImageDataUrl] = useState("");
|
||||
const lastArchivedReplayKeyRef = useRef("");
|
||||
const [sidebarOpen, setSidebarOpen] = useState(false);
|
||||
const [currentThumbnail, setCurrentThumbnail] = useState<string | null>(null);
|
||||
@@ -422,6 +432,11 @@ export default function Page() {
|
||||
if (String(e?.source || "") === "user_rewrite" && typeof e?.text === "string" && e.text.trim()) {
|
||||
return e.text.trim();
|
||||
}
|
||||
// Steering opening: the backend overwrites text/source with the enhanced
|
||||
// prompt once ready, so fall back to the stable rawText record.
|
||||
if (e?.steeringUserPrompt && typeof e?.rawText === "string" && e.rawText.trim()) {
|
||||
return e.rawText.trim();
|
||||
}
|
||||
}
|
||||
return "Untitled project";
|
||||
}, [selectedPreset, promptEvents]);
|
||||
@@ -432,11 +447,39 @@ export default function Page() {
|
||||
const canDownloadVideo = useMemo(() => {
|
||||
const currentActiveClip = activeClip as Record<string, any> | null;
|
||||
if (currentActiveClip?.blob instanceof Blob) return true;
|
||||
return (completedClips as Record<string, any>[]).some((clip) => clip?.blob instanceof Blob);
|
||||
}, [activeClip, completedClips]);
|
||||
if ((completedClips as Record<string, any>[]).some((clip) => clip?.blob instanceof Blob)) return true;
|
||||
// Steering only: once playback has started the live AV pipeline holds playable segments, so
|
||||
// the user can download the in-progress video at any time (handleDownloadVideo remuxes live
|
||||
// segments). Auto mode keeps its original blob-gated behavior.
|
||||
return Boolean(manualContinuationMode) && Boolean(avPlaybackStarted);
|
||||
}, [activeClip, completedClips, avPlaybackStarted, manualContinuationMode]);
|
||||
|
||||
// Steering mode scene list (oldest first). Primary source is the user's own words, captured
|
||||
// stably at submit time as `rawText` (the backend later overwrites text/source with the
|
||||
// enhanced prompt, so we never read those). A segment with no user prompt — e.g. a preset's
|
||||
// opening scene — falls back to promptHistory (the actual prompt that drove that segment).
|
||||
const steeringScenes = useMemo(() => {
|
||||
if (!manualContinuationMode) return [] as Record<string, any>[];
|
||||
const userScenes = (promptEvents as Record<string, any>[])
|
||||
.filter((e) => e?.steeringUserPrompt && !e?.steeringFailed && typeof e?.rawText === "string" && e.rawText.trim())
|
||||
.slice()
|
||||
.reverse() // oldest -> newest
|
||||
.map((e) => ({ id: e.promptId, prompt: e.rawText as string }));
|
||||
const scenes: Record<string, any>[] = [];
|
||||
// Preset opening segments: curated seeds with no user prompt of their own.
|
||||
const curatedHists = (promptHistory as Record<string, any>[])
|
||||
.slice()
|
||||
.reverse() // oldest first
|
||||
.filter((h) => !USER_PROMPT_SOURCES.has(String(h?.source || "")) && typeof h?.prompt === "string" && (h.prompt as string).trim());
|
||||
scenes.push(...curatedHists.map((h) => ({ id: h.id || "scene_open", prompt: h.prompt })));
|
||||
scenes.push(...userScenes);
|
||||
return scenes;
|
||||
}, [manualContinuationMode, promptEvents, promptHistory]);
|
||||
|
||||
const hasEdits = useMemo(
|
||||
() => Boolean(sessionStarted) && (promptEvents as Record<string, any>[]).some((e) => typeof e?.text === "string" && e.text.trim() && String(e?.source || "").trim() === "user_rewrite"),
|
||||
() => Boolean(sessionStarted) && (
|
||||
(promptEvents as Record<string, any>[]).some((e) => typeof e?.text === "string" && e.text.trim() && String(e?.source || "").trim() === "user_rewrite")
|
||||
),
|
||||
[sessionStarted, promptEvents],
|
||||
);
|
||||
|
||||
@@ -1421,7 +1464,8 @@ export default function Page() {
|
||||
if (!prompt) return;
|
||||
lastSubmitTimeRef.current = now;
|
||||
const rCAR = !uiStore.get().devtoolsMode && !uiStore.get().demoMode;
|
||||
if (rCAR || (sessionStore.get().livePromptRewriteMode && !uiStore.get().demoMode)) {
|
||||
const inManualMode = sessionStore.get().manualContinuationMode;
|
||||
if (!inManualMode && (rCAR || (sessionStore.get().livePromptRewriteMode && !uiStore.get().demoMode))) {
|
||||
if (rewriteStore.get().rewritingSeedPrompts) return;
|
||||
const rewriteSourcePromptWindowPrompts = getActivePromptWindowPrompts();
|
||||
const nextPendingClip = {
|
||||
@@ -1462,6 +1506,11 @@ export default function Page() {
|
||||
status: "submitted",
|
||||
source: "user_raw",
|
||||
text: prompt,
|
||||
// Stable record of the user's own words for the steering scene list. The backend
|
||||
// later overwrites `text`/`source` with the enhanced prompt via prompt/ready, but
|
||||
// these two fields are never touched by trackPromptEvent.
|
||||
steeringUserPrompt: true,
|
||||
rawText: prompt,
|
||||
});
|
||||
ws.send(
|
||||
JSON.stringify({
|
||||
@@ -1475,7 +1524,14 @@ export default function Page() {
|
||||
activeClipId: shouldUseArchivedPlaybackFallback() ? streamStore.get().activeClipId : "",
|
||||
activePlaybackStartTime: shouldUseArchivedPlaybackFallback() ? streamStore.get().activePlaybackStartTime : 0,
|
||||
});
|
||||
sessionStore.patch({ livePromptDraft: "" });
|
||||
sessionStore.patch({
|
||||
livePromptDraft: "",
|
||||
waitingForSegmentPrompt: false,
|
||||
sessionNotice: "",
|
||||
// Light the "Generating next scene" overlay immediately on a real submit;
|
||||
// stream/media_init (or a fallback/error) clears it.
|
||||
...(inManualMode ? { generatingNextScene: true } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
function setLivePromptRewriteMode(enabled: boolean) {
|
||||
@@ -1550,6 +1606,14 @@ export default function Page() {
|
||||
);
|
||||
}
|
||||
|
||||
// Steering (manual continuation) vs the automatic 6-segment rollout — a pre-session
|
||||
// preference honored when the session starts.
|
||||
function handleManualContinuationToggle(event: any) {
|
||||
sessionStore.patch({
|
||||
manualContinuationMode: Boolean(event.currentTarget.checked),
|
||||
});
|
||||
}
|
||||
|
||||
function handleLoopGenerationToggle(event: any) {
|
||||
sessionStore.patch({
|
||||
loopGenerationEnabled: Boolean(event.currentTarget.checked),
|
||||
@@ -1725,6 +1789,9 @@ export default function Page() {
|
||||
sessionNotice: preserveSessionNotice ? sessionStore.get().sessionNotice : "",
|
||||
sessionExpired: preserveSessionNotice ? sessionStore.get().sessionExpired : false,
|
||||
projectResetPending: false,
|
||||
manualContinuationMode: true,
|
||||
waitingForSegmentPrompt: false,
|
||||
generatingNextScene: false,
|
||||
});
|
||||
rewriteStore.resetSessionState();
|
||||
streamStore.resetSessionState();
|
||||
@@ -1737,6 +1804,7 @@ export default function Page() {
|
||||
function resetToProjectLobbyState() {
|
||||
setVideoMuted(true);
|
||||
pendingInitialPromptRef.current = "";
|
||||
pendingInitialPromptIdRef.current = "";
|
||||
sessionStore.patch({
|
||||
sessionStarted: false,
|
||||
livePromptDraft: "",
|
||||
@@ -1750,6 +1818,9 @@ export default function Page() {
|
||||
sessionNotice: "",
|
||||
sessionExpired: false,
|
||||
projectResetPending: false,
|
||||
manualContinuationMode: true,
|
||||
waitingForSegmentPrompt: false,
|
||||
generatingNextScene: false,
|
||||
});
|
||||
rewriteStore.resetSessionState();
|
||||
streamStore.resetSessionState();
|
||||
@@ -1760,7 +1831,14 @@ export default function Page() {
|
||||
}
|
||||
|
||||
function buildProjectInitPayload(type: "session_init_v2" | "project_init_v1") {
|
||||
const segmentPrompts = getSessionInitPrompts();
|
||||
const manualMode = Boolean(sessionStore.get().manualContinuationMode);
|
||||
let segmentPrompts = getSessionInitPrompts();
|
||||
// Steering mode: seed the first 2 segments from the preset so there's no
|
||||
// gap between segment 1 and 2; the user drives every subsequent segment.
|
||||
// Force auto/loop off so the backend waits after the seeded prompts run out.
|
||||
if (manualMode) {
|
||||
segmentPrompts = segmentPrompts.slice(0, 2);
|
||||
}
|
||||
setSeedPrompts(segmentPrompts);
|
||||
return {
|
||||
type,
|
||||
@@ -1768,11 +1846,17 @@ export default function Page() {
|
||||
preset_label: getInitialPresetLabel(),
|
||||
curated_prompts: segmentPrompts,
|
||||
initial_rollout_prompt: normalizeInitialPrompt(pendingInitialPromptRef.current),
|
||||
initial_image: null,
|
||||
// Ties the backend's pre-seeded opening PromptSubmission to the prompt event
|
||||
// recorded in beginProjectLocally so its status updates land on that record.
|
||||
initial_rollout_prompt_id: pendingInitialPromptIdRef.current,
|
||||
initial_image: initialImageDataUrl
|
||||
? { data_url: initialImageDataUrl, mime_type: initialImageDataUrl.split(";")[0].split(":")[1] || "image/png", name: "upload.png" }
|
||||
: null,
|
||||
single_clip_mode: false,
|
||||
enhancement_enabled: sessionStore.get().enhancementEnabled,
|
||||
auto_extension_enabled: sessionStore.get().autoExtensionEnabled,
|
||||
loop_generation_enabled: sessionStore.get().loopGenerationEnabled,
|
||||
auto_extension_enabled: manualMode ? false : sessionStore.get().autoExtensionEnabled,
|
||||
loop_generation_enabled: manualMode ? false : sessionStore.get().loopGenerationEnabled,
|
||||
manual_continuation_mode: manualMode,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -1951,7 +2035,11 @@ export default function Page() {
|
||||
setCurrentThumbnail(null);
|
||||
const initialPrompt = normalizeInitialPrompt(sessionStore.get().livePromptDraft as string);
|
||||
pendingInitialPromptRef.current = initialPrompt;
|
||||
setInitialImageDataUrl("");
|
||||
const rCAR = !uiStore.get().devtoolsMode && !uiStore.get().demoMode;
|
||||
// The "Steering mode" toggle is authoritative: checked = manual per-segment steering,
|
||||
// unchecked = automatic 6-segment rollout (default).
|
||||
const nextManualContinuationMode = Boolean(sessionStore.get().manualContinuationMode);
|
||||
sessionStore.patch({
|
||||
sessionNotice: "",
|
||||
sessionExpired: false,
|
||||
@@ -1966,6 +2054,9 @@ export default function Page() {
|
||||
autoExtensionTimeoutHint: "",
|
||||
generationPaused: false,
|
||||
projectResetPending: false,
|
||||
manualContinuationMode: nextManualContinuationMode,
|
||||
waitingForSegmentPrompt: false,
|
||||
generatingNextScene: false,
|
||||
});
|
||||
resetPlaybackState();
|
||||
streamStore.patch({
|
||||
@@ -1979,12 +2070,15 @@ export default function Page() {
|
||||
selectedHistoryId: "",
|
||||
});
|
||||
rewriteStore.resetSessionState();
|
||||
pendingInitialPromptIdRef.current = initialPrompt ? makePromptId() : "";
|
||||
if (initialPrompt) {
|
||||
addPromptEvent({
|
||||
promptId: makePromptId(),
|
||||
promptId: pendingInitialPromptIdRef.current,
|
||||
status: "rewrite_requested",
|
||||
source: "user_rewrite",
|
||||
text: initialPrompt,
|
||||
// In steering mode the typed opening is the user's Scene 1 — record it stably.
|
||||
...(nextManualContinuationMode ? { steeringUserPrompt: true, rawText: initialPrompt } : {}),
|
||||
});
|
||||
}
|
||||
setSeedPrompts(getSessionInitPrompts());
|
||||
@@ -2500,6 +2594,7 @@ export default function Page() {
|
||||
selectedPresetId={selectedPresetId as string}
|
||||
enhancementEnabled={enhancementEnabled as boolean}
|
||||
autoExtensionEnabled={autoExtensionEnabled as boolean}
|
||||
manualContinuationEnabled={manualContinuationMode as boolean}
|
||||
loopGenerationEnabled={loopGenerationEnabled as boolean}
|
||||
canJoinSession={canJoinSession as boolean}
|
||||
canSubmitContinuation={canSubmitContinuation}
|
||||
@@ -2512,6 +2607,7 @@ export default function Page() {
|
||||
onEnhancementToggle={handleEnhancementToggle}
|
||||
onCuratedPromptLimitChange={handleCuratedPromptLimitChange}
|
||||
onAutoExtensionToggle={handleAutoExtensionToggle}
|
||||
onManualContinuationToggle={handleManualContinuationToggle}
|
||||
onLoopToggle={handleLoopGenerationToggle}
|
||||
onJoin={joinSession}
|
||||
onLeave={leaveSession}
|
||||
@@ -2641,6 +2737,7 @@ export default function Page() {
|
||||
/>
|
||||
<Header timeLeft={headerTimeLeft} formatTime={formatTime} onToggleSidebar={() => setSidebarOpen((prev) => !prev)} />
|
||||
|
||||
<div className={cn("flex flex-1 min-h-0 flex-col", sessionStarted && "pb-16 sm:pb-28")}>
|
||||
<div className="relative flex flex-1 min-h-0 flex-col justify-center px-4 pb-2 sm:px-6 sm:pb-12">
|
||||
{isViewingMode && (
|
||||
<>
|
||||
@@ -2724,6 +2821,8 @@ export default function Page() {
|
||||
showLivePlayback={showLivePlayback}
|
||||
defaultMuted={videoMuted}
|
||||
canDownload={canDownloadVideo}
|
||||
waitingForSegmentPrompt={waitingForSegmentPrompt as boolean}
|
||||
generatingNextScene={generatingNextScene as boolean}
|
||||
onPlaying={markFirstFrameRendered}
|
||||
onDownload={handleDownloadVideo}
|
||||
/>
|
||||
@@ -2734,6 +2833,7 @@ export default function Page() {
|
||||
<section className={cn("mx-auto w-full max-w-2xl", hasEdits && "flex-1 min-h-0 overflow-y-auto")}>
|
||||
<Workspace
|
||||
promptEvents={promptEvents as any[]}
|
||||
manualMode={manualContinuationMode as boolean}
|
||||
currentThumbnail={currentThumbnail}
|
||||
originalLabel={pendingInitialPromptRef.current || (selectedPreset as Record<string, any>)?.label || ""}
|
||||
sessionStarted={sessionStarted as boolean}
|
||||
@@ -2780,11 +2880,17 @@ export default function Page() {
|
||||
isGenerating={loadingAnimation as boolean}
|
||||
storyPresets={storyPresets as any[]}
|
||||
continuationDraft={livePromptDraft as string}
|
||||
manualContinuationEnabled={manualContinuationMode as boolean}
|
||||
onModeChange={(manual) => sessionStore.patch({ manualContinuationMode: manual })}
|
||||
initialImageDataUrl={initialImageDataUrl}
|
||||
onImageUpload={(dataUrl) => setInitialImageDataUrl(dataUrl)}
|
||||
onImageClear={() => setInitialImageDataUrl("")}
|
||||
canJoinSession={canStartSession}
|
||||
canSubmitContinuation={canSubmitContinuation}
|
||||
sessionExpired={sessionExpired as boolean}
|
||||
sessionNotice={sessionNotice as string}
|
||||
projectResetPending={projectResetPending as boolean}
|
||||
waitingForSegmentPrompt={waitingForSegmentPrompt as boolean}
|
||||
onPresetGenerate={handlePresetGenerate}
|
||||
onContinuationInput={handleLivePromptInput}
|
||||
onContinuationKeydown={handleLivePromptKeydown}
|
||||
@@ -2798,6 +2904,12 @@ export default function Page() {
|
||||
</motion.div>
|
||||
</div>
|
||||
</div>
|
||||
{manualContinuationMode && (
|
||||
<div className="px-4 sm:px-6">
|
||||
<SceneHistoryList sceneHistory={steeringScenes as any[]} />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</main>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2,13 +2,16 @@
|
||||
|
||||
import React, { useRef, useState, useCallback, useEffect } from "react";
|
||||
import Image from "next/image";
|
||||
import { Film, ArrowUp, X, Loader2, ArrowLeft } from "lucide-react";
|
||||
import { Film, ArrowUp, X, Loader2, ArrowLeft, ImagePlus } from "lucide-react";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import LeaveSessionModal, { shouldShowLeaveWarning } from "@/components/LeaveSessionModal";
|
||||
import SpeechToTextButton from "@/components/SpeechToTextButton";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
const PROMPT_MAX_LENGTH = 500;
|
||||
// Must match backend session_init_image.py: MAX_SESSION_INIT_IMAGE_BYTES / SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES.
|
||||
const IMAGE_MAX_BYTES = 15 * 1024 * 1024;
|
||||
const IMAGE_ALLOWED_TYPES = ["image/png", "image/jpeg", "image/webp"];
|
||||
|
||||
interface Props {
|
||||
sessionStarted?: boolean;
|
||||
@@ -22,6 +25,12 @@ interface Props {
|
||||
sessionNotice?: string;
|
||||
projectResetPending?: boolean;
|
||||
viewingReadOnly?: boolean;
|
||||
waitingForSegmentPrompt?: boolean;
|
||||
manualContinuationEnabled?: boolean;
|
||||
onModeChange?: (manual: boolean) => void;
|
||||
initialImageDataUrl?: string;
|
||||
onImageUpload?: (dataUrl: string, mimeType: string, name: string) => void;
|
||||
onImageClear?: () => void;
|
||||
onPresetGenerate?: (presetId: string) => void;
|
||||
onContinuationInput?: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
|
||||
onContinuationKeydown?: (e: React.KeyboardEvent<HTMLTextAreaElement>) => void;
|
||||
@@ -46,6 +55,12 @@ export default function ChatBar({
|
||||
sessionNotice = "",
|
||||
projectResetPending = false,
|
||||
viewingReadOnly = false,
|
||||
waitingForSegmentPrompt = false,
|
||||
manualContinuationEnabled = false,
|
||||
onModeChange = () => {},
|
||||
initialImageDataUrl = "",
|
||||
onImageUpload = () => {},
|
||||
onImageClear = () => {},
|
||||
onPresetGenerate = () => {},
|
||||
onContinuationInput = () => {},
|
||||
onContinuationKeydown = () => {},
|
||||
@@ -59,15 +74,51 @@ export default function ChatBar({
|
||||
}: Props) {
|
||||
const [sttBusy, setSttBusy] = useState(false);
|
||||
const [leaveModalOpen, setLeaveModalOpen] = useState(false);
|
||||
const [imageError, setImageError] = useState("");
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
|
||||
const processImageFile = useCallback((file: File) => {
|
||||
if (!IMAGE_ALLOWED_TYPES.includes(file.type)) {
|
||||
setImageError("Unsupported image type. Use a PNG, JPEG, or WebP image.");
|
||||
return;
|
||||
}
|
||||
if (file.size > IMAGE_MAX_BYTES) {
|
||||
setImageError("Image is too large. The maximum size is 15MB.");
|
||||
return;
|
||||
}
|
||||
const reader = new FileReader();
|
||||
reader.onload = (e) => {
|
||||
const dataUrl = e.target?.result as string;
|
||||
if (dataUrl) {
|
||||
setImageError("");
|
||||
onImageUpload(dataUrl, file.type, file.name);
|
||||
}
|
||||
};
|
||||
reader.onerror = () => {
|
||||
setImageError("Could not read the image file. Please try again.");
|
||||
};
|
||||
reader.readAsDataURL(file);
|
||||
}, [onImageUpload]);
|
||||
|
||||
const handleImagePaste = useCallback((e: React.ClipboardEvent) => {
|
||||
if (sessionStarted) return;
|
||||
const items = Array.from(e.clipboardData?.items ?? []);
|
||||
const imageItem = items.find((item) => item.type.startsWith("image/"));
|
||||
if (!imageItem) return;
|
||||
const file = imageItem.getAsFile();
|
||||
if (file) processImageFile(file);
|
||||
}, [sessionStarted, processImageFile]);
|
||||
const showSpinner = isGenerating || rewritingSeedPrompts;
|
||||
const isBusy = isGenerating || rewritingSeedPrompts || projectResetPending;
|
||||
const messagePlaceholder = projectResetPending
|
||||
? "Starting new project\u2026"
|
||||
: isBusy
|
||||
? "Generating video\u2026"
|
||||
: !sessionStarted
|
||||
? "What video are you imagining?"
|
||||
: "What do you want to edit?";
|
||||
: waitingForSegmentPrompt
|
||||
? "Describe the next scene\u2026"
|
||||
: !sessionStarted
|
||||
? "What video are you imagining?"
|
||||
: "What do you want to edit?";
|
||||
const actionLabel = !sessionStarted ? "Generate" : "Rewrite rollout";
|
||||
|
||||
const inputRef = useRef<HTMLTextAreaElement>(null);
|
||||
@@ -267,9 +318,9 @@ export default function ChatBar({
|
||||
<Button onClick={onStartNewProject} size="sm" className="rounded-full px-5">
|
||||
New Project
|
||||
</Button>
|
||||
<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
|
||||
<a href="https://haoailab.com/blogs/dreamverse/" target="_blank" rel="noopener noreferrer">
|
||||
<Button variant="outline" size="sm" className="rounded-full px-5">
|
||||
Join Waitlist
|
||||
Blog
|
||||
</Button>
|
||||
</a>
|
||||
</div>
|
||||
@@ -346,6 +397,31 @@ export default function ChatBar({
|
||||
</div>
|
||||
)}
|
||||
|
||||
|
||||
{!sessionStarted && imageError && (
|
||||
<div className="rounded-xl border border-rose-500/20 bg-rose-500/10 px-4 py-2.5 text-center text-xs text-rose-700 dark:text-rose-300">
|
||||
{imageError}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{!sessionStarted && initialImageDataUrl && (
|
||||
<div className="flex items-center gap-2 rounded-2xl border border-input bg-card/65 px-3 py-2">
|
||||
<img src={initialImageDataUrl} alt="Initial frame" className="h-12 w-12 rounded-lg object-cover" />
|
||||
<span className="flex-1 truncate text-xs text-muted-foreground">Starting image set</span>
|
||||
<button type="button" onClick={() => { setImageError(""); onImageClear(); }} className="text-muted-foreground hover:text-foreground transition-colors">
|
||||
<X className="size-4" />
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<input
|
||||
ref={fileInputRef}
|
||||
type="file"
|
||||
accept={IMAGE_ALLOWED_TYPES.join(",")}
|
||||
className="hidden"
|
||||
onChange={(e) => { const f = e.target.files?.[0]; if (f) processImageFile(f); e.target.value = ""; }}
|
||||
/>
|
||||
|
||||
<div
|
||||
className={cn(
|
||||
"flex min-w-0 items-center gap-1.5 rounded-4xl border py-2.5 pl-5 pr-2.5 shadow-md backdrop-blur-sm transition-all duration-200",
|
||||
@@ -359,6 +435,7 @@ export default function ChatBar({
|
||||
value={continuationDraft}
|
||||
onChange={onContinuationInput}
|
||||
onKeyDown={handleKeyDown}
|
||||
onPaste={handleImagePaste}
|
||||
placeholder={sttBusy ? "Listening\u2026" : messagePlaceholder}
|
||||
maxLength={PROMPT_MAX_LENGTH}
|
||||
disabled={isBusy || sttBusy}
|
||||
@@ -369,6 +446,19 @@ export default function ChatBar({
|
||||
)}
|
||||
/>
|
||||
{onSpeechTranscript && <SpeechToTextButton disabled={isBusy} onTranscript={onSpeechTranscript} onInterimChange={onSpeechInterimChange} onBusyChange={setSttBusy} />}
|
||||
{!sessionStarted && (
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon-sm"
|
||||
title="Add starting image"
|
||||
onClick={() => fileInputRef.current?.click()}
|
||||
disabled={isBusy}
|
||||
className="shrink-0 rounded-full text-muted-foreground hover:text-foreground"
|
||||
>
|
||||
<ImagePlus className="size-4" />
|
||||
</Button>
|
||||
)}
|
||||
{!sessionStarted ? (
|
||||
<Button
|
||||
aria-label={actionLabel}
|
||||
|
||||
@@ -8,7 +8,6 @@ import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { ThemeToggle } from "@/components/ui/theme-toggle";
|
||||
|
||||
const FASTVIDEO_REPO_URL = "https://haoailab.com/blogs/dreamverse/";
|
||||
const FASTVIDEO_BLOG_URL = "https://haoailab.com/blogs/dreamverse/";
|
||||
|
||||
interface Props {
|
||||
@@ -29,13 +28,13 @@ export default function Header({ timeLeft = null, formatTime = (seconds) => `${s
|
||||
<SidePanelOpenFilled size={20} />
|
||||
</Button>
|
||||
)}
|
||||
<a href={FASTVIDEO_REPO_URL} target="_blank" rel="noopener noreferrer" title="FastVideo on GitHub">
|
||||
<a href="/" title="FastVideo home">
|
||||
<Image src="/logo.svg" alt="FastVideo" width={32} height={32} className="h-8 w-auto sm:h-9 transition-opacity hover:opacity-70" />
|
||||
</a>
|
||||
<div className="hidden sm:flex items-center gap-3">
|
||||
<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
|
||||
<a href={FASTVIDEO_BLOG_URL} target="_blank" rel="noopener noreferrer">
|
||||
<Button variant="outline" size="sm" className="gap-1.5 rounded-full px-3 text-xs">
|
||||
Join Waitlist
|
||||
Blog
|
||||
<ExternalLink className="size-3 opacity-60" />
|
||||
</Button>
|
||||
</a>
|
||||
@@ -53,9 +52,9 @@ export default function Header({ timeLeft = null, formatTime = (seconds) => `${s
|
||||
</div>
|
||||
|
||||
<div className="flex sm:hidden items-center gap-2 px-4 pb-3">
|
||||
<a href="https://docs.google.com/forms/d/e/1FAIpQLSe5zpO1iD8Ds-Ih-fOLm64qd7YZVvuvAyHuJaAfw1hkRHTe_A/viewform?usp=publish-editor" target="_blank" rel="noopener noreferrer">
|
||||
<a href={FASTVIDEO_BLOG_URL} target="_blank" rel="noopener noreferrer">
|
||||
<Button variant="outline" size="sm" className="gap-1.5 rounded-full px-3 text-xs">
|
||||
Join Waitlist
|
||||
Blog
|
||||
<ExternalLink className="size-3 opacity-60" />
|
||||
</Button>
|
||||
</a>
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
import React, { useState, useEffect, useRef, useCallback } from "react";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { PlayFilledAlt } from "@carbon/icons-react";
|
||||
import { Download, Loader2, Share } from "lucide-react";
|
||||
import { Check, ChevronDown, Download, Loader2, Share } from "lucide-react";
|
||||
import { Button } from "@/components/ui/button";
|
||||
interface VideoPlayerProps {
|
||||
videoRef?: React.RefCallback<HTMLVideoElement>;
|
||||
@@ -21,6 +21,8 @@ interface VideoPlayerProps {
|
||||
showLivePlayback?: boolean;
|
||||
defaultMuted?: boolean;
|
||||
rewritePending?: boolean;
|
||||
waitingForSegmentPrompt?: boolean;
|
||||
generatingNextScene?: boolean;
|
||||
onPlaying?: () => void;
|
||||
onDownload?: () => void;
|
||||
}
|
||||
@@ -57,6 +59,8 @@ export default function VideoPlayer({
|
||||
showLivePlayback = true,
|
||||
defaultMuted = true,
|
||||
rewritePending = false,
|
||||
waitingForSegmentPrompt = false,
|
||||
generatingNextScene = false,
|
||||
onPlaying = () => {},
|
||||
onDownload,
|
||||
}: VideoPlayerProps) {
|
||||
@@ -88,6 +92,137 @@ export default function VideoPlayer({
|
||||
setCanShare(typeof navigator.canShare === "function" && window.matchMedia("(pointer: coarse)").matches);
|
||||
}, []);
|
||||
|
||||
// Steering mode: the backend may signal "waiting for next prompt" while the current
|
||||
// segment is still PLAYING (it generates ahead). Only surface the "Segment complete"
|
||||
// overlay once the playhead actually reaches the end of the buffered segment.
|
||||
const [playbackReachedEnd, setPlaybackReachedEnd] = useState(false);
|
||||
// Steering mode: when the user submits the next scene, the segment is generated
|
||||
// (a few seconds of latency) before frames stream. Show a "Generating next scene…"
|
||||
// indicator across that gap so the frozen frame isn't silent. Driven by the explicit
|
||||
// generatingNextScene state (set on scene submit / prompt selection), never inferred
|
||||
// from waitingForSegmentPrompt edges.
|
||||
const [generatingNext, setGeneratingNext] = useState(false);
|
||||
// End of the buffered timeline captured the moment generation starts; the freshly
|
||||
// generated segment extends the buffer past this, which is how we know it landed.
|
||||
const genBoundaryRef = useRef(0);
|
||||
|
||||
// Steering mode: the backend may signal "waiting for next prompt" while the current
|
||||
// segment is still PLAYING (it generates ahead). Track whether the playhead has reached
|
||||
// the end of the buffered segment so end-overlays only show there. Keep tracking through
|
||||
// the generating phase too, so scrubbing back and replaying to the end re-shows them.
|
||||
useEffect(() => {
|
||||
if (!waitingForSegmentPrompt && !generatingNext) {
|
||||
setPlaybackReachedEnd(false);
|
||||
return;
|
||||
}
|
||||
const el = liveVideoEl.current;
|
||||
if (!el) return;
|
||||
const check = () => {
|
||||
try {
|
||||
const buffered = el.buffered;
|
||||
if (buffered.length === 0) return;
|
||||
const end = buffered.end(buffered.length - 1);
|
||||
// Track proximity both ways: scrubbing back off the end hides the overlay,
|
||||
// playing forward to the end re-shows it.
|
||||
setPlaybackReachedEnd(el.ended || end - el.currentTime <= 0.2);
|
||||
} catch {
|
||||
/* buffered access can throw mid-append */
|
||||
}
|
||||
};
|
||||
check();
|
||||
el.addEventListener("timeupdate", check);
|
||||
el.addEventListener("ended", check);
|
||||
el.addEventListener("waiting", check);
|
||||
el.addEventListener("stalled", check);
|
||||
el.addEventListener("pause", check);
|
||||
el.addEventListener("seeking", check);
|
||||
el.addEventListener("seeked", check);
|
||||
el.addEventListener("playing", check);
|
||||
el.addEventListener("progress", check);
|
||||
return () => {
|
||||
el.removeEventListener("timeupdate", check);
|
||||
el.removeEventListener("ended", check);
|
||||
el.removeEventListener("waiting", check);
|
||||
el.removeEventListener("stalled", check);
|
||||
el.removeEventListener("pause", check);
|
||||
el.removeEventListener("seeking", check);
|
||||
el.removeEventListener("seeked", check);
|
||||
el.removeEventListener("playing", check);
|
||||
el.removeEventListener("progress", check);
|
||||
};
|
||||
}, [waitingForSegmentPrompt, generatingNext]);
|
||||
|
||||
useEffect(() => {
|
||||
if (waitingForSegmentPrompt || !sessionStarted) {
|
||||
// Back to waiting (or session over): nothing is generating.
|
||||
setGeneratingNext(false);
|
||||
return;
|
||||
}
|
||||
if (!generatingNextScene) return;
|
||||
// Snapshot the current end of the buffered timeline; the generated segment will
|
||||
// extend the buffer past this boundary.
|
||||
const el = liveVideoEl.current;
|
||||
let boundary = el?.currentTime ?? 0;
|
||||
try {
|
||||
const b = el?.buffered;
|
||||
if (b && b.length) boundary = Math.max(boundary, b.end(b.length - 1));
|
||||
} catch {
|
||||
/* buffered access can throw mid-append */
|
||||
}
|
||||
genBoundaryRef.current = boundary;
|
||||
setGeneratingNext(true);
|
||||
}, [generatingNextScene, waitingForSegmentPrompt, sessionStarted]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!generatingNext) return;
|
||||
if (!sessionStarted) {
|
||||
setGeneratingNext(false);
|
||||
return;
|
||||
}
|
||||
const el = liveVideoEl.current;
|
||||
if (!el) return;
|
||||
// Clear the instant the freshly generated segment lands: the buffer grows past the
|
||||
// boundary captured at generation start (or the playhead advances into the new
|
||||
// frames). Deliberately NOT a bare "playing" handler — scrubbing back and replaying
|
||||
// the EXISTING segment must keep "Generating" up until the new frames actually arrive.
|
||||
const check = () => {
|
||||
try {
|
||||
const b = el.buffered;
|
||||
const end = b.length ? b.end(b.length - 1) : 0;
|
||||
if (end > genBoundaryRef.current + 0.25 || el.currentTime > genBoundaryRef.current + 0.1) {
|
||||
setGeneratingNext(false);
|
||||
}
|
||||
} catch {
|
||||
/* buffered access can throw mid-append */
|
||||
}
|
||||
};
|
||||
check();
|
||||
el.addEventListener("progress", check);
|
||||
el.addEventListener("timeupdate", check);
|
||||
el.addEventListener("durationchange", check);
|
||||
return () => {
|
||||
el.removeEventListener("progress", check);
|
||||
el.removeEventListener("timeupdate", check);
|
||||
el.removeEventListener("durationchange", check);
|
||||
};
|
||||
}, [generatingNext, sessionStarted]);
|
||||
|
||||
// Drive a ~10.5s progress bar during generation so the wait has a visible ETA.
|
||||
const GEN_DURATION_MS = 10500;
|
||||
const [genProgress, setGenProgress] = useState(0);
|
||||
useEffect(() => {
|
||||
if (!generatingNext) {
|
||||
setGenProgress(0);
|
||||
return;
|
||||
}
|
||||
const start = performance.now();
|
||||
setGenProgress(0);
|
||||
const id = setInterval(() => {
|
||||
setGenProgress(Math.min((performance.now() - start) / GEN_DURATION_MS, 1));
|
||||
}, 50);
|
||||
return () => clearInterval(id);
|
||||
}, [generatingNext]);
|
||||
|
||||
return (
|
||||
<div className="mx-auto w-full max-w-3xl mb-2 sm:mb-6">
|
||||
<div className="rounded-2xl border border-border bg-card/50 p-2 shadow-lg backdrop-blur-md">
|
||||
@@ -104,7 +239,7 @@ export default function VideoPlayer({
|
||||
<PlayFilledAlt className="size-10 text-white/25" />
|
||||
<p className="text-sm text-white/50">Your video will appear here</p>
|
||||
</div>
|
||||
) : !avPlaybackStarted && !mediaAppendError && !inQueue && loadingAnimation ? (
|
||||
) : !avPlaybackStarted && !mediaAppendError && !inQueue && !waitingForSegmentPrompt && loadingAnimation ? (
|
||||
<div className="absolute inset-0 flex flex-col items-center justify-center gap-4 bg-slate-900/60 p-4 backdrop-blur-[2px]">
|
||||
<div className="pointer-events-none absolute inset-0 overflow-hidden">
|
||||
<div className="absolute inset-0 -translate-x-full animate-[shimmer_3s_ease-in-out_infinite] bg-gradient-to-r from-transparent via-white/[0.04] to-transparent" />
|
||||
@@ -114,6 +249,35 @@ export default function VideoPlayer({
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
{/* Steering mode: this segment finished — wait gracefully for the user's next scene
|
||||
instead of spinning. The last frame stays visible behind a soft bottom gradient. */}
|
||||
{sessionStarted && waitingForSegmentPrompt && playbackReachedEnd && !mediaAppendError && !inQueue && (
|
||||
<div className="pointer-events-none absolute inset-0 flex flex-col items-center justify-end gap-2 bg-gradient-to-t from-slate-950/85 via-slate-950/15 to-transparent p-5 pb-6 text-center">
|
||||
<div className="flex size-9 items-center justify-center rounded-full border border-white/25 bg-white/10 shadow-lg backdrop-blur-md">
|
||||
<Check className="size-4 text-white/90" />
|
||||
</div>
|
||||
<div className="space-y-0.5">
|
||||
<p className="text-sm font-medium text-white/95">Segment complete</p>
|
||||
<p className="text-xs text-white/65">Describe the next scene below to keep going</p>
|
||||
</div>
|
||||
<ChevronDown className="size-4 animate-bounce text-white/45" />
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Steering mode: generating the next segment — show a ~10.5s progress bar so the wait has an ETA.
|
||||
Gated on playbackReachedEnd like "Segment complete": scrubbing back hides it, playing to the end re-shows it. */}
|
||||
{sessionStarted && generatingNext && playbackReachedEnd && !mediaAppendError && !inQueue && (
|
||||
<div className="pointer-events-none absolute inset-0 flex flex-col items-center justify-end gap-3 bg-gradient-to-t from-slate-950/85 via-slate-950/15 to-transparent p-5 pb-7 text-center">
|
||||
<p className="text-sm font-medium text-white/95">Generating next scene…</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,11 +3,90 @@ import React, { useRef, useMemo, useEffect, useCallback, useState } from "react"
|
||||
import { motion, useAnimationControls } from "framer-motion";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { Check, Lightbulb, Pencil } from "lucide-react";
|
||||
import { Check, Clapperboard, Lightbulb, Pencil } from "lucide-react";
|
||||
|
||||
export const WORKSPACE_ORIGINAL_SELECTION_KEY = "original";
|
||||
export const WORKSPACE_CURRENT_SELECTION_KEY = "current";
|
||||
|
||||
export function SceneHistoryList({ sceneHistory = [] }: { sceneHistory?: Record<string, any>[] }) {
|
||||
const bottomSentinelRef = useRef<HTMLDivElement>(null);
|
||||
const topSentinelRef = useRef<HTMLDivElement>(null);
|
||||
const [showTopFade, setShowTopFade] = useState(false);
|
||||
|
||||
const scenes = useMemo(
|
||||
() => (sceneHistory || []).filter((s) => normalizeText(s?.prompt)),
|
||||
[sceneHistory],
|
||||
);
|
||||
|
||||
const scrollToBottom = useCallback(() => {
|
||||
setTimeout(() => {
|
||||
bottomSentinelRef.current?.scrollIntoView({ block: "end", behavior: "smooth" });
|
||||
}, 60);
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
if (scenes.length > 0) scrollToBottom();
|
||||
}, [scenes.length, scrollToBottom]);
|
||||
|
||||
useEffect(() => {
|
||||
const sentinel = bottomSentinelRef.current;
|
||||
if (!sentinel || typeof ResizeObserver === "undefined") return;
|
||||
let container: HTMLElement | null = sentinel.parentElement;
|
||||
while (container) {
|
||||
const oy = getComputedStyle(container).overflowY;
|
||||
if (oy === "auto" || oy === "scroll") break;
|
||||
container = container.parentElement;
|
||||
}
|
||||
if (!container) return;
|
||||
const ro = new ResizeObserver(() => {
|
||||
const nearBottom = container!.scrollHeight - container!.scrollTop - container!.clientHeight < 96;
|
||||
if (nearBottom) scrollToBottom();
|
||||
});
|
||||
ro.observe(container);
|
||||
return () => ro.disconnect();
|
||||
}, [scenes.length, scrollToBottom]);
|
||||
|
||||
useEffect(() => {
|
||||
const el = topSentinelRef.current;
|
||||
if (!el) return;
|
||||
const observer = new IntersectionObserver(([entry]) => setShowTopFade(!entry.isIntersecting), { threshold: 0.1 });
|
||||
observer.observe(el);
|
||||
return () => observer.disconnect();
|
||||
}, [scenes.length]);
|
||||
|
||||
if (scenes.length === 0) return null;
|
||||
|
||||
return (
|
||||
<section className="relative z-10 flex flex-col mx-auto w-full max-w-2xl max-h-32 overflow-y-auto">
|
||||
<div
|
||||
className={cn(
|
||||
"pointer-events-none sticky top-0 z-20 -mb-12 h-12 bg-linear-to-b from-background to-transparent transition-opacity duration-200",
|
||||
showTopFade ? "opacity-100" : "opacity-0",
|
||||
)}
|
||||
aria-hidden="true"
|
||||
/>
|
||||
<div ref={topSentinelRef} className="h-0 w-0" aria-hidden="true" />
|
||||
<div className="flex flex-col gap-2 pt-12 pb-4">
|
||||
{scenes.map((scene, index) => (
|
||||
<div
|
||||
key={scene.id || index}
|
||||
className="flex items-start gap-3 rounded-xl p-3 transition-colors duration-200 hover:bg-slate-200/50 hover:dark:bg-slate-800/30"
|
||||
>
|
||||
<div className="flex min-w-0 flex-1 flex-col gap-2">
|
||||
<Badge variant="secondary" className="horizontal gap-2 items-center w-fit">
|
||||
<Clapperboard className="size-3 opacity-70" />
|
||||
{`Scene ${index + 1}`}
|
||||
</Badge>
|
||||
<p className="line-clamp-2 text-sm leading-5 text-muted-foreground">{scene.prompt}</p>
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
<div ref={bottomSentinelRef} className="h-0 w-0" aria-hidden="true" />
|
||||
</section>
|
||||
);
|
||||
}
|
||||
|
||||
interface WorkspaceProps {
|
||||
promptEvents?: Record<string, any>[];
|
||||
currentThumbnail?: string | null;
|
||||
@@ -19,6 +98,7 @@ interface WorkspaceProps {
|
||||
selectedClipId?: string;
|
||||
selectedEntryKey?: string;
|
||||
originalClipId?: string;
|
||||
manualMode?: boolean;
|
||||
}
|
||||
|
||||
function normalizeText(value: any): string {
|
||||
@@ -218,7 +298,7 @@ function ChromaGradient({ sessionStarted = false }: { sessionStarted?: boolean }
|
||||
);
|
||||
}
|
||||
|
||||
export default function Workspace({ promptEvents = [], currentThumbnail = null, originalLabel = "", sessionStarted = false, onSelectOriginal, onSelectEvent, onSelectCurrent, selectedClipId, selectedEntryKey: selectedEntryKeyProp, originalClipId = "" }: WorkspaceProps) {
|
||||
export default function Workspace({ promptEvents = [], currentThumbnail = null, originalLabel = "", sessionStarted = false, onSelectOriginal, onSelectEvent, onSelectCurrent, selectedClipId, selectedEntryKey: selectedEntryKeyProp, originalClipId = "", manualMode = false }: WorkspaceProps) {
|
||||
const bottomSentinelRef = useRef<HTMLDivElement>(null);
|
||||
const topSentinelRef = useRef<HTMLDivElement>(null);
|
||||
const [showTopFade, setShowTopFade] = useState(false);
|
||||
@@ -259,6 +339,14 @@ export default function Workspace({ promptEvents = [], currentThumbnail = null,
|
||||
return () => observer.disconnect();
|
||||
}, [conversationEvents.length]);
|
||||
|
||||
if (manualMode) {
|
||||
return (
|
||||
<div className="mt-auto flex flex-col">
|
||||
<ChromaGradient sessionStarted={sessionStarted} />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="mt-auto flex flex-col">
|
||||
<ChromaGradient sessionStarted={sessionStarted} />
|
||||
|
||||
@@ -34,6 +34,7 @@ interface DevtoolsComposerProps {
|
||||
demoMode?: boolean;
|
||||
enhancementEnabled?: boolean;
|
||||
autoExtensionEnabled?: boolean;
|
||||
manualContinuationEnabled?: boolean;
|
||||
loopGenerationEnabled?: boolean;
|
||||
curatedPromptLimit?: number;
|
||||
maxCuratedPromptCount?: number;
|
||||
@@ -49,6 +50,7 @@ interface DevtoolsComposerProps {
|
||||
onEnhancementToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onCuratedPromptLimitChange?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onAutoExtensionToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onManualContinuationToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onLoopToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onLivePromptModeToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onSpeechTranscript?: (text: string) => void;
|
||||
@@ -71,6 +73,7 @@ export default function DevtoolsComposer({
|
||||
demoMode = false,
|
||||
enhancementEnabled = true,
|
||||
autoExtensionEnabled = false,
|
||||
manualContinuationEnabled = false,
|
||||
loopGenerationEnabled = false,
|
||||
curatedPromptLimit = 0,
|
||||
maxCuratedPromptCount = 0,
|
||||
@@ -86,6 +89,7 @@ export default function DevtoolsComposer({
|
||||
onEnhancementToggle = () => {},
|
||||
onCuratedPromptLimitChange = () => {},
|
||||
onAutoExtensionToggle = () => {},
|
||||
onManualContinuationToggle = () => {},
|
||||
onLoopToggle = () => {},
|
||||
onLivePromptModeToggle = () => {},
|
||||
onSpeechTranscript,
|
||||
@@ -328,6 +332,28 @@ export default function DevtoolsComposer({
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex items-start gap-3">
|
||||
<Checkbox
|
||||
id="devtools-steering-mode"
|
||||
checked={manualContinuationEnabled}
|
||||
onCheckedChange={(checked) =>
|
||||
onManualContinuationToggle({
|
||||
target: { checked: Boolean(checked) },
|
||||
currentTarget: { checked: Boolean(checked) },
|
||||
} as React.ChangeEvent<HTMLInputElement>)
|
||||
}
|
||||
/>
|
||||
<div className="space-y-1">
|
||||
<Label htmlFor="devtools-steering-mode">
|
||||
Steering mode
|
||||
</Label>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Drive each segment manually — type the next scene to
|
||||
continue (vs the automatic 6-segment rollout).
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex items-start gap-3">
|
||||
<Checkbox
|
||||
id="devtools-loop-generation"
|
||||
|
||||
@@ -21,6 +21,7 @@ interface DevtoolsShellProps {
|
||||
selectedPresetId?: string;
|
||||
enhancementEnabled?: boolean;
|
||||
autoExtensionEnabled?: boolean;
|
||||
manualContinuationEnabled?: boolean;
|
||||
loopGenerationEnabled?: boolean;
|
||||
canJoinSession?: boolean;
|
||||
canSubmitContinuation?: boolean;
|
||||
@@ -34,6 +35,7 @@ interface DevtoolsShellProps {
|
||||
onEnhancementToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onCuratedPromptLimitChange?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onAutoExtensionToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onManualContinuationToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onLoopToggle?: (e: React.ChangeEvent<HTMLInputElement>) => void;
|
||||
onJoin?: () => void;
|
||||
onLeave?: () => void;
|
||||
@@ -128,6 +130,7 @@ export default function DevtoolsShell({
|
||||
selectedPresetId = '',
|
||||
enhancementEnabled = true,
|
||||
autoExtensionEnabled = false,
|
||||
manualContinuationEnabled = false,
|
||||
loopGenerationEnabled = false,
|
||||
canJoinSession = false,
|
||||
canSubmitContinuation = false,
|
||||
@@ -141,6 +144,7 @@ export default function DevtoolsShell({
|
||||
onEnhancementToggle = () => {},
|
||||
onCuratedPromptLimitChange = () => {},
|
||||
onAutoExtensionToggle = () => {},
|
||||
onManualContinuationToggle = () => {},
|
||||
onLoopToggle = () => {},
|
||||
onJoin = () => {},
|
||||
onLeave = () => {},
|
||||
@@ -284,6 +288,7 @@ export default function DevtoolsShell({
|
||||
demoMode={demoMode}
|
||||
enhancementEnabled={enhancementEnabled}
|
||||
autoExtensionEnabled={autoExtensionEnabled}
|
||||
manualContinuationEnabled={manualContinuationEnabled}
|
||||
loopGenerationEnabled={loopGenerationEnabled}
|
||||
curatedPromptLimit={curatedPromptLimit}
|
||||
maxCuratedPromptCount={maxCuratedPromptCount}
|
||||
@@ -299,6 +304,7 @@ export default function DevtoolsShell({
|
||||
onEnhancementToggle={onEnhancementToggle}
|
||||
onCuratedPromptLimitChange={onCuratedPromptLimitChange}
|
||||
onAutoExtensionToggle={onAutoExtensionToggle}
|
||||
onManualContinuationToggle={onManualContinuationToggle}
|
||||
onLoopToggle={onLoopToggle}
|
||||
onLivePromptModeToggle={onLivePromptModeToggle}
|
||||
onSpeechTranscript={onSpeechTranscript}
|
||||
|
||||
@@ -57,4 +57,32 @@ describe('prependPromptEvent', () => {
|
||||
expect(next[0].promptId).toBe('new');
|
||||
expect(next.some((item: any) => item.promptId === 'p-23')).toBe(false);
|
||||
});
|
||||
|
||||
it('never drops steering scene events when capping', () => {
|
||||
// 30 scenes interleaved with 30 other events — well past the cap.
|
||||
let events: Record<string, any>[] = [];
|
||||
for (let i = 0; i < 30; i += 1) {
|
||||
events = prependPromptEvent(events, {
|
||||
promptId: `scene-${i}`,
|
||||
status: 'submitted',
|
||||
steeringUserPrompt: true,
|
||||
rawText: `scene ${i}`,
|
||||
});
|
||||
events = prependPromptEvent(events, {
|
||||
promptId: `other-${i}`,
|
||||
status: 'submitted',
|
||||
});
|
||||
}
|
||||
|
||||
const scenes = events.filter((e) => e.steeringUserPrompt);
|
||||
expect(scenes).toHaveLength(30);
|
||||
// Oldest-first scene order (and therefore numbering) is stable and complete.
|
||||
expect(scenes[scenes.length - 1].promptId).toBe('scene-0');
|
||||
expect(scenes[0].promptId).toBe('scene-29');
|
||||
// Non-scene events are still capped, oldest dropped first.
|
||||
const others = events.filter((e) => !e.steeringUserPrompt);
|
||||
expect(others.length).toBeLessThanOrEqual(24);
|
||||
expect(others.some((e) => e.promptId === 'other-0')).toBe(false);
|
||||
expect(others[0].promptId).toBe('other-29');
|
||||
});
|
||||
});
|
||||
|
||||
@@ -16,5 +16,19 @@ export function prependPromptEvent(
|
||||
events: Record<string, any>[],
|
||||
event: Record<string, any>,
|
||||
): Record<string, any>[] {
|
||||
return [event, ...events].slice(0, MAX_PROMPT_EVENTS);
|
||||
const next = [event, ...events];
|
||||
if (next.length <= MAX_PROMPT_EVENTS) {
|
||||
return next;
|
||||
}
|
||||
// Steering scene events (steeringUserPrompt) are exempt from the cap: the
|
||||
// scene list is derived from them and must stay complete and stably numbered
|
||||
// for long sessions. Only the oldest non-scene events are dropped.
|
||||
let nonSceneKept = 0;
|
||||
return next.filter((e) => {
|
||||
if (e?.steeringUserPrompt) {
|
||||
return true;
|
||||
}
|
||||
nonSceneKept += 1;
|
||||
return nonSceneKept <= MAX_PROMPT_EVENTS;
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { resolveSessionErrorMessage } from './reducer';
|
||||
import { applyNormalizedSocketEvent, resolveSessionErrorMessage } from './reducer';
|
||||
import { createSessionStore } from '../../stores/session';
|
||||
import { createRewriteStore } from '../../stores/rewrite';
|
||||
import { createStreamStore } from '../../stores/stream';
|
||||
import { createUiStore } from '../../stores/ui';
|
||||
import { createPromptWindowStore } from '../../stores/promptWindow';
|
||||
|
||||
describe('resolveSessionErrorMessage', () => {
|
||||
it('returns a dedicated message for IP session limit errors', () => {
|
||||
@@ -19,3 +24,162 @@ describe('resolveSessionErrorMessage', () => {
|
||||
})).toBe('Backend replica unavailable. Rejoin session.');
|
||||
});
|
||||
});
|
||||
|
||||
function buildContext(overrides: Record<string, unknown> = {}) {
|
||||
const sessionStore = createSessionStore();
|
||||
const rewriteStore = createRewriteStore();
|
||||
const streamStore = createStreamStore();
|
||||
const uiStore = createUiStore();
|
||||
const promptWindowStore = createPromptWindowStore();
|
||||
const avPipeline = {
|
||||
reset: () => {},
|
||||
setStreamCompleted: () => {},
|
||||
noteSegmentInit: () => {},
|
||||
noteSegmentComplete: () => {},
|
||||
maybeStartPlayback: () => {},
|
||||
ensurePipeline: async () => {},
|
||||
};
|
||||
return {
|
||||
sessionStore,
|
||||
promptWindowStore,
|
||||
rewriteStore,
|
||||
streamStore,
|
||||
uiStore,
|
||||
avPipeline,
|
||||
tick: async () => {},
|
||||
defaultAvMime: 'video/mp4',
|
||||
fixedRewriteModel: 'model',
|
||||
parseLatencyMs: () => null,
|
||||
formatPromptWindowEventText: () => '',
|
||||
makePromptId: () => 'generated-id',
|
||||
buildClipLabel: () => 'clip',
|
||||
startSessionCountdown: () => {},
|
||||
clearCountdownInterval: () => {},
|
||||
resetTtffTimer: () => {},
|
||||
startTtffTimer: () => {},
|
||||
preserveArchivedPlaybackSelection: false,
|
||||
finalizeStreamCompletion: async () => {},
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe('steering generatingNextScene flow', () => {
|
||||
it('sets generatingNextScene on prompt/sources_resumed in manual mode', async () => {
|
||||
const context = buildContext();
|
||||
await applyNormalizedSocketEvent(
|
||||
{ type: 'prompt/sources_resumed', payload: { segment_idx: 2 } },
|
||||
context,
|
||||
);
|
||||
expect(context.sessionStore.get().generatingNextScene).toBe(true);
|
||||
expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(false);
|
||||
});
|
||||
|
||||
it('does NOT set generatingNextScene on session/auto_extension_updated', async () => {
|
||||
const context = buildContext();
|
||||
await applyNormalizedSocketEvent(
|
||||
{ type: 'session/auto_extension_updated', payload: { enabled: true } },
|
||||
context,
|
||||
);
|
||||
expect(context.sessionStore.get().generatingNextScene).toBe(false);
|
||||
expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(false);
|
||||
});
|
||||
|
||||
it('clears generatingNextScene when segment media arrives', async () => {
|
||||
const context = buildContext();
|
||||
context.sessionStore.patch({ generatingNextScene: true });
|
||||
await applyNormalizedSocketEvent(
|
||||
{
|
||||
type: 'stream/media_init',
|
||||
payload: { segment_idx: 2, stream_id: 's', mime: 'video/mp4' },
|
||||
},
|
||||
context,
|
||||
);
|
||||
expect(context.sessionStore.get().generatingNextScene).toBe(false);
|
||||
});
|
||||
|
||||
it('clears generatingNextScene and returns to waiting on prompt/sources_blocked', async () => {
|
||||
const context = buildContext();
|
||||
context.sessionStore.patch({ generatingNextScene: true });
|
||||
await applyNormalizedSocketEvent(
|
||||
{ type: 'prompt/sources_blocked', payload: { segment_idx: 3 } },
|
||||
context,
|
||||
);
|
||||
expect(context.sessionStore.get().generatingNextScene).toBe(false);
|
||||
expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(true);
|
||||
});
|
||||
|
||||
it('clears generatingNextScene when the opening prompt falls back', async () => {
|
||||
const context = buildContext();
|
||||
context.sessionStore.patch({ generatingNextScene: true });
|
||||
await applyNormalizedSocketEvent(
|
||||
{
|
||||
type: 'prompt/fallback_used',
|
||||
payload: { prompt_id: 'p1', prompt: '', source: 'user_enhancement_failed' },
|
||||
},
|
||||
context,
|
||||
);
|
||||
expect(context.sessionStore.get().generatingNextScene).toBe(false);
|
||||
expect(context.sessionStore.get().waitingForSegmentPrompt).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('opening prompt id tracking', () => {
|
||||
it('routes prompt lifecycle updates to the frontend-recorded opening event', async () => {
|
||||
// The frontend records the opening scene under its own prompt id and sends it
|
||||
// as initial_rollout_prompt_id; the backend echoes it in status updates.
|
||||
const context = buildContext();
|
||||
context.rewriteStore.addPromptEvent({
|
||||
promptId: 'opening-id',
|
||||
status: 'rewrite_requested',
|
||||
source: 'user_rewrite',
|
||||
text: 'a castle at dawn',
|
||||
steeringUserPrompt: true,
|
||||
rawText: 'a castle at dawn',
|
||||
});
|
||||
|
||||
await applyNormalizedSocketEvent(
|
||||
{ type: 'prompt/enhancing', payload: { prompt_id: 'opening-id' } },
|
||||
context,
|
||||
);
|
||||
let opening = (context.rewriteStore.get().promptEvents as Record<string, any>[])
|
||||
.find((e) => e.promptId === 'opening-id');
|
||||
expect(opening?.status).toBe('enhancing');
|
||||
|
||||
await applyNormalizedSocketEvent(
|
||||
{
|
||||
type: 'prompt/fallback_used',
|
||||
payload: { prompt_id: 'opening-id', prompt: '', source: 'user_enhancement_failed' },
|
||||
},
|
||||
context,
|
||||
);
|
||||
opening = (context.rewriteStore.get().promptEvents as Record<string, any>[])
|
||||
.find((e) => e.promptId === 'opening-id');
|
||||
expect(opening?.status).toBe('ready_fallback');
|
||||
// A failed opening is dropped from the steering scene list instead of
|
||||
// lingering as a ghost "Scene 1".
|
||||
expect(opening?.steeringFailed).toBe(true);
|
||||
});
|
||||
|
||||
it('marks a prompt-scoped session/error (e.g. safety block) as steeringFailed', async () => {
|
||||
const context = buildContext();
|
||||
context.rewriteStore.addPromptEvent({
|
||||
promptId: 'blocked-id',
|
||||
status: 'queued',
|
||||
source: 'user_raw',
|
||||
text: 'a blocked prompt',
|
||||
steeringUserPrompt: true,
|
||||
rawText: 'a blocked prompt',
|
||||
});
|
||||
|
||||
await applyNormalizedSocketEvent(
|
||||
{
|
||||
type: 'session/error',
|
||||
payload: { message: 'Prompt blocked by safety filter.', prompt_id: 'blocked-id' },
|
||||
},
|
||||
context,
|
||||
);
|
||||
const blocked = (context.rewriteStore.get().promptEvents as Record<string, any>[])
|
||||
.find((e) => e.promptId === 'blocked-id');
|
||||
expect(blocked?.steeringFailed).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -99,7 +99,19 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
|
||||
status: "ready_fallback",
|
||||
source: payload.source || "user_raw",
|
||||
text: payload.prompt,
|
||||
// Steering: this prompt produced no segment — drop it from the scene list.
|
||||
steeringFailed: true,
|
||||
});
|
||||
// Steering recovery: enhancement failed so the backend enqueued nothing AND won't
|
||||
// re-emit prompt_sources_blocked (its drained flag is still set). Put the user back to
|
||||
// "describe the next scene" ourselves so the generating overlay clears and they can retry.
|
||||
if (!uiStore.get().simpleMode && sessionStore.get().manualContinuationMode) {
|
||||
sessionStore.patch({
|
||||
waitingForSegmentPrompt: true,
|
||||
generatingNextScene: false,
|
||||
sessionNotice: "Couldn't continue from that prompt — try rephrasing the next scene.",
|
||||
});
|
||||
}
|
||||
console.warn("[PromptEnhanceFallback] Prompt extension failed for this request.");
|
||||
return;
|
||||
|
||||
@@ -243,19 +255,32 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
|
||||
}
|
||||
|
||||
case "prompt/sources_blocked":
|
||||
sessionStore.patch({
|
||||
autoExtensionTimeoutHint: uiStore.get().simpleMode ? "" : "blocked on user input, increase prompt count for smoother experience",
|
||||
});
|
||||
if (sessionStore.get().manualContinuationMode) {
|
||||
sessionStore.patch({ waitingForSegmentPrompt: true, generatingNextScene: false, autoExtensionTimeoutHint: "" });
|
||||
} else {
|
||||
sessionStore.patch({
|
||||
autoExtensionTimeoutHint: uiStore.get().simpleMode ? "" : "blocked on user input, increase prompt count for smoother experience",
|
||||
});
|
||||
}
|
||||
return;
|
||||
|
||||
case "prompt/sources_resumed":
|
||||
sessionStore.patch({
|
||||
autoExtensionTimeoutHint: "",
|
||||
waitingForSegmentPrompt: false,
|
||||
// A real prompt was just selected for the next segment; media arriving
|
||||
// (stream/media_init) clears this again.
|
||||
...(sessionStore.get().manualContinuationMode ? { generatingNextScene: true } : {}),
|
||||
});
|
||||
return;
|
||||
|
||||
case "session/auto_extension_updated":
|
||||
sessionStore.patch({ autoExtensionTimeoutHint: "" });
|
||||
if (event.type === "session/auto_extension_updated") {
|
||||
console.log("[AutoExtensionUpdated]", {
|
||||
enabled: sessionStore.get().autoExtensionEnabled,
|
||||
});
|
||||
}
|
||||
// Deliberately does NOT touch generatingNextScene: toggling auto extension
|
||||
// starts no generation.
|
||||
sessionStore.patch({ autoExtensionTimeoutHint: "", waitingForSegmentPrompt: false });
|
||||
console.log("[AutoExtensionUpdated]", {
|
||||
enabled: sessionStore.get().autoExtensionEnabled,
|
||||
});
|
||||
return;
|
||||
|
||||
case "segment/step_complete":
|
||||
@@ -277,6 +302,7 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
|
||||
projectResetPending: false,
|
||||
sessionExpired: true,
|
||||
sessionNotice: "",
|
||||
generatingNextScene: false,
|
||||
});
|
||||
console.log("Session timed out");
|
||||
clearCountdownInterval();
|
||||
@@ -354,6 +380,8 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
|
||||
return;
|
||||
|
||||
case "stream/media_init":
|
||||
// Segment media is arriving — the "Generating next scene" phase is over.
|
||||
sessionStore.patch({ generatingNextScene: false });
|
||||
streamStore.patch({
|
||||
mediaAppendError: null,
|
||||
loadingAnimation: streamStore.get().avPlaybackStarted ? streamStore.get().loadingAnimation : true,
|
||||
@@ -471,17 +499,30 @@ export async function applyNormalizedSocketEvent(event: any, context: any): Prom
|
||||
sessionStore.patch({
|
||||
generationCapReached: false,
|
||||
sessionNotice: "",
|
||||
generatingNextScene: false,
|
||||
});
|
||||
await finalizeStreamCompletion();
|
||||
return;
|
||||
|
||||
case "session/error": {
|
||||
const errorMessage = resolveSessionErrorMessage(payload);
|
||||
if (payload?.prompt_id) {
|
||||
// Prompt-scoped error (e.g. safety-blocked): the prompt produced no
|
||||
// segment, so drop it from the steering scene list.
|
||||
rewriteStore.trackPromptEvent(payload.prompt_id, {
|
||||
steeringFailed: true,
|
||||
});
|
||||
}
|
||||
sessionStore.patch({
|
||||
generationCapReached: false,
|
||||
preservePlaybackOnClose: false,
|
||||
promptExtensionError: "",
|
||||
sessionNotice: errorMessage,
|
||||
// Steering: a blocked/failed prompt produced no segment and the backend won't re-emit
|
||||
// prompt_sources_blocked, so recover the "describe the next scene" state ourselves.
|
||||
...(!uiStore.get().simpleMode && sessionStore.get().manualContinuationMode
|
||||
? { waitingForSegmentPrompt: true, generatingNextScene: false }
|
||||
: {}),
|
||||
});
|
||||
rewriteStore.patch({
|
||||
rewritingSeedPrompts: false,
|
||||
|
||||
@@ -24,6 +24,9 @@ export interface SessionState {
|
||||
livePromptRewriteMode: boolean;
|
||||
sessionExpired: boolean;
|
||||
projectResetPending: boolean;
|
||||
manualContinuationMode: boolean;
|
||||
waitingForSegmentPrompt: boolean;
|
||||
generatingNextScene: boolean;
|
||||
}
|
||||
|
||||
const DEFAULT_SESSION_STATE: SessionState = {
|
||||
@@ -49,6 +52,12 @@ const DEFAULT_SESSION_STATE: SessionState = {
|
||||
livePromptRewriteMode: false,
|
||||
sessionExpired: false,
|
||||
projectResetPending: false,
|
||||
// Steering-only product: every session drives scenes manually. There is no
|
||||
// auto-rollout mode and no UI selector, so this stays true throughout.
|
||||
manualContinuationMode: true,
|
||||
waitingForSegmentPrompt: false,
|
||||
// True from scene submit / prompt selection until the segment's media arrives.
|
||||
generatingNextScene: false,
|
||||
};
|
||||
|
||||
export type SessionStore = ManagedStore<SessionState> & {
|
||||
|
||||
@@ -1,5 +1,18 @@
|
||||
import { createManagedStore, type ManagedStore } from "./createManagedStore";
|
||||
|
||||
// Prompt-history sources driven by the user's own submissions. Curated (non-user)
|
||||
// entries — e.g. a preset's opening scene — feed the steering scene list and are
|
||||
// exempt from the history cap so Scene 1 survives long sessions.
|
||||
export const USER_PROMPT_SOURCES = new Set([
|
||||
"user_raw",
|
||||
"user",
|
||||
"user_enhanced",
|
||||
"user_rewrite",
|
||||
"user_enhancement_failed",
|
||||
]);
|
||||
|
||||
const PROMPT_HISTORY_CAP = 120;
|
||||
|
||||
export interface StreamState {
|
||||
playingSeedPromptIndex: number | null;
|
||||
generatingSeedPromptIndex: number | null;
|
||||
@@ -159,10 +172,14 @@ export function createStreamStore(initialState: Partial<StreamState> = {}): Stre
|
||||
loopIteration: typeof loopIteration === "number" ? loopIteration : null,
|
||||
};
|
||||
|
||||
const nextHistory = [entry, ...state.promptHistory];
|
||||
return {
|
||||
...state,
|
||||
promptHistoryCounter: nextCounter,
|
||||
promptHistory: [entry, ...state.promptHistory].slice(0, 120),
|
||||
promptHistory: nextHistory.length > PROMPT_HISTORY_CAP
|
||||
? nextHistory.filter((item, index) =>
|
||||
index < PROMPT_HISTORY_CAP || !USER_PROMPT_SOURCES.has(String(item.source || "")))
|
||||
: nextHistory,
|
||||
selectedHistoryId: state.selectedHistoryId || (entry.id as string),
|
||||
};
|
||||
});
|
||||
|
||||
@@ -104,6 +104,11 @@ and does not override stored status.
|
||||
- `GET /api/performance/trends?days=90&run_source=scheduled_main`
|
||||
- `GET /api/performance/records?days=90&run_source=local`
|
||||
|
||||
The current v1 grouping key is `(model_id, gpu_type)`. Baselines are computed
|
||||
from the latest five previous successful records in each group for dashboard
|
||||
context. CI gating uses only records marked `baseline_eligible=true`.
|
||||
V2 records use the same comparison cohort as CI: `workload_id`, `variant_id`,
|
||||
`benchmark_version`, `recipe_fingerprint`, `hardware_profile_id`, and
|
||||
`software_profile_id`. `model_id` and `gpu_type` remain display/filter
|
||||
metadata, so renaming either does not split history. Legacy records still group
|
||||
by `(model_id, gpu_type)`. Dashboard baselines use the latest five previous
|
||||
successful, baseline-eligible records in each group. Summary and trend filters
|
||||
match the latest display metadata after grouping, while the raw records endpoint
|
||||
continues to filter individual records.
|
||||
|
||||
@@ -12,7 +12,8 @@ It serves three audiences:
|
||||
* **Maintainers** — surfaces regressions in a Markdown summary on every
|
||||
performance build and a long-form Plotly dashboard.
|
||||
* **Local developers** — lets you run the same benchmark on your own machine,
|
||||
then compare against the historical baseline for the same model and GPU.
|
||||
then compare against the historical baseline for the same comparable
|
||||
identity.
|
||||
|
||||
## Quick start (local)
|
||||
|
||||
@@ -81,9 +82,9 @@ fastvideo/performance/
|
||||
The HF dataset (`FastVideo/performance-tracking` by default) holds one
|
||||
normalized JSON per run. For v2 records, the rolling baseline is the median of
|
||||
the last 5 successful, baseline-eligible records in the same comparison cohort:
|
||||
`model_id`, `gpu_type`, `workload_id`, `variant_id`, `benchmark_version`,
|
||||
`recipe_fingerprint`, `hardware_profile_id`, and `software_profile_id`. PR and
|
||||
local records are visible in the dashboard but are not baseline eligible.
|
||||
`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
|
||||
`hardware_profile_id`, and `software_profile_id`. PR and local records are
|
||||
visible in the dashboard but are not baseline eligible.
|
||||
|
||||
## Planned Coverage
|
||||
|
||||
@@ -163,12 +164,11 @@ headroom and almost never need touching.
|
||||
|
||||
`compare_baseline.py` loads the last 5 successful, baseline-eligible records
|
||||
for the same comparison cohort from the HF dataset, computes the median for
|
||||
each available metric, and evaluates the current run with the metric's
|
||||
rolling regression policy. For v2 records, that cohort is `model_id`,
|
||||
`gpu_type`, `workload_id`, `variant_id`, `benchmark_version`,
|
||||
`recipe_fingerprint`, `hardware_profile_id`, and `software_profile_id`. For
|
||||
latency, memory, and component times, higher values are regressions. For
|
||||
throughput, lower values are regressions.
|
||||
each available metric, and evaluates the current run with the metric's rolling
|
||||
regression policy. For v2 records, that cohort is `workload_id`, `variant_id`,
|
||||
`benchmark_version`, `recipe_fingerprint`, `hardware_profile_id`, and
|
||||
`software_profile_id`. For latency, memory, and component times, higher values
|
||||
are regressions. For throughput, lower values are regressions.
|
||||
|
||||
A metric exceeds its rolling threshold when both of these are true:
|
||||
|
||||
@@ -188,10 +188,33 @@ slowly add up. Only scheduled-main successful records are baseline eligible.
|
||||
Local and pull-request runs can upload dashboard-visible records, but they do
|
||||
not update future gating baselines.
|
||||
|
||||
Comparator summaries and normalized artifacts include an explicit
|
||||
`comparison_status`:
|
||||
|
||||
| Status | Meaning | CI behavior |
|
||||
|---|---|---|
|
||||
| `PASS` | Comparable baseline exists and no gated metric regressed. Legacy records with no baseline also keep the historical initialization behavior. | Passes |
|
||||
| `REGRESSION` | The record exceeds one of its static thresholds or at least one gated metric regressed against a comparable baseline. | Fails |
|
||||
| `CALIBRATION_NEEDED` | A v2 record has no exact comparable baseline. | Passes, may upload when `PERF_UPLOAD_POLICY=pass`, never seeds a baseline |
|
||||
| `RECIPE_MISMATCH` | The same workload, variant, and benchmark version has trusted successful records under another recipe fingerprint, including records from other hardware or software profiles. | Fails |
|
||||
| `INFRA_ERROR` | The comparator cannot safely classify the record, such as a v2 record missing required identity fields. | Fails |
|
||||
|
||||
`QUALITY_BLOCKED` is reserved for a future variant-promotion workflow and is
|
||||
not emitted by normal rolling-baseline comparison.
|
||||
|
||||
For `RECIPE_MISMATCH`, trusted records are scheduled-main records or records
|
||||
already marked `baseline_eligible`. PR and local calibration uploads remain
|
||||
visible but do not authorize future CI-gating recipe mismatch failures.
|
||||
|
||||
When the baseline shifts for a legitimate reason (torch upgrade, kernel
|
||||
change, etc.) and CI starts failing, use the
|
||||
[`reseed-performance-baseline`](https://github.com/hao-ai-lab/FastVideo/blob/main/.agents/skills/reseed-performance-baseline/SKILL.md)
|
||||
agent skill to advance the rolling median.
|
||||
agent skill to advance the rolling median. To approve the first baseline for
|
||||
a new v2 exact identity, follow that skill with a reviewed scheduled-main
|
||||
full-suite `CALIBRATION_NEEDED` normalized artifact. Its prepare step uses
|
||||
`fastvideo/tests/performance/seed_baseline.py`; after a separate human gate,
|
||||
the skill rechecks the current HF revision and conditionally uploads the whole
|
||||
seed batch in one commit.
|
||||
|
||||
## Schemas
|
||||
|
||||
@@ -207,14 +230,14 @@ configs and remain loadable. New or migrated configs should use
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2
|
||||
"benchmark_version": 3
|
||||
}
|
||||
```
|
||||
|
||||
`benchmark_id` is still required in this phase because raw artifact names,
|
||||
generated-video directories, normalized record paths, and the current rolling
|
||||
baseline comparator still depend on it. The v2 identity fields are config
|
||||
metadata that make the measured workload explicit:
|
||||
`benchmark_id` is still required because raw artifact names, generated-video
|
||||
directories, normalized record paths, and legacy storage directories depend on
|
||||
it. The v2 comparator does not use it as part of the comparison cohort. The v2
|
||||
identity fields make the measured workload explicit:
|
||||
|
||||
| Field | Purpose |
|
||||
|---|---|
|
||||
@@ -225,20 +248,26 @@ metadata that make the measured workload explicit:
|
||||
If a config declares `config_schema_version: 2`, loading fails clearly when any
|
||||
required v2 identity field is missing. If v2 identity or metadata fields are
|
||||
added without `config_schema_version: 2`, loading also fails so partial
|
||||
migrations do not silently run as v1 configs. Optional v2 metadata fields
|
||||
reserved for follow-up work, such as `metric_threshold_policy` and
|
||||
`quality_metadata`, must be JSON objects when present. (`recipe` is emitted
|
||||
by the harness and is not config-declarable.)
|
||||
migrations do not silently run as v1 configs. Optional v2 `quality_metadata`
|
||||
and the v1/v2 `regression_thresholds` policy must be JSON objects when present.
|
||||
(`recipe` is emitted by the harness and is not config-declarable.)
|
||||
|
||||
Recipe fingerprinting, hardware/software profile IDs, exact-identity
|
||||
comparison, and dashboard cohort grouping land with this change: v2 records
|
||||
compare only within their identity cohort, and a record that opens a NEW
|
||||
cohort is marked `baseline_status: "initialized_new_cohort"` (regression
|
||||
gating starts once that cohort accumulates history). Legacy v1 configs still
|
||||
V2 records compare only within their exact identity cohort. A record that opens
|
||||
a new cohort is marked `baseline_status: "initialized_new_cohort"` and
|
||||
`comparison_status: "CALIBRATION_NEEDED"`; it remains ineligible until a
|
||||
reviewed scheduled-main artifact is seeded explicitly. Legacy v1 configs still
|
||||
run and are normalized for reporting, but their records skip rolling-baseline
|
||||
comparison entirely (`baseline_status: "skipped_missing_identity"`, never
|
||||
baseline eligible); only static thresholds gate them. Metric-specific
|
||||
threshold policies and promoted baselines remain separate follow-ups.
|
||||
baseline eligible); only static thresholds gate them. Metric-specific threshold
|
||||
policies are active. `QUALITY_BLOCKED` remains reserved for future variant
|
||||
promotion policy.
|
||||
|
||||
The shipped Wan benchmark uses `benchmark_version: 3` because recipe schema 2
|
||||
changed the recipe fingerprint by removing the legacy `benchmark_id` display
|
||||
name. This intentionally opens a new comparison cohort: after deployment, a
|
||||
reviewed scheduled-main full-suite `CALIBRATION_NEEDED` artifact must be seeded
|
||||
once before rolling regression gating resumes for that exact identity. Static
|
||||
thresholds remain active during calibration.
|
||||
|
||||
### Raw record (`results/perf_*.json`)
|
||||
|
||||
@@ -250,7 +279,7 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"benchmark_version": 3,
|
||||
"model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
|
||||
"device": "NVIDIA L40S",
|
||||
"num_gpus": 2,
|
||||
@@ -289,12 +318,12 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
"dit_time_s": 8.437,
|
||||
"vae_decode_time_s": 3.208,
|
||||
"recipe": {
|
||||
"recipe_schema_version": 1,
|
||||
"recipe_schema_version": 2,
|
||||
"benchmark": {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2
|
||||
"benchmark_version": 3
|
||||
},
|
||||
"model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
|
||||
"init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
|
||||
@@ -314,6 +343,9 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
"python": "3.12",
|
||||
"pytorch": "2.12",
|
||||
"cuda": "13.0",
|
||||
"attention_backend": "FLASH_ATTN",
|
||||
"flash_attention_4_enabled": true,
|
||||
"container_image_version": "py3.12-cuda13.0.0",
|
||||
"packages": {
|
||||
"fastvideo_kernel": "0.3.2",
|
||||
"flashinfer": "0.2.11",
|
||||
@@ -322,7 +354,12 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
}
|
||||
},
|
||||
"software_profile_id": "sw-<sha256-prefix>",
|
||||
"environment_metadata": { "env": { "IMAGE_VERSION": "py3.12-cuda13.0.0" } },
|
||||
"environment_metadata": {
|
||||
"env": {
|
||||
"IMAGE_VERSION": "py3.12-cuda13.0.0",
|
||||
"FASTVIDEO_CONTAINER_IMAGE_REF": "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0@sha256:<digest>"
|
||||
}
|
||||
},
|
||||
"environment_fingerprint": "env-<sha256-prefix>"
|
||||
}
|
||||
```
|
||||
@@ -338,7 +375,7 @@ result, used as the rolling-baseline source of truth.
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"benchmark_version": 3,
|
||||
"timestamp": "2026-05-08T22:00:00+00:00",
|
||||
"commit_sha": "<full sha>",
|
||||
"gpu_type": "NVIDIA L40S",
|
||||
@@ -367,6 +404,10 @@ result, used as the rolling-baseline source of truth.
|
||||
"build_id": "<buildkite-build-id>",
|
||||
"job_id": "<buildkite-job-id>",
|
||||
"quality_metadata": { "quality_status": "canonical" },
|
||||
"baseline_status": "compared",
|
||||
"comparison_status": "PASS",
|
||||
"comparison_status_reason": "Comparable baseline found with no gated regressions",
|
||||
"baseline_eligible": false,
|
||||
"success": true
|
||||
}
|
||||
```
|
||||
@@ -384,9 +425,12 @@ Current `perf_*.json` artifacts that lack the v2 comparison identity are
|
||||
normalized for reporting but skip rolling-baseline comparison and are not marked
|
||||
baseline eligible.
|
||||
|
||||
New records compare only against the same `model_id`, `gpu_type`,
|
||||
`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
|
||||
`hardware_profile_id`, and `software_profile_id` cohort.
|
||||
New v2 records compare only against the same `workload_id`, `variant_id`,
|
||||
`benchmark_version`, `recipe_fingerprint`, `hardware_profile_id`, and
|
||||
`software_profile_id` cohort, independent of the legacy `model_id` directory
|
||||
and `gpu_type` display string. Historical v1 records remain readable for
|
||||
reporting, but current legacy artifacts do not perform a `(model_id, gpu_type)`
|
||||
rolling comparison or seed new rolling baselines.
|
||||
`environment_metadata` and `environment_fingerprint` are audit data and are not
|
||||
part of the comparison key.
|
||||
The recipe prompt digests describe the prompts actually measured by the
|
||||
@@ -406,11 +450,16 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
|
||||
| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `fastvideo/performance/hf_store.py` | Required for upload or private dataset reads. |
|
||||
| `PERF_RUN_SOURCE` | inferred | `compare_baseline.py`, `test_inference_performance.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
|
||||
| `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
|
||||
| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
|
||||
| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Performance pytest exit code. Measured static-threshold failures are attributed per record; otherwise a nonzero code reports an infrastructure error. |
|
||||
| `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
|
||||
| `BUILDKITE_BRANCH`, `BUILDKITE_COMMIT`, `BUILDKITE_PULL_REQUEST` | unset | `compare_baseline.py`, `test_inference_performance.py` | CI metadata stamped into records. |
|
||||
| `DASHBOARD_DAYS` | `30` | `dashboard.py` | Lookback window for the Plotly trend pages. |
|
||||
| `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `fastvideo/performance/hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |
|
||||
| `FASTVIDEO_ATTENTION_BACKEND` | `auto` | `test_inference_performance.py` | Requested attention backend included in `software_profile_id`. |
|
||||
| `FASTVIDEO_FA4` | `0` | `test_inference_performance.py` | FlashAttention-4 toggle included in `software_profile_id`. |
|
||||
| `FASTVIDEO_PERFORMANCE_PROFILE_VERSION` | unset | `test_inference_performance.py` | Optional explicit software cohort/profile version included in `software_profile_id`. |
|
||||
| `IMAGE_VERSION` | unset | `test_inference_performance.py` | CI container image/profile version included in `software_profile_id` when available. |
|
||||
| `FASTVIDEO_CONTAINER_IMAGE_REF` | unset | `pr_test.py`, `launch_l40s_job.py`, `test_inference_performance.py` | Resolved CI container image ref or digest recorded in `environment_metadata` for audit without changing `software_profile_id`. |
|
||||
| `FASTVIDEO_STAGE_LOGGING` | set by the pytest test | `test_inference_performance.py` | Enables pipeline stage timing capture for component metrics during benchmark runs. |
|
||||
|
||||
## CI integration
|
||||
@@ -425,9 +474,11 @@ Each performance build runs pytest first. PR and direct runs only continue to
|
||||
`compare_baseline.py` when that fixed-threshold phase passes; if pytest fails,
|
||||
Markdown summaries and normalized JSON artifacts are not emitted. Scheduled
|
||||
main runs set `PERF_UPLOAD_POLICY=always`, so they still run
|
||||
`compare_baseline.py` (with `PERF_PYTEST_RC` set) after a fixed-threshold
|
||||
failure. Those failed scheduled main runs emit summaries and normalized
|
||||
records, upload records with `success=false`, and are excluded from future
|
||||
`compare_baseline.py` (with `PERF_PYTEST_RC` set) after pytest fails. Each raw
|
||||
record is checked against its own static thresholds: a measured breach reports
|
||||
`REGRESSION`, while unaffected records retain their rolling-baseline status. A
|
||||
nonzero pytest exit with no attributable static-threshold breach reports
|
||||
`INFRA_ERROR`. Failed records have `success=false` and are excluded from future
|
||||
rolling baselines. The dashboard still runs best-effort for observability.
|
||||
When the rolling-baseline phase runs, it emits:
|
||||
|
||||
@@ -487,10 +538,13 @@ When the rolling-baseline phase runs, it emits:
|
||||
2. The pytest test auto-discovers all configs — no test code needed. CI
|
||||
picks it up on the next `/test performance` run.
|
||||
|
||||
3. The first persisted main-branch run with no HF history initializes the
|
||||
baseline (passes automatically). Subsequent runs compare against it. Local
|
||||
and pull-request runs with no HF history also pass, but they do not seed the
|
||||
shared baseline.
|
||||
3. Legacy benchmarks are gated only by their static thresholds. Their current
|
||||
records skip rolling-baseline comparison and are never baseline eligible.
|
||||
V2 benchmarks with no exact comparable baseline report
|
||||
`CALIBRATION_NEEDED`; the record remains visible but does not become
|
||||
baseline eligible until a comparable scheduled-main full-suite run is
|
||||
reviewed and seeded through the prepare, review, and conditional-upload
|
||||
steps in the `reseed-performance-baseline` skill.
|
||||
|
||||
4. If the benchmark targets a GPU not currently in `thresholds`, either add
|
||||
that GPU as a key or rely on the `default` block. Note that `default` is
|
||||
@@ -510,8 +564,12 @@ When the rolling-baseline phase runs, it emits:
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**"No baseline for ... Initializing"** — first run for this comparison cohort.
|
||||
Run will pass and (if persisting) seed the first record.
|
||||
**`CALIBRATION_NEEDED` / "No baseline found for exact comparable identity"** —
|
||||
the v2 run passes, but its normalized record remains
|
||||
`baseline_eligible=false`. Review a successful scheduled-main full-suite
|
||||
normalized artifact, then follow the `reseed-performance-baseline` skill. The
|
||||
utility only prepares a digest-protected manifest; the separately confirmed
|
||||
upload rechecks remote state and commits the batch atomically.
|
||||
|
||||
**Persistent failure right after a torch / kernel / image upgrade** —
|
||||
genuine regression *or* baseline drift. Compare the failing normalized record
|
||||
|
||||
@@ -83,7 +83,6 @@ class TextEncoderConfig(EncoderConfig):
|
||||
arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
|
||||
is_chat_model: bool = False
|
||||
treat_empty_as_dot: bool = False
|
||||
chat_template_enable_thinking: bool = field(default=False, kw_only=True)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -191,10 +191,13 @@ def run_server(serve_config: ServeConfig, *, generator: _GeneratorProto | None =
|
||||
|
||||
generator = VideoGenerator.from_pretrained(config=serve_config.generator)
|
||||
app = build_app(serve_config, generator)
|
||||
# A 15MB init image is ~20MB as a base64 ws message, above uvicorn's
|
||||
# default 16MiB frame cap.
|
||||
uvicorn.run(
|
||||
app,
|
||||
host=serve_config.server.host,
|
||||
port=serve_config.server.port,
|
||||
ws_max_size=32 * 1024 * 1024,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -75,6 +75,25 @@ if not is_flash_attn_2_available():
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def _resolve_torch_dtype(dtype, default: torch.dtype = torch.float32) -> torch.dtype:
|
||||
if isinstance(dtype, torch.dtype):
|
||||
return dtype
|
||||
if isinstance(dtype, str):
|
||||
normalized = dtype.strip().removeprefix("torch.").lower()
|
||||
return {
|
||||
"bfloat16": torch.bfloat16,
|
||||
"bf16": torch.bfloat16,
|
||||
"float16": torch.float16,
|
||||
"fp16": torch.float16,
|
||||
"half": torch.float16,
|
||||
"float32": torch.float32,
|
||||
"fp32": torch.float32,
|
||||
"float": torch.float32,
|
||||
}.get(normalized, default)
|
||||
return default
|
||||
|
||||
|
||||
class Qwen2_5_VLMLP(nn.Module):
|
||||
def __init__(self, config, bias: bool = False):
|
||||
super().__init__()
|
||||
@@ -304,10 +323,13 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module):
|
||||
config_class = Qwen2_5_VLVisionConfig
|
||||
_no_split_modules = ["Qwen2_5_VLVisionBlock"]
|
||||
|
||||
def __init__(self, config) -> None:
|
||||
def __init__(self, config, parent_torch_dtype=None) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.dtype = torch.bfloat16 if config.torch_dtype == "bfloat16" else torch.float32
|
||||
config_torch_dtype = getattr(config, "torch_dtype", None)
|
||||
self.dtype = _resolve_torch_dtype(
|
||||
config_torch_dtype if config_torch_dtype is not None else parent_torch_dtype
|
||||
)
|
||||
|
||||
self.spatial_merge_size = config.spatial_merge_size
|
||||
self.patch_size = config.patch_size
|
||||
@@ -1458,7 +1480,10 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
|
||||
super().__init__()
|
||||
config = _flatten_text_config(config)
|
||||
self.config = config
|
||||
self.visual = Qwen2_5_VisionTransformerPretrainedModel(config.vision_config)
|
||||
self.visual = Qwen2_5_VisionTransformerPretrainedModel(
|
||||
config.vision_config,
|
||||
parent_torch_dtype=getattr(config, "torch_dtype", None),
|
||||
)
|
||||
|
||||
self.model = Qwen2_5_VLModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
@@ -327,12 +327,14 @@ class Qwen3ForCausalLM(TextEncoder):
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> nn.Module:
|
||||
from transformers import AutoModel
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
# FastVideo uses Qwen3 only as a text encoder. Loading the body avoids
|
||||
# materializing an unused LM head and full-vocabulary logits, including
|
||||
# for checkpoints whose metadata names Qwen3ForCausalLM.
|
||||
return AutoModel.from_pretrained(
|
||||
if device.type == "cpu" and torch.cuda.is_available():
|
||||
from fastvideo.distributed import get_local_torch_device
|
||||
|
||||
device = get_local_torch_device()
|
||||
|
||||
return AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
local_files_only=True,
|
||||
torch_dtype=dtype,
|
||||
@@ -366,18 +368,9 @@ class Qwen3ForCausalLM(TextEncoder):
|
||||
residual = None
|
||||
|
||||
if position_ids is None:
|
||||
# Expand to [batch_size, seq_len]: the rotary layer flattens
|
||||
# positions to ``num_tokens`` and reshapes q/k to
|
||||
# ``(num_tokens, -1, head_dim)``. A bare [1, seq_len] only matches
|
||||
# ``num_tokens`` when batch_size == 1; for batched inputs it folds
|
||||
# the batch dim into the head dim and misaligns RoPE. Expanding to
|
||||
# ``batch_size * seq_len`` tokens keeps the layout correct.
|
||||
position_ids = torch.arange(
|
||||
0,
|
||||
hidden_states.shape[1],
|
||||
device=hidden_states.device,
|
||||
dtype=torch.long,
|
||||
).unsqueeze(0).expand(hidden_states.shape[0], -1)
|
||||
0, hidden_states.shape[1], device=hidden_states.device
|
||||
).unsqueeze(0)
|
||||
|
||||
all_hidden_states: tuple[Any, ...] | None = (
|
||||
() if output_hidden_states else None
|
||||
@@ -412,20 +405,6 @@ class Qwen3ForCausalLM(TextEncoder):
|
||||
) -> set[str]:
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: set[str] = set()
|
||||
stacked_params_mapping = self.config.arch_config.stacked_params_mapping
|
||||
# A fused destination is initialized either by one already-fused tensor
|
||||
# or after every split source projection has been loaded. Include
|
||||
# auxiliary quantization parameters (for example scale_weight) rather
|
||||
# than limiting completeness checks to weight/bias tensors.
|
||||
expected_stacked_shards = {
|
||||
(name, shard_id)
|
||||
for name in params_dict
|
||||
for param_name, _, shard_id in stacked_params_mapping
|
||||
if param_name in name
|
||||
}
|
||||
fused_param_names = {name for name, _ in expected_stacked_shards}
|
||||
loaded_stacked_shards: set[tuple[str, str | int]] = set()
|
||||
loaded_fused_params: set[str] = set()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
if name.startswith("model."):
|
||||
@@ -444,66 +423,37 @@ class Qwen3ForCausalLM(TextEncoder):
|
||||
continue
|
||||
name = kv_scale_name
|
||||
|
||||
matched_stacked_param = False
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
for (
|
||||
param_name,
|
||||
weight_name,
|
||||
shard_id,
|
||||
) in self.config.arch_config.stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
matched_stacked_param = True
|
||||
target_name = name.replace(weight_name, param_name)
|
||||
name = name.replace(weight_name, param_name)
|
||||
|
||||
if target_name.endswith(".bias") and target_name not in params_dict:
|
||||
break
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
|
||||
if target_name not in params_dict:
|
||||
break
|
||||
if name not in params_dict:
|
||||
continue
|
||||
|
||||
param = params_dict[target_name]
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_params.add(target_name)
|
||||
loaded_stacked_shards.add((target_name, shard_id))
|
||||
break
|
||||
|
||||
if matched_stacked_param:
|
||||
continue
|
||||
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
|
||||
if name not in params_dict:
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
if name in fused_param_names and name.endswith(".scale_weight"):
|
||||
# Merged scale loaders interpret a missing shard id as shard 0.
|
||||
# An exact fused key is already a complete vector, so copy it
|
||||
# atomically and retain the default loader's shape validation.
|
||||
default_weight_loader(param, loaded_weight)
|
||||
else:
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
|
||||
if name not in params_dict:
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
if name in fused_param_names:
|
||||
loaded_fused_params.add(name)
|
||||
|
||||
required_split_shards = {
|
||||
(name, shard_id)
|
||||
for name, shard_id in expected_stacked_shards
|
||||
if name not in loaded_fused_params
|
||||
}
|
||||
missing_stacked_shards = required_split_shards - loaded_stacked_shards
|
||||
if missing_stacked_shards:
|
||||
formatted_missing = ", ".join(
|
||||
f"{name}[{shard_id}]"
|
||||
for name, shard_id in sorted(
|
||||
missing_stacked_shards,
|
||||
key=lambda item: (item[0], str(item[1])),
|
||||
)
|
||||
)
|
||||
raise ValueError(
|
||||
"Missing required stacked checkpoint shards: "
|
||||
f"{formatted_missing}"
|
||||
)
|
||||
loaded_params.add(name)
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
@@ -401,10 +401,6 @@ class TextEncoderLoader(ComponentLoader):
|
||||
dtype=PRECISION_TO_TYPE[dtype],
|
||||
device=target_device,
|
||||
)
|
||||
# HF passthrough encoders return before FastVideo's FSDP
|
||||
# wrapping path, so the text stage needs their placement to
|
||||
# put token tensors on the same device.
|
||||
model._fastvideo_input_device = target_device
|
||||
return model.eval()
|
||||
|
||||
model = model_cls(model_config) # type: ignore
|
||||
|
||||
@@ -78,9 +78,6 @@ _TEXT_ENCODER_MODELS = {
|
||||
"Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"),
|
||||
"Qwen2_5_VLForConditionalGeneration":
|
||||
("encoders", "reason1", "Reason1TextEncoder"),
|
||||
# Z-Image-Turbo's text_encoder/config.json declares architecture
|
||||
# "Qwen3Model"; route it to the shared Qwen3 encoder (added for Flux2 Klein).
|
||||
"Qwen3Model": ("encoders", "qwen3", "Qwen3ForCausalLM"),
|
||||
"LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
|
||||
"Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
|
||||
"Mistral3ForConditionalGeneration":
|
||||
|
||||
@@ -96,14 +96,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
|
||||
The minimum sigma value for the noise schedule.
|
||||
sigma_data (`float`, *optional*):
|
||||
The sigma data value for scaling.
|
||||
use_reference_discrete_timesteps (`bool`, defaults to False):
|
||||
Some reference schedulers (e.g. Z-Image) construct the timestep
|
||||
schedule by linspacing `num_inference_steps + 1` points from
|
||||
`t_max` to `t_min` and dropping the terminal point. Default
|
||||
(`False`) preserves the original `np.linspace(t_max, t_min,
|
||||
num_inference_steps)` (float64) behaviour used by every existing
|
||||
model. Enable this flag only when matching a reference scheduler
|
||||
that expects the +1 + drop-terminal construction.
|
||||
"""
|
||||
|
||||
_compatibles: list[Any] = []
|
||||
@@ -130,7 +122,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
|
||||
sigma_max: float | None = None,
|
||||
sigma_min: float | None = None,
|
||||
sigma_data: float | None = None,
|
||||
use_reference_discrete_timesteps: bool = False,
|
||||
):
|
||||
if sum([
|
||||
self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
|
||||
@@ -164,7 +155,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
|
||||
|
||||
self.sigmas = sigmas.to(
|
||||
"cpu") # to avoid too much CPU/GPU communication
|
||||
self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item()
|
||||
self.sigma_min = self.sigmas[-1].item()
|
||||
self.sigma_max = self.sigmas[0].item()
|
||||
|
||||
BaseScheduler.__init__(self)
|
||||
@@ -359,19 +350,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
|
||||
if timesteps_array is None:
|
||||
t_max = self._sigma_to_t(self.sigma_max)
|
||||
t_min = self._sigma_to_t(self.sigma_min)
|
||||
if self.config.use_reference_discrete_timesteps:
|
||||
# Some reference schedulers (for example Z-Image) build a
|
||||
# float64 num_steps+1 linspace and drop the terminal point.
|
||||
timesteps_array = np.linspace(
|
||||
t_max,
|
||||
t_min,
|
||||
num_inference_steps + 1,
|
||||
)[:-1]
|
||||
else:
|
||||
# Preserve the original numpy default (float64) here —
|
||||
# casting to float32 silently shifts rounded timestep
|
||||
# values for every existing model that uses this branch.
|
||||
timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
|
||||
timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
|
||||
sigmas_array = timesteps_array / self.config.num_train_timesteps
|
||||
else:
|
||||
sigmas_array = np.array(sigmas).astype(np.float32)
|
||||
|
||||
@@ -18,6 +18,7 @@ from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from huggingface_hub.constants import ENDPOINT
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration — read once at import time, shared across both consumers
|
||||
@@ -35,7 +36,8 @@ SYNC_REUSE_TTL_SECONDS = int(os.environ.get("PERFORMANCE_TRACKING_SYNC_REUSE_TTL
|
||||
|
||||
def sanitize(value: str) -> str:
|
||||
"""Return a filesystem- and HF-path-safe version of *value*."""
|
||||
return re.sub(r"[^A-Za-z0-9._-]", "_", value)
|
||||
sanitized = re.sub(r"[^A-Za-z0-9._-]", "_", value)
|
||||
return f"_{sanitized}" if not sanitized or sanitized.startswith(".") else sanitized
|
||||
|
||||
|
||||
def safe_float(value: Any) -> float | None:
|
||||
@@ -60,6 +62,19 @@ def is_baseline_eligible_record(record: dict[str, Any]) -> bool:
|
||||
return "baseline_eligible" not in record and "run_source" not in record
|
||||
|
||||
|
||||
def _parse_record_timestamp(record: dict[str, Any]) -> datetime | None:
|
||||
raw_ts = record.get("timestamp")
|
||||
if not raw_ts:
|
||||
return None
|
||||
try:
|
||||
ts = datetime.fromisoformat(str(raw_ts))
|
||||
except ValueError:
|
||||
return None
|
||||
if ts.tzinfo is None:
|
||||
ts = ts.replace(tzinfo=timezone.utc)
|
||||
return ts
|
||||
|
||||
|
||||
def resolve_hf_token() -> str | None:
|
||||
"""Return the first configured Hugging Face token env var."""
|
||||
for env_var in HF_TOKEN_ENV_VARS:
|
||||
@@ -95,11 +110,23 @@ def _sync_marker_is_fresh(marker_path: str) -> bool:
|
||||
return age.total_seconds() <= SYNC_REUSE_TTL_SECONDS
|
||||
|
||||
|
||||
def _sync_marker_matches_request(marker_path: str, revision: str | None) -> bool:
|
||||
try:
|
||||
with open(marker_path, encoding="utf-8") as marker:
|
||||
marker_data = json.load(marker)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return False
|
||||
if marker_data.get("endpoint") != ENDPOINT or marker_data.get("repo_id") != HF_REPO_ID:
|
||||
return False
|
||||
return marker_data.get("revision") == revision
|
||||
|
||||
|
||||
def sync_from_hf(
|
||||
local_dir: str,
|
||||
*,
|
||||
strict: bool = False,
|
||||
reuse_existing: bool = False,
|
||||
revision: str | None = None,
|
||||
) -> str:
|
||||
"""Download the HF dataset repo snapshot to *local_dir*.
|
||||
|
||||
@@ -116,14 +143,18 @@ def sync_from_hf(
|
||||
snapshot checks when compare and dashboard scripts run sequentially in the
|
||||
same CI job, without silently reusing stale data in persistent local or
|
||||
long-lived runner environments.
|
||||
|
||||
Pass ``revision`` to pin the snapshot to a previously read Hub commit. This
|
||||
is used by conditional writers that must validate one exact remote state
|
||||
before committing with that revision as their parent.
|
||||
"""
|
||||
marker_path = _sync_marker_path(local_dir)
|
||||
if reuse_existing and os.path.exists(marker_path):
|
||||
if _sync_marker_is_fresh(marker_path):
|
||||
if _sync_marker_is_fresh(marker_path) and _sync_marker_matches_request(marker_path, revision):
|
||||
print(f"hf_store: reusing existing sync at {local_dir}")
|
||||
return local_dir
|
||||
os.remove(marker_path)
|
||||
print(f"hf_store: existing sync at {local_dir} is stale; refreshing")
|
||||
print(f"hf_store: existing sync at {local_dir} is stale or mismatched; refreshing")
|
||||
|
||||
if not reuse_existing and os.path.exists(marker_path):
|
||||
os.remove(marker_path)
|
||||
@@ -143,13 +174,17 @@ def sync_from_hf(
|
||||
local_dir=local_dir,
|
||||
token=resolve_hf_token(),
|
||||
allow_patterns="*.json",
|
||||
revision=revision,
|
||||
)
|
||||
os.makedirs(local_dir, exist_ok=True)
|
||||
with open(marker_path, "w", encoding="utf-8") as marker:
|
||||
json.dump({
|
||||
"repo_id": HF_REPO_ID,
|
||||
"synced_at": datetime.now(timezone.utc).isoformat(),
|
||||
}, marker)
|
||||
json.dump(
|
||||
{
|
||||
"endpoint": ENDPOINT,
|
||||
"repo_id": HF_REPO_ID,
|
||||
"synced_at": datetime.now(timezone.utc).isoformat(),
|
||||
"revision": revision,
|
||||
}, marker)
|
||||
except Exception as exc:
|
||||
if strict:
|
||||
raise
|
||||
@@ -231,7 +266,7 @@ def load_records(
|
||||
if days is not None:
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=days)
|
||||
|
||||
records: list[dict[str, Any]] = []
|
||||
records: list[tuple[datetime, str, dict[str, Any]]] = []
|
||||
|
||||
for path in sorted(glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True)):
|
||||
try:
|
||||
@@ -246,21 +281,17 @@ def load_records(
|
||||
if baseline_eligible_only and not is_baseline_eligible_record(data):
|
||||
continue
|
||||
|
||||
if cutoff is not None:
|
||||
raw_ts = data.get("timestamp")
|
||||
if raw_ts:
|
||||
try:
|
||||
ts = datetime.fromisoformat(raw_ts)
|
||||
if ts.tzinfo is None:
|
||||
ts = ts.replace(tzinfo=timezone.utc)
|
||||
if ts < cutoff:
|
||||
continue
|
||||
except ValueError:
|
||||
pass # keep records with unparsable timestamps
|
||||
ts = _parse_record_timestamp(data)
|
||||
if cutoff is not None and ts is not None and ts < cutoff:
|
||||
continue
|
||||
|
||||
records.append(data)
|
||||
records.append((
|
||||
ts or datetime.min.replace(tzinfo=timezone.utc),
|
||||
path,
|
||||
data,
|
||||
))
|
||||
|
||||
return records
|
||||
return [data for _ts, _path, data in sorted(records)]
|
||||
|
||||
|
||||
def load_records_for_model(
|
||||
@@ -332,6 +363,36 @@ def load_records_for_model(
|
||||
return records
|
||||
|
||||
|
||||
def load_records_for_identity(
|
||||
local_dir: str,
|
||||
identity_filters: dict[str, str],
|
||||
*,
|
||||
last_n: int | None = None,
|
||||
successful_only: bool = True,
|
||||
baseline_eligible_only: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return records matching comparable identity fields.
|
||||
|
||||
V2 performance comparison is intentionally independent from the legacy
|
||||
``model_id`` directory and ``gpu_type`` display string. The comparable
|
||||
identity filters usually contain the full v2 comparison key, and may also
|
||||
contain a subset when looking for same-cohort recipe mismatches.
|
||||
"""
|
||||
records = load_records(
|
||||
local_dir,
|
||||
successful_only=successful_only,
|
||||
baseline_eligible_only=baseline_eligible_only,
|
||||
)
|
||||
|
||||
for key, expected in identity_filters.items():
|
||||
records = [r for r in records if str(r.get(key)) == str(expected)]
|
||||
|
||||
if last_n is not None:
|
||||
records = records[-last_n:]
|
||||
|
||||
return records
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DataFrame helpers (dashboard / analytics consumers)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -23,6 +23,10 @@ FRONTEND_DIST = os.path.abspath(
|
||||
os.path.join(os.path.dirname(__file__), "..", "..", "performance_dashboard", "frontend", "dist"))
|
||||
|
||||
|
||||
def _matches_display_filters(record: dict[str, Any], model_id: str | None, gpu_type: str | None) -> bool:
|
||||
return (not model_id or record.get("model_id") == model_id) and (not gpu_type or record.get("gpu_type") == gpu_type)
|
||||
|
||||
|
||||
class PerformanceDataStore:
|
||||
|
||||
def __init__(self, tracking_root: str | None = None) -> None:
|
||||
@@ -147,11 +151,10 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
|
||||
# query is kept only so the frontend can share filter state across
|
||||
# endpoints without affecting the summary semantics.
|
||||
loaded = data_store.load_records(days=None)
|
||||
filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type)
|
||||
rows = build_latest_summary(
|
||||
filtered,
|
||||
run_source=run_source,
|
||||
)
|
||||
rows = [
|
||||
row for row in build_latest_summary(loaded, run_source=run_source)
|
||||
if _matches_display_filters(row, model_id, gpu_type)
|
||||
]
|
||||
return {
|
||||
"rows": rows,
|
||||
"count": len(rows),
|
||||
@@ -177,8 +180,10 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
|
||||
run_source: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
loaded = data_store.load_records(days=days)
|
||||
filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type, run_source=run_source)
|
||||
groups = build_trends(filtered)
|
||||
source_filtered = filter_records(loaded, run_source=run_source)
|
||||
groups = [
|
||||
group for group in build_trends(source_filtered) if _matches_display_filters(group, model_id, gpu_type)
|
||||
]
|
||||
return {
|
||||
"groups": groups,
|
||||
"count": len(groups),
|
||||
|
||||
@@ -17,7 +17,7 @@ from fastvideo.performance.hf_store import is_baseline_eligible_record, safe_flo
|
||||
from fastvideo.performance.metric_policy import regression_delta, resolve_metric_policies
|
||||
|
||||
Record = dict[str, Any]
|
||||
CohortKey = tuple[str, str, str, str, str, str, str, str]
|
||||
CohortKey = tuple[str, ...]
|
||||
COMPARISON_COHORT_KEYS = (
|
||||
"workload_id",
|
||||
"variant_id",
|
||||
@@ -102,16 +102,18 @@ def record_comparison_metadata(record: Record) -> Record:
|
||||
return {key: _cohort_metadata_value(record.get(key)) for key in COMPARISON_COHORT_KEYS}
|
||||
|
||||
|
||||
def _record_has_complete_v2_identity(record: Record) -> bool:
|
||||
return all(_cohort_metadata_value(record.get(key)) != "" for key in COMPARISON_COHORT_KEYS)
|
||||
|
||||
|
||||
def comparison_cohort_key(record: Record) -> CohortKey:
|
||||
if _record_has_complete_v2_identity(record):
|
||||
return ("v2", *(_cohort_key_value(record.get(key)) for key in COMPARISON_COHORT_KEYS))
|
||||
return (
|
||||
"legacy",
|
||||
str(record.get("model_id") or "unknown"),
|
||||
str(record.get("gpu_type") or "unknown"),
|
||||
_cohort_key_value(record.get("workload_id")),
|
||||
_cohort_key_value(record.get("variant_id")),
|
||||
_cohort_key_value(record.get("benchmark_version")),
|
||||
_cohort_key_value(record.get("recipe_fingerprint")),
|
||||
_cohort_key_value(record.get("hardware_profile_id")),
|
||||
_cohort_key_value(record.get("software_profile_id")),
|
||||
*(_cohort_key_value(record.get(key)) for key in COMPARISON_COHORT_KEYS),
|
||||
)
|
||||
|
||||
|
||||
@@ -122,20 +124,18 @@ def group_by_comparison_cohort(records: list[Record]) -> dict[CohortKey, list[Re
|
||||
return {key: sorted(value, key=record_sort_key) for key, value in groups.items()}
|
||||
|
||||
|
||||
def comparison_sort_key(record: Record) -> CohortKey:
|
||||
return comparison_cohort_key(record)
|
||||
def comparison_sort_key(record: Record) -> tuple[str, ...]:
|
||||
return (
|
||||
str(record.get("model_id") or "unknown"),
|
||||
str(record.get("gpu_type") or "unknown"),
|
||||
*(_cohort_key_value(record.get(key)) for key in COMPARISON_COHORT_KEYS),
|
||||
)
|
||||
|
||||
|
||||
def latest_row_sort_key(row: Record) -> tuple[Any, ...]:
|
||||
return (row["status"] != "fail", *comparison_sort_key(row))
|
||||
|
||||
|
||||
def group_identity(key: CohortKey) -> tuple[str, str]:
|
||||
model_id = key[0]
|
||||
gpu_type = key[1]
|
||||
return model_id, gpu_type
|
||||
|
||||
|
||||
def baseline_value(records: list[Record], metric_key: str) -> float | None:
|
||||
values = [safe_float(record.get(metric_key)) for record in records]
|
||||
values = [value for value in values if value is not None]
|
||||
@@ -149,8 +149,7 @@ def build_latest_summary(records: list[Record],
|
||||
baseline_window: int = 5,
|
||||
run_source: str | None = None) -> list[Record]:
|
||||
rows: list[Record] = []
|
||||
for key, group in group_by_comparison_cohort(records).items():
|
||||
model_id, gpu_type = group_identity(key)
|
||||
for group in group_by_comparison_cohort(records).values():
|
||||
latest_candidates = group
|
||||
if run_source:
|
||||
latest_candidates = [record for record in group if record_run_source(record) == run_source]
|
||||
@@ -158,6 +157,8 @@ def build_latest_summary(records: list[Record],
|
||||
continue
|
||||
|
||||
latest = latest_candidates[-1]
|
||||
model_id = str(latest.get("model_id") or "unknown")
|
||||
gpu_type = str(latest.get("gpu_type") or "unknown")
|
||||
latest_index = next(index for index, record in enumerate(group) if record is latest)
|
||||
baseline_pool = [
|
||||
record for record in group[:latest_index]
|
||||
@@ -224,8 +225,10 @@ def build_latest_summary(records: list[Record],
|
||||
|
||||
def build_trends(records: list[Record]) -> list[Record]:
|
||||
trends: list[Record] = []
|
||||
for key, group in group_by_comparison_cohort(records).items():
|
||||
model_id, gpu_type = group_identity(key)
|
||||
for group in group_by_comparison_cohort(records).values():
|
||||
latest = group[-1]
|
||||
model_id = str(latest.get("model_id") or "unknown")
|
||||
gpu_type = str(latest.get("gpu_type") or "unknown")
|
||||
points = []
|
||||
for record in group:
|
||||
metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
|
||||
@@ -244,7 +247,7 @@ def build_trends(records: list[Record]) -> list[Record]:
|
||||
trends.append({
|
||||
"model_id": model_id,
|
||||
"gpu_type": gpu_type,
|
||||
**record_comparison_metadata(group[-1]),
|
||||
**record_comparison_metadata(latest),
|
||||
"points": points,
|
||||
})
|
||||
return sorted(trends, key=comparison_sort_key)
|
||||
|
||||
@@ -73,36 +73,25 @@ class DecodingStage(PipelineStage):
|
||||
|
||||
cfg = getattr(self.vae, "config", None)
|
||||
|
||||
# Matrix-Game 2.0-style: z = z * std + mean. Some configs declare these
|
||||
# fields as None, so test the values rather than mere attribute presence
|
||||
# -- otherwise torch.tensor(None) raises instead of falling through.
|
||||
latents_mean_value = getattr(cfg, "latents_mean", None) if cfg is not None else None
|
||||
latents_std_value = getattr(cfg, "latents_std", None) if cfg is not None else None
|
||||
if latents_mean_value is not None and latents_std_value is not None:
|
||||
latents_mean = torch.tensor(latents_mean_value, device=latents.device,
|
||||
# Matrix-Game 2.0-style: z = z * std + mean
|
||||
if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")):
|
||||
latents_mean = torch.tensor(cfg.latents_mean, device=latents.device,
|
||||
dtype=latents.dtype).view(1, -1, 1, 1, 1)
|
||||
latents_std = torch.tensor(latents_std_value, device=latents.device,
|
||||
dtype=latents.dtype).view(1, -1, 1, 1, 1)
|
||||
latents_std = torch.tensor(cfg.latents_std, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
|
||||
return latents * latents_std + latents_mean
|
||||
|
||||
# Diffusers-style: scaling_factor (+ optional shift_factor), read from
|
||||
# the module only. Nearly every VAE *config* also carries a
|
||||
# scaling_factor whose decode() already accounts for it, so sourcing it
|
||||
# from cfg here would newly divide latents for VAEs that have always
|
||||
# passed through untouched (flux2, hunyuanvideo, hunyuanvideo15, ltx2).
|
||||
scaling_factor = getattr(self.vae, "scaling_factor", None)
|
||||
if scaling_factor is not None:
|
||||
if isinstance(scaling_factor, torch.Tensor):
|
||||
latents = latents / scaling_factor.to(latents.device, latents.dtype)
|
||||
# Diffusers-style: scaling_factor (+ optional shift_factor)
|
||||
if hasattr(self.vae, "scaling_factor"):
|
||||
if isinstance(self.vae.scaling_factor, torch.Tensor):
|
||||
latents = latents / self.vae.scaling_factor.to(latents.device, latents.dtype)
|
||||
else:
|
||||
latents = latents / scaling_factor
|
||||
latents = latents / self.vae.scaling_factor
|
||||
|
||||
shift_factor = getattr(self.vae, "shift_factor", None)
|
||||
if shift_factor is not None:
|
||||
if isinstance(shift_factor, torch.Tensor):
|
||||
latents = latents + shift_factor.to(latents.device, latents.dtype)
|
||||
if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
|
||||
if isinstance(self.vae.shift_factor, torch.Tensor):
|
||||
latents = latents + self.vae.shift_factor.to(latents.device, latents.dtype)
|
||||
else:
|
||||
latents = latents + shift_factor
|
||||
latents = latents + self.vae.shift_factor
|
||||
|
||||
return latents
|
||||
|
||||
|
||||
@@ -221,20 +221,6 @@ class TextEncodingStage(PipelineStage):
|
||||
encoder_device = torch.device(target_device)
|
||||
moved_for_forward = True
|
||||
|
||||
# An explicit `device=` wins. Otherwise follow the encoder's real
|
||||
# param device. Once it has been moved for the forward that is the
|
||||
# target device, and an HF-passthrough encoder's
|
||||
# _fastvideo_input_device (stamped at load, e.g. "cpu" under
|
||||
# text_encoder_cpu_offload) is stale -- honouring it would feed cpu
|
||||
# token ids to cuda weights. The marker only speaks when nothing
|
||||
# moved and the module has no parameters to speak for it.
|
||||
if device is not None:
|
||||
input_device = torch.device(target_device)
|
||||
elif moved_for_forward:
|
||||
input_device = encoder_device
|
||||
else:
|
||||
input_device = getattr(text_encoder, "_fastvideo_input_device", encoder_device)
|
||||
|
||||
tok_kwargs = dict(encoder_config.tokenizer_kwargs)
|
||||
if max_length is not None:
|
||||
tok_kwargs["max_length"] = max_length
|
||||
@@ -278,7 +264,7 @@ class TextEncodingStage(PipelineStage):
|
||||
# pre-format prompts into message lists upstream and rely on
|
||||
# the inner tokenizer + full tokenizer_kwargs (which include
|
||||
# add_generation_prompt). Preserve that original path exactly.
|
||||
text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(input_device)
|
||||
text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(encoder_device)
|
||||
else:
|
||||
# Two-step approach matching Diffusers: format with chat
|
||||
# template first, then tokenize the resulting strings.
|
||||
@@ -289,12 +275,12 @@ class TextEncodingStage(PipelineStage):
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=encoder_config.chat_template_enable_thinking,
|
||||
enable_thinking=False,
|
||||
)
|
||||
formatted_texts.append(formatted)
|
||||
text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(input_device)
|
||||
text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(encoder_device)
|
||||
else:
|
||||
text_inputs = tok(processed_texts, **tok_kwargs).to(input_device)
|
||||
text_inputs = tok(processed_texts, **tok_kwargs).to(encoder_device)
|
||||
|
||||
input_ids = text_inputs["input_ids"]
|
||||
attention_mask = text_inputs["attention_mask"]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Guard the Modal FA4 defaults that keep CI lanes on their intended backend.
|
||||
"""Guard Modal runtime policies that keep CI lanes on their intended backend.
|
||||
|
||||
Pure text/AST analysis: no fastvideo imports, no torch, no Modal client.
|
||||
"""
|
||||
@@ -40,6 +40,22 @@ def test_ssim_launcher_keeps_fa4_enabled_by_default():
|
||||
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1")' in source
|
||||
|
||||
|
||||
def test_performance_identity_env_reaches_modal_runtime():
|
||||
pr_source = PR_TEST.read_text(encoding="utf-8")
|
||||
launch_source = LAUNCH_L40S_JOB.read_text(encoding="utf-8")
|
||||
runtime_secret = pr_source.split("ci_env_secret =", 1)[1].split("hf_secret =", 1)[0]
|
||||
|
||||
for key in ("FASTVIDEO_ATTENTION_BACKEND", "FASTVIDEO_PERFORMANCE_PROFILE_VERSION"):
|
||||
assert key in runtime_secret
|
||||
assert "FASTVIDEO_PERFORMANCE_PROFILE_VERSION" in launch_source.split(".env({", 1)[1]
|
||||
|
||||
|
||||
def test_performance_lane_classifies_pull_requests_before_main():
|
||||
function_strings = _function_strings(PR_TEST, "run_performance_tests")
|
||||
assert function_strings.index(
|
||||
"BUILDKITE_PULL_REQUEST") < function_strings.index("BUILDKITE_BRANCH")
|
||||
|
||||
|
||||
def test_pr_model_load_and_training_lanes_disable_fa4():
|
||||
lanes = {
|
||||
"run_transformer_tests": "pytest ./fastvideo/tests/transformers -vs",
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.models.encoders.qwen2_5_vl_custom import (
|
||||
Qwen2_5_VisionTransformerPretrainedModel,
|
||||
Qwen2_5_VLForConditionalGenerationSimple,
|
||||
)
|
||||
|
||||
|
||||
def _vision_config(torch_dtype=None):
|
||||
return SimpleNamespace(
|
||||
torch_dtype=torch_dtype,
|
||||
spatial_merge_size=2,
|
||||
patch_size=14,
|
||||
fullatt_block_indexes=[],
|
||||
window_size=112,
|
||||
temporal_patch_size=2,
|
||||
in_channels=3,
|
||||
hidden_size=8,
|
||||
num_heads=2,
|
||||
depth=0,
|
||||
_attn_implementation="sdpa",
|
||||
out_hidden_size=8,
|
||||
)
|
||||
|
||||
|
||||
def _full_config(torch_dtype):
|
||||
return SimpleNamespace(
|
||||
vision_config=_vision_config(),
|
||||
torch_dtype=torch_dtype,
|
||||
vocab_size=16,
|
||||
hidden_size=8,
|
||||
num_hidden_layers=0,
|
||||
num_attention_heads=2,
|
||||
pad_token_id=0,
|
||||
_attn_implementation="sdpa",
|
||||
rms_norm_eps=1e-6,
|
||||
max_position_embeddings=32,
|
||||
rope_theta=1_000_000.0,
|
||||
rope_scaling=None,
|
||||
)
|
||||
|
||||
|
||||
def test_vision_dtype_uses_parent_bfloat16_string_when_vision_dtype_missing():
|
||||
model = Qwen2_5_VisionTransformerPretrainedModel(
|
||||
_vision_config(),
|
||||
parent_torch_dtype="bfloat16",
|
||||
)
|
||||
|
||||
assert model.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_vision_dtype_accepts_parent_torch_dtype_object():
|
||||
model = Qwen2_5_VisionTransformerPretrainedModel(
|
||||
_vision_config(),
|
||||
parent_torch_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
assert model.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_vision_dtype_strips_parent_dtype_string_whitespace():
|
||||
model = Qwen2_5_VisionTransformerPretrainedModel(
|
||||
_vision_config(),
|
||||
parent_torch_dtype=" torch.bfloat16 ",
|
||||
)
|
||||
|
||||
assert model.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_vision_dtype_prefers_explicit_vision_dtype_over_parent_dtype():
|
||||
model = Qwen2_5_VisionTransformerPretrainedModel(
|
||||
_vision_config(torch_dtype="float16"),
|
||||
parent_torch_dtype="bfloat16",
|
||||
)
|
||||
|
||||
assert model.dtype == torch.float16
|
||||
|
||||
|
||||
def test_vision_dtype_falls_back_to_float32_for_missing_or_unknown_dtype():
|
||||
missing = Qwen2_5_VisionTransformerPretrainedModel(_vision_config())
|
||||
unknown = Qwen2_5_VisionTransformerPretrainedModel(
|
||||
_vision_config(),
|
||||
parent_torch_dtype="not-a-real-dtype",
|
||||
)
|
||||
|
||||
assert missing.dtype == torch.float32
|
||||
assert unknown.dtype == torch.float32
|
||||
|
||||
|
||||
def test_conditional_generation_passes_parent_dtype_to_visual_tower():
|
||||
model = Qwen2_5_VLForConditionalGenerationSimple(_full_config("torch.bfloat16"))
|
||||
|
||||
assert model.visual.dtype == torch.bfloat16
|
||||
@@ -95,8 +95,13 @@ image = (
|
||||
"PATH": "/root/.cargo/bin:$PATH",
|
||||
"HF_HOME": "/root/data/.cache",
|
||||
"TOKENIZERS_PARALLELISM": "false",
|
||||
"IMAGE_VERSION": IMAGE_VERSION,
|
||||
"FASTVIDEO_CONTAINER_IMAGE_REF": IMAGE_REF,
|
||||
**({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
|
||||
"FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
|
||||
**({
|
||||
"FASTVIDEO_PERFORMANCE_PROFILE_VERSION": os.environ["FASTVIDEO_PERFORMANCE_PROFILE_VERSION"]
|
||||
} if os.environ.get("FASTVIDEO_PERFORMANCE_PROFILE_VERSION") else {}),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4). Generic ad hoc jobs should follow the
|
||||
# product default unless a caller opts in through the local env or
|
||||
# --env-vars.
|
||||
|
||||
@@ -65,11 +65,21 @@ ci_env_secret = modal.Secret.from_dict({
|
||||
"BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
"BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"BUILDKITE_SOURCE": os.environ.get("BUILDKITE_SOURCE", ""),
|
||||
"BUILDKITE_BUILD_URL": os.environ.get("BUILDKITE_BUILD_URL", ""),
|
||||
"BUILDKITE_BUILD_ID": os.environ.get("BUILDKITE_BUILD_ID", ""),
|
||||
"BUILDKITE_JOB_ID": os.environ.get("BUILDKITE_JOB_ID", ""),
|
||||
"TEST_SCOPE": os.environ.get("TEST_SCOPE", ""),
|
||||
"IMAGE_VERSION": image_version,
|
||||
"FASTVIDEO_CONTAINER_IMAGE_REF": image_ref,
|
||||
**{
|
||||
key: os.environ[key]
|
||||
for key in (
|
||||
"FASTVIDEO_ATTENTION_BACKEND",
|
||||
"FASTVIDEO_PERFORMANCE_PROFILE_VERSION",
|
||||
)
|
||||
if os.environ.get(key)
|
||||
},
|
||||
**({
|
||||
"UV_TORCH_BACKEND": uv_torch_backend_override
|
||||
} if uv_torch_backend_override else {}),
|
||||
@@ -408,12 +418,12 @@ def run_performance_tests():
|
||||
"export HF_HOME='/root/data/.cache' && "
|
||||
"export PERFORMANCE_TRACKING_ROOT='/tmp/perf-tracking' && "
|
||||
"hf auth login --token $HF_API_KEY && "
|
||||
"if [ \"${BUILDKITE_BRANCH:-}\" = 'main' ] && [ \"${TEST_SCOPE:-}\" = 'full' ]; then "
|
||||
"export PERF_RUN_SOURCE='scheduled_main'; "
|
||||
"export PERF_UPLOAD_POLICY='always'; "
|
||||
"elif [ -n \"${BUILDKITE_PULL_REQUEST:-}\" ] && [ \"${BUILDKITE_PULL_REQUEST:-false}\" != 'false' ]; then "
|
||||
"if [ -n \"${BUILDKITE_PULL_REQUEST:-}\" ] && [ \"${BUILDKITE_PULL_REQUEST:-false}\" != 'false' ]; then "
|
||||
"export PERF_RUN_SOURCE='pr'; "
|
||||
"export PERF_UPLOAD_POLICY='pass'; "
|
||||
"elif [ \"${BUILDKITE_BRANCH:-}\" = 'main' ] && ( [ \"${BUILDKITE_SOURCE:-}\" = 'schedule' ] || [ \"${TEST_SCOPE:-}\" = 'full' ] ); then "
|
||||
"export PERF_RUN_SOURCE='scheduled_main'; "
|
||||
"export PERF_UPLOAD_POLICY='always'; "
|
||||
"elif [ \"${TEST_SCOPE:-}\" = 'direct' ]; then "
|
||||
"export PERF_RUN_SOURCE='unknown'; "
|
||||
"export PERF_UPLOAD_POLICY='pass'; "
|
||||
|
||||
@@ -23,7 +23,7 @@ from typing import Any
|
||||
|
||||
try:
|
||||
from fastvideo.performance.hf_store import (
|
||||
load_records_for_model,
|
||||
load_records_for_identity,
|
||||
safe_float,
|
||||
sanitize,
|
||||
sync_from_hf,
|
||||
@@ -40,7 +40,7 @@ except ImportError:
|
||||
if repo_root not in sys.path:
|
||||
sys.path.insert(0, repo_root)
|
||||
from fastvideo.performance.hf_store import (
|
||||
load_records_for_model,
|
||||
load_records_for_identity,
|
||||
safe_float,
|
||||
sanitize,
|
||||
sync_from_hf,
|
||||
@@ -89,6 +89,20 @@ COMPARISON_IDENTITY_KEYS = (
|
||||
"hardware_profile_id",
|
||||
"software_profile_id",
|
||||
)
|
||||
STATUS_PASS = "PASS"
|
||||
STATUS_REGRESSION = "REGRESSION"
|
||||
STATUS_CALIBRATION_NEEDED = "CALIBRATION_NEEDED"
|
||||
STATUS_RECIPE_MISMATCH = "RECIPE_MISMATCH"
|
||||
STATUS_INFRA_ERROR = "INFRA_ERROR"
|
||||
# Reserved for the promoted-baseline workflow; never emitted here.
|
||||
STATUS_QUALITY_BLOCKED = "QUALITY_BLOCKED"
|
||||
STATIC_THRESHOLD_FIELDS = (
|
||||
("avg_generation_time_s", "max_generation_time_s"),
|
||||
("max_peak_memory_mb", "max_peak_memory_mb"),
|
||||
("text_encoder_time_s", "max_text_encoder_time_s"),
|
||||
("dit_time_s", "max_dit_time_s"),
|
||||
("vae_decode_time_s", "max_vae_decode_time_s"),
|
||||
)
|
||||
|
||||
|
||||
def _should_persist_tracking() -> bool:
|
||||
@@ -122,7 +136,13 @@ def _detect_run_source() -> str:
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _is_baseline_eligible(run_source: str, success: bool) -> bool:
|
||||
def _is_baseline_eligible(
|
||||
run_source: str,
|
||||
success: bool,
|
||||
comparison_status: str | None = None,
|
||||
) -> bool:
|
||||
if comparison_status is not None and comparison_status != STATUS_PASS:
|
||||
return False
|
||||
return run_source == "scheduled_main" and success
|
||||
|
||||
|
||||
@@ -135,7 +155,7 @@ def _upload_allowed(record: dict[str, Any]) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _result_failed_static_thresholds() -> bool:
|
||||
def _performance_pytest_failed() -> bool:
|
||||
value = os.environ.get("PERF_PYTEST_RC", "")
|
||||
if not value:
|
||||
return False
|
||||
@@ -193,18 +213,52 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
|
||||
}
|
||||
|
||||
|
||||
def _record_uses_v2_identity(record: dict[str, Any]) -> bool:
|
||||
schema_version = record.get("result_schema_version")
|
||||
return str(schema_version) == "2" or any(key in record for key in COMPARISON_IDENTITY_KEYS)
|
||||
|
||||
|
||||
def _recipe_cohort_filters(record: dict[str, Any]) -> dict[str, str]:
|
||||
identity = _comparison_identity_filters(record)
|
||||
return {
|
||||
key: identity[key]
|
||||
for key in ("workload_id", "variant_id", "benchmark_version")
|
||||
}
|
||||
|
||||
|
||||
def _comparison_identity_filters_or_none(record: dict[str, Any]) -> dict[str, str] | None:
|
||||
try:
|
||||
return _comparison_identity_filters(record)
|
||||
except ValueError as exc:
|
||||
print("Skipping rolling baseline comparison for "
|
||||
f"{record.get('model_id', 'unknown')} on "
|
||||
f"{record.get('gpu_type', 'unknown')}: {exc}. "
|
||||
"The normalized artifact will still be written, but this record "
|
||||
"will not be baseline eligible.")
|
||||
if _record_uses_v2_identity(record):
|
||||
print(f"Invalid v2 comparison identity for {record.get('model_id', 'unknown')}: {exc}")
|
||||
else:
|
||||
print("Skipping rolling baseline comparison for "
|
||||
f"{record.get('model_id', 'unknown')} on "
|
||||
f"{record.get('gpu_type', 'unknown')}: {exc}. "
|
||||
"The normalized artifact will still be written, but this record "
|
||||
"will not be baseline eligible.")
|
||||
return None
|
||||
|
||||
|
||||
def _recipe_mismatch_records(
|
||||
record: dict[str, Any],
|
||||
cohort_records: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
current_recipe = str(record.get("recipe_fingerprint"))
|
||||
return [
|
||||
item for item in cohort_records
|
||||
if _record_can_authorize_recipe_mismatch(item)
|
||||
and item.get("recipe_fingerprint") is not None
|
||||
and str(item.get("recipe_fingerprint")).strip()
|
||||
and str(item.get("recipe_fingerprint")) != current_recipe
|
||||
]
|
||||
|
||||
|
||||
def _record_can_authorize_recipe_mismatch(record: dict[str, Any]) -> bool:
|
||||
return record.get("run_source") == "scheduled_main" or record.get("baseline_eligible") is True
|
||||
|
||||
|
||||
def _format_identity_filters(filters: dict[str, str]) -> str:
|
||||
if not filters:
|
||||
return ""
|
||||
@@ -343,12 +397,99 @@ def _check_regressions(
|
||||
return failures
|
||||
|
||||
|
||||
def _fixed_threshold_failures(raw_result: dict[str, Any]) -> list[str]:
|
||||
thresholds = raw_result.get("thresholds")
|
||||
if not isinstance(thresholds, dict):
|
||||
return []
|
||||
|
||||
failures = []
|
||||
for result_key, threshold_key in STATIC_THRESHOLD_FIELDS:
|
||||
current = safe_float(raw_result.get(result_key))
|
||||
threshold = safe_float(thresholds.get(threshold_key))
|
||||
if current is None or threshold is None or current <= threshold:
|
||||
continue
|
||||
failures.append(
|
||||
f"{raw_result.get('benchmark_id', 'unknown')} {result_key} exceeded fixed threshold "
|
||||
f"(current={current:.3f}, threshold={threshold:.3f})"
|
||||
)
|
||||
return failures
|
||||
|
||||
|
||||
def _pytest_infra_failure(record: dict[str, Any]) -> str:
|
||||
return (f"{record['model_id']} performance pytest failed without an "
|
||||
"attributable static-threshold regression "
|
||||
f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
|
||||
|
||||
|
||||
def _recipe_mismatch_failure(
|
||||
record: dict[str, Any],
|
||||
recipe_mismatch_records: list[dict[str, Any]],
|
||||
) -> str:
|
||||
seen = sorted({
|
||||
str(item.get("recipe_fingerprint"))
|
||||
for item in recipe_mismatch_records
|
||||
if item.get("recipe_fingerprint") is not None
|
||||
})
|
||||
suffix = f"; existing recipe_fingerprint values: {', '.join(seen)}" if seen else ""
|
||||
return (f"{record['model_id']} recipe_fingerprint={record.get('recipe_fingerprint')} "
|
||||
"does not match baseline records for the same workload, variant, "
|
||||
"and benchmark version across hardware and software profiles"
|
||||
f"{suffix}")
|
||||
|
||||
|
||||
def _evaluate_record_comparison(
|
||||
record: dict[str, Any],
|
||||
baseline_records: list[dict[str, Any]],
|
||||
recipe_mismatch_records: list[dict[str, Any]],
|
||||
metric_policies: tuple[MetricPolicy, ...],
|
||||
static_threshold_failures: list[str],
|
||||
unattributed_pytest_failure: bool,
|
||||
) -> tuple[list[str], str, str]:
|
||||
identity_filters = None
|
||||
if _record_uses_v2_identity(record):
|
||||
try:
|
||||
identity_filters = _comparison_identity_filters(record)
|
||||
except ValueError as exc:
|
||||
failure = f"{record['model_id']} cannot compare v2 record: {exc}"
|
||||
return [failure], STATUS_INFRA_ERROR, failure
|
||||
if not baseline_records and recipe_mismatch_records:
|
||||
failure = _recipe_mismatch_failure(record, recipe_mismatch_records)
|
||||
return [failure], STATUS_RECIPE_MISMATCH, failure
|
||||
|
||||
if static_threshold_failures:
|
||||
return static_threshold_failures, STATUS_REGRESSION, "; ".join(static_threshold_failures)
|
||||
|
||||
if unattributed_pytest_failure:
|
||||
failure = _pytest_infra_failure(record)
|
||||
return [failure], STATUS_INFRA_ERROR, failure
|
||||
|
||||
if not baseline_records:
|
||||
if identity_filters is not None:
|
||||
return [], STATUS_CALIBRATION_NEEDED, (
|
||||
"No baseline found for exact comparable identity"
|
||||
f"{_format_identity_filters(identity_filters)}"
|
||||
)
|
||||
|
||||
return [], STATUS_PASS, f"No legacy baseline for {record['model_id']} on {record['gpu_type']}"
|
||||
|
||||
failures = _check_regressions(record, baseline_records, metric_policies)
|
||||
if failures:
|
||||
return failures, STATUS_REGRESSION, "; ".join(failures)
|
||||
|
||||
return [], STATUS_PASS, "Comparable baseline found with no gated regressions"
|
||||
|
||||
|
||||
def _compact_value(value: float | None, precision: int = 3) -> str:
|
||||
if value is None:
|
||||
return "n/a"
|
||||
return f"{value:.{precision}f}"
|
||||
|
||||
|
||||
def _markdown_cell(value: Any) -> str:
|
||||
text = str(value) if value is not None else ""
|
||||
return text.replace("\n", " ").replace("|", "/")
|
||||
|
||||
|
||||
def _build_summary_row(
|
||||
record: dict[str, Any],
|
||||
baseline_records: list[dict[str, Any]],
|
||||
@@ -400,6 +541,8 @@ def _build_summary_row(
|
||||
"threshold_exceeded_metrics": threshold_exceeded_metrics,
|
||||
"failing_metrics": failing_metrics,
|
||||
"failed": has_failed,
|
||||
"status": record.get("comparison_status", STATUS_REGRESSION if has_failed else STATUS_PASS),
|
||||
"status_reason": record.get("comparison_status_reason", ""),
|
||||
}
|
||||
|
||||
|
||||
@@ -417,8 +560,8 @@ def _build_markdown_summary(
|
||||
"Throughput (curr/base) | Memory (curr/base) | "
|
||||
"Text Enc (curr/base) | DiT (curr/base) | "
|
||||
"VAE Decode (curr/base) | Worst Regression | Exceeded Metrics | "
|
||||
"Failing Metrics | Status |"),
|
||||
"|---|---|---:|---|---|---|---|---|---|---:|---|---|---|",
|
||||
"Failing Metrics | Status | Status Detail |"),
|
||||
"|---|---|---:|---|---|---|---|---|---|---:|---|---|---|---|",
|
||||
]
|
||||
|
||||
for row in summary_rows:
|
||||
@@ -435,12 +578,14 @@ def _build_markdown_summary(
|
||||
else "none"
|
||||
)
|
||||
failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none"
|
||||
status = "FAIL" if row["failed"] else "PASS"
|
||||
status = row["status"]
|
||||
status_detail = row["status_reason"]
|
||||
|
||||
lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
|
||||
f"{row['baseline_n']} | "
|
||||
f"{' | '.join(metric_cells)} | "
|
||||
f"{worst_reg} | {exceeded_metrics} | {failing_metrics} | {status} |")
|
||||
f"{worst_reg} | {exceeded_metrics} | {failing_metrics} | "
|
||||
f"{_markdown_cell(status)} | {_markdown_cell(status_detail)} |")
|
||||
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
@@ -470,7 +615,7 @@ def _emit_markdown_summary(markdown: str, commit_sha: str) -> None:
|
||||
def main() -> int:
|
||||
persist_tracking = _should_persist_tracking()
|
||||
upload_policy = _normalized_upload_policy()
|
||||
static_threshold_failed = _result_failed_static_thresholds()
|
||||
performance_pytest_failed = _performance_pytest_failed()
|
||||
|
||||
# Strict on upload-enabled runs: silent sync failure would make comparison
|
||||
# and upload state ambiguous.
|
||||
@@ -481,6 +626,9 @@ def main() -> int:
|
||||
print(f"No performance result files found in {RESULTS_DIR}")
|
||||
return 0
|
||||
|
||||
static_threshold_failures = [_fixed_threshold_failures(raw) for raw in current_results]
|
||||
unattributed_pytest_failure = performance_pytest_failed and not any(static_threshold_failures)
|
||||
|
||||
all_failures: list[str] = []
|
||||
summary_rows: list[dict[str, Any]] = []
|
||||
|
||||
@@ -489,61 +637,75 @@ def main() -> int:
|
||||
else:
|
||||
print("Tracking persistence disabled: PERF_UPLOAD_POLICY=never")
|
||||
|
||||
if static_threshold_failed:
|
||||
print(f"Static-threshold phase failed: PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
|
||||
if performance_pytest_failed:
|
||||
if unattributed_pytest_failure:
|
||||
print("Performance pytest failed without a measured static-threshold regression: "
|
||||
f"PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
|
||||
else:
|
||||
print("Performance pytest failure attributed to per-record static thresholds: "
|
||||
f"PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
|
||||
|
||||
for raw in current_results:
|
||||
for raw, fixed_threshold_failures in zip(current_results, static_threshold_failures, strict=True):
|
||||
record = _normalize_record(raw)
|
||||
metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
|
||||
identity_filters = _comparison_identity_filters_or_none(record)
|
||||
|
||||
baseline_records: list[dict[str, Any]] = []
|
||||
if identity_filters is None:
|
||||
# Records without the full v2 identity block skip rolling-baseline
|
||||
# comparison entirely: only the static thresholds gate them and
|
||||
# they never become baseline eligible.
|
||||
record["baseline_status"] = "skipped_missing_identity"
|
||||
failures: list[str] = []
|
||||
else:
|
||||
baseline_records = load_records_for_model(
|
||||
TRACKING_ROOT,
|
||||
record["model_id"],
|
||||
record["gpu_type"],
|
||||
**identity_filters,
|
||||
last_n=5,
|
||||
successful_only=True,
|
||||
baseline_eligible_only=True,
|
||||
)
|
||||
|
||||
if not baseline_records:
|
||||
# A brand-new cohort has NO regression gating until history
|
||||
# accumulates — make that loud and machine-readable instead of
|
||||
# an indistinguishable pass, so a cohort shift (intended or
|
||||
# accidental, e.g. an identity-field change) never silently
|
||||
# blinds the comparison.
|
||||
record["baseline_status"] = "initialized_new_cohort"
|
||||
print("=" * 72)
|
||||
print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
|
||||
f"{record['model_id']} on {record['gpu_type']}"
|
||||
f"{_format_identity_filters(identity_filters)}")
|
||||
print("Regression gating is INACTIVE for this cohort until "
|
||||
"baseline history accumulates. If this cohort shift is "
|
||||
"unexpected, check the identity fields above.")
|
||||
print("=" * 72)
|
||||
failures = []
|
||||
recipe_mismatch_records: list[dict[str, Any]] = []
|
||||
try:
|
||||
if identity_filters is None:
|
||||
record["baseline_status"] = (
|
||||
"invalid_identity" if _record_uses_v2_identity(record) else "skipped_missing_identity"
|
||||
)
|
||||
else:
|
||||
record["baseline_status"] = "compared"
|
||||
failures = _check_regressions(record, baseline_records, metric_policies)
|
||||
if static_threshold_failed:
|
||||
failures.append(f"{record['model_id']} fixed-threshold phase failed "
|
||||
f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
|
||||
baseline_records = load_records_for_identity(
|
||||
TRACKING_ROOT,
|
||||
identity_filters,
|
||||
last_n=5,
|
||||
successful_only=True,
|
||||
baseline_eligible_only=True,
|
||||
)
|
||||
if baseline_records:
|
||||
record["baseline_status"] = "compared"
|
||||
else:
|
||||
recipe_cohort_records = load_records_for_identity(
|
||||
TRACKING_ROOT,
|
||||
_recipe_cohort_filters(record),
|
||||
successful_only=True,
|
||||
)
|
||||
recipe_mismatch_records = _recipe_mismatch_records(record, recipe_cohort_records)
|
||||
record["baseline_status"] = (
|
||||
"recipe_mismatch" if recipe_mismatch_records else "initialized_new_cohort"
|
||||
)
|
||||
|
||||
failures, comparison_status, comparison_status_reason = _evaluate_record_comparison(
|
||||
record,
|
||||
baseline_records,
|
||||
recipe_mismatch_records,
|
||||
metric_policies,
|
||||
fixed_threshold_failures,
|
||||
unattributed_pytest_failure,
|
||||
)
|
||||
except Exception as exc:
|
||||
comparison_status = STATUS_INFRA_ERROR
|
||||
comparison_status_reason = f"{record['model_id']} baseline comparison hit an infra error: {exc}"
|
||||
failures = [comparison_status_reason]
|
||||
baseline_records = []
|
||||
record["baseline_status"] = "infra_error"
|
||||
|
||||
record["comparison_status"] = comparison_status
|
||||
record["comparison_status_reason"] = comparison_status_reason
|
||||
|
||||
record["success"] = not failures
|
||||
record["baseline_eligible"] = (
|
||||
identity_filters is not None and _is_baseline_eligible(record["run_source"], record["success"])
|
||||
identity_filters is not None
|
||||
and _is_baseline_eligible(record["run_source"], record["success"], comparison_status)
|
||||
)
|
||||
all_failures.extend(failures)
|
||||
|
||||
print(f"{record['model_id']} comparison status: "
|
||||
f"{comparison_status} - {comparison_status_reason}")
|
||||
|
||||
_write_normalized_artifact(record)
|
||||
|
||||
if _upload_allowed(record):
|
||||
|
||||
@@ -37,7 +37,7 @@ COMPARISON_COHORT_KEYS = (
|
||||
"hardware_profile_id",
|
||||
"software_profile_id",
|
||||
)
|
||||
GROUP_KEYS = ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)
|
||||
GROUP_KEYS = ("_cohort_schema", "_cohort_model_id", "_cohort_gpu_type", *COMPARISON_COHORT_KEYS)
|
||||
|
||||
|
||||
def _cohort_value(value: object) -> str:
|
||||
@@ -62,35 +62,43 @@ def _short_value(value: object) -> str:
|
||||
return text[:12]
|
||||
|
||||
|
||||
def _cohort_title(group_key: tuple[object, ...]) -> str:
|
||||
workload = _display_value(group_key[2])
|
||||
variant = _display_value(group_key[3])
|
||||
version = _display_value(group_key[4])
|
||||
def _cohort_title(record: object) -> str:
|
||||
workload = _display_value(record.get("workload_id"))
|
||||
variant = _display_value(record.get("variant_id"))
|
||||
version = _display_value(record.get("benchmark_version"))
|
||||
version_label = version if version == "legacy" else f"v{version}"
|
||||
return f"{workload} / {variant} / {version_label}"
|
||||
|
||||
|
||||
def _cohort_detail(group_key: tuple[object, ...]) -> str:
|
||||
def _cohort_detail(record: object) -> str:
|
||||
return " | ".join((
|
||||
f"recipe {_short_value(group_key[5])}",
|
||||
_short_value(group_key[6]),
|
||||
_short_value(group_key[7]),
|
||||
f"recipe {_short_value(record.get('recipe_fingerprint'))}",
|
||||
_short_value(record.get("hardware_profile_id")),
|
||||
_short_value(record.get("software_profile_id")),
|
||||
))
|
||||
|
||||
|
||||
def _group_record(group_key: tuple[object, ...]) -> dict[str, str]:
|
||||
def _group_record(record: object) -> dict[str, str]:
|
||||
return {
|
||||
key: _cohort_value(value)
|
||||
for key, value in zip(GROUP_KEYS, group_key)
|
||||
key: _cohort_value(record.get(key))
|
||||
for key in ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)
|
||||
}
|
||||
|
||||
|
||||
def _has_complete_v2_identity(record: object) -> bool:
|
||||
return all(_cohort_value(record.get(key)) for key in COMPARISON_COHORT_KEYS)
|
||||
|
||||
|
||||
def _dashboard_frame(df: pd.DataFrame) -> pd.DataFrame:
|
||||
dashboard_df = df.copy()
|
||||
for key in GROUP_KEYS:
|
||||
for key in ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS):
|
||||
if key not in dashboard_df.columns:
|
||||
dashboard_df[key] = ""
|
||||
dashboard_df[key] = dashboard_df[key].map(_cohort_value)
|
||||
uses_v2_identity = dashboard_df.apply(_has_complete_v2_identity, axis=1)
|
||||
dashboard_df["_cohort_schema"] = uses_v2_identity.map({True: "v2", False: "legacy"})
|
||||
dashboard_df["_cohort_model_id"] = dashboard_df["model_id"].where(~uses_v2_identity, "")
|
||||
dashboard_df["_cohort_gpu_type"] = dashboard_df["gpu_type"].where(~uses_v2_identity, "")
|
||||
return dashboard_df
|
||||
|
||||
# -----------------------------
|
||||
@@ -107,12 +115,14 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
|
||||
figs = []
|
||||
skipped_metrics: list[dict[str, object]] = []
|
||||
|
||||
for group_key, g in group_data(df):
|
||||
model_id, gpu_type = group_key[:2]
|
||||
group_record = _group_record(group_key)
|
||||
cohort_title = _cohort_title(group_key)
|
||||
cohort_detail = _cohort_detail(group_key)
|
||||
for _group_key, g in group_data(df):
|
||||
g = g.sort_values("timestamp")
|
||||
display_record = g.iloc[-1]
|
||||
model_id = display_record["model_id"]
|
||||
gpu_type = display_record["gpu_type"]
|
||||
group_record = _group_record(display_record)
|
||||
cohort_title = _cohort_title(display_record)
|
||||
cohort_detail = _cohort_detail(display_record)
|
||||
|
||||
# One chart per metric so the y-axes aren't on wildly different scales
|
||||
for metric in METRICS:
|
||||
@@ -146,7 +156,7 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
|
||||
x="timestamp",
|
||||
y=metric,
|
||||
markers=True,
|
||||
hover_data=["config_id", "commit_sha", *COMPARISON_COHORT_KEYS],
|
||||
hover_data=["model_id", "gpu_type", "config_id", "commit_sha", *COMPARISON_COHORT_KEYS],
|
||||
title=f"{model_id} | {gpu_type} | {cohort_title} | {cohort_detail} | {metric}",
|
||||
labels={"timestamp": "Time", metric: metric},
|
||||
)
|
||||
|
||||
@@ -17,7 +17,7 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
RECIPE_SCHEMA_VERSION = 1
|
||||
RECIPE_SCHEMA_VERSION = 2
|
||||
PROFILE_ID_LENGTH = 16
|
||||
_PATH_EXCLUDED_GENERATION_KEYS = {"output_path", "output_video_name"}
|
||||
_PROMPT_KEYS = {"prompt", "negative_prompt", "neg_prompt"}
|
||||
@@ -32,6 +32,7 @@ _PROFILE_ENV_VARS = (
|
||||
"FASTVIDEO_ATTENTION_BACKEND",
|
||||
"IMAGE_VERSION",
|
||||
"UV_TORCH_BACKEND",
|
||||
"FASTVIDEO_CONTAINER_IMAGE_REF",
|
||||
)
|
||||
_PACKAGE_DISTRIBUTIONS = {
|
||||
"fastvideo_kernel": ("fastvideo-kernel", "fastvideo_kernel"),
|
||||
@@ -87,6 +88,12 @@ def recipe_fingerprint(recipe: Mapping[str, Any]) -> str:
|
||||
# Intentional backend changes are declared via requested_backend,
|
||||
# which stays in the hash.
|
||||
attention.pop("resolved_backend", None)
|
||||
benchmark = pruned.get("benchmark")
|
||||
if isinstance(benchmark, dict):
|
||||
# benchmark_id names artifacts and storage directories, but is not one
|
||||
# of the six v2 comparison fields. Keep it in the recipe for audit
|
||||
# while allowing display-only renames to stay in the same cohort.
|
||||
benchmark.pop("benchmark_id", None)
|
||||
return sha256_hexdigest(canonical_json(pruned))
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,760 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Prepare reviewed v2 calibration artifacts as baseline seed records."""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import tempfile
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from huggingface_hub import CommitOperationAdd, HfApi
|
||||
from huggingface_hub.constants import ENDPOINT
|
||||
|
||||
try:
|
||||
from fastvideo.tests.performance.compare_baseline import (
|
||||
STATUS_CALIBRATION_NEEDED,
|
||||
STATUS_PASS,
|
||||
_comparison_identity_filters,
|
||||
_recipe_cohort_filters,
|
||||
_recipe_mismatch_records,
|
||||
_record_uses_v2_identity,
|
||||
)
|
||||
from fastvideo.performance.hf_store import (
|
||||
HF_REPO_ID,
|
||||
load_records_for_identity,
|
||||
resolve_hf_token,
|
||||
safe_float,
|
||||
sanitize,
|
||||
sync_from_hf,
|
||||
)
|
||||
from fastvideo.performance.metric_policy import DEFAULT_METRIC_POLICIES
|
||||
except ImportError:
|
||||
repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
|
||||
if repo_root not in sys.path:
|
||||
sys.path.insert(0, repo_root)
|
||||
from fastvideo.tests.performance.compare_baseline import (
|
||||
STATUS_CALIBRATION_NEEDED,
|
||||
STATUS_PASS,
|
||||
_comparison_identity_filters,
|
||||
_recipe_cohort_filters,
|
||||
_recipe_mismatch_records,
|
||||
_record_uses_v2_identity,
|
||||
)
|
||||
from fastvideo.performance.hf_store import (
|
||||
HF_REPO_ID,
|
||||
load_records_for_identity,
|
||||
resolve_hf_token,
|
||||
safe_float,
|
||||
sanitize,
|
||||
sync_from_hf,
|
||||
)
|
||||
from fastvideo.performance.metric_policy import DEFAULT_METRIC_POLICIES
|
||||
|
||||
TRACKING_ROOT = os.environ.get("PERFORMANCE_TRACKING_ROOT", "/tmp/perf-tracking")
|
||||
STAGING_ROOT = os.environ.get("PERFORMANCE_RESEED_STAGING_ROOT", "/tmp/performance_reseed_prepared")
|
||||
DEFAULT_MAX_INTRA_BATCH_REGRESSION = 0.05
|
||||
_CORE_MEASUREMENT_FIELDS = frozenset({"latency", "throughput", "memory"})
|
||||
|
||||
|
||||
def _now_utc_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _load_json(path: str) -> dict[str, Any]:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _write_json(path: str, payload: dict[str, Any]) -> None:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, indent=2)
|
||||
|
||||
|
||||
def _file_sha256(path: str) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _source_identity(record: dict[str, Any]) -> dict[str, str]:
|
||||
if not _record_uses_v2_identity(record):
|
||||
raise ValueError("baseline seeds require a v2 record with exact comparable identity fields")
|
||||
return _comparison_identity_filters(record)
|
||||
|
||||
|
||||
def _require_nonempty_string(record: dict[str, Any], field: str) -> str:
|
||||
value = record.get(field)
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"baseline seed records require a non-empty string {field}")
|
||||
return value
|
||||
|
||||
|
||||
def _source_provenance_identity(record: dict[str, Any]) -> tuple[str, ...] | None:
|
||||
job_id = str(record.get("job_id") or "")
|
||||
if job_id:
|
||||
return ("job_id", job_id)
|
||||
build_id = str(record.get("build_id") or "")
|
||||
if build_id:
|
||||
return ("build_id", build_id)
|
||||
commit_sha = str(record.get("commit_sha") or "")
|
||||
timestamp = str(record.get("timestamp") or "")
|
||||
if commit_sha or timestamp:
|
||||
return ("commit_timestamp", commit_sha, timestamp)
|
||||
return None
|
||||
|
||||
|
||||
def _validate_unique_sources(source_paths: list[str], records: list[dict[str, Any]]) -> None:
|
||||
seen_paths: set[str] = set()
|
||||
seen_contents: set[str] = set()
|
||||
seen_provenance: set[tuple[str, ...]] = set()
|
||||
|
||||
for index, (source_path, record) in enumerate(zip(source_paths, records), start=1):
|
||||
resolved_path = os.path.realpath(os.path.abspath(source_path))
|
||||
if resolved_path in seen_paths:
|
||||
raise ValueError(f"source artifact {index} duplicates a resolved source path")
|
||||
seen_paths.add(resolved_path)
|
||||
|
||||
content_identity = json.dumps(record, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
|
||||
if content_identity in seen_contents:
|
||||
raise ValueError(f"source artifact {index} duplicates source content")
|
||||
seen_contents.add(content_identity)
|
||||
|
||||
provenance_identity = _source_provenance_identity(record)
|
||||
if provenance_identity is not None:
|
||||
if provenance_identity in seen_provenance:
|
||||
raise ValueError(f"source artifact {index} duplicates source provenance")
|
||||
seen_provenance.add(provenance_identity)
|
||||
|
||||
|
||||
def _truthy_pr_number(value: Any) -> bool:
|
||||
return bool(value and str(value) not in {"false", "0", "None", "none"})
|
||||
|
||||
|
||||
def _validate_source_measurements(record: dict[str, Any]) -> None:
|
||||
for policy in DEFAULT_METRIC_POLICIES:
|
||||
raw_value = record.get(policy.key)
|
||||
if raw_value is None:
|
||||
if policy.key in _CORE_MEASUREMENT_FIELDS:
|
||||
raise ValueError(f"baseline seed records require a finite positive {policy.key} measurement")
|
||||
continue
|
||||
|
||||
value = None if isinstance(raw_value, bool) else safe_float(raw_value)
|
||||
if value is None or not math.isfinite(value):
|
||||
raise ValueError(f"baseline seed records require a finite {policy.key} measurement")
|
||||
if policy.key in _CORE_MEASUREMENT_FIELDS:
|
||||
if value <= 0:
|
||||
raise ValueError(f"baseline seed records require a positive {policy.key} measurement")
|
||||
elif value < 0:
|
||||
raise ValueError(f"baseline seed records require a non-negative {policy.key} measurement")
|
||||
|
||||
|
||||
def _validate_calibration_source(record: dict[str, Any]) -> None:
|
||||
_require_nonempty_string(record, "model_id")
|
||||
_source_identity(record)
|
||||
_validate_source_measurements(record)
|
||||
if record.get("result_schema_version") != 2:
|
||||
raise ValueError("baseline seeds require a normalized v2 source artifact")
|
||||
if record.get("comparison_status") != STATUS_CALIBRATION_NEEDED:
|
||||
raise ValueError("baseline seeds require a CALIBRATION_NEEDED source artifact")
|
||||
if record.get("success") is not True:
|
||||
raise ValueError("baseline seeds require a successful source artifact")
|
||||
if record.get("baseline_eligible") is not False:
|
||||
raise ValueError("baseline seeds require a baseline_eligible=false source artifact")
|
||||
if record.get("run_source") != "scheduled_main":
|
||||
raise ValueError("baseline seeds require a scheduled_main source artifact")
|
||||
if _truthy_pr_number(record.get("pr_number")):
|
||||
raise ValueError("baseline seeds require a non-PR source artifact")
|
||||
if record.get("branch") != "main" or record.get("test_scope") != "full":
|
||||
raise ValueError("baseline seeds require a main-branch full-suite source artifact")
|
||||
|
||||
|
||||
def build_baseline_seed_record(
|
||||
source_record: dict[str, Any],
|
||||
*,
|
||||
reason: str,
|
||||
source_result: str | None = None,
|
||||
operator: str | None = None,
|
||||
timestamp: str | None = None,
|
||||
batch_size: int = 1,
|
||||
batch_index: int = 1,
|
||||
) -> dict[str, Any]:
|
||||
"""Return an approved v2 baseline seed derived from a calibration artifact."""
|
||||
if not reason.strip():
|
||||
raise ValueError("baseline seed reason must be a non-empty string")
|
||||
_validate_calibration_source(source_record)
|
||||
|
||||
seed = dict(source_record)
|
||||
seed.update({
|
||||
"timestamp": timestamp or _now_utc_iso(),
|
||||
"success": True,
|
||||
"baseline_eligible": True,
|
||||
"comparison_status": STATUS_PASS,
|
||||
"comparison_status_reason": "Approved baseline seed from CALIBRATION_NEEDED source artifact",
|
||||
"baseline_seed": True,
|
||||
"baseline_seed_reason": reason,
|
||||
"baseline_seed_source_status": source_record.get("comparison_status"),
|
||||
"baseline_seed_source_timestamp": source_record.get("timestamp"),
|
||||
"baseline_seed_source_success": source_record.get("success"),
|
||||
"baseline_seed_source_run_source": source_record.get("run_source"),
|
||||
"baseline_seed_source_branch": source_record.get("branch"),
|
||||
"baseline_seed_source_test_scope": source_record.get("test_scope"),
|
||||
"baseline_seed_source_pr_number": source_record.get("pr_number"),
|
||||
"baseline_seed_batch_size": batch_size,
|
||||
"baseline_seed_batch_index": batch_index,
|
||||
})
|
||||
if source_result:
|
||||
seed["baseline_seed_source_result"] = source_result
|
||||
if operator:
|
||||
seed["baseline_seed_operator"] = operator
|
||||
return seed
|
||||
|
||||
|
||||
def write_seed_record(
|
||||
local_dir: str,
|
||||
record: dict[str, Any],
|
||||
*,
|
||||
suffix: str | None = None,
|
||||
) -> str:
|
||||
"""Write *record* under the staging root and return the local JSON path."""
|
||||
out_path = _seed_record_path(local_dir, record, suffix=suffix)
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
_write_json(out_path, record)
|
||||
return out_path
|
||||
|
||||
|
||||
def _seed_record_path(
|
||||
local_dir: str,
|
||||
record: dict[str, Any],
|
||||
*,
|
||||
suffix: str | None = None,
|
||||
) -> str:
|
||||
model_id = _require_nonempty_string(record, "model_id")
|
||||
timestamp = _require_nonempty_string(record, "timestamp")
|
||||
model_dir = os.path.join(local_dir, sanitize(model_id))
|
||||
commit = sanitize(record.get("commit_sha") or "unknown")
|
||||
suffix_part = f"_{sanitize(suffix)}" if suffix else ""
|
||||
return os.path.join(model_dir, f"{sanitize(timestamp)}_{commit}_seed{suffix_part}.json")
|
||||
|
||||
|
||||
def _validate_same_identity(records: list[dict[str, Any]]) -> dict[str, str]:
|
||||
first = _source_identity(records[0])
|
||||
for index, record in enumerate(records[1:], start=2):
|
||||
identity = _source_identity(record)
|
||||
if identity != first:
|
||||
raise ValueError(f"source artifact {index} does not match the first exact comparable identity")
|
||||
return first
|
||||
|
||||
|
||||
def _validate_separate_roots(tracking_root: str, staging_root: str) -> None:
|
||||
tracking_root = os.path.realpath(os.path.abspath(tracking_root))
|
||||
staging_root = os.path.realpath(os.path.abspath(staging_root))
|
||||
try:
|
||||
common_root = os.path.commonpath((tracking_root, staging_root))
|
||||
except ValueError:
|
||||
return
|
||||
if common_root in {tracking_root, staging_root}:
|
||||
raise ValueError("baseline seed staging and canonical tracking roots must be separate and non-nested")
|
||||
|
||||
|
||||
def _validate_no_existing_seed(
|
||||
records_root: str,
|
||||
record: dict[str, Any],
|
||||
*,
|
||||
source: str,
|
||||
) -> None:
|
||||
identity = _source_identity(record)
|
||||
if load_records_for_identity(
|
||||
records_root,
|
||||
identity,
|
||||
last_n=1,
|
||||
successful_only=True,
|
||||
baseline_eligible_only=True,
|
||||
):
|
||||
raise ValueError(
|
||||
f"{source} already has a baseline-eligible record for the exact comparable identity; "
|
||||
"the CALIBRATION_NEEDED source artifact is stale"
|
||||
)
|
||||
|
||||
cohort_records = load_records_for_identity(
|
||||
records_root,
|
||||
_recipe_cohort_filters(record),
|
||||
successful_only=True,
|
||||
)
|
||||
mismatches = _recipe_mismatch_records(record, cohort_records)
|
||||
if mismatches:
|
||||
recipes = sorted({str(item["recipe_fingerprint"]) for item in mismatches})
|
||||
raise ValueError(
|
||||
f"{source} already has trusted records for another recipe in this workload, "
|
||||
f"variant, and benchmark version: {', '.join(recipes)}"
|
||||
)
|
||||
|
||||
|
||||
def _repository_descriptor() -> dict[str, str]:
|
||||
return {
|
||||
"endpoint": ENDPOINT,
|
||||
"repo_id": HF_REPO_ID,
|
||||
"repo_type": "dataset",
|
||||
}
|
||||
|
||||
|
||||
def _reservation_path(staging_root: str, identity: dict[str, str]) -> str:
|
||||
payload = json.dumps({
|
||||
"identity": identity,
|
||||
"repository": _repository_descriptor(),
|
||||
}, sort_keys=True, separators=(",", ":"))
|
||||
digest = hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||||
return os.path.join(staging_root, ".seed-reservations", digest)
|
||||
|
||||
|
||||
def _reserve_staging_identity(staging_root: str, identity: dict[str, str]) -> str:
|
||||
reservation = _reservation_path(staging_root, identity)
|
||||
os.makedirs(os.path.dirname(reservation), exist_ok=True)
|
||||
try:
|
||||
os.mkdir(reservation, 0o700)
|
||||
except FileExistsError as exc:
|
||||
raise ValueError(
|
||||
"staging already has a reservation for the exact comparable identity; "
|
||||
"reuse or explicitly clean the existing preparation"
|
||||
) from exc
|
||||
return reservation
|
||||
|
||||
|
||||
def _release_failed_reservation(reservation: str, prepared_paths: list[str]) -> None:
|
||||
reservation = os.path.realpath(os.path.abspath(reservation))
|
||||
prepared_dirs: set[str] = set()
|
||||
for path in prepared_paths:
|
||||
path = os.path.realpath(os.path.abspath(path))
|
||||
if os.path.commonpath((reservation, path)) != reservation:
|
||||
continue
|
||||
try:
|
||||
os.unlink(path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
prepared_dirs.add(os.path.dirname(path))
|
||||
for name in ("manifest", "manifest.tmp"):
|
||||
try:
|
||||
os.unlink(os.path.join(reservation, name))
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
for directory in sorted(prepared_dirs, key=len, reverse=True):
|
||||
try:
|
||||
os.rmdir(directory)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
os.rmdir(reservation)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _write_preparation_manifest(
|
||||
reservation: str,
|
||||
staging_root: str,
|
||||
identity: dict[str, str],
|
||||
source_paths: list[str],
|
||||
source_records: list[dict[str, Any]],
|
||||
prepared_paths: list[str],
|
||||
) -> str:
|
||||
manifest_path = os.path.join(reservation, "manifest")
|
||||
temp_path = f"{manifest_path}.tmp"
|
||||
source_manifest_entries = []
|
||||
for path, expected_record in zip(source_paths, source_records, strict=True):
|
||||
with open(path, "rb") as f:
|
||||
blob = f.read()
|
||||
if json.loads(blob) != expected_record:
|
||||
raise ValueError(f"source record changed during preparation: {path}")
|
||||
source_manifest_entries.append({
|
||||
"path": os.path.realpath(os.path.abspath(path)),
|
||||
"sha256": hashlib.sha256(blob).hexdigest(),
|
||||
})
|
||||
manifest = {
|
||||
"manifest_schema_version": 1,
|
||||
"repository": _repository_descriptor(),
|
||||
"identity": identity,
|
||||
"recipe_cohort": {
|
||||
key: identity[key]
|
||||
for key in ("workload_id", "variant_id", "benchmark_version")
|
||||
},
|
||||
"staging_root": os.path.realpath(os.path.abspath(staging_root)),
|
||||
"source_records": source_manifest_entries,
|
||||
"prepared_records": [{
|
||||
"path": os.path.realpath(os.path.abspath(path)),
|
||||
"sha256": _file_sha256(path),
|
||||
} for path in prepared_paths],
|
||||
}
|
||||
_write_json(temp_path, manifest)
|
||||
os.replace(temp_path, manifest_path)
|
||||
return manifest_path
|
||||
|
||||
|
||||
def _validate_prepared_seed(record: dict[str, Any]) -> None:
|
||||
_source_identity(record)
|
||||
_require_nonempty_string(record, "model_id")
|
||||
_validate_source_measurements(record)
|
||||
if record.get("baseline_seed") is not True:
|
||||
raise ValueError("prepared records must have baseline_seed=true")
|
||||
if record.get("success") is not True or record.get("baseline_eligible") is not True:
|
||||
raise ValueError("prepared records must be successful and baseline eligible")
|
||||
if record.get("comparison_status") != STATUS_PASS:
|
||||
raise ValueError("prepared records must have comparison_status=PASS")
|
||||
if record.get("baseline_seed_source_status") != STATUS_CALIBRATION_NEEDED:
|
||||
raise ValueError("prepared records must retain CALIBRATION_NEEDED source provenance")
|
||||
if not str(record.get("baseline_seed_reason") or "").strip():
|
||||
raise ValueError("prepared records must retain a non-empty approval rationale")
|
||||
if (
|
||||
record.get("baseline_seed_source_success") is not True
|
||||
or record.get("baseline_seed_source_run_source") != "scheduled_main"
|
||||
or record.get("baseline_seed_source_branch") != "main"
|
||||
or record.get("baseline_seed_source_test_scope") != "full"
|
||||
or _truthy_pr_number(record.get("baseline_seed_source_pr_number"))
|
||||
):
|
||||
raise ValueError("prepared records must retain trusted scheduled-main source provenance")
|
||||
|
||||
|
||||
def upload_prepared_seed_manifest(manifest_path: str) -> str:
|
||||
"""Conditionally upload one reviewed manifest in a single Hub commit."""
|
||||
manifest_path = os.path.realpath(os.path.abspath(manifest_path))
|
||||
manifest = _load_json(manifest_path)
|
||||
if manifest.get("manifest_schema_version") != 1:
|
||||
raise ValueError("unsupported or missing prepared-seed manifest schema")
|
||||
expected_repository = _repository_descriptor()
|
||||
if manifest.get("repository") != expected_repository:
|
||||
raise ValueError("prepared-seed manifest repository does not match the configured Hub destination")
|
||||
|
||||
manifest_staging_root = manifest.get("staging_root")
|
||||
if not isinstance(manifest_staging_root, str) or not manifest_staging_root:
|
||||
raise ValueError("prepared-seed manifest is missing its staging root")
|
||||
staging_root = os.path.realpath(os.path.abspath(manifest_staging_root))
|
||||
source_entries = manifest.get("source_records")
|
||||
if not isinstance(source_entries, list) or not source_entries:
|
||||
raise ValueError("prepared-seed manifest must retain its reviewed source records")
|
||||
source_paths: list[str] = []
|
||||
source_records: list[dict[str, Any]] = []
|
||||
for entry in source_entries:
|
||||
if not isinstance(entry, dict):
|
||||
raise ValueError("prepared-seed manifest source entries must be objects")
|
||||
path = os.path.realpath(os.path.abspath(str(entry.get("path") or "")))
|
||||
with open(path, "rb") as f:
|
||||
blob = f.read()
|
||||
if hashlib.sha256(blob).hexdigest() != entry.get("sha256"):
|
||||
raise ValueError(f"source record changed after review: {path}")
|
||||
record = json.loads(blob)
|
||||
if not isinstance(record, dict):
|
||||
raise ValueError(f"source record must contain a JSON object: {path}")
|
||||
_validate_calibration_source(record)
|
||||
source_paths.append(path)
|
||||
source_records.append(record)
|
||||
_validate_unique_sources(source_paths, source_records)
|
||||
source_identity = _validate_same_identity(source_records)
|
||||
reservation = _reservation_path(staging_root, source_identity)
|
||||
if os.path.dirname(manifest_path) != reservation:
|
||||
raise ValueError("prepared-seed manifest is outside its identity reservation")
|
||||
|
||||
prepared_entries = manifest.get("prepared_records")
|
||||
if not isinstance(prepared_entries, list) or not prepared_entries:
|
||||
raise ValueError("prepared-seed manifest must contain at least one record")
|
||||
|
||||
prepared_paths: list[str] = []
|
||||
prepared_blobs: list[bytes] = []
|
||||
records: list[dict[str, Any]] = []
|
||||
for entry in prepared_entries:
|
||||
if not isinstance(entry, dict):
|
||||
raise ValueError("prepared-seed manifest record entries must be objects")
|
||||
path = os.path.realpath(os.path.abspath(str(entry.get("path") or "")))
|
||||
if os.path.commonpath((reservation, path)) != reservation:
|
||||
raise ValueError(f"prepared record is outside its identity reservation: {path}")
|
||||
with open(path, "rb") as f:
|
||||
blob = f.read()
|
||||
if hashlib.sha256(blob).hexdigest() != entry.get("sha256"):
|
||||
raise ValueError(f"prepared record changed after review: {path}")
|
||||
record = json.loads(blob)
|
||||
if not isinstance(record, dict):
|
||||
raise ValueError(f"prepared record must contain a JSON object: {path}")
|
||||
prepared_paths.append(path)
|
||||
prepared_blobs.append(blob)
|
||||
records.append(record)
|
||||
|
||||
for record in records:
|
||||
_validate_prepared_seed(record)
|
||||
if len(set(prepared_paths)) != len(prepared_paths):
|
||||
raise ValueError("prepared-seed manifest repeats a prepared record path")
|
||||
identity = _validate_same_identity(records)
|
||||
if source_identity != identity:
|
||||
raise ValueError("prepared records do not match their reviewed source identity")
|
||||
if manifest.get("identity") != identity:
|
||||
raise ValueError("prepared records no longer match the manifest identity")
|
||||
if manifest.get("recipe_cohort") != _recipe_cohort_filters(records[0]):
|
||||
raise ValueError("prepared records no longer match the manifest recipe cohort")
|
||||
|
||||
batch_size = len(records)
|
||||
if len(source_records) != batch_size:
|
||||
raise ValueError("prepared-seed manifest source and prepared batch sizes differ")
|
||||
if len({str(record.get("baseline_seed_reason")) for record in records}) != 1:
|
||||
raise ValueError("prepared records have inconsistent approval rationales")
|
||||
raw_batch_indices = [record.get("baseline_seed_batch_index") for record in records]
|
||||
if any(record.get("baseline_seed_batch_size") != batch_size for record in records):
|
||||
raise ValueError("prepared records have inconsistent baseline seed batch sizes")
|
||||
if any(isinstance(index, bool) or not isinstance(index, int) for index in raw_batch_indices):
|
||||
raise ValueError("prepared records have invalid baseline seed batch indices")
|
||||
batch_indices = sorted(raw_batch_indices)
|
||||
if batch_indices != list(range(1, batch_size + 1)):
|
||||
raise ValueError("prepared records have inconsistent baseline seed batch indices")
|
||||
for path, record in zip(prepared_paths, records, strict=True):
|
||||
index = record["baseline_seed_batch_index"]
|
||||
suffix = f"{index:02d}" if batch_size > 1 else None
|
||||
if path != _seed_record_path(reservation, record, suffix=suffix):
|
||||
raise ValueError(f"prepared record path does not match its reserved identity: {path}")
|
||||
|
||||
ordered_sources = _order_sources_by_timestamp(source_paths, source_records)
|
||||
records_by_index = sorted(records, key=lambda record: record["baseline_seed_batch_index"])
|
||||
for index, ((source_path, source_record), record) in enumerate(
|
||||
zip(ordered_sources, records_by_index, strict=True),
|
||||
start=1,
|
||||
):
|
||||
expected = build_baseline_seed_record(
|
||||
source_record,
|
||||
reason=str(record["baseline_seed_reason"]),
|
||||
source_result=source_path,
|
||||
operator=record.get("baseline_seed_operator"),
|
||||
timestamp=str(record["timestamp"]),
|
||||
batch_size=batch_size,
|
||||
batch_index=index,
|
||||
)
|
||||
if record != expected:
|
||||
raise ValueError(f"prepared record {index} no longer matches its reviewed source")
|
||||
|
||||
token = resolve_hf_token()
|
||||
if not token:
|
||||
raise RuntimeError("a Hugging Face write token is required to upload baseline seeds")
|
||||
api = HfApi(token=token)
|
||||
repo_info = api.repo_info(repo_id=HF_REPO_ID, repo_type="dataset", revision="main")
|
||||
parent_commit = getattr(repo_info, "sha", None)
|
||||
if not parent_commit:
|
||||
raise RuntimeError("could not resolve the current performance-tracking repository revision")
|
||||
|
||||
operations = []
|
||||
with tempfile.TemporaryDirectory(prefix="performance-seed-remote-") as remote_root:
|
||||
sync_from_hf(remote_root, strict=True, revision=parent_commit)
|
||||
_validate_no_existing_seed(remote_root, records[0], source="remote tracking history")
|
||||
|
||||
seen_destinations: set[str] = set()
|
||||
for path, blob, record in zip(prepared_paths, prepared_blobs, records, strict=True):
|
||||
destination = f"{sanitize(record['model_id'])}/{os.path.basename(path)}"
|
||||
if destination in seen_destinations:
|
||||
raise ValueError(f"prepared records collide at Hub path {destination}")
|
||||
if os.path.exists(os.path.join(remote_root, destination)):
|
||||
raise ValueError(f"Hub path already exists: {destination}")
|
||||
seen_destinations.add(destination)
|
||||
operations.append(CommitOperationAdd(path_in_repo=destination, path_or_fileobj=blob))
|
||||
|
||||
commit = api.create_commit(
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
revision="main",
|
||||
parent_commit=parent_commit,
|
||||
create_pr=False,
|
||||
operations=operations,
|
||||
commit_message=f"Perf: seed {records[0]['model_id']} baseline",
|
||||
)
|
||||
|
||||
commit_id = getattr(commit, "oid", None) or getattr(commit, "commit_id", None)
|
||||
commit_id = str(commit_id or commit)
|
||||
try:
|
||||
_write_json(os.path.join(os.path.dirname(manifest_path), "uploaded"), {
|
||||
"commit_id": commit_id,
|
||||
"parent_commit": parent_commit,
|
||||
})
|
||||
except OSError as exc:
|
||||
print(f"Warning: seed commit {commit_id} succeeded but local upload marker failed: {exc}")
|
||||
return commit_id
|
||||
|
||||
|
||||
def _order_sources_by_timestamp(
|
||||
source_paths: list[str],
|
||||
records: list[dict[str, Any]],
|
||||
) -> list[tuple[str, dict[str, Any]]]:
|
||||
timestamped = []
|
||||
undated = []
|
||||
for index, (source_path, record) in enumerate(zip(source_paths, records), start=1):
|
||||
try:
|
||||
timestamp = datetime.fromisoformat(str(record.get("timestamp")))
|
||||
except (TypeError, ValueError):
|
||||
timestamp = None
|
||||
|
||||
if timestamp is None:
|
||||
print(
|
||||
f"Warning: source artifact {index} has a missing or unparsable timestamp; "
|
||||
"preserving its input order after timestamped sources."
|
||||
)
|
||||
undated.append((source_path, record))
|
||||
continue
|
||||
|
||||
if timestamp.tzinfo is None:
|
||||
timestamp = timestamp.replace(tzinfo=timezone.utc)
|
||||
timestamped.append((timestamp, index, source_path, record))
|
||||
|
||||
timestamped.sort(key=lambda item: (item[0], item[1]))
|
||||
return [(source_path, record) for _, _, source_path, record in timestamped] + undated
|
||||
|
||||
|
||||
def _nonnegative_finite_float(value: str) -> float:
|
||||
try:
|
||||
parsed = float(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(f"expected a number, got {value!r}") from exc
|
||||
if not math.isfinite(parsed) or parsed < 0:
|
||||
raise argparse.ArgumentTypeError("value must be a finite non-negative number")
|
||||
return parsed
|
||||
|
||||
|
||||
def _validate_batch_consistency(
|
||||
records: list[dict[str, Any]],
|
||||
max_regression: float,
|
||||
) -> None:
|
||||
print(f"Source batch consistency (maximum regression {max_regression * 100:.1f}%):")
|
||||
failures = []
|
||||
for policy in DEFAULT_METRIC_POLICIES:
|
||||
values = [
|
||||
(index, value)
|
||||
for index, record in enumerate(records, start=1)
|
||||
if (value := safe_float(record.get(policy.key))) is not None
|
||||
]
|
||||
if len(values) < 2:
|
||||
continue
|
||||
|
||||
batch_median = statistics.median(value for _, value in values)
|
||||
if batch_median <= 0:
|
||||
raise ValueError(f"cannot validate {policy.key} consistency against non-positive batch median")
|
||||
|
||||
regressions = []
|
||||
for source_index, value in values:
|
||||
if policy.lower_is_better:
|
||||
regression = (value - batch_median) / batch_median
|
||||
else:
|
||||
regression = (batch_median - value) / batch_median
|
||||
regressions.append((source_index, regression))
|
||||
print(
|
||||
f" {policy.key} source={source_index} value={value:.6f} "
|
||||
f"median={batch_median:.6f} regression={regression * 100:.1f}%")
|
||||
|
||||
worst_source, worst_regression = max(regressions, key=lambda item: item[1])
|
||||
print(f" {policy.key} worst regression: {worst_regression * 100:.1f}% (source={worst_source})")
|
||||
if worst_regression > max_regression:
|
||||
failures.append(
|
||||
f"source artifact {worst_source} {policy.key} regresses by "
|
||||
f"{worst_regression * 100:.1f}% against the batch median")
|
||||
|
||||
if failures:
|
||||
raise ValueError("source batch exceeds maximum intra-batch regression: " + "; ".join(failures))
|
||||
|
||||
|
||||
def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Create reviewed baseline seed records from v2 CALIBRATION_NEEDED artifacts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source-result",
|
||||
dest="source_results",
|
||||
action="append",
|
||||
required=True,
|
||||
help="Path to a normalized_perf_*.json artifact. Repeat for multiple reviewed source artifacts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--intent-rationale",
|
||||
required=True,
|
||||
help="Reviewed reason this calibration artifact should seed the baseline.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tracking-root",
|
||||
default=TRACKING_ROOT,
|
||||
help=(
|
||||
"Operator tracking-mirror path used only to enforce staging separation; "
|
||||
"remote validation uses a fresh temporary snapshot."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--staging-root",
|
||||
default=STAGING_ROOT,
|
||||
help="Separate local root for prepared seed records.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-intra-batch-regression",
|
||||
type=_nonnegative_finite_float,
|
||||
default=DEFAULT_MAX_INTRA_BATCH_REGRESSION,
|
||||
help="Maximum regression of any source against the batch median (default: 0.05).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--operator",
|
||||
default=os.environ.get("USER"),
|
||||
help="Operator name recorded in seed provenance.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = _parse_args(argv)
|
||||
_validate_separate_roots(args.tracking_root, args.staging_root)
|
||||
source_records = [_load_json(path) for path in args.source_results]
|
||||
for source_record in source_records:
|
||||
_validate_calibration_source(source_record)
|
||||
_validate_unique_sources(args.source_results, source_records)
|
||||
identity = _validate_same_identity(source_records)
|
||||
with tempfile.TemporaryDirectory(prefix="performance-seed-prepare-remote-") as remote_root:
|
||||
sync_from_hf(remote_root, strict=True)
|
||||
_validate_no_existing_seed(remote_root, source_records[0], source="remote tracking history")
|
||||
_validate_no_existing_seed(args.staging_root, source_records[0], source="local staging")
|
||||
_validate_batch_consistency(source_records, args.max_intra_batch_regression)
|
||||
ordered_sources = _order_sources_by_timestamp(args.source_results, source_records)
|
||||
print("Seeding exact comparable identity:")
|
||||
for key, value in identity.items():
|
||||
print(f" {key}: {value}")
|
||||
|
||||
seed_records = []
|
||||
for index, (source_path, source_record) in enumerate(ordered_sources, start=1):
|
||||
seed_record = build_baseline_seed_record(
|
||||
source_record,
|
||||
reason=args.intent_rationale,
|
||||
source_result=os.path.realpath(os.path.abspath(source_path)),
|
||||
operator=args.operator,
|
||||
batch_size=len(source_records),
|
||||
batch_index=index,
|
||||
)
|
||||
suffix = f"{index:02d}" if len(source_records) > 1 else None
|
||||
seed_records.append((seed_record, suffix))
|
||||
|
||||
reservation = _reserve_staging_identity(args.staging_root, identity)
|
||||
prepared_seeds = [
|
||||
(_seed_record_path(reservation, seed_record, suffix=suffix), seed_record, suffix)
|
||||
for seed_record, suffix in seed_records
|
||||
]
|
||||
prepared_paths = [seed_path for seed_path, _seed_record, _suffix in prepared_seeds]
|
||||
try:
|
||||
for seed_path, seed_record, suffix in prepared_seeds:
|
||||
write_seed_record(reservation, seed_record, suffix=suffix)
|
||||
print(f"Prepared baseline seed: {seed_path}")
|
||||
manifest_path = _write_preparation_manifest(
|
||||
reservation,
|
||||
args.staging_root,
|
||||
identity,
|
||||
args.source_results,
|
||||
source_records,
|
||||
prepared_paths,
|
||||
)
|
||||
except Exception:
|
||||
_release_failed_reservation(reservation, prepared_paths)
|
||||
raise
|
||||
print(f"Reserved exact identity in staging: {reservation}")
|
||||
print(f"Prepared upload manifest: {manifest_path}")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -102,15 +102,22 @@ def test_partial_v2_identity_requires_schema_version():
|
||||
|
||||
|
||||
def test_optional_v2_metadata_fields_must_be_objects():
|
||||
cfg = {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"quality_metadata": ["not", "an", "object"],
|
||||
}
|
||||
cfg = _v2_config()
|
||||
cfg["quality_metadata"] = ["not", "an", "object"]
|
||||
|
||||
expected = "wan.json: optional v2 metadata field 'quality_metadata' must be an object"
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
|
||||
|
||||
def test_regression_thresholds_validate_without_forcing_v2_schema():
|
||||
cfg = {
|
||||
"benchmark_id": "legacy-benchmark",
|
||||
"regression_thresholds": {"latency": {"threshold_percent": 0.1}},
|
||||
}
|
||||
|
||||
_validate_benchmark_config(cfg, "legacy.json")
|
||||
|
||||
cfg["regression_thresholds"] = []
|
||||
with pytest.raises(ValueError, match="benchmark config field 'regression_thresholds' must be an object"):
|
||||
_validate_benchmark_config(cfg, "legacy.json")
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -177,6 +177,55 @@ def test_records_and_trends_endpoints_filter_by_model_and_gpu():
|
||||
assert trends["groups"][0]["gpu_type"] == "NVIDIA L40S"
|
||||
|
||||
|
||||
def test_v2_display_filters_keep_renamed_cohort_history():
|
||||
identity = {
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"recipe_fingerprint": "recipe-a",
|
||||
"hardware_profile_id": "hw-l40s",
|
||||
"software_profile_id": "sw-cu130",
|
||||
}
|
||||
app = create_app(FakeStore([
|
||||
_record(
|
||||
"old-display",
|
||||
"NVIDIA L40S old label",
|
||||
"2026-01-01T00:00:00+00:00",
|
||||
"a" * 40,
|
||||
10.0,
|
||||
10.0,
|
||||
run_source="scheduled_main",
|
||||
baseline_eligible=True,
|
||||
**identity,
|
||||
),
|
||||
_record(
|
||||
"new-display",
|
||||
"NVIDIA L40S",
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
11.0,
|
||||
9.0,
|
||||
**identity,
|
||||
),
|
||||
]))
|
||||
client = TestClient(app)
|
||||
|
||||
filters = {"model_id": "new-display", "gpu_type": "NVIDIA L40S"}
|
||||
summary = client.get("/api/performance/summary", params=filters).json()
|
||||
trends = client.get("/api/performance/trends", params=filters).json()
|
||||
raw_old_records = client.get("/api/performance/records", params={"model_id": "old-display"}).json()
|
||||
stale_display = client.get("/api/performance/summary", params={"model_id": "old-display"}).json()
|
||||
|
||||
assert summary["count"] == 1
|
||||
assert summary["rows"][0]["baseline_n"] == 1
|
||||
assert summary["rows"][0]["metrics"]["latency"]["baseline"] == 10.0
|
||||
assert trends["count"] == 1
|
||||
assert len(trends["groups"][0]["points"]) == 2
|
||||
assert raw_old_records["count"] == 1
|
||||
assert stale_display["count"] == 0
|
||||
|
||||
|
||||
def test_refresh_endpoint_reports_sync_metadata():
|
||||
app = create_app(FakeStore([]))
|
||||
client = TestClient(app)
|
||||
|
||||
@@ -38,24 +38,69 @@ def test_group_data_uses_full_comparison_cohort():
|
||||
groups = list(dashboard.group_data(df))
|
||||
|
||||
assert len(groups) == 2
|
||||
assert {key[5] for key, _group in groups} == {"recipe-a", "recipe-b"}
|
||||
assert {key[7] for key, _group in groups} == {"sw-a", "sw-b"}
|
||||
assert {key[6] for key, _group in groups} == {"recipe-a", "recipe-b"}
|
||||
assert {key[8] for key, _group in groups} == {"sw-a", "sw-b"}
|
||||
|
||||
|
||||
def test_group_data_ignores_v2_display_name_changes():
|
||||
df = pd.DataFrame([
|
||||
_record(model_id="old-display", gpu_type="NVIDIA L40S old label"),
|
||||
_record(
|
||||
model_id="new-display",
|
||||
gpu_type="NVIDIA L40S",
|
||||
timestamp="2026-01-02T00:00:00+00:00",
|
||||
),
|
||||
])
|
||||
|
||||
groups = list(dashboard.group_data(df))
|
||||
figs, _skipped_metrics = dashboard.build_plots(df)
|
||||
|
||||
assert len(groups) == 1
|
||||
assert all("new-display | NVIDIA L40S" in fig.layout.title.text for fig in figs)
|
||||
|
||||
|
||||
def test_group_data_fills_legacy_cohort_columns():
|
||||
df = pd.DataFrame([{
|
||||
"model_id": "legacy-wan",
|
||||
"gpu_type": "NVIDIA L40S",
|
||||
"timestamp": "2026-01-01T00:00:00+00:00",
|
||||
"commit_sha": "a" * 40,
|
||||
"config_id": "aaaaaaa",
|
||||
"latency": 10.0,
|
||||
}])
|
||||
df = pd.DataFrame([
|
||||
{
|
||||
"model_id": model_id,
|
||||
"gpu_type": "NVIDIA L40S",
|
||||
"timestamp": "2026-01-01T00:00:00+00:00",
|
||||
"commit_sha": "a" * 40,
|
||||
"config_id": "aaaaaaa",
|
||||
"latency": 10.0,
|
||||
}
|
||||
for model_id in ("legacy-wan", "legacy-ltx")
|
||||
])
|
||||
|
||||
groups = list(dashboard.group_data(df))
|
||||
|
||||
assert len(groups) == 1
|
||||
assert groups[0][0][2:] == ("", "", "", "", "", "")
|
||||
assert len(groups) == 2
|
||||
assert {key[1] for key, _group in groups} == {"legacy-wan", "legacy-ltx"}
|
||||
assert all(key[3:] == ("", "", "", "", "", "") for key, _group in groups)
|
||||
|
||||
|
||||
def test_group_data_keeps_partial_v2_identity_scoped_by_display_metadata():
|
||||
partial_identity = {
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"recipe_fingerprint": "",
|
||||
"hardware_profile_id": "",
|
||||
"software_profile_id": "",
|
||||
}
|
||||
df = pd.DataFrame([
|
||||
_record(model_id="wan", gpu_type="NVIDIA L40S", **partial_identity),
|
||||
_record(model_id="ltx", gpu_type="NVIDIA H100", **partial_identity),
|
||||
])
|
||||
|
||||
groups = list(dashboard.group_data(df))
|
||||
|
||||
assert len(groups) == 2
|
||||
assert {(key[1], key[2]) for key, _group in groups} == {
|
||||
("wan", "NVIDIA L40S"),
|
||||
("ltx", "NVIDIA H100"),
|
||||
}
|
||||
|
||||
|
||||
def test_build_plots_labels_distinct_cohorts():
|
||||
|
||||
@@ -290,6 +290,103 @@ def test_build_latest_summary_keeps_variant_versions_separate():
|
||||
assert version_2_row["metrics"]["latency"]["baseline"] == 20.0
|
||||
|
||||
|
||||
def test_v2_dashboard_cohort_spans_display_name_changes():
|
||||
identity = {
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"recipe_fingerprint": "recipe-a",
|
||||
"hardware_profile_id": "hw-l40s",
|
||||
"software_profile_id": "sw-cu130",
|
||||
}
|
||||
records = [
|
||||
_record(
|
||||
"2026-01-01T00:00:00+00:00",
|
||||
"a" * 40,
|
||||
10.0,
|
||||
10.0,
|
||||
model_id="old-display-name",
|
||||
gpu_type="NVIDIA L40S old label",
|
||||
run_source="scheduled_main",
|
||||
baseline_eligible=True,
|
||||
**identity,
|
||||
),
|
||||
_record(
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
11.0,
|
||||
9.0,
|
||||
model_id="new-display-name",
|
||||
gpu_type="NVIDIA L40S",
|
||||
**identity,
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
trends = build_trends(records)
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["model_id"] == "new-display-name"
|
||||
assert rows[0]["gpu_type"] == "NVIDIA L40S"
|
||||
assert rows[0]["baseline_n"] == 1
|
||||
assert rows[0]["metrics"]["latency"]["baseline"] == 10.0
|
||||
assert len(trends) == 1
|
||||
assert trends[0]["model_id"] == "new-display-name"
|
||||
assert len(trends[0]["points"]) == 2
|
||||
|
||||
|
||||
def test_legacy_dashboard_cohorts_still_use_model_and_gpu():
|
||||
records = [
|
||||
_record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0, model_id="wan"),
|
||||
_record("2026-01-02T00:00:00+00:00", "b" * 40, 20.0, 5.0, model_id="ltx"),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
trends = build_trends(records)
|
||||
|
||||
assert {row["model_id"] for row in rows} == {"wan", "ltx"}
|
||||
assert len(trends) == 2
|
||||
|
||||
|
||||
def test_partial_v2_identity_does_not_cross_model_or_gpu():
|
||||
partial_identity = {
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
}
|
||||
records = [
|
||||
_record(
|
||||
"2026-01-01T00:00:00+00:00",
|
||||
"a" * 40,
|
||||
10.0,
|
||||
10.0,
|
||||
model_id="wan",
|
||||
gpu_type="NVIDIA L40S",
|
||||
**partial_identity,
|
||||
),
|
||||
_record(
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
20.0,
|
||||
5.0,
|
||||
model_id="ltx",
|
||||
gpu_type="NVIDIA H100",
|
||||
**partial_identity,
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
trends = build_trends(records)
|
||||
|
||||
assert {(row["model_id"], row["gpu_type"]) for row in rows} == {
|
||||
("wan", "NVIDIA L40S"),
|
||||
("ltx", "NVIDIA H100"),
|
||||
}
|
||||
assert len(trends) == 2
|
||||
|
||||
|
||||
def test_dashboard_identity_preserves_zero_benchmark_version():
|
||||
records = [
|
||||
_record(
|
||||
|
||||
@@ -79,6 +79,7 @@ def test_same_config_produces_same_recipe_fingerprint():
|
||||
def test_recipe_includes_first_class_benchmark_identity():
|
||||
recipe = build_recipe_from_benchmark_config(_benchmark_config())
|
||||
|
||||
assert recipe["recipe_schema_version"] == identity_module.RECIPE_SCHEMA_VERSION == 2
|
||||
assert recipe["benchmark"] == {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v",
|
||||
|
||||
@@ -52,9 +52,9 @@ V2_REQUIRED_IDENTITY_FIELDS = (
|
||||
# declaring it now fails validation loudly instead of being silently
|
||||
# overwritten by the generated one.
|
||||
V2_OPTIONAL_METADATA_FIELDS = (
|
||||
"metric_threshold_policy",
|
||||
"quality_metadata",
|
||||
)
|
||||
COMMON_OBJECT_FIELDS = ("regression_thresholds",)
|
||||
RESULT_SCHEMA_VERSION = 2
|
||||
VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
|
||||
OPTIONAL_RESULT_METADATA_FIELDS = ("quality_metadata", "variant_metadata")
|
||||
@@ -96,6 +96,10 @@ def _validate_benchmark_config(cfg, path="<memory>"):
|
||||
if missing_common:
|
||||
raise ValueError(f"{path}: missing required benchmark config fields: {', '.join(missing_common)}")
|
||||
|
||||
for field in COMMON_OBJECT_FIELDS:
|
||||
if field in cfg and not isinstance(cfg[field], Mapping):
|
||||
raise ValueError(f"{path}: benchmark config field {field!r} must be an object")
|
||||
|
||||
schema_version = cfg.get("config_schema_version")
|
||||
if schema_version is None:
|
||||
if _has_v2_fields(cfg):
|
||||
@@ -245,6 +249,45 @@ def _validate_run_counts(run_config: Mapping[str, Any], benchmark_id: str) -> tu
|
||||
return num_warmup, num_measure
|
||||
|
||||
|
||||
def _resolve_num_gpus(
|
||||
init_kwargs: Mapping[str, Any],
|
||||
run_config: Mapping[str, Any],
|
||||
benchmark_id: str,
|
||||
) -> int:
|
||||
init_num_gpus = init_kwargs.get("num_gpus")
|
||||
required_gpus = run_config.get("required_gpus")
|
||||
for field, value in (
|
||||
("init_kwargs.num_gpus", init_num_gpus),
|
||||
("run_config.required_gpus", required_gpus),
|
||||
):
|
||||
if value is not None and (
|
||||
isinstance(value, bool) or not isinstance(value, int) or value < 1):
|
||||
raise ValueError(f"{benchmark_id}: {field} must be a positive integer")
|
||||
parallel_sizes = []
|
||||
for field in ("tp_size", "sp_size"):
|
||||
value = init_kwargs.get(field)
|
||||
if value is None or value == -1:
|
||||
continue
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
|
||||
raise ValueError(f"{benchmark_id}: init_kwargs.{field} must be -1 or a positive integer")
|
||||
parallel_sizes.append(value)
|
||||
if (
|
||||
init_num_gpus is not None
|
||||
and required_gpus is not None
|
||||
and init_num_gpus != required_gpus
|
||||
):
|
||||
raise ValueError(
|
||||
f"{benchmark_id}: init_kwargs.num_gpus ({init_num_gpus}) must match "
|
||||
f"run_config.required_gpus ({required_gpus})")
|
||||
declared_num_gpus = init_num_gpus or required_gpus
|
||||
parallel_num_gpus = max(parallel_sizes, default=1)
|
||||
if declared_num_gpus is not None and declared_num_gpus < parallel_num_gpus:
|
||||
raise ValueError(
|
||||
f"{benchmark_id}: declared GPU count ({declared_num_gpus}) must be at least "
|
||||
f"max(tp_size, sp_size) ({parallel_num_gpus})")
|
||||
return declared_num_gpus or parallel_num_gpus
|
||||
|
||||
|
||||
def _write_results(results):
|
||||
"""Write JSON results to the results directory."""
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
@@ -339,15 +382,31 @@ def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity):
|
||||
# identity fields, so v2 records always get the full identity block.
|
||||
if cfg.get("config_schema_version") is None:
|
||||
return {}
|
||||
run_config = cfg.get("run_config") or {}
|
||||
num_gpus = _resolve_num_gpus(init_kwargs, run_config, cfg["benchmark_id"])
|
||||
recipe_cfg = dict(cfg)
|
||||
recipe_cfg["init_kwargs"] = {
|
||||
**dict(init_kwargs),
|
||||
"num_gpus": num_gpus,
|
||||
}
|
||||
recipe = build_recipe_from_benchmark_config(
|
||||
cfg,
|
||||
recipe_cfg,
|
||||
resolved_attention_backend=runtime_identity.get("resolved_attention_backend"),
|
||||
resolved_model_revision=runtime_identity.get("resolved_model_revision"),
|
||||
measured_prompts=[prompt],
|
||||
)
|
||||
num_gpus = init_kwargs.get("num_gpus", cfg.get("run_config", {}).get("required_gpus", 1))
|
||||
hw_profile = hardware_profile(num_gpus=num_gpus)
|
||||
sw_profile = software_profile()
|
||||
sw_profile.update({
|
||||
"attention_backend": os.environ.get("FASTVIDEO_ATTENTION_BACKEND") or "auto",
|
||||
"flash_attention_4_enabled": os.environ.get("FASTVIDEO_FA4", "0") != "0",
|
||||
})
|
||||
performance_profile_version = os.environ.get("FASTVIDEO_PERFORMANCE_PROFILE_VERSION")
|
||||
if performance_profile_version:
|
||||
sw_profile["performance_profile_version"] = performance_profile_version
|
||||
image_version = os.environ.get("IMAGE_VERSION")
|
||||
if image_version:
|
||||
sw_profile["container_image_version"] = image_version
|
||||
env_metadata = environment_metadata(
|
||||
hardware=hw_profile,
|
||||
software=sw_profile,
|
||||
@@ -424,6 +483,7 @@ def _build_result_record(
|
||||
) -> dict[str, Any]:
|
||||
if not times or not peak_memories:
|
||||
raise ValueError("Cannot build a performance result record without measurement runs")
|
||||
num_gpus = _resolve_num_gpus(init_kwargs, cfg.get("run_config") or {}, cfg["benchmark_id"])
|
||||
avg_time = sum(times) / len(times)
|
||||
max_peak_memory = max(peak_memories)
|
||||
num_frames = gen_kwargs.get("num_frames")
|
||||
@@ -431,12 +491,18 @@ def _build_result_record(
|
||||
if isinstance(num_frames, (int, float)) and avg_time > 0:
|
||||
throughput_fps = num_frames / avg_time
|
||||
|
||||
result_schema_fields = (
|
||||
{"result_schema_version": RESULT_SCHEMA_VERSION}
|
||||
if _is_v2_config(cfg)
|
||||
else {}
|
||||
)
|
||||
|
||||
return {
|
||||
"benchmark_id": cfg["benchmark_id"],
|
||||
"result_schema_version": RESULT_SCHEMA_VERSION,
|
||||
**result_schema_fields,
|
||||
"model_short_name": model_info.get("model_short_name", ""),
|
||||
"device": device_name,
|
||||
"num_gpus": init_kwargs.get("num_gpus", 1),
|
||||
"num_gpus": num_gpus,
|
||||
"num_warmup_runs": num_warmup,
|
||||
"num_measurement_runs": num_measure,
|
||||
"avg_generation_time_s": round(avg_time, 3),
|
||||
@@ -463,14 +529,16 @@ def _build_result_record(
|
||||
# -- Test -------------------------------------------------------------------
|
||||
|
||||
def _run_benchmark(cfg):
|
||||
run_config = cfg.get("run_config", {})
|
||||
required_gpus = run_config.get("required_gpus", 1)
|
||||
available = torch.cuda.device_count()
|
||||
if available < required_gpus:
|
||||
pytest.skip(f"Need {required_gpus} GPUs, only {available} available")
|
||||
|
||||
run_config = cfg.get("run_config") or {}
|
||||
model_info = cfg["model"]
|
||||
init_kwargs = dict(cfg.get("init_kwargs", {}))
|
||||
num_gpus = _resolve_num_gpus(init_kwargs, run_config, cfg["benchmark_id"])
|
||||
init_kwargs["num_gpus"] = num_gpus
|
||||
|
||||
available = torch.cuda.device_count()
|
||||
if available < num_gpus:
|
||||
pytest.skip(f"Need {num_gpus} GPUs, only {available} available")
|
||||
|
||||
gen_kwargs = dict(cfg.get("generation_kwargs", {}))
|
||||
prompts = cfg.get("test_prompts", ["A cinematic video."])
|
||||
prompt = prompts[0]
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import copy
|
||||
|
||||
from fastvideo.tests.performance import compare_baseline
|
||||
from fastvideo.tests.performance import test_inference_performance as perf
|
||||
|
||||
|
||||
def _benchmark_config():
|
||||
return {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": perf.V2_CONFIG_SCHEMA_VERSION,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"model": {
|
||||
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
"model_short_name": "Wan2.1-T2V-1.3B",
|
||||
},
|
||||
"init_kwargs": {
|
||||
"num_gpus": 2,
|
||||
"flow_shift": 7.0,
|
||||
"sp_size": 2,
|
||||
"tp_size": 1,
|
||||
"vae_sp": True,
|
||||
"vae_tiling": True,
|
||||
"text_encoder_precisions": ["fp32"],
|
||||
},
|
||||
"generation_kwargs": {
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 45,
|
||||
"num_inference_steps": 4,
|
||||
"guidance_scale": 3,
|
||||
"embedded_cfg_scale": 6,
|
||||
"seed": 1024,
|
||||
"fps": 24,
|
||||
"neg_prompt": "low quality",
|
||||
},
|
||||
"run_config": {
|
||||
"num_warmup_runs": 2,
|
||||
"num_measurement_runs": 5,
|
||||
"required_gpus": 2,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_performance_producer_emits_v2_identity_from_raw_result_shape(monkeypatch):
|
||||
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
cfg = _benchmark_config()
|
||||
init_kwargs = dict(cfg["init_kwargs"])
|
||||
cfg["generation_kwargs"]["output_path"] = "/tmp/generated"
|
||||
prompt = "A cinematic video."
|
||||
|
||||
identity_fields = perf._build_identity_fields(
|
||||
cfg,
|
||||
init_kwargs,
|
||||
prompt,
|
||||
{
|
||||
"resolved_attention_backend": "FLASH_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
raw_result = {
|
||||
"benchmark_id": cfg["benchmark_id"],
|
||||
"model_short_name": cfg["model"]["model_short_name"],
|
||||
"device": "NVIDIA L40S PCIe",
|
||||
"num_gpus": init_kwargs["num_gpus"],
|
||||
"num_warmup_runs": 2,
|
||||
"num_measurement_runs": 5,
|
||||
"avg_generation_time_s": 10.0,
|
||||
"individual_times_s": [10.0],
|
||||
"throughput_fps": 4.5,
|
||||
"max_peak_memory_mb": 10000.0,
|
||||
"individual_peak_memories_mb": [10000.0],
|
||||
"thresholds": {},
|
||||
"commit": "a" * 40,
|
||||
"pr_number": "123",
|
||||
"timestamp": "2026-06-16T00:00:00+00:00",
|
||||
"text_encoder_time_s": 1.0,
|
||||
"dit_time_s": 2.0,
|
||||
"vae_decode_time_s": 3.0,
|
||||
**identity_fields,
|
||||
}
|
||||
|
||||
record = compare_baseline.normalize_performance_result(raw_result)
|
||||
|
||||
assert record["workload_id"] == "wan-t2v-1.3b"
|
||||
assert record["variant_id"] == "canonical"
|
||||
assert record["benchmark_version"] == 1
|
||||
assert record["hardware_profile_id"].startswith("hw-")
|
||||
assert record["software_profile_id"].startswith("sw-")
|
||||
assert record["software_profile_id"] == perf.software_profile_id(record["software_profile"])
|
||||
assert len(record["recipe_fingerprint"]) == 64
|
||||
assert "output_path" not in record["recipe"]["generation_kwargs"]
|
||||
|
||||
|
||||
def test_display_benchmark_id_rename_preserves_producer_fingerprint_and_cohort(monkeypatch):
|
||||
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
original = _benchmark_config()
|
||||
renamed = copy.deepcopy(original)
|
||||
renamed["benchmark_id"] = "wan-t2v-renamed-display-id"
|
||||
|
||||
def build_identity(cfg):
|
||||
return perf._build_identity_fields(
|
||||
cfg,
|
||||
dict(cfg["init_kwargs"]),
|
||||
"A cinematic video.",
|
||||
{
|
||||
"resolved_attention_backend": "FLASH_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
|
||||
original_identity = build_identity(original)
|
||||
renamed_identity = build_identity(renamed)
|
||||
|
||||
assert original_identity["recipe"]["benchmark"]["benchmark_id"] == original["benchmark_id"]
|
||||
assert renamed_identity["recipe"]["benchmark"]["benchmark_id"] == renamed["benchmark_id"]
|
||||
assert original_identity["recipe_fingerprint"] == renamed_identity["recipe_fingerprint"]
|
||||
assert compare_baseline._comparison_identity_filters(
|
||||
original_identity) == compare_baseline._comparison_identity_filters(renamed_identity)
|
||||
|
||||
|
||||
def test_v2_identity_tolerates_null_run_config(monkeypatch):
|
||||
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
cfg = _benchmark_config()
|
||||
cfg["run_config"] = None
|
||||
init_kwargs = dict(cfg["init_kwargs"])
|
||||
|
||||
identity_fields = perf._build_identity_fields(
|
||||
cfg,
|
||||
init_kwargs,
|
||||
"A cinematic video.",
|
||||
{
|
||||
"resolved_attention_backend": "FLASH_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
|
||||
assert identity_fields["hardware_profile"]["gpu_count"] == 2
|
||||
|
||||
|
||||
def test_effective_gpu_count_is_recorded_in_recipe_and_hardware(monkeypatch):
|
||||
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
cfg = _benchmark_config()
|
||||
del cfg["init_kwargs"]["num_gpus"]
|
||||
init_kwargs = dict(cfg["init_kwargs"])
|
||||
|
||||
identity_fields = perf._build_identity_fields(
|
||||
cfg,
|
||||
init_kwargs,
|
||||
"A cinematic video.",
|
||||
{
|
||||
"resolved_attention_backend": "FLASH_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
|
||||
assert identity_fields["recipe"]["init_kwargs"]["num_gpus"] == 2
|
||||
assert identity_fields["hardware_profile"]["gpu_count"] == 2
|
||||
|
||||
|
||||
def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypatch):
|
||||
base_profile = {
|
||||
"python": "3.12",
|
||||
"pytorch": "2.7",
|
||||
"cuda": "12.8",
|
||||
"packages": {
|
||||
"fastvideo_kernel": "0.3.2",
|
||||
"triton": "3.2.1",
|
||||
},
|
||||
}
|
||||
monkeypatch.setattr(perf, "software_profile", lambda: dict(base_profile))
|
||||
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
||||
monkeypatch.setenv("FASTVIDEO_FA4", "1")
|
||||
monkeypatch.setenv("FASTVIDEO_PERFORMANCE_PROFILE_VERSION", "perf-profile-v2")
|
||||
monkeypatch.setenv("IMAGE_VERSION", "py3.12-cuda13.0")
|
||||
monkeypatch.setenv(
|
||||
"FASTVIDEO_CONTAINER_IMAGE_REF",
|
||||
"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev@sha256:abc",
|
||||
)
|
||||
|
||||
cfg = _benchmark_config()
|
||||
identity_fields = perf._build_identity_fields(
|
||||
cfg,
|
||||
dict(cfg["init_kwargs"]),
|
||||
"A cinematic video.",
|
||||
{
|
||||
"resolved_attention_backend": "SAGE_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
|
||||
expected_profile = {
|
||||
**base_profile,
|
||||
"attention_backend": "SAGE_ATTN",
|
||||
"flash_attention_4_enabled": True,
|
||||
"performance_profile_version": "perf-profile-v2",
|
||||
"container_image_version": "py3.12-cuda13.0",
|
||||
}
|
||||
assert identity_fields["software_profile"] == expected_profile
|
||||
assert identity_fields["software_profile_id"] == perf.software_profile_id(expected_profile)
|
||||
assert identity_fields["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith(
|
||||
"sha256:abc"
|
||||
)
|
||||
|
||||
first_profile_id = identity_fields["software_profile_id"]
|
||||
monkeypatch.setenv(
|
||||
"FASTVIDEO_CONTAINER_IMAGE_REF",
|
||||
"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev@sha256:def",
|
||||
)
|
||||
changed_audit = perf._build_identity_fields(
|
||||
cfg,
|
||||
dict(cfg["init_kwargs"]),
|
||||
"A cinematic video.",
|
||||
{
|
||||
"resolved_attention_backend": "SAGE_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
assert changed_audit["software_profile_id"] == first_profile_id
|
||||
assert changed_audit["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith(
|
||||
"sha256:def"
|
||||
)
|
||||
|
||||
monkeypatch.setenv("FASTVIDEO_FA4", "0")
|
||||
changed_runtime = perf._build_identity_fields(
|
||||
cfg,
|
||||
dict(cfg["init_kwargs"]),
|
||||
"A cinematic video.",
|
||||
{
|
||||
"resolved_attention_backend": "SAGE_ATTN",
|
||||
"resolved_model_revision": None,
|
||||
},
|
||||
)
|
||||
assert changed_runtime["software_profile_id"] != first_profile_id
|
||||
@@ -8,6 +8,7 @@ from fastvideo.tests.performance import test_inference_performance as perf_test
|
||||
def _benchmark_config():
|
||||
return {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": perf_test.V2_CONFIG_SCHEMA_VERSION,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
@@ -155,6 +156,33 @@ def test_validate_run_counts_returns_defaults():
|
||||
assert perf_test._validate_run_counts({}, "wan-t2v-1.3b-2gpu") == (1, 3)
|
||||
|
||||
|
||||
def test_resolve_num_gpus_uses_one_count_and_rejects_conflicts():
|
||||
assert perf_test._resolve_num_gpus(
|
||||
{},
|
||||
{"required_gpus": 2},
|
||||
"wan-t2v-1.3b-2gpu",
|
||||
) == 2
|
||||
assert perf_test._resolve_num_gpus(
|
||||
{"tp_size": 2},
|
||||
{},
|
||||
"wan-t2v-1.3b-tp2",
|
||||
) == 2
|
||||
|
||||
with pytest.raises(ValueError, match="must match"):
|
||||
perf_test._resolve_num_gpus(
|
||||
{"num_gpus": 1},
|
||||
{"required_gpus": 2},
|
||||
"wan-t2v-1.3b-2gpu",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"at least max\(tp_size, sp_size\)"):
|
||||
perf_test._resolve_num_gpus(
|
||||
{"num_gpus": 1, "tp_size": 2},
|
||||
{"required_gpus": 1},
|
||||
"wan-t2v-1.3b-tp2",
|
||||
)
|
||||
|
||||
|
||||
def test_build_result_record_rejects_empty_measurements(monkeypatch):
|
||||
monkeypatch.setattr(perf_test, "_build_identity_fields", lambda *_args: _identity_fields())
|
||||
|
||||
|
||||
@@ -2,7 +2,6 @@ import torch
|
||||
import types
|
||||
import pytest
|
||||
|
||||
from fastvideo.configs.models.encoders.clip import CLIPTextArchConfig, CLIPTextConfig
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
@@ -14,9 +13,6 @@ class TensorDict(dict):
|
||||
return TensorDict({k: v.to(device) for k, v in self.items()})
|
||||
|
||||
class FakeTokenizer:
|
||||
def __init__(self):
|
||||
self.last_chat_template_kwargs = None
|
||||
|
||||
def __call__(self, texts, **kwargs):
|
||||
B = len(texts)
|
||||
seq_len = int(kwargs.get("max_length", 4))
|
||||
@@ -25,10 +21,6 @@ class FakeTokenizer:
|
||||
"attention_mask": torch.ones(B, seq_len, dtype=torch.long),
|
||||
})
|
||||
|
||||
def apply_chat_template(self, messages, **kwargs):
|
||||
self.last_chat_template_kwargs = kwargs
|
||||
return "formatted prompt"
|
||||
|
||||
|
||||
class FakeChatTokenizer:
|
||||
def __init__(self):
|
||||
@@ -48,22 +40,17 @@ class FakeChatTokenizer:
|
||||
"attention_mask": torch.ones(B, seq_len, dtype=torch.long),
|
||||
})
|
||||
|
||||
|
||||
class FakeTextEncoder(torch.nn.Module):
|
||||
def __init__(self, hidden_size=8):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.last_input_device = None
|
||||
self.last_output_hidden_states = None
|
||||
|
||||
def forward(self, input_ids, attention_mask, output_hidden_states=False):
|
||||
self.last_input_device = input_ids.device
|
||||
self.last_output_hidden_states = bool(output_hidden_states)
|
||||
B, T = input_ids.shape
|
||||
last_hidden_state = torch.arange(
|
||||
B * T * self.hidden_size,
|
||||
dtype=torch.float32,
|
||||
device=input_ids.device,
|
||||
).view(B, T, self.hidden_size)
|
||||
last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size)
|
||||
hidden_states = (last_hidden_state, ) if output_hidden_states else None
|
||||
return types.SimpleNamespace(last_hidden_state=last_hidden_state,
|
||||
hidden_states=hidden_states)
|
||||
@@ -226,104 +213,3 @@ def test_chat_list_preprocess_output_is_not_stripped():
|
||||
{"role": "user", "content": "a robotic arm welding a metal structure"},
|
||||
]]
|
||||
assert tokenizer.last_kwargs["return_tensors"] == "pt"
|
||||
|
||||
|
||||
def test_chat_template_thinking_is_keyword_only_and_preserves_subclass_positions():
|
||||
arch = CLIPTextArchConfig()
|
||||
config = CLIPTextConfig(
|
||||
arch,
|
||||
"legacy-prefix",
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
False,
|
||||
7,
|
||||
False,
|
||||
True,
|
||||
False,
|
||||
)
|
||||
|
||||
assert config.num_hidden_layers_override == 7
|
||||
assert config.require_post_norm is False
|
||||
assert config.enable_scale is True
|
||||
assert config.is_causal is False
|
||||
assert config.chat_template_enable_thinking is False
|
||||
|
||||
configured = CLIPTextConfig(chat_template_enable_thinking=True)
|
||||
assert configured.chat_template_enable_thinking is True
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
CLIPTextConfig(
|
||||
arch,
|
||||
"legacy-prefix",
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
False,
|
||||
7,
|
||||
False,
|
||||
True,
|
||||
False,
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_thinking", [False, True])
|
||||
def test_encode_text_forwards_chat_template_thinking_config(enable_thinking):
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
encoder_config = fastvideo_args.pipeline_config.text_encoder_configs[0]
|
||||
encoder_config.is_chat_model = True
|
||||
encoder_config.chat_template_enable_thinking = enable_thinking
|
||||
|
||||
stage = make_stage(num_encoders=1, hidden_size=hidden)
|
||||
stage.encode_text("a", fastvideo_args, encoder_index=[0])
|
||||
|
||||
assert stage.tokenizers[0].last_chat_template_kwargs == {
|
||||
"tokenize": False,
|
||||
"add_generation_prompt": True,
|
||||
"enable_thinking": enable_thinking,
|
||||
}
|
||||
|
||||
|
||||
def test_encode_text_uses_hf_passthrough_input_device(monkeypatch):
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
stage = make_stage(num_encoders=1, hidden_size=hidden)
|
||||
stage.text_encoders[0]._fastvideo_input_device = torch.device("cpu")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.pipelines.stages.text_encoding.get_local_torch_device",
|
||||
lambda: torch.device("meta"),
|
||||
)
|
||||
|
||||
batch = ForwardBatch(
|
||||
data_type="video",
|
||||
prompt="a",
|
||||
do_classifier_free_guidance=False,
|
||||
prompt_embeds=[],
|
||||
negative_prompt_embeds=None,
|
||||
prompt_attention_mask=[],
|
||||
negative_attention_mask=None,
|
||||
)
|
||||
output = stage.forward(batch, fastvideo_args)
|
||||
|
||||
# The marker governs where the encoder *receives* its tokens...
|
||||
assert stage.text_encoders[0].last_input_device == torch.device("cpu")
|
||||
# ...while the stage still normalizes the embeds it returns onto the
|
||||
# caller's target device.
|
||||
assert output.prompt_embeds[0].device.type == "meta"
|
||||
|
||||
|
||||
def test_encode_text_explicit_device_overrides_hf_passthrough_marker():
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
stage = make_stage(num_encoders=1, hidden_size=hidden)
|
||||
stage.text_encoders[0]._fastvideo_input_device = torch.device("cpu")
|
||||
|
||||
output = stage.encode_text(
|
||||
"a",
|
||||
fastvideo_args,
|
||||
encoder_index=[0],
|
||||
device=torch.device("meta"),
|
||||
)
|
||||
|
||||
assert stage.text_encoders[0].last_input_device == torch.device("meta")
|
||||
assert output[0].device.type == "meta"
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
# Z-Image Port Status
|
||||
|
||||
## Summary
|
||||
|
||||
- model_family: `zimage`
|
||||
- variant: `Z-Image-Turbo`
|
||||
- workload_types: `T2I`
|
||||
- PR scope: component-only (`text_encoder`, `tokenizer`, `vae`, `scheduler`)
|
||||
- official_ref: `https://github.com/Tongyi-MAI/Z-Image`
|
||||
- official_ref_dir: `<repo_root>/Z-Image/src`
|
||||
- official_ref_commit: `26f23eda626ffadda020b04ff79488e1d72004cd`
|
||||
- hf_weights_path: `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6`
|
||||
- local_weights_dir: `<repo_root>/official_weights/Z-Image/`
|
||||
- source_layout: `diffusers`
|
||||
- local_tests_readme: `tests/local_tests/zimage/README.md`
|
||||
|
||||
The pinned text encoder has 36 decoder blocks, hidden size 2560,
|
||||
intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension
|
||||
128. FastVideo obtains those values from the pinned `config.json`.
|
||||
|
||||
## Current Phase
|
||||
|
||||
- phase: Phase 6 (component parity revalidation)
|
||||
- status: in_progress
|
||||
- owner: parity
|
||||
- last_updated: 2026-07-04
|
||||
|
||||
## Component Matrix
|
||||
|
||||
| Component | Type | Reuse/Port | Official Definition / Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
|
||||
|---|---|---|---|---|---|---|---|---|
|
||||
| Text encoder (Qwen3 body) | text_encoder | reused production body + native implementation | Pinned `Qwen3ForCausalLM` checkpoint; Transformers `AutoModel` body; official pipeline consumes `hidden_states[-2]` | `TextEncoderLoader` -> `Qwen3ForCausalLM.from_pretrained_local` -> independent body-only `AutoModel`, plus native `fastvideo.models.encoders.qwen3.Qwen3ForCausalLM` | DONE | not_needed for current HF subfolder | REVALIDATION REQUIRED (historical unpinned native fp32/bf16 PASS only) | I001, I006, I007 |
|
||||
| Tokenizer | tokenizer | reused | `AutoTokenizer` from `<weights>/tokenizer/`; official thinking template and length 512 | `TokenizerLoader` | DONE | not_needed | REVALIDATION REQUIRED (historical unpinned PASS) | I003, I010 |
|
||||
| VAE | vae | reused with accepted wrapper exception | Pinned `zimage.AutoencoderKL`; official decode uses `(latents / scaling_factor) + shift_factor` | Existing Diffusers-backed `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` through production `VAELoader`, plus direct implementation coverage | DONE | not_needed | REVALIDATION REQUIRED (historical unpinned direct-decode PASS only) | I008 |
|
||||
| Scheduler | scheduler | reused with Z-Image extension | Pinned `zimage.FlowMatchEulerDiscreteScheduler`; official pipeline sets `sigma_min=0.0` | `FlowMatchEulerDiscreteScheduler(use_reference_discrete_timesteps=True, sigma_min=0.0)` | DONE | not_needed | REVALIDATION REQUIRED; asset-free regressions remain separately runnable | I002, I009 |
|
||||
| Transformer (`ZImageTransformer2DModel`) | dit | planned native port | `<weights>/transformer/` and pinned official source | future PR | NOT_STARTED | unknown_until_prototype | NOT_STARTED | I004 |
|
||||
| Pipeline | pipeline | planned | Pinned `ZImagePipeline` | future PR | NOT_STARTED | depends_on_transformer | NOT_STARTED | I005 |
|
||||
|
||||
## Conversion State
|
||||
|
||||
- conversion_script: none for the reused components in this component-only PR
|
||||
- converted_weights_dir: none
|
||||
- source_layout: diffusers
|
||||
- future_transformer_conversion: unknown until native key/shape prototype exists
|
||||
- strict_load_status: revalidation_required
|
||||
- production_passthrough: body-only Qwen3 `AutoModel`, tokenizer assets
|
||||
- direct_load_components: native Qwen3 parity target, VAE, scheduler
|
||||
- native_qwen_exclusion: `lm_head.weight` only; the native encoder owns the body, not an LM head
|
||||
- native_qwen_fused_contract: every Q/K/V and gate/up shard is required unless the exact fused destination is present
|
||||
- native_qwen_auxiliary_contract: mapped quantization auxiliaries, including scale parameters, must load and may not be silently discarded
|
||||
- retry_history: none
|
||||
|
||||
## Parity Commands
|
||||
|
||||
| Scope | Command | Last Result | Notes |
|
||||
|---|---|---|---|
|
||||
| Scheduler | `pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v -s` | RERUN REQUIRED; historical unpinned 2/2 PASS on A40, 2026-05-12 | Now enforces reference SHA/import origin and exact zero endpoint; includes asset-free positional/default regressions. |
|
||||
| Tokenizer | `pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -s` | RERUN REQUIRED; historical unpinned 2/2 PASS on A40, 2026-05-12 | Now uses `enable_thinking=True`, `max_length=512`, and fails if chat templating is unavailable. |
|
||||
| VAE | `pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v -s` | RERUN REQUIRED; historical unpinned 1/1 direct-decode PASS on A40, 2026-05-12 | Scope is now `both`: direct decode plus `VAELoader` and official latent normalization. |
|
||||
| Text encoder fp32 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | RERUN REQUIRED; historical unpinned native PASS on L40S, 2026-06-21 | New run must cover independent reference, body-only production path, and native implementation. |
|
||||
| Text encoder bf16 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16] -v -s` | RERUN REQUIRED; historical unpinned native PASS on L40S, 2026-06-21 | Historical thresholds retained pending pinned-snapshot validation. |
|
||||
| Per-layer bf16 diagnostic | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | RERUN REQUIRED; historical unpinned informational run on L40S, 2026-06-21 | Pinned architecture has 36 blocks and 37 hidden-state entries under the tested contract. |
|
||||
|
||||
## Open Questions
|
||||
|
||||
| ID | Question | Owner | Needed By Phase | Status | Resolution |
|
||||
|---|---|---|---|---|---|
|
||||
| Q001 | Which exact Z-Image source revision is the parity oracle? | prep | Phase 1 | resolved | Pinned `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`; scheduler/VAE tests fail on a wrong HEAD or import origin. |
|
||||
| Q002 | Which immutable HF snapshot supplies component weights? | prep | Phase 1 | resolved | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6`. |
|
||||
| Q003 | What prompt/tokenization/hidden-state contract does the official pipeline use? | pipeline | Phase 6 | resolved | `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)`, tokenization length 512, and valid-token `hidden_states[-2]`. |
|
||||
| Q004 | Does the future native transformer require conversion? | conversion | Phase 5 of future transformer PR | open | Decide from official/native key-shape dumps; Diffusers layout alone is not proof. |
|
||||
|
||||
## Issues And Blockers
|
||||
|
||||
| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
|
||||
|---|---|---|---|---|---|---|---|---|
|
||||
| I001 | Phase 6 | text_encoder | medium | The full checkpoint includes `lm_head.weight`, while the native FastVideo encoder owns only the body. | Native parity allowlists exactly `lm_head.weight`; production uses body-only Transformers `AutoModel`. | parity | in_progress | Implementation is complete; resolve only after the pinned fp32/bf16 rerun. |
|
||||
| I002 | Phase 7 | scheduler/pipeline | high | Published `scheduler_config.json` omits the reference-schedule flag and official zero endpoint. | Pinned pipeline mutates `sigma_min=0.0`; parity requires both scheduler values. | pipeline | open | Future pipeline config must serialize `use_reference_discrete_timesteps=True` and `sigma_min=0.0`. |
|
||||
| I003 | Phase 6 | tokenizer/text_encoder | medium | Prompt formatting and target length previously differed from the official call. | Pinned source uses thinking chat template, 512 tokens, and `hidden_states[-2]`. | parity | in_progress | Test/config behavior corrected; pinned asset rerun required before resolving. |
|
||||
| I004 | Phase 4 | transformer | high | `ZImageTransformer2DModel` is not ported. | No FastVideo target or parity test exists. | port | open | Separate future PR. |
|
||||
| I005 | Phase 7 | pipeline | high | No FastVideo pipeline/config/preset/registry/example or quality regression exists. | Component-only PR scope. | pipeline | open | Separate future PR after component gates pass. |
|
||||
| I006 | Phase 6 | text_encoder | high | Production loading previously used a causal-LM object and could override CPU placement, while parity exercised only direct native construction. | Production path is now body-only `AutoModel`; loader records `_fastvideo_input_device`; stage moves tokens to that device; parity uses a distinct reference instance. | encoder | in_progress | Implementation complete; pinned production-loader rerun required. |
|
||||
| I007 | Phase 6 | text_encoder | high | Lenient native loading could miss one fused source shard or an auxiliary quantization scale. | Loader completeness now covers every destination, required split shard, exact fused destination, and mapped scale parameter. | encoder | in_progress | Asset-free strictness regressions and pinned checkpoint load must pass before resolving. |
|
||||
| I008 | Phase 6 | vae | medium | Direct decode alone did not cover production registry/config/device/strict-load behavior; the shared target is a runtime Diffusers wrapper. | VAE parity scope is now `both` and includes `VAELoader` plus official normalization. | parity | in_progress | Existing Diffusers `AutoencoderKL` wrapper exception is accepted only for this component-only PR; pinned rerun required. |
|
||||
| I009 | Phase 2 | reference assets | high | Adding a nonexistent clone path to `sys.path` could import another installed `zimage` package and yield false parity. | Tests now verify exact clone HEAD and module paths under `Z-Image/src`; `/Z-Image/` is ignored. | prep | resolved | Wrong SHA/path/import is a failure; absent local assets may skip. |
|
||||
| I010 | Phase 6 | tokenizer | medium | Missing `apply_chat_template` previously skipped and could hide an incompatible tokenizer. | Tokenizer parity now asserts the API and exact official kwargs. | parity | resolved | Missing chat-template behavior fails. Asset-backed numerical rerun remains I003. |
|
||||
|
||||
## Escape Hatches
|
||||
|
||||
| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
|
||||
|---|---|---|---|---|---|---|
|
||||
|
||||
## Decisions
|
||||
|
||||
| Date | Decision | Rationale | Impact |
|
||||
|---|---|---|---|
|
||||
| 2026-07-04 | Pin both external sources and treat all earlier GPU results as historical until rerun. | Earlier evidence did not record an immutable HF revision, and source imports were not origin-checked. | README has executable clone/download commands; tests enforce the source SHA; all asset-backed parity rows are revalidation-required. |
|
||||
| 2026-07-04 | Use a body-only Transformers `AutoModel` for production and a separate `AutoModel` object as the parity oracle. | Every FastVideo consumer needs hidden states, not full-vocabulary logits or an LM head; comparing an object with itself is not independent evidence. | Production avoids the LM head while parity covers production loading and native implementation separately. |
|
||||
| 2026-07-04 | Make fused-source and quantization-scale completeness part of native Qwen strict loading. | A loaded destination name alone did not prove that every required Q/K/V or gate/up shard and auxiliary scale arrived. | Missing split shards or mapped scale parameters fail; an exact fused destination remains valid. |
|
||||
| 2026-07-04 | Accept the existing Diffusers-backed AutoencoderKL wrapper only for this component-only PR. | The target is established shared code and this PR adds production-loader evidence rather than a new VAE implementation. | Exception is limited to the VAE here and is not precedent for the transformer or future full pipeline. |
|
||||
| 2026-07-04 | Treat absent chat-template behavior as incompatibility, not an optional test environment. | Thinking-format prompt construction is part of the official numerical contract. | Tokenizer parity fails when `apply_chat_template` is unavailable. |
|
||||
| 2026-07-04 | Record the pinned Qwen architecture as 36 blocks, hidden size 2560, intermediate size 9728, 32 attention heads, and 8 KV heads. | Earlier notes described a different architecture and could mis-size the native target. | Config/review evidence now matches `f332072aa78be7aecdf3ee76d5c247082da564a6`; expected hidden-state length is 37. |
|
||||
| 2026-06-21 | Expand generated Qwen `position_ids` to `[batch_size, seq_len]`. | The rotary layer flattens positions; `[1, seq_len]` misaligned RoPE for batch sizes greater than one. | Batch-one behavior is unchanged; pinned batch-two parity still requires rerun. |
|
||||
| 2026-06-21 | Reuse the shared native `Qwen3ForCausalLM` implementation rather than maintain a second Z-Image encoder. | The shared implementation supports the same config-driven GQA, fused QKV/gate-up, and hidden-state output contract. | Registry/native parity uses the shared class; architecture values come from the pinned config. |
|
||||
| 2026-05-12 | Use mean and median drift checks for deep bf16 encoder parity. | Cross-kernel fused/unfused bf16 arithmetic produces a long max-error tail; fp32 and per-layer diagnostics distinguish this from structural divergence. | Existing thresholds remain provisional until the pinned 36-block snapshot is rerun. |
|
||||
| 2026-05-12 | Forward the complete scheduler config except loader metadata. | Hand-picking three fields could silently omit future scheduler behavior. | Scheduler parity consumes all on-disk constructor fields and explicitly supplies Z-Image runtime overrides. |
|
||||
|
||||
## Handoff Notes
|
||||
|
||||
- No asset-backed component currently has pinned-snapshot `non_skip_pass`
|
||||
evidence. Do not promote a historical row to PASS without recording the exact
|
||||
command and non-skip result against both immutable pins.
|
||||
- Text-encoder revalidation must cover the independent `AutoModel` oracle,
|
||||
body-only production loader, native implementation, requested input device,
|
||||
fused/split source completeness, and auxiliary quantization scales.
|
||||
- VAE revalidation must cover both direct implementation decode and production
|
||||
`VAELoader` behavior with official scaling/shift normalization.
|
||||
- Next separate-PR work: transformer prototype/parity, conversion decision,
|
||||
pipeline/config/preset/registry/example, pipeline parity, and image-quality
|
||||
regression.
|
||||
@@ -1,142 +0,0 @@
|
||||
# Z-Image Local Tests
|
||||
|
||||
Local-only component coverage for the `zimage` FastVideo port (`T2I`,
|
||||
Z-Image-Turbo only). This component-only PR covers the Qwen3 text encoder,
|
||||
tokenizer, AutoencoderKL VAE, and FlowMatchEulerDiscreteScheduler. The native
|
||||
`ZImageTransformer2DModel`, pipeline, pipeline parity, example, and image-quality
|
||||
regression remain out of scope.
|
||||
|
||||
> **Status:** DRAFT port, in progress. Every test that consumes the reference
|
||||
> clone or HF assets requires a fresh non-skip rerun against the immutable pins
|
||||
> below. Earlier A40/L40S results used an unrecorded HF snapshot and are retained
|
||||
> only as historical diagnostics, not current verification evidence. See
|
||||
> [`PORT_STATUS.md`](./PORT_STATUS.md) for the live state.
|
||||
|
||||
## Reference assets and scope
|
||||
|
||||
| Field | Value |
|
||||
|---|---|
|
||||
| Model family / variant | `zimage` / `Z-Image-Turbo` |
|
||||
| Workload | `T2I` |
|
||||
| Component-only PR scope | text encoder, tokenizer, VAE, scheduler |
|
||||
| Official reference | `https://github.com/Tongyi-MAI/Z-Image` |
|
||||
| Local reference dir | `<repo_root>/Z-Image/src` |
|
||||
| Official commit | `26f23eda626ffadda020b04ff79488e1d72004cd` |
|
||||
| HF weights | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` |
|
||||
| Local weights dir | `<repo_root>/official_weights/Z-Image/` (`text_encoder/`, `tokenizer/`, `vae/`, `scheduler/`) |
|
||||
| Source layout | Diffusers-style per-component subfolders |
|
||||
| Conversion | not needed for the reused components in this PR; the future native transformer/full-pipeline decision will come from native and official key/shape prototypes |
|
||||
| HF token env | `HF_TOKEN` (the pinned repository is public; never record a token value) |
|
||||
|
||||
The pinned text-encoder config declares `Qwen3ForCausalLM` with 36 decoder
|
||||
blocks, hidden size 2560, intermediate size 9728, 32 attention heads, 8 KV
|
||||
heads, and head dimension 128. FastVideo loads those values from
|
||||
`text_encoder/config.json`; they are not hard-coded in the Z-Image integration.
|
||||
|
||||
## Shared environment setup
|
||||
|
||||
Run from the FastVideo repository root in the same environment used for
|
||||
FastVideo. Keep the reference clone and weights inside the ignored paths shown
|
||||
below.
|
||||
|
||||
Clone and pin the official reference:
|
||||
|
||||
```bash
|
||||
python .agents/skills/add-model-01-prep/scripts/clone_reference_repo.py \
|
||||
https://github.com/Tongyi-MAI/Z-Image.git \
|
||||
Z-Image \
|
||||
--commit 26f23eda626ffadda020b04ff79488e1d72004cd \
|
||||
--update-gitignore
|
||||
git -C Z-Image rev-parse HEAD
|
||||
```
|
||||
|
||||
The final command must print
|
||||
`26f23eda626ffadda020b04ff79488e1d72004cd`. The scheduler and VAE tests also
|
||||
enforce that exact HEAD and verify that imported `zimage` modules resolve under
|
||||
`<repo_root>/Z-Image/src`. A missing clone may skip local parity; a wrong SHA or
|
||||
wrong import origin fails rather than silently testing another installation.
|
||||
|
||||
Download the immutable HF snapshot:
|
||||
|
||||
```bash
|
||||
python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py \
|
||||
Tongyi-MAI/Z-Image-Turbo \
|
||||
official_weights/Z-Image \
|
||||
--revision f332072aa78be7aecdf3ee76d5c247082da564a6 \
|
||||
--allow-pattern 'text_encoder/*' \
|
||||
--allow-pattern 'tokenizer/*' \
|
||||
--allow-pattern 'vae/*' \
|
||||
--allow-pattern 'scheduler/*'
|
||||
```
|
||||
|
||||
Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
|
||||
`flash-attn`, `triton`, or CUDA packages) without explicit approval.
|
||||
|
||||
```text
|
||||
dependency_changes: none
|
||||
official_env_status: not_verified
|
||||
private_dep_stubs: none
|
||||
blocked_on: pinned-snapshot non-skip component reruns
|
||||
```
|
||||
|
||||
## Run the tests
|
||||
|
||||
```bash
|
||||
pytest tests/local_tests/zimage/ -v -s
|
||||
```
|
||||
|
||||
| Component | Coverage scope and contract | Status |
|
||||
|---|---|---|
|
||||
| Scheduler ([`test_zimage_scheduler_parity.py`](./test_zimage_scheduler_parity.py)) | `implementation_subcomponent`; exact official `sigma_min=0.0` plus `use_reference_discrete_timesteps=True`; positional/default-path regressions remain asset-free | REVALIDATION REQUIRED |
|
||||
| Tokenizer ([`test_zimage_tokenizer_parity.py`](./test_zimage_tokenizer_parity.py)) | `production_loader`; exact `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)` and `max_length=512`; absence of `apply_chat_template` is a failure, not a skip | REVALIDATION REQUIRED (historical unpinned PASS) |
|
||||
| VAE ([`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py)) | `both`; direct implementation decode plus production `VAELoader`; production check applies `(latents / scaling_factor) + shift_factor` before decode | REVALIDATION REQUIRED (historical unpinned direct-decode PASS only) |
|
||||
| Text encoder ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `both`; independent Transformers `AutoModel` reference, body-only production loader, and FastVideo-native implementation; fp32/bf16 output checks plus fused/split/quant-scale strictness | REVALIDATION REQUIRED (historical unpinned native PASS only) |
|
||||
|
||||
## Component contracts
|
||||
|
||||
### Text encoder and tokenizer
|
||||
|
||||
- The production `TextEncoderLoader` path deliberately loads Transformers
|
||||
`AutoModel`, even when the checkpoint advertises `Qwen3ForCausalLM`. FastVideo
|
||||
consumes hidden states only, so the production path is body-only and does not
|
||||
materialize or execute an LM head.
|
||||
- The parity oracle is a separate `AutoModel` instance; it is not the object
|
||||
returned by the production loader. The native `Qwen3ForCausalLM` implementation
|
||||
is compared separately.
|
||||
- CPU/MPS text-encoder offload remains on the requested device. The loader records
|
||||
`_fastvideo_input_device`, and `TextEncodingStage` sends token tensors there.
|
||||
- Native loading must account for every destination parameter and every required
|
||||
fused source shard. Q/K/V and gate/up shards are all required unless the exact
|
||||
destination is supplied already fused; quantized auxiliary parameters such as
|
||||
`scale_weight` are not silently ignored.
|
||||
- The pinned official prompt path enables thinking, tokenizes to 512, and consumes
|
||||
`hidden_states[-2]` at valid-token positions. A 36-block Qwen model exposes 37
|
||||
hidden-state entries in this contract (embedding/intermediate entries plus the
|
||||
final normalized state).
|
||||
- `TextEncoderConfig.chat_template_enable_thinking` is keyword-only and defaults
|
||||
to `False` for existing model families. A future Z-Image pipeline config must
|
||||
opt in with `True`.
|
||||
|
||||
### VAE
|
||||
|
||||
The test covers both direct implementation parity and production `VAELoader`
|
||||
resolution/strict loading. This component-only PR explicitly accepts the
|
||||
existing shared `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` wrapper,
|
||||
which subclasses Diffusers `AutoencoderKL`, as a narrowly scoped exception to
|
||||
the native-component boundary. This is not precedent for the future transformer
|
||||
or full Z-Image pipeline, and the exception must be revisited if that scope grows.
|
||||
|
||||
### Scheduler
|
||||
|
||||
The pinned official pipeline mutates `scheduler.sigma_min = 0.0` before building
|
||||
its `num_steps + 1` schedule. A future Z-Image pipeline config must set both
|
||||
`use_reference_discrete_timesteps=True` and `sigma_min=0.0`; the published
|
||||
`scheduler_config.json` alone does not encode the complete runtime contract.
|
||||
|
||||
## Remaining work
|
||||
|
||||
- Run every asset-backed component test non-skip against the pinned clone and HF
|
||||
snapshot before claiming component parity.
|
||||
- Port and validate `ZImageTransformer2DModel` in a separate PR.
|
||||
- Add the Z-Image pipeline/config/preset/registry/example, pipeline parity, and
|
||||
image-quality regression after all component gates pass.
|
||||
@@ -1 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -1,743 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Parity test for Z-Image Qwen3 text encoder support in FastVideo.
|
||||
|
||||
This compares:
|
||||
1) direct transformers AutoModel output, and
|
||||
2) FastVideo's production ``TextEncoderLoader`` passthrough, and
|
||||
3) FastVideo's shared native Qwen3 encoder (``Qwen3ForCausalLM``, added for Flux2
|
||||
Klein and reused here — Z-Image-Turbo's ``Qwen3Model`` checkpoint routes
|
||||
to it via the model registry),
|
||||
|
||||
using identical local checkpoint, tokenization, and inputs.
|
||||
|
||||
Usage:
|
||||
pytest tests/local_tests/zimage/test_zimage_encoder_parity.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def _reclaim_vram() -> None:
|
||||
"""Actually reclaim VRAM after the caller has ``del``'d its model refs.
|
||||
|
||||
HF transformer models hold reference cycles, so ``del`` alone does not free
|
||||
them — a ``gc.collect()`` is required before the CUDA caching allocator will
|
||||
release the blocks. Without this the parity test keeps two full encoder
|
||||
copies resident and OOMs on a 44 GB GPU.
|
||||
"""
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _gpu_mem_mb() -> float:
|
||||
if torch.cuda.is_available():
|
||||
return torch.cuda.memory_allocated() / (1024 ** 2)
|
||||
return 0.0
|
||||
from torch.testing import assert_close
|
||||
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
|
||||
from safetensors.torch import safe_open
|
||||
|
||||
from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
|
||||
from fastvideo.distributed.parallel_state import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
|
||||
from fastvideo.layers.quantization.absmax_fp8 import AbsMaxFP8MergedParameter
|
||||
from fastvideo.models.encoders.qwen3 import Qwen3ForCausalLM
|
||||
from fastvideo.models.loader.component_loader import TextEncoderLoader
|
||||
from fastvideo.models.registry import ModelRegistry
|
||||
|
||||
PARITY_SCOPE = "both"
|
||||
|
||||
# Strict-load contract: the shared Qwen3 encoder is body-only (embed_tokens +
|
||||
# layers + norm, no lm_head). Z-Image-Turbo ships a full Qwen3 checkpoint, so
|
||||
# `lm_head.weight` is the only key that may go unmatched; anything else means a
|
||||
# real silent drop. Enforced here in the test since the shared encoder's loader
|
||||
# is intentionally lenient (it's used by multiple models).
|
||||
_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"}
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
ZIMAGE_TEXT_ENCODER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "text_encoder"
|
||||
ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def _init_dist_and_tp_groups():
|
||||
os.environ.setdefault("MASTER_ADDR", "localhost")
|
||||
os.environ.setdefault("MASTER_PORT", "29531")
|
||||
os.environ.setdefault("RANK", "0")
|
||||
os.environ.setdefault("WORLD_SIZE", "1")
|
||||
os.environ.setdefault("LOCAL_RANK", "0")
|
||||
|
||||
maybe_init_distributed_environment_and_model_parallel(1, 1)
|
||||
yield
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def _load_json(path: Path) -> dict:
|
||||
with path.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _iter_safetensors(path: Path):
|
||||
with safe_open(str(path), framework="pt", device="cpu") as f:
|
||||
for k in f.keys():
|
||||
yield k, f.get_tensor(k)
|
||||
|
||||
|
||||
def _iter_pretrained_safetensors(model_dir: Path):
|
||||
single = model_dir / "model.safetensors"
|
||||
if single.exists():
|
||||
yield from _iter_safetensors(single)
|
||||
return
|
||||
|
||||
index = model_dir / "model.safetensors.index.json"
|
||||
if index.exists():
|
||||
idx = _load_json(index)
|
||||
shard_names = sorted(set(idx["weight_map"].values()))
|
||||
for shard in shard_names:
|
||||
yield from _iter_safetensors(model_dir / shard)
|
||||
return
|
||||
|
||||
raise FileNotFoundError(
|
||||
f"Missing safetensors checkpoint in {model_dir} (expected model.safetensors or model.safetensors.index.json)"
|
||||
)
|
||||
|
||||
|
||||
def _pretrained_safetensor_keys(model_dir: Path) -> set[str]:
|
||||
return {name for name, _ in _iter_pretrained_safetensors(model_dir)}
|
||||
|
||||
|
||||
def _load_qwen3_config() -> Qwen3TextConfig:
|
||||
cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
|
||||
for key in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
|
||||
cfg_raw.pop(key, None)
|
||||
config = Qwen3TextConfig()
|
||||
config.update_model_arch(cfg_raw)
|
||||
assert config.num_hidden_layers == 36
|
||||
assert config.hidden_size == 2560
|
||||
assert config.intermediate_size == 9728
|
||||
assert config.num_attention_heads == 32
|
||||
assert config.num_key_value_heads == 8
|
||||
return config
|
||||
|
||||
|
||||
def _loader_args(cpu_offload: bool) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
text_encoder_cpu_offload=cpu_offload,
|
||||
override_text_encoder_quant=None,
|
||||
override_text_encoder_safetensors=None,
|
||||
pin_cpu_memory=False,
|
||||
)
|
||||
|
||||
|
||||
def _precision_name(dtype: torch.dtype) -> str:
|
||||
if dtype == torch.float32:
|
||||
return "fp32"
|
||||
if dtype == torch.bfloat16:
|
||||
return "bf16"
|
||||
raise ValueError(f"Unsupported parity dtype: {dtype}")
|
||||
|
||||
|
||||
def _tokenize_official_prompts(tokenizer, prompts: list[str]):
|
||||
formatted = [
|
||||
tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": prompt}],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=True,
|
||||
)
|
||||
for prompt in prompts
|
||||
]
|
||||
return tokenizer(
|
||||
formatted,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
max_length=512,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
|
||||
def _assert_native_load_surface(
|
||||
model,
|
||||
checkpoint_keys: set[str],
|
||||
loaded_params: set[str],
|
||||
) -> None:
|
||||
param_names = {name for name, _ in model.named_parameters()}
|
||||
missing = param_names - loaded_params
|
||||
assert not missing, f"FastVideo parameters not loaded: {sorted(missing)}"
|
||||
|
||||
normalized_keys = {
|
||||
name[len("model."):] if name.startswith("model.") else name
|
||||
for name in checkpoint_keys
|
||||
}
|
||||
mapped_keys: set[str] = set()
|
||||
for checkpoint_name in normalized_keys:
|
||||
for param_name, weight_name, _ in model.config.arch_config.stacked_params_mapping:
|
||||
if weight_name in checkpoint_name:
|
||||
mapped_keys.add(checkpoint_name.replace(weight_name, param_name))
|
||||
break
|
||||
else:
|
||||
mapped_keys.add(checkpoint_name)
|
||||
|
||||
unexpected = mapped_keys - param_names - {
|
||||
"rotary_emb.inv_freq",
|
||||
"rotary_emb.cos_cached",
|
||||
"rotary_emb.sin_cached",
|
||||
}
|
||||
assert unexpected <= _ALLOWED_UNEXPECTED_KEYS, (
|
||||
"Unexpected checkpoint keys not in allowlist: "
|
||||
f"{sorted(unexpected - _ALLOWED_UNEXPECTED_KEYS)}"
|
||||
)
|
||||
|
||||
|
||||
def _fake_qwen3_weight_target(include_quant_scales: bool = False):
|
||||
stacked = Qwen3TextConfig().arch_config.stacked_params_mapping
|
||||
params = {
|
||||
"layers.0.self_attn.qkv_proj.weight": torch.nn.Parameter(torch.empty(1)),
|
||||
"layers.0.mlp.gate_up_proj.weight": torch.nn.Parameter(torch.empty(1)),
|
||||
}
|
||||
if include_quant_scales:
|
||||
qkv_scale = AbsMaxFP8MergedParameter(torch.zeros(3), requires_grad=False)
|
||||
qkv_scale.output_partition_sizes = [1, 1, 1]
|
||||
gate_up_scale = AbsMaxFP8MergedParameter(torch.zeros(2), requires_grad=False)
|
||||
gate_up_scale.output_partition_sizes = [1, 1]
|
||||
params.update({
|
||||
"layers.0.self_attn.qkv_proj.scale_weight": qkv_scale,
|
||||
"layers.0.mlp.gate_up_proj.scale_weight": gate_up_scale,
|
||||
})
|
||||
for param in params.values():
|
||||
if not hasattr(param, "weight_loader"):
|
||||
param.weight_loader = (
|
||||
lambda target, loaded, *args: target.data.copy_(loaded)
|
||||
)
|
||||
|
||||
model = SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
arch_config=SimpleNamespace(stacked_params_mapping=stacked)
|
||||
),
|
||||
named_parameters=lambda: params.items(),
|
||||
)
|
||||
return model, params, stacked
|
||||
|
||||
|
||||
# The provisional bf16 thresholds below come from the historical unpinned
|
||||
# 35-layer run. Per add-model-02-parity's calibration block, bf16 max-based
|
||||
# asserts are meaningless for deep encoders — fused-vs-unfused linear order
|
||||
# alone produces a long tail, so the test uses mean + median instead. The
|
||||
# pinned snapshot has 36 blocks and must be rerun before these thresholds are
|
||||
# treated as current evidence.
|
||||
#
|
||||
# Historical empirical numbers (unrecorded snapshot, 35 layers, CUDA):
|
||||
# last_hidden_state (post-norm): mean=0.0152, median=0.0117
|
||||
# hidden_states[-2] (pre-norm): mean=0.0754, median=0.0625
|
||||
# Thresholds below add ~1.6x headroom over observed.
|
||||
_BF16_MEAN_DRIFT_POST_NORM = 0.025
|
||||
_BF16_MEAN_DRIFT_PRE_NORM = 0.120
|
||||
_BF16_MEDIAN_DRIFT_POST_NORM = 0.020
|
||||
_BF16_MEDIAN_DRIFT_PRE_NORM = 0.100
|
||||
_FP32_ATOL = 1e-4
|
||||
_FP32_RTOL = 1e-4
|
||||
|
||||
|
||||
def _print_diag(label: str, ref: torch.Tensor, fv: torch.Tensor) -> torch.Tensor:
|
||||
diff = (ref.float() - fv.float()).abs()
|
||||
flat = diff.flatten()
|
||||
p99 = flat.kthvalue(max(1, int(0.99 * flat.numel()))).values
|
||||
print(
|
||||
f"[{label}] max_diff={diff.max():.4f} mean_diff={diff.mean():.4f} "
|
||||
f"median_diff={diff.median():.4f} p99_diff={p99:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
return diff
|
||||
|
||||
|
||||
def test_qwen3_production_loader_uses_body_only_model_and_honors_cpu_offload(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""The production passthrough must avoid the LM head and preserve CPU placement."""
|
||||
import fastvideo.platforms as platforms
|
||||
|
||||
placements: list[torch.device] = []
|
||||
auto_model_calls: list[dict] = []
|
||||
|
||||
class FakeHFModel:
|
||||
def eval(self):
|
||||
return self
|
||||
|
||||
def to(self, device: torch.device):
|
||||
placements.append(torch.device(device))
|
||||
return self
|
||||
|
||||
fake_model = FakeHFModel()
|
||||
|
||||
def fake_auto_model_from_pretrained(*args, **kwargs):
|
||||
auto_model_calls.append(kwargs)
|
||||
return fake_model
|
||||
|
||||
def fail_causal_lm_from_pretrained(*args, **kwargs):
|
||||
raise AssertionError("Qwen text encoding must not load AutoModelForCausalLM")
|
||||
|
||||
monkeypatch.setattr(
|
||||
AutoModel,
|
||||
"from_pretrained",
|
||||
staticmethod(fake_auto_model_from_pretrained),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
AutoModelForCausalLM,
|
||||
"from_pretrained",
|
||||
staticmethod(fail_causal_lm_from_pretrained),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
platforms,
|
||||
"_current_platform",
|
||||
SimpleNamespace(
|
||||
is_mps=lambda: False,
|
||||
verify_model_arch=lambda _arch: None,
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.distributed.get_local_torch_device",
|
||||
lambda: torch.device("cuda"),
|
||||
)
|
||||
|
||||
config = Qwen3TextConfig()
|
||||
# The pinned checkpoint names the causal wrapper even though the official
|
||||
# Z-Image loader intentionally asks Transformers for the body-only model.
|
||||
config.arch_config.architectures = ["Qwen3ForCausalLM"]
|
||||
loaded = TextEncoderLoader().load_model(
|
||||
"unused-by-mock",
|
||||
config,
|
||||
torch.device("cuda"),
|
||||
_loader_args(cpu_offload=True),
|
||||
dtype="fp32",
|
||||
use_text_encoder_override=True,
|
||||
)
|
||||
|
||||
assert loaded is fake_model
|
||||
assert auto_model_calls == [{
|
||||
"local_files_only": True,
|
||||
"torch_dtype": torch.float32,
|
||||
"low_cpu_mem_usage": True,
|
||||
}]
|
||||
assert placements == [torch.device("cpu")]
|
||||
assert loaded._fastvideo_input_device == torch.device("cpu")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_prefix", ["", "model."])
|
||||
def test_qwen3_native_load_accepts_already_fused_weights_and_scales(
|
||||
checkpoint_prefix: str,
|
||||
):
|
||||
"""Direct FastVideo/state-dict keys satisfy each fused target atomically."""
|
||||
fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True)
|
||||
expected = {
|
||||
name: torch.arange(1, param.numel() + 1, dtype=param.dtype)
|
||||
for name, param in params.items()
|
||||
}
|
||||
weights = [
|
||||
(f"{checkpoint_prefix}{name}", expected[name])
|
||||
for name in params
|
||||
]
|
||||
|
||||
loaded = Qwen3ForCausalLM.load_weights(fake_model, weights)
|
||||
|
||||
assert loaded == set(params)
|
||||
for name, param in params.items():
|
||||
assert torch.equal(param.data, expected[name])
|
||||
_assert_native_load_surface(
|
||||
fake_model,
|
||||
{name for name, _ in weights},
|
||||
loaded,
|
||||
)
|
||||
|
||||
|
||||
def test_qwen3_native_load_rejects_incompatible_fused_quant_scale_shape():
|
||||
fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True)
|
||||
weights = [
|
||||
(name, torch.ones_like(param))
|
||||
for name, param in params.items()
|
||||
if name != "layers.0.self_attn.qkv_proj.scale_weight"
|
||||
]
|
||||
weights.append(("layers.0.self_attn.qkv_proj.scale_weight", torch.ones(1)))
|
||||
|
||||
with pytest.raises(AssertionError, match="Attempted to load weight"):
|
||||
Qwen3ForCausalLM.load_weights(fake_model, weights)
|
||||
|
||||
|
||||
def test_qwen3_native_load_uses_real_shard_loader_for_split_quant_scales():
|
||||
fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True)
|
||||
weights = [
|
||||
(name, torch.ones_like(param))
|
||||
for name, param in params.items()
|
||||
if name.endswith(".weight")
|
||||
]
|
||||
weights.extend(
|
||||
(f"model.layers.0.self_attn{weight_name}.scale_weight", torch.tensor(float(index + 1)))
|
||||
for index, (_, weight_name, _) in enumerate(stacked[:3])
|
||||
)
|
||||
weights.extend(
|
||||
(f"model.layers.0.mlp{weight_name}.scale_weight", torch.tensor(float(index + 4)))
|
||||
for index, (_, weight_name, _) in enumerate(stacked[3:])
|
||||
)
|
||||
|
||||
loaded = Qwen3ForCausalLM.load_weights(fake_model, weights)
|
||||
|
||||
assert loaded == set(params)
|
||||
assert torch.equal(
|
||||
params["layers.0.self_attn.qkv_proj.scale_weight"].data,
|
||||
torch.tensor([1.0, 2.0, 3.0]),
|
||||
)
|
||||
assert torch.equal(
|
||||
params["layers.0.mlp.gate_up_proj.scale_weight"].data,
|
||||
torch.tensor([4.0, 5.0]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("missing_weight_name", "missing_shard_id"),
|
||||
[
|
||||
(".q_proj", "q"),
|
||||
(".k_proj", "k"),
|
||||
(".v_proj", "v"),
|
||||
(".gate_proj", 0),
|
||||
(".up_proj", 1),
|
||||
],
|
||||
)
|
||||
def test_qwen3_native_load_rejects_missing_fused_weight_shard(
|
||||
missing_weight_name: str,
|
||||
missing_shard_id: str | int,
|
||||
):
|
||||
fake_model, _, stacked = _fake_qwen3_weight_target()
|
||||
weights = [
|
||||
(f"model.layers.0.self_attn{weight_name}.weight", torch.empty(1))
|
||||
for _, weight_name, _ in stacked[:3]
|
||||
if weight_name != missing_weight_name
|
||||
]
|
||||
weights.extend(
|
||||
(f"model.layers.0.mlp{weight_name}.weight", torch.empty(1))
|
||||
for _, weight_name, _ in stacked[3:]
|
||||
if weight_name != missing_weight_name
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info:
|
||||
Qwen3ForCausalLM.load_weights(fake_model, weights)
|
||||
|
||||
destination = (
|
||||
"layers.0.self_attn.qkv_proj.weight"
|
||||
if isinstance(missing_shard_id, str)
|
||||
else "layers.0.mlp.gate_up_proj.weight"
|
||||
)
|
||||
assert f"{destination}[{missing_shard_id}]" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("missing_weight_name", "missing_shard_id"),
|
||||
[
|
||||
(".q_proj", "q"),
|
||||
(".k_proj", "k"),
|
||||
(".v_proj", "v"),
|
||||
(".gate_proj", 0),
|
||||
(".up_proj", 1),
|
||||
],
|
||||
)
|
||||
def test_qwen3_native_load_rejects_missing_fused_quant_scale_shard(
|
||||
missing_weight_name: str,
|
||||
missing_shard_id: str | int,
|
||||
):
|
||||
fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True)
|
||||
weights = [
|
||||
(name, torch.empty(1))
|
||||
for name in params
|
||||
if name.endswith(".weight")
|
||||
]
|
||||
weights.extend(
|
||||
(f"model.layers.0.self_attn{weight_name}.scale_weight", torch.empty(1))
|
||||
for _, weight_name, _ in stacked[:3]
|
||||
if weight_name != missing_weight_name
|
||||
)
|
||||
weights.extend(
|
||||
(f"model.layers.0.mlp{weight_name}.scale_weight", torch.empty(1))
|
||||
for _, weight_name, _ in stacked[3:]
|
||||
if weight_name != missing_weight_name
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info:
|
||||
Qwen3ForCausalLM.load_weights(fake_model, weights)
|
||||
|
||||
destination = (
|
||||
"layers.0.self_attn.qkv_proj.scale_weight"
|
||||
if isinstance(missing_shard_id, str)
|
||||
else "layers.0.mlp.gate_up_proj.scale_weight"
|
||||
)
|
||||
assert f"{destination}[{missing_shard_id}]" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required")
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
pytest.param(torch.float32, id="fp32"),
|
||||
pytest.param(
|
||||
torch.bfloat16,
|
||||
id="bf16",
|
||||
marks=pytest.mark.skipif(
|
||||
not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
|
||||
reason="bf16 parity requires a bf16-capable CUDA device",
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype):
|
||||
if not ZIMAGE_TOKENIZER_DIR.exists():
|
||||
pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
torch.manual_seed(11)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
|
||||
|
||||
prompts = [
|
||||
"A cinematic shot of a rainy neon street at night.",
|
||||
"A watercolor illustration of a fox in autumn leaves.",
|
||||
]
|
||||
toks = _tokenize_official_prompts(tokenizer, prompts)
|
||||
input_ids = toks["input_ids"].to(device)
|
||||
attention_mask = toks["attention_mask"].to(device)
|
||||
config = _load_qwen3_config()
|
||||
|
||||
ref = AutoModel.from_pretrained(
|
||||
str(ZIMAGE_TEXT_ENCODER_DIR),
|
||||
local_files_only=True,
|
||||
trust_remote_code=True,
|
||||
dtype=dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
).eval().to(device=device, dtype=dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
ref_out = ref(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
return_dict=True,
|
||||
)
|
||||
ref_last = ref_out.last_hidden_state.detach().float().cpu()
|
||||
assert ref_out.hidden_states is not None
|
||||
ref_hs_m2 = ref_out.hidden_states[-2].detach().float().cpu()
|
||||
|
||||
print(f"[mem] after HF ref forward: {_gpu_mem_mb():.0f} MiB", flush=True)
|
||||
# Keep only CPU outputs so each full encoder is resident one at a time.
|
||||
del ref, ref_out
|
||||
_reclaim_vram()
|
||||
print(f"[mem] after freeing HF ref: {_gpu_mem_mb():.0f} MiB", flush=True)
|
||||
|
||||
production = TextEncoderLoader().load_model(
|
||||
str(ZIMAGE_TEXT_ENCODER_DIR),
|
||||
config,
|
||||
device,
|
||||
_loader_args(cpu_offload=False),
|
||||
dtype=_precision_name(dtype),
|
||||
use_text_encoder_override=True,
|
||||
)
|
||||
assert production.__class__.__name__ == "Qwen3Model"
|
||||
assert not hasattr(production, "lm_head")
|
||||
assert production._fastvideo_input_device == device
|
||||
|
||||
with torch.no_grad():
|
||||
production_out = production(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
return_dict=True,
|
||||
)
|
||||
production_last = production_out.last_hidden_state.detach().float().cpu()
|
||||
assert production_out.hidden_states is not None
|
||||
production_hs_m2 = production_out.hidden_states[-2].detach().float().cpu()
|
||||
|
||||
# Production loading must preserve the independent Transformers oracle
|
||||
# exactly; native-kernel drift is assessed separately below.
|
||||
assert_close(ref_last, production_last, atol=0.0, rtol=0.0)
|
||||
assert_close(ref_hs_m2, production_hs_m2, atol=0.0, rtol=0.0)
|
||||
del production, production_out, production_last, production_hs_m2
|
||||
_reclaim_vram()
|
||||
|
||||
fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model")
|
||||
fv = fv_cls(config).eval()
|
||||
loaded = fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
|
||||
assert loaded, "No Qwen3 weights were loaded into FastVideo model"
|
||||
fv = fv.to(device=device, dtype=dtype)
|
||||
print(f"[mem] after FastVideo load: {_gpu_mem_mb():.0f} MiB", flush=True)
|
||||
|
||||
_assert_native_load_surface(
|
||||
fv,
|
||||
_pretrained_safetensor_keys(ZIMAGE_TEXT_ENCODER_DIR),
|
||||
loaded,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
fv_out = fv(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
fv_last = fv_out.last_hidden_state.detach().float().cpu()
|
||||
assert fv_out.hidden_states is not None
|
||||
fv_hs_m2 = fv_out.hidden_states[-2].detach().float().cpu()
|
||||
|
||||
# Z-Image uses only valid token embeddings selected by attention mask.
|
||||
# Compare parity on that exact path to avoid padded-token artifacts.
|
||||
mask_cpu = attention_mask.detach().bool().cpu()
|
||||
is_bf16 = dtype == torch.bfloat16
|
||||
|
||||
for i in range(mask_cpu.shape[0]):
|
||||
valid = mask_cpu[i]
|
||||
ref_last_v = ref_last[i][valid]
|
||||
fv_last_v = fv_last[i][valid]
|
||||
ref_hs_m2_v = ref_hs_m2[i][valid]
|
||||
fv_hs_m2_v = fv_hs_m2[i][valid]
|
||||
|
||||
last_diff = _print_diag(f"batch{i} last_hidden_state {dtype}", ref_last_v, fv_last_v)
|
||||
hs_m2_diff = _print_diag(f"batch{i} hidden_states[-2] {dtype}", ref_hs_m2_v, fv_hs_m2_v)
|
||||
|
||||
if is_bf16:
|
||||
# Distribution checks instead of element-wise assert_close: cross-
|
||||
# kernel bf16 produces a long max tail (fused vs unfused linears),
|
||||
# but a real op bug pushes mean *and* median high. Mean catches
|
||||
# systematic drift; median catches "wrong weights / swapped
|
||||
# layers" where most elements diverge. Pre-norm tensors get a
|
||||
# looser bound because the final RMSNorm compresses drift ~5x.
|
||||
assert last_diff.mean().item() < _BF16_MEAN_DRIFT_POST_NORM, (
|
||||
f"batch{i} last_hidden_state bf16 mean drift "
|
||||
f"{last_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_POST_NORM}"
|
||||
)
|
||||
assert last_diff.median().item() < _BF16_MEDIAN_DRIFT_POST_NORM, (
|
||||
f"batch{i} last_hidden_state bf16 median drift "
|
||||
f"{last_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_POST_NORM}"
|
||||
)
|
||||
assert hs_m2_diff.mean().item() < _BF16_MEAN_DRIFT_PRE_NORM, (
|
||||
f"batch{i} hidden_states[-2] bf16 mean drift "
|
||||
f"{hs_m2_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_PRE_NORM}"
|
||||
)
|
||||
assert hs_m2_diff.median().item() < _BF16_MEDIAN_DRIFT_PRE_NORM, (
|
||||
f"batch{i} hidden_states[-2] bf16 median drift "
|
||||
f"{hs_m2_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_PRE_NORM}"
|
||||
)
|
||||
else:
|
||||
assert_close(ref_last_v, fv_last_v, atol=_FP32_ATOL, rtol=_FP32_RTOL)
|
||||
assert_close(ref_hs_m2_v, fv_hs_m2_v, atol=_FP32_ATOL, rtol=_FP32_RTOL)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required")
|
||||
@pytest.mark.skipif(
|
||||
not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
|
||||
reason="bf16 per-layer diagnostic requires a bf16-capable CUDA device",
|
||||
)
|
||||
def test_zimage_qwen3_encoder_per_layer_bf16_diagnostic():
|
||||
"""Diagnostic-only: prints per-layer FastVideo-vs-HF drift for the bf16
|
||||
forward to distinguish monotonic accumulation (bf16 tail) from a
|
||||
single-layer spike (real op mismatch).
|
||||
|
||||
This test always passes; review the captured stdout to interpret.
|
||||
Expected pattern for healthy bf16 accumulation: max/mean/median grow
|
||||
smoothly with layer depth. Suspect pattern: one layer's diff is
|
||||
>>2x the previous, suggesting a real bug at that block.
|
||||
"""
|
||||
if not ZIMAGE_TOKENIZER_DIR.exists():
|
||||
pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
|
||||
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
torch.manual_seed(11)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
|
||||
toks = _tokenize_official_prompts(
|
||||
tokenizer,
|
||||
["A cinematic shot of a rainy neon street at night."],
|
||||
)
|
||||
input_ids = toks["input_ids"].to(device)
|
||||
attention_mask = toks["attention_mask"].to(device)
|
||||
|
||||
ref = AutoModel.from_pretrained(
|
||||
str(ZIMAGE_TEXT_ENCODER_DIR),
|
||||
local_files_only=True,
|
||||
trust_remote_code=True,
|
||||
dtype=dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
).eval().to(device=device, dtype=dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
ref_out = ref(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
return_dict=True,
|
||||
)
|
||||
# Cache the reference hidden states on CPU, then free the HF model so
|
||||
# only one encoder is resident when FastVideo runs (avoids OOM).
|
||||
ref_hidden_cpu = [h.detach().float().cpu() for h in ref_out.hidden_states]
|
||||
del ref, ref_out
|
||||
_reclaim_vram()
|
||||
|
||||
fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model")
|
||||
cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
|
||||
for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
|
||||
cfg_raw.pop(k, None)
|
||||
cfg = Qwen3TextConfig()
|
||||
cfg.update_model_arch(cfg_raw)
|
||||
fv = fv_cls(cfg).eval()
|
||||
fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
|
||||
fv = fv.to(device=device, dtype=dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
fv_out = fv(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
fv_hidden_cpu = [h.detach().float().cpu() for h in fv_out.hidden_states]
|
||||
|
||||
assert fv_out.hidden_states is not None
|
||||
assert len(ref_hidden_cpu) == 37, (
|
||||
f"expected embeddings plus 36 block outputs, got {len(ref_hidden_cpu)} entries"
|
||||
)
|
||||
assert len(ref_hidden_cpu) == len(fv_hidden_cpu), (
|
||||
f"len mismatch: HF={len(ref_hidden_cpu)} FV={len(fv_hidden_cpu)}"
|
||||
)
|
||||
|
||||
mask_cpu = attention_mask.detach().bool().cpu()[0]
|
||||
print(
|
||||
f"\n[per-layer bf16 diagnostic] num_hidden_states={len(ref_hidden_cpu)} "
|
||||
f"valid_tokens={mask_cpu.sum().item()}/{mask_cpu.numel()}",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f"{'idx':>3} {'kind':<10} {'max':>8} {'mean':>8} {'median':>8} {'p99':>8}",
|
||||
flush=True,
|
||||
)
|
||||
for idx, (ref_h, fv_h) in enumerate(zip(ref_hidden_cpu, fv_hidden_cpu)):
|
||||
ref_v = ref_h[0][mask_cpu]
|
||||
fv_v = fv_h[0][mask_cpu]
|
||||
diff = (ref_v - fv_v).abs()
|
||||
flat = diff.flatten()
|
||||
p99 = torch.quantile(flat, 0.99).item()
|
||||
if idx == 0:
|
||||
kind = "embedding"
|
||||
elif idx == len(ref_hidden_cpu) - 1:
|
||||
kind = "post-norm"
|
||||
else:
|
||||
kind = f"layer{idx - 1}-out"
|
||||
print(
|
||||
f"{idx:>3} {kind:<10} "
|
||||
f"{diff.max().item():>8.4f} {diff.mean().item():>8.4f} "
|
||||
f"{diff.median().item():>8.4f} {p99:>8.4f}",
|
||||
flush=True,
|
||||
)
|
||||
@@ -1,227 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Parity tests for Z-Image FlowMatchEulerDiscreteScheduler.
|
||||
|
||||
These tests compare FastVideo's scheduler implementation against the
|
||||
reference scheduler shipped in the local Z-Image repository.
|
||||
|
||||
Usage:
|
||||
pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch.testing import assert_close
|
||||
|
||||
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
|
||||
FlowMatchEulerDiscreteScheduler as FastVideoScheduler,
|
||||
)
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
ZIMAGE_REPO = REPO_ROOT / "Z-Image"
|
||||
ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src"
|
||||
ZIMAGE_SCHEDULER_CFG = (
|
||||
REPO_ROOT / "official_weights" / "Z-Image" / "scheduler" / "scheduler_config.json"
|
||||
)
|
||||
ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd"
|
||||
PARITY_SCOPE = "implementation_subcomponent"
|
||||
|
||||
|
||||
def _require_pinned_reference_module(module_name: str, source_file: Path):
|
||||
if not ZIMAGE_REPO.exists():
|
||||
pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}")
|
||||
if not source_file.is_file():
|
||||
pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
except (FileNotFoundError, subprocess.CalledProcessError) as exc:
|
||||
pytest.fail(f"Cannot verify Z-Image reference revision: {exc}")
|
||||
|
||||
actual_revision = result.stdout.strip()
|
||||
assert actual_revision == ZIMAGE_REFERENCE_REVISION, (
|
||||
"Z-Image reference clone is not at the pinned revision: "
|
||||
f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}"
|
||||
)
|
||||
|
||||
if str(ZIMAGE_SRC) not in sys.path:
|
||||
sys.path.insert(0, str(ZIMAGE_SRC))
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
except Exception as exc:
|
||||
pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}")
|
||||
|
||||
module_file = Path(module.__file__ or "").resolve()
|
||||
assert module_file.is_relative_to(ZIMAGE_SRC.resolve()), (
|
||||
f"{module_name} resolved outside the pinned clone: {module_file}"
|
||||
)
|
||||
return module
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def reference_scheduler_cls():
|
||||
module = _require_pinned_reference_module(
|
||||
"zimage.scheduler",
|
||||
ZIMAGE_SRC / "zimage" / "scheduler.py",
|
||||
)
|
||||
return module.FlowMatchEulerDiscreteScheduler
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def scheduler_config() -> dict:
|
||||
if not ZIMAGE_SCHEDULER_CFG.exists():
|
||||
pytest.skip(f"Z-Image scheduler config not found: {ZIMAGE_SCHEDULER_CFG}")
|
||||
with ZIMAGE_SCHEDULER_CFG.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
# Diffusers-style scheduler-config keys we do NOT want to forward — they are
|
||||
# loader/state metadata, not constructor kwargs. Anything else in
|
||||
# scheduler_config.json (including future keys like `time_shift_type`,
|
||||
# `invert_sigmas`, etc.) must be forwarded so parity actually exercises the
|
||||
# config-on-disk and not a hand-curated subset.
|
||||
_SCHEDULER_CONFIG_LOADER_KEYS = frozenset({"_class_name", "_diffusers_version", "_name_or_path"})
|
||||
|
||||
|
||||
def _scheduler_kwargs_from_config(scheduler_config: dict) -> dict:
|
||||
return {k: v for k, v in scheduler_config.items() if k not in _SCHEDULER_CONFIG_LOADER_KEYS}
|
||||
|
||||
|
||||
def _make_scheduler_pair(reference_scheduler_cls, scheduler_config: dict, **overrides):
|
||||
kwargs = _scheduler_kwargs_from_config(scheduler_config)
|
||||
kwargs.update(overrides)
|
||||
|
||||
ref = reference_scheduler_cls(**kwargs)
|
||||
# The pinned Z-Image pipeline mutates this immediately before scheduling.
|
||||
ref.sigma_min = 0.0
|
||||
|
||||
# These values must be serialized by the future production pipeline config.
|
||||
# setdefault also allows this test to consume that config once it lands.
|
||||
fv_kwargs = dict(kwargs)
|
||||
fv_kwargs.setdefault("use_reference_discrete_timesteps", True)
|
||||
fv_kwargs.setdefault("sigma_min", 0.0)
|
||||
fv = FastVideoScheduler(**fv_kwargs)
|
||||
return ref, fv
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def scheduler_pair(reference_scheduler_cls, scheduler_config: dict):
|
||||
return _make_scheduler_pair(reference_scheduler_cls, scheduler_config)
|
||||
|
||||
|
||||
def _run_step_loop(ref_scheduler, fv_scheduler, sample_shape=(2, 4, 32, 32)):
|
||||
torch.manual_seed(123)
|
||||
|
||||
ref_sample = torch.randn(sample_shape, dtype=torch.float32)
|
||||
fv_sample = ref_sample.clone()
|
||||
|
||||
for t in ref_scheduler.timesteps:
|
||||
model_output = torch.randn_like(ref_sample)
|
||||
|
||||
ref_next = ref_scheduler.step(
|
||||
model_output,
|
||||
t,
|
||||
ref_sample,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
fv_next = fv_scheduler.step(
|
||||
model_output,
|
||||
t,
|
||||
fv_sample,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
assert_close(ref_next, fv_next, atol=1e-6, rtol=1e-6)
|
||||
|
||||
ref_sample = ref_next
|
||||
fv_sample = fv_next
|
||||
|
||||
|
||||
def test_zimage_scheduler_parity_default_schedule_and_step(scheduler_pair):
|
||||
ref, fv = scheduler_pair
|
||||
|
||||
num_inference_steps = 8
|
||||
ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu")
|
||||
fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu")
|
||||
|
||||
assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6)
|
||||
assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6)
|
||||
|
||||
_run_step_loop(ref, fv)
|
||||
|
||||
|
||||
def test_scheduler_positional_args_keep_existing_bindings_and_default_schedule():
|
||||
scheduler = FastVideoScheduler(1000, 1.0, False, 0.7)
|
||||
|
||||
assert scheduler.config.base_shift == 0.7
|
||||
assert scheduler.config.use_reference_discrete_timesteps is False
|
||||
|
||||
scheduler.set_timesteps(num_inference_steps=8, device="cpu")
|
||||
assert_close(scheduler.sigmas[-2], torch.tensor(scheduler.sigma_min))
|
||||
|
||||
|
||||
def test_scheduler_honors_explicit_zero_sigma_min():
|
||||
scheduler = FastVideoScheduler(
|
||||
sigma_min=0.0,
|
||||
use_reference_discrete_timesteps=True,
|
||||
)
|
||||
|
||||
assert scheduler.sigma_min == 0.0
|
||||
scheduler.set_timesteps(num_inference_steps=8, device="cpu")
|
||||
assert_close(scheduler.timesteps[-1], torch.tensor(125.0))
|
||||
|
||||
|
||||
def test_reference_schedule_preserves_float64_linspace_rounding():
|
||||
scheduler = FastVideoScheduler(
|
||||
sigma_min=0.0,
|
||||
use_reference_discrete_timesteps=True,
|
||||
)
|
||||
scheduler.set_timesteps(num_inference_steps=9, device="cpu")
|
||||
|
||||
expected = torch.from_numpy(
|
||||
np.linspace(1000.0, 0.0, 10)[:-1] / 1000.0,
|
||||
).to(torch.float32) * 1000.0
|
||||
float32_regression = torch.from_numpy(
|
||||
np.linspace(1000.0, 0.0, 10, dtype=np.float32)[:-1] / np.float32(1000.0),
|
||||
) * 1000.0
|
||||
|
||||
# The fifth timestep differs by one float32 ULP depending on where the
|
||||
# linspace is rounded, so this assertion fails if dtype=np.float32 returns.
|
||||
assert expected[4].item() != float32_regression[4].item()
|
||||
assert scheduler.timesteps[4].item() == expected[4].item()
|
||||
|
||||
|
||||
def test_zimage_scheduler_parity_dynamic_shifting_with_mu(
|
||||
reference_scheduler_cls,
|
||||
scheduler_config: dict,
|
||||
):
|
||||
ref, fv = _make_scheduler_pair(
|
||||
reference_scheduler_cls,
|
||||
scheduler_config,
|
||||
use_dynamic_shifting=True,
|
||||
)
|
||||
|
||||
mu = 0.75
|
||||
num_inference_steps = 8
|
||||
ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu)
|
||||
fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu)
|
||||
|
||||
assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6)
|
||||
assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6)
|
||||
|
||||
_run_step_loop(ref, fv)
|
||||
@@ -1,130 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Parity test for Z-Image tokenizer loading and tokenization behavior.
|
||||
|
||||
This compares:
|
||||
1) direct transformers AutoTokenizer loading from local tokenizer dir, and
|
||||
2) FastVideo TokenizerLoader loading path,
|
||||
|
||||
using the exact chat-template and tokenization settings from the pinned
|
||||
Z-Image pipeline.
|
||||
|
||||
Usage:
|
||||
pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.testing import assert_close
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from fastvideo.models.loader.component_loader import TokenizerLoader
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer"
|
||||
PARITY_SCOPE = "production_loader"
|
||||
OFFICIAL_MAX_SEQUENCE_LENGTH = 512
|
||||
OFFICIAL_CHAT_TEMPLATE_KWARGS = {
|
||||
"tokenize": False,
|
||||
"add_generation_prompt": True,
|
||||
"enable_thinking": True,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DummyPipelineConfig:
|
||||
text_encoder_configs: tuple = field(default_factory=tuple)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DummyFastVideoArgs:
|
||||
pipeline_config: _DummyPipelineConfig = field(default_factory=_DummyPipelineConfig)
|
||||
|
||||
|
||||
def _load_reference_tokenizer():
|
||||
if not ZIMAGE_TOKENIZER_DIR.exists():
|
||||
pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
|
||||
return AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
|
||||
|
||||
|
||||
def _load_fastvideo_tokenizer():
|
||||
if not ZIMAGE_TOKENIZER_DIR.exists():
|
||||
pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
|
||||
loader = TokenizerLoader()
|
||||
return loader.load(str(ZIMAGE_TOKENIZER_DIR), _DummyFastVideoArgs())
|
||||
|
||||
|
||||
def test_zimage_tokenizer_loader_and_tokenization_parity():
|
||||
ref_tok = _load_reference_tokenizer()
|
||||
fv_tok = _load_fastvideo_tokenizer()
|
||||
|
||||
prompts = [
|
||||
"A close-up portrait of a cat in warm studio lighting.",
|
||||
"An oil painting of a lighthouse at sunset.",
|
||||
]
|
||||
|
||||
tok_kwargs = {
|
||||
"padding": "max_length",
|
||||
"max_length": OFFICIAL_MAX_SEQUENCE_LENGTH,
|
||||
"truncation": True,
|
||||
"return_tensors": "pt",
|
||||
}
|
||||
|
||||
ref_out = ref_tok(prompts, **tok_kwargs)
|
||||
fv_out = fv_tok(prompts, **tok_kwargs)
|
||||
|
||||
assert_close(ref_out["input_ids"], fv_out["input_ids"], atol=0, rtol=0)
|
||||
assert_close(ref_out["attention_mask"], fv_out["attention_mask"], atol=0, rtol=0)
|
||||
|
||||
# Compare key tokenizer attributes that affect generation-time behavior.
|
||||
assert ref_tok.padding_side == fv_tok.padding_side
|
||||
assert ref_tok.pad_token_id == fv_tok.pad_token_id
|
||||
assert ref_tok.eos_token_id == fv_tok.eos_token_id
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ZIMAGE_TOKENIZER_DIR.exists(), reason="Z-Image tokenizer assets are required")
|
||||
def test_zimage_tokenizer_chat_template_parity():
|
||||
ref_tok = _load_reference_tokenizer()
|
||||
fv_tok = _load_fastvideo_tokenizer()
|
||||
|
||||
assert callable(getattr(ref_tok, "apply_chat_template", None)), (
|
||||
"Pinned Z-Image tokenizer must expose apply_chat_template"
|
||||
)
|
||||
assert callable(getattr(fv_tok, "apply_chat_template", None)), (
|
||||
"FastVideo TokenizerLoader dropped the required apply_chat_template API"
|
||||
)
|
||||
|
||||
messages = [{"role": "user", "content": "Describe a futuristic city skyline."}]
|
||||
|
||||
ref_text = ref_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS)
|
||||
fv_text = fv_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS)
|
||||
|
||||
assert ref_text == fv_text
|
||||
|
||||
# Z-Image deliberately enables Qwen's thinking template. Guard against a
|
||||
# pipeline silently using the generic non-thinking prompt path.
|
||||
non_thinking_text = fv_tok.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
assert fv_text != non_thinking_text
|
||||
|
||||
tokenization_kwargs = {
|
||||
"padding": "max_length",
|
||||
"max_length": OFFICIAL_MAX_SEQUENCE_LENGTH,
|
||||
"truncation": True,
|
||||
"return_tensors": "pt",
|
||||
}
|
||||
ref_tokens = ref_tok(ref_text, **tokenization_kwargs)
|
||||
fv_tokens = fv_tok(fv_text, **tokenization_kwargs)
|
||||
|
||||
assert_close(ref_tokens["input_ids"], fv_tokens["input_ids"], atol=0, rtol=0)
|
||||
assert_close(ref_tokens["attention_mask"], fv_tokens["attention_mask"], atol=0, rtol=0)
|
||||
@@ -1,226 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Parity test for Z-Image VAE (AutoencoderKL decode path).
|
||||
|
||||
This compares FastVideo's AutoencoderKL wrapper against the Z-Image
|
||||
reference AutoencoderKL implementation using the same local weights.
|
||||
|
||||
Usage:
|
||||
pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import load_file as safetensors_load_file
|
||||
from torch.testing import assert_close
|
||||
|
||||
from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
|
||||
from fastvideo.models.loader.component_loader import VAELoader
|
||||
from fastvideo.models.vaes.autoencoder_kl import AutoencoderKL as FastVideoAutoencoderKL
|
||||
from fastvideo.pipelines.stages.decoding import DecodingStage
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
ZIMAGE_REPO = REPO_ROOT / "Z-Image"
|
||||
ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src"
|
||||
ZIMAGE_VAE_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "vae"
|
||||
ZIMAGE_VAE_CFG = ZIMAGE_VAE_DIR / "config.json"
|
||||
ZIMAGE_VAE_WEIGHTS = ZIMAGE_VAE_DIR / "diffusion_pytorch_model.safetensors"
|
||||
ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd"
|
||||
ZIMAGE_VAE_SCALING_FACTOR = 0.3611
|
||||
ZIMAGE_VAE_SHIFT_FACTOR = 0.1159
|
||||
PARITY_SCOPE = "both"
|
||||
|
||||
|
||||
def _require_pinned_reference_module(module_name: str, source_file: Path):
|
||||
if not ZIMAGE_REPO.exists():
|
||||
pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}")
|
||||
if not source_file.is_file():
|
||||
pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
except (FileNotFoundError, subprocess.CalledProcessError) as exc:
|
||||
pytest.fail(f"Cannot verify Z-Image reference revision: {exc}")
|
||||
|
||||
actual_revision = result.stdout.strip()
|
||||
assert actual_revision == ZIMAGE_REFERENCE_REVISION, (
|
||||
"Z-Image reference clone is not at the pinned revision: "
|
||||
f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}"
|
||||
)
|
||||
|
||||
if str(ZIMAGE_SRC) not in sys.path:
|
||||
sys.path.insert(0, str(ZIMAGE_SRC))
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
except Exception as exc:
|
||||
pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}")
|
||||
|
||||
module_file = Path(module.__file__ or "").resolve()
|
||||
assert module_file.is_relative_to(ZIMAGE_SRC.resolve()), (
|
||||
f"{module_name} resolved outside the pinned clone: {module_file}"
|
||||
)
|
||||
return module
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def reference_autoencoder_cls():
|
||||
module = _require_pinned_reference_module(
|
||||
"zimage.autoencoder",
|
||||
ZIMAGE_SRC / "zimage" / "autoencoder.py",
|
||||
)
|
||||
return module.AutoencoderKL
|
||||
|
||||
|
||||
def _load_cfg() -> dict:
|
||||
if not ZIMAGE_VAE_CFG.exists():
|
||||
pytest.skip(f"Z-Image VAE config not found: {ZIMAGE_VAE_CFG}")
|
||||
with ZIMAGE_VAE_CFG.open("r", encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
cfg.pop("_class_name", None)
|
||||
cfg.pop("_diffusers_version", None)
|
||||
cfg.pop("_name_or_path", None)
|
||||
return cfg
|
||||
|
||||
|
||||
def _load_weights() -> dict[str, torch.Tensor]:
|
||||
if not ZIMAGE_VAE_WEIGHTS.exists():
|
||||
pytest.skip(f"Z-Image VAE weights not found: {ZIMAGE_VAE_WEIGHTS}")
|
||||
return safetensors_load_file(str(ZIMAGE_VAE_WEIGHTS), device="cpu")
|
||||
|
||||
|
||||
def _build_reference(reference_autoencoder_cls, cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module:
|
||||
ref = reference_autoencoder_cls(
|
||||
in_channels=cfg["in_channels"],
|
||||
out_channels=cfg["out_channels"],
|
||||
down_block_types=tuple(cfg["down_block_types"]),
|
||||
up_block_types=tuple(cfg["up_block_types"]),
|
||||
block_out_channels=tuple(cfg["block_out_channels"]),
|
||||
layers_per_block=cfg["layers_per_block"],
|
||||
latent_channels=cfg["latent_channels"],
|
||||
norm_num_groups=cfg["norm_num_groups"],
|
||||
scaling_factor=cfg["scaling_factor"],
|
||||
shift_factor=cfg["shift_factor"],
|
||||
use_quant_conv=cfg["use_quant_conv"],
|
||||
use_post_quant_conv=cfg["use_post_quant_conv"],
|
||||
mid_block_add_attention=cfg["mid_block_add_attention"],
|
||||
).eval()
|
||||
ref.load_state_dict(weights, strict=True)
|
||||
return ref
|
||||
|
||||
|
||||
def _build_fastvideo(cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module:
|
||||
fv_cfg = AutoencoderKLVAEConfig()
|
||||
fv_cfg.update_model_arch(cfg)
|
||||
fv = FastVideoAutoencoderKL(fv_cfg).eval()
|
||||
fv.load_state_dict(weights, strict=True)
|
||||
return fv
|
||||
|
||||
|
||||
def _load_fastvideo_production_vae(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.models.loader.component_loader.get_local_torch_device",
|
||||
lambda: torch.device("cpu"),
|
||||
)
|
||||
pipeline_config = SimpleNamespace(
|
||||
vae_config=AutoencoderKLVAEConfig(),
|
||||
vae_precision="fp32",
|
||||
)
|
||||
fastvideo_args = SimpleNamespace(
|
||||
model_paths={},
|
||||
pipeline_config=pipeline_config,
|
||||
vae_cpu_offload=False,
|
||||
)
|
||||
vae = VAELoader().load(str(ZIMAGE_VAE_DIR), fastvideo_args)
|
||||
assert isinstance(vae, FastVideoAutoencoderKL)
|
||||
assert fastvideo_args.model_paths["vae"] == str(ZIMAGE_VAE_DIR)
|
||||
return vae
|
||||
|
||||
|
||||
def _official_decode_latents(latents: torch.Tensor, vae: torch.nn.Module) -> torch.Tensor:
|
||||
scaling_factor = vae.config.scaling_factor
|
||||
shift_factor = vae.config.shift_factor or 0.0
|
||||
dtype = next(vae.parameters()).dtype
|
||||
return (latents.to(dtype=dtype) / scaling_factor) + shift_factor
|
||||
|
||||
|
||||
def test_zimage_decode_normalization_uses_autoencoder_config_values():
|
||||
vae = SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
latents_mean=None,
|
||||
latents_std=None,
|
||||
scaling_factor=ZIMAGE_VAE_SCALING_FACTOR,
|
||||
shift_factor=ZIMAGE_VAE_SHIFT_FACTOR,
|
||||
)
|
||||
)
|
||||
latents = torch.tensor([[[[0.0, 0.5], [-0.5, 1.0]]]], dtype=torch.float32)
|
||||
|
||||
actual = DecodingStage(vae)._denormalize_latents(latents)
|
||||
expected = latents / ZIMAGE_VAE_SCALING_FACTOR + ZIMAGE_VAE_SHIFT_FACTOR
|
||||
|
||||
assert_close(actual, expected, atol=0, rtol=0)
|
||||
|
||||
|
||||
def test_zimage_vae_decode_parity(reference_autoencoder_cls):
|
||||
torch.manual_seed(7)
|
||||
|
||||
cfg = _load_cfg()
|
||||
weights = _load_weights()
|
||||
|
||||
ref = _build_reference(reference_autoencoder_cls, cfg, weights)
|
||||
fv = _build_fastvideo(cfg, weights)
|
||||
|
||||
# Decode-only parity is the critical path for Z-Image pipeline integration.
|
||||
latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32)
|
||||
|
||||
with torch.no_grad():
|
||||
ref_out = ref.decode(latents, return_dict=False)[0].detach().float()
|
||||
fv_out = fv.decode(latents, return_dict=False)[0].detach().float()
|
||||
|
||||
assert ref_out.shape == fv_out.shape
|
||||
assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4)
|
||||
|
||||
|
||||
def test_zimage_vae_production_loader_and_decode_normalization(
|
||||
reference_autoencoder_cls,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
torch.manual_seed(17)
|
||||
|
||||
cfg = _load_cfg()
|
||||
weights = _load_weights()
|
||||
assert cfg["scaling_factor"] == ZIMAGE_VAE_SCALING_FACTOR
|
||||
assert cfg["shift_factor"] == ZIMAGE_VAE_SHIFT_FACTOR
|
||||
ref = _build_reference(reference_autoencoder_cls, cfg, weights)
|
||||
fv = _load_fastvideo_production_vae(monkeypatch)
|
||||
|
||||
assert ref.config.scaling_factor == cfg["scaling_factor"]
|
||||
assert fv.config.scaling_factor == cfg["scaling_factor"]
|
||||
assert ref.config.shift_factor == cfg["shift_factor"]
|
||||
assert fv.config.shift_factor == cfg["shift_factor"]
|
||||
|
||||
pipeline_latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32)
|
||||
ref_decode_latents = _official_decode_latents(pipeline_latents, ref)
|
||||
fv_decode_latents = DecodingStage(fv)._denormalize_latents(pipeline_latents)
|
||||
assert_close(ref_decode_latents, fv_decode_latents, atol=0, rtol=0)
|
||||
|
||||
with torch.no_grad():
|
||||
ref_out = ref.decode(ref_decode_latents, return_dict=False)[0].detach().float()
|
||||
fv_out = fv.decode(fv_decode_latents, return_dict=False)[0].detach().float()
|
||||
|
||||
assert ref_out.shape == fv_out.shape
|
||||
assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4)
|
||||
Reference in New Issue
Block a user