Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
af4a8d093f | ||
|
|
0f1f6caea1 | ||
|
|
acc8128878 | ||
|
|
066d4118c0 | ||
|
|
716ce80a6f | ||
|
|
1746d866bc | ||
|
|
91e211b192 |
@@ -69,7 +69,7 @@ approval, then upload reviewed accepted baseline records.
|
||||
| `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
|
||||
| `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
|
||||
| `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
|
||||
| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
|
||||
| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `PERF_MAX_REGRESSION` if set, otherwise `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. |
|
||||
|
||||
Hardcoded defaults:
|
||||
@@ -148,7 +148,8 @@ For each metric with at least two non-null source values:
|
||||
4. Stop if any source record regresses against the batch median by more than
|
||||
`max_intra_batch_regression`.
|
||||
|
||||
Default `max_intra_batch_regression` to `0.05`. Print a table with per-source values, batch median, and
|
||||
Default `max_intra_batch_regression` to `PERF_MAX_REGRESSION` when set,
|
||||
otherwise `0.05`. Print a table with per-source values, batch median, and
|
||||
worst intra-batch regression.
|
||||
|
||||
This check prevents uploading a mixed batch where one JSON is materially
|
||||
@@ -182,7 +183,7 @@ present, that run is not a valid source for baseline reseeding.
|
||||
|
||||
### 2. Sync and back up existing HF records under /tmp
|
||||
|
||||
Use `fastvideo/performance/hf_store.py` helpers directly. Do **not** use
|
||||
Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
|
||||
`compare_baseline.py` as a sync shortcut; on full main runs it can persist
|
||||
records, while this step must only fetch and back up existing history.
|
||||
|
||||
@@ -191,7 +192,7 @@ The sync command pattern is:
|
||||
```bash
|
||||
export PERFORMANCE_TRACKING_ROOT="${PERFORMANCE_TRACKING_ROOT:-/tmp/perf-tracking}"
|
||||
export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
|
||||
python -c 'from fastvideo.performance.hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
|
||||
PYTHONPATH=fastvideo/tests/performance python -c 'from hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
|
||||
```
|
||||
|
||||
Then back up only the sanitized model directory under `/tmp`:
|
||||
@@ -199,8 +200,8 @@ Then back up only the sanitized model directory under `/tmp`:
|
||||
```bash
|
||||
SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
|
||||
TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)
|
||||
MODEL_SAFE=$(python - <<'PY'
|
||||
from fastvideo.performance.hf_store import sanitize
|
||||
MODEL_SAFE=$(PYTHONPATH=fastvideo/tests/performance python - <<'PY'
|
||||
from hf_store import sanitize
|
||||
print(sanitize("<model_id>"))
|
||||
PY
|
||||
)
|
||||
@@ -234,7 +235,7 @@ first baseline seed. Continue, but report that baseline history was empty.
|
||||
Load the last 5 successful records for the target:
|
||||
|
||||
```python
|
||||
from fastvideo.performance.hf_store import load_records_for_model
|
||||
from hf_store import load_records_for_model
|
||||
|
||||
records = load_records_for_model(
|
||||
"/tmp/perf-tracking",
|
||||
@@ -371,7 +372,7 @@ prepared records plus backup on disk.
|
||||
Use the shared storage helper so the path and repo type match CI:
|
||||
|
||||
```python
|
||||
from fastvideo.performance.hf_store import upload_record
|
||||
from hf_store import upload_record
|
||||
|
||||
upload_record("<local_record_path>", record, strict=True)
|
||||
```
|
||||
@@ -459,7 +460,7 @@ 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,
|
||||
- `fastvideo/tests/performance/hf_store.py` — HF sync, record loading,
|
||||
`sanitize()`, and `upload_record()`.
|
||||
- `fastvideo/tests/performance/test_inference_performance.py` — source result
|
||||
JSON schema.
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
{
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"description": "Wan2.1 T2V 1.3B inference performance",
|
||||
"model": {
|
||||
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
|
||||
@@ -1,11 +1,7 @@
|
||||
name: pre-commit
|
||||
|
||||
on:
|
||||
# pull_request_target instead of pull_request: the workflow definition and
|
||||
# the hook config are always taken from the BASE branch, so fork /
|
||||
# first-time-contributor PRs run immediately without a maintainer clicking
|
||||
# "Approve and run". The PR head is checked out as data only.
|
||||
pull_request_target:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
workflow_call:
|
||||
inputs:
|
||||
@@ -19,25 +15,12 @@ permissions:
|
||||
|
||||
jobs:
|
||||
pre-commit:
|
||||
if: github.event.pull_request.draft != true
|
||||
if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.ref || '' }}
|
||||
# For PR events, lint the PR head — but keep the hook definitions from
|
||||
# the base branch so an untrusted PR cannot alter what gets executed.
|
||||
- name: Save trusted hook config
|
||||
if: github.event_name == 'pull_request_target'
|
||||
run: cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
|
||||
- uses: actions/checkout@v4
|
||||
if: github.event_name == 'pull_request_target'
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
persist-credentials: false
|
||||
- name: Restore trusted hook config
|
||||
if: github.event_name == 'pull_request_target'
|
||||
run: cp "$RUNNER_TEMP/trusted-pre-commit-config.yaml" .pre-commit-config.yaml
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
+2
-2
@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
|
||||
# cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
|
||||
# flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
|
||||
# upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
|
||||
ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
|
||||
ARG FA4_CUTE_REF=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
|
||||
|
||||
# Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
|
||||
# select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
|
||||
@@ -161,7 +161,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
|
||||
fi
|
||||
|
||||
# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
|
||||
# Overlay the cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
|
||||
# wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
|
||||
# runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
|
||||
# the same deps the [dreamverse] extra already installs in CI; the installed torch
|
||||
|
||||
@@ -79,11 +79,10 @@ fastvideo/performance/
|
||||
```
|
||||
|
||||
The HF dataset (`FastVideo/performance-tracking` by default) holds one
|
||||
normalized JSON per run. For v2 records, the rolling baseline is the median of
|
||||
the last 5 successful, baseline-eligible records in the same comparison cohort:
|
||||
`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.
|
||||
normalized JSON per `(model_id, gpu_type, run)` tuple. The rolling baseline is
|
||||
the median of the last 5 successful, baseline-eligible records for that
|
||||
model+GPU. PR and local records are visible in the dashboard but are not
|
||||
baseline eligible.
|
||||
|
||||
## Planned Coverage
|
||||
|
||||
@@ -113,11 +112,10 @@ per-metric policy with direction, percent threshold, absolute threshold, and a
|
||||
while it runs so pipeline stage execution times are available in
|
||||
`generate_video(...).logging_info`. Stage logs use pipeline-unique keys such as
|
||||
`prompt_encoding_stage` so duplicate stage classes do not collide. For
|
||||
`PipelineStage` entries, shared component stage bases emit a stable
|
||||
`component_metric`: text encoding stages map to `text_encoder_time_s`,
|
||||
denoising stages and subclasses map to `dit_time_s`, and decoding stages map to
|
||||
`vae_decode_time_s`. The extractor falls back to known `stage_class` names for
|
||||
older logs that do not include `component_metric` or that used the class name as
|
||||
`PipelineStage` entries, the extractor maps the `stage_class` field:
|
||||
`TextEncodingStage` maps to `text_encoder_time_s`, `DenoisingStage` and
|
||||
`DmdDenoisingStage` map to `dit_time_s`, and `DecodingStage` maps to
|
||||
`vae_decode_time_s`, with a fallback for older logs that used the class name as
|
||||
the stage key. Generator-side timings such as `PostDecodeFrameProcessStage`,
|
||||
`VideoSaveStage`, and `AudioMuxStage` are intentionally ignored. If a pipeline
|
||||
does not report one of the mapped stages, that component metric is stored as
|
||||
@@ -159,16 +157,13 @@ unrealistic memory growth, and optionally large component-specific slowdowns
|
||||
even when the rolling baseline is empty. They are hand-set with generous
|
||||
headroom and almost never need touching.
|
||||
|
||||
### Rolling baseline (per comparison cohort)
|
||||
### Rolling baseline (per `(model_id, gpu_type)`)
|
||||
|
||||
`compare_baseline.py` loads the last 5 successful, baseline-eligible records
|
||||
for the same comparison cohort from the HF dataset, computes the median for
|
||||
each available metric, and evaluates the current run with the metric's
|
||||
rolling regression policy. For v2 records, that cohort is `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.
|
||||
for the same `(model_id, gpu_type)` from the HF dataset, computes the median
|
||||
for each available metric, and evaluates the current run with the metric's
|
||||
rolling regression policy. For latency, memory, and component times, higher
|
||||
values are regressions. For throughput, lower values are regressions.
|
||||
|
||||
A metric exceeds its rolling threshold when both of these are true:
|
||||
|
||||
@@ -205,9 +200,9 @@ configs and remain loadable. New or migrated configs should use
|
||||
{
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1
|
||||
}
|
||||
```
|
||||
|
||||
@@ -218,8 +213,8 @@ metadata that make the measured workload explicit:
|
||||
|
||||
| Field | Purpose |
|
||||
|---|---|
|
||||
| `workload_id` | Stable benchmark family, such as `wan-t2v`. |
|
||||
| `variant_id` | Intentional recipe family, including model size and parallelism config, such as `1.3b-sp2`. |
|
||||
| `workload_id` | Stable benchmark family, such as `wan-t2v-1.3b`. |
|
||||
| `variant_id` | Intentional recipe family, such as `canonical`. |
|
||||
| `benchmark_version` | Version of the measurement protocol and comparison policy. |
|
||||
|
||||
If a config declares `config_schema_version: 2`, loading fails clearly when any
|
||||
@@ -234,10 +229,8 @@ 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
|
||||
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
|
||||
gating starts once that cohort accumulates history). Legacy v1 configs keep
|
||||
the `(model_id, gpu_type)` rolling-baseline comparison. Metric-specific
|
||||
threshold policies and promoted baselines remain separate follow-ups.
|
||||
|
||||
### Raw record (`results/perf_*.json`)
|
||||
@@ -247,10 +240,10 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
```jsonc
|
||||
{
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
|
||||
"device": "NVIDIA L40S",
|
||||
"num_gpus": 2,
|
||||
@@ -276,15 +269,8 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
}
|
||||
},
|
||||
"commit": "<full sha>",
|
||||
"run_source": "pr",
|
||||
"branch": "feature/perf-change",
|
||||
"pr_number": "1234",
|
||||
"test_scope": "direct",
|
||||
"build_url": "https://buildkite.example/build",
|
||||
"build_id": "<buildkite-build-id>",
|
||||
"job_id": "<buildkite-job-id>",
|
||||
"timestamp": "2026-05-08T22:00:00+00:00",
|
||||
"quality_metadata": { "quality_status": "canonical" },
|
||||
"text_encoder_time_s": 2.141,
|
||||
"dit_time_s": 8.437,
|
||||
"vae_decode_time_s": 3.208,
|
||||
@@ -292,9 +278,9 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
"recipe_schema_version": 1,
|
||||
"benchmark": {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1
|
||||
},
|
||||
"model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
|
||||
"init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
|
||||
@@ -335,10 +321,9 @@ result, used as the rolling-baseline source of truth.
|
||||
```jsonc
|
||||
{
|
||||
"model_id": "wan-t2v-1.3b-2gpu",
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"timestamp": "2026-05-08T22:00:00+00:00",
|
||||
"commit_sha": "<full sha>",
|
||||
"gpu_type": "NVIDIA L40S",
|
||||
@@ -359,30 +344,17 @@ result, used as the rolling-baseline source of truth.
|
||||
"hardware_profile_id": "hw-<sha256-prefix>",
|
||||
"software_profile_id": "sw-<sha256-prefix>",
|
||||
"environment_fingerprint": "env-<sha256-prefix>",
|
||||
"run_source": "pr",
|
||||
"branch": "feature/perf-change",
|
||||
"pr_number": "1234",
|
||||
"test_scope": "direct",
|
||||
"build_url": "https://buildkite.example/build",
|
||||
"build_id": "<buildkite-build-id>",
|
||||
"job_id": "<buildkite-job-id>",
|
||||
"quality_metadata": { "quality_status": "canonical" },
|
||||
"success": true
|
||||
}
|
||||
```
|
||||
|
||||
### Compatibility with legacy records
|
||||
|
||||
Older records in the HF dataset may not have `result_schema_version`,
|
||||
component timing fields, or v2 identity/profile fields. Records without
|
||||
`result_schema_version` are treated as v1. The comparator ignores missing or
|
||||
`null` metrics when computing a median, and the dashboard lists skipped plots
|
||||
for metric series that have no non-null values. Records missing both
|
||||
`run_source` and `baseline_eligible` are treated as legacy successful
|
||||
main/full-suite uploads and remain eligible for rolling baselines.
|
||||
Current `perf_*.json` artifacts that lack the v2 comparison identity are
|
||||
normalized for reporting but skip rolling-baseline comparison and are not marked
|
||||
baseline eligible.
|
||||
Older records in the HF dataset may not have component timing fields. The
|
||||
comparator ignores missing or `null` metrics when computing a median, and the
|
||||
dashboard lists skipped plots for metric series that have no non-null values.
|
||||
Records missing both `run_source` and `baseline_eligible` are treated as legacy
|
||||
successful main/full-suite uploads and remain eligible for rolling baselines.
|
||||
|
||||
New records compare only against the same `model_id`, `gpu_type`,
|
||||
`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
|
||||
@@ -404,7 +376,7 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
|
||||
| `PERF_REPORTS_DIR` | `/root/data/perf_reports` | `compare_baseline.py`, `dashboard.py` | Where the Markdown summary and Plotly HTML get written for Buildkite to pick up. |
|
||||
| `HF_REPO_ID` | `FastVideo/performance-tracking` | `fastvideo/performance/hf_store.py` | HF dataset repo holding rolling-baseline records. |
|
||||
| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `fastvideo/performance/hf_store.py` | Required for upload or private dataset reads. |
|
||||
| `PERF_RUN_SOURCE` | inferred | `compare_baseline.py`, `test_inference_performance.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
|
||||
| `PERF_RUN_SOURCE` | inferred | `compare_baseline.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
|
||||
| `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
|
||||
| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
|
||||
| `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
|
||||
@@ -421,14 +393,11 @@ point is `fastvideo/tests/modal/pr_test.py:run_performance_tests` and the
|
||||
Buildkite artifact upload is in
|
||||
`.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.
|
||||
|
||||
Each performance build runs pytest first. PR and direct runs only continue to
|
||||
`compare_baseline.py` when that fixed-threshold phase passes; if pytest fails,
|
||||
Markdown summaries and normalized JSON artifacts are not emitted. Scheduled
|
||||
main runs set `PERF_UPLOAD_POLICY=always`, so they still run
|
||||
`compare_baseline.py` (with `PERF_PYTEST_RC` set) after a fixed-threshold
|
||||
failure. Those failed scheduled main runs emit summaries and normalized
|
||||
records, upload records with `success=false`, and are excluded from future
|
||||
rolling baselines. The dashboard still runs best-effort for observability.
|
||||
Each performance build runs pytest first. If that fixed-threshold phase fails,
|
||||
PR/direct runs skip `compare_baseline.py` because they only upload passing
|
||||
records. Scheduled-main runs still execute `compare_baseline.py` with
|
||||
`PERF_PYTEST_RC` set so the failed canonical attempt is visible in normalized
|
||||
JSON and dashboard history. The dashboard runs best-effort for observability.
|
||||
When the rolling-baseline phase runs, it emits:
|
||||
|
||||
* **Markdown summary** — appended to `$GITHUB_STEP_SUMMARY` when that variable
|
||||
@@ -436,7 +405,7 @@ When the rolling-baseline phase runs, it emits:
|
||||
per-benchmark row with current vs. baseline values for latency, throughput,
|
||||
memory, text encoder time, DiT time, and VAE decode time.
|
||||
* **Plotly dashboard** — `dashboard_<sha>_<ts>.html` showing time-series for
|
||||
each metric grouped by comparison cohort.
|
||||
each metric grouped by `(model_id, gpu_type)`.
|
||||
* **Normalized records** — `normalized_perf_*.json`, one per benchmark.
|
||||
Useful as input to the
|
||||
[`reseed-performance-baseline`](https://github.com/hao-ai-lab/FastVideo/blob/main/.agents/skills/reseed-performance-baseline/SKILL.md)
|
||||
@@ -453,7 +422,7 @@ When the rolling-baseline phase runs, it emits:
|
||||
"benchmark_id": "<unique-id>",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "<stable-workload-id>",
|
||||
"variant_id": "<variant, e.g. 1.3b-sp2>",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"model": { "model_path": "...", "model_short_name": "..." },
|
||||
"init_kwargs": { "num_gpus": 1, ... },
|
||||
@@ -510,8 +479,8 @@ 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.
|
||||
**"No baseline for ... Initializing"** — first run for this `(model_id,
|
||||
gpu_type)`. Run will pass and (if persisting) seed the first record.
|
||||
|
||||
**Persistent failure right after a torch / kernel / image upgrade** —
|
||||
genuine regression *or* baseline drift. Compare the failing normalized record
|
||||
@@ -524,5 +493,5 @@ pipelines that did not report a mapped component stage.
|
||||
|
||||
**Component timing is `null`** — the generated result did not include a mapped
|
||||
stage in `logging_info.stages`. Check that the pipeline emits stage logging
|
||||
and that the stage emits `component_metric` or is covered by the legacy
|
||||
`STAGE_METRIC_MAP` fallback in `test_inference_performance.py`.
|
||||
and that the stage name is listed in `STAGE_METRIC_MAP` in
|
||||
`test_inference_performance.py`.
|
||||
|
||||
@@ -191,9 +191,6 @@ surfaces:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.gen3c.Gen3CConfig
|
||||
- fastvideo.configs.pipelines.gen3c.Gen3CInferenceConfig
|
||||
color_correction_strength:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.dreamx_world.DreamXWorld5BARPipelineConfig
|
||||
default_camera_rotation:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.gen3c.Gen3CConfig
|
||||
|
||||
@@ -58,8 +58,6 @@ pipeline initialization and sampling.
|
||||
| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
|
||||
| DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
|
||||
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
|
||||
|
||||
@@ -1,91 +0,0 @@
|
||||
# Training Trackers
|
||||
|
||||
FastVideo can send training metrics and validation media to Weights & Biases
|
||||
or SwanLab. Tracking runs only on global rank 0, and local tracker files are
|
||||
stored under `<output_dir>/tracker`.
|
||||
|
||||
## Supported Trackers
|
||||
|
||||
| Value | Backend | Installation |
|
||||
|-------|---------|--------------|
|
||||
| `wandb` | Weights & Biases | Included with FastVideo |
|
||||
| `swanlab` | SwanLab | Install the optional `swanlab` dependency |
|
||||
| `none` | Disable external tracking | No additional package |
|
||||
|
||||
You can enable more than one backend, for example `trackers: [wandb, swanlab]`.
|
||||
Metrics and validation media are converted to the artifact type required by
|
||||
each backend.
|
||||
|
||||
## Install SwanLab
|
||||
|
||||
For a published FastVideo installation, install the SwanLab extra:
|
||||
|
||||
```bash
|
||||
uv pip install "fastvideo[swanlab]"
|
||||
```
|
||||
|
||||
For an editable source checkout, include the same extra during installation:
|
||||
|
||||
```bash
|
||||
uv pip install -e ".[swanlab]"
|
||||
```
|
||||
|
||||
If FastVideo is already installed, you can install the compatible SDK directly:
|
||||
|
||||
```bash
|
||||
uv pip install "swanlab>=0.6.7"
|
||||
```
|
||||
|
||||
Authenticate once before starting a training run:
|
||||
|
||||
```bash
|
||||
swanlab login
|
||||
```
|
||||
|
||||
See the [SwanLab login documentation](https://docs.swanlab.cn/en/api/cli-swanlab-login.html)
|
||||
for non-interactive and self-hosted setups.
|
||||
|
||||
## Configure Tracking
|
||||
|
||||
Select SwanLab in the YAML config used by the modular training framework:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
checkpoint:
|
||||
output_dir: outputs/my_run
|
||||
tracker:
|
||||
trackers: [swanlab]
|
||||
project_name: my_project
|
||||
run_name: my_run
|
||||
```
|
||||
|
||||
To log to both supported services:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
tracker:
|
||||
trackers: [wandb, swanlab]
|
||||
project_name: my_project
|
||||
run_name: my_run
|
||||
```
|
||||
|
||||
An empty or omitted `trackers` list selects W&B when `project_name` is set.
|
||||
Use an explicit `none` entry to disable external tracking:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
tracker:
|
||||
trackers: [none]
|
||||
```
|
||||
|
||||
## Validation Videos
|
||||
|
||||
SwanLab currently accepts GIF video artifacts. FastVideo converts validation
|
||||
MP4 files and in-memory video arrays to GIF automatically before logging them.
|
||||
For video files, FastVideo uses the sampling frame rate supplied by the caller,
|
||||
or the source file's frame rate when no value is supplied. In-memory arrays use
|
||||
the frame rate supplied by the caller. Both forms fall back to 16 FPS when no
|
||||
frame rate is available.
|
||||
|
||||
For details about configuring validation callbacks, see
|
||||
[Training Infrastructure](train_infra.md#callbacks-pluggable-hooks).
|
||||
@@ -161,9 +161,6 @@ training:
|
||||
decay_interval_steps: 0
|
||||
```
|
||||
|
||||
See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
|
||||
including SwanLab installation and authentication.
|
||||
|
||||
### `callbacks` — Pluggable hooks
|
||||
|
||||
Callbacks run at specific points in the training loop (before/after optimizer
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
import os
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
|
||||
OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world")
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
return int(os.getenv(name, str(default)))
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
return float(os.getenv(name, str(default)))
|
||||
|
||||
|
||||
def main():
|
||||
model_name = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
override_pipeline_cls_name="DreamXWorldPipeline",
|
||||
)
|
||||
|
||||
prompt = os.getenv(
|
||||
"DREAMX_WORLD_PROMPT",
|
||||
"A cinematic first-person drive through a futuristic coastal city at "
|
||||
"sunrise, reflective glass towers, clean streets, soft volumetric light.",
|
||||
)
|
||||
image_path = os.getenv(
|
||||
"DREAMX_WORLD_IMAGE_PATH",
|
||||
"https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG",
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"output_path": OUTPUT_PATH,
|
||||
"save_video": os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
|
||||
"height": _env_int("DREAMX_WORLD_HEIGHT", 480),
|
||||
"width": _env_int("DREAMX_WORLD_WIDTH", 832),
|
||||
"num_frames": _env_int("DREAMX_WORLD_NUM_FRAMES", 161),
|
||||
"num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30),
|
||||
"guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0),
|
||||
"action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
|
||||
"action_speed_list": [
|
||||
float(value)
|
||||
for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")
|
||||
],
|
||||
}
|
||||
if image_path:
|
||||
kwargs["image_path"] = image_path
|
||||
|
||||
try:
|
||||
generator.generate_video(prompt, **kwargs)
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,37 +0,0 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
OUTPUT_PATH = "video_samples_kandinsky5_i2v"
|
||||
|
||||
IMAGE_PATH = "assets/girl.png"
|
||||
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers",
|
||||
# "kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"A woman stands up and walks away"
|
||||
)
|
||||
_ = generator.generate_video(
|
||||
prompt,
|
||||
image_path=IMAGE_PATH,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
height=1024,
|
||||
width=1024,
|
||||
num_frames=121,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,37 +0,0 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
OUTPUT_PATH = "video_samples_kandinsky5_t2v"
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers"
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
_ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True,height=512, width=768, num_frames=121)
|
||||
|
||||
prompt2 = (
|
||||
"A majestic lion strides across the golden savanna, its powerful frame "
|
||||
"glistening under the warm afternoon sun. The tall grass ripples gently in "
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
_ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=512, width=768, num_frames=121)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,69 +0,0 @@
|
||||
# LTX-2.3 distilled inference configs
|
||||
|
||||
Ready-to-run `fastvideo generate` run configs for the LTX-2.3
|
||||
distilled model (`FastVideo/LTX-2.3-Distilled-Diffusers`), covering both
|
||||
workloads (t2v / i2v), both two-stage step schedules (`5+2`, `8+3` = denoise
|
||||
+ refine), and four resolutions.
|
||||
|
||||
```bash
|
||||
fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml
|
||||
```
|
||||
|
||||
Each config is self-contained (no preset registry needed): the two-stage
|
||||
refine is wired via `generator.pipeline.preset_overrides.refine`, and the
|
||||
base sampling knobs live under `request.sampling`. The refine upsampler
|
||||
auto-resolves from the model's `spatial_upscaler`.
|
||||
|
||||
## Configs
|
||||
|
||||
| workload | schedule | resolution (HxW) | file |
|
||||
|---|---|---|---|
|
||||
| t2v | 5+2 | 1280x832 | `t2v_5s2_1280x832.yaml` |
|
||||
| t2v | 5+2 | 1024x1536 | `t2v_5s2_1024x1536.yaml` |
|
||||
| t2v | 5+2 | 768x1280 | `t2v_5s2_768x1280.yaml` |
|
||||
| t2v | 5+2 | 512x768 | `t2v_5s2_512x768.yaml` |
|
||||
| t2v | 8+3 | 1280x832 | `t2v_8s3_1280x832.yaml` |
|
||||
| t2v | 8+3 | 1024x1536 | `t2v_8s3_1024x1536.yaml` |
|
||||
| t2v | 8+3 | 768x1280 | `t2v_8s3_768x1280.yaml` |
|
||||
| t2v | 8+3 | 512x768 | `t2v_8s3_512x768.yaml` |
|
||||
| i2v | 5+2 | 1280x832 | `i2v_5s2_1280x832.yaml` |
|
||||
| i2v | 5+2 | 1024x1536 | `i2v_5s2_1024x1536.yaml` |
|
||||
| i2v | 5+2 | 768x1280 | `i2v_5s2_768x1280.yaml` |
|
||||
| i2v | 5+2 | 512x768 | `i2v_5s2_512x768.yaml` |
|
||||
| i2v | 8+3 | 1280x832 | `i2v_8s3_1280x832.yaml` |
|
||||
| i2v | 8+3 | 1024x1536 | `i2v_8s3_1024x1536.yaml` |
|
||||
| i2v | 8+3 | 768x1280 | `i2v_8s3_768x1280.yaml` |
|
||||
| i2v | 8+3 | 512x768 | `i2v_8s3_512x768.yaml` |
|
||||
|
||||
## Overriding without editing a file
|
||||
|
||||
Dotted overrides (prefixes `generator.` / `request.`) let you tweak any field:
|
||||
|
||||
```bash
|
||||
# swap prompt
|
||||
fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml \
|
||||
--request.prompt "a red fox running through fresh snow"
|
||||
|
||||
# change output path / gpu count
|
||||
fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_512x768.yaml \
|
||||
--request.output.output_path outputs/preview.mp4 \
|
||||
--generator.engine.num_gpus 4
|
||||
```
|
||||
|
||||
## i2v
|
||||
|
||||
The `i2v_*` configs take a first-frame image via
|
||||
`request.extensions.ltx2_images` (`[[path, frame_offset, weight]]`). Edit the
|
||||
path in the file, or override it:
|
||||
|
||||
```bash
|
||||
fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1280x832.yaml \
|
||||
--request.extensions.ltx2_images '[["/data/portrait.jpg", 0, 1.0]]'
|
||||
```
|
||||
|
||||
## Schedules
|
||||
|
||||
`5+2` is the fast preview schedule; `8+3` is the higher-quality distilled
|
||||
recipe. Refine (`preset_overrides.refine.num_inference_steps`) only accepts 2
|
||||
or 3 steps.
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 5+2 two-stage at 1024x1536.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_1024x1536.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 1024
|
||||
width: 1536
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_5s2_1024x1536.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 5+2 two-stage at 1280x832.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_1280x832.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 1280
|
||||
width: 832
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_5s2_1280x832.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 5+2 two-stage at 512x768.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_512x768.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 512
|
||||
width: 768
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_5s2_512x768.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 5+2 two-stage at 768x1280.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_768x1280.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 768
|
||||
width: 1280
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_5s2_768x1280.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 8+3 two-stage at 1024x1536.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1024x1536.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 1024
|
||||
width: 1536
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_8s3_1024x1536.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 8+3 two-stage at 1280x832.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1280x832.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 1280
|
||||
width: 832
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_8s3_1280x832.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 8+3 two-stage at 512x768.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_512x768.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 512
|
||||
width: 768
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_8s3_512x768.mp4
|
||||
save_video: true
|
||||
@@ -1,38 +0,0 @@
|
||||
# LTX-2.3 distilled i2v — 8+3 two-stage at 768x1280.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_768x1280.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: i2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
|
||||
sampling:
|
||||
height: 768
|
||||
width: 1280
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
# i2v conditioning — replace the path with your own first-frame image.
|
||||
extensions:
|
||||
ltx2_images:
|
||||
- ["/path/to/your/first_frame.jpg", 0, 1.0]
|
||||
ltx2_image_crf: 0.0
|
||||
output:
|
||||
output_path: outputs/ltx2_3_i2v_8s3_768x1280.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 5+2 two-stage at 1024x1536.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_1024x1536.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 1024
|
||||
width: 1536
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_5s2_1024x1536.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 5+2 two-stage at 1280x832.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_1280x832.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 1280
|
||||
width: 832
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_5s2_1280x832.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 5+2 two-stage at 512x768.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_512x768.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 512
|
||||
width: 768
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_5s2_512x768.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 5+2 two-stage at 768x1280.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_768x1280.yaml
|
||||
#
|
||||
# Stage 1 denoises for 5 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 2
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 768
|
||||
width: 1280
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 5
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_5s2_768x1280.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 8+3 two-stage at 1024x1536.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1024x1536.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 1024
|
||||
width: 1536
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_8s3_1024x1536.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 8+3 two-stage at 1280x832.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 1280
|
||||
width: 832
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_8s3_1280x832.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 8+3 two-stage at 512x768.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_512x768.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 512
|
||||
width: 768
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_8s3_512x768.mp4
|
||||
save_video: true
|
||||
@@ -1,33 +0,0 @@
|
||||
# LTX-2.3 distilled t2v — 8+3 two-stage at 768x1280.
|
||||
#
|
||||
# Run:
|
||||
# fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_768x1280.yaml
|
||||
#
|
||||
# Stage 1 denoises for 8 steps at half resolution; the latents are then
|
||||
# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
|
||||
# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
|
||||
generator:
|
||||
model_path: FastVideo/LTX-2.3-Distilled-Diffusers
|
||||
engine:
|
||||
num_gpus: 1
|
||||
pipeline:
|
||||
workload_type: t2v
|
||||
preset_overrides:
|
||||
refine:
|
||||
enabled: true
|
||||
num_inference_steps: 3
|
||||
guidance_scale: 1.0
|
||||
add_noise: true
|
||||
request:
|
||||
prompt: >-
|
||||
A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
|
||||
sampling:
|
||||
height: 768
|
||||
width: 1280
|
||||
num_frames: 121
|
||||
fps: 24
|
||||
guidance_scale: 1.0
|
||||
num_inference_steps: 8
|
||||
output:
|
||||
output_path: outputs/ltx2_3_t2v_8s3_768x1280.mp4
|
||||
save_video: true
|
||||
@@ -1,78 +0,0 @@
|
||||
# Causal Consistency Distillation: Wan 2.1 T2V 1.3B Causal
|
||||
#
|
||||
# ODE-data-free distillation. A frozen teacher takes a single CFG Euler step;
|
||||
# the student matches an EMA copy of itself at the next timestep, all under
|
||||
# clean-history teacher forcing.
|
||||
#
|
||||
# All three roles initialize from the SAME checkpoint (the teacher-forcing
|
||||
# AR-diffusion model). Point init_from at that checkpoint for a real run.
|
||||
|
||||
models:
|
||||
student:
|
||||
_target_: fastvideo.train.models.wan.WanCausalModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: true
|
||||
teacher:
|
||||
_target_: fastvideo.train.models.wan.WanCausalModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: false
|
||||
ema:
|
||||
_target_: fastvideo.train.models.wan.WanCausalModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: false
|
||||
|
||||
method:
|
||||
_target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
|
||||
discrete_cd_N: 48
|
||||
guidance_scale: 3.0
|
||||
ema_decay: 0.99
|
||||
ema_start_step: 200
|
||||
|
||||
training:
|
||||
distributed:
|
||||
num_gpus: 8
|
||||
sp_size: 1
|
||||
tp_size: 1
|
||||
hsdp_replicate_dim: 1
|
||||
hsdp_shard_dim: 8
|
||||
|
||||
data:
|
||||
data_path: data/Wan-Syn_77x448x832_600k
|
||||
dataloader_num_workers: 4
|
||||
train_batch_size: 1
|
||||
training_cfg_rate: 0.0
|
||||
seed: 1000
|
||||
num_latent_t: 18
|
||||
num_height: 448
|
||||
num_width: 832
|
||||
num_frames: 69
|
||||
|
||||
optimizer:
|
||||
learning_rate: 2.0e-6
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.01
|
||||
lr_scheduler: constant
|
||||
lr_warmup_steps: 0
|
||||
|
||||
loop:
|
||||
max_train_steps: 3000
|
||||
gradient_accumulation_steps: 1
|
||||
|
||||
checkpoint:
|
||||
output_dir: outputs/wan2.1_causal_cd
|
||||
training_state_checkpointing_steps: 1000
|
||||
checkpoints_total_limit: 3
|
||||
|
||||
tracker:
|
||||
project_name: distillation_wan_r
|
||||
run_name: wan2.1_causal_cd_shift5
|
||||
|
||||
model:
|
||||
enable_gradient_checkpointing_type: full
|
||||
|
||||
callbacks:
|
||||
grad_clip:
|
||||
max_grad_norm: 1.0
|
||||
|
||||
pipeline:
|
||||
flow_shift: 5
|
||||
@@ -1,73 +0,0 @@
|
||||
# DFSFT (Diffusion-Forcing SFT), frame-wise: Wan 2.1 T2V 1.3B Causal
|
||||
#
|
||||
# - Student: trainable causal Wan model with a block size of 1 frame
|
||||
# - Training: each frame gets its own independent noise level (frame-wise
|
||||
# diffusion forcing), versus the chunk-wise variant that shares one noise
|
||||
# level across num_frames_per_block frames.
|
||||
|
||||
models:
|
||||
student:
|
||||
_target_: fastvideo.train.models.wan.WanCausalModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: true
|
||||
num_frames_per_block: 1
|
||||
|
||||
method:
|
||||
_target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
|
||||
chunk_size: 1
|
||||
|
||||
training:
|
||||
distributed:
|
||||
num_gpus: 8
|
||||
sp_size: 1
|
||||
tp_size: 1
|
||||
hsdp_replicate_dim: 1
|
||||
hsdp_shard_dim: 8
|
||||
|
||||
data:
|
||||
data_path: data/Wan-Syn_77x448x832_600k
|
||||
dataloader_num_workers: 4
|
||||
train_batch_size: 1
|
||||
training_cfg_rate: 0.0
|
||||
seed: 1000
|
||||
num_latent_t: 18
|
||||
num_height: 448
|
||||
num_width: 832
|
||||
num_frames: 69
|
||||
|
||||
optimizer:
|
||||
learning_rate: 2.0e-6
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.01
|
||||
lr_scheduler: constant
|
||||
lr_warmup_steps: 0
|
||||
|
||||
loop:
|
||||
max_train_steps: 4000
|
||||
gradient_accumulation_steps: 1
|
||||
|
||||
checkpoint:
|
||||
output_dir: outputs/wan2.1_causal_dfsft_framewise
|
||||
training_state_checkpointing_steps: 1000
|
||||
checkpoints_total_limit: 3
|
||||
|
||||
tracker:
|
||||
project_name: distillation_wan_r
|
||||
run_name: wan2.1_causal_dfsft_framewise_shift5_gauss_weight
|
||||
|
||||
model:
|
||||
enable_gradient_checkpointing_type: full
|
||||
|
||||
callbacks:
|
||||
grad_clip:
|
||||
max_grad_norm: 1.0
|
||||
validation:
|
||||
pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
|
||||
dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
|
||||
every_steps: 50
|
||||
sampling_steps: [40]
|
||||
guidance_scale: 6.0
|
||||
num_frames: 69
|
||||
|
||||
pipeline:
|
||||
flow_shift: 5
|
||||
@@ -1,72 +0,0 @@
|
||||
# TFSFT (Teacher-Forcing SFT): Wan 2.1 T2V 1.3B Causal
|
||||
#
|
||||
# - Student: trainable causal Wan model
|
||||
# - Training: inhomogeneous timesteps per chunk, but the causal transformer
|
||||
# denoises the current block while attending to *clean* history (clean_x),
|
||||
# not its own noisy rollout.
|
||||
|
||||
models:
|
||||
student:
|
||||
_target_: fastvideo.train.models.wan.WanCausalModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: true
|
||||
|
||||
method:
|
||||
_target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
|
||||
chunk_size: 3
|
||||
|
||||
training:
|
||||
distributed:
|
||||
num_gpus: 8
|
||||
sp_size: 1
|
||||
tp_size: 1
|
||||
hsdp_replicate_dim: 1
|
||||
hsdp_shard_dim: 8
|
||||
|
||||
data:
|
||||
data_path: data/Wan-Syn_77x448x832_600k
|
||||
dataloader_num_workers: 4
|
||||
train_batch_size: 1
|
||||
training_cfg_rate: 0.0
|
||||
seed: 1000
|
||||
num_latent_t: 18
|
||||
num_height: 448
|
||||
num_width: 832
|
||||
num_frames: 69
|
||||
|
||||
optimizer:
|
||||
learning_rate: 2.0e-6
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.01
|
||||
lr_scheduler: constant
|
||||
lr_warmup_steps: 0
|
||||
|
||||
loop:
|
||||
max_train_steps: 4000
|
||||
gradient_accumulation_steps: 1
|
||||
|
||||
checkpoint:
|
||||
output_dir: outputs/wan2.1_causal_tfsft
|
||||
training_state_checkpointing_steps: 1000
|
||||
checkpoints_total_limit: 3
|
||||
|
||||
tracker:
|
||||
project_name: distillation_wan_r
|
||||
run_name: wan2.1_causal_tfsft_shift5_gauss_weight
|
||||
|
||||
model:
|
||||
enable_gradient_checkpointing_type: full
|
||||
|
||||
callbacks:
|
||||
grad_clip:
|
||||
max_grad_norm: 1.0
|
||||
validation:
|
||||
pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
|
||||
dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
|
||||
every_steps: 50
|
||||
sampling_steps: [40]
|
||||
guidance_scale: 6.0
|
||||
num_frames: 69
|
||||
|
||||
pipeline:
|
||||
flow_shift: 5
|
||||
@@ -338,11 +338,7 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
)
|
||||
|
||||
qkv = torch.stack([query, key, value], dim=2)
|
||||
key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, attn_mask.shape[-1]).to(device=query.device)
|
||||
if key_padding_mask.shape[-1] > qkv.shape[1]:
|
||||
raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
|
||||
f"expected at most {qkv.shape[1]}, got {key_padding_mask.shape[-1]}")
|
||||
attn_mask_padded = F.pad(key_padding_mask, (qkv.shape[1] - key_padding_mask.shape[-1], 0), value=True)
|
||||
attn_mask_padded = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True)
|
||||
output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
|
||||
elif self.nvfp4_fa4:
|
||||
output = self._forward_nvfp4(query, key, value)
|
||||
|
||||
@@ -1,147 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""NABLA block-sparse flex-attention backend (Kandinsky5 "nabla" checkpoints).
|
||||
|
||||
The block mask is data-dependent: nablaT_v2 mean-pools 64-token blocks of the
|
||||
fractal-ordered sequence, thresholds the softmaxed block map, and ORs it with a
|
||||
precomputed spatio-temporal-window (STA) mask carried on the attention
|
||||
metadata. The mask spans the full sequence, so this backend does not support
|
||||
sequence parallelism — use it via LocalAttention only.
|
||||
"""
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from torch.nn.attention.flex_attention import BlockMask, flex_attention
|
||||
flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
|
||||
CAN_USE_FLEX_ATTN = True
|
||||
except ImportError:
|
||||
CAN_USE_FLEX_ATTN = False
|
||||
|
||||
from fastvideo.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
|
||||
|
||||
def nablaT_v2(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
sta: torch.Tensor,
|
||||
thr: float = 0.9,
|
||||
) -> "BlockMask":
|
||||
q = q.transpose(1, 2).contiguous()
|
||||
k = k.transpose(1, 2).contiguous()
|
||||
|
||||
# Map estimation
|
||||
B, h, S, D = q.shape
|
||||
s1 = S // 64
|
||||
qa = q.reshape(B, h, s1, 64, D).mean(-2)
|
||||
ka = k.reshape(B, h, s1, 64, D).mean(-2).transpose(-2, -1)
|
||||
map = qa @ ka
|
||||
|
||||
map = torch.softmax(map / math.sqrt(D), dim=-1)
|
||||
# Map binarization
|
||||
vals, inds = map.sort(-1)
|
||||
cvals = vals.cumsum_(-1)
|
||||
mask = (cvals >= 1 - thr).int()
|
||||
mask = mask.gather(-1, inds.argsort(-1))
|
||||
|
||||
mask = torch.logical_or(mask, sta)
|
||||
|
||||
# BlockMask creation
|
||||
kv_nb = mask.sum(-1).to(torch.int32)
|
||||
kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32)
|
||||
return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None)
|
||||
|
||||
|
||||
class NablaAttentionBackend(AttentionBackend):
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "NABLA_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["NablaAttentionImpl"]:
|
||||
return NablaAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["NablaAttentionMetadata"]:
|
||||
return NablaAttentionMetadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["NablaAttentionMetadataBuilder"]:
|
||||
return NablaAttentionMetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class NablaAttentionMetadata(AttentionMetadata):
|
||||
# Block-level STA window mask [1, 1, S/64, S/64], precomputed once per run.
|
||||
sta_mask: torch.Tensor = None # type: ignore[assignment]
|
||||
# Cumulative-probability threshold for block-map binarization.
|
||||
P: float = 0.9
|
||||
visual_shape: tuple[int, int, int] = (0, 0, 0)
|
||||
|
||||
|
||||
class NablaAttentionMetadataBuilder(AttentionMetadataBuilder):
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def build(
|
||||
self,
|
||||
current_timestep: int,
|
||||
sta_mask: torch.Tensor,
|
||||
P: float,
|
||||
visual_shape: tuple[int, int, int],
|
||||
**kwargs: Any,
|
||||
) -> NablaAttentionMetadata:
|
||||
return NablaAttentionMetadata(
|
||||
current_timestep=current_timestep,
|
||||
sta_mask=sta_mask,
|
||||
P=P,
|
||||
visual_shape=visual_shape,
|
||||
)
|
||||
|
||||
|
||||
class NablaAttentionImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
softmax_scale: float,
|
||||
causal: bool = False,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
if not CAN_USE_FLEX_ATTN:
|
||||
raise RuntimeError("NABLA attention requires torch.nn.attention.flex_attention, "
|
||||
"which is unavailable in this PyTorch build.")
|
||||
if causal:
|
||||
raise ValueError("NABLA attention does not support causal masking.")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: NablaAttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
# q/k/v: [B, S, heads, head_dim], fractal-ordered by the model; S % 64 == 0.
|
||||
block_mask = nablaT_v2(query, key, attn_metadata.sta_mask, thr=attn_metadata.P)
|
||||
return flex_attention(
|
||||
query=query.transpose(1, 2),
|
||||
key=key.transpose(1, 2),
|
||||
value=value.transpose(1, 2),
|
||||
block_mask=block_mask,
|
||||
).transpose(1, 2)
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from torch.nn import functional as F
|
||||
from fastvideo.attention.backends.abstract import ( # FlashAttentionMetadata,
|
||||
AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder)
|
||||
from fastvideo.logger import init_logger
|
||||
@@ -20,7 +19,7 @@ class SDPABackend(AttentionBackend):
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "TORCH_SDPA"
|
||||
return "SDPA"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["SDPAImpl"]:
|
||||
@@ -50,51 +49,9 @@ class SDPAMetadataBuilder(AttentionMetadataBuilder):
|
||||
current_timestep: int,
|
||||
attn_mask: torch.Tensor,
|
||||
) -> SDPAMetadata:
|
||||
# Store the mask exactly as passed. The metadata is cross-backend:
|
||||
# call sites (HYWorld, HunyuanVideo15) build SDPAMetadata while the
|
||||
# layer's selector may pick FLASH_ATTN, and the shared convention for
|
||||
# padding masks is the tokenizer-style 2D [batch, key_len]. Any
|
||||
# reshaping for torch.sdpa happens inside the SDPA impl
|
||||
# (_normalize_attn_mask_for_sdpa).
|
||||
return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask)
|
||||
|
||||
|
||||
def _normalize_attn_mask_for_sdpa(
|
||||
attn_mask: torch.Tensor | None,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
) -> torch.Tensor | None:
|
||||
if attn_mask is None:
|
||||
return None
|
||||
|
||||
attn_mask = attn_mask.to(device=query.device)
|
||||
# F.scaled_dot_product_attention only accepts bool or float masks;
|
||||
# tokenizers commonly produce int64 0/1 padding masks.
|
||||
if attn_mask.dtype != torch.bool and not attn_mask.dtype.is_floating_point:
|
||||
attn_mask = attn_mask != 0
|
||||
|
||||
key_len = key.shape[-2]
|
||||
if attn_mask.shape[-1] > key_len:
|
||||
raise ValueError("Invalid attention mask length for SDPA: "
|
||||
f"expected at most {key_len}, got {attn_mask.shape[-1]}")
|
||||
if attn_mask.shape[-1] < key_len:
|
||||
# Front-pad as "attend": double-stream layouts (HYWorld) prepend
|
||||
# non-text tokens the tokenizer mask does not cover.
|
||||
valid_value = True if attn_mask.dtype == torch.bool else 0.0
|
||||
attn_mask = F.pad(attn_mask, (key_len - attn_mask.shape[-1], 0), value=valid_value)
|
||||
|
||||
if attn_mask.dim() == 2:
|
||||
# In-tree producers pass 2D [batch, key_len] padding masks; lift to a
|
||||
# broadcastable [batch, 1, 1, key_len] here so torch.sdpa does not
|
||||
# reinterpret 2D as its documented [query_len, key_len] broadcast.
|
||||
return attn_mask[:, None, None, :]
|
||||
if attn_mask.dim() == 3:
|
||||
return attn_mask[:, None, :, :]
|
||||
if attn_mask.dim() == 4:
|
||||
return attn_mask
|
||||
raise ValueError(f"Unsupported attention mask shape for SDPA: {attn_mask.shape}")
|
||||
|
||||
|
||||
class SDPAImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
@@ -125,7 +82,6 @@ class SDPAImpl(AttentionImpl):
|
||||
|
||||
attn_mask = attn_metadata.attn_mask if (attn_metadata is not None
|
||||
and hasattr(attn_metadata, "attn_mask")) else None
|
||||
attn_mask = _normalize_attn_mask_for_sdpa(attn_mask, query, key)
|
||||
attn_kwargs = {
|
||||
"attn_mask": attn_mask,
|
||||
"dropout_p": self.dropout,
|
||||
|
||||
@@ -252,7 +252,6 @@ class LocalAttention(nn.Module):
|
||||
causal: bool = False,
|
||||
supported_attention_backends: tuple[AttentionBackendEnum, ...]
|
||||
| None = None,
|
||||
default_backend: AttentionBackendEnum | None = None,
|
||||
**extra_impl_args) -> None:
|
||||
super().__init__()
|
||||
if softmax_scale is None:
|
||||
@@ -263,10 +262,7 @@ class LocalAttention(nn.Module):
|
||||
num_kv_heads = num_heads
|
||||
|
||||
dtype = get_compute_dtype()
|
||||
attn_backend = get_attn_backend(head_size,
|
||||
dtype,
|
||||
supported_attention_backends=supported_attention_backends,
|
||||
default_backend=default_backend)
|
||||
attn_backend = get_attn_backend(head_size, dtype, supported_attention_backends=supported_attention_backends)
|
||||
impl_cls = attn_backend.get_impl_cls()
|
||||
self.attn_impl = impl_cls(num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
|
||||
@@ -84,9 +84,8 @@ def get_attn_backend(
|
||||
dtype: torch.dtype,
|
||||
supported_attention_backends: tuple[AttentionBackendEnum, ...]
|
||||
| None = None,
|
||||
default_backend: AttentionBackendEnum | None = None,
|
||||
) -> type[AttentionBackend]:
|
||||
return _cached_get_attn_backend(head_size, dtype, supported_attention_backends, default_backend)
|
||||
return _cached_get_attn_backend(head_size, dtype, supported_attention_backends)
|
||||
|
||||
|
||||
@cache
|
||||
@@ -95,7 +94,6 @@ def _cached_get_attn_backend(
|
||||
dtype: torch.dtype,
|
||||
supported_attention_backends: tuple[AttentionBackendEnum, ...]
|
||||
| None = None,
|
||||
default_backend: AttentionBackendEnum | None = None,
|
||||
) -> type[AttentionBackend]:
|
||||
# Check whether a particular choice of backend was
|
||||
# previously forced.
|
||||
@@ -114,12 +112,6 @@ def _cached_get_attn_backend(
|
||||
if backend_by_env_var is not None:
|
||||
selected_backend = backend_name_to_enum(backend_by_env_var)
|
||||
|
||||
# Layer-level default (e.g. a checkpoint that requires a specific sparse
|
||||
# backend). Lower precedence than the global force and the env var, so
|
||||
# users can still override it.
|
||||
if selected_backend is None and default_backend is not None:
|
||||
selected_backend = default_backend
|
||||
|
||||
# get device-specific attn_backend
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
|
||||
from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
|
||||
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
|
||||
from fastvideo.configs.models.dits.flux_2 import Flux2Config
|
||||
from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
|
||||
from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
|
||||
@@ -14,7 +13,7 @@ from fastvideo.configs.models.dits.hyworld import HYWorldConfig
|
||||
from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
|
||||
|
||||
__all__ = [
|
||||
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
|
||||
"DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
|
||||
"HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
|
||||
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
|
||||
"Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
|
||||
"MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
|
||||
]
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXWorldArchConfig(WanVideoArchConfig):
|
||||
"""DreamX-World DiT config with camera PRoPE control fields."""
|
||||
|
||||
add_control_adapter: bool = True
|
||||
cam_method: str | None = "prope"
|
||||
attn_compress: int = 1
|
||||
cam_self_attn_layers: tuple[int, ...] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXWorldConfig(DiTConfig):
|
||||
arch_config: DiTArchConfig = field(default_factory=DreamXWorldArchConfig)
|
||||
|
||||
prefix: str = "Wan"
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXWorldARArchConfig(DreamXWorldArchConfig):
|
||||
"""DreamX-World-5B autoregressive causal DiT config."""
|
||||
|
||||
model_type: str = "ti2v"
|
||||
patch_size: tuple[int, int, int] = (1, 2, 2)
|
||||
text_len: int = 512
|
||||
text_dim: int = 4096
|
||||
freq_dim: int = 256
|
||||
attn_compress: int = 4
|
||||
cam_self_attn_layers: tuple[int, ...] | None = tuple(range(30))
|
||||
local_attn_size: int = 12
|
||||
sink_size: int = 3
|
||||
num_frames_per_block: int = 3
|
||||
rope_cache_policy: str = "block_relativistic"
|
||||
# The official AR checkpoint (AMAP-ML/DreamX-World ``model.safetensors``)
|
||||
# already uses FastVideo's native key names and the converter copies the
|
||||
# tensors verbatim, so every rule is an identity. The rules enumerate the
|
||||
# full state-dict surface of ``DreamXWorldARTransformer3DModel`` (norm1 /
|
||||
# norm2 / head.norm are affine-free and have no parameters).
|
||||
param_names_mapping: dict = field(
|
||||
default_factory=lambda: {
|
||||
r"^patch_embedding\.(.*)$": r"patch_embedding.\1",
|
||||
r"^text_embedding\.([02])\.(.*)$": r"text_embedding.\1.\2",
|
||||
r"^time_embedding\.([02])\.(.*)$": r"time_embedding.\1.\2",
|
||||
r"^time_projection\.1\.(.*)$": r"time_projection.1.\1",
|
||||
r"^blocks\.(\d+)\.self_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.self_attn.\2.\3",
|
||||
r"^blocks\.(\d+)\.self_attn\.norm_(q|k)\.weight$": r"blocks.\1.self_attn.norm_\2.weight",
|
||||
r"^blocks\.(\d+)\.cross_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.cross_attn.\2.\3",
|
||||
r"^blocks\.(\d+)\.cross_attn\.norm_(q|k)\.weight$": r"blocks.\1.cross_attn.norm_\2.weight",
|
||||
r"^blocks\.(\d+)\.cam_self_attn\.(q_proj|k_proj|v_proj|out_proj)\.(.*)$": r"blocks.\1.cam_self_attn.\2.\3",
|
||||
r"^blocks\.(\d+)\.cam_self_attn\.norm_(q|k)\.weight$": r"blocks.\1.cam_self_attn.norm_\2.weight",
|
||||
r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm3.\2",
|
||||
r"^blocks\.(\d+)\.ffn\.([02])\.(.*)$": r"blocks.\1.ffn.\2.\3",
|
||||
r"^blocks\.(\d+)\.modulation$": r"blocks.\1.modulation",
|
||||
r"^head\.head\.(.*)$": r"head.head.\1",
|
||||
r"^head\.modulation$": r"head.modulation",
|
||||
})
|
||||
reverse_param_names_mapping: dict = field(default_factory=dict)
|
||||
lora_param_names_mapping: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXWorldARConfig(DreamXWorldConfig):
|
||||
arch_config: DiTArchConfig = field(default_factory=DreamXWorldARArchConfig)
|
||||
|
||||
prefix: str = "Wan"
|
||||
@@ -2,24 +2,14 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
from fastvideo.platforms import AttentionBackendEnum
|
||||
|
||||
|
||||
def _is_kandinsky5_transformer_block(n: str, m) -> bool:
|
||||
return ("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
|
||||
|
||||
|
||||
@dataclass
|
||||
class Kandinsky5ArchConfig(DiTArchConfig):
|
||||
_fsdp_shard_conditions: list = field(default_factory=lambda: [_is_kandinsky5_transformer_block])
|
||||
|
||||
# NABLA block-sparse attention for attention_type="nabla" checkpoints, plus
|
||||
# the dense backends every DiT supports.
|
||||
_supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
|
||||
AttentionBackendEnum.NABLA_ATTN,
|
||||
AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.TORCH_SDPA,
|
||||
)
|
||||
_fsdp_shard_conditions: list = field(default_factory=lambda: [
|
||||
lambda n, m:
|
||||
("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
|
||||
])
|
||||
|
||||
# Native FastVideo implementation uses the same parameter names as diffusers
|
||||
# except FFN internals: Diffusers FFN uses `in_layer/out_layer`, while
|
||||
|
||||
@@ -43,16 +43,11 @@ class TextEncoderArchConfig(EncoderArchConfig):
|
||||
require_processor: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# update_model_arch re-runs __post_init__ after pipeline configs may
|
||||
# have customized tokenizer_kwargs (e.g. kandinsky5/gen3c/longcat set
|
||||
# "padding"); rebuilding the dict here would silently wipe those
|
||||
# customizations, so only fill in defaults for keys not already set.
|
||||
defaults = {
|
||||
self.tokenizer_kwargs = {
|
||||
"truncation": True,
|
||||
"max_length": self.text_len,
|
||||
"return_tensors": "pt",
|
||||
}
|
||||
self.tokenizer_kwargs = defaults | self.tokenizer_kwargs
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.configs.pipelines.cosmos import CosmosConfig
|
||||
from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
|
||||
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
|
||||
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
|
||||
from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
|
||||
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
|
||||
from fastvideo.configs.pipelines.hyworld import HYWorldConfig
|
||||
from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
|
||||
from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
|
||||
from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
|
||||
from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
|
||||
@@ -18,6 +16,5 @@ __all__ = [
|
||||
"HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
|
||||
"Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
|
||||
"SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
|
||||
"DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
|
||||
"Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
|
||||
"HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
|
||||
]
|
||||
|
||||
@@ -1,127 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World-5B-Cam FastVideo model configuration helpers."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
|
||||
from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldARArchConfig, DreamXWorldARConfig,
|
||||
DreamXWorldArchConfig, DreamXWorldConfig)
|
||||
from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
|
||||
from fastvideo.configs.models.encoders.t5 import T5ArchConfig
|
||||
from fastvideo.configs.models.vaes import WanVAEConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.configs.pipelines.wan import LucyEditDevConfig, t5_postprocess_text
|
||||
|
||||
|
||||
def make_dreamx_world_5b_cam_dit_config() -> DreamXWorldConfig:
|
||||
"""Return the DreamX-World DiT config matching DreamX-World-5B-Cam."""
|
||||
return DreamXWorldConfig(arch_config=DreamXWorldArchConfig(
|
||||
num_attention_heads=24,
|
||||
attention_head_dim=128,
|
||||
in_channels=48,
|
||||
out_channels=48,
|
||||
ffn_dim=14336,
|
||||
num_layers=30,
|
||||
cross_attn_norm=True,
|
||||
qk_norm="rms_norm_across_heads",
|
||||
add_control_adapter=True,
|
||||
cam_method="prope",
|
||||
attn_compress=1,
|
||||
cam_self_attn_layers=None,
|
||||
))
|
||||
|
||||
|
||||
def make_dreamx_world_5b_ar_dit_config() -> DreamXWorldARConfig:
|
||||
"""Return the DreamX-World-5B autoregressive causal DiT config."""
|
||||
return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig(
|
||||
model_type="ti2v",
|
||||
num_attention_heads=24,
|
||||
attention_head_dim=128,
|
||||
in_channels=48,
|
||||
out_channels=48,
|
||||
ffn_dim=14336,
|
||||
num_layers=30,
|
||||
cross_attn_norm=True,
|
||||
qk_norm=True,
|
||||
add_control_adapter=True,
|
||||
cam_method="prope",
|
||||
attn_compress=4,
|
||||
cam_self_attn_layers=tuple(range(30)),
|
||||
local_attn_size=12,
|
||||
sink_size=3,
|
||||
num_frames_per_block=3,
|
||||
))
|
||||
|
||||
|
||||
def make_dreamx_world_5b_cam_vae_config() -> WanVAEConfig:
|
||||
"""Return the Wan2.2 48-channel VAE config used by DreamX-World-5B-Cam."""
|
||||
return LucyEditDevConfig().vae_config
|
||||
|
||||
|
||||
def make_dreamx_world_5b_cam_text_encoder_config() -> T5Config:
|
||||
"""Return the UMT5-XXL text encoder config used by DreamX-World-5B-Cam."""
|
||||
return T5Config(
|
||||
arch_config=T5ArchConfig(
|
||||
vocab_size=256384,
|
||||
d_model=4096,
|
||||
d_kv=64,
|
||||
d_ff=10240,
|
||||
num_layers=24,
|
||||
num_decoder_layers=None,
|
||||
num_heads=64,
|
||||
relative_attention_num_buckets=32,
|
||||
dropout_rate=0.0,
|
||||
text_len=512,
|
||||
feed_forward_proj="gelu",
|
||||
is_encoder_decoder=False,
|
||||
),
|
||||
prefix="umt5",
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXWorld5BCamPipelineConfig(PipelineConfig):
|
||||
"""Pipeline config for the first-scope DreamX-World-5B-Cam mode."""
|
||||
|
||||
dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_cam_dit_config)
|
||||
vae_config: VAEConfig = field(default_factory=make_dreamx_world_5b_cam_vae_config)
|
||||
text_encoder_configs: tuple[EncoderConfig,
|
||||
...] = field(default_factory=lambda: (make_dreamx_world_5b_cam_text_encoder_config(), ))
|
||||
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
|
||||
...] = field(default_factory=lambda: (t5_postprocess_text, ))
|
||||
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
|
||||
flow_shift: float | None = 3.0
|
||||
ti2v_task: bool = True
|
||||
expand_timesteps: bool = True
|
||||
vae_tiling: bool = False
|
||||
vae_sp: bool = False
|
||||
vae_precision: str = "fp32"
|
||||
vae_decode_precision: str | None = "bf16"
|
||||
dit_precision: str = "bf16"
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.vae_config.load_encoder = True
|
||||
self.vae_config.load_decoder = True
|
||||
self.dit_config.expand_timesteps = self.expand_timesteps
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXWorld5BARPipelineConfig(DreamXWorld5BCamPipelineConfig):
|
||||
"""Pipeline config for DreamX-World-5B autoregressive forcing."""
|
||||
|
||||
dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_ar_dit_config)
|
||||
flow_shift: float | None = 5.0
|
||||
ti2v_task: bool = True
|
||||
is_causal: bool = True
|
||||
dmd_denoising_steps: tuple[int, ...] = (1000, 750, 500, 250)
|
||||
warp_denoising_step: bool = True
|
||||
context_noise: float = 0.1
|
||||
num_frames_per_block: int = 3
|
||||
color_correction_strength: float = 1.0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
self.dit_config.expand_timesteps = True
|
||||
@@ -1,122 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
|
||||
from fastvideo.configs.models.dits import Kandinsky5VideoConfig
|
||||
from fastvideo.configs.models.encoders import BaseEncoderOutput, CLIPTextConfig
|
||||
from fastvideo.configs.models.encoders.reason1 import Reason1Config
|
||||
from fastvideo.configs.models.vaes import HunyuanVAEConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
|
||||
|
||||
# Byte-exact copy of the upstream Kandinsky5/diffusers template, including the
|
||||
# "promt"/"scren" typos: the checkpoints were trained with this exact system # codespell:ignore promt,scren
|
||||
# prompt, and ENCODE_START_IDX below is the tokenized length of everything
|
||||
# before the user prompt. Fixing the typos shifts user content to index 127
|
||||
# and mis-conditions every generation.
|
||||
KANDINSKY5_PROMPT_TEMPLATE = "\n".join([
|
||||
"<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", # codespell:ignore promt
|
||||
"Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.",
|
||||
"Describe the location of the video, main characters or objects and their action.",
|
||||
"Describe the dynamism of the video and presented actions.",
|
||||
"Name the visual style of the video: whether it is a professional footage, user generated content, some kind of animation, video game or scren content.", # codespell:ignore scren
|
||||
"Describe the visual effects, postprocessing and transitions if they are presented in the video.",
|
||||
"Pay attention to the order of key actions shown in the scene.<|im_end|>",
|
||||
"<|im_start|>user\n{}<|im_end|>",
|
||||
])
|
||||
KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX = 129
|
||||
|
||||
|
||||
def kandinsky5_qwen_preprocess_text(prompt: str) -> str:
|
||||
if not prompt.strip():
|
||||
prompt = "."
|
||||
return KANDINSKY5_PROMPT_TEMPLATE.format(prompt)
|
||||
|
||||
|
||||
def kandinsky5_qwen_postprocess_text(outputs: BaseEncoderOutput,
|
||||
mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if outputs.hidden_states is None:
|
||||
raise RuntimeError("Kandinsky5 Qwen prompt embeddings require hidden_states.")
|
||||
hidden_states = outputs.hidden_states[-1]
|
||||
prompt_embeds = hidden_states[:, KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX:]
|
||||
mask = mask[:, KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX:]
|
||||
if prompt_embeds.shape[1] == 0:
|
||||
prompt_embeds = hidden_states[:, -1:]
|
||||
mask = torch.ones((mask.shape[0], 1), dtype=mask.dtype, device=mask.device)
|
||||
return prompt_embeds, mask
|
||||
|
||||
|
||||
def kandinsky5_clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
|
||||
if outputs.pooler_output is None:
|
||||
raise RuntimeError("Kandinsky5 CLIP pooled output is required.")
|
||||
return outputs.pooler_output
|
||||
|
||||
|
||||
@dataclass
|
||||
class Kandinsky5T2VConfig(PipelineConfig):
|
||||
"""Kandinsky-5.0 Lite text-to-video pipeline configuration."""
|
||||
|
||||
dit_config: DiTConfig = field(default_factory=Kandinsky5VideoConfig)
|
||||
vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig)
|
||||
|
||||
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Reason1Config(), CLIPTextConfig()))
|
||||
preprocess_text_funcs: tuple[Callable[[str], Any], ...] = field(
|
||||
default_factory=lambda: (kandinsky5_qwen_preprocess_text, preprocess_text))
|
||||
postprocess_text_funcs: tuple[Callable[..., Any], ...] = field(
|
||||
default_factory=lambda: (kandinsky5_qwen_postprocess_text, kandinsky5_clip_postprocess_text))
|
||||
|
||||
dit_precision: str = "bf16"
|
||||
vae_precision: str = "bf16"
|
||||
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", "bf16"))
|
||||
text_encoder_max_lengths: tuple[int, ...] = field(
|
||||
default_factory=lambda: (KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX + 512, 77))
|
||||
|
||||
flow_shift: float | None = 5.0
|
||||
vae_tiling: bool = True
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if len(self.text_encoder_configs) != 2:
|
||||
raise ValueError(f"Kandinsky5 pipeline requires exactly 2 text encoders (qwen and clip), "
|
||||
f"but got {len(self.text_encoder_configs)} encoder(s).")
|
||||
if len(self.text_encoder_precisions) != 2:
|
||||
raise ValueError("Kandinsky5 pipeline requires exactly 2 text encoder precisions, "
|
||||
f"but got {len(self.text_encoder_precisions)}.")
|
||||
if len(self.text_encoder_max_lengths) != 2:
|
||||
raise ValueError("Kandinsky5 pipeline requires exactly 2 text encoder max lengths, "
|
||||
f"but got {len(self.text_encoder_max_lengths)}.")
|
||||
|
||||
self.vae_config.load_encoder = False
|
||||
self.vae_config.load_decoder = True
|
||||
|
||||
qwen_cfg = self.text_encoder_configs[0]
|
||||
qwen_cfg.arch_config.output_hidden_states = True
|
||||
qwen_cfg.arch_config.tokenizer_kwargs.update({
|
||||
"padding": True,
|
||||
"truncation": True,
|
||||
"return_tensors": "pt",
|
||||
})
|
||||
|
||||
clip_cfg = self.text_encoder_configs[1]
|
||||
clip_cfg.arch_config.tokenizer_kwargs.update({
|
||||
"padding": "max_length",
|
||||
"max_length": 77,
|
||||
"truncation": True,
|
||||
"add_special_tokens": True,
|
||||
"return_tensors": "pt",
|
||||
})
|
||||
|
||||
|
||||
@dataclass
|
||||
class Kandinsky5I2VConfig(Kandinsky5T2VConfig):
|
||||
"""Kandinsky-5.0 image-to-video pipeline configuration."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
# I2V needs the VAE encoder to encode the conditioning image.
|
||||
self.vae_config.load_encoder = True
|
||||
@@ -23,7 +23,6 @@ If you only need to use the distributed environment without model parallelism,
|
||||
you can skip the model parallel initialization and destruction steps.
|
||||
"""
|
||||
import contextlib
|
||||
import gc
|
||||
import os
|
||||
import pickle
|
||||
import weakref
|
||||
@@ -1007,13 +1006,6 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
|
||||
if shutdown_ray:
|
||||
import ray # Lazy import Ray
|
||||
ray.shutdown()
|
||||
# Actually free GPU memory, as the name promises. FSDP-wrapped modules sit
|
||||
# in reference cycles (module <-> FSDP state <-> hooks), so dropping the
|
||||
# last user reference does not free their parameters until a gc pass runs.
|
||||
# Without this, back-to-back model-load tests in one pytest session OOM.
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0) -> list[int]:
|
||||
|
||||
@@ -437,7 +437,6 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
|
||||
# Causal-specific
|
||||
self.block_mask = None
|
||||
self.teacher_forcing_block_mask = None
|
||||
self.num_frame_per_block = config.arch_config.num_frames_per_block
|
||||
assert self.num_frame_per_block <= 3
|
||||
self.independent_first_frame = False
|
||||
@@ -501,70 +500,6 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
|
||||
return block_mask
|
||||
|
||||
@staticmethod
|
||||
def _prepare_teacher_forcing_mask(
|
||||
device: torch.device | str, num_frames: int = 21,
|
||||
frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1
|
||||
) -> BlockMask:
|
||||
"""Attention mask for the teacher-forcing ``[clean | noisy]`` sequence.
|
||||
|
||||
A noisy token attends to its own block plus the clean context of all
|
||||
strictly previous blocks; clean tokens are block-wise causal.
|
||||
"""
|
||||
if local_attn_size != -1:
|
||||
raise NotImplementedError(
|
||||
f"Teacher forcing ignores local_attn_size={local_attn_size}: "
|
||||
"unlike the block-wise causal mask, this mask always attends "
|
||||
"to the full clean context. Windowed teacher forcing is not "
|
||||
"implemented; use local_attn_size=-1 for teacher-forcing "
|
||||
"training.")
|
||||
total_length = num_frames * frame_seqlen * 2
|
||||
padded_length = math.ceil(total_length / 128) * 128 - total_length
|
||||
|
||||
clean_ends = num_frames * frame_seqlen
|
||||
context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
|
||||
noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
|
||||
noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
|
||||
noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
|
||||
noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
|
||||
|
||||
attention_block_size = frame_seqlen * num_frame_per_block
|
||||
frame_indices = torch.arange(
|
||||
start=0, end=num_frames * frame_seqlen,
|
||||
step=attention_block_size, device=device, dtype=torch.long
|
||||
)
|
||||
for start in frame_indices:
|
||||
context_ends[start:start + attention_block_size] = start + attention_block_size
|
||||
|
||||
noisy_image_start_list = torch.arange(
|
||||
num_frames * frame_seqlen, total_length,
|
||||
step=attention_block_size, device=device, dtype=torch.long
|
||||
)
|
||||
noisy_image_end_list = noisy_image_start_list + attention_block_size
|
||||
for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):
|
||||
noise_noise_starts[start:end] = start
|
||||
noise_noise_ends[start:end] = end
|
||||
noise_context_ends[start:end] = block_index * attention_block_size
|
||||
|
||||
def attention_mask(b, h, q_idx, kv_idx):
|
||||
clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])
|
||||
c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])
|
||||
c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])
|
||||
noise_mask = (q_idx >= clean_ends) & (c1 | c2)
|
||||
eye_mask = q_idx == kv_idx
|
||||
return eye_mask | clean_mask | noise_mask
|
||||
|
||||
block_mask = create_block_mask(
|
||||
attention_mask, B=None, H=None,
|
||||
Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length,
|
||||
_compile=False, device=device)
|
||||
|
||||
if not dist.is_initialized() or dist.get_rank() == 0:
|
||||
print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames")
|
||||
print(block_mask)
|
||||
|
||||
return block_mask
|
||||
|
||||
def _forward_inference(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -693,12 +628,9 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
|
||||
| None = None,
|
||||
start_frame: int = 0,
|
||||
clean_x: torch.Tensor | None = None,
|
||||
aug_t: torch.Tensor | None = None,
|
||||
**kwargs) -> torch.Tensor:
|
||||
|
||||
orig_dtype = hidden_states.dtype
|
||||
teacher_forcing = clean_x is not None
|
||||
if not isinstance(encoder_hidden_states, torch.Tensor):
|
||||
encoder_hidden_states = encoder_hidden_states[0]
|
||||
if isinstance(encoder_hidden_states_image,
|
||||
@@ -731,26 +663,15 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
freqs_cis = (freqs_cos,
|
||||
freqs_sin) if freqs_cos is not None else None
|
||||
|
||||
if teacher_forcing:
|
||||
if self.teacher_forcing_block_mask is None:
|
||||
self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask(
|
||||
device=hidden_states.device,
|
||||
num_frames=num_frames,
|
||||
frame_seqlen=post_patch_height * post_patch_width,
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
local_attn_size=self.local_attn_size,
|
||||
)
|
||||
block_mask = self.teacher_forcing_block_mask
|
||||
else:
|
||||
if self.block_mask is None:
|
||||
self.block_mask = self._prepare_blockwise_causal_attn_mask(
|
||||
device=hidden_states.device,
|
||||
num_frames=num_frames,
|
||||
frame_seqlen=post_patch_height * post_patch_width,
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
local_attn_size=self.local_attn_size
|
||||
)
|
||||
block_mask = self.block_mask
|
||||
# Construct blockwise causal attn mask
|
||||
if self.block_mask is None:
|
||||
self.block_mask = self._prepare_blockwise_causal_attn_mask(
|
||||
device=hidden_states.device,
|
||||
num_frames=num_frames,
|
||||
frame_seqlen=post_patch_height * post_patch_width,
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
local_attn_size=self.local_attn_size
|
||||
)
|
||||
|
||||
hidden_states = self.patch_embedding(hidden_states)
|
||||
grid_sizes = torch.stack(
|
||||
@@ -758,7 +679,6 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
|
||||
encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
|
||||
encoder_hidden_states_text = encoder_hidden_states
|
||||
|
||||
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
|
||||
timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
|
||||
@@ -774,35 +694,18 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
|
||||
assert encoder_hidden_states.dtype == orig_dtype
|
||||
|
||||
if teacher_forcing:
|
||||
# Tile RoPE/modulation so clean frame i and noisy frame i share a position.
|
||||
clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2)
|
||||
hidden_states = torch.cat([clean_tokens, hidden_states], dim=1)
|
||||
if aug_t is None:
|
||||
aug_t = torch.zeros_like(timestep)
|
||||
_, timestep_proj_clean, _, _ = self.condition_embedder(
|
||||
aug_t.flatten(), encoder_hidden_states_text, None)
|
||||
timestep_proj_clean = timestep_proj_clean.unflatten(
|
||||
1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
|
||||
timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1)
|
||||
freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0),
|
||||
torch.cat([freqs_sin, freqs_sin], dim=0))
|
||||
|
||||
# 4. Transformer blocks
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
for block in self.blocks:
|
||||
hidden_states = self._gradient_checkpointing_func(
|
||||
block, hidden_states, encoder_hidden_states,
|
||||
timestep_proj, freqs_cis,
|
||||
block_mask=block_mask)
|
||||
block_mask=self.block_mask)
|
||||
else:
|
||||
for block in self.blocks:
|
||||
hidden_states = block(hidden_states, encoder_hidden_states,
|
||||
timestep_proj, freqs_cis,
|
||||
block_mask=block_mask)
|
||||
|
||||
if teacher_forcing:
|
||||
hidden_states = hidden_states[:, hidden_states.shape[1] // 2:]
|
||||
block_mask=self.block_mask)
|
||||
|
||||
# 5. Output norm, projection & unpatchify
|
||||
temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
|
||||
|
||||
@@ -1,511 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig
|
||||
from fastvideo.distributed.communication_op import (
|
||||
sequence_model_parallel_all_gather_with_unpad,
|
||||
sequence_model_parallel_shard)
|
||||
from fastvideo.distributed.parallel_state import get_sp_world_size
|
||||
from fastvideo.layers.layernorm import RMSNorm
|
||||
from fastvideo.layers.linear import ReplicatedLinear
|
||||
from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
|
||||
from fastvideo.models.dits.wanvideo import (LayerNormScaleShift,
|
||||
PatchEmbed,
|
||||
WanTimeTextImageEmbedding,
|
||||
WanTransformer3DModel,
|
||||
WanTransformerBlock)
|
||||
from fastvideo.platforms import AttentionBackendEnum, current_platform
|
||||
from fastvideo.attention import LocalAttention
|
||||
from fastvideo.distributed.parallel_state import get_sp_world_size
|
||||
from fastvideo.layers.quantization import QuantizationConfig
|
||||
from fastvideo.models.dits.base import BaseDiT
|
||||
|
||||
|
||||
def _dreamx_invert_se3(transforms: torch.Tensor) -> torch.Tensor:
|
||||
assert transforms.shape[-2:] == (4, 4)
|
||||
rot_inv = transforms[..., :3, :3].transpose(-1, -2)
|
||||
out = torch.zeros_like(transforms)
|
||||
out[..., :3, :3] = rot_inv
|
||||
out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv,
|
||||
transforms[..., :3, 3])
|
||||
out[..., 3, 3] = 1.0
|
||||
return out.to(dtype=transforms.dtype)
|
||||
|
||||
|
||||
def _dreamx_lift_k(intrinsics: torch.Tensor) -> torch.Tensor:
|
||||
assert intrinsics.shape[-2:] == (3, 3)
|
||||
out = torch.zeros(intrinsics.shape[:-2] + (4, 4),
|
||||
device=intrinsics.device,
|
||||
dtype=intrinsics.dtype)
|
||||
out[..., :3, :3] = intrinsics
|
||||
out[..., 3, 3] = 1.0
|
||||
return out
|
||||
|
||||
|
||||
def _dreamx_invert_k(intrinsics: torch.Tensor) -> torch.Tensor:
|
||||
assert intrinsics.shape[-2:] == (3, 3)
|
||||
out = torch.zeros_like(intrinsics)
|
||||
out[..., 0, 0] = 1.0 / intrinsics[..., 0, 0]
|
||||
out[..., 1, 1] = 1.0 / intrinsics[..., 1, 1]
|
||||
out[..., 0, 2] = -intrinsics[..., 0, 2] / intrinsics[..., 0, 0]
|
||||
out[..., 1, 2] = -intrinsics[..., 1, 2] / intrinsics[..., 1, 1]
|
||||
out[..., 2, 2] = 1.0
|
||||
return out.to(dtype=intrinsics.dtype)
|
||||
|
||||
|
||||
def _dreamx_apply_tiled_projmat(feats: torch.Tensor,
|
||||
matrix: torch.Tensor) -> torch.Tensor:
|
||||
batch, num_heads, seq_len, feat_dim = feats.shape
|
||||
proj_dim = matrix.shape[-1]
|
||||
assert feat_dim % proj_dim == 0
|
||||
|
||||
if matrix.shape[1] == seq_len:
|
||||
feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim,
|
||||
proj_dim)
|
||||
out = torch.einsum("btij,bntpj->bntpi", matrix, feats)
|
||||
return out.reshape(batch, num_heads, seq_len, feat_dim)
|
||||
|
||||
cameras = matrix.shape[1]
|
||||
assert seq_len > cameras and seq_len % cameras == 0
|
||||
feats = feats.reshape(batch, num_heads, cameras, -1,
|
||||
feat_dim // proj_dim, proj_dim)
|
||||
out = torch.einsum("bcij,bncpkj->bncpki", matrix, feats)
|
||||
return out.reshape(batch, num_heads, seq_len, feat_dim)
|
||||
|
||||
|
||||
def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||||
viewmats: torch.Tensor, intrinsics: torch.Tensor):
|
||||
batch, num_heads, seq_len, head_dim = q.shape
|
||||
cameras = viewmats.shape[1]
|
||||
assert q.shape == k.shape == v.shape
|
||||
assert viewmats.shape == (batch, cameras, 4, 4)
|
||||
assert intrinsics.shape == (batch, cameras, 3, 3)
|
||||
assert head_dim % 4 == 0
|
||||
|
||||
intrinsics_norm = torch.zeros_like(intrinsics)
|
||||
intrinsics_norm[..., 0, 0] = intrinsics[..., 0, 0]
|
||||
intrinsics_norm[..., 1, 1] = intrinsics[..., 1, 1]
|
||||
intrinsics_norm[..., 2, 2] = 1.0
|
||||
|
||||
proj = torch.einsum("...ij,...jk->...ik",
|
||||
_dreamx_lift_k(intrinsics_norm), viewmats)
|
||||
proj_t = proj.transpose(-1, -2).to(dtype=viewmats.dtype)
|
||||
proj_inv = torch.einsum(
|
||||
"...ij,...jk->...ik",
|
||||
_dreamx_invert_se3(viewmats),
|
||||
_dreamx_lift_k(_dreamx_invert_k(intrinsics_norm)),
|
||||
).to(dtype=viewmats.dtype)
|
||||
|
||||
q = _dreamx_apply_tiled_projmat(q, proj_t)
|
||||
k = _dreamx_apply_tiled_projmat(k, proj_inv)
|
||||
v = _dreamx_apply_tiled_projmat(v, proj_inv)
|
||||
return q, k, v, proj
|
||||
|
||||
|
||||
class DreamXPropeSelfAttention(nn.Module):
|
||||
"""DreamX-World parallel PRoPE camera self-attention branch."""
|
||||
|
||||
def __init__(self,
|
||||
dim: int,
|
||||
attn_dim: int,
|
||||
num_heads: int,
|
||||
qk_norm: str | bool = True,
|
||||
eps: float = 1e-6,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = ""):
|
||||
super().__init__()
|
||||
assert attn_dim % num_heads == 0
|
||||
self.attn_dim = attn_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = attn_dim // num_heads
|
||||
self.qk_norm = qk_norm
|
||||
|
||||
self.q_proj = ReplicatedLinear(dim,
|
||||
attn_dim,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.q_proj")
|
||||
self.k_proj = ReplicatedLinear(dim,
|
||||
attn_dim,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.k_proj")
|
||||
self.v_proj = ReplicatedLinear(dim,
|
||||
attn_dim,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.v_proj")
|
||||
self.out_proj = ReplicatedLinear(attn_dim,
|
||||
dim,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.out_proj")
|
||||
|
||||
if qk_norm == "rms_norm":
|
||||
self.norm_q = RMSNorm(self.head_dim, eps=eps)
|
||||
self.norm_k = RMSNorm(self.head_dim, eps=eps)
|
||||
elif qk_norm in (True, "rms_norm_across_heads"):
|
||||
self.norm_q = RMSNorm(attn_dim, eps=eps)
|
||||
self.norm_k = RMSNorm(attn_dim, eps=eps)
|
||||
elif qk_norm is False:
|
||||
self.norm_q = nn.Identity()
|
||||
self.norm_k = nn.Identity()
|
||||
else:
|
||||
raise ValueError(f"Unsupported qk_norm for DreamX PRoPE: {qk_norm}")
|
||||
|
||||
nn.init.zeros_(self.out_proj.weight)
|
||||
if self.out_proj.bias is not None:
|
||||
nn.init.zeros_(self.out_proj.bias)
|
||||
|
||||
self.attn = LocalAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=self.head_dim,
|
||||
dropout_rate=0,
|
||||
softmax_scale=None,
|
||||
causal=False,
|
||||
supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.TORCH_SDPA))
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor,
|
||||
y_camera: dict[str, torch.Tensor]) -> torch.Tensor:
|
||||
if get_sp_world_size() > 1:
|
||||
# The transformer shards the sequence before the block loop and
|
||||
# this branch uses LocalAttention (no all-to-all): under
|
||||
# sequence parallelism each rank would attend only within its
|
||||
# own shard — silently wrong output. Fail loudly until this
|
||||
# path is ported to DistributedAttention and validated.
|
||||
raise NotImplementedError(
|
||||
"DreamXPropeSelfAttention does not support sequence "
|
||||
"parallelism yet (LocalAttention on a sharded sequence "
|
||||
"corrupts output). Run with sp_size=1.")
|
||||
batch_size, seq_len, _ = hidden_states.shape
|
||||
|
||||
query, _ = self.q_proj(hidden_states)
|
||||
key, _ = self.k_proj(hidden_states)
|
||||
value, _ = self.v_proj(hidden_states)
|
||||
|
||||
if self.qk_norm == "rms_norm":
|
||||
query = query.view(batch_size, seq_len, self.num_heads,
|
||||
self.head_dim)
|
||||
key = key.view(batch_size, seq_len, self.num_heads, self.head_dim)
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
else:
|
||||
query = self.norm_q(query).view(batch_size, seq_len,
|
||||
self.num_heads, self.head_dim)
|
||||
key = self.norm_k(key).view(batch_size, seq_len, self.num_heads,
|
||||
self.head_dim)
|
||||
|
||||
value = value.view(batch_size, seq_len, self.num_heads, self.head_dim)
|
||||
query = query.transpose(1, 2)
|
||||
key = key.transpose(1, 2)
|
||||
value = value.transpose(1, 2)
|
||||
|
||||
query, key, value, output_projection = _dreamx_prope_qkv(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
viewmats=y_camera["viewmats"],
|
||||
intrinsics=y_camera["K"],
|
||||
)
|
||||
|
||||
out = self.attn(query.transpose(1, 2), key.transpose(1, 2),
|
||||
value.transpose(1, 2))
|
||||
out = _dreamx_apply_tiled_projmat(out.transpose(1, 2),
|
||||
output_projection).transpose(1, 2)
|
||||
out = out.flatten(2)
|
||||
out, _ = self.out_proj(out)
|
||||
return out
|
||||
|
||||
|
||||
class DreamXWorldTransformerBlock(WanTransformerBlock):
|
||||
|
||||
def __init__(self,
|
||||
dim: int,
|
||||
ffn_dim: int,
|
||||
num_heads: int,
|
||||
qk_norm: str = "rms_norm_across_heads",
|
||||
cross_attn_norm: bool = False,
|
||||
eps: float = 1e-6,
|
||||
added_kv_proj_dim: int | None = None,
|
||||
supported_attention_backends: tuple[AttentionBackendEnum, ...]
|
||||
| None = None,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
add_control_adapter: bool = True,
|
||||
cam_method: str | None = "prope",
|
||||
attn_compress: int = 1,
|
||||
cam_self_attn_layers: tuple[int, ...] | None = None,
|
||||
layer_idx: int | None = None):
|
||||
super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm,
|
||||
eps, added_kv_proj_dim,
|
||||
supported_attention_backends, quant_config, prefix)
|
||||
self.cam_self_attn = None
|
||||
add_cam_attn = add_control_adapter and cam_method == "prope"
|
||||
if add_cam_attn and cam_self_attn_layers is not None:
|
||||
add_cam_attn = layer_idx in cam_self_attn_layers
|
||||
if add_cam_attn:
|
||||
if num_heads % attn_compress != 0 or dim % attn_compress != 0:
|
||||
raise ValueError("DreamX attn_compress must divide dim and num_heads")
|
||||
self.cam_self_attn = DreamXPropeSelfAttention(
|
||||
dim,
|
||||
dim // attn_compress,
|
||||
num_heads // attn_compress,
|
||||
qk_norm=qk_norm,
|
||||
eps=eps,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.cam_self_attn")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
freqs_cis: tuple[torch.Tensor, torch.Tensor],
|
||||
original_seq_len: int,
|
||||
y_camera: dict[str, torch.Tensor] | None = None,
|
||||
) -> torch.Tensor:
|
||||
if hidden_states.dim() == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
orig_dtype = hidden_states.dtype
|
||||
|
||||
if temb.dim() == 4:
|
||||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
|
||||
self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(
|
||||
6, dim=2)
|
||||
shift_msa = shift_msa.squeeze(2)
|
||||
scale_msa = scale_msa.squeeze(2)
|
||||
gate_msa = gate_msa.squeeze(2)
|
||||
c_shift_msa = c_shift_msa.squeeze(2)
|
||||
c_scale_msa = c_scale_msa.squeeze(2)
|
||||
c_gate_msa = c_gate_msa.squeeze(2)
|
||||
else:
|
||||
e = self.scale_shift_table + temb.float()
|
||||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
|
||||
6, dim=1)
|
||||
assert shift_msa.dtype == torch.float32
|
||||
|
||||
norm_hidden_states = (self.norm1(hidden_states.float()) *
|
||||
(1 + scale_msa) + shift_msa).to(orig_dtype)
|
||||
query, _ = self.to_q(norm_hidden_states)
|
||||
key, _ = self.to_k(norm_hidden_states)
|
||||
value, _ = self.to_v(norm_hidden_states)
|
||||
|
||||
if self.norm_q is not None:
|
||||
query = self.norm_q(query)
|
||||
if self.norm_k is not None:
|
||||
key = self.norm_k(key)
|
||||
|
||||
query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
|
||||
attn_output, _ = self.attn1(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
original_seq_len,
|
||||
freqs_cis=freqs_cis,
|
||||
)
|
||||
attn_output = attn_output.flatten(2)
|
||||
attn_output, _ = self.to_out(attn_output)
|
||||
attn_output = attn_output.squeeze(1)
|
||||
if self.cam_self_attn is not None and y_camera is not None:
|
||||
attn_output = attn_output + self.cam_self_attn(
|
||||
norm_hidden_states, y_camera)
|
||||
|
||||
null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
|
||||
norm_hidden_states, hidden_states = self.self_attn_residual_norm(
|
||||
hidden_states, attn_output, gate_msa, null_shift, null_scale)
|
||||
norm_hidden_states, hidden_states = norm_hidden_states.to(
|
||||
orig_dtype), hidden_states.to(orig_dtype)
|
||||
|
||||
attn_output = self.attn2(norm_hidden_states,
|
||||
context=encoder_hidden_states,
|
||||
context_lens=None)
|
||||
norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
|
||||
hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
|
||||
norm_hidden_states, hidden_states = norm_hidden_states.to(
|
||||
orig_dtype), hidden_states.to(orig_dtype)
|
||||
|
||||
ff_output = self.ffn(norm_hidden_states)
|
||||
hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
|
||||
hidden_states = hidden_states.to(orig_dtype)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class DreamXWorldTransformer3DModel(WanTransformer3DModel):
|
||||
_fsdp_shard_conditions = DreamXWorldConfig()._fsdp_shard_conditions
|
||||
_compile_conditions = DreamXWorldConfig()._compile_conditions
|
||||
_supported_attention_backends = DreamXWorldConfig(
|
||||
)._supported_attention_backends
|
||||
param_names_mapping = DreamXWorldConfig().param_names_mapping
|
||||
reverse_param_names_mapping = DreamXWorldConfig().reverse_param_names_mapping
|
||||
lora_param_names_mapping = DreamXWorldConfig().lora_param_names_mapping
|
||||
|
||||
def __init__(self, config: DreamXWorldConfig, hf_config: dict[str,
|
||||
Any]) -> None:
|
||||
BaseDiT.__init__(self, config=config, hf_config=hf_config)
|
||||
self.quant_config = config.quant_config
|
||||
|
||||
inner_dim = config.num_attention_heads * config.attention_head_dim
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.in_channels = config.in_channels
|
||||
self.out_channels = config.out_channels
|
||||
self.num_channels_latents = config.num_channels_latents
|
||||
self.patch_size = config.patch_size
|
||||
self.text_len = config.text_len
|
||||
|
||||
assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
|
||||
|
||||
self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
|
||||
embed_dim=inner_dim,
|
||||
patch_size=config.patch_size,
|
||||
flatten=False)
|
||||
self.condition_embedder = WanTimeTextImageEmbedding(
|
||||
dim=inner_dim,
|
||||
time_freq_dim=config.freq_dim,
|
||||
text_embed_dim=config.text_dim,
|
||||
image_embed_dim=config.image_dim,
|
||||
)
|
||||
self.blocks = nn.ModuleList([
|
||||
DreamXWorldTransformerBlock(
|
||||
inner_dim,
|
||||
config.ffn_dim,
|
||||
config.num_attention_heads,
|
||||
config.qk_norm,
|
||||
config.cross_attn_norm,
|
||||
config.eps,
|
||||
config.added_kv_proj_dim,
|
||||
self._supported_attention_backends,
|
||||
quant_config=config.quant_config,
|
||||
prefix=f"{config.prefix}.blocks.{i}",
|
||||
add_control_adapter=config.add_control_adapter,
|
||||
cam_method=config.cam_method,
|
||||
attn_compress=config.attn_compress,
|
||||
cam_self_attn_layers=config.cam_self_attn_layers,
|
||||
layer_idx=i)
|
||||
for i in range(config.num_layers)
|
||||
])
|
||||
self.norm_out = LayerNormScaleShift(inner_dim,
|
||||
norm_type="layer",
|
||||
eps=config.eps,
|
||||
elementwise_affine=False,
|
||||
dtype=torch.float32,
|
||||
compute_dtype=torch.float32)
|
||||
self.proj_out = nn.Linear(
|
||||
inner_dim, config.out_channels * math.prod(config.patch_size))
|
||||
self.scale_shift_table = nn.Parameter(
|
||||
torch.randn(1, 2, inner_dim) / inner_dim**0.5)
|
||||
self.gradient_checkpointing = False
|
||||
self.__post_init__()
|
||||
|
||||
def forward(self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor | list[torch.Tensor],
|
||||
timestep: torch.LongTensor,
|
||||
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
|
||||
| None = None,
|
||||
guidance=None,
|
||||
y_camera: dict[str, torch.Tensor] | None = None,
|
||||
**kwargs) -> torch.Tensor:
|
||||
orig_dtype = hidden_states.dtype
|
||||
if encoder_hidden_states is not None and not isinstance(
|
||||
encoder_hidden_states, torch.Tensor):
|
||||
encoder_hidden_states = encoder_hidden_states[0]
|
||||
if isinstance(encoder_hidden_states_image,
|
||||
list) and len(encoder_hidden_states_image) > 0:
|
||||
encoder_hidden_states_image = encoder_hidden_states_image[0]
|
||||
else:
|
||||
encoder_hidden_states_image = None
|
||||
|
||||
batch_size, _, num_frames, height, width = hidden_states.shape
|
||||
p_t, p_h, p_w = self.patch_size
|
||||
post_patch_num_frames = num_frames // p_t
|
||||
post_patch_height = height // p_h
|
||||
post_patch_width = width // p_w
|
||||
|
||||
d = self.hidden_size // self.num_attention_heads
|
||||
rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
|
||||
freqs_cos, freqs_sin = get_rotary_pos_embed(
|
||||
(post_patch_num_frames, post_patch_height, post_patch_width),
|
||||
self.hidden_size,
|
||||
self.num_attention_heads,
|
||||
rope_dim_list,
|
||||
dtype=torch.float32 if current_platform.is_mps() else torch.float64,
|
||||
rope_theta=10000)
|
||||
freqs_cis = (freqs_cos.to(hidden_states.device).float(),
|
||||
freqs_sin.to(hidden_states.device).float())
|
||||
|
||||
hidden_states = self.patch_embedding(hidden_states)
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
hidden_states, original_seq_len = sequence_model_parallel_shard(
|
||||
hidden_states, dim=1)
|
||||
|
||||
if timestep.dim() == 2:
|
||||
ts_seq_len = timestep.shape[1]
|
||||
timestep = timestep.flatten()
|
||||
else:
|
||||
ts_seq_len = None
|
||||
|
||||
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
encoder_hidden_states_image,
|
||||
timestep_seq_len=ts_seq_len)
|
||||
if ts_seq_len is not None:
|
||||
timestep_proj = timestep_proj.unflatten(2, (6, -1))
|
||||
else:
|
||||
timestep_proj = timestep_proj.unflatten(1, (6, -1))
|
||||
|
||||
if encoder_hidden_states_image is not None:
|
||||
if encoder_hidden_states is not None:
|
||||
encoder_hidden_states = torch.concat(
|
||||
[encoder_hidden_states_image, encoder_hidden_states], dim=1)
|
||||
else:
|
||||
encoder_hidden_states = encoder_hidden_states_image
|
||||
|
||||
if current_platform.is_mps() or current_platform.is_npu():
|
||||
encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
|
||||
|
||||
assert encoder_hidden_states.dtype == orig_dtype
|
||||
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
for block in self.blocks:
|
||||
hidden_states = self._gradient_checkpointing_func(
|
||||
block, hidden_states, encoder_hidden_states,
|
||||
timestep_proj, freqs_cis, original_seq_len, y_camera)
|
||||
else:
|
||||
for block in self.blocks:
|
||||
hidden_states = block(hidden_states, encoder_hidden_states,
|
||||
timestep_proj, freqs_cis, original_seq_len,
|
||||
y_camera=y_camera)
|
||||
|
||||
if temb.dim() == 3:
|
||||
shift, scale = (self.scale_shift_table.unsqueeze(0) +
|
||||
temb.unsqueeze(2)).chunk(2, dim=2)
|
||||
shift = shift.squeeze(2)
|
||||
scale = scale.squeeze(2)
|
||||
else:
|
||||
shift, scale = (self.scale_shift_table +
|
||||
temb.unsqueeze(1)).chunk(2, dim=1)
|
||||
|
||||
hidden_states = self.norm_out(hidden_states, shift, scale)
|
||||
hidden_states = sequence_model_parallel_all_gather_with_unpad(
|
||||
hidden_states, original_seq_len, dim=1)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
|
||||
post_patch_height,
|
||||
post_patch_width, p_t, p_h, p_w,
|
||||
-1)
|
||||
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
|
||||
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
EntryClass = DreamXWorldTransformer3DModel
|
||||
@@ -1,920 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World autoregressive causal DiT.
|
||||
|
||||
Adapted from DreamX-World's Apache-2.0
|
||||
``wan/modules/causal_camera_model_2_2_prope_infinity.py``. The implementation is
|
||||
kept native to FastVideo: no production import from DreamX, Diffusers, or
|
||||
Transformers is required.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig
|
||||
from fastvideo.layers.linear import ReplicatedLinear
|
||||
from fastvideo.models.dits.base import BaseDiT
|
||||
from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat,
|
||||
_dreamx_prope_qkv)
|
||||
|
||||
|
||||
def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
|
||||
# Deliberately raw SDPA rather than fastvideo.attention.LocalAttention:
|
||||
# (1) LocalAttention dispatches through the attention-backend registry, so
|
||||
# FLASH_ATTN could be selected and its kernel is not bit-identical to
|
||||
# torch SDPA — the AR KV-cache rollout must stay numerically frozen;
|
||||
# (2) LocalAttention requires an active ForwardContext, which direct
|
||||
# transformer invocations (parity tests) do not set;
|
||||
# (3) the sibling causal model keeps raw SDPA in the same KV-cache window
|
||||
# path (matrixgame2/causal_model.py).
|
||||
# Sequence-parallel gap: this model never shards the sequence; run with
|
||||
# sp_size=1 (see fastvideo/layers/AGENTS.md on documenting raw SDPA).
|
||||
q_bhld = q.transpose(1, 2)
|
||||
k_bhld = k.transpose(1, 2)
|
||||
v_bhld = v.transpose(1, 2)
|
||||
out = F.scaled_dot_product_attention(q_bhld, k_bhld, v_bhld, dropout_p=0.0)
|
||||
return out.transpose(1, 2)
|
||||
|
||||
|
||||
def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||||
viewmats: torch.Tensor, Ks: torch.Tensor):
|
||||
q, k, v, output_projection = _dreamx_prope_qkv(q, k, v, viewmats, Ks)
|
||||
|
||||
def apply_fn_o(x: torch.Tensor) -> torch.Tensor:
|
||||
return _dreamx_apply_tiled_projmat(x, output_projection)
|
||||
|
||||
return q, k, v, apply_fn_o
|
||||
|
||||
|
||||
def sinusoidal_embedding_1d(dim, position):
|
||||
assert dim % 2 == 0
|
||||
half = dim // 2
|
||||
position = position.type(torch.float64)
|
||||
sinusoid = torch.outer(
|
||||
position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
|
||||
return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
|
||||
|
||||
|
||||
def rope_params(max_seq_len, dim, theta=10000):
|
||||
assert dim % 2 == 0
|
||||
freqs = torch.outer(
|
||||
torch.arange(max_seq_len),
|
||||
1.0 / torch.pow(theta,
|
||||
torch.arange(0, dim, 2).to(torch.float64).div(dim)))
|
||||
return torch.polar(torch.ones_like(freqs), freqs)
|
||||
|
||||
|
||||
class WanRMSNorm(nn.Module):
|
||||
"""Kept private instead of fastvideo.layers.layernorm.RMSNorm.
|
||||
|
||||
The official DreamX-World ``model_2_2.py`` computes the RMS statistics in
|
||||
the *input* dtype — the upstream code has the fp32 upcast explicitly
|
||||
commented out (``# return self._norm(x.float())...``). FastVideo's RMSNorm
|
||||
always normalizes in fp32, which is not bit-identical under bf16, so the
|
||||
verbatim implementation stays.
|
||||
"""
|
||||
|
||||
def __init__(self, dim, eps=1e-5):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x):
|
||||
return self._norm(x).type_as(x) * self.weight
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
|
||||
|
||||
|
||||
class WanLayerNorm(nn.LayerNorm):
|
||||
"""Kept private instead of fastvideo.layers.layernorm.FP32LayerNorm.
|
||||
|
||||
The official DreamX-World ``model_2_2.py`` normalizes in the *input* dtype
|
||||
(no ``x.float()`` upcast, unlike Wan2.1). FP32LayerNorm casts input and
|
||||
affine params to fp32, which is not bit-identical under bf16, so the
|
||||
verbatim implementation stays.
|
||||
"""
|
||||
|
||||
def __init__(self, dim, eps=1e-6, elementwise_affine=False):
|
||||
super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
|
||||
|
||||
def forward(self, x):
|
||||
return super().forward(x).type_as(x)
|
||||
|
||||
|
||||
class WanCrossAttention(nn.Module):
|
||||
def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.q = ReplicatedLinear(dim, dim)
|
||||
self.k = ReplicatedLinear(dim, dim)
|
||||
self.v = ReplicatedLinear(dim, dim)
|
||||
self.o = ReplicatedLinear(dim, dim)
|
||||
self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
def _kv(self, context, b, n, d):
|
||||
k, _ = self.k(context)
|
||||
k = self.norm_k(k).view(b, -1, n, d)
|
||||
v, _ = self.v(context)
|
||||
v = v.view(b, -1, n, d)
|
||||
return k, v
|
||||
|
||||
def forward(self, x, context, context_lens, crossattn_cache=None):
|
||||
b, n, d = x.size(0), self.num_heads, self.head_dim
|
||||
q, _ = self.q(x)
|
||||
q = self.norm_q(q).view(b, -1, n, d)
|
||||
|
||||
if crossattn_cache is not None:
|
||||
if not crossattn_cache["is_init"]:
|
||||
crossattn_cache["is_init"] = True
|
||||
k, v = self._kv(context, b, n, d)
|
||||
crossattn_cache["k"] = k
|
||||
crossattn_cache["v"] = v
|
||||
else:
|
||||
k = crossattn_cache["k"]
|
||||
v = crossattn_cache["v"]
|
||||
else:
|
||||
k, v = self._kv(context, b, n, d)
|
||||
|
||||
x = attention(q, k, v)
|
||||
x = x.flatten(2)
|
||||
out, _ = self.o(x)
|
||||
return out
|
||||
|
||||
|
||||
def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_indices=None):
|
||||
"""
|
||||
Apply Block-Relativistic RoPE to input tensor.
|
||||
Adapted from Infinity-RoPE (https://arxiv.org/abs/2511.20649).
|
||||
|
||||
Args:
|
||||
x: Input tensor [B, L, num_heads, head_dim]
|
||||
grid_sizes: Tensor [B, 3] containing (F, H, W)
|
||||
freqs: RoPE frequencies
|
||||
start_frame: Starting frame index for sequential RoPE
|
||||
relative_frame_indices: Optional tensor [F] specifying explicit frame indices
|
||||
for Block-Relativistic RoPE. Overrides start_frame if provided.
|
||||
"""
|
||||
n, c = x.size(2), x.size(3) // 2
|
||||
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
||||
|
||||
output = []
|
||||
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
||||
seq_len = f * h * w
|
||||
x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
|
||||
seq_len, n, -1, 2))
|
||||
|
||||
if relative_frame_indices is not None:
|
||||
frame_indices = relative_frame_indices.long()
|
||||
freqs_temporal = freqs[0][frame_indices].view(f, 1, 1, -1).expand(f, h, w, -1)
|
||||
else:
|
||||
freqs_temporal = freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1)
|
||||
|
||||
freqs_i = torch.cat([
|
||||
freqs_temporal,
|
||||
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
|
||||
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
|
||||
], dim=-1).reshape(seq_len, 1, -1)
|
||||
|
||||
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
||||
x_i = torch.cat([x_i, x[i, seq_len:]])
|
||||
output.append(x_i)
|
||||
|
||||
return torch.stack(output).type_as(x)
|
||||
|
||||
|
||||
class CausalWanSelfAttention(nn.Module):
|
||||
"""Self-attention with KV cache and Block-Relativistic RoPE for causal inference."""
|
||||
|
||||
def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1,
|
||||
qk_norm=True, eps=1e-6):
|
||||
assert dim % num_heads == 0
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.local_attn_size = local_attn_size
|
||||
self.sink_size = sink_size
|
||||
self.qk_norm = qk_norm
|
||||
self.eps = eps
|
||||
self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
|
||||
|
||||
self.q = ReplicatedLinear(dim, dim)
|
||||
self.k = ReplicatedLinear(dim, dim)
|
||||
self.v = ReplicatedLinear(dim, dim)
|
||||
self.o = ReplicatedLinear(dim, dim)
|
||||
self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache,
|
||||
current_start=0, cache_start=None, sink_recache_after_switch=False):
|
||||
"""
|
||||
Args:
|
||||
x: Shape [B, L, C]
|
||||
seq_lens: Shape [B]
|
||||
grid_sizes: Shape [B, 3] containing (F, H, W)
|
||||
freqs: RoPE frequencies [1024, head_dim / 2]
|
||||
kv_cache: Dict with 'k', 'v', 'global_end_index', 'local_end_index'
|
||||
current_start: Current position in the global token sequence
|
||||
"""
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
if cache_start is None:
|
||||
cache_start = current_start
|
||||
|
||||
q, _ = self.q(x)
|
||||
q = self.norm_q(q).view(b, s, n, d)
|
||||
k, _ = self.k(x)
|
||||
k = self.norm_k(k).view(b, s, n, d)
|
||||
v, _ = self.v(x)
|
||||
v = v.view(b, s, n, d)
|
||||
|
||||
frame_seqlen = math.prod(grid_sizes[0][1:]).item()
|
||||
num_new_frames = grid_sizes[0][0].item()
|
||||
current_end = current_start + q.shape[1]
|
||||
sink_tokens = self.sink_size * frame_seqlen
|
||||
kv_cache_size = kv_cache["k"].shape[1]
|
||||
num_new_tokens = q.shape[1]
|
||||
|
||||
cache_update_info = None
|
||||
is_recompute = current_end <= kv_cache["global_end_index"].item() and current_start > 0
|
||||
|
||||
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
|
||||
num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
|
||||
# === ROLLING MODE: cache full, evict oldest non-sink tokens ===
|
||||
num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
|
||||
num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
|
||||
local_end_index = kv_cache["local_end_index"].item() + current_end - \
|
||||
kv_cache["global_end_index"].item() - num_evicted_tokens
|
||||
local_start_index = local_end_index - num_new_tokens
|
||||
|
||||
temp_k = kv_cache["k"].detach().clone()
|
||||
temp_v = kv_cache["v"].detach().clone()
|
||||
temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
|
||||
temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
|
||||
temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
|
||||
temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
|
||||
|
||||
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
|
||||
roped_offset = max(0, write_start_index - local_start_index)
|
||||
write_len = max(0, local_end_index - write_start_index)
|
||||
if write_len > 0:
|
||||
temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
|
||||
temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
|
||||
|
||||
# Block-Relativistic RoPE: query uses window-relative indices
|
||||
query_relative_indices = torch.arange(
|
||||
self.local_attn_size - num_new_frames, self.local_attn_size, device=q.device)
|
||||
roped_query = block_relativistic_rope(
|
||||
q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
|
||||
|
||||
# Block-Relativistic RoPE: cached K uses position-in-window indices
|
||||
num_cache_frames = local_end_index // frame_seqlen
|
||||
cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
|
||||
cache_grid_sizes = grid_sizes.clone()
|
||||
cache_grid_sizes[0, 0] = num_cache_frames
|
||||
roped_temp_k = block_relativistic_rope(
|
||||
temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
|
||||
cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
|
||||
|
||||
cache_update_info = {
|
||||
"action": "roll_and_insert",
|
||||
"sink_tokens": sink_tokens,
|
||||
"num_rolled_tokens": num_rolled_tokens,
|
||||
"num_evicted_tokens": num_evicted_tokens,
|
||||
"local_start_index": local_start_index,
|
||||
"local_end_index": local_end_index,
|
||||
"write_start_index": write_start_index,
|
||||
"write_end_index": local_end_index,
|
||||
"new_k": k[:, roped_offset:roped_offset + write_len],
|
||||
"new_v": v[:, roped_offset:roped_offset + write_len],
|
||||
"current_end": current_end,
|
||||
"is_recompute": is_recompute
|
||||
}
|
||||
else:
|
||||
# === DIRECT INSERT MODE: cache not yet full ===
|
||||
local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
|
||||
local_start_index = local_end_index - num_new_tokens
|
||||
|
||||
temp_k = kv_cache["k"].detach().clone()
|
||||
temp_v = kv_cache["v"].detach().clone()
|
||||
|
||||
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
|
||||
if sink_recache_after_switch:
|
||||
write_start_index = local_start_index
|
||||
roped_offset = max(0, write_start_index - local_start_index)
|
||||
write_len = max(0, local_end_index - write_start_index)
|
||||
if write_len > 0:
|
||||
temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
|
||||
temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
|
||||
|
||||
# RoPE with relative indices (growing sequentially before cache fills)
|
||||
current_frame_in_window = local_start_index // frame_seqlen
|
||||
query_relative_indices = torch.arange(
|
||||
current_frame_in_window, current_frame_in_window + num_new_frames, device=q.device)
|
||||
roped_query = block_relativistic_rope(
|
||||
q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
|
||||
|
||||
num_cache_frames = local_end_index // frame_seqlen
|
||||
cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
|
||||
cache_grid_sizes = grid_sizes.clone()
|
||||
cache_grid_sizes[0, 0] = num_cache_frames
|
||||
roped_temp_k = block_relativistic_rope(
|
||||
temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
|
||||
cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
|
||||
|
||||
cache_update_info = {
|
||||
"action": "direct_insert",
|
||||
"local_start_index": local_start_index,
|
||||
"local_end_index": local_end_index,
|
||||
"write_start_index": write_start_index,
|
||||
"write_end_index": local_end_index,
|
||||
"new_k": k[:, roped_offset:roped_offset + write_len],
|
||||
"new_v": v[:, roped_offset:roped_offset + write_len],
|
||||
"current_end": current_end,
|
||||
"is_recompute": is_recompute
|
||||
}
|
||||
|
||||
# Attention: sink tokens + local window
|
||||
if sink_tokens > 0:
|
||||
local_budget = self.max_attention_size - sink_tokens
|
||||
k_sink = roped_temp_k[:, :sink_tokens]
|
||||
v_sink = temp_v[:, :sink_tokens]
|
||||
if local_budget > 0:
|
||||
local_start_for_window = max(sink_tokens, local_end_index - local_budget)
|
||||
k_local = roped_temp_k[:, local_start_for_window:local_end_index]
|
||||
v_local = temp_v[:, local_start_for_window:local_end_index]
|
||||
k_cat = torch.cat([k_sink, k_local], dim=1)
|
||||
v_cat = torch.cat([v_sink, v_local], dim=1)
|
||||
else:
|
||||
k_cat = k_sink
|
||||
v_cat = v_sink
|
||||
x = attention(roped_query, k_cat, v_cat)
|
||||
else:
|
||||
window_start = max(0, local_end_index - self.max_attention_size)
|
||||
x = attention(
|
||||
roped_query,
|
||||
roped_temp_k[:, window_start:local_end_index],
|
||||
temp_v[:, window_start:local_end_index])
|
||||
|
||||
x = x.flatten(2)
|
||||
x, _ = self.o(x)
|
||||
return x, (current_end, local_end_index, cache_update_info)
|
||||
|
||||
|
||||
class CausalPropeSelfAttention(nn.Module):
|
||||
"""PRoPE self-attention with optional KV cache for camera-controlled inference."""
|
||||
|
||||
def __init__(self, dim, attn_dim, num_heads, window_size=(-1, -1),
|
||||
local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6):
|
||||
assert dim % num_heads == 0
|
||||
assert attn_dim % num_heads == 0
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.attn_dim = attn_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = attn_dim // num_heads
|
||||
self.local_attn_size = local_attn_size
|
||||
self.sink_size = sink_size
|
||||
self.qk_norm = qk_norm
|
||||
self.eps = eps
|
||||
self.window_size = window_size
|
||||
self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
|
||||
|
||||
self.q_proj = ReplicatedLinear(dim, attn_dim)
|
||||
self.k_proj = ReplicatedLinear(dim, attn_dim)
|
||||
self.v_proj = ReplicatedLinear(dim, attn_dim)
|
||||
self.out_proj = ReplicatedLinear(attn_dim, dim)
|
||||
|
||||
self.norm_q = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
nn.init.zeros_(self.out_proj.weight)
|
||||
nn.init.zeros_(self.out_proj.bias)
|
||||
|
||||
def forward(self, x, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
|
||||
kv_cache=None, current_start=0, cache_start=None,
|
||||
sink_recache_after_switch=False, cache_update_policy="commit_detached"):
|
||||
"""
|
||||
Args:
|
||||
x: Shape [B, L, C]
|
||||
cam_viewmats: Camera view matrices
|
||||
cam_K: Camera intrinsics
|
||||
kv_cache: Optional KV cache dict. When None, runs full attention over current chunk.
|
||||
"""
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
if cache_start is None:
|
||||
cache_start = current_start
|
||||
|
||||
q, _ = self.q_proj(x)
|
||||
q = self.norm_q(q).view(b, s, n, d)
|
||||
k, _ = self.k_proj(x)
|
||||
k = self.norm_k(k).view(b, s, n, d)
|
||||
v, _ = self.v_proj(x)
|
||||
v = v.view(b, s, n, d)
|
||||
|
||||
# Apply PRoPE (Positional Rotary Position Embedding from camera parameters)
|
||||
q_t, k_t, v_t, apply_fn_o = prope_qkv(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
|
||||
viewmats=cam_viewmats, Ks=cam_K)
|
||||
proped_q = q_t.transpose(1, 2)
|
||||
proped_k = k_t.transpose(1, 2)
|
||||
proped_v = v_t.transpose(1, 2)
|
||||
|
||||
if kv_cache is None:
|
||||
# No cache: full attention over current chunk
|
||||
x_out = attention(proped_q, proped_k, proped_v)
|
||||
else:
|
||||
# KV cache mode with rolling cache support
|
||||
frame_seqlen = math.prod(grid_sizes[0][1:]).item()
|
||||
num_new_tokens = s
|
||||
current_end = current_start + num_new_tokens
|
||||
sink_tokens = self.sink_size * frame_seqlen
|
||||
kv_cache_size = kv_cache["k"].shape[1]
|
||||
is_recompute = (current_end <= kv_cache["global_end_index"].item()) and (current_start > 0)
|
||||
|
||||
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
|
||||
num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
|
||||
# === ROLLING MODE ===
|
||||
num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
|
||||
num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
|
||||
local_end_index = kv_cache["local_end_index"].item() + current_end - \
|
||||
kv_cache["global_end_index"].item() - num_evicted_tokens
|
||||
local_start_index = local_end_index - num_new_tokens
|
||||
|
||||
if cache_update_policy != "none":
|
||||
with torch.no_grad():
|
||||
kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
|
||||
kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
|
||||
kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
|
||||
kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
|
||||
|
||||
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
|
||||
roped_offset = max(0, write_start_index - local_start_index)
|
||||
write_len = max(0, local_end_index - write_start_index)
|
||||
if write_len > 0:
|
||||
with torch.no_grad():
|
||||
kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
|
||||
kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
|
||||
else:
|
||||
# === DIRECT INSERT MODE ===
|
||||
local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
|
||||
local_start_index = local_end_index - num_new_tokens
|
||||
|
||||
if cache_update_policy != "none":
|
||||
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
|
||||
if sink_recache_after_switch:
|
||||
write_start_index = local_start_index
|
||||
roped_offset = max(0, write_start_index - local_start_index)
|
||||
write_len = max(0, local_end_index - write_start_index)
|
||||
if write_len > 0:
|
||||
with torch.no_grad():
|
||||
kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
|
||||
kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
|
||||
|
||||
# Attention: sink tokens + local window
|
||||
if sink_tokens > 0:
|
||||
local_budget = self.max_attention_size - sink_tokens
|
||||
k_sink = kv_cache["k"][:, :sink_tokens].detach()
|
||||
v_sink = kv_cache["v"][:, :sink_tokens].detach()
|
||||
if local_budget > 0:
|
||||
local_start_for_window = max(sink_tokens, local_end_index - local_budget)
|
||||
k_local = kv_cache["k"][:, local_start_for_window:local_end_index].detach()
|
||||
v_local = kv_cache["v"][:, local_start_for_window:local_end_index].detach()
|
||||
k_cat = torch.cat([k_sink, k_local], dim=1)
|
||||
v_cat = torch.cat([v_sink, v_local], dim=1)
|
||||
else:
|
||||
k_cat = k_sink
|
||||
v_cat = v_sink
|
||||
x_out = attention(proped_q, k_cat, v_cat)
|
||||
else:
|
||||
window_start = max(0, local_end_index - self.max_attention_size)
|
||||
x_out = attention(
|
||||
proped_q,
|
||||
kv_cache["k"][:, window_start:local_end_index].detach(),
|
||||
kv_cache["v"][:, window_start:local_end_index].detach())
|
||||
|
||||
if not is_recompute and cache_update_policy != "none":
|
||||
kv_cache["global_end_index"].fill_(current_end)
|
||||
kv_cache["local_end_index"].fill_(local_end_index)
|
||||
|
||||
# Apply inverse PRoPE
|
||||
x = apply_fn_o(x_out.transpose(1, 2)).transpose(1, 2)
|
||||
x = x.flatten(2)
|
||||
x, _ = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class CausalWanAttentionBlock(nn.Module):
|
||||
|
||||
def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, sink_size=0,
|
||||
qk_norm=True, cross_attn_norm=False, eps=1e-6, **kwargs):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.num_heads = num_heads
|
||||
self.local_attn_size = local_attn_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
self.add_control_adapter = kwargs.get('add_control_adapter', False)
|
||||
self.cam_method = kwargs.get('cam_method')
|
||||
self.attn_compress = kwargs.get('attn_compress', 1)
|
||||
self.layer_idx = kwargs.get('layer_idx')
|
||||
cam_self_attn_layers = kwargs.get('cam_self_attn_layers')
|
||||
|
||||
# layers
|
||||
self.norm1 = WanLayerNorm(dim, eps)
|
||||
self.self_attn = CausalWanSelfAttention(
|
||||
dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
|
||||
self.norm3 = WanLayerNorm(
|
||||
dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
|
||||
self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
|
||||
self.norm2 = WanLayerNorm(dim, eps)
|
||||
# nn.Linear (not ReplicatedLinear) on purpose: the official checkpoint
|
||||
# stores these as positional Sequential keys (ffn.0 / ffn.2) that the
|
||||
# copy-only converter and the strict-load tests require verbatim, and
|
||||
# ReplicatedLinear's (out, bias) tuple return cannot compose inside
|
||||
# nn.Sequential without renaming the state-dict surface.
|
||||
self.ffn = nn.Sequential(
|
||||
nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
|
||||
nn.Linear(ffn_dim, dim))
|
||||
|
||||
# PRoPE self-attention branch for camera control
|
||||
add_cam_attn = self.add_control_adapter and self.cam_method == 'prope'
|
||||
if add_cam_attn and cam_self_attn_layers is not None:
|
||||
add_cam_attn = self.layer_idx in cam_self_attn_layers
|
||||
if add_cam_attn:
|
||||
self.cam_self_attn = CausalPropeSelfAttention(
|
||||
dim, dim // self.attn_compress, num_heads,
|
||||
local_attn_size=local_attn_size, sink_size=sink_size,
|
||||
qk_norm=qk_norm, eps=eps)
|
||||
|
||||
# modulation
|
||||
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
|
||||
|
||||
def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens,
|
||||
kv_cache, crossattn_cache=None, current_start=0, cache_start=None,
|
||||
cam_viewmats=None, cam_K=None, sink_recache_after_switch=False,
|
||||
cache_update_policy="commit_detached"):
|
||||
num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
|
||||
e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
|
||||
|
||||
# self-attention
|
||||
attn_input = (self.norm1(x).unflatten(
|
||||
dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2)
|
||||
y, cache_update_info = self.self_attn(
|
||||
attn_input, seq_lens, grid_sizes, freqs, kv_cache,
|
||||
current_start, cache_start, sink_recache_after_switch)
|
||||
|
||||
# PRoPE camera attention (parallel branch)
|
||||
if hasattr(self, 'cam_self_attn') and cam_viewmats is not None and cam_K is not None:
|
||||
prope_kv_cache = None
|
||||
if kv_cache is not None and "prope_k" in kv_cache:
|
||||
prope_kv_cache = {
|
||||
"k": kv_cache["prope_k"],
|
||||
"v": kv_cache["prope_v"],
|
||||
"global_end_index": kv_cache["prope_global_end_index"],
|
||||
"local_end_index": kv_cache["prope_local_end_index"],
|
||||
}
|
||||
y = y + self.cam_self_attn(
|
||||
attn_input, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
|
||||
kv_cache=prope_kv_cache, current_start=current_start,
|
||||
cache_start=cache_start, cache_update_policy=cache_update_policy)
|
||||
|
||||
x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)
|
||||
|
||||
# cross-attention & FFN
|
||||
x = x + self.cross_attn(self.norm3(x), context, context_lens,
|
||||
crossattn_cache=crossattn_cache)
|
||||
y = self.ffn(
|
||||
(self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
|
||||
* (1 + e[4]) + e[3]).flatten(1, 2))
|
||||
x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2)
|
||||
|
||||
return x, cache_update_info
|
||||
|
||||
|
||||
class CausalHead(nn.Module):
|
||||
|
||||
def __init__(self, dim, out_dim, patch_size, eps=1e-6):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.out_dim = out_dim
|
||||
self.patch_size = patch_size
|
||||
self.eps = eps
|
||||
|
||||
out_dim = math.prod(patch_size) * out_dim
|
||||
self.norm = WanLayerNorm(dim, eps)
|
||||
self.head = ReplicatedLinear(dim, out_dim)
|
||||
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
|
||||
|
||||
def forward(self, x, e):
|
||||
num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
|
||||
e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
|
||||
x, _ = self.head(
|
||||
self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
|
||||
* (1 + e[1]) + e[0])
|
||||
return x
|
||||
|
||||
|
||||
class DreamXWorldARTransformer3DModel(BaseDiT):
|
||||
"""DreamX-World-5B autoregressive causal transformer."""
|
||||
|
||||
_fsdp_shard_conditions = DreamXWorldARConfig()._fsdp_shard_conditions
|
||||
_compile_conditions = DreamXWorldARConfig()._compile_conditions
|
||||
_supported_attention_backends = DreamXWorldARConfig()._supported_attention_backends
|
||||
param_names_mapping = DreamXWorldARConfig().param_names_mapping
|
||||
reverse_param_names_mapping = DreamXWorldARConfig().reverse_param_names_mapping
|
||||
lora_param_names_mapping = DreamXWorldARConfig().lora_param_names_mapping
|
||||
_no_split_modules = ["CausalWanAttentionBlock"]
|
||||
|
||||
def __init__(self, config: DreamXWorldARConfig, hf_config: dict[str, Any]) -> None:
|
||||
super().__init__(config=config, hf_config=hf_config)
|
||||
|
||||
model_type = config.model_type
|
||||
patch_size = config.patch_size
|
||||
text_len = config.text_len
|
||||
in_dim = config.in_channels
|
||||
dim = config.hidden_size
|
||||
ffn_dim = config.ffn_dim
|
||||
freq_dim = config.freq_dim
|
||||
text_dim = config.text_dim
|
||||
out_dim = config.out_channels
|
||||
num_heads = config.num_attention_heads
|
||||
num_layers = config.num_layers
|
||||
local_attn_size = config.local_attn_size
|
||||
sink_size = config.sink_size
|
||||
qk_norm = bool(config.qk_norm)
|
||||
cross_attn_norm = config.cross_attn_norm
|
||||
eps = config.eps
|
||||
add_control_adapter = config.add_control_adapter
|
||||
cam_method = config.cam_method
|
||||
attn_compress = config.attn_compress
|
||||
cam_self_attn_layers = config.cam_self_attn_layers
|
||||
|
||||
assert model_type in ['t2v', 'i2v', 'ti2v']
|
||||
self.model_type = model_type
|
||||
self.patch_size = patch_size
|
||||
self.text_len = text_len
|
||||
self.in_dim = in_dim
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.freq_dim = freq_dim
|
||||
self.text_dim = text_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.local_attn_size = local_attn_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
# embeddings — nn.Linear inside nn.Sequential on purpose: the official
|
||||
# checkpoint keys are positional (text_embedding.0/.2, time_embedding.0/.2,
|
||||
# time_projection.1) and must load verbatim (see ffn comment above).
|
||||
self.patch_embedding = nn.Conv3d(
|
||||
in_dim, dim, kernel_size=patch_size, stride=patch_size)
|
||||
self.text_embedding = nn.Sequential(
|
||||
nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
|
||||
nn.Linear(dim, dim))
|
||||
self.time_embedding = nn.Sequential(
|
||||
nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
|
||||
self.time_projection = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(dim, dim * 6))
|
||||
|
||||
# transformer blocks
|
||||
self.blocks = nn.ModuleList([
|
||||
CausalWanAttentionBlock(
|
||||
dim, ffn_dim, num_heads, local_attn_size, sink_size,
|
||||
qk_norm, cross_attn_norm, eps,
|
||||
add_control_adapter=add_control_adapter,
|
||||
cam_method=cam_method,
|
||||
attn_compress=attn_compress,
|
||||
layer_idx=layer_idx,
|
||||
cam_self_attn_layers=cam_self_attn_layers)
|
||||
for layer_idx in range(num_layers)
|
||||
])
|
||||
for layer_idx, block in enumerate(self.blocks):
|
||||
block.self_attn.layer_idx = layer_idx
|
||||
block.self_attn.num_layers = self.num_layers
|
||||
|
||||
# head
|
||||
self.head = CausalHead(dim, out_dim, patch_size, eps)
|
||||
|
||||
# RoPE frequencies
|
||||
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
|
||||
d = dim // num_heads
|
||||
self.freqs = torch.cat([
|
||||
rope_params(1024, d - 4 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6))
|
||||
], dim=1)
|
||||
|
||||
self.num_attention_heads = num_heads
|
||||
self.attention_head_dim = dim // num_heads
|
||||
self.hidden_size = dim
|
||||
self.in_channels = in_dim
|
||||
self.out_channels = out_dim
|
||||
self.num_channels_latents = out_dim
|
||||
self.init_weights()
|
||||
self.num_frame_per_block = config.arch_config.num_frames_per_block
|
||||
self.__post_init__()
|
||||
|
||||
def forward(self, x=None, t=None, context=None, seq_len=None, y=None, y_camera=None,
|
||||
kv_cache=None, crossattn_cache=None, current_start=0,
|
||||
cache_start=0, cache_update_policy="commit_detached",
|
||||
hidden_states=None, encoder_hidden_states=None, timestep=None, **kwargs):
|
||||
"""
|
||||
Causal inference with KV caching.
|
||||
See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772).
|
||||
|
||||
Args:
|
||||
x: List of input video tensors [C_in, F, H, W]
|
||||
t: Timestep tensor [B, L]
|
||||
context: List of text embeddings [L, C]
|
||||
seq_len: Maximum sequence length for positional encoding
|
||||
y: Optional conditional video inputs (I2V mode)
|
||||
y_camera: Camera parameters dict {'viewmats': ..., 'K': ...}
|
||||
kv_cache: List of KV cache dicts per transformer block
|
||||
crossattn_cache: List of cross-attention cache dicts
|
||||
current_start: Current position in global token sequence
|
||||
cache_start: Cache start position
|
||||
cache_update_policy: Cache update strategy ('commit_detached' or 'none')
|
||||
|
||||
Returns:
|
||||
Stacked output tensors [B, C_out, F, H/8, W/8]
|
||||
"""
|
||||
if x is None and hidden_states is not None:
|
||||
x = [sample for sample in hidden_states]
|
||||
if t is None and timestep is not None:
|
||||
t = timestep
|
||||
if context is None and encoder_hidden_states is not None:
|
||||
if isinstance(encoder_hidden_states, torch.Tensor):
|
||||
context = [sample for sample in encoder_hidden_states]
|
||||
else:
|
||||
context = encoder_hidden_states
|
||||
if seq_len is None:
|
||||
if torch.is_tensor(t):
|
||||
seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel())
|
||||
elif x is not None:
|
||||
sample = x[0]
|
||||
seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (sample.shape[3] // self.patch_size[2])
|
||||
if x is None or t is None or context is None or seq_len is None:
|
||||
raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases")
|
||||
|
||||
device = self.patch_embedding.weight.device
|
||||
if self.freqs.is_meta or self.freqs.device != device:
|
||||
d = self.dim // self.num_heads
|
||||
self.freqs = torch.cat([
|
||||
rope_params(1024, d - 4 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
], dim=1).to(device)
|
||||
|
||||
if y is not None:
|
||||
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]
|
||||
|
||||
# patch embedding
|
||||
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
|
||||
grid_sizes = torch.stack(
|
||||
[torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
|
||||
x = [u.flatten(2).transpose(1, 2) for u in x]
|
||||
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
|
||||
assert seq_lens.max() <= seq_len
|
||||
x = torch.cat(x)
|
||||
|
||||
# time embedding
|
||||
e = self.time_embedding(
|
||||
sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
|
||||
e0 = self.time_projection(e).unflatten(
|
||||
1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
|
||||
|
||||
# text embedding
|
||||
context_lens = None
|
||||
context = self.text_embedding(
|
||||
torch.stack([
|
||||
torch.cat(
|
||||
[u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
|
||||
for u in context
|
||||
]))
|
||||
|
||||
# camera parameters
|
||||
if y_camera is not None and isinstance(y_camera, dict):
|
||||
cam_viewmats = y_camera['viewmats']
|
||||
cam_K = y_camera['K']
|
||||
else:
|
||||
cam_viewmats = None
|
||||
cam_K = None
|
||||
|
||||
block_kwargs = dict(
|
||||
e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs,
|
||||
context=context, context_lens=context_lens,
|
||||
cam_viewmats=cam_viewmats, cam_K=cam_K,
|
||||
cache_update_policy=cache_update_policy,
|
||||
)
|
||||
|
||||
cache_update_infos = []
|
||||
for block_index, block in enumerate(self.blocks):
|
||||
block_kwargs.update({
|
||||
"kv_cache": kv_cache[block_index] if kv_cache is not None else None,
|
||||
"crossattn_cache": crossattn_cache[block_index] if crossattn_cache is not None else None,
|
||||
"current_start": current_start,
|
||||
"cache_start": cache_start,
|
||||
})
|
||||
x, block_cache_update_info = block(x, **block_kwargs)
|
||||
if kv_cache is not None:
|
||||
cache_update_infos.append((block_index, block_cache_update_info))
|
||||
|
||||
# Apply deferred cache updates
|
||||
if kv_cache is not None and cache_update_infos and cache_update_policy != "none":
|
||||
self._apply_cache_updates(kv_cache, cache_update_infos)
|
||||
|
||||
# head & unpatchify
|
||||
x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
|
||||
x = self.unpatchify(x, grid_sizes)
|
||||
return torch.stack(x)
|
||||
|
||||
def _apply_cache_updates(self, kv_cache, cache_update_infos):
|
||||
"""Apply deferred cache updates collected from all transformer blocks.
|
||||
|
||||
For Block-Relativistic RoPE, this stores un-roped K values in the cache.
|
||||
RoPE is applied dynamically during attention based on each token's current
|
||||
relative position in the sliding window.
|
||||
"""
|
||||
with torch.no_grad():
|
||||
for block_index, (current_end, local_end_index, update_info) in cache_update_infos:
|
||||
if update_info is not None:
|
||||
cache = kv_cache[block_index]
|
||||
|
||||
if update_info["action"] == "roll_and_insert":
|
||||
sink_tokens = update_info["sink_tokens"]
|
||||
num_rolled_tokens = update_info["num_rolled_tokens"]
|
||||
num_evicted_tokens = update_info["num_evicted_tokens"]
|
||||
write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
|
||||
write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
|
||||
new_k = update_info["new_k"].detach()
|
||||
new_v = update_info["new_v"].detach()
|
||||
|
||||
cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
|
||||
cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
|
||||
cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
|
||||
cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
|
||||
|
||||
if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
|
||||
cache["k"][:, write_start_index:write_end_index] = new_k
|
||||
cache["v"][:, write_start_index:write_end_index] = new_v
|
||||
|
||||
elif update_info["action"] == "direct_insert":
|
||||
write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
|
||||
write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
|
||||
new_k = update_info["new_k"].detach()
|
||||
new_v = update_info["new_v"].detach()
|
||||
|
||||
if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
|
||||
cache["k"][:, write_start_index:write_end_index] = new_k
|
||||
cache["v"][:, write_start_index:write_end_index] = new_v
|
||||
|
||||
is_recompute = False if update_info is None else update_info.get("is_recompute", False)
|
||||
if not is_recompute:
|
||||
kv_cache[block_index]["global_end_index"].fill_(current_end)
|
||||
kv_cache[block_index]["local_end_index"].fill_(local_end_index)
|
||||
|
||||
def unpatchify(self, x, grid_sizes):
|
||||
"""Reconstruct video tensors from patch embeddings."""
|
||||
c = self.out_dim
|
||||
out = []
|
||||
for u, v in zip(x, grid_sizes.tolist(), strict=True):
|
||||
u = u[:math.prod(v)].view(*v, *self.patch_size, c)
|
||||
u = torch.einsum('fhwpqrc->cfphqwr', u)
|
||||
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
|
||||
out.append(u)
|
||||
return out
|
||||
|
||||
def init_weights(self):
|
||||
"""Initialize model parameters using Xavier initialization."""
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Linear, ReplicatedLinear)):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
|
||||
for m in self.text_embedding.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=.02)
|
||||
for m in self.time_embedding.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=.02)
|
||||
|
||||
nn.init.zeros_(self.head.head.weight)
|
||||
|
||||
|
||||
EntryClass = DreamXWorldARTransformer3DModel
|
||||
@@ -699,8 +699,8 @@ class IndividualTokenRefinerBlock(nn.Module):
|
||||
if mask is None:
|
||||
attn_output = self.attn(q, k, v)
|
||||
else:
|
||||
from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
|
||||
attn_metadata = SDPAMetadataBuilder().build(
|
||||
from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
|
||||
attn_metadata = FlashAttnMetadataBuilder().build(
|
||||
current_timestep=0,
|
||||
attn_mask=mask,
|
||||
)
|
||||
|
||||
@@ -173,11 +173,8 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
|
||||
img_v_prope = img_v_prope.permute(0, 2, 1, 3) # [batch, seqlen, num_heads, head_dim]
|
||||
# end hyworld
|
||||
|
||||
# The metadata only carries the text padding mask; the executing
|
||||
# attention kernel is still whatever the layer's selector picked
|
||||
# (flash-attn when installed), not necessarily torch SDPA.
|
||||
from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
|
||||
attn_metadata = SDPAMetadataBuilder().build(
|
||||
from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
|
||||
attn_metadata = FlashAttnMetadataBuilder().build(
|
||||
current_timestep=0,
|
||||
attn_mask=encoder_attention_mask,
|
||||
)
|
||||
@@ -195,9 +192,14 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
|
||||
)
|
||||
|
||||
# begin hyworld
|
||||
# attention with prope (same text mask, so reuse the metadata)
|
||||
# attention with prope
|
||||
from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
|
||||
attn_metadata_prope = FlashAttnMetadataBuilder().build(
|
||||
current_timestep=0,
|
||||
attn_mask=encoder_attention_mask,
|
||||
)
|
||||
# NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope
|
||||
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
|
||||
with set_forward_context(current_timestep=0, attn_metadata=attn_metadata_prope):
|
||||
img_attn_prope, _ = self.attn(
|
||||
img_q_prope,
|
||||
img_k_prope,
|
||||
|
||||
@@ -10,21 +10,13 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fastvideo.attention import LocalAttention
|
||||
from fastvideo.attention.backends.nabla import CAN_USE_FLEX_ATTN, flex_attention, nablaT_v2
|
||||
from fastvideo.configs.models.dits import Kandinsky5VideoConfig
|
||||
from fastvideo.layers.layernorm import LayerNormScaleShift
|
||||
from fastvideo.layers.linear import ReplicatedLinear
|
||||
from fastvideo.layers.mlp import MLP
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.dits.base import BaseDiT
|
||||
from fastvideo.platforms import AttentionBackendEnum
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
if not CAN_USE_FLEX_ATTN:
|
||||
logger.warning("torch.nn.attention.flex_attention is unavailable in this PyTorch build; "
|
||||
"Kandinsky5 NABLA sparse attention (Pro checkpoints) cannot be used.")
|
||||
|
||||
FRACTAL_PIXEL_SIZE = 8
|
||||
_ARCH_CONFIG_DEFAULTS = Kandinsky5VideoConfig().arch_config
|
||||
|
||||
@@ -271,9 +263,10 @@ class Kandinsky5Modulation(nn.Module):
|
||||
|
||||
|
||||
def _apply_rotary(x: torch.Tensor, rope: torch.Tensor) -> torch.Tensor:
|
||||
orig_dtype = x.dtype
|
||||
x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32)
|
||||
x_out = (rope * x_).sum(dim=-1)
|
||||
return x_out.reshape(*x.shape).to(x.dtype)
|
||||
return x_out.reshape(*x.shape).to(orig_dtype)
|
||||
|
||||
|
||||
class Kandinsky5Attention(nn.Module):
|
||||
@@ -284,7 +277,6 @@ class Kandinsky5Attention(nn.Module):
|
||||
head_dim: int,
|
||||
supported_attention_backends: tuple[AttentionBackendEnum, ...] | None,
|
||||
prefix: str = "",
|
||||
use_nabla: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
assert num_channels % head_dim == 0
|
||||
@@ -314,17 +306,6 @@ class Kandinsky5Attention(nn.Module):
|
||||
causal=False,
|
||||
supported_attention_backends=supported_attention_backends,
|
||||
)
|
||||
# NABLA checkpoints get a second attention layer whose backend defaults
|
||||
# to NABLA_ATTN; FASTVIDEO_ATTENTION_BACKEND still overrides it.
|
||||
self.nabla_attention = None
|
||||
if use_nabla:
|
||||
self.nabla_attention = LocalAttention(
|
||||
num_heads=self.num_heads,
|
||||
head_size=head_dim,
|
||||
causal=False,
|
||||
supported_attention_backends=supported_attention_backends,
|
||||
default_backend=AttentionBackendEnum.NABLA_ATTN,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -358,54 +339,27 @@ class Kandinsky5Attention(nn.Module):
|
||||
key = _apply_rotary(key, rotary_emb).type_as(key)
|
||||
|
||||
if sparse_params is not None:
|
||||
if self.nabla_attention is None:
|
||||
raise RuntimeError("sparse_params passed to an attention layer built without use_nabla; "
|
||||
"this checkpoint/config combination is inconsistent.")
|
||||
try:
|
||||
# Backend impl reads sta_mask/P from the forward-context
|
||||
# attention metadata built by the denoising stage.
|
||||
hidden_states = self.nabla_attention(query, key, value)
|
||||
except AssertionError as exc:
|
||||
# Standalone parity tests call the model without a pipeline
|
||||
# forward context; run the NABLA kernel directly.
|
||||
if "Forward context is not set" not in str(exc):
|
||||
raise
|
||||
attn_mask = nablaT_v2(query, key, sparse_params["sta_mask"], thr=sparse_params["P"])
|
||||
hidden_states = flex_attention(
|
||||
query=query.transpose(1, 2),
|
||||
key=key.transpose(1, 2),
|
||||
value=value.transpose(1, 2),
|
||||
block_mask=attn_mask,
|
||||
).transpose(1, 2)
|
||||
else:
|
||||
try:
|
||||
hidden_states = self.local_attention(query, key, value)
|
||||
raise NotImplementedError(
|
||||
"Sparse attention is not yet supported for Kandinsky5 in FastVideo."
|
||||
)
|
||||
|
||||
except AssertionError as exc:
|
||||
# LocalAttention requires pipeline forward context. Standalone
|
||||
# parity tests call the model directly, so fallback to Torch SDPA.
|
||||
if "Forward context is not set" not in str(exc):
|
||||
raise
|
||||
|
||||
query_shape = query.shape[:-2]
|
||||
key_shape = key.shape[:-2]
|
||||
query = query.reshape(query_shape[0], -1, self.num_heads,
|
||||
query.shape[-1]).transpose(1, 2)
|
||||
key = key.reshape(key_shape[0], -1, self.num_heads,
|
||||
key.shape[-1]).transpose(1, 2)
|
||||
value = value.reshape(key_shape[0], -1, self.num_heads,
|
||||
value.shape[-1]).transpose(1, 2)
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=None,
|
||||
is_causal=False,
|
||||
)
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(
|
||||
*query_shape, self.num_heads, -1)
|
||||
|
||||
hidden_states = hidden_states.flatten(-2, -1)
|
||||
try:
|
||||
hidden_states = self.local_attention(query, key, value)
|
||||
except AssertionError as exc:
|
||||
# LocalAttention requires pipeline forward context. Standalone
|
||||
# parity tests call the model directly, so fallback to Torch SDPA.
|
||||
if "Forward context is not set" not in str(exc):
|
||||
raise
|
||||
query = query.transpose(1, 2)
|
||||
key = key.transpose(1, 2)
|
||||
value = value.transpose(1, 2)
|
||||
hidden_states = F.scaled_dot_product_attention(query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=None,
|
||||
is_causal=False)
|
||||
hidden_states = hidden_states.transpose(1, 2)
|
||||
hidden_states = hidden_states.flatten(2)
|
||||
|
||||
hidden_states, _ = self.out_layer(hidden_states)
|
||||
return hidden_states
|
||||
@@ -522,8 +476,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
|
||||
head_dim: int,
|
||||
supported_attention_backends: tuple[AttentionBackendEnum, ...]
|
||||
| None = None,
|
||||
prefix: str = "",
|
||||
use_nabla: bool = False):
|
||||
prefix: str = ""):
|
||||
super().__init__()
|
||||
self.visual_modulation = Kandinsky5Modulation(time_dim, model_dim, 9)
|
||||
|
||||
@@ -538,8 +491,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
|
||||
model_dim,
|
||||
head_dim,
|
||||
supported_attention_backends=supported_attention_backends,
|
||||
prefix=f"{prefix}.self_attention",
|
||||
use_nabla=use_nabla)
|
||||
prefix=f"{prefix}.self_attention")
|
||||
|
||||
self.cross_attention_norm = LayerNormScaleShift(
|
||||
model_dim,
|
||||
@@ -672,8 +624,7 @@ class Kandinsky5Transformer3DModel(BaseDiT):
|
||||
arch.ff_dim,
|
||||
head_dim,
|
||||
self._supported_attention_backends,
|
||||
prefix=f"{config.prefix}.visual_transformer_blocks.{i}",
|
||||
use_nabla=arch.attention_type == "nabla")
|
||||
prefix=f"{config.prefix}.visual_transformer_blocks.{i}")
|
||||
for i in range(arch.num_visual_blocks)
|
||||
])
|
||||
|
||||
@@ -743,7 +694,6 @@ class Kandinsky5Transformer3DModel(BaseDiT):
|
||||
scale_factor)
|
||||
to_fractal = sparse_params[
|
||||
"to_fractal"] if sparse_params is not None else False
|
||||
|
||||
visual_embed, visual_rope = fractal_flatten(visual_embed, visual_rope,
|
||||
visual_shape,
|
||||
block_mask=to_fractal)
|
||||
@@ -774,7 +724,6 @@ class Kandinsky5Transformer3DModel(BaseDiT):
|
||||
|
||||
if return_dict:
|
||||
return Kandinsky5TransformerOutput(sample=x)
|
||||
|
||||
return x
|
||||
|
||||
def materialize_non_persistent_buffers(self, device: torch.device,
|
||||
|
||||
@@ -467,18 +467,6 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module):
|
||||
return hidden_states
|
||||
|
||||
|
||||
def _compute_default_rope_parameters(config, device=None, seq_len=None, **kwargs):
|
||||
# transformers>=5 removes the "default" entry from ROPE_INIT_FUNCTIONS and
|
||||
# moves rope_theta inside rope_parameters; replicate the 4.x default init.
|
||||
rope_params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
|
||||
base = rope_params.get("rope_theta", getattr(config, "rope_theta", 10000.0))
|
||||
head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
||||
partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
|
||||
dim = int(head_dim * partial_rotary_factor)
|
||||
inv_freq = 1.0 / (base**(torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
|
||||
return inv_freq, 1.0
|
||||
|
||||
|
||||
class Qwen2_5_VLRotaryEmbedding(nn.Module):
|
||||
def __init__(self, config: Qwen2_5_VLConfig, device=None):
|
||||
super().__init__()
|
||||
@@ -491,14 +479,7 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module):
|
||||
self.original_max_seq_len = config.max_position_embeddings
|
||||
|
||||
self.config = config
|
||||
if self.rope_type in ROPE_INIT_FUNCTIONS:
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
elif self.rope_type == "default":
|
||||
# transformers>=5 drops the "default" entry from ROPE_INIT_FUNCTIONS.
|
||||
self.rope_init_fn = _compute_default_rope_parameters
|
||||
else:
|
||||
raise KeyError(f"Unsupported rope_type '{self.rope_type}'; available: "
|
||||
f"{['default', *ROPE_INIT_FUNCTIONS]}")
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
@@ -967,16 +948,6 @@ QWEN2_5_VL_ATTENTION_CLASSES = {
|
||||
# If FlashAttention2 is not available, transparently fall back to SDPA.
|
||||
if not is_flash_attn_2_available():
|
||||
QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention
|
||||
else:
|
||||
# transformers>=5 only resolves the flash-attn functions when the model
|
||||
# preloads them via its attention interface; this module bypasses
|
||||
# PreTrainedModel, so _flash_attention_forward(implementation=None) raises
|
||||
# unless we preload here.
|
||||
try:
|
||||
from transformers.modeling_flash_attention_utils import lazy_import_flash_attention
|
||||
lazy_import_flash_attention("flash_attention_2")
|
||||
except (ImportError, ValueError):
|
||||
QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention
|
||||
|
||||
class Qwen2_5_VLDecoderLayer(nn.Module):
|
||||
def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int):
|
||||
@@ -1064,7 +1035,7 @@ class Qwen2_5_VLModel(nn.Module):
|
||||
def __init__(self, config: Qwen2_5_VLConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.padding_idx = getattr(config, "pad_token_id", None)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||||
@@ -1437,18 +1408,6 @@ class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput):
|
||||
rope_deltas: Optional[torch.LongTensor] = None
|
||||
|
||||
|
||||
def _flatten_text_config(config):
|
||||
# transformers>=5 stops forwarding text-model attributes (hidden_size,
|
||||
# vocab_size, rope_scaling, ...) from the composite Qwen2_5_VLConfig to
|
||||
# config.text_config; this module reads them from the top level.
|
||||
text_config = getattr(config, "text_config", None)
|
||||
if text_config is not None:
|
||||
for key, value in text_config.to_dict().items():
|
||||
if not hasattr(config, key):
|
||||
setattr(config, key, value)
|
||||
return config
|
||||
|
||||
|
||||
class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
|
||||
_tied_weights_keys = ["lm_head.weight"]
|
||||
config_class = Qwen2_5_VLConfig
|
||||
@@ -1456,7 +1415,6 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
config = _flatten_text_config(config)
|
||||
self.config = config
|
||||
self.visual = Qwen2_5_VisionTransformerPretrainedModel(config.vision_config)
|
||||
|
||||
|
||||
@@ -320,19 +320,6 @@ class TextEncoderLoader(ComponentLoader):
|
||||
f"text encoder index {idx} out of range for text_encoder_configs (len={len(encoder_configs)}), model_path={model_path}"
|
||||
)
|
||||
encoder_config = encoder_configs[idx]
|
||||
if (
|
||||
model_config.get("architectures") == ["CLIPModel"]
|
||||
and isinstance(model_config.get("text_config"), dict)
|
||||
):
|
||||
valid_arch_fields = {
|
||||
f.name for f in dataclasses.fields(encoder_config.arch_config)
|
||||
}
|
||||
model_config = {
|
||||
key: value
|
||||
for key, value in deepcopy(model_config["text_config"]).items()
|
||||
if key in valid_arch_fields
|
||||
}
|
||||
model_config["architectures"] = ["CLIPTextModel"]
|
||||
encoder_config.update_model_arch(model_config)
|
||||
if idx < 0 or idx >= len(encoder_precisions):
|
||||
raise IndexError(
|
||||
@@ -417,9 +404,7 @@ class TextEncoderLoader(ComponentLoader):
|
||||
else:
|
||||
loaded_weights: set[str] = model.load_weights(
|
||||
self._get_all_weights(
|
||||
model,
|
||||
model_path,
|
||||
to_cpu=fastvideo_args.text_encoder_cpu_offload,
|
||||
model, model_path, to_cpu=use_cpu_offload
|
||||
)
|
||||
) # type: ignore
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ import tempfile
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Callable, Set
|
||||
from dataclasses import dataclass, field
|
||||
from functools import cache, lru_cache
|
||||
from functools import lru_cache
|
||||
from typing import NoReturn, TypeVar, cast
|
||||
|
||||
import cloudpickle
|
||||
@@ -32,8 +32,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
|
||||
"HYWorldTransformer3DModel":
|
||||
("dits", "hyworld", "HYWorldTransformer3DModel"),
|
||||
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
|
||||
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
|
||||
"DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
|
||||
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
|
||||
"CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
|
||||
"Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
|
||||
@@ -50,8 +48,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
|
||||
_IMAGE_TO_VIDEO_DIT_MODELS = {
|
||||
# "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
|
||||
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
|
||||
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
|
||||
"DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
|
||||
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
|
||||
"MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
|
||||
"CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
|
||||
@@ -145,7 +141,7 @@ _LEGACY_FAST_VIDEO_MODELS = {
|
||||
MODELS_PATH = os.path.dirname(__file__)
|
||||
|
||||
|
||||
@cache
|
||||
@lru_cache(maxsize=None)
|
||||
def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
|
||||
discovered_models: dict[str, tuple[str, str, str]] = {}
|
||||
for root, dirs, files in os.walk(MODELS_PATH):
|
||||
@@ -160,7 +156,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
|
||||
|
||||
filepath = os.path.join(root, filename)
|
||||
try:
|
||||
with open(filepath, encoding="utf-8") as f:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
source = f.read()
|
||||
tree = ast.parse(source, filename=filename)
|
||||
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
from urllib.parse import unquote, urlparse
|
||||
@@ -18,10 +16,6 @@ import torch.nn.functional as F
|
||||
import torchvision.transforms.functional as TF
|
||||
from packaging import version
|
||||
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
if version.parse(version.parse(
|
||||
PIL.__version__).base_version) >= version.parse("9.1.0"):
|
||||
PIL_INTERPOLATION = {
|
||||
@@ -95,39 +89,6 @@ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
|
||||
return 2.0 * images - 1.0
|
||||
|
||||
|
||||
def _fetch_image_bytes(url: str,
|
||||
attempts: int = 3,
|
||||
timeout: float = 30.0) -> bytes:
|
||||
"""Download ``url`` fully, retrying transient network errors.
|
||||
|
||||
Reads the whole body via ``response.content`` instead of handing PIL a
|
||||
``.raw`` stream: a connection dropped mid-body (e.g. an HF-CDN
|
||||
``IncompleteRead``) then fails here, where it can be retried, instead of
|
||||
surfacing as a truncated image inside PIL.
|
||||
"""
|
||||
for attempt in range(attempts):
|
||||
try:
|
||||
response = requests.get(url, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
except requests.RequestException as e:
|
||||
status = getattr(getattr(e, "response", None), "status_code",
|
||||
None)
|
||||
# Permanent client errors (4xx except 408/429) won't heal on
|
||||
# retry — fail fast.
|
||||
if (status is not None and 400 <= status < 500
|
||||
and status not in (408, 429)):
|
||||
raise
|
||||
if attempt == attempts - 1:
|
||||
raise
|
||||
backoff = 2**attempt
|
||||
logger.warning(
|
||||
"Failed to download image %s (attempt %d/%d): %s. "
|
||||
"Retrying in %ds.", url, attempt + 1, attempts, e, backoff)
|
||||
time.sleep(backoff)
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
# adapted from diffusers.utils import load_image
|
||||
def load_image(
|
||||
image: str | PIL.Image.Image,
|
||||
@@ -149,7 +110,7 @@ def load_image(
|
||||
"""
|
||||
if isinstance(image, str):
|
||||
if image.startswith("http://") or image.startswith("https://"):
|
||||
image = PIL.Image.open(io.BytesIO(_fetch_image_bytes(image)))
|
||||
image = PIL.Image.open(requests.get(image, stream=True).raw)
|
||||
elif os.path.isfile(image):
|
||||
image = PIL.Image.open(image)
|
||||
else:
|
||||
|
||||
@@ -337,12 +337,8 @@ def load_records_for_model(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_NUMERIC_COLS = (
|
||||
"latency",
|
||||
"throughput",
|
||||
"memory",
|
||||
"text_encoder_time_s",
|
||||
"dit_time_s",
|
||||
"vae_decode_time_s",
|
||||
"latency", "throughput", "memory",
|
||||
"text_encoder_time_s", "dit_time_s", "vae_decode_time_s",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -36,7 +36,6 @@ DEFAULT_METRIC_POLICIES: tuple[MetricPolicy, ...] = (
|
||||
MetricPolicy("vae_decode_time_s", "VAE Decode", 3, True, 0.05, 0.25),
|
||||
)
|
||||
|
||||
|
||||
def _optional_float(value: Any) -> float | None:
|
||||
if value is None or isinstance(value, bool):
|
||||
return None
|
||||
@@ -58,7 +57,9 @@ def _optional_bool(value: Any) -> bool | None:
|
||||
return None
|
||||
|
||||
|
||||
def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) -> tuple[MetricPolicy, ...]:
|
||||
def resolve_metric_policies(
|
||||
threshold_overrides: Mapping[str, Any] | None,
|
||||
) -> tuple[MetricPolicy, ...]:
|
||||
"""Return default metric policies with optional per-metric overrides."""
|
||||
|
||||
if not isinstance(threshold_overrides, Mapping):
|
||||
@@ -79,15 +80,25 @@ def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) ->
|
||||
label=base_policy.label,
|
||||
precision=base_policy.precision,
|
||||
lower_is_better=base_policy.lower_is_better,
|
||||
threshold_percent=(base_policy.threshold_percent if threshold_percent is None else threshold_percent),
|
||||
threshold_absolute=(base_policy.threshold_absolute
|
||||
if threshold_absolute is None else threshold_absolute),
|
||||
threshold_percent=(
|
||||
base_policy.threshold_percent
|
||||
if threshold_percent is None
|
||||
else threshold_percent
|
||||
),
|
||||
threshold_absolute=(
|
||||
base_policy.threshold_absolute
|
||||
if threshold_absolute is None
|
||||
else threshold_absolute
|
||||
),
|
||||
gated=base_policy.gated if gated is None else gated,
|
||||
))
|
||||
)
|
||||
)
|
||||
return tuple(policies)
|
||||
|
||||
|
||||
def serialize_metric_thresholds(policies: tuple[MetricPolicy, ...], ) -> dict[str, dict[str, float | bool]]:
|
||||
def serialize_metric_thresholds(
|
||||
policies: tuple[MetricPolicy, ...],
|
||||
) -> dict[str, dict[str, float | bool]]:
|
||||
return {
|
||||
policy.key: {
|
||||
"threshold_percent": policy.threshold_percent,
|
||||
@@ -107,7 +118,10 @@ def regression_delta(
|
||||
return None
|
||||
absolute_delta = current - baseline if policy.lower_is_better else baseline - current
|
||||
percent_delta = absolute_delta / baseline
|
||||
threshold_exceeded = (percent_delta > policy.threshold_percent and absolute_delta > policy.threshold_absolute)
|
||||
threshold_exceeded = (
|
||||
percent_delta > policy.threshold_percent
|
||||
and absolute_delta > policy.threshold_absolute
|
||||
)
|
||||
return MetricDelta(
|
||||
absolute=absolute_delta,
|
||||
percent=percent_delta,
|
||||
|
||||
@@ -86,32 +86,20 @@ def record_metadata(record: Record) -> Record:
|
||||
}
|
||||
|
||||
|
||||
def _cohort_metadata_value(value: Any) -> Any:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, str) and not value.strip():
|
||||
return ""
|
||||
return value
|
||||
|
||||
|
||||
def _cohort_key_value(value: Any) -> str:
|
||||
return str(_cohort_metadata_value(value))
|
||||
|
||||
|
||||
def record_comparison_metadata(record: Record) -> Record:
|
||||
return {key: _cohort_metadata_value(record.get(key)) for key in COMPARISON_COHORT_KEYS}
|
||||
return {key: record.get(key) or "" for key in COMPARISON_COHORT_KEYS}
|
||||
|
||||
|
||||
def comparison_cohort_key(record: Record) -> CohortKey:
|
||||
return (
|
||||
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")),
|
||||
str(record.get("workload_id") or ""),
|
||||
str(record.get("variant_id") or ""),
|
||||
str(record.get("benchmark_version") or ""),
|
||||
str(record.get("recipe_fingerprint") or ""),
|
||||
str(record.get("hardware_profile_id") or ""),
|
||||
str(record.get("software_profile_id") or ""),
|
||||
)
|
||||
|
||||
|
||||
@@ -158,10 +146,9 @@ def build_latest_summary(records: list[Record],
|
||||
continue
|
||||
|
||||
latest = latest_candidates[-1]
|
||||
latest_index = next(index for index, record in enumerate(group) if record is latest)
|
||||
baseline_pool = [
|
||||
record for record in group[:latest_index]
|
||||
if record.get("success", True) and is_baseline_eligible_record(record)
|
||||
record for record in group
|
||||
if record is not latest and record.get("success", True) and is_baseline_eligible_record(record)
|
||||
]
|
||||
baseline_records = baseline_pool[-baseline_window:]
|
||||
metric_policies = resolve_metric_policies(latest.get("regression_thresholds"))
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from fastvideo.configs.pipelines.dreamx_world import (
|
||||
DreamXWorld5BARPipelineConfig,
|
||||
DreamXWorld5BCamPipelineConfig,
|
||||
make_dreamx_world_5b_ar_dit_config,
|
||||
make_dreamx_world_5b_cam_dit_config,
|
||||
make_dreamx_world_5b_cam_text_encoder_config,
|
||||
make_dreamx_world_5b_cam_vae_config,
|
||||
)
|
||||
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
|
||||
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import (
|
||||
DREAMX_Y_CAMERA_KEY,
|
||||
DreamXWorldCameraConditioningStage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DREAMX_Y_CAMERA_KEY",
|
||||
"DreamXWorld5BARPipelineConfig",
|
||||
"DreamXWorld5BCamPipelineConfig",
|
||||
"DreamXWorldCameraConditioningStage",
|
||||
"DreamXWorldARPipeline",
|
||||
"DreamXWorldPipeline",
|
||||
"make_dreamx_world_5b_ar_dit_config",
|
||||
"make_dreamx_world_5b_cam_dit_config",
|
||||
"make_dreamx_world_5b_cam_text_encoder_config",
|
||||
"make_dreamx_world_5b_cam_vae_config",
|
||||
]
|
||||
@@ -1,219 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World autoregressive causal denoising stage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import DREAMX_Y_CAMERA_KEY
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.denoising import DenoisingStage
|
||||
|
||||
|
||||
class DreamXWorldARCausalDenoisingStage(DenoisingStage):
|
||||
"""Official DreamX AR-forcing denoising loop with KV cache."""
|
||||
|
||||
_AR_NOISE_SEED_OFFSET = 1_000_003
|
||||
|
||||
def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
|
||||
super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
|
||||
self.num_transformer_blocks = len(self.transformer.blocks)
|
||||
self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
|
||||
self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
assert batch.latents is not None, "latents must be prepared before DreamX AR denoising"
|
||||
assert batch.prompt_embeds, "prompt embeds must be prepared before DreamX AR denoising"
|
||||
latents = batch.latents
|
||||
device = latents.device
|
||||
target_dtype = torch.bfloat16
|
||||
autocast_enabled = device.type == "cuda" and not fastvideo_args.disable_autocast
|
||||
|
||||
frame_seq_length = (latents.shape[-2] // self.transformer.patch_size[1]) * (latents.shape[-1] //
|
||||
self.transformer.patch_size[2])
|
||||
timesteps = torch.tensor(
|
||||
tuple(getattr(fastvideo_args.pipeline_config, "dmd_denoising_steps", (1000, 750, 500, 250))),
|
||||
dtype=torch.long,
|
||||
).cpu()
|
||||
if getattr(fastvideo_args.pipeline_config, "warp_denoising_step", True):
|
||||
self.scheduler.set_timesteps(1000)
|
||||
scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
|
||||
timesteps = scheduler_timesteps[1000 - timesteps]
|
||||
timesteps = timesteps.to(device)
|
||||
|
||||
if latents.shape[2] % self.num_frame_per_block != 0:
|
||||
raise ValueError("DreamX AR latent frames must be divisible by num_frame_per_block")
|
||||
|
||||
y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY, batch.extra.get("y_camera"))
|
||||
if isinstance(y_camera, dict):
|
||||
y_camera = {
|
||||
k: v.to(device=device, dtype=target_dtype) if torch.is_tensor(v) else v
|
||||
for k, v in y_camera.items()
|
||||
}
|
||||
|
||||
if batch.image_latent is not None and batch.image_latent.shape[1] == latents.shape[1]:
|
||||
latents[:, :, :batch.image_latent.shape[2]] = batch.image_latent.to(device=device, dtype=latents.dtype)
|
||||
|
||||
kv_cache = self._initialize_kv_cache(latents.shape[0], target_dtype, device, frame_seq_length)
|
||||
crossattn_cache = self._initialize_crossattn_cache(latents.shape[0], target_dtype, device)
|
||||
prompt = batch.prompt_embeds[0]
|
||||
if torch.is_tensor(prompt):
|
||||
prompt = prompt.to(device=device, dtype=target_dtype)
|
||||
context = [sample for sample in prompt]
|
||||
else:
|
||||
context = prompt
|
||||
|
||||
num_blocks = latents.shape[2] // self.num_frame_per_block
|
||||
start = 0
|
||||
first_frame_mask = torch.ones_like(latents)
|
||||
first_frame_mask[:, :, 0] = 0
|
||||
base_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
|
||||
noise_generator = self._make_noise_generator(base_generator, device)
|
||||
|
||||
with tqdm(total=num_blocks * len(timesteps), desc="DreamX AR denoising", leave=False) as progress:
|
||||
for _ in range(num_blocks):
|
||||
current_num_frames = self.num_frame_per_block
|
||||
block_latents = latents[:, :, start:start + current_num_frames]
|
||||
noisy_input = block_latents.clone()
|
||||
mask_block = first_frame_mask[:, :, start:start + current_num_frames]
|
||||
camera_block = self._slice_camera(y_camera, start, current_num_frames)
|
||||
|
||||
for idx, current_timestep in enumerate(timesteps):
|
||||
timestep = torch.full(
|
||||
(latents.shape[0], current_num_frames * frame_seq_length),
|
||||
int(current_timestep.item()),
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
if start == 0:
|
||||
timestep[:, :frame_seq_length] = 0
|
||||
with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
|
||||
denoised = self.transformer(
|
||||
hidden_states=block_latents.to(target_dtype),
|
||||
encoder_hidden_states=torch.stack(context).to(target_dtype),
|
||||
timestep=timestep,
|
||||
y_camera=camera_block,
|
||||
kv_cache=kv_cache,
|
||||
crossattn_cache=crossattn_cache,
|
||||
current_start=start * frame_seq_length,
|
||||
)
|
||||
denoised = denoised.to(latents.dtype)
|
||||
if idx < len(timesteps) - 1:
|
||||
next_timestep = torch.full((latents.shape[0], current_num_frames),
|
||||
int(timesteps[idx + 1].item()),
|
||||
device=device,
|
||||
dtype=torch.long)
|
||||
noise_kwargs = {"device": device, "dtype": denoised.dtype}
|
||||
if noise_generator is not None:
|
||||
noise_kwargs["generator"] = noise_generator
|
||||
noise = torch.randn(denoised.permute(0, 2, 1, 3, 4).shape, **noise_kwargs)
|
||||
block_btchw = self.scheduler.add_noise(
|
||||
denoised.permute(0, 2, 1, 3, 4).flatten(0, 1),
|
||||
noise.flatten(0, 1),
|
||||
next_timestep.flatten(),
|
||||
).unflatten(0, (latents.shape[0], current_num_frames))
|
||||
block_latents = block_btchw.permute(0, 2, 1, 3, 4)
|
||||
block_latents = block_latents * mask_block + noisy_input * (1 - mask_block)
|
||||
else:
|
||||
block_latents = denoised * mask_block + noisy_input * (1 - mask_block)
|
||||
progress.update()
|
||||
|
||||
latents[:, :, start:start + current_num_frames] = block_latents
|
||||
self._update_context_cache(block_latents, context, camera_block, kv_cache, crossattn_cache, start,
|
||||
frame_seq_length, target_dtype, autocast_enabled,
|
||||
float(getattr(fastvideo_args.pipeline_config, "context_noise", 0.1)))
|
||||
start += current_num_frames
|
||||
|
||||
batch.latents = latents
|
||||
return batch
|
||||
|
||||
def _make_noise_generator(self, generator: torch.Generator | None, device: torch.device) -> torch.Generator | None:
|
||||
if generator is None:
|
||||
return None
|
||||
if getattr(generator, "device", None) == device:
|
||||
return generator
|
||||
seed = int(generator.initial_seed()) + self._AR_NOISE_SEED_OFFSET
|
||||
return torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
@staticmethod
|
||||
def _context_noise_timestep(context_noise: float) -> int:
|
||||
if 0.0 < context_noise <= 1.0:
|
||||
return int(context_noise * 1000)
|
||||
return int(context_noise)
|
||||
|
||||
def _slice_camera(self, y_camera: Any, start: int, num_frames: int):
|
||||
if not isinstance(y_camera, dict):
|
||||
return y_camera
|
||||
return {
|
||||
"viewmats": y_camera["viewmats"][:, start:start + num_frames],
|
||||
"K": y_camera["K"][:, start:start + num_frames],
|
||||
}
|
||||
|
||||
def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device,
|
||||
frame_seq_length: int) -> list[dict[str, Any]]:
|
||||
size = self.local_attn_size * frame_seq_length if self.local_attn_size != -1 else 18480
|
||||
heads = self.transformer.num_attention_heads
|
||||
head_dim = self.transformer.attention_head_dim
|
||||
cam_self_attn = next(
|
||||
(getattr(block, "cam_self_attn", None)
|
||||
for block in self.transformer.blocks if getattr(block, "cam_self_attn", None) is not None),
|
||||
None,
|
||||
)
|
||||
|
||||
caches = []
|
||||
for _ in range(self.num_transformer_blocks):
|
||||
cache = {
|
||||
"k": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
|
||||
"v": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
|
||||
"global_end_index": torch.tensor([0], dtype=torch.long, device=device),
|
||||
"local_end_index": torch.tensor([0], dtype=torch.long, device=device),
|
||||
}
|
||||
if cam_self_attn is not None:
|
||||
cam_heads = int(cam_self_attn.num_heads)
|
||||
cam_head_dim = int(cam_self_attn.head_dim)
|
||||
cache.update({
|
||||
"prope_k":
|
||||
torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
|
||||
"prope_v":
|
||||
torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
|
||||
"prope_global_end_index":
|
||||
torch.tensor([0], dtype=torch.long, device=device),
|
||||
"prope_local_end_index":
|
||||
torch.tensor([0], dtype=torch.long, device=device),
|
||||
})
|
||||
caches.append(cache)
|
||||
return caches
|
||||
|
||||
def _initialize_crossattn_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
|
||||
heads = self.transformer.num_attention_heads
|
||||
head_dim = self.transformer.attention_head_dim
|
||||
return [{
|
||||
"k": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
|
||||
"v": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
|
||||
"is_init": False,
|
||||
} for _ in range(self.num_transformer_blocks)]
|
||||
|
||||
def _update_context_cache(self, block_latents: torch.Tensor, context: Any, camera_block: Any,
|
||||
kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], start: int,
|
||||
frame_seq_length: int, target_dtype: torch.dtype, autocast_enabled: bool,
|
||||
context_noise: float) -> None:
|
||||
timestep = torch.full(
|
||||
(block_latents.shape[0], block_latents.shape[2] * frame_seq_length),
|
||||
self._context_noise_timestep(context_noise),
|
||||
device=block_latents.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
|
||||
self.transformer(
|
||||
hidden_states=block_latents.to(target_dtype),
|
||||
encoder_hidden_states=torch.stack(context).to(target_dtype),
|
||||
timestep=timestep,
|
||||
y_camera=camera_block,
|
||||
kv_cache=kv_cache,
|
||||
crossattn_cache=crossattn_cache,
|
||||
current_start=start * frame_seq_length,
|
||||
)
|
||||
@@ -1,228 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.interpolate import interp1d
|
||||
from scipy.spatial.transform import Rotation, Slerp
|
||||
|
||||
_ACTION_TO_MOTION = {
|
||||
"w": "forward",
|
||||
"a": "left",
|
||||
"d": "right",
|
||||
"s": "backward",
|
||||
"j": "left_rot",
|
||||
"l": "right_rot",
|
||||
"i": "up_rot",
|
||||
"k": "down_rot",
|
||||
}
|
||||
_TRANSLATION_BASE_UNIT = 1.0
|
||||
_ROTATION_BASE_UNIT = 10.0
|
||||
_INTRINSIC_ROW = [0.8, 0.5, 0.5, 0.5]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXCamera:
|
||||
fx: float
|
||||
fy: float
|
||||
cx: float
|
||||
cy: float
|
||||
w2c_mat: np.ndarray
|
||||
|
||||
@property
|
||||
def c2w_mat(self) -> np.ndarray:
|
||||
return np.linalg.inv(self.w2c_mat)
|
||||
|
||||
@classmethod
|
||||
def from_pose_row(cls, row: list[float]) -> DreamXCamera:
|
||||
w2c_mat = np.eye(4, dtype=np.float64)
|
||||
w2c_mat[:3, :] = np.asarray(row[7:], dtype=np.float64).reshape(3, 4)
|
||||
return cls(
|
||||
fx=float(row[1]),
|
||||
fy=float(row[2]),
|
||||
cx=float(row[3]),
|
||||
cy=float(row[4]),
|
||||
w2c_mat=w2c_mat,
|
||||
)
|
||||
|
||||
|
||||
def _translation_step(motion_type: str, current_pose: dict[str, np.ndarray], value: float, duration: int) -> np.ndarray:
|
||||
if motion_type in ("forward", "backward"):
|
||||
yaw = np.radians(current_pose["rotation"][1])
|
||||
pitch = np.radians(current_pose["rotation"][0])
|
||||
forward = np.array([-math.sin(yaw) * math.cos(pitch), math.sin(pitch), math.cos(yaw) * math.cos(pitch)])
|
||||
direction = 1 if motion_type == "forward" else -1
|
||||
return forward * value * direction / duration
|
||||
if motion_type in ("left", "right"):
|
||||
yaw = np.radians(current_pose["rotation"][1])
|
||||
right = np.array([math.cos(yaw), 0.0, math.sin(yaw)])
|
||||
direction = -1 if motion_type == "left" else 1
|
||||
return right * value * direction / duration
|
||||
return np.zeros(3)
|
||||
|
||||
|
||||
def _rotation_step(motion_type: str, value: float, duration: int) -> np.ndarray:
|
||||
if not motion_type.endswith("rot"):
|
||||
return np.zeros(3)
|
||||
axis = motion_type.split("_")[0]
|
||||
rotation = np.zeros(3)
|
||||
if axis == "left":
|
||||
rotation[1] = value
|
||||
elif axis == "right":
|
||||
rotation[1] = -value
|
||||
elif axis == "up":
|
||||
rotation[0] = -value
|
||||
elif axis == "down":
|
||||
rotation[0] = value
|
||||
return rotation / duration
|
||||
|
||||
|
||||
def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
|
||||
pitch, yaw, roll = np.radians(angles)
|
||||
cy = math.cos(yaw * 0.5)
|
||||
sy = math.sin(yaw * 0.5)
|
||||
cp = math.cos(pitch * 0.5)
|
||||
sp = math.sin(pitch * 0.5)
|
||||
cr = math.cos(roll * 0.5)
|
||||
sr = math.sin(roll * 0.5)
|
||||
return [
|
||||
cy * cp * cr + sy * sp * sr,
|
||||
cy * sp * cr + sy * cp * sr,
|
||||
sy * cp * cr - cy * sp * sr,
|
||||
cy * cp * sr - sy * sp * cr,
|
||||
]
|
||||
|
||||
|
||||
def _quaternion_to_rotation_matrix(quaternion: list[float]) -> np.ndarray:
|
||||
qw, qx, qy, qz = quaternion
|
||||
return np.array([
|
||||
[1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
|
||||
[2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
|
||||
[2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
|
||||
])
|
||||
|
||||
|
||||
def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float], duration: int) -> list[list[float]]:
|
||||
if len(action_seq) != len(action_speed_list):
|
||||
raise ValueError("action_seq and action_speed_list must have the same length")
|
||||
|
||||
positions: list[np.ndarray] = []
|
||||
rotations: list[np.ndarray] = []
|
||||
current_pose = {
|
||||
"position": np.array([0.0, 0.0, 0.0]),
|
||||
"rotation": np.array([0.0, 0.0, 0.0]),
|
||||
}
|
||||
|
||||
for action_id, speed in zip(action_seq, action_speed_list, strict=True):
|
||||
motion_types = [_ACTION_TO_MOTION[key] for key in list(action_id)]
|
||||
translation_step = np.zeros(3)
|
||||
rotation_step = np.zeros(3)
|
||||
for motion_type in motion_types:
|
||||
translation_step += _translation_step(motion_type, current_pose,
|
||||
float(speed) * _TRANSLATION_BASE_UNIT, duration)
|
||||
rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
|
||||
|
||||
segment_positions = []
|
||||
segment_rotations = []
|
||||
for index in range(1, duration + 1):
|
||||
segment_positions.append(current_pose["position"] + translation_step * index)
|
||||
segment_rotations.append(current_pose["rotation"] + rotation_step * index)
|
||||
current_pose["position"] = segment_positions[-1].copy()
|
||||
current_pose["rotation"] = segment_rotations[-1].copy()
|
||||
positions.extend(segment_positions)
|
||||
rotations.extend(segment_rotations)
|
||||
|
||||
rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] +
|
||||
[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
|
||||
for index, (position, rotation) in enumerate(zip(positions, rotations, strict=False)):
|
||||
rotation_matrix = _quaternion_to_rotation_matrix(_euler_to_quaternion(rotation))
|
||||
translation = -rotation_matrix @ position
|
||||
extrinsic = np.hstack([rotation_matrix, translation.reshape(3, 1)])
|
||||
rows.append([float(index)] + _INTRINSIC_ROW + [0.0, 0.0] + extrinsic.flatten().tolist())
|
||||
return rows
|
||||
|
||||
|
||||
def _interpolate_camera_poses(
|
||||
cameras: list[DreamXCamera],
|
||||
src_indices: np.ndarray,
|
||||
tgt_indices: np.ndarray,
|
||||
) -> list[DreamXCamera]:
|
||||
if len(cameras) <= 1:
|
||||
return [cameras[0]] * len(tgt_indices) if cameras else []
|
||||
src_rot_mat = np.array([camera.w2c_mat[:3, :3] for camera in cameras])
|
||||
src_trans_vec = np.array([camera.w2c_mat[:3, 3] for camera in cameras])
|
||||
|
||||
dets = np.linalg.det(src_rot_mat)
|
||||
flip_handedness = dets.size > 0 and np.median(dets) < 0.0
|
||||
if flip_handedness:
|
||||
flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
|
||||
src_rot_mat = src_rot_mat @ flip_mat
|
||||
|
||||
trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False,
|
||||
fill_value="extrapolate")(tgt_indices)
|
||||
quats = Rotation.from_matrix(src_rot_mat).as_quat().copy()
|
||||
for index in range(1, len(quats)):
|
||||
if np.dot(quats[index], quats[index - 1]) < 0:
|
||||
quats[index] = -quats[index]
|
||||
rot = Slerp(src_indices, Rotation.from_quat(quats))(tgt_indices).as_matrix()
|
||||
if flip_handedness:
|
||||
rot = rot @ flip_mat
|
||||
|
||||
ref = cameras[0]
|
||||
result = []
|
||||
for index in range(len(tgt_indices)):
|
||||
w2c_mat = np.eye(4, dtype=np.float64)
|
||||
w2c_mat[:3, :] = np.hstack([rot[index], trans[index].reshape(3, 1)])
|
||||
result.append(DreamXCamera(ref.fx, ref.fy, ref.cx, ref.cy, w2c_mat))
|
||||
return result
|
||||
|
||||
|
||||
def _relative_c2w_poses(cameras: list[DreamXCamera]) -> np.ndarray:
|
||||
abs_w2cs = [camera.w2c_mat for camera in cameras]
|
||||
abs_c2ws = [camera.c2w_mat for camera in cameras]
|
||||
target_cam_c2w = np.eye(4, dtype=np.float64)
|
||||
abs2rel = target_cam_c2w @ abs_w2cs[0]
|
||||
poses = [target_cam_c2w] + [abs2rel @ c2w for c2w in abs_c2ws[1:]]
|
||||
return np.asarray(poses, dtype=np.float32)
|
||||
|
||||
|
||||
def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
|
||||
rotation_inv = transforms[..., :3, :3].transpose(-1, -2)
|
||||
output = torch.zeros_like(transforms)
|
||||
output[..., :3, :3] = rotation_inv
|
||||
output[..., :3, 3] = -torch.einsum("...ij,...j->...i", rotation_inv, transforms[..., :3, 3])
|
||||
output[..., 3, 3] = 1.0
|
||||
return output
|
||||
|
||||
|
||||
def build_dreamx_camera_condition(
|
||||
action_seq: list[str],
|
||||
action_speed_list: list[float],
|
||||
*,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
device: torch.device | str = "cpu",
|
||||
) -> dict[str, torch.Tensor]:
|
||||
del height, width # DreamX-World-5B-Cam uses fixed normalized intrinsics.
|
||||
duration = math.ceil(num_frames / len(action_seq))
|
||||
rows = _pose_rows_from_actions(action_seq, action_speed_list, duration)[:num_frames]
|
||||
cameras = [DreamXCamera.from_pose_row(row) for row in rows]
|
||||
|
||||
latent_frame_count = 1 + (len(cameras) - 1) // 4
|
||||
src_indices = np.arange(len(cameras), dtype=np.float64)
|
||||
tgt_indices = np.linspace(0, len(cameras) - 1, latent_frame_count)
|
||||
cameras = _interpolate_camera_poses(cameras, src_indices, tgt_indices)
|
||||
|
||||
c2ws = torch.as_tensor(_relative_c2w_poses(cameras), dtype=dtype, device=device)
|
||||
viewmats = _invert_se3(c2ws)
|
||||
|
||||
intrinsics = torch.zeros((latent_frame_count, 3, 3), dtype=dtype, device=device)
|
||||
intrinsics[:, 0, 0] = 969.6969696969696 / (960.0 * 2)
|
||||
intrinsics[:, 1, 1] = 969.6969696969696 / (540.0 * 2)
|
||||
intrinsics[:, 2, 2] = 1.0
|
||||
return {"viewmats": viewmats, "K": intrinsics}
|
||||
@@ -1,20 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Compatibility exports for DreamX-World pipeline configs."""
|
||||
|
||||
from fastvideo.configs.pipelines.dreamx_world import (
|
||||
DreamXWorld5BARPipelineConfig,
|
||||
DreamXWorld5BCamPipelineConfig,
|
||||
make_dreamx_world_5b_ar_dit_config,
|
||||
make_dreamx_world_5b_cam_dit_config,
|
||||
make_dreamx_world_5b_cam_text_encoder_config,
|
||||
make_dreamx_world_5b_cam_vae_config,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DreamXWorld5BARPipelineConfig",
|
||||
"DreamXWorld5BCamPipelineConfig",
|
||||
"make_dreamx_world_5b_ar_dit_config",
|
||||
"make_dreamx_world_5b_cam_dit_config",
|
||||
"make_dreamx_world_5b_cam_text_encoder_config",
|
||||
"make_dreamx_world_5b_cam_vae_config",
|
||||
]
|
||||
@@ -1,67 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World-5B autoregressive pipeline entrypoint."""
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
||||
from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import (
|
||||
DreamXWorldCameraConditioningStage,
|
||||
DreamXWorldImageVAEEncodingStage,
|
||||
)
|
||||
from fastvideo.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DreamXWorldARPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
"""DreamX-World-5B autoregressive causal camera pipeline."""
|
||||
|
||||
_required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
|
||||
pipeline_config_cls = DreamXWorld5BARPipelineConfig
|
||||
sampling_params_cls = SamplingParam
|
||||
|
||||
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
|
||||
self.modules["scheduler"].set_timesteps(1000)
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
|
||||
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
|
||||
self.add_stage(stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
))
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
self.add_stage(stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
|
||||
self.add_stage(stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
))
|
||||
self.add_stage(stage_name="image_latent_preparation_stage",
|
||||
stage=DreamXWorldImageVAEEncodingStage(vae=self.get_module("vae")))
|
||||
self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
|
||||
self.add_stage(stage_name="denoising_stage",
|
||||
stage=DreamXWorldARCausalDenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
vae=self.get_module("vae"),
|
||||
pipeline=self,
|
||||
))
|
||||
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
|
||||
logger.info("DreamXWorldARPipeline initialized with autoregressive causal denoising")
|
||||
|
||||
|
||||
EntryClass = DreamXWorldARPipeline
|
||||
@@ -1,78 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World video pipeline entrypoint."""
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
||||
from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraConditioningStage
|
||||
from fastvideo.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DreamXWorldPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
"""DreamX-World-5B-Cam pipeline with native FastVideo camera conditioning."""
|
||||
|
||||
_required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
|
||||
pipeline_config_cls = DreamXWorld5BCamPipelineConfig
|
||||
sampling_params_cls = SamplingParam
|
||||
|
||||
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
|
||||
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
transformer_2=self.get_module("transformer_2", None),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
vae=self.get_module("vae"),
|
||||
pipeline=self,
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
|
||||
|
||||
logger.info("DreamXWorldPipeline initialized with native camera conditioning")
|
||||
|
||||
|
||||
EntryClass = DreamXWorldPipeline
|
||||
@@ -1,58 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World model family pipeline presets."""
|
||||
|
||||
from fastvideo.api.presets import InferencePreset, PresetStageSpec
|
||||
|
||||
_NEGATIVE_PROMPT_CN = ("色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,"
|
||||
"静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,"
|
||||
"多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,"
|
||||
"形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,"
|
||||
"背景人很多,倒着走")
|
||||
|
||||
_DENOISE_STAGE = PresetStageSpec(
|
||||
name="denoise",
|
||||
kind="denoising",
|
||||
description="DreamX-World camera-conditioned denoising pass",
|
||||
allowed_overrides=frozenset({
|
||||
"num_inference_steps",
|
||||
"guidance_scale",
|
||||
}),
|
||||
)
|
||||
|
||||
DREAMX_WORLD_5B_CAM = InferencePreset(
|
||||
name="dreamx_world_5b_cam",
|
||||
version=1,
|
||||
model_family="dreamx_world",
|
||||
description="DreamX-World 5B camera-control video generation",
|
||||
workload_type="i2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 161,
|
||||
"fps": 16,
|
||||
"guidance_scale": 5.0,
|
||||
"num_inference_steps": 30,
|
||||
"negative_prompt": _NEGATIVE_PROMPT_CN,
|
||||
},
|
||||
)
|
||||
|
||||
DREAMX_WORLD_5B_AR = InferencePreset(
|
||||
name="dreamx_world_5b_ar",
|
||||
version=1,
|
||||
model_family="dreamx_world",
|
||||
description="DreamX-World 5B autoregressive camera-control generation",
|
||||
workload_type="i2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 704,
|
||||
"width": 1280,
|
||||
"num_frames": 1005,
|
||||
"fps": 16,
|
||||
"guidance_scale": 1.0,
|
||||
"num_inference_steps": 4,
|
||||
"negative_prompt": _NEGATIVE_PROMPT_CN,
|
||||
},
|
||||
)
|
||||
|
||||
ALL_PRESETS = (DREAMX_WORLD_5B_CAM, DREAMX_WORLD_5B_AR)
|
||||
@@ -1,156 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World pipeline stages."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.base import PipelineStage
|
||||
from fastvideo.pipelines.stages.validators import VerificationResult
|
||||
|
||||
from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (
|
||||
build_dreamx_camera_condition, )
|
||||
|
||||
DREAMX_Y_CAMERA_KEY = "dreamx_y_camera"
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DreamXWorldCameraConditioningStage(PipelineStage):
|
||||
"""Build PRoPE camera conditioning for DreamX-World-5B-Cam."""
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> ForwardBatch:
|
||||
del fastvideo_args
|
||||
if DREAMX_Y_CAMERA_KEY in batch.extra:
|
||||
return batch
|
||||
|
||||
action_seq = batch.extra.get("dreamx_action_seq", batch.action_list)
|
||||
action_speed_list = batch.extra.get("dreamx_action_speed_list", batch.action_speed_list)
|
||||
if action_seq is None:
|
||||
action_seq = ["w"]
|
||||
if action_speed_list is None:
|
||||
action_speed_list = [4]
|
||||
|
||||
if isinstance(action_seq, str):
|
||||
action_seq = [action_seq]
|
||||
if isinstance(action_speed_list, int | float):
|
||||
action_speed_list = [action_speed_list]
|
||||
if len(action_speed_list) == 1 and len(action_seq) > 1:
|
||||
action_speed_list = list(action_speed_list) * len(action_seq)
|
||||
action_speed_list = [float(speed) for speed in action_speed_list]
|
||||
|
||||
height = int(batch.height) if batch.height is not None else 704
|
||||
width = int(batch.width) if batch.width is not None else 1280
|
||||
num_frames = int(batch.num_frames)
|
||||
dtype = batch.latents.dtype if torch.is_tensor(batch.latents) else torch.float32
|
||||
device = batch.latents.device if torch.is_tensor(batch.latents) else "cpu"
|
||||
|
||||
y_camera = build_dreamx_camera_condition(
|
||||
list(action_seq),
|
||||
action_speed_list,
|
||||
num_frames=num_frames,
|
||||
height=height,
|
||||
width=width,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
batch.extra[DREAMX_Y_CAMERA_KEY] = {key: value.unsqueeze(0) for key, value in y_camera.items()}
|
||||
return batch
|
||||
|
||||
def verify_output(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> VerificationResult:
|
||||
del fastvideo_args
|
||||
result = VerificationResult()
|
||||
y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY)
|
||||
result.add_check("dreamx_y_camera", y_camera, lambda value: isinstance(value, dict))
|
||||
if isinstance(y_camera, dict):
|
||||
result.add_check("dreamx_y_camera.viewmats", y_camera.get("viewmats"), torch.is_tensor)
|
||||
result.add_check("dreamx_y_camera.K", y_camera.get("K"), torch.is_tensor)
|
||||
return result
|
||||
|
||||
|
||||
class DreamXWorldImageVAEEncodingStage(PipelineStage):
|
||||
"""Encode the conditioning image into the first-frame latent.
|
||||
|
||||
Official AR-forcing flow (AMAP-ML/DreamX-World inference_ar_forcing.py):
|
||||
the input image is resized, normalized to [-1, 1], VAE-encoded
|
||||
deterministically, and written into frame 0 of the noise — the causal
|
||||
denoiser then treats frame 0 as clean context. This stage produces
|
||||
``batch.image_latent`` ([B, C, 1, H_lat, W_lat]); the injection into
|
||||
the latents happens in DreamXWorldARCausalDenoisingStage.
|
||||
"""
|
||||
|
||||
def __init__(self, vae) -> None:
|
||||
super().__init__()
|
||||
self.vae = vae
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> ForwardBatch:
|
||||
if batch.pil_image is None:
|
||||
# No conditioning image: the causal denoiser falls back to
|
||||
# running from pure noise (frame 0 uninitialized). Warn loudly —
|
||||
# this pipeline is registered I2V and the official flow always
|
||||
# forces from a frame.
|
||||
logger.warning("DreamXWorldARPipeline called without an input image; "
|
||||
"first-frame context will be noise (T2V-style). Pass an "
|
||||
"image for the official AR-forcing behavior.")
|
||||
return batch
|
||||
|
||||
from fastvideo.platforms import get_local_torch_device
|
||||
from fastvideo.utils import PRECISION_TO_TYPE
|
||||
|
||||
device = get_local_torch_device()
|
||||
image = batch.pil_image
|
||||
if not isinstance(image, torch.Tensor):
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
assert isinstance(image, PIL.Image.Image)
|
||||
width = batch.width if isinstance(batch.width, int) else batch.width[0]
|
||||
height = batch.height if isinstance(batch.height, int) else batch.height[0]
|
||||
image = image.convert("RGB").resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
arr = torch.from_numpy(np.asarray(image)).float().permute(2, 0, 1) / 255.0
|
||||
image = (arr - 0.5) / 0.5 # official Normalize([0.5], [0.5])
|
||||
image = image.unsqueeze(0) # [1, C, H, W]
|
||||
if image.dim() == 4:
|
||||
image = image.unsqueeze(2) # [B, C, 1, H, W]
|
||||
elif image.dim() == 5:
|
||||
image = image[:, :, :1]
|
||||
image = image.to(device=device, dtype=torch.float32)
|
||||
|
||||
vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
|
||||
vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
|
||||
self.vae = self.vae.to(device)
|
||||
with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
|
||||
if not vae_autocast_enabled:
|
||||
image = image.to(vae_dtype)
|
||||
encoder_output = self.vae.encode(image)
|
||||
|
||||
# Official encode_to_latent is deterministic ((mean - mu) / sigma per
|
||||
# channel); the posterior mean + shift/scale is the FastVideo
|
||||
# equivalent of that normalization.
|
||||
latent = encoder_output.mean
|
||||
if getattr(self.vae, "shift_factor", None) is not None:
|
||||
shift = self.vae.shift_factor
|
||||
latent = latent - (shift.to(latent.device, latent.dtype) if isinstance(shift, torch.Tensor) else shift)
|
||||
scale = self.vae.scaling_factor
|
||||
latent = latent * (scale.to(latent.device, latent.dtype) if isinstance(scale, torch.Tensor) else scale)
|
||||
|
||||
batch.image_latent = latent
|
||||
return batch
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
result = VerificationResult()
|
||||
return result
|
||||
@@ -1 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -1,89 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
|
||||
from fastvideo.pipelines.stages.input_validation import InputValidationStage
|
||||
from fastvideo.pipelines.stages.kandinsky5 import (
|
||||
Kandinsky5DecodingStage,
|
||||
Kandinsky5DenoisingStage,
|
||||
Kandinsky5ImageEncodingStage,
|
||||
Kandinsky5LatentPreparationStage,
|
||||
Kandinsky5NormalizationStage,
|
||||
)
|
||||
from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
|
||||
from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
|
||||
|
||||
|
||||
class Kandinsky5I2VPipeline(ComposedPipelineBase):
|
||||
"""Kandinsky-5.0 image-to-video pipeline."""
|
||||
|
||||
_required_config_modules = [
|
||||
"scheduler",
|
||||
"text_encoder",
|
||||
"text_encoder_2",
|
||||
"tokenizer",
|
||||
"tokenizer_2",
|
||||
"transformer",
|
||||
"vae",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
|
||||
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="text_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[
|
||||
self.get_module("text_encoder"),
|
||||
self.get_module("text_encoder_2"),
|
||||
],
|
||||
tokenizers=[
|
||||
self.get_module("tokenizer"),
|
||||
self.get_module("tokenizer_2"),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=Kandinsky5LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
# Runs AFTER latent preparation so the initial noise is the seeded
|
||||
# generator's first draw (official kandinsky-5 RNG order); this stage
|
||||
# then samples the image latent and places it into the latents.
|
||||
self.add_stage(
|
||||
stage_name="image_encoding_stage",
|
||||
stage=Kandinsky5ImageEncodingStage(vae=self.get_module("vae")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=Kandinsky5DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="normalization_stage",
|
||||
stage=Kandinsky5NormalizationStage(),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage",
|
||||
stage=Kandinsky5DecodingStage(vae=self.get_module("vae"), pipeline=self),
|
||||
)
|
||||
|
||||
|
||||
EntryClass = Kandinsky5I2VPipeline
|
||||
@@ -1,74 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
|
||||
from fastvideo.pipelines.stages.input_validation import InputValidationStage
|
||||
from fastvideo.pipelines.stages.kandinsky5 import (
|
||||
Kandinsky5DecodingStage,
|
||||
Kandinsky5DenoisingStage,
|
||||
Kandinsky5LatentPreparationStage,
|
||||
)
|
||||
from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
|
||||
from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
|
||||
|
||||
|
||||
class Kandinsky5T2VPipeline(ComposedPipelineBase):
|
||||
"""Kandinsky-5.0 Lite text-to-video pipeline."""
|
||||
|
||||
_required_config_modules = [
|
||||
"scheduler",
|
||||
"text_encoder",
|
||||
"text_encoder_2",
|
||||
"tokenizer",
|
||||
"tokenizer_2",
|
||||
"transformer",
|
||||
"vae",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
|
||||
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="text_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[
|
||||
self.get_module("text_encoder"),
|
||||
self.get_module("text_encoder_2"),
|
||||
],
|
||||
tokenizers=[
|
||||
self.get_module("tokenizer"),
|
||||
self.get_module("tokenizer_2"),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=Kandinsky5LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=Kandinsky5DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage",
|
||||
stage=Kandinsky5DecodingStage(vae=self.get_module("vae"), pipeline=self),
|
||||
)
|
||||
|
||||
|
||||
EntryClass = Kandinsky5T2VPipeline
|
||||
@@ -1,165 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Kandinsky-5 model family pipeline presets."""
|
||||
|
||||
from fastvideo.api.presets import InferencePreset, PresetStageSpec
|
||||
|
||||
_NEGATIVE_PROMPT = ("Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, "
|
||||
"deformed, walking backwards")
|
||||
|
||||
_DENOISE_STAGE = PresetStageSpec(
|
||||
name="denoise",
|
||||
kind="denoising",
|
||||
description="Main Kandinsky-5 denoising pass",
|
||||
allowed_overrides=frozenset({
|
||||
"num_inference_steps",
|
||||
"guidance_scale",
|
||||
}),
|
||||
)
|
||||
|
||||
KANDINSKY5_T2V_LITE_5S = InferencePreset(
|
||||
name="kandinsky5_t2v_lite_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Lite T2V 5s",
|
||||
workload_type="t2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 5.0,
|
||||
"num_inference_steps": 50,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_T2V_LITE_DISTILLED_5S = InferencePreset(
|
||||
name="kandinsky5_t2v_lite_distilled_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Lite T2V Distilled 5s",
|
||||
workload_type="t2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 1.0,
|
||||
"num_inference_steps": 16,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_T2V_PRO_5S = InferencePreset(
|
||||
name="kandinsky5_t2v_pro_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Pro T2V 5s",
|
||||
workload_type="t2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 5.0,
|
||||
"num_inference_steps": 50,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_T2V_PRO_DISTILLED_5S = InferencePreset(
|
||||
name="kandinsky5_t2v_pro_distilled_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Pro T2V Distilled 5s",
|
||||
workload_type="t2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 1.0,
|
||||
"num_inference_steps": 16,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_I2V_LITE_5S = InferencePreset(
|
||||
name="kandinsky5_i2v_lite_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Lite I2V 5s",
|
||||
workload_type="i2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 5.0,
|
||||
"num_inference_steps": 50,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_I2V_PRO_5S = InferencePreset(
|
||||
name="kandinsky5_i2v_pro_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Pro I2V 5s",
|
||||
workload_type="i2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 5.0,
|
||||
"num_inference_steps": 50,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_I2V_LITE_DISTILLED_5S = InferencePreset(
|
||||
name="kandinsky5_i2v_lite_distilled_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Lite I2V Distilled 5s",
|
||||
workload_type="i2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 1.0,
|
||||
"num_inference_steps": 16,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
KANDINSKY5_I2V_PRO_DISTILLED_5S = InferencePreset(
|
||||
name="kandinsky5_i2v_pro_distilled_5s",
|
||||
version=1,
|
||||
model_family="kandinsky5",
|
||||
description="Kandinsky-5.0 Pro I2V Distilled 5s",
|
||||
workload_type="i2v",
|
||||
stage_schemas=(_DENOISE_STAGE, ),
|
||||
defaults={
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"fps": 24,
|
||||
"guidance_scale": 1.0,
|
||||
"num_inference_steps": 16,
|
||||
"negative_prompt": _NEGATIVE_PROMPT,
|
||||
},
|
||||
)
|
||||
|
||||
ALL_PRESETS = (KANDINSKY5_T2V_LITE_5S, KANDINSKY5_T2V_LITE_DISTILLED_5S, KANDINSKY5_T2V_PRO_5S,
|
||||
KANDINSKY5_T2V_PRO_DISTILLED_5S, KANDINSKY5_I2V_LITE_5S, KANDINSKY5_I2V_LITE_DISTILLED_5S,
|
||||
KANDINSKY5_I2V_PRO_5S, KANDINSKY5_I2V_PRO_DISTILLED_5S)
|
||||
@@ -19,39 +19,6 @@ from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage, LTX
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _resolve_refine_path(model_path: str, value: str | None) -> str | None:
|
||||
"""Resolve a refine component path, preferring paths inside the local
|
||||
model snapshot when the value is relative."""
|
||||
if value is None:
|
||||
return None
|
||||
if os.path.isabs(value):
|
||||
return value
|
||||
candidate = os.path.join(model_path, value)
|
||||
if os.path.exists(candidate):
|
||||
return candidate
|
||||
return value
|
||||
|
||||
|
||||
def _resolve_refine_upsampler_path(model_path: str, model_index: dict[str, Any]) -> str | None:
|
||||
"""Resolve the refine upsampler directory for a local model snapshot.
|
||||
|
||||
model_index.json keys ("fastvideo_refine_upsampler_path", then
|
||||
"spatial_upsampler") are the documented override and win when present.
|
||||
Otherwise probe the snapshot root: published LTX-2 distilled checkpoints
|
||||
bundle the upsampler weights on disk (spelled "spatial_upscaler")
|
||||
without declaring them in model_index.json.
|
||||
"""
|
||||
path = _resolve_refine_path(model_path, model_index.get("fastvideo_refine_upsampler_path"))
|
||||
if path is None and "spatial_upsampler" in model_index:
|
||||
path = _resolve_refine_path(model_path, "spatial_upsampler")
|
||||
if path is None:
|
||||
for dirname in ("spatial_upscaler", "spatial_upsampler"):
|
||||
candidate = os.path.join(model_path, dirname)
|
||||
if os.path.isdir(candidate):
|
||||
return candidate
|
||||
return path
|
||||
|
||||
|
||||
class LTX2Pipeline(LoRAPipeline):
|
||||
|
||||
_required_config_modules = [
|
||||
@@ -222,16 +189,28 @@ class LTX2Pipeline(LoRAPipeline):
|
||||
# Apply optional FastVideo-specific refine defaults embedded in
|
||||
# model_index.json. These are bundled with distilled checkpoints
|
||||
# so the pipeline can self-configure without explicit user kwargs.
|
||||
def _resolve_refine_path(value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
if os.path.isabs(value):
|
||||
return value
|
||||
candidate = os.path.join(self.model_path, value)
|
||||
if os.path.exists(candidate):
|
||||
return candidate
|
||||
return value
|
||||
|
||||
if (model_index.get("fastvideo_refine_enabled") is True and fastvideo_args.refine_enabled is None):
|
||||
fastvideo_args.ltx2_refine_enabled = True
|
||||
if (fastvideo_args.refine_upsampler_path is None and fastvideo_args.ltx2_refine_upsampler_path is None):
|
||||
fastvideo_args.ltx2_refine_upsampler_path = (_resolve_refine_upsampler_path(self.model_path, model_index))
|
||||
fastvideo_args.ltx2_refine_upsampler_path = _resolve_refine_path(
|
||||
model_index.get("fastvideo_refine_upsampler_path"))
|
||||
if (fastvideo_args.ltx2_refine_upsampler_path is None and "spatial_upsampler" in model_index):
|
||||
fastvideo_args.ltx2_refine_upsampler_path = (_resolve_refine_path("spatial_upsampler"))
|
||||
if (fastvideo_args.refine_transformer_path is None and fastvideo_args.ltx2_refine_transformer_path is None):
|
||||
fastvideo_args.ltx2_refine_transformer_path = (_resolve_refine_path(
|
||||
self.model_path, model_index.get("fastvideo_refine_transformer_path")))
|
||||
model_index.get("fastvideo_refine_transformer_path")))
|
||||
if (fastvideo_args.refine_lora_path is None and fastvideo_args.ltx2_refine_lora_path is None):
|
||||
fastvideo_args.ltx2_refine_lora_path = _resolve_refine_path(self.model_path,
|
||||
model_index.get("fastvideo_refine_lora_path"))
|
||||
fastvideo_args.ltx2_refine_lora_path = _resolve_refine_path(model_index.get("fastvideo_refine_lora_path"))
|
||||
if (fastvideo_args.refine_num_inference_steps is None
|
||||
and fastvideo_args.ltx2_refine_num_inference_steps == FastVideoArgs.ltx2_refine_num_inference_steps
|
||||
and model_index.get("fastvideo_refine_num_inference_steps") is not None):
|
||||
@@ -247,11 +226,10 @@ class LTX2Pipeline(LoRAPipeline):
|
||||
if (fastvideo_args.refine_add_noise is None and model_index.get("fastvideo_refine_add_noise") is not None):
|
||||
fastvideo_args.ltx2_refine_add_noise = bool(model_index["fastvideo_refine_add_noise"])
|
||||
if (fastvideo_args.refine_noise_path is None and fastvideo_args.ltx2_refine_noise_path is None):
|
||||
fastvideo_args.ltx2_refine_noise_path = _resolve_refine_path(self.model_path,
|
||||
model_index.get("fastvideo_refine_noise_path"))
|
||||
fastvideo_args.ltx2_refine_noise_path = _resolve_refine_path(model_index.get("fastvideo_refine_noise_path"))
|
||||
if (fastvideo_args.refine_audio_noise_path is None and fastvideo_args.ltx2_refine_audio_noise_path is None):
|
||||
fastvideo_args.ltx2_refine_audio_noise_path = (_resolve_refine_path(
|
||||
self.model_path, model_index.get("fastvideo_refine_audio_noise_path")))
|
||||
model_index.get("fastvideo_refine_audio_noise_path")))
|
||||
|
||||
model_index.pop("_class_name")
|
||||
model_index.pop("_diffusers_version")
|
||||
|
||||
@@ -33,8 +33,6 @@ from fastvideo.pipelines.basic.ltx2.stages import ( # noqa: F401
|
||||
from fastvideo.pipelines.stages.matrixgame2_denoising import MatrixGame2CausalDenoisingStage
|
||||
from fastvideo.pipelines.stages.matrixgame3_denoising import MatrixGame3DenoisingStage
|
||||
from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage
|
||||
from fastvideo.pipelines.stages.kandinsky5 import (Kandinsky5DecodingStage, Kandinsky5DenoisingStage,
|
||||
Kandinsky5LatentPreparationStage)
|
||||
from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage
|
||||
from fastvideo.pipelines.stages.gen3c_stages import (Gen3CCFGPolicyStage, Gen3CConditioningStage, Gen3CDenoisingStage,
|
||||
Gen3CLatentPreparationStage)
|
||||
@@ -67,9 +65,6 @@ __all__ = [
|
||||
"MatrixGame2CausalDenoisingStage",
|
||||
"MatrixGame3DenoisingStage",
|
||||
"HYWorldDenoisingStage",
|
||||
"Kandinsky5DecodingStage",
|
||||
"Kandinsky5DenoisingStage",
|
||||
"Kandinsky5LatentPreparationStage",
|
||||
"GameCraftDenoisingStage",
|
||||
"Gen3CCFGPolicyStage",
|
||||
"Gen3CConditioningStage",
|
||||
|
||||
@@ -34,7 +34,6 @@ class PipelineStage(ABC):
|
||||
composed with other stages to create a complete pipeline. Each stage is responsible
|
||||
for a specific part of the process, such as prompt encoding, latent preparation, etc.
|
||||
"""
|
||||
performance_component_metric: str | None = None
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
"""
|
||||
@@ -156,9 +155,6 @@ class PipelineStage(ABC):
|
||||
logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000)
|
||||
batch.logging_info.add_stage_execution_time(stage_key, execution_time)
|
||||
batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name)
|
||||
component_metric = self.performance_component_metric
|
||||
if component_metric is not None:
|
||||
batch.logging_info.add_stage_metric(stage_key, "component_metric", component_metric)
|
||||
except Exception as e:
|
||||
torch.cuda.synchronize()
|
||||
execution_time = time.perf_counter() - start_time
|
||||
|
||||
@@ -28,7 +28,6 @@ class DecodingStage(PipelineStage):
|
||||
This stage handles the decoding of latent representations into the final
|
||||
output format (e.g., pixel values).
|
||||
"""
|
||||
performance_component_metric = "vae_decode_time_s"
|
||||
|
||||
def __init__(self, vae, pipeline=None) -> None:
|
||||
self.vae: ParallelTiledVAE = vae
|
||||
|
||||
@@ -51,7 +51,6 @@ class DenoisingStage(PipelineStage):
|
||||
This stage handles the iterative denoising process that transforms
|
||||
the initial noise into the final output.
|
||||
"""
|
||||
performance_component_metric = "dit_time_s"
|
||||
|
||||
def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
|
||||
super().__init__()
|
||||
@@ -191,19 +190,6 @@ class DenoisingStage(PipelineStage):
|
||||
},
|
||||
)
|
||||
|
||||
dreamx_y_camera = batch.extra.get("dreamx_y_camera", batch.extra.get("y_camera"))
|
||||
if isinstance(dreamx_y_camera, dict):
|
||||
dreamx_y_camera = {
|
||||
key: value.to(device=local_device, dtype=target_dtype) if torch.is_tensor(value) else value
|
||||
for key, value in dreamx_y_camera.items()
|
||||
}
|
||||
dreamx_camera_kwargs = self.prepare_extra_func_kwargs(
|
||||
self.transformer.forward,
|
||||
{
|
||||
"y_camera": dreamx_y_camera,
|
||||
},
|
||||
)
|
||||
|
||||
for key in ("flux2_txt_ids", "flux2_img_ids"):
|
||||
value = batch.extra.get(key)
|
||||
if torch.is_tensor(value):
|
||||
@@ -255,11 +241,7 @@ class DenoisingStage(PipelineStage):
|
||||
# the image latent instead of appending along the channel dim
|
||||
assert batch.image_latent is None, "TI2V task should not have image latents"
|
||||
assert self.vae is not None, "VAE is not provided for TI2V task"
|
||||
vae_device = next(self.vae.parameters()).device
|
||||
self.vae = self.vae.to(local_device)
|
||||
z = self.vae.encode(batch.pil_image).mean.float()
|
||||
if getattr(fastvideo_args, "vae_cpu_offload", False):
|
||||
self.vae = self.vae.to(vae_device)
|
||||
if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
|
||||
if isinstance(self.vae.shift_factor, torch.Tensor):
|
||||
z -= self.vae.shift_factor.to(z.device, z.dtype)
|
||||
@@ -512,7 +494,6 @@ class DenoisingStage(PipelineStage):
|
||||
**pos_cond_kwargs,
|
||||
**action_kwargs,
|
||||
**camera_kwargs,
|
||||
**dreamx_camera_kwargs,
|
||||
**timesteps_r_kwarg,
|
||||
**flux2_id_kwargs,
|
||||
)
|
||||
@@ -555,7 +536,6 @@ class DenoisingStage(PipelineStage):
|
||||
**neg_cond_kwargs,
|
||||
**action_kwargs,
|
||||
**camera_kwargs,
|
||||
**dreamx_camera_kwargs,
|
||||
**timesteps_r_kwarg,
|
||||
**flux2_id_kwargs,
|
||||
)
|
||||
@@ -1207,7 +1187,6 @@ class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
|
||||
|
||||
class Cosmos25AutoDenoisingStage(PipelineStage):
|
||||
"""Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
|
||||
performance_component_metric = "dit_time_s"
|
||||
|
||||
def __init__(self, transformer, scheduler) -> None:
|
||||
super().__init__()
|
||||
|
||||
@@ -1,545 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from typing import Any
|
||||
|
||||
import PIL
|
||||
import torch
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder
|
||||
from fastvideo.distributed import get_local_torch_device
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.forward_context import set_forward_context
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.loader.component_loader import TransformerLoader, VAELoader
|
||||
from fastvideo.models.vaes.common import ParallelTiledVAE
|
||||
from fastvideo.models.vision_utils import normalize, numpy_to_pt, pil_to_numpy, resize
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.base import PipelineStage
|
||||
from fastvideo.pipelines.stages.decoding import DecodingStage
|
||||
from fastvideo.pipelines.stages.encoding import EncodingStage
|
||||
from fastvideo.pipelines.stages.validators import StageValidators as V
|
||||
from fastvideo.pipelines.stages.validators import VerificationResult
|
||||
from fastvideo.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Kandinsky5LatentPreparationStage(PipelineStage):
|
||||
|
||||
def __init__(self, scheduler, transformer) -> None:
|
||||
super().__init__()
|
||||
self.scheduler = scheduler
|
||||
self.transformer = transformer
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
if batch.height is None or batch.width is None:
|
||||
raise ValueError("height and width must be provided for Kandinsky5.")
|
||||
height = int(batch.height)
|
||||
width = int(batch.width)
|
||||
num_frames = int(batch.num_frames)
|
||||
|
||||
temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
|
||||
spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
|
||||
patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
|
||||
|
||||
if num_frames % temporal_ratio != 1:
|
||||
num_frames = num_frames // temporal_ratio * temporal_ratio + 1
|
||||
batch.num_frames = num_frames
|
||||
|
||||
required_divisor_h = spatial_ratio * patch_size[1]
|
||||
required_divisor_w = spatial_ratio * patch_size[2]
|
||||
if height % required_divisor_h != 0 or width % required_divisor_w != 0:
|
||||
raise ValueError(f"Kandinsky5 height must be divisible by {required_divisor_h} and width by "
|
||||
f"{required_divisor_w}; "
|
||||
f"got height={height}, width={width}.")
|
||||
|
||||
# NABLA sparse attention (Pro checkpoints) reshapes the post-patch grid
|
||||
# into 8x8 blocks; validate here instead of crashing mid-denoise after
|
||||
# all the encoding work is done.
|
||||
arch_cfg = getattr(self.transformer, "config", None) or fastvideo_args.pipeline_config.dit_config.arch_config
|
||||
if getattr(arch_cfg, "attention_type", "regular") == "nabla":
|
||||
nabla_divisor_h = required_divisor_h * 8
|
||||
nabla_divisor_w = required_divisor_w * 8
|
||||
if height % nabla_divisor_h != 0 or width % nabla_divisor_w != 0:
|
||||
raise ValueError(f"Kandinsky5 NABLA checkpoints require height divisible by {nabla_divisor_h} and "
|
||||
f"width divisible by {nabla_divisor_w}; "
|
||||
f"got height={height}, width={width}.")
|
||||
|
||||
if isinstance(batch.prompt, list):
|
||||
batch_size = len(batch.prompt)
|
||||
elif batch.prompt is not None:
|
||||
batch_size = 1
|
||||
else:
|
||||
batch_size = batch.prompt_embeds[0].shape[0]
|
||||
batch_size *= batch.num_videos_per_prompt
|
||||
|
||||
dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
|
||||
device = get_local_torch_device()
|
||||
num_latent_frames = (num_frames - 1) // temporal_ratio + 1
|
||||
num_channels = getattr(
|
||||
self.transformer,
|
||||
"in_visual_dim",
|
||||
fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
|
||||
)
|
||||
shape = (
|
||||
batch_size,
|
||||
num_latent_frames,
|
||||
height // spatial_ratio,
|
||||
width // spatial_ratio,
|
||||
num_channels,
|
||||
)
|
||||
|
||||
if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
|
||||
raise ValueError(f"generator list length {len(batch.generator)} does not match batch size {batch_size}.")
|
||||
|
||||
visual_cond = getattr(self.transformer, "visual_cond", False)
|
||||
|
||||
if batch.latents is None:
|
||||
latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
|
||||
if hasattr(self.scheduler, "init_noise_sigma"):
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
else:
|
||||
valid_shapes = [shape]
|
||||
if visual_cond:
|
||||
valid_shapes.append((*shape[:-1], num_channels * 2 + 1))
|
||||
if tuple(batch.latents.shape) not in valid_shapes:
|
||||
raise ValueError(f"Provided latents shape {list(batch.latents.shape)} does not match expected "
|
||||
f"Kandinsky5 latent shape(s): {[list(s) for s in valid_shapes]}.")
|
||||
latents = batch.latents.to(device=device, dtype=dtype)
|
||||
|
||||
if visual_cond and latents.shape[-1] == num_channels:
|
||||
cond = torch.zeros_like(latents)
|
||||
cond_mask = torch.zeros(
|
||||
(*latents.shape[:-1], 1),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
latents = torch.cat([latents, cond, cond_mask], dim=-1)
|
||||
|
||||
# I2V image conditioning is placed by Kandinsky5ImageEncodingStage,
|
||||
# which runs AFTER this stage so the initial noise is the generator's
|
||||
# first draw (matching the official kandinskylab/kandinsky-5 order).
|
||||
batch.latents = latents
|
||||
batch.raw_latent_shape = (
|
||||
batch_size,
|
||||
num_channels,
|
||||
num_latent_frames,
|
||||
height // spatial_ratio,
|
||||
width // spatial_ratio,
|
||||
)
|
||||
return batch
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
result = VerificationResult()
|
||||
result.add_check("num_frames", batch.num_frames, V.positive_int)
|
||||
result.add_check("height", batch.height, V.positive_int)
|
||||
result.add_check("width", batch.width, V.positive_int)
|
||||
result.add_check("latents", batch.latents, V.none_or_tensor)
|
||||
return result
|
||||
|
||||
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
result = VerificationResult()
|
||||
result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
|
||||
return result
|
||||
|
||||
|
||||
class Kandinsky5DenoisingStage(PipelineStage):
|
||||
|
||||
def __init__(self, transformer, scheduler) -> None:
|
||||
super().__init__()
|
||||
self.transformer = transformer
|
||||
self.scheduler = scheduler
|
||||
|
||||
@staticmethod
|
||||
def _scale_factor(height: int, width: int) -> tuple[float, float, float]:
|
||||
if 480 <= height <= 854 and 480 <= width <= 854:
|
||||
return (1.0, 2.0, 2.0)
|
||||
return (1.0, 3.16, 3.16)
|
||||
|
||||
@staticmethod
|
||||
def _text_rope_pos(mask: torch.Tensor, device: torch.device) -> torch.Tensor:
|
||||
seq_len = int(mask.sum(1).max().item())
|
||||
return torch.arange(seq_len, device=device)
|
||||
|
||||
@staticmethod
|
||||
def fast_sta_nabla(
|
||||
T: int,
|
||||
H: int,
|
||||
W: int,
|
||||
wT: int = 3,
|
||||
wH: int = 3,
|
||||
wW: int = 3,
|
||||
device: torch.device | str = "cuda",
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Create a sparse temporal attention (STA) mask for efficient video generation.
|
||||
|
||||
This method generates a mask that limits attention to nearby frames and spatial positions, reducing
|
||||
computational complexity for video generation.
|
||||
|
||||
Args:
|
||||
T (int): Number of temporal frames
|
||||
H (int): Height in latent space
|
||||
W (int): Width in latent space
|
||||
wT (int): Temporal attention window size
|
||||
wH (int): Height attention window size
|
||||
wW (int): Width attention window size
|
||||
device (str): Device to create tensor on
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Sparse attention mask of shape (T*H*W, T*H*W)
|
||||
"""
|
||||
max_extent = int(torch.tensor([T, H, W], device=device).amax().item())
|
||||
r = torch.arange(0, max_extent, 1, dtype=torch.int16, device=device)
|
||||
mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs()
|
||||
sta_t, sta_h, sta_w = (
|
||||
mat[:T, :T].flatten(),
|
||||
mat[:H, :H].flatten(),
|
||||
mat[:W, :W].flatten(),
|
||||
)
|
||||
sta_t = sta_t <= wT // 2
|
||||
sta_h = sta_h <= wH // 2
|
||||
sta_w = sta_w <= wW // 2
|
||||
sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten()
|
||||
sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2)
|
||||
return sta.reshape(T * H * W, T * H * W)
|
||||
|
||||
def get_sparse_params(self, sample: torch.Tensor, device: torch.device) -> dict[str, Any] | None:
|
||||
"""
|
||||
Generate sparse attention parameters for the transformer based on sample dimensions.
|
||||
|
||||
This method computes the sparse attention configuration needed for efficient video processing in the
|
||||
transformer model.
|
||||
|
||||
Args:
|
||||
sample (torch.Tensor): Input sample tensor
|
||||
device (torch.device): Device to place tensors on
|
||||
|
||||
Returns:
|
||||
Dict: Dictionary containing sparse attention parameters
|
||||
"""
|
||||
assert self.transformer.config.patch_size[0] == 1
|
||||
_, T, H, W, _ = sample.shape
|
||||
T, H, W = (
|
||||
T // self.transformer.config.patch_size[0],
|
||||
H // self.transformer.config.patch_size[1],
|
||||
W // self.transformer.config.patch_size[2],
|
||||
)
|
||||
if self.transformer.config.attention_type == "nabla":
|
||||
sta_mask = self.fast_sta_nabla(
|
||||
T,
|
||||
H // 8,
|
||||
W // 8,
|
||||
self.transformer.config.attention_wT,
|
||||
self.transformer.config.attention_wH,
|
||||
self.transformer.config.attention_wW,
|
||||
device=device,
|
||||
)
|
||||
|
||||
sparse_params = {
|
||||
"sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0),
|
||||
"attention_type": self.transformer.config.attention_type,
|
||||
"to_fractal": True,
|
||||
"P": self.transformer.config.attention_P,
|
||||
"wT": self.transformer.config.attention_wT,
|
||||
"wW": self.transformer.config.attention_wW,
|
||||
"wH": self.transformer.config.attention_wH,
|
||||
"add_sta": self.transformer.config.attention_add_sta,
|
||||
"visual_shape": (T, H, W),
|
||||
"method": self.transformer.config.attention_method,
|
||||
}
|
||||
else:
|
||||
sparse_params = None
|
||||
|
||||
return sparse_params
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
if batch.timesteps is None:
|
||||
raise ValueError("timesteps must be prepared before Kandinsky5 denoising.")
|
||||
if batch.latents is None:
|
||||
raise ValueError("latents must be prepared before Kandinsky5 denoising.")
|
||||
if not fastvideo_args.model_loaded["transformer"]:
|
||||
loader = TransformerLoader()
|
||||
self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
|
||||
fastvideo_args.model_loaded["transformer"] = True
|
||||
|
||||
device = get_local_torch_device()
|
||||
target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
|
||||
autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
|
||||
latents = batch.latents
|
||||
num_channels = getattr(
|
||||
self.transformer,
|
||||
"in_visual_dim",
|
||||
fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
|
||||
)
|
||||
|
||||
prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
|
||||
pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
|
||||
if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
|
||||
raise ValueError("Kandinsky5 requires Qwen prompt attention masks.")
|
||||
text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)
|
||||
|
||||
neg_prompt_embeds = None
|
||||
neg_pooled = None
|
||||
negative_text_rope_pos = None
|
||||
if batch.do_classifier_free_guidance and batch.negative_prompt_embeds:
|
||||
neg_prompt_embeds = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
|
||||
neg_pooled = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
|
||||
if batch.negative_attention_mask is None or not batch.negative_attention_mask:
|
||||
raise ValueError("Kandinsky5 requires Qwen negative attention masks for CFG.")
|
||||
negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask[0].to(device), device)
|
||||
|
||||
height = int(batch.height)
|
||||
width = int(batch.width)
|
||||
temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
|
||||
spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
|
||||
num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
|
||||
visual_rope_pos = [
|
||||
torch.arange(num_latent_frames, device=device),
|
||||
torch.arange(height // spatial_ratio // 2, device=device),
|
||||
torch.arange(width // spatial_ratio // 2, device=device),
|
||||
]
|
||||
scale_factor = self._scale_factor(height, width)
|
||||
|
||||
sparse_params = self.get_sparse_params(latents, device)
|
||||
|
||||
# I2V keeps the first (conditioning) frame fixed during denoising.
|
||||
# Key off the actual image conditioning, not transformer.visual_cond:
|
||||
# official T2V checkpoints also ship visual_cond=True, and skipping
|
||||
# frame 0 for them leaves it as undenoised noise.
|
||||
cond_frames = 1 if batch.image_latent is not None else 0
|
||||
|
||||
trajectory_timesteps: list[torch.Tensor] = []
|
||||
trajectory_latents: list[torch.Tensor] = []
|
||||
|
||||
with tqdm(total=batch.num_inference_steps, desc="Kandinsky5 Denoising") as progress_bar:
|
||||
for i, timestep in enumerate(batch.timesteps):
|
||||
if hasattr(self, "interrupt") and self.interrupt:
|
||||
break
|
||||
|
||||
t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
|
||||
attn_metadata = None
|
||||
if sparse_params is not None:
|
||||
attn_metadata = NablaAttentionMetadataBuilder().build(
|
||||
current_timestep=i,
|
||||
sta_mask=sparse_params["sta_mask"],
|
||||
P=sparse_params["P"],
|
||||
visual_shape=sparse_params["visual_shape"],
|
||||
)
|
||||
autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
|
||||
if device.type == "cuda" else contextlib.nullcontext())
|
||||
with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
|
||||
forward_batch=batch), autocast_ctx:
|
||||
pred_velocity = self.transformer(
|
||||
hidden_states=latents.to(dtype=target_dtype),
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
pooled_projections=pooled,
|
||||
timestep=t_expand,
|
||||
visual_rope_pos=visual_rope_pos,
|
||||
text_rope_pos=text_rope_pos,
|
||||
scale_factor=scale_factor,
|
||||
sparse_params=sparse_params,
|
||||
return_dict=True,
|
||||
).sample
|
||||
|
||||
if neg_prompt_embeds is not None and neg_pooled is not None:
|
||||
uncond_pred_velocity = self.transformer(
|
||||
hidden_states=latents.to(dtype=target_dtype),
|
||||
encoder_hidden_states=neg_prompt_embeds,
|
||||
pooled_projections=neg_pooled,
|
||||
timestep=t_expand,
|
||||
visual_rope_pos=visual_rope_pos,
|
||||
text_rope_pos=negative_text_rope_pos,
|
||||
scale_factor=scale_factor,
|
||||
sparse_params=sparse_params,
|
||||
return_dict=True,
|
||||
).sample
|
||||
pred_velocity = uncond_pred_velocity + batch.guidance_scale * (pred_velocity -
|
||||
uncond_pred_velocity)
|
||||
|
||||
latents[:, cond_frames:, :, :, :num_channels] = self.scheduler.step(
|
||||
pred_velocity[:, cond_frames:],
|
||||
timestep,
|
||||
latents[:, cond_frames:, :, :, :num_channels],
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if batch.return_trajectory_latents:
|
||||
trajectory_timesteps.append(timestep)
|
||||
# latents is mutated in place, so snapshot a channels-first copy.
|
||||
trajectory_latents.append(latents[..., :num_channels].permute(0, 4, 1, 2, 3).cpu())
|
||||
|
||||
if i == len(batch.timesteps) - 1 or (i + 1) % self.scheduler.order == 0:
|
||||
progress_bar.update()
|
||||
|
||||
if trajectory_latents:
|
||||
batch.trajectory_latents = torch.stack(trajectory_latents, dim=1)
|
||||
batch.trajectory_timesteps = torch.stack(trajectory_timesteps, dim=0).cpu()
|
||||
|
||||
batch.latents = latents[:, :, :, :, :num_channels]
|
||||
return batch
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
result = VerificationResult()
|
||||
result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
|
||||
result.add_check("prompt_embeds", batch.prompt_embeds, V.min_list_length(2))
|
||||
return result
|
||||
|
||||
|
||||
class Kandinsky5DecodingStage(DecodingStage):
|
||||
|
||||
def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:
|
||||
super().__init__(vae=vae, pipeline=pipeline)
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
if batch.latents is None:
|
||||
raise ValueError("latents must be available before Kandinsky5 decoding.")
|
||||
# Kandinsky5 latents are channels-last [B, T, H, W, C]; the base stage
|
||||
# (and the trajectory latents recorded by the denoising stage) work
|
||||
# channels-first.
|
||||
batch.latents = batch.latents.permute(0, 4, 1, 2, 3).contiguous()
|
||||
return super().forward(batch, fastvideo_args)
|
||||
|
||||
|
||||
class Kandinsky5ImageEncodingStage(EncodingStage):
|
||||
"""Encode the conditioning image into a VAE latent for I2V."""
|
||||
|
||||
def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:
|
||||
super().__init__(vae=vae)
|
||||
|
||||
@staticmethod
|
||||
def _preprocess(image, height: int, width: int) -> torch.Tensor:
|
||||
if isinstance(image, PIL.Image.Image):
|
||||
image = resize(image, height, width)
|
||||
image = numpy_to_pt(pil_to_numpy(image)) # always lands in [0, 1]
|
||||
return normalize(image) # [0, 1] -> [-1, 1]
|
||||
# Tensor input: no reliable way to tell [0, 1] from an already
|
||||
# normalized [-1, 1] tensor whose values happen to be non-negative,
|
||||
# so mirror diffusers' heuristic and say what we assumed.
|
||||
if image.min() >= 0:
|
||||
logger.warning("Kandinsky5 conditioning image tensor has no negative values; "
|
||||
"assuming range [0, 1] and normalizing to [-1, 1]. "
|
||||
"Pass a [-1, 1] tensor with negative values to skip normalization.")
|
||||
image = normalize(image) # [0, 1] -> [-1, 1]
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
if image.shape[-2:] != (height, width):
|
||||
image = torch.nn.functional.interpolate(image.float(),
|
||||
size=(height, width),
|
||||
mode="bilinear",
|
||||
antialias=True)
|
||||
return image
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
if batch.pil_image is None:
|
||||
raise ValueError("Kandinsky5 I2V requires an input image.")
|
||||
|
||||
if not fastvideo_args.model_loaded["vae"]:
|
||||
vae = getattr(self, "vae", None)
|
||||
if vae is None:
|
||||
loader = VAELoader()
|
||||
vae = loader.load(fastvideo_args.model_paths["vae"], fastvideo_args)
|
||||
self.vae = vae
|
||||
fastvideo_args.model_loaded["vae"] = True
|
||||
|
||||
device = get_local_torch_device()
|
||||
vae = self.vae.to(device)
|
||||
self.vae = vae
|
||||
vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
|
||||
vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast
|
||||
|
||||
# [B, C, H, W] -> [B, C, 1, H, W]
|
||||
image = self._preprocess(batch.pil_image, int(batch.height), int(batch.width))
|
||||
image = image.to(device=device, dtype=torch.float32).unsqueeze(2)
|
||||
|
||||
# Encode the single conditioning frame without tiling (matches diffusers).
|
||||
# The untested causal-VAE spatial_tiled_encode path corrupts the latent.
|
||||
prev_use_tiling = vae.use_tiling
|
||||
vae.use_tiling = False
|
||||
try:
|
||||
with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
|
||||
if not vae_autocast_enabled:
|
||||
image = image.to(vae_dtype)
|
||||
# Sample with the batch generator (diffusers parity); mode()
|
||||
# would make seed-for-seed reproduction of the reference
|
||||
# pipeline impossible.
|
||||
generator = batch.generator
|
||||
if isinstance(generator, list) and len(generator) != image.shape[0]:
|
||||
generator = generator[0]
|
||||
image_latent = vae.encode(image).sample(generator=generator)
|
||||
finally:
|
||||
vae.use_tiling = prev_use_tiling
|
||||
|
||||
image_latent = image_latent * vae.scaling_factor
|
||||
# [B, C, 1, H, W] -> [B, 1, H, W, C] to match channels-last latents
|
||||
batch.image_latent = image_latent.permute(0, 2, 3, 4, 1).contiguous()
|
||||
|
||||
# Place the conditioning latent into the prepared latents: frame 0 of
|
||||
# the main channels, the visual_cond channel block, and the mask.
|
||||
# NOTE: the official kandinsky-5 repo leaves the visual_cond block
|
||||
# zeros (generation_utils.py generate()), while the diffusers port
|
||||
# copies the image latent into it. A same-seed A/B on
|
||||
# Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers showed the Diffusers
|
||||
# export requires the copy: zeroing the block produces smeared faces
|
||||
# mid-video. Keep the diffusers semantics for Diffusers-format
|
||||
# checkpoints.
|
||||
latents = batch.latents
|
||||
image_latent = batch.image_latent.to(device=latents.device, dtype=latents.dtype)
|
||||
num_channels = image_latent.shape[-1]
|
||||
latents[:, 0:1, :, :, :num_channels] = image_latent
|
||||
if latents.shape[-1] > num_channels:
|
||||
latents[:, 0:1, :, :, num_channels:2 * num_channels] = image_latent
|
||||
latents[:, 0:1, :, :, 2 * num_channels:] = 1.0
|
||||
batch.latents = latents
|
||||
|
||||
if fastvideo_args.vae_cpu_offload:
|
||||
vae.to("cpu")
|
||||
return batch
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
result = VerificationResult()
|
||||
result.add_check("pil_image", batch.pil_image, V.not_none)
|
||||
result.add_check("height", batch.height, V.positive_int)
|
||||
result.add_check("width", batch.width, V.positive_int)
|
||||
# This stage runs after latent preparation and writes into its output.
|
||||
result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
|
||||
return result
|
||||
|
||||
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
result = VerificationResult()
|
||||
result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)])
|
||||
return result
|
||||
|
||||
|
||||
class Kandinsky5NormalizationStage(PipelineStage):
|
||||
"""Normalize the first latent frames to reduce I2V conditioning artifacts."""
|
||||
|
||||
COND_FRAMES = 4
|
||||
REFERENCE_FRAMES = 5
|
||||
|
||||
@staticmethod
|
||||
def _adaptive_mean_std(source: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
|
||||
source_mean = source.mean(dim=(1, 2, 3, 4), keepdim=True)
|
||||
source_std = source.std(dim=(1, 2, 3, 4), keepdim=True)
|
||||
# Magic constants limit how far the first frames may drift.
|
||||
ref_mean = torch.clamp(reference.mean(dim=(1, 2, 3, 4), keepdim=True), source_mean - 0.05, source_mean + 0.1)
|
||||
ref_std = torch.clamp(reference.std(dim=(1, 2, 3, 4), keepdim=True), source_std - 0.1, source_std + 0.25)
|
||||
normalized = (source - source_mean) / source_std
|
||||
return normalized * ref_std + ref_mean
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
latents = batch.latents
|
||||
n = self.COND_FRAMES
|
||||
if latents is None or latents.shape[1] <= n:
|
||||
return batch
|
||||
|
||||
reference = latents[:, n:n + min(self.REFERENCE_FRAMES, latents.shape[1] - 1)]
|
||||
latents[:, :n] = self._adaptive_mean_std(latents[:, :n].clone(), reference)
|
||||
batch.latents = latents
|
||||
return batch
|
||||
@@ -8,8 +8,6 @@ This module contains implementations of prompt encoding stages for diffusion pip
|
||||
import torch
|
||||
from typing import Any
|
||||
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
from fastvideo.distributed import get_local_torch_device
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.forward_context import set_forward_context
|
||||
@@ -26,7 +24,6 @@ class TextEncodingStage(PipelineStage):
|
||||
This stage handles the encoding of text prompts into the embedding space
|
||||
expected by the diffusion model.
|
||||
"""
|
||||
performance_component_metric = "text_encoder_time_s"
|
||||
|
||||
def __init__(self, text_encoders, tokenizers) -> None:
|
||||
"""
|
||||
@@ -205,21 +202,6 @@ class TextEncodingStage(PipelineStage):
|
||||
encoder_config = encoder_cfgs[i]
|
||||
preprocess_func = preprocess_funcs[i]
|
||||
postprocess_func = postprocess_funcs[i]
|
||||
# cpu_offload semantics: params rest on CPU between calls but the
|
||||
# forward computes on GPU. FSDP2-wrapped encoders (CPUOffloadPolicy;
|
||||
# DTensor params) stream themselves per-layer — leave inputs on the
|
||||
# param device and let FSDP's root pre-forward move them. A plain
|
||||
# module parked on CPU by text_encoder_cpu_offload is swapped to the
|
||||
# target device for the forward and back afterwards, mirroring the
|
||||
# image-encoder/VAE offload pattern.
|
||||
first_param = next(text_encoder.parameters(), None)
|
||||
encoder_device = first_param.device if first_param is not None else torch.device(target_device)
|
||||
moved_for_forward = False
|
||||
if (first_param is not None and not isinstance(first_param, DTensor)
|
||||
and encoder_device.type != torch.device(target_device).type):
|
||||
text_encoder = text_encoder.to(target_device)
|
||||
encoder_device = torch.device(target_device)
|
||||
moved_for_forward = True
|
||||
|
||||
tok_kwargs = dict(encoder_config.tokenizer_kwargs)
|
||||
if max_length is not None:
|
||||
@@ -264,7 +246,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(encoder_device)
|
||||
text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(target_device)
|
||||
else:
|
||||
# Two-step approach matching Diffusers: format with chat
|
||||
# template first, then tokenize the resulting strings.
|
||||
@@ -278,9 +260,9 @@ class TextEncodingStage(PipelineStage):
|
||||
enable_thinking=False,
|
||||
)
|
||||
formatted_texts.append(formatted)
|
||||
text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(encoder_device)
|
||||
text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(target_device)
|
||||
else:
|
||||
text_inputs = tok(processed_texts, **tok_kwargs).to(encoder_device)
|
||||
text_inputs = tok(processed_texts, **tok_kwargs).to(target_device)
|
||||
|
||||
input_ids = text_inputs["input_ids"]
|
||||
attention_mask = text_inputs["attention_mask"]
|
||||
@@ -300,18 +282,15 @@ class TextEncodingStage(PipelineStage):
|
||||
except Exception:
|
||||
prompt_embeds, attention_mask = postprocess_func(outputs, attention_mask)
|
||||
if is_ltx2 and getattr(outputs, "hidden_states", None):
|
||||
audio_embed = outputs.hidden_states[0].to(device=target_device)
|
||||
audio_embed = outputs.hidden_states[0]
|
||||
if dtype is not None:
|
||||
audio_embed = audio_embed.to(dtype=dtype)
|
||||
audio_embeds_list.append(audio_embed)
|
||||
prompt_embeds = prompt_embeds.to(device=target_device)
|
||||
if dtype is not None:
|
||||
prompt_embeds = prompt_embeds.to(dtype=dtype)
|
||||
embeds_list.append(prompt_embeds)
|
||||
if return_attention_mask:
|
||||
attn_masks_list.append(attention_mask.to(device=target_device))
|
||||
if moved_for_forward and fastvideo_args.text_encoder_cpu_offload:
|
||||
text_encoder.to("cpu")
|
||||
attn_masks_list.append(attention_mask)
|
||||
self._last_audio_embeds = audio_embeds_list if is_ltx2 else None
|
||||
return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices)
|
||||
|
||||
@@ -371,7 +350,6 @@ class Cosmos25TextEncodingStage(PipelineStage):
|
||||
Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's
|
||||
`compute_text_embeddings_online()`.
|
||||
"""
|
||||
performance_component_metric = "text_encoder_time_s"
|
||||
|
||||
def __init__(self, text_encoder) -> None:
|
||||
super().__init__()
|
||||
|
||||
@@ -219,16 +219,6 @@ class StageValidators:
|
||||
|
||||
return validator
|
||||
|
||||
@staticmethod
|
||||
def min_list_length(min_length: int) -> Callable[[Any], bool]:
|
||||
"""Return a validator that checks if a list has at least min_length items."""
|
||||
|
||||
def validator(value: Any) -> bool:
|
||||
return StageValidators.list_min_length(value, min_length)
|
||||
|
||||
validator.__name__ = f"list_min_length_{min_length}"
|
||||
return validator
|
||||
|
||||
@staticmethod
|
||||
def min_dims(min_dims: int) -> Callable[[Any], bool]:
|
||||
"""Return a validator that checks if tensor has at least min_dims dimensions and no NaN values."""
|
||||
|
||||
@@ -157,14 +157,6 @@ class CudaPlatformBase(Platform):
|
||||
"ATTN_QAT_TRAIN selected but fastvideo_kernel.triton_kernels.attn_qat_train is not built. "
|
||||
"Silent fallback would produce a non-QAT training run; refusing to proceed. "
|
||||
"Install the training kernel or pick a different FASTVIDEO_ATTENTION_BACKEND.")
|
||||
elif selected_backend == AttentionBackendEnum.NABLA_ATTN:
|
||||
from fastvideo.attention.backends.nabla import CAN_USE_FLEX_ATTN
|
||||
if CAN_USE_FLEX_ATTN:
|
||||
logger.info("Using NABLA block-sparse flex-attention backend.")
|
||||
return "fastvideo.attention.backends.nabla.NablaAttentionBackend"
|
||||
raise ImportError("NABLA_ATTN selected but torch.nn.attention.flex_attention is unavailable in this "
|
||||
"PyTorch build. Silent fallback to dense attention would be orders of magnitude "
|
||||
"slower and diverge from the reference; upgrade PyTorch or pick a different backend.")
|
||||
elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
|
||||
try:
|
||||
from fastvideo_kernel import video_sparse_attn # noqa: F401
|
||||
|
||||
@@ -22,7 +22,6 @@ class AttentionBackendEnum(enum.Enum):
|
||||
VMOBA_ATTN = enum.auto()
|
||||
SLA_ATTN = enum.auto()
|
||||
SAGE_SLA_ATTN = enum.auto()
|
||||
NABLA_ATTN = enum.auto()
|
||||
NO_ATTENTION = enum.auto()
|
||||
|
||||
|
||||
|
||||
+4
-217
@@ -20,7 +20,6 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
|
||||
Cosmos25Config,
|
||||
Cosmos25_14BConfig,
|
||||
)
|
||||
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
|
||||
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
|
||||
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
|
||||
from fastvideo.configs.pipelines.gen3c import Gen3CConfig
|
||||
@@ -28,7 +27,6 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu
|
||||
Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig,
|
||||
Hunyuan15SR1080PConfig)
|
||||
from fastvideo.configs.pipelines.hyworld import HYWorldConfig
|
||||
from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
|
||||
from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
|
||||
from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
|
||||
from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
|
||||
@@ -119,10 +117,6 @@ class ConfigInfo:
|
||||
workload_types: tuple[WorkloadType, ...]
|
||||
model_family: str | None = None
|
||||
default_preset: str | None = None
|
||||
# When set, overrides the model_index `_class_name` for pipeline resolution.
|
||||
# Lets a model family map to a specific pipeline class by path/detector
|
||||
# (e.g. a T2V and I2V checkpoint that share a `_class_name`).
|
||||
pipeline_cls_name: str | None = None
|
||||
|
||||
|
||||
# The central registry mapping a model name to its configuration information
|
||||
@@ -143,7 +137,6 @@ def register_configs(
|
||||
model_detectors: list[Callable[[str], bool]] | None = None,
|
||||
model_family: str | None = None,
|
||||
default_preset: str | None = None,
|
||||
pipeline_cls_name: str | None = None,
|
||||
) -> None:
|
||||
"""Register config classes for a model family.
|
||||
|
||||
@@ -158,7 +151,6 @@ def register_configs(
|
||||
workload_types=workload_types,
|
||||
model_family=model_family,
|
||||
default_preset=default_preset,
|
||||
pipeline_cls_name=pipeline_cls_name,
|
||||
)
|
||||
|
||||
if hf_model_paths:
|
||||
@@ -498,174 +490,18 @@ def _register_configs() -> None:
|
||||
default_preset="lingbotworld_i2v",
|
||||
)
|
||||
|
||||
def _kandinsky5_detector(require: tuple[str, ...] = (), exclude: tuple[str, ...] = ()) -> Callable[[str], bool]:
|
||||
|
||||
def detect(path: str) -> bool:
|
||||
path_lower = path.lower()
|
||||
if "kandinsky5" not in path_lower and "kandinsky-5" not in path_lower:
|
||||
return False
|
||||
return (all(token in path_lower for token in require) and not any(token in path_lower for token in exclude))
|
||||
|
||||
return detect
|
||||
|
||||
# t2v/i2v exclude each other so a checkpoint stored under a directory
|
||||
# containing the other token (e.g. ~/i2v_experiments/kandinsky5-t2v-ft)
|
||||
# falls through to the model_index _class_name fallback detectors below
|
||||
# instead of being misrouted.
|
||||
_is_kandinsky5_t2v = _kandinsky5_detector(require=("t2v", ), exclude=("i2v", ))
|
||||
_is_kandinsky5_i2v = _kandinsky5_detector(require=("i2v", ), exclude=("t2v", ))
|
||||
_is_kandinsky5_t2v_lite = _kandinsky5_detector(require=("t2v", "lite"), exclude=("i2v", "distilled"))
|
||||
_is_kandinsky5_t2v_pro = _kandinsky5_detector(require=("t2v", "pro"), exclude=("i2v", "distilled"))
|
||||
_is_kandinsky5_t2v_lite_distilled = _kandinsky5_detector(require=("t2v", "lite", "distilled"), exclude=("i2v", ))
|
||||
_is_kandinsky5_t2v_pro_distilled = _kandinsky5_detector(require=("t2v", "pro", "distilled"), exclude=("i2v", ))
|
||||
_is_kandinsky5_i2v_lite = _kandinsky5_detector(require=("i2v", "lite"), exclude=("t2v", "distilled"))
|
||||
_is_kandinsky5_i2v_pro = _kandinsky5_detector(require=("i2v", "pro"), exclude=("t2v", "distilled"))
|
||||
_is_kandinsky5_i2v_lite_distilled = _kandinsky5_detector(require=("i2v", "lite", "distilled"), exclude=("t2v", ))
|
||||
_is_kandinsky5_i2v_pro_distilled = _kandinsky5_detector(require=("i2v", "pro", "distilled"), exclude=("t2v", ))
|
||||
|
||||
# Kandinsky5 Lite T2V
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5T2VConfig,
|
||||
pipeline_config_cls=PipelineConfig,
|
||||
workload_types=(WorkloadType.T2V, ),
|
||||
hf_model_paths=[
|
||||
"kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
|
||||
],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_t2v_lite,
|
||||
lambda path: any(token in path.lower() for token in ("kandinsky5", "kandinsky-5")),
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_t2v_lite_5s",
|
||||
pipeline_cls_name="Kandinsky5T2VPipeline",
|
||||
)
|
||||
|
||||
# Kandinsky5 Pro T2V
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5T2VConfig,
|
||||
workload_types=(WorkloadType.T2V, ),
|
||||
hf_model_paths=[
|
||||
"kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers",
|
||||
],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_t2v_pro,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_t2v_pro_5s",
|
||||
)
|
||||
|
||||
# Kandinsky5 Lite T2V Distilled
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5T2VConfig,
|
||||
workload_types=(WorkloadType.T2V, ),
|
||||
hf_model_paths=[
|
||||
"kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers",
|
||||
],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_t2v_lite_distilled,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_t2v_lite_distilled_5s",
|
||||
)
|
||||
|
||||
# Kandinsky5 Pro T2V Distilled
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5T2VConfig,
|
||||
workload_types=(WorkloadType.T2V, ),
|
||||
hf_model_paths=[
|
||||
"kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers",
|
||||
],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_t2v_pro_distilled,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_t2v_pro_distilled_5s",
|
||||
)
|
||||
|
||||
# Kandinsky5 Lite I2V
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5I2VConfig,
|
||||
workload_types=(WorkloadType.I2V, ),
|
||||
hf_model_paths=["kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_i2v_lite,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_i2v_lite_5s",
|
||||
pipeline_cls_name="Kandinsky5I2VPipeline",
|
||||
)
|
||||
|
||||
# Kandinsky5 Pro I2V
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5I2VConfig,
|
||||
workload_types=(WorkloadType.I2V, ),
|
||||
hf_model_paths=["kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_i2v_pro,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_i2v_pro_5s",
|
||||
)
|
||||
|
||||
# Kandinsky5 Pro I2V Distilled
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5I2VConfig,
|
||||
workload_types=(WorkloadType.I2V, ),
|
||||
hf_model_paths=["kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers"],
|
||||
model_detectors=[
|
||||
_is_kandinsky5_i2v_pro_distilled,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_i2v_pro_distilled_5s",
|
||||
)
|
||||
|
||||
# Kandinsky5 Lite I2V Distilled (no official hub repo yet; local
|
||||
# conversions get distilled sampling defaults instead of the sft ones the
|
||||
# fallback would apply).
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5I2VConfig,
|
||||
workload_types=(),
|
||||
model_detectors=[
|
||||
_is_kandinsky5_i2v_lite_distilled,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_i2v_lite_distilled_5s",
|
||||
pipeline_cls_name="Kandinsky5I2VPipeline",
|
||||
)
|
||||
|
||||
# Kandinsky5 fallbacks — registered AFTER the variant detectors so those
|
||||
# win first-match. Catch checkpoints the variant detectors cannot resolve:
|
||||
# token-less local paths matched via the model_index _class_name
|
||||
# ("kandinsky5t2vpipeline" carries no lite/pro marker), variant combos
|
||||
# without a dedicated entry (e.g. I2V Lite distilled), and t2v+i2v
|
||||
# ambiguous paths resolved by the checkpoint's _class_name.
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5T2VConfig,
|
||||
workload_types=(),
|
||||
model_detectors=[
|
||||
_is_kandinsky5_t2v,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_t2v_lite_5s",
|
||||
pipeline_cls_name="Kandinsky5T2VPipeline",
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=Kandinsky5I2VConfig,
|
||||
workload_types=(),
|
||||
model_detectors=[
|
||||
_is_kandinsky5_i2v,
|
||||
],
|
||||
model_family="kandinsky5",
|
||||
default_preset="kandinsky5_i2v_lite_5s",
|
||||
pipeline_cls_name="Kandinsky5I2VPipeline",
|
||||
)
|
||||
|
||||
# LongCat (T2V, I2V, VC use same config; workload varies by path)
|
||||
@@ -714,7 +550,7 @@ def _register_configs() -> None:
|
||||
"FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers",
|
||||
"FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers",
|
||||
"FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers",
|
||||
# Legacy HF paths (kept for backward compat - pre-rename names):
|
||||
# Legacy HF paths (kept for backward compat — pre-rename names):
|
||||
"FastVideo/Matrix-Game-2.0-Base-Diffusers",
|
||||
"FastVideo/Matrix-Game-2.0-GTA-Diffusers",
|
||||
"FastVideo/Matrix-Game-2.0-TempleRun-Diffusers",
|
||||
@@ -937,40 +773,6 @@ def _register_configs() -> None:
|
||||
model_family="wan",
|
||||
default_preset="wan_2_2_ti2v_5b",
|
||||
)
|
||||
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=DreamXWorld5BCamPipelineConfig,
|
||||
workload_types=(WorkloadType.I2V, ),
|
||||
hf_model_paths=[
|
||||
"FastVideo/DreamX-World-5B-Cam-Diffusers",
|
||||
],
|
||||
model_detectors=[
|
||||
# Mutually exclusive with the AR detector below: Cam requires an
|
||||
# explicit "cam" marker so hyphenated AR local paths (e.g.
|
||||
# /ckpts/dreamx-world-5b-converted) don't first-match here —
|
||||
# detector resolution is first-match in registration order.
|
||||
lambda path:
|
||||
("dreamx-world" in path.lower() and "cam" in path.lower()) or "dreamxworldpipeline" in path.lower()
|
||||
],
|
||||
model_family="dreamx_world",
|
||||
default_preset="dreamx_world_5b_cam",
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=DreamXWorld5BARPipelineConfig,
|
||||
workload_types=(WorkloadType.I2V, ),
|
||||
hf_model_paths=[
|
||||
"FastVideo/DreamX-World-5B-Diffusers",
|
||||
],
|
||||
model_detectors=[
|
||||
lambda path:
|
||||
("dreamx-world-5b" in path.lower() and "cam" not in path.lower()) or "dreamxworldarpipeline" in path.lower(
|
||||
)
|
||||
],
|
||||
model_family="dreamx_world",
|
||||
default_preset="dreamx_world_5b_ar",
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=None,
|
||||
pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
|
||||
@@ -1100,10 +902,7 @@ def get_model_info(
|
||||
assert config_info is not None, "config_info must be resolved"
|
||||
|
||||
if override_pipeline_cls_name:
|
||||
# Explicit override: skip config resolution entirely so checkpoints
|
||||
# without a diffusers model_index.json keep working (and no download
|
||||
# is triggered just to log the replaced name).
|
||||
pipeline_name: str | None = override_pipeline_cls_name
|
||||
pipeline_name = override_pipeline_cls_name
|
||||
logger.info("Using override pipeline class name %s", pipeline_name)
|
||||
else:
|
||||
if os.path.exists(model_path):
|
||||
@@ -1112,12 +911,6 @@ def get_model_info(
|
||||
config = maybe_download_model_index(model_path)
|
||||
|
||||
pipeline_name = config.get("_class_name")
|
||||
if config_info.pipeline_cls_name is not None:
|
||||
# The resolved (path/detector-based) config pins the pipeline class,
|
||||
# e.g. an I2V checkpoint whose `_class_name` would otherwise resolve
|
||||
# to the T2V pipeline.
|
||||
logger.info("Pinning pipeline class name from %s to %s", pipeline_name, config_info.pipeline_cls_name)
|
||||
pipeline_name = config_info.pipeline_cls_name
|
||||
|
||||
if pipeline_name is None:
|
||||
raise ValueError("Model config does not contain a _class_name attribute. "
|
||||
@@ -1158,8 +951,6 @@ def _register_presets() -> None:
|
||||
from fastvideo.api.presets import register_preset
|
||||
from fastvideo.pipelines.basic.cosmos.presets import (
|
||||
ALL_PRESETS as COSMOS_PRESETS, )
|
||||
from fastvideo.pipelines.basic.dreamx_world.presets import (
|
||||
ALL_PRESETS as DREAMX_WORLD_PRESETS, )
|
||||
from fastvideo.pipelines.basic.gamecraft.presets import (
|
||||
ALL_PRESETS as GAMECRAFT_PRESETS, )
|
||||
from fastvideo.pipelines.basic.gen3c.presets import (
|
||||
@@ -1170,8 +961,6 @@ def _register_presets() -> None:
|
||||
ALL_PRESETS as HUNYUAN15_PRESETS, )
|
||||
from fastvideo.pipelines.basic.hyworld.presets import (
|
||||
ALL_PRESETS as HYWORLD_PRESETS, )
|
||||
from fastvideo.pipelines.basic.kandinsky5.presets import (
|
||||
ALL_PRESETS as KANDINSKY5_PRESETS, )
|
||||
from fastvideo.pipelines.basic.lingbotworld.presets import (
|
||||
ALL_PRESETS as LINGBOTWORLD_PRESETS, )
|
||||
from fastvideo.pipelines.basic.longcat.presets import (
|
||||
@@ -1195,14 +984,12 @@ def _register_presets() -> None:
|
||||
|
||||
all_preset_groups = (
|
||||
COSMOS_PRESETS,
|
||||
DREAMX_WORLD_PRESETS,
|
||||
FLUX2_PRESETS,
|
||||
GAMECRAFT_PRESETS,
|
||||
GEN3C_PRESETS,
|
||||
HUNYUAN_PRESETS,
|
||||
HUNYUAN15_PRESETS,
|
||||
HYWORLD_PRESETS,
|
||||
KANDINSKY5_PRESETS,
|
||||
LINGBOTWORLD_PRESETS,
|
||||
LONGCAT_PRESETS,
|
||||
LTX2_PRESETS,
|
||||
|
||||
@@ -17,6 +17,46 @@ class TestLegacyTorchCompileKwargsTranslation:
|
||||
first-class :class:`CompileConfig` fields and anything unknown falls
|
||||
into ``extras``."""
|
||||
|
||||
def test_all_typed_keys_promoted(self) -> None:
|
||||
config = legacy_from_pretrained_to_config(
|
||||
"/models/ltx2",
|
||||
{
|
||||
"enable_torch_compile": True,
|
||||
"torch_compile_kwargs": {
|
||||
"backend": "inductor",
|
||||
"fullgraph": True,
|
||||
"mode": "max-autotune-no-cudagraphs",
|
||||
"dynamic": False,
|
||||
},
|
||||
},
|
||||
)
|
||||
compile_config = config.engine.compile
|
||||
assert compile_config.enabled is True
|
||||
assert compile_config.backend == "inductor"
|
||||
assert compile_config.fullgraph is True
|
||||
assert compile_config.mode == "max-autotune-no-cudagraphs"
|
||||
assert compile_config.dynamic is False
|
||||
assert compile_config.extras == {}
|
||||
|
||||
def test_unknown_keys_land_in_extras(self) -> None:
|
||||
config = legacy_from_pretrained_to_config(
|
||||
"/models/ltx2",
|
||||
{
|
||||
"enable_torch_compile": True,
|
||||
"torch_compile_kwargs": {
|
||||
"backend": "inductor",
|
||||
"options": {"triton.cudagraphs": False},
|
||||
"disable": False,
|
||||
},
|
||||
},
|
||||
)
|
||||
compile_config = config.engine.compile
|
||||
assert compile_config.backend == "inductor"
|
||||
assert compile_config.extras == {
|
||||
"options": {"triton.cudagraphs": False},
|
||||
"disable": False,
|
||||
}
|
||||
|
||||
def test_empty_kwargs_produces_empty_extras(self) -> None:
|
||||
config = legacy_from_pretrained_to_config(
|
||||
"/models/ltx2",
|
||||
@@ -77,6 +117,13 @@ class TestLegacyLtx2VaeTilingTranslation:
|
||||
``generator.pipeline.vae_tiling``; reverse direction emits the
|
||||
legacy name back to FastVideoArgs."""
|
||||
|
||||
def test_forward_routes_to_pipeline_vae_tiling(self) -> None:
|
||||
config = legacy_from_pretrained_to_config(
|
||||
"/models/ltx2",
|
||||
{"ltx2_vae_tiling": False},
|
||||
)
|
||||
assert config.pipeline.vae_tiling is False
|
||||
|
||||
def test_true_round_trips(self) -> None:
|
||||
config = legacy_from_pretrained_to_config(
|
||||
"/models/ltx2",
|
||||
|
||||
@@ -1,43 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CPU-only regression test for the SDPA metadata mask contract.
|
||||
|
||||
SDPAMetadata is cross-backend: HYWorld and HunyuanVideo15 build it while the
|
||||
layer's attention selector may pick FLASH_ATTN, and the shared convention
|
||||
for padding masks is the tokenizer-style 2D [batch, key_len]. The builder
|
||||
must store 2D masks unchanged; only the torch-sdpa impl lifts them to 4D
|
||||
internally.
|
||||
|
||||
Regression guard for the CI break where the builder lifted 2D -> 4D and the
|
||||
FLASH_ATTN equal-length branch padded ``attn_mask.shape[1]`` (== 1 for 4D)
|
||||
into a garbage-length mask (CUDA device-side assert, SIGABRT).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.attention.backends.sdpa import (SDPAMetadataBuilder, _normalize_attn_mask_for_sdpa)
|
||||
|
||||
|
||||
def test_builder_keeps_2d_padding_mask_2d() -> None:
|
||||
"""A tokenizer-style 2D int mask must be stored as-is (cross-backend contract)."""
|
||||
mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.int64)
|
||||
md = SDPAMetadataBuilder().build(current_timestep=0, attn_mask=mask)
|
||||
assert md.attn_mask is not None
|
||||
assert md.attn_mask.dim() == 2
|
||||
assert md.attn_mask.shape == (2, 3)
|
||||
assert torch.equal(md.attn_mask, mask)
|
||||
|
||||
|
||||
def test_normalize_lifts_2d_int_mask_to_4d_bool_for_torch_sdpa() -> None:
|
||||
"""The torch-sdpa consumer coerces int -> bool and lifts 2D -> [B,1,1,K]."""
|
||||
batch, q_len, k_len, heads, head_dim = 2, 4, 3, 1, 8
|
||||
query = torch.zeros(batch, heads, q_len, head_dim)
|
||||
key = torch.zeros(batch, heads, k_len, head_dim)
|
||||
mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.int64)
|
||||
|
||||
out = _normalize_attn_mask_for_sdpa(mask, query, key)
|
||||
assert out is not None
|
||||
assert out.dtype == torch.bool
|
||||
assert out.shape == (batch, 1, 1, k_len)
|
||||
assert torch.equal(out[:, 0, 0, :], mask != 0)
|
||||
@@ -1,84 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Guard: every test directory must be collected by some CI lane or be on
|
||||
the explicit allowlist below.
|
||||
|
||||
Three separate incidents on 2026-07-05 found test files that no CI lane
|
||||
ever collects (fastvideo/tests/stages/, tests/local_tests/ additions in
|
||||
PR #1509, and this sweep found seven dark directories in total): the tests
|
||||
pass review, merge, and then silently never run. This test makes going
|
||||
dark an explicit, reviewed decision instead of an accident: adding a new
|
||||
test directory fails CI until it is either wired into a lane or
|
||||
allowlisted here with a reason.
|
||||
|
||||
Pure text analysis — no fastvideo imports, no GPU, no torch.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
TESTS_ROOT = REPO_ROOT / "fastvideo" / "tests"
|
||||
|
||||
# Files whose text constitutes "a CI lane references this directory".
|
||||
CI_SOURCES = [
|
||||
TESTS_ROOT / "modal" / "pr_test.py",
|
||||
TESTS_ROOT / "modal" / "ssim_test.py",
|
||||
*sorted((REPO_ROOT / ".buildkite").rglob("*.yml")),
|
||||
*sorted((REPO_ROOT / ".buildkite").rglob("*.sh")),
|
||||
]
|
||||
|
||||
# Directories that intentionally have no CI lane today. Every entry needs a
|
||||
# reason; remove the entry when the directory gets wired into a lane.
|
||||
# State as found on 2026-07-05 — these SHOULD shrink over time, not grow.
|
||||
ALLOWLIST = {
|
||||
"attention": "no lane yet — GPU attention-backend tests, run manually",
|
||||
"audio": "no lane yet — audio encoder tests, run manually",
|
||||
"distributed": "no lane yet — multi-GPU torchrun tests, run manually",
|
||||
"hooks": "no lane yet — run manually",
|
||||
"layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
|
||||
"nightly": "by design: nightly cadence, not per-PR",
|
||||
"modal": "CI infrastructure itself, not a test suite",
|
||||
}
|
||||
|
||||
|
||||
def _dirs_with_tests() -> list[str]:
|
||||
dirs = []
|
||||
for child in sorted(TESTS_ROOT.iterdir()):
|
||||
if child.is_dir() and any(child.rglob("test_*.py")):
|
||||
dirs.append(child.name)
|
||||
return dirs
|
||||
|
||||
|
||||
def _ci_text() -> str:
|
||||
return "\n".join(
|
||||
src.read_text(errors="replace") for src in CI_SOURCES if src.exists())
|
||||
|
||||
|
||||
def test_every_test_directory_is_collected_or_allowlisted():
|
||||
ci_text = _ci_text()
|
||||
dark = [
|
||||
name for name in _dirs_with_tests()
|
||||
if f"tests/{name}" not in ci_text and name not in ALLOWLIST
|
||||
]
|
||||
assert not dark, (
|
||||
f"Test directories not referenced by any CI lane and not "
|
||||
f"allowlisted: {dark}. Wire them into a lane in "
|
||||
f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an "
|
||||
f"allowlist entry with a reason in {__file__}.")
|
||||
|
||||
|
||||
def test_local_tests_stays_out_of_ci():
|
||||
# tests/local_tests/ (repo root) is developer-local by design (author
|
||||
# decision, 2026-07-05): parity scaffolds and machine-specific checks
|
||||
# that must never gate CI. Fail if any CI source starts collecting it.
|
||||
assert "tests/local_tests" not in _ci_text(), (
|
||||
"tests/local_tests/ is local-only by design; remove the CI "
|
||||
"reference or move the tests into a fastvideo/tests/ lane.")
|
||||
|
||||
|
||||
def test_allowlist_entries_are_still_real_directories():
|
||||
# A stale allowlist hides regressions; entries must track reality.
|
||||
missing = [
|
||||
name for name in ALLOWLIST
|
||||
if name != "modal" and not (TESTS_ROOT / name).is_dir()
|
||||
]
|
||||
assert not missing, (
|
||||
f"Allowlisted directories no longer exist — remove them: {missing}")
|
||||
@@ -1,58 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Guard the Modal FA4 defaults that keep CI lanes on their intended backend.
|
||||
|
||||
Pure text/AST analysis: no fastvideo imports, no torch, no Modal client.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
MODAL_ROOT = REPO_ROOT / "fastvideo" / "tests" / "modal"
|
||||
PR_TEST = MODAL_ROOT / "pr_test.py"
|
||||
LAUNCH_L40S_JOB = MODAL_ROOT / "launch_l40s_job.py"
|
||||
SSIM_TEST = MODAL_ROOT / "ssim_test.py"
|
||||
|
||||
|
||||
def _function_strings(path: Path, function_name: str) -> str:
|
||||
source = path.read_text(encoding="utf-8")
|
||||
tree = ast.parse(source)
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.FunctionDef) and node.name == function_name:
|
||||
return "\n".join(
|
||||
child.value
|
||||
for child in ast.walk(node)
|
||||
if isinstance(child, ast.Constant)
|
||||
and isinstance(child.value, str)
|
||||
)
|
||||
raise AssertionError(f"{function_name} not found in {path}")
|
||||
|
||||
|
||||
def test_generic_l40s_launcher_defaults_fa4_off():
|
||||
source = LAUNCH_L40S_JOB.read_text(encoding="utf-8")
|
||||
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0")' in source
|
||||
|
||||
|
||||
def test_ssim_launcher_keeps_fa4_enabled_by_default():
|
||||
source = SSIM_TEST.read_text(encoding="utf-8")
|
||||
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1")' in source
|
||||
|
||||
|
||||
def test_pr_model_load_and_training_lanes_disable_fa4():
|
||||
lanes = {
|
||||
"run_transformer_tests": "pytest ./fastvideo/tests/transformers -vs",
|
||||
"run_training_tests": "pytest ./fastvideo/tests/training/Vanilla -srP",
|
||||
"run_training_lora_tests": "pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP",
|
||||
"run_training_tests_VSA": "pytest ./fastvideo/tests/training/VSA -srP",
|
||||
"run_distill_dmd_tests": "pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs",
|
||||
"run_self_forcing_tests": "pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs",
|
||||
"run_train_framework_tests": "pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs",
|
||||
"seed_grad_norm_references": "pytest ./fastvideo/tests/train/methods -vs -rs",
|
||||
}
|
||||
|
||||
for function_name, pytest_command in lanes.items():
|
||||
function_strings = _function_strings(PR_TEST, function_name)
|
||||
assert "FASTVIDEO_FA4=0" in function_strings
|
||||
assert pytest_command in function_strings
|
||||
@@ -1,80 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CPU-only tests for load_image URL retry behavior.
|
||||
|
||||
A HF-CDN IncompleteRead mid-body previously killed SSIM runs (build 4202,
|
||||
test_wan_i2v) because load_image handed PIL a partial ``.raw`` stream with no
|
||||
retry. These tests mock requests.get: a transient failure must be retried and
|
||||
a full-body success must yield a usable PIL image.
|
||||
"""
|
||||
import io
|
||||
from unittest import mock
|
||||
|
||||
import PIL.Image
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from fastvideo.models import vision_utils
|
||||
|
||||
URL = "https://example.com/image.png"
|
||||
|
||||
|
||||
def _png_bytes() -> bytes:
|
||||
buf = io.BytesIO()
|
||||
PIL.Image.new("RGB", (4, 4), (255, 0, 0)).save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _good_response() -> mock.Mock:
|
||||
response = mock.Mock()
|
||||
response.content = _png_bytes()
|
||||
response.raise_for_status = mock.Mock()
|
||||
return response
|
||||
|
||||
|
||||
def test_load_image_retries_transient_error_then_succeeds(monkeypatch):
|
||||
monkeypatch.setattr(vision_utils.time, "sleep", lambda _s: None)
|
||||
get = mock.Mock(side_effect=[
|
||||
# The shape that broke build 4202: connection dropped mid-body.
|
||||
requests.exceptions.ChunkedEncodingError(
|
||||
"IncompleteRead(1024 bytes read, 2048 more expected)"),
|
||||
_good_response(),
|
||||
])
|
||||
monkeypatch.setattr(vision_utils.requests, "get", get)
|
||||
|
||||
image = vision_utils.load_image(URL)
|
||||
|
||||
assert get.call_count == 2
|
||||
assert get.call_args.kwargs["timeout"] == 30.0
|
||||
assert isinstance(image, PIL.Image.Image)
|
||||
assert image.size == (4, 4)
|
||||
assert image.mode == "RGB"
|
||||
|
||||
|
||||
def test_load_image_raises_after_exhausting_retries(monkeypatch):
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(vision_utils.time, "sleep", sleeps.append)
|
||||
get = mock.Mock(
|
||||
side_effect=requests.exceptions.ConnectionError("connection reset"))
|
||||
monkeypatch.setattr(vision_utils.requests, "get", get)
|
||||
|
||||
with pytest.raises(requests.exceptions.ConnectionError):
|
||||
vision_utils.load_image(URL)
|
||||
|
||||
assert get.call_count == 3
|
||||
assert sleeps == [1, 2]
|
||||
|
||||
|
||||
def test_load_image_does_not_retry_on_client_error(monkeypatch):
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(vision_utils.time, "sleep", sleeps.append)
|
||||
response = mock.Mock(status_code=404)
|
||||
error = requests.HTTPError("404 Client Error", response=response)
|
||||
response.raise_for_status.side_effect = error
|
||||
get = mock.Mock(return_value=response)
|
||||
monkeypatch.setattr(vision_utils.requests, "get", get)
|
||||
|
||||
with pytest.raises(requests.HTTPError):
|
||||
vision_utils.load_image("https://example.com/missing.png")
|
||||
|
||||
assert get.call_count == 1
|
||||
assert sleeps == []
|
||||
@@ -12,7 +12,6 @@ from fastvideo.configs.pipelines import CosmosConfig, PipelineConfig, WanT2V480P
|
||||
from fastvideo.forward_context import set_forward_context
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.loader.component_loader import TextEncoderLoader
|
||||
from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
|
||||
from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.configs.models.encoders import T5Config, T5LargeConfig
|
||||
@@ -37,11 +36,9 @@ def t5_model_paths_and_config():
|
||||
@pytest.fixture
|
||||
def t5_large_model_paths_and_config():
|
||||
base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World"
|
||||
local_dir = os.path.join('data', base_model_path)
|
||||
skip_if_gated_repo_inaccessible(base_model_path,
|
||||
local_path=local_dir,
|
||||
test_name="Cosmos T5-large encoder test")
|
||||
model_path = maybe_download_model(base_model_path, local_dir=local_dir)
|
||||
model_path = maybe_download_model(base_model_path,
|
||||
local_dir=os.path.join(
|
||||
'data', base_model_path))
|
||||
text_encoder_path = os.path.join(model_path, "text_encoder")
|
||||
tokenizer_path = os.path.join(model_path, "tokenizer")
|
||||
return text_encoder_path, tokenizer_path, CosmosConfig()
|
||||
|
||||
@@ -97,10 +97,9 @@ image = (
|
||||
"TOKENIZERS_PARALLELISM": "false",
|
||||
**({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
|
||||
"FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4). Generic ad hoc jobs should follow the
|
||||
# product default unless a caller opts in through the local env or
|
||||
# --env-vars.
|
||||
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0"),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4); keep CI parity with the seeded
|
||||
# references. Caller override wins.
|
||||
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
|
||||
})
|
||||
)
|
||||
|
||||
|
||||
@@ -30,15 +30,6 @@ print(f"Using image: {image_ref}")
|
||||
# the baked setting, and a caller override always wins.
|
||||
uv_torch_backend_override = resolve_uv_torch_backend(image_tag)
|
||||
|
||||
# INVARIANT: this image definition must be byte-identical for every CI job at
|
||||
# a given base image digest -- one build, shared cache across all concurrent
|
||||
# lanes. Never put a per-job/per-commit value (BUILDKITE_*, TEST_SCOPE,
|
||||
# env-derived overrides) into the image via .env()/run_commands: it becomes an
|
||||
# image layer, so whenever the base digest changes every concurrent job
|
||||
# rebuilds its own image variant (~15-20 min each), blowing the Buildkite job
|
||||
# budget. `image_ref` is the only env-derived input allowed here, because it
|
||||
# *selects* the base digest. Per-job values arrive at runtime via
|
||||
# `ci_env_secret` below.
|
||||
image = (modal.Image.from_registry(
|
||||
image_ref, add_python="3.12"
|
||||
).run_commands("rm -rf /FastVideo").apt_install(
|
||||
@@ -46,45 +37,41 @@ image = (modal.Image.from_registry(
|
||||
).run_commands(
|
||||
"curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable"
|
||||
).run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({
|
||||
"PATH": "/root/.cargo/bin:$PATH",
|
||||
"HF_REPO_ID": "FastVideo/performance-tracking",
|
||||
"PATH":
|
||||
"/root/.cargo/bin:$PATH",
|
||||
"BUILDKITE_REPO":
|
||||
os.environ.get("BUILDKITE_REPO", ""),
|
||||
"BUILDKITE_COMMIT":
|
||||
os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
"BUILDKITE_PULL_REQUEST":
|
||||
os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"BUILDKITE_BRANCH":
|
||||
os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"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,
|
||||
**({
|
||||
"UV_TORCH_BACKEND": uv_torch_backend_override
|
||||
} if uv_torch_backend_override else {}),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4); CI lanes keep it enabled to match the
|
||||
# SSIM/perf baselines. Caller override wins.
|
||||
"FASTVIDEO_FA4":
|
||||
os.environ.get("FASTVIDEO_FA4", "1"),
|
||||
"HF_REPO_ID":
|
||||
"FastVideo/performance-tracking",
|
||||
}))
|
||||
|
||||
dreamverse_image = (image.run_commands(
|
||||
"curl -fsSL https://deb.nodesource.com/setup_22.x | bash -"
|
||||
).apt_install("nodejs").run_commands("node --version && npm --version"))
|
||||
|
||||
# Per-job/per-invocation values are injected into the container environment at
|
||||
# RUNTIME via this secret (attached to every function below), so the image
|
||||
# stays identical across jobs. Consumers (run_test_command's checkout,
|
||||
# fastvideo/tests/performance/{compare_baseline,identity}.py, the FA4
|
||||
# resolver, `uv pip install`) all read os.environ at runtime, so nothing else
|
||||
# changes.
|
||||
ci_env_secret = modal.Secret.from_dict({
|
||||
"BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
|
||||
"BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
"BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"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,
|
||||
**({
|
||||
"UV_TORCH_BACKEND": uv_torch_backend_override
|
||||
} if uv_torch_backend_override else {}),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4). Keep the default enabled for
|
||||
# inference/perf parity; model-load and training lanes that do not exercise
|
||||
# FA4 explicitly set FASTVIDEO_FA4=0 in their command strings below.
|
||||
# Caller override wins.
|
||||
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
|
||||
})
|
||||
|
||||
hf_secret = modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
wandb_secret = modal.Secret.from_dict(
|
||||
{"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
|
||||
|
||||
|
||||
def run_test(pytest_command: str):
|
||||
"""Helper function to run a test suite with custom pytest command"""
|
||||
@@ -160,7 +147,10 @@ def run_test_command(test_command: str,
|
||||
@app.function(gpu="H100:1",
|
||||
image=image,
|
||||
timeout=1200,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_encoder_tests():
|
||||
run_test(
|
||||
@@ -171,7 +161,10 @@ def run_encoder_tests():
|
||||
@app.function(gpu="L40S:1",
|
||||
image=image,
|
||||
timeout=1200,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_vae_tests():
|
||||
run_test(
|
||||
@@ -182,109 +175,113 @@ def run_vae_tests():
|
||||
@app.function(gpu="L40S:1",
|
||||
image=image,
|
||||
timeout=900,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_transformer_tests():
|
||||
run_test(
|
||||
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
|
||||
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/transformers -vs"
|
||||
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="L40S:4",
|
||||
cpu=8.0,
|
||||
memory=32768,
|
||||
image=image,
|
||||
timeout=900,
|
||||
secrets=[wandb_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_training_tests():
|
||||
run_test(
|
||||
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
|
||||
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/Vanilla -srP"
|
||||
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="L40S:2",
|
||||
cpu=8.0,
|
||||
memory=32768,
|
||||
image=image,
|
||||
timeout=900,
|
||||
secrets=[wandb_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_training_lora_tests():
|
||||
run_test(
|
||||
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
|
||||
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP"
|
||||
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="H100:2",
|
||||
image=image,
|
||||
timeout=900,
|
||||
secrets=[wandb_secret, ci_env_secret])
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
|
||||
])
|
||||
def run_training_tests_VSA():
|
||||
run_test(
|
||||
"wandb login $WANDB_API_KEY && FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/VSA -srP"
|
||||
"wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/VSA -srP"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
@app.function(gpu="H100:1", image=image, timeout=900)
|
||||
def run_kernel_tests():
|
||||
run_test("pytest fastvideo-kernel/tests/ -vs")
|
||||
|
||||
|
||||
# @app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
# @app.function(gpu="H100:1", image=image, timeout=900)
|
||||
# def run_precision_tests_VSA():
|
||||
# # VSA correctness is covered by the same file now
|
||||
# run_test("pytest fastvideo-kernel/tests/test_correctness.py")
|
||||
|
||||
# @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
# @app.function(gpu="L40S:1", image=image, timeout=900)
|
||||
# def run_precision_tests_vmoba():
|
||||
# run_test("pytest fastvideo-kernel/tests/test_vmoba_correctness.py")
|
||||
|
||||
|
||||
@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
@app.function(gpu="L40S:1", image=image, timeout=900)
|
||||
def run_inference_tests_vmoba():
|
||||
run_test('python fastvideo/tests/inference/vmoba/test_vmoba_inference.py')
|
||||
|
||||
|
||||
@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[ci_env_secret])
|
||||
@app.function(gpu="L40S:1", image=image, timeout=1200)
|
||||
def run_inference_lora_tests():
|
||||
run_test(
|
||||
"pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
@app.function(gpu="L40S:2", image=image, timeout=900)
|
||||
def run_distill_dmd_tests():
|
||||
run_test(
|
||||
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
|
||||
"pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
|
||||
|
||||
|
||||
@app.function(gpu="L40S:2",
|
||||
image=image,
|
||||
timeout=900,
|
||||
secrets=[wandb_secret, ci_env_secret])
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
|
||||
])
|
||||
def run_self_forcing_tests():
|
||||
run_test(
|
||||
"wandb login $WANDB_API_KEY && "
|
||||
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
|
||||
"wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
@app.function(gpu="L40S:1", image=image, timeout=900)
|
||||
def run_unit_test():
|
||||
run_test(
|
||||
"pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
|
||||
"pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
|
||||
)
|
||||
|
||||
|
||||
# TODO: David: GPU only used to resolve import time requirement (not needed for this test). Maybe make those imports lazy?
|
||||
@app.function(gpu="L40S:1",
|
||||
image=dreamverse_image,
|
||||
timeout=1800,
|
||||
secrets=[ci_env_secret])
|
||||
@app.function(gpu="L40S:1", image=dreamverse_image, timeout=1800)
|
||||
def run_dreamverse_app_tests():
|
||||
run_test_command(
|
||||
install_command="",
|
||||
@@ -321,23 +318,26 @@ def run_dreamverse_app_tests():
|
||||
|
||||
|
||||
@app.function(gpu="L40S:1",
|
||||
cpu=8.0,
|
||||
memory=32768,
|
||||
image=image,
|
||||
timeout=1800,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_train_framework_tests():
|
||||
run_test(
|
||||
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
|
||||
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs"
|
||||
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="L40S:1",
|
||||
image=image,
|
||||
timeout=1800,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def seed_grad_norm_references():
|
||||
"""Record the per-method grad-norm reference for the **CI GPU (L40S only)**.
|
||||
@@ -356,15 +356,17 @@ def seed_grad_norm_references():
|
||||
the local command and the ``_DEVICE_MAPPINGS`` table.
|
||||
"""
|
||||
run_test(
|
||||
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
|
||||
"FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs"
|
||||
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs"
|
||||
)
|
||||
|
||||
|
||||
@app.function(gpu="L40S:1",
|
||||
image=image,
|
||||
timeout=3600,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_eval_tests():
|
||||
# Eval metric regression: drives the high-level fastvideo.eval API on a
|
||||
@@ -387,7 +389,10 @@ def run_eval_tests():
|
||||
@app.function(gpu="L40S:1",
|
||||
image=image,
|
||||
timeout=3600,
|
||||
secrets=[hf_secret, ci_env_secret])
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
])
|
||||
def run_lora_extraction_tests():
|
||||
run_test(
|
||||
"hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py"
|
||||
@@ -395,11 +400,12 @@ def run_lora_extraction_tests():
|
||||
|
||||
|
||||
@app.function(gpu="L40S:2",
|
||||
cpu=8.0,
|
||||
memory=32768,
|
||||
image=image,
|
||||
timeout=1800,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_performance_tests():
|
||||
# PR/direct records are uploaded only on pass; scheduled main uploads pass
|
||||
@@ -421,8 +427,6 @@ def run_performance_tests():
|
||||
"export PERF_RUN_SOURCE='unknown'; "
|
||||
"export PERF_UPLOAD_POLICY='never'; "
|
||||
"fi; "
|
||||
"(nvidia-smi --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu "
|
||||
"--format=csv -l 10 > /tmp/gpu_telemetry.csv 2>/dev/null &); "
|
||||
"pytest ./fastvideo/tests/performance -vs; "
|
||||
"PYTEST_RC=$?; "
|
||||
"PERF_RC=0; "
|
||||
@@ -431,8 +435,6 @@ def run_performance_tests():
|
||||
"PERF_RC=$?; "
|
||||
"fi; "
|
||||
"python ./fastvideo/tests/performance/dashboard.py || true; "
|
||||
"echo '--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---'; "
|
||||
"cat /tmp/gpu_telemetry.csv || true; "
|
||||
"FINAL_RC=$PYTEST_RC; "
|
||||
"if [ $FINAL_RC -eq 0 ]; then FINAL_RC=$PERF_RC; fi; "
|
||||
"exit $FINAL_RC")
|
||||
@@ -441,7 +443,10 @@ def run_performance_tests():
|
||||
@app.function(gpu="L40S:1",
|
||||
image=image,
|
||||
timeout=1800,
|
||||
secrets=[hf_secret, ci_env_secret],
|
||||
secrets=[
|
||||
modal.Secret.from_dict(
|
||||
{"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
|
||||
],
|
||||
volumes={"/root/data": model_vol})
|
||||
def run_api_server_tests():
|
||||
run_test(
|
||||
|
||||
@@ -57,12 +57,10 @@ image = (
|
||||
.run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
|
||||
.env(
|
||||
{
|
||||
# INVARIANT: keep this image identical for every CI job at a given
|
||||
# base digest -- per-job/per-commit values (BUILDKITE_*) must never
|
||||
# be baked here as image layers or every job rebuilds its own image
|
||||
# variant. repo/commit/PR already reach the container as
|
||||
# run_ssim_partition() arguments.
|
||||
"PATH": "/root/.cargo/bin:$PATH",
|
||||
"BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
|
||||
"BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
"BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"IMAGE_VERSION": image_version,
|
||||
**({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4); the SSIM references were seeded
|
||||
|
||||
@@ -34,9 +34,6 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
|
||||
method._convert_scale(scale)
|
||||
|
||||
def test_create_weights_rejects_invalid_dtype(self):
|
||||
# float64 is outside the supported set (bf16/fp16/fp32). The test
|
||||
# originally passed float32, which create_weights has accepted since
|
||||
# PR #981 — it never ran in CI, so the mismatch went unnoticed.
|
||||
method = AbsMaxFP8LinearMethod()
|
||||
layer = nn.Module()
|
||||
with self.assertRaisesRegex(AssertionError, "only supports"):
|
||||
@@ -46,7 +43,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
|
||||
output_partition_sizes=[3],
|
||||
input_size=2,
|
||||
output_size=3,
|
||||
params_dtype=torch.float64,
|
||||
params_dtype=torch.float32,
|
||||
)
|
||||
|
||||
def test_absmax_fp8_parameter_weight_loader(self):
|
||||
|
||||
@@ -66,7 +66,6 @@ UPLOAD_POLICY = os.environ.get("PERF_UPLOAD_POLICY", "never").strip().lower()
|
||||
VALID_UPLOAD_POLICIES = {"never", "pass", "always"}
|
||||
VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
|
||||
IDENTITY_KEYS = (
|
||||
"result_schema_version",
|
||||
"workload_id",
|
||||
"variant_id",
|
||||
"benchmark_version",
|
||||
@@ -78,8 +77,6 @@ IDENTITY_KEYS = (
|
||||
"software_profile_id",
|
||||
"environment_metadata",
|
||||
"environment_fingerprint",
|
||||
"quality_metadata",
|
||||
"variant_metadata",
|
||||
)
|
||||
COMPARISON_IDENTITY_KEYS = (
|
||||
"workload_id",
|
||||
@@ -146,22 +143,18 @@ def _result_failed_static_thresholds() -> bool:
|
||||
|
||||
|
||||
def _record_metadata(run_source: str, result: dict[str, Any]) -> dict[str, Any]:
|
||||
raw_run_source = str(result.get("run_source") or "").strip().lower()
|
||||
if raw_run_source in VALID_RUN_SOURCES:
|
||||
run_source = raw_run_source
|
||||
|
||||
pr_number = result.get("pr_number") or os.environ.get("BUILDKITE_PULL_REQUEST", "")
|
||||
if not _truthy_pr_number(str(pr_number)):
|
||||
pr_number = ""
|
||||
return {
|
||||
"run_source": run_source,
|
||||
"baseline_eligible": False,
|
||||
"branch": result.get("branch") or os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"branch": os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"pr_number": pr_number,
|
||||
"test_scope": result.get("test_scope") or os.environ.get("TEST_SCOPE", ""),
|
||||
"build_url": result.get("build_url") or os.environ.get("BUILDKITE_BUILD_URL", ""),
|
||||
"build_id": result.get("build_id") or os.environ.get("BUILDKITE_BUILD_ID", ""),
|
||||
"job_id": result.get("job_id") or os.environ.get("BUILDKITE_JOB_ID", ""),
|
||||
"test_scope": os.environ.get("TEST_SCOPE", ""),
|
||||
"build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
|
||||
"build_id": os.environ.get("BUILDKITE_BUILD_ID", ""),
|
||||
"job_id": os.environ.get("BUILDKITE_JOB_ID", ""),
|
||||
}
|
||||
|
||||
|
||||
@@ -185,6 +178,11 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
|
||||
key for key in COMPARISON_IDENTITY_KEYS
|
||||
if key not in record or record[key] is None or (isinstance(record[key], str) and not record[key].strip())
|
||||
]
|
||||
if len(missing) == len(COMPARISON_IDENTITY_KEYS):
|
||||
# Legacy v1 record: no identity fields at all. Keep the pre-v2
|
||||
# (model_id, gpu_type) rolling-baseline comparison instead of
|
||||
# crashing — v1 configs are documented as loadable.
|
||||
return {}
|
||||
if missing:
|
||||
raise ValueError("Performance record missing required comparison identity fields: " + ", ".join(missing))
|
||||
return {
|
||||
@@ -193,18 +191,6 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
|
||||
}
|
||||
|
||||
|
||||
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.")
|
||||
return None
|
||||
|
||||
|
||||
def _format_identity_filters(filters: dict[str, str]) -> str:
|
||||
if not filters:
|
||||
return ""
|
||||
@@ -495,53 +481,43 @@ def main() -> int:
|
||||
for raw in current_results:
|
||||
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"
|
||||
baseline_records = load_records_for_model(
|
||||
TRACKING_ROOT,
|
||||
record["model_id"],
|
||||
record["gpu_type"],
|
||||
**_comparison_identity_filters(record),
|
||||
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(_comparison_identity_filters(record))}")
|
||||
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: list[str] = []
|
||||
record["success"] = True
|
||||
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 = []
|
||||
else:
|
||||
record["baseline_status"] = "compared"
|
||||
failures = _check_regressions(record, baseline_records, metric_policies)
|
||||
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')})")
|
||||
|
||||
record["success"] = not failures
|
||||
record["baseline_eligible"] = (
|
||||
identity_filters is not None and _is_baseline_eligible(record["run_source"], record["success"])
|
||||
)
|
||||
record["baseline_eligible"] = _is_baseline_eligible(record["run_source"], record["success"])
|
||||
all_failures.extend(failures)
|
||||
|
||||
_write_normalized_artifact(record)
|
||||
|
||||
@@ -14,9 +14,9 @@ def _v2_config():
|
||||
return {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
}
|
||||
|
||||
|
||||
@@ -41,9 +41,9 @@ def test_v2_benchmark_config_identity_validates_and_is_preserved():
|
||||
assert _is_v2_config(cfg) is True
|
||||
assert _config_identity_metadata(cfg) == {
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"quality_metadata": {"some": "data"},
|
||||
}
|
||||
|
||||
@@ -91,9 +91,9 @@ def test_v2_benchmark_config_rejects_invalid_benchmark_version_values(value):
|
||||
def test_partial_v2_identity_requires_schema_version():
|
||||
cfg = {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
}
|
||||
|
||||
expected = "wan.json: v2 benchmark identity fields require config_schema_version=2"
|
||||
@@ -105,9 +105,9 @@ 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,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"quality_metadata": ["not", "an", "object"],
|
||||
}
|
||||
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo.tests.performance import compare_baseline
|
||||
@@ -127,7 +125,6 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch):
|
||||
monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
|
||||
raw = _raw_result()
|
||||
raw.update({
|
||||
"result_schema_version": 2,
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
@@ -149,14 +146,10 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch):
|
||||
},
|
||||
},
|
||||
"environment_fingerprint": "env-1",
|
||||
"quality_metadata": {
|
||||
"quality_status": "canonical",
|
||||
},
|
||||
})
|
||||
|
||||
record = compare_baseline.normalize_performance_result(raw)
|
||||
|
||||
assert record["result_schema_version"] == 2
|
||||
assert record["workload_id"] == "wan-t2v"
|
||||
assert record["variant_id"] == "1.3b-sp2"
|
||||
assert record["benchmark_version"] == 2
|
||||
@@ -168,87 +161,6 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch):
|
||||
assert record["software_profile_id"] == "sw-1"
|
||||
assert record["environment_metadata"] == {"env": {"IMAGE_VERSION": "latest"}}
|
||||
assert record["environment_fingerprint"] == "env-1"
|
||||
assert record["quality_metadata"] == {"quality_status": "canonical"}
|
||||
|
||||
|
||||
def test_normalized_record_prefers_raw_v2_provenance(monkeypatch):
|
||||
monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
|
||||
monkeypatch.setenv("BUILDKITE_BRANCH", "feature/from-env")
|
||||
monkeypatch.setenv("TEST_SCOPE", "direct")
|
||||
monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/env")
|
||||
monkeypatch.setenv("BUILDKITE_BUILD_ID", "env-build")
|
||||
monkeypatch.setenv("BUILDKITE_JOB_ID", "env-job")
|
||||
raw = _raw_result()
|
||||
raw.update({
|
||||
"result_schema_version": 2,
|
||||
"run_source": "scheduled_main",
|
||||
"branch": "main",
|
||||
"test_scope": "full",
|
||||
"build_url": "https://buildkite.example/raw",
|
||||
"build_id": "raw-build",
|
||||
"job_id": "raw-job",
|
||||
"pr_number": "false",
|
||||
})
|
||||
|
||||
record = compare_baseline.normalize_performance_result(raw)
|
||||
|
||||
assert record["result_schema_version"] == 2
|
||||
assert record["run_source"] == "scheduled_main"
|
||||
assert record["branch"] == "main"
|
||||
assert record["test_scope"] == "full"
|
||||
assert record["build_url"] == "https://buildkite.example/raw"
|
||||
assert record["build_id"] == "raw-build"
|
||||
assert record["job_id"] == "raw-job"
|
||||
assert record["pr_number"] == ""
|
||||
|
||||
|
||||
def test_v1_normalized_record_has_no_result_schema_version(monkeypatch):
|
||||
monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
|
||||
|
||||
record = compare_baseline.normalize_performance_result(_raw_result())
|
||||
|
||||
assert "result_schema_version" not in record
|
||||
|
||||
|
||||
def test_main_writes_v1_current_artifact_without_comparison_identity(monkeypatch, tmp_path, capsys):
|
||||
results_dir = tmp_path / "results"
|
||||
reports_dir = tmp_path / "reports"
|
||||
tracking_root = tmp_path / "tracking"
|
||||
results_dir.mkdir()
|
||||
(results_dir / "perf_legacy.json").write_text(json.dumps(_raw_result()), encoding="utf-8")
|
||||
uploaded_records = []
|
||||
|
||||
monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main")
|
||||
monkeypatch.delenv("PERF_PYTEST_RC", raising=False)
|
||||
monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir))
|
||||
monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir))
|
||||
monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tracking_root))
|
||||
monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "always")
|
||||
monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir)
|
||||
|
||||
def fail_load_records_for_model(*_args, **_kwargs):
|
||||
raise AssertionError("legacy current records should skip v2 baseline lookup")
|
||||
|
||||
def fake_upload_record(_path, record, *, strict=False):
|
||||
uploaded_records.append(record.copy())
|
||||
|
||||
monkeypatch.setattr(compare_baseline, "load_records_for_model", fail_load_records_for_model)
|
||||
monkeypatch.setattr(compare_baseline, "upload_record", fake_upload_record)
|
||||
|
||||
assert compare_baseline.main() == 0
|
||||
|
||||
output = capsys.readouterr().out
|
||||
normalized_files = list((reports_dir / "results").glob("normalized_perf_*.json"))
|
||||
assert "Skipping rolling baseline comparison" in output
|
||||
assert "workload_id" in output
|
||||
assert len(normalized_files) == 1
|
||||
|
||||
normalized = json.loads(normalized_files[0].read_text(encoding="utf-8"))
|
||||
assert "result_schema_version" not in normalized
|
||||
assert normalized["success"] is True
|
||||
assert normalized["baseline_eligible"] is False
|
||||
assert len(uploaded_records) == 1
|
||||
assert uploaded_records[0]["baseline_eligible"] is False
|
||||
|
||||
|
||||
def test_normalized_record_reads_identity_labels_from_recipe(monkeypatch):
|
||||
|
||||
@@ -154,8 +154,8 @@ def test_dashboard_endpoints_filter_and_return_run_source_metadata():
|
||||
assert summary["count"] == 1
|
||||
assert summary["rows"][0]["run_source"] == "pr"
|
||||
assert summary["rows"][0]["pr_number"] == "123"
|
||||
assert summary["rows"][0]["baseline_n"] == 0
|
||||
assert summary["rows"][0]["metrics"]["latency"]["baseline"] is None
|
||||
assert summary["rows"][0]["baseline_n"] == 1
|
||||
assert summary["rows"][0]["metrics"]["latency"]["baseline"] == 11.0
|
||||
assert summary["filters"]["run_source"] == "pr"
|
||||
assert trends["count"] == 1
|
||||
assert trends["groups"][0]["points"][0]["run_source"] == "pr"
|
||||
|
||||
@@ -142,47 +142,6 @@ def test_build_latest_summary_separates_informational_threshold_crossing():
|
||||
assert rows[0]["threshold_exceeded_metrics"] == ["latency"]
|
||||
assert rows[0]["failing_metrics"] == []
|
||||
assert rows[0]["computed_regression_status"] == "pass"
|
||||
|
||||
|
||||
def test_build_latest_summary_run_source_filter_excludes_future_baselines():
|
||||
records = [
|
||||
_record(
|
||||
"2026-01-01T00:00:00+00:00",
|
||||
"a" * 40,
|
||||
10.0,
|
||||
10.0,
|
||||
run_source="scheduled_main",
|
||||
baseline_eligible=True,
|
||||
),
|
||||
_record(
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
11.0,
|
||||
9.0,
|
||||
run_source="pr",
|
||||
baseline_eligible=False,
|
||||
pr_number="123",
|
||||
),
|
||||
_record(
|
||||
"2026-01-03T00:00:00+00:00",
|
||||
"c" * 40,
|
||||
30.0,
|
||||
3.0,
|
||||
run_source="scheduled_main",
|
||||
baseline_eligible=True,
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records, run_source="pr")
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["run_source"] == "pr"
|
||||
assert rows[0]["baseline_n"] == 1
|
||||
assert rows[0]["metrics"]["latency"]["baseline"] == 10.0
|
||||
assert rows[0]["metrics"]["throughput"]["baseline"] == 10.0
|
||||
|
||||
|
||||
|
||||
def test_build_latest_summary_keeps_identity_cohorts_separate():
|
||||
records = [
|
||||
_record(
|
||||
@@ -227,7 +186,7 @@ def test_build_latest_summary_keeps_identity_cohorts_separate():
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
rows = build_latest_summary(records, max_regression=0.05)
|
||||
recipe_b_row = next(row for row in rows if row["recipe_fingerprint"] == "recipe-b")
|
||||
|
||||
assert len(rows) == 2
|
||||
@@ -282,7 +241,7 @@ def test_build_latest_summary_keeps_variant_versions_separate():
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
rows = build_latest_summary(records, max_regression=0.05)
|
||||
version_2_row = next(row for row in rows if row["benchmark_version"] == 2)
|
||||
|
||||
assert len(rows) == 2
|
||||
@@ -290,58 +249,6 @@ def test_build_latest_summary_keeps_variant_versions_separate():
|
||||
assert version_2_row["metrics"]["latency"]["baseline"] == 20.0
|
||||
|
||||
|
||||
def test_dashboard_identity_preserves_zero_benchmark_version():
|
||||
records = [
|
||||
_record(
|
||||
"2026-01-01T00:00:00+00:00",
|
||||
"a" * 40,
|
||||
10.0,
|
||||
10.0,
|
||||
workload_id="wan-t2v",
|
||||
variant_id="1.3b-sp2",
|
||||
benchmark_version=0,
|
||||
recipe_fingerprint="recipe-a",
|
||||
hardware_profile_id="hw-l40s",
|
||||
software_profile_id="sw-cu130",
|
||||
run_source="scheduled_main",
|
||||
baseline_eligible=True,
|
||||
),
|
||||
_record(
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
11.0,
|
||||
9.0,
|
||||
workload_id="wan-t2v",
|
||||
variant_id="1.3b-sp2",
|
||||
benchmark_version=0,
|
||||
recipe_fingerprint="recipe-a",
|
||||
hardware_profile_id="hw-l40s",
|
||||
software_profile_id="sw-cu130",
|
||||
),
|
||||
_record(
|
||||
"2026-01-03T00:00:00+00:00",
|
||||
"c" * 40,
|
||||
20.0,
|
||||
5.0,
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
version_zero_row = next(row for row in rows if row["benchmark_version"] == 0)
|
||||
legacy_row = next(row for row in rows if row["benchmark_version"] == "")
|
||||
trends = build_trends(records)
|
||||
version_zero_trend = next(trend for trend in trends if trend["benchmark_version"] == 0)
|
||||
legacy_trend = next(trend for trend in trends if trend["benchmark_version"] == "")
|
||||
|
||||
assert len(rows) == 2
|
||||
assert version_zero_row["baseline_n"] == 1
|
||||
assert version_zero_row["metrics"]["latency"]["baseline"] == 10.0
|
||||
assert legacy_row["baseline_n"] == 0
|
||||
assert len(trends) == 2
|
||||
assert version_zero_trend["points"][0]["benchmark_version"] == 0
|
||||
assert legacy_trend["points"][0]["benchmark_version"] == ""
|
||||
|
||||
|
||||
def test_filter_records_and_trends_preserve_metric_points():
|
||||
records = [
|
||||
_record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
|
||||
|
||||
@@ -55,9 +55,6 @@ V2_OPTIONAL_METADATA_FIELDS = (
|
||||
"metric_threshold_policy",
|
||||
"quality_metadata",
|
||||
)
|
||||
RESULT_SCHEMA_VERSION = 2
|
||||
VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
|
||||
OPTIONAL_RESULT_METADATA_FIELDS = ("quality_metadata", "variant_metadata")
|
||||
|
||||
# -- Config discovery -------------------------------------------------------
|
||||
|
||||
@@ -198,11 +195,7 @@ def _extract_component_times(result: dict) -> dict[str, float | None]:
|
||||
logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
|
||||
continue
|
||||
stage_class = stage_data.get("stage_class", stage_name)
|
||||
component_metric = stage_data.get("component_metric")
|
||||
if isinstance(component_metric, str) and component_metric in component_times:
|
||||
metric_key = component_metric
|
||||
else:
|
||||
metric_key = STAGE_METRIC_MAP.get(stage_class)
|
||||
metric_key = STAGE_METRIC_MAP.get(stage_class)
|
||||
if metric_key is None:
|
||||
logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
|
||||
stage_name,
|
||||
@@ -234,17 +227,6 @@ def _run_generation(generator, prompt, generation_kwargs):
|
||||
component_times = _extract_component_times(result)
|
||||
return elapsed, peak_memory_mb, component_times
|
||||
|
||||
|
||||
def _validate_run_counts(run_config: Mapping[str, Any], benchmark_id: str) -> tuple[int, int]:
|
||||
num_warmup = run_config.get("num_warmup_runs", 1)
|
||||
num_measure = run_config.get("num_measurement_runs", 3)
|
||||
if isinstance(num_warmup, bool) or not isinstance(num_warmup, int) or num_warmup < 0:
|
||||
raise ValueError(f"{benchmark_id}: run_config.num_warmup_runs must be a non-negative integer")
|
||||
if isinstance(num_measure, bool) or not isinstance(num_measure, int) or num_measure < 1:
|
||||
raise ValueError(f"{benchmark_id}: run_config.num_measurement_runs must be a positive integer")
|
||||
return num_warmup, num_measure
|
||||
|
||||
|
||||
def _write_results(results):
|
||||
"""Write JSON results to the results directory."""
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
@@ -365,101 +347,6 @@ def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity):
|
||||
}
|
||||
|
||||
|
||||
def _truthy_pr_number(value: str | None) -> bool:
|
||||
return bool(value and value not in {"false", "0", "None", "none"})
|
||||
|
||||
|
||||
def _detect_run_source() -> str:
|
||||
explicit = os.environ.get("PERF_RUN_SOURCE", "").strip().lower()
|
||||
if explicit in VALID_RUN_SOURCES:
|
||||
return explicit
|
||||
if _truthy_pr_number(os.environ.get("BUILDKITE_PULL_REQUEST")):
|
||||
return "pr"
|
||||
if os.environ.get("BUILDKITE_BRANCH") == "main" and os.environ.get("TEST_SCOPE") == "full":
|
||||
return "scheduled_main"
|
||||
if not os.environ.get("BUILDKITE_COMMIT"):
|
||||
return "local"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _ci_provenance_fields() -> dict[str, str]:
|
||||
pr_number = os.environ.get("BUILDKITE_PULL_REQUEST", "")
|
||||
if not _truthy_pr_number(pr_number):
|
||||
pr_number = ""
|
||||
return {
|
||||
"run_source": _detect_run_source(),
|
||||
"branch": os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"pr_number": pr_number,
|
||||
"test_scope": os.environ.get("TEST_SCOPE", ""),
|
||||
"build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
|
||||
"build_id": os.environ.get("BUILDKITE_BUILD_ID", ""),
|
||||
"job_id": os.environ.get("BUILDKITE_JOB_ID", ""),
|
||||
}
|
||||
|
||||
|
||||
def _configured_result_metadata(cfg: Mapping[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
field: cfg[field]
|
||||
for field in OPTIONAL_RESULT_METADATA_FIELDS
|
||||
if field in cfg and cfg[field] is not None
|
||||
}
|
||||
|
||||
|
||||
def _build_result_record(
|
||||
*,
|
||||
cfg: Mapping[str, Any],
|
||||
model_info: Mapping[str, Any],
|
||||
init_kwargs: Mapping[str, Any],
|
||||
gen_kwargs: Mapping[str, Any],
|
||||
num_warmup: int,
|
||||
num_measure: int,
|
||||
thresholds: Mapping[str, Any],
|
||||
times: list[float],
|
||||
peak_memories: list[float],
|
||||
all_component_times: list[dict],
|
||||
prompt: str,
|
||||
runtime_identity: Mapping[str, Any],
|
||||
device_name: str,
|
||||
timestamp: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if not times or not peak_memories:
|
||||
raise ValueError("Cannot build a performance result record without measurement runs")
|
||||
avg_time = sum(times) / len(times)
|
||||
max_peak_memory = max(peak_memories)
|
||||
num_frames = gen_kwargs.get("num_frames")
|
||||
throughput_fps = (1.0 / avg_time) if avg_time > 0 else None
|
||||
if isinstance(num_frames, (int, float)) and avg_time > 0:
|
||||
throughput_fps = num_frames / avg_time
|
||||
|
||||
return {
|
||||
"benchmark_id": cfg["benchmark_id"],
|
||||
"result_schema_version": RESULT_SCHEMA_VERSION,
|
||||
"model_short_name": model_info.get("model_short_name", ""),
|
||||
"device": device_name,
|
||||
"num_gpus": init_kwargs.get("num_gpus", 1),
|
||||
"num_warmup_runs": num_warmup,
|
||||
"num_measurement_runs": num_measure,
|
||||
"avg_generation_time_s": round(avg_time, 3),
|
||||
"individual_times_s": [round(t, 3) for t in times],
|
||||
"throughput_fps": round(throughput_fps, 3)
|
||||
if throughput_fps is not None else None,
|
||||
"max_peak_memory_mb": round(max_peak_memory, 1),
|
||||
"individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
|
||||
"thresholds": dict(thresholds),
|
||||
"regression_thresholds": cfg.get("regression_thresholds", {}),
|
||||
"commit": os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
**_ci_provenance_fields(),
|
||||
"timestamp": timestamp or datetime.now(timezone.utc).isoformat(),
|
||||
**_configured_result_metadata(cfg),
|
||||
"text_encoder_time_s": _avg_component(all_component_times,
|
||||
"text_encoder_time_s"),
|
||||
"dit_time_s": _avg_component(all_component_times, "dit_time_s"),
|
||||
"vae_decode_time_s": _avg_component(all_component_times,
|
||||
"vae_decode_time_s"),
|
||||
**_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity),
|
||||
}
|
||||
|
||||
|
||||
# -- Test -------------------------------------------------------------------
|
||||
|
||||
def _run_benchmark(cfg):
|
||||
@@ -475,7 +362,8 @@ def _run_benchmark(cfg):
|
||||
prompts = cfg.get("test_prompts", ["A cinematic video."])
|
||||
prompt = prompts[0]
|
||||
|
||||
num_warmup, num_measure = _validate_run_counts(run_config, cfg["benchmark_id"])
|
||||
num_warmup = run_config.get("num_warmup_runs", 1)
|
||||
num_measure = run_config.get("num_measurement_runs", 3)
|
||||
thresholds = _get_thresholds(cfg)
|
||||
|
||||
# Remap JSON keys to VideoGenerator kwargs
|
||||
@@ -522,21 +410,37 @@ def _run_benchmark(cfg):
|
||||
avg_time = sum(times) / len(times)
|
||||
max_peak_memory = max(peak_memories)
|
||||
device_name = torch.cuda.get_device_name()
|
||||
results = _build_result_record(
|
||||
cfg=cfg,
|
||||
model_info=model_info,
|
||||
init_kwargs=init_kwargs,
|
||||
gen_kwargs=gen_kwargs,
|
||||
num_warmup=num_warmup,
|
||||
num_measure=num_measure,
|
||||
thresholds=thresholds,
|
||||
times=times,
|
||||
peak_memories=peak_memories,
|
||||
all_component_times=all_component_times,
|
||||
prompt=prompt,
|
||||
runtime_identity=runtime_identity,
|
||||
device_name=device_name,
|
||||
)
|
||||
num_frames = gen_kwargs.get("num_frames")
|
||||
throughput_fps = (1.0 / avg_time) if avg_time > 0 else None
|
||||
if isinstance(num_frames, (int, float)) and avg_time > 0:
|
||||
throughput_fps = num_frames / avg_time
|
||||
|
||||
results = {
|
||||
"benchmark_id": cfg["benchmark_id"],
|
||||
**_config_identity_metadata(cfg),
|
||||
"model_short_name": model_info.get("model_short_name", ""),
|
||||
"device": device_name,
|
||||
"num_gpus": init_kwargs.get("num_gpus", 1),
|
||||
"num_warmup_runs": num_warmup,
|
||||
"num_measurement_runs": num_measure,
|
||||
"avg_generation_time_s": round(avg_time, 3),
|
||||
"individual_times_s": [round(t, 3) for t in times],
|
||||
"throughput_fps": round(throughput_fps, 3)
|
||||
if throughput_fps is not None else None,
|
||||
"max_peak_memory_mb": round(max_peak_memory, 1),
|
||||
"individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
|
||||
"thresholds": thresholds,
|
||||
"regression_thresholds": cfg.get("regression_thresholds", {}),
|
||||
"commit": os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
"pr_number": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"text_encoder_time_s": _avg_component(all_component_times,
|
||||
"text_encoder_time_s"),
|
||||
"dit_time_s": _avg_component(all_component_times, "dit_time_s"),
|
||||
"vae_decode_time_s": _avg_component(all_component_times,
|
||||
"vae_decode_time_s"),
|
||||
**_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity),
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Performance results: avg_time=%.2fs, "
|
||||
|
||||
@@ -1,20 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from fastvideo.pipelines.pipeline_batch_info import PipelineLoggingInfo
|
||||
from fastvideo.pipelines.stages.denoising import Cosmos25AutoDenoisingStage, DenoisingStage
|
||||
from fastvideo.pipelines.stages.text_encoding import Cosmos25TextEncodingStage
|
||||
from fastvideo.tests.performance.test_inference_performance import _extract_component_times
|
||||
|
||||
|
||||
class SubclassStyleDenoisingStage(DenoisingStage):
|
||||
pass
|
||||
|
||||
|
||||
def test_extract_component_times_handles_pipeline_logging_info_object():
|
||||
logging_info = PipelineLoggingInfo()
|
||||
logging_info.add_stage_execution_time("prompt_encoding_stage", 1.25)
|
||||
logging_info.add_stage_metric("prompt_encoding_stage", "stage_class", "TextEncodingStage")
|
||||
logging_info.add_stage_metric("prompt_encoding_stage", "component_metric", "text_encoder_time_s")
|
||||
|
||||
assert _extract_component_times({"logging_info": logging_info}) == {
|
||||
"text_encoder_time_s": 1.25,
|
||||
@@ -52,35 +45,6 @@ def test_extract_component_times_uses_stage_class_for_pipeline_stage_keys():
|
||||
}
|
||||
|
||||
|
||||
def test_denoising_stage_subclasses_inherit_component_metric():
|
||||
assert SubclassStyleDenoisingStage.performance_component_metric == "dit_time_s"
|
||||
|
||||
|
||||
def test_cosmos25_direct_pipeline_stages_define_component_metrics():
|
||||
assert Cosmos25TextEncodingStage.performance_component_metric == "text_encoder_time_s"
|
||||
assert Cosmos25AutoDenoisingStage.performance_component_metric == "dit_time_s"
|
||||
|
||||
|
||||
def test_extract_component_times_uses_component_metric_for_stage_subclasses():
|
||||
result = {
|
||||
"logging_info": {
|
||||
"stages": {
|
||||
"denoising_stage": {
|
||||
"execution_time": 4.2,
|
||||
"stage_class": "CosmosDenoisingStage",
|
||||
"component_metric": "dit_time_s",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
assert _extract_component_times(result) == {
|
||||
"text_encoder_time_s": None,
|
||||
"dit_time_s": 4.2,
|
||||
"vae_decode_time_s": None,
|
||||
}
|
||||
|
||||
|
||||
def test_extract_component_times_keeps_legacy_class_name_keys():
|
||||
# Backward-compatibility check for logs produced before pipeline-unique
|
||||
# stage keys carried a separate stage_class field.
|
||||
|
||||
@@ -1,176 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo.tests.performance import test_inference_performance as perf_test
|
||||
|
||||
|
||||
def _benchmark_config():
|
||||
return {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"quality_metadata": {
|
||||
"quality_status": "canonical",
|
||||
},
|
||||
"model": {
|
||||
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
"model_short_name": "Wan2.1-T2V-1.3B",
|
||||
},
|
||||
"generation_kwargs": {
|
||||
"num_frames": 45,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _identity_fields():
|
||||
return {
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
"recipe": {
|
||||
"recipe_schema_version": 1,
|
||||
"benchmark": {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v",
|
||||
"variant_id": "1.3b-sp2",
|
||||
"benchmark_version": 2,
|
||||
},
|
||||
},
|
||||
"recipe_fingerprint": "recipe-1",
|
||||
"hardware_profile": {
|
||||
"device_type": "cuda",
|
||||
"gpu_count": 2,
|
||||
},
|
||||
"hardware_profile_id": "hw-l40s",
|
||||
"software_profile": {
|
||||
"python": "3.12",
|
||||
"pytorch": "2.12",
|
||||
"cuda": "13.0",
|
||||
},
|
||||
"software_profile_id": "sw-cu130",
|
||||
"environment_metadata": {
|
||||
"env": {
|
||||
"IMAGE_VERSION": "py3.12-cuda13.0.0",
|
||||
},
|
||||
},
|
||||
"environment_fingerprint": "env-ci",
|
||||
}
|
||||
|
||||
|
||||
def test_build_result_record_emits_v2_wan_shape(monkeypatch):
|
||||
monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main")
|
||||
monkeypatch.setenv("BUILDKITE_COMMIT", "a" * 40)
|
||||
monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false")
|
||||
monkeypatch.setenv("BUILDKITE_BRANCH", "main")
|
||||
monkeypatch.setenv("TEST_SCOPE", "full")
|
||||
monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/build")
|
||||
monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-1")
|
||||
monkeypatch.setenv("BUILDKITE_JOB_ID", "job-1")
|
||||
monkeypatch.setattr(perf_test, "_build_identity_fields", lambda *_args: _identity_fields())
|
||||
|
||||
record = perf_test._build_result_record(
|
||||
cfg=_benchmark_config(),
|
||||
model_info={"model_short_name": "Wan2.1-T2V-1.3B"},
|
||||
init_kwargs={"num_gpus": 2},
|
||||
gen_kwargs={"num_frames": 45},
|
||||
num_warmup=2,
|
||||
num_measure=3,
|
||||
thresholds={
|
||||
"max_generation_time_s": 34.0,
|
||||
"max_peak_memory_mb": 11000.0,
|
||||
},
|
||||
times=[30.0, 31.0, 32.0],
|
||||
peak_memories=[10000.0, 10100.0, 10050.0],
|
||||
all_component_times=[
|
||||
{
|
||||
"text_encoder_time_s": 1.0,
|
||||
"dit_time_s": 8.0,
|
||||
"vae_decode_time_s": 3.0,
|
||||
},
|
||||
{
|
||||
"text_encoder_time_s": 1.2,
|
||||
"dit_time_s": 8.2,
|
||||
"vae_decode_time_s": 3.2,
|
||||
},
|
||||
{
|
||||
"text_encoder_time_s": None,
|
||||
"dit_time_s": 8.4,
|
||||
"vae_decode_time_s": 3.4,
|
||||
},
|
||||
],
|
||||
prompt="A cinematic video.",
|
||||
runtime_identity={
|
||||
"resolved_attention_backend": "FLASH_ATTN",
|
||||
},
|
||||
device_name="NVIDIA L40S",
|
||||
timestamp="2026-07-05T00:00:00+00:00",
|
||||
)
|
||||
|
||||
assert record["result_schema_version"] == perf_test.RESULT_SCHEMA_VERSION
|
||||
assert record["benchmark_id"] == "wan-t2v-1.3b-2gpu"
|
||||
assert record["workload_id"] == "wan-t2v"
|
||||
assert record["variant_id"] == "1.3b-sp2"
|
||||
assert record["benchmark_version"] == 2
|
||||
assert record["recipe_fingerprint"] == "recipe-1"
|
||||
assert record["hardware_profile_id"] == "hw-l40s"
|
||||
assert record["software_profile_id"] == "sw-cu130"
|
||||
assert record["environment_fingerprint"] == "env-ci"
|
||||
assert record["environment_metadata"]["env"]["IMAGE_VERSION"] == "py3.12-cuda13.0.0"
|
||||
assert record["quality_metadata"] == {"quality_status": "canonical"}
|
||||
assert record["run_source"] == "scheduled_main"
|
||||
assert record["branch"] == "main"
|
||||
assert record["test_scope"] == "full"
|
||||
assert record["build_url"] == "https://buildkite.example/build"
|
||||
assert record["build_id"] == "build-1"
|
||||
assert record["job_id"] == "job-1"
|
||||
assert record["pr_number"] == ""
|
||||
assert record["commit"] == "a" * 40
|
||||
assert record["avg_generation_time_s"] == 31.0
|
||||
assert record["throughput_fps"] == 1.452
|
||||
assert record["max_peak_memory_mb"] == 10100.0
|
||||
assert record["text_encoder_time_s"] == 1.1
|
||||
assert record["dit_time_s"] == 8.2
|
||||
assert record["vae_decode_time_s"] == 3.2
|
||||
|
||||
|
||||
def test_validate_run_counts_rejects_zero_measurement_runs():
|
||||
with pytest.raises(ValueError, match="num_measurement_runs"):
|
||||
perf_test._validate_run_counts({
|
||||
"num_warmup_runs": 1,
|
||||
"num_measurement_runs": 0,
|
||||
}, "wan-t2v-1.3b-2gpu")
|
||||
|
||||
|
||||
def test_validate_run_counts_rejects_negative_warmup_runs():
|
||||
with pytest.raises(ValueError, match="num_warmup_runs"):
|
||||
perf_test._validate_run_counts({
|
||||
"num_warmup_runs": -1,
|
||||
"num_measurement_runs": 3,
|
||||
}, "wan-t2v-1.3b-2gpu")
|
||||
|
||||
|
||||
def test_validate_run_counts_returns_defaults():
|
||||
assert perf_test._validate_run_counts({}, "wan-t2v-1.3b-2gpu") == (1, 3)
|
||||
|
||||
|
||||
def test_build_result_record_rejects_empty_measurements(monkeypatch):
|
||||
monkeypatch.setattr(perf_test, "_build_identity_fields", lambda *_args: _identity_fields())
|
||||
|
||||
with pytest.raises(ValueError, match="measurement runs"):
|
||||
perf_test._build_result_record(
|
||||
cfg=_benchmark_config(),
|
||||
model_info={},
|
||||
init_kwargs={},
|
||||
gen_kwargs={},
|
||||
num_warmup=0,
|
||||
num_measure=0,
|
||||
thresholds={},
|
||||
times=[],
|
||||
peak_memories=[],
|
||||
all_component_times=[],
|
||||
prompt="A cinematic video.",
|
||||
runtime_identity={},
|
||||
device_name="NVIDIA L40S",
|
||||
)
|
||||
@@ -3,9 +3,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from logging import Logger
|
||||
from typing import Iterator
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.tests.ssim.reference_utils import (
|
||||
@@ -67,12 +67,6 @@ def _find_reference_video(reference_folder: str, prompt: str) -> str:
|
||||
raise FileNotFoundError("Reference video missing")
|
||||
|
||||
|
||||
def _remove_stale_generated_video(output_dir: str, output_video_name: str) -> None:
|
||||
stale_path = os.path.join(output_dir, output_video_name)
|
||||
if os.path.exists(stale_path):
|
||||
os.remove(stale_path)
|
||||
|
||||
|
||||
def _assert_similarity(
|
||||
*,
|
||||
logger: Logger,
|
||||
@@ -220,7 +214,6 @@ def run_text_to_video_similarity_test(
|
||||
)
|
||||
output_video_name = f"{prompt[:100].strip()}.mp4"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
_remove_stale_generated_video(output_dir, output_video_name)
|
||||
|
||||
params_map = select_ssim_params(
|
||||
default_params_map,
|
||||
@@ -296,7 +289,6 @@ def run_image_to_video_similarity_test(
|
||||
)
|
||||
output_video_name = f"{prompt[:100].strip()}.mp4"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
_remove_stale_generated_video(output_dir, output_video_name)
|
||||
|
||||
params_map = select_ssim_params(
|
||||
default_params_map,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user