SeedVR2: bounded long-video restoration (#8102)
Signed-off-by: 0z5a <Dezhen.lu@student.uni-tuebingen.de> Signed-off-by: 0z5a <dezhen.lu@uni-tuebingen.de> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de> Signed-off-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: 0z5a <dezhen.lu@uni-tuebingen.de> Co-authored-by: princepride <wangzhipeng628@gmail.com>
This commit is contained in:
20 files changed
+4684
-9
No files matched your search
@@ -47,6 +47,7 @@ nav:
|
||||
- configuration/*
|
||||
- Models:
|
||||
- models/supported_models.md
|
||||
- SeedVR2: models/seedvr2.md
|
||||
- Features:
|
||||
- Overview: features/README.md
|
||||
- Runtime and Stage Execution:
|
||||
@@ -151,6 +152,7 @@ nav:
|
||||
- Hybrid Sharded Data Parallel (HSDP): design/feature/hsdp.md
|
||||
- Pipeline Parallel: design/feature/pipeline_parallel.md
|
||||
- Sequence Parallel: design/feature/sequence_parallel.md
|
||||
- Window Sequence Parallel: design/feature/window_sequence_parallel.md
|
||||
- Tensor Parallel: design/feature/tensor_parallel.md
|
||||
- VAE Patch Parallelism: design/feature/vae_parallel.md
|
||||
- Attention Optimization:
|
||||
|
||||
@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.
|
||||
|
||||
## Model-specific variables
|
||||
|
||||
The audit found 58 variables read by a single model or pipeline family. They are
|
||||
The inventory includes 66 variables read by a single model or pipeline family. They are
|
||||
not listed as public usage options here because doing so would turn implementation
|
||||
escape hatches into an accidental compatibility contract.
|
||||
|
||||
@@ -223,11 +223,11 @@ Every audited model-specific name has a migration disposition in the
|
||||
|
||||
| Disposition | Count | Required outcome |
|
||||
| --- | ---: | --- |
|
||||
| Promote | 35 | Move a stable setting into typed stage or model configuration. |
|
||||
| Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
|
||||
| Promote | 38 | Move a stable setting into typed stage or model configuration. |
|
||||
| Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
|
||||
| External | 0 | Retain only when a supported third-party library owns the contract. |
|
||||
| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
|
||||
| Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. |
|
||||
| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
|
||||
| Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
|
||||
|
||||
The disposition is a migration target, not a statement that the existing
|
||||
environment switch is stable. Promote or request-scope work should land in
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
# Window-Aligned Sequence Parallelism (SeedVR2)
|
||||
|
||||
## Motivation
|
||||
|
||||
SeedVR2's diffusion transformer (`NaDiT`) has no global attention: every layer
|
||||
partitions the video tokens into 3D windows and attends only inside a window.
|
||||
The window size is normalised against a fixed 720p reference area, so raising the
|
||||
output resolution grows the *number* of windows, never their size.
|
||||
|
||||
Ulysses redistributes the sequence and head dimensions around attention.
|
||||
For SeedVR2, window-aligned SP offers a different communication tradeoff: one
|
||||
hidden-state exchange per layout boundary versus Q/K/V and output exchanges
|
||||
for Ulysses. The actual performance crossover requires measurement.
|
||||
Window-aligned SP instead assigns **whole windows** to ranks:
|
||||
|
||||
* every rank keeps **all** attention heads, so no head-count divisibility is
|
||||
required,
|
||||
* attention inside a layer needs **zero** communication,
|
||||
* regular and shifted layouts alternate at every layer, so the released
|
||||
32-layer model has 31 inter-layer transitions, plus one text reduction per
|
||||
layer and a final video gather. A transition can be a local permutation
|
||||
when ownership does not change.
|
||||
|
||||
## Window layouts
|
||||
|
||||
`nadit.py` ports the reference windowing verbatim (Apache-2.0,
|
||||
ByteDance SeedVR2). For a post-patch token grid `(T, H, W)`:
|
||||
|
||||
```text
|
||||
scale = sqrt(45 * 80 / (H * W)) # 720p reference area
|
||||
wh, ww = ceil(round(H * scale) / nh), ceil(round(W * scale) / nw)
|
||||
wt = ceil(min(T, 30) / nt) # BYTEDANCE_MAX_TEMPORAL_WINDOW
|
||||
```
|
||||
|
||||
The regular method tiles with `slice(it * wt, min((it + 1) * wt, T))`; the
|
||||
shifted method moves every cut by half a window and **clips at the boundary**
|
||||
(`max(..., 0)` / `min(..., size)`) -- it is not a cyclic `torch.roll`. Both
|
||||
methods produce an exact partition of the grid; the shifted layout is genuinely
|
||||
different, so an `A -> B -> A -> B` layer schedule needs both directions of
|
||||
re-sharding.
|
||||
|
||||
Windows are enumerated `for iw ... for ih ... for it ...` (w outermost, t
|
||||
innermost) and tokens inside a window are row-major in `(t, h, w)`. The
|
||||
canonical token id is `(t * H + h) * W + w`, which is exactly what the reference
|
||||
obtains by windowing `torch.arange`.
|
||||
|
||||
Consequences that matter (all verified against the reference):
|
||||
|
||||
* Window token counts are **ragged** (edge windows are clipped), so a shard must
|
||||
be balanced by token count, not by window count.
|
||||
* The number of layout transitions follows the actual layer schedule. For the
|
||||
released 3B checkpoint (`num_layers=32`, alternating methods) there are 31
|
||||
inter-layer transitions, not 16; `ensure_layout(current, required)` derives
|
||||
them instead of assuming `layer % 2`.
|
||||
|
||||
## Planner
|
||||
|
||||
`nadit.py` implements a deterministic LPT assignment:
|
||||
|
||||
```text
|
||||
windows sorted by (-video_token_count, original_window_id)
|
||||
each window is placed on the rank with the smallest (assigned_tokens, rank_id)
|
||||
each rank then lists its windows in ascending original order
|
||||
```
|
||||
|
||||
The cache key covers the post-patch token grid, the full window method and
|
||||
parameters (via a SHA-1 geometry fingerprint), the geometry version, the **SP
|
||||
world size** and the planner version, so a different degree can never reuse a
|
||||
stale assignment. Caches are LRU-bounded.
|
||||
|
||||
## Plan A redistribution
|
||||
|
||||
At a layout boundary the activations must move from the source assignment to the
|
||||
destination assignment. With `Z_src[r]` / `Z_dst[d]` the canonical token ids held
|
||||
by each rank and `C[r, d]` the number of rows rank `d` needs from rank `r`:
|
||||
|
||||
* `input_split_sizes` of rank `r` is row `r` of `C`, `output_split_sizes` is
|
||||
column `r`,
|
||||
* each peer segment is packed in **destination-local order**, so the receiver can
|
||||
reorder with a single `argsort` of the concatenated destination positions,
|
||||
* `network_exchange_required` is computed from the global count matrix, so every
|
||||
rank takes the same branch; a rank that happens to exchange nothing never
|
||||
skips the collective alone.
|
||||
|
||||
The runtime is a variable-split `torch.distributed.all_to_all_single` over an
|
||||
regular SP process group (`get_sp_group().device_group`), plus a local-permutation fast
|
||||
path when the whole group stays in place (which is the SP=1 case). The first
|
||||
version is deliberately synchronous: no extra streams, no per-layer
|
||||
`cuda.synchronize()` or barrier "fixups", and no hidden all-gather.
|
||||
|
||||
## Text stream
|
||||
|
||||
Text is replicated on every rank and stays identical: each layer's attention
|
||||
produces one text output per window, ranks sum their local window outputs,
|
||||
one SUM all-reduce runs on the window-SP group, and every rank divides by the
|
||||
**global** window count of the current layout.
|
||||
|
||||
```text
|
||||
T_next = (sum_r sum_{w in Omega_r} T_w) / W_global
|
||||
```
|
||||
|
||||
This is neither a token-weighted mean nor a mean of rank means; a rank with no
|
||||
window contributes a zero tensor (with the *same dtype as its peers* -- a dtype
|
||||
mismatch changes the collective and deadlocks NCCL, which is why the empty-rank
|
||||
path is covered by a GPU test).
|
||||
|
||||
## Existing SP configuration
|
||||
|
||||
Configure `ulysses_degree=N` to create the existing N-rank SP group. Pass
|
||||
`DiffusionParallelConfig` to `SeedVR2NaDiT.build_runtime(parallel_config=...)`;
|
||||
the model validates the configuration and obtains `get_sp_group().device_group`.
|
||||
This adds no framework parallel option or group-construction branch.
|
||||
|
||||
For SeedVR2 this selects **window-aligned Plan A**, not head-sharded Ulysses.
|
||||
Every rank retains all 20 attention heads, so SP=8 does not require the head
|
||||
count to divide by eight. The model has no `_sp_plan`, and its attention uses
|
||||
`skip_sequence_parallel=True`; generic sequence hooks and Ulysses attention
|
||||
exchanges are not installed. Reports name the actual execution path explicitly.
|
||||
|
||||
SeedVR2 rejects ring, allgather, `advanced_uaa`, and `ulysses_a2a_permute`.
|
||||
Other diffusion models keep their existing Ulysses/Ring/AllGather behavior.
|
||||
|
||||
## SeedVR2 integration
|
||||
|
||||
`SeedVR2WindowRuntime` owns, per request: the layer schedule, the entry
|
||||
distribution, `ensure_layout` transitions, the text reduction and the final
|
||||
reconstruction back to canonical token order. The DiT forward keeps only this
|
||||
rank's video rows from the patch embedding onwards; the patch projection,
|
||||
norms, ADA modulation, residual and MLP are all tokenwise, so they run locally.
|
||||
|
||||
The block never calls a parallel attention strategy: the shared `Attention`
|
||||
layer is created with `skip_sequence_parallel=True`, and window attention is a
|
||||
packed-varlen call over `[video window, replicated text]` segments (with a
|
||||
grouped-SDPA fallback for backends without a packed-varlen entry point).
|
||||
|
||||
## Reference quirks reproduced on purpose
|
||||
|
||||
* `vid_out_ada` is declared with `layers=["out"]`, whose embedding regrouping is
|
||||
arithmetically incompatible with its 1-D parameters. The reference never uses
|
||||
that value: it asks the forward-wide cache for `emb_repeat_0_vid`, a key the
|
||||
first block already wrote, and therefore applies the *block* attention
|
||||
modulation together with the learned `out_scale` / `out_shift`. The port
|
||||
applies that effective modulation directly; changing it would break parity.
|
||||
* `mm_layers=10` controls weight sharing only. Text is consumed and produced by
|
||||
all 32 blocks; only the last block freezes its text MLP path.
|
||||
* RoPE positions are window-local: a video token at window-local `(dt, dh, dw)`
|
||||
uses axial indices `(text_len + dt, dh, dw)`, and every window restarts at 0.
|
||||
* The rotary span is `3 * (rope_dim // 3)` (126 of 128 channels for the 3B), so
|
||||
the last two channels pass through unrotated.
|
||||
|
||||
## Checkpoint compatibility
|
||||
|
||||
The parameter layout of the port mirrors the reference, with one explicit
|
||||
exception: the released checkpoint stores the per-block RoPE table as
|
||||
`blocks.<i>.attn.rope.rope.freqs`, while this port registers it as
|
||||
`blocks.<i>.attn.rope.freqs`. The validation loader normalizes exactly that
|
||||
suffix (and only that suffix); every other key must already match.
|
||||
|
||||
> The validation loader explicitly normalizes the reference RoPE buffer suffix
|
||||
> from `.rope.rope.freqs` to `.rope.freqs`. Checkpoint validation fails on any
|
||||
> remaining missing or unexpected key after normalization.
|
||||
|
||||
Normalization is restricted and collision-checked: a key that is already
|
||||
normalized is left untouched, and two source keys mapping onto the same target
|
||||
key raise instead of silently overwriting a tensor. Loading then fails on any
|
||||
remaining missing, unexpected or shape-mismatched key, so a checkpoint that does
|
||||
not match the port cannot reach inference. Full 32-layer loading is the
|
||||
acceptance path; truncating the block list is a development fixture and requires
|
||||
an explicit flag.
|
||||
|
||||
The native serving loader applies the same suffix normalization and strict
|
||||
checkpoint check.
|
||||
|
||||
## Integration limits
|
||||
|
||||
The native serving pipeline passes its parallel configuration to the DiT runtime
|
||||
and returns restored video through the standard media path. The causal VAE can
|
||||
run replicated, tiled, or height-sharded across the window-SP group.
|
||||
|
||||
## Limitations
|
||||
|
||||
* Plan A re-shards the full activation at every layout boundary. Plan B halo
|
||||
exchange is a follow-up with a separate crossover evaluation.
|
||||
* The packed-varlen attention kernel needs a backend that accepts
|
||||
`cu_seqlens`; the grouped-SDPA path is the portable fallback and is what the
|
||||
pinned L20 environment exercises.
|
||||
@@ -0,0 +1,280 @@
|
||||
# SeedVR2 3B video restoration
|
||||
|
||||
The native `SeedVR2Pipeline` restores an input video using the released 3B FP16
|
||||
NaDiT and s8/c16/t4 causal VAE. It uses fixed checkpoint conditioning, CFG=1,
|
||||
and one Euler step. Text prompts do not change conditioning.
|
||||
|
||||
## Code ownership
|
||||
|
||||
The model package groups each computation stage in one file: `nadit.py` owns
|
||||
window geometry, routing, RoPE and the DiT; `vae.py` owns the causal VAE and its
|
||||
temporal/spatial tiling; `pipeline_seedvr2.py` owns input validation, decoding,
|
||||
conditioning, the Euler step and colour transfer. `config.py` owns model limits,
|
||||
and `long_video.py` owns the optional HTTP job/window lifecycle. Framework code
|
||||
only registers the model, its pre/postprocessors, metadata and HTTP routes.
|
||||
|
||||
`color_correction_method` is model-specific: online requests use the existing
|
||||
`extra_params` JSON object, and offline sampling uses
|
||||
`extra_args={"color_correction_method": "wavelet"}`. The default remains `lab`;
|
||||
invalid values are rejected by model-side validation before inference. It is
|
||||
not a top-level field in the common video API.
|
||||
|
||||
## Model directory
|
||||
|
||||
Place these files together, retaining their names:
|
||||
|
||||
| File | SHA-256 of the validated release |
|
||||
| --- | --- |
|
||||
| `seedvr2_ema_3b_fp16.safetensors` | `2fd0e03a3dad24e07086750360727ca437de4ecd456f769856e960ae93e2b304` |
|
||||
| `ema_vae_fp16.safetensors` | `20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1` |
|
||||
| `pos_emb.pt` | `fa07a14844314772266b66c3b95deb0027696d8fe7065721263db5176f45d799` |
|
||||
|
||||
Add `model_index.json` containing:
|
||||
|
||||
```json
|
||||
{"_class_name": "SeedVR2Pipeline"}
|
||||
```
|
||||
|
||||
The loader checks every checkpoint tensor, including RoPE buffers, and rejects
|
||||
missing, duplicate, unexpected, or incompatible tensors. Install PyAV for video
|
||||
decoding. The validated environment uses PyAV 18.1, PyTorch 2.13/CUDA 13, and L20 GPUs.
|
||||
|
||||
## Serve and restore a video
|
||||
|
||||
```bash
|
||||
vllm serve "$MODEL_DIR" --omni \
|
||||
--model-class-name SeedVR2Pipeline --dtype float16 --enforce-eager \
|
||||
--num-gpus 1 --host 127.0.0.1 --port 8098
|
||||
|
||||
curl --fail-with-body http://127.0.0.1:8098/v1/videos/sync \
|
||||
-F 'prompt= ' \
|
||||
-F 'input_references=@input.mp4;type=video/mp4' \
|
||||
-F 'size=224x128' -F 'num_inference_steps=1' \
|
||||
-F 'guidance_scale=1' -F 'seed=7723' \
|
||||
--output restored.mp4
|
||||
```
|
||||
|
||||
For two-rank window sequence parallelism, expose two GPUs and replace `--num-gpus 1` with:
|
||||
|
||||
```bash
|
||||
--num-gpus 2 --distributed-executor-backend mp \
|
||||
--stage-overrides '{"0":{"ulysses_degree":2}}'
|
||||
```
|
||||
|
||||
Use degree 4 and four visible GPUs for SP4. Window-SP alone replicates the VAE.
|
||||
To shard VAE activations across the same ranks, use:
|
||||
|
||||
```bash
|
||||
--vae-use-tiling --num-gpus 2 --distributed-executor-backend mp \
|
||||
--stage-overrides '{"0":{"ulysses_degree":2,"vae_patch_parallel_size":2,"vae_parallel_mode":"spatial_shard_height"}}'
|
||||
```
|
||||
|
||||
The VAE patch degree must match the window-SP degree.
|
||||
For a 16-frame 1280×720 input restored at its original size, use four GPUs,
|
||||
VAE tiling, and height sharding with both degrees set to four:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "$MODEL_DIR" --omni \
|
||||
--model-class-name SeedVR2Pipeline --dtype float16 --enforce-eager \
|
||||
--vae-use-tiling --num-gpus 4 --distributed-executor-backend mp \
|
||||
--stage-overrides '{"0":{"ulysses_degree":4,"vae_patch_parallel_size":4,"vae_parallel_mode":"spatial_shard_height"}}' \
|
||||
--host 127.0.0.1 --port 8098
|
||||
```
|
||||
|
||||
Request `size=1280x720` for this landscape input or `size=720x1280` for a
|
||||
portrait input of the same dimensions. Keep the requested output size equal to
|
||||
the input size when testing restoration without upscaling.
|
||||
|
||||
The multipart API requires a prompt field; a single space supplies a blank
|
||||
prompt. Nonblank text is rejected. Omit `fps` to retain the source frame rate;
|
||||
an explicit rate must match the source. Output dimensions are explicit multiples
|
||||
of 16. Input frames are resized with antialiased bicubic interpolation. The final
|
||||
frame is repeated internally to reach 4n+1, and the decoded output is cropped back
|
||||
to the original frame count. Five frames remain five; six frames are internally
|
||||
padded to nine and return six.
|
||||
|
||||
## Colour correction
|
||||
|
||||
Restoration reproduces detail faithfully but shifts global colour, so the
|
||||
resized input is used to transfer colour back onto the restored frames. This is
|
||||
on by default; pass `extra_params={"color_correction_method":"wavelet"}` to choose how:
|
||||
|
||||
| Value | Behavior |
|
||||
| --- | --- |
|
||||
| `lab` (default) | Swaps the lowest frequency band, then histogram-matches CIELAB chroma. Most faithful colour. |
|
||||
| `wavelet` | Swaps only the lowest frequency band, keeping every higher band from the restoration. |
|
||||
| `adain` | Matches per-channel mean and standard deviation. Cheapest, and corrects only a global tint. |
|
||||
| `none` | Skips the transfer and returns raw restored colour. |
|
||||
|
||||
```bash
|
||||
curl --fail-with-body http://127.0.0.1:8098/v1/videos/sync \
|
||||
-F 'prompt= ' -F 'input_references=@input.mp4;type=video/mp4' \
|
||||
-F 'size=224x128' -F 'num_inference_steps=1' -F 'guidance_scale=1' \
|
||||
-F 'seed=7723' -F 'extra_params={"color_correction_method":"wavelet"}' \
|
||||
--output restored.mp4
|
||||
```
|
||||
|
||||
`lab` and `wavelet` retain the restored high-frequency detail; `adain` measurably
|
||||
softens it because it rescales every frequency. The transfer runs in FP32 after
|
||||
the VAE decode and costs a fraction of a second per frame. Unrecognized values
|
||||
are rejected with 400 before the clip is admitted. Correction cannot recover
|
||||
colour lost to 4:2:0 chroma subsampling in the returned MP4, which dominates the
|
||||
remaining deviation once the transfer is applied.
|
||||
|
||||
## Request admission
|
||||
|
||||
The default 3B limit is 848×480 pixels per input or output frame and 2,035,200
|
||||
pixels across the requested output clip. With VAE tiling and window SP of at
|
||||
least four ranks whose `vae_patch_parallel_size` matches `ulysses_degree`, the
|
||||
per-frame limit rises to 2560×1472 and the clip budget to 18,841,600 pixels.
|
||||
Raising the degree never lowers the budget. Both budgets count temporal
|
||||
padding to 4n+1 frames: the default admits up to five 848×480 frames or 93
|
||||
192×112 frames; the sharded profile admits up to 45 848×480 frames or five
|
||||
2560×1472 frames. A 16-frame 1280×720 or 720×1280 clip pads to 17 frames and
|
||||
uses 15,667,200 of the sharded clip's 18,841,600 pixels. Both orientations passed
|
||||
original-size HTTP restoration on four RTX 5090 GPUs with VAE tiling and height
|
||||
sharding. An independent 257-frame cap bounds per-frame decoder work for tiny
|
||||
inputs. The other longer combinations above are admission bounds, not completed
|
||||
GPU validation. The decoder checks
|
||||
declared duration, frame count, and input dimensions when available, then
|
||||
enforces the limits as frames arrive. Requests outside these budgets return 400
|
||||
before building the resized whole-clip tensor.
|
||||
|
||||
The per-frame cap is calibrated for the smallest qualified device. The clip cap
|
||||
scales with the smallest visible device's memory: it stays at the calibrated
|
||||
five 2560×1472 frames on a 32 GB device and grows on larger ones from a
|
||||
per-rank memory model measured at SP4 (see `sharded_budget` in
|
||||
`vllm_omni/diffusion/models/seedvr2/pipeline_seedvr2.py`), which keeps 15% of the device
|
||||
free and leaves room for allocator fragmentation. It stops at 2,118,057,984
|
||||
padded pixels (513 frames at 1536×2688), the largest clip validated on four
|
||||
B300 ranks; 80 GB devices get about 1.08 billion and 141 GB or larger devices
|
||||
reach the cap. Clips at the cap in both memory regimes (561 frames at 2560×1472
|
||||
and 2,049 frames at 768×1344) peaked at 54 GiB of activations per rank. Set
|
||||
these before starting the server to override the caps; each must be a positive
|
||||
integer, and validating an override on the target hardware is the operator's
|
||||
responsibility:
|
||||
|
||||
| Variable | Default | Bounds |
|
||||
| --- | --- | --- |
|
||||
| `VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS` | 3,768,320 | Pixels per frame on the sharded profile |
|
||||
| `VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS` | 18,841,600, scaled up with device memory | Padded pixels per clip on the sharded profile |
|
||||
| `VLLM_OMNI_SEEDVR2_MAX_FRAMES` | 257 | Decoder-work frame cap, all profiles |
|
||||
|
||||
The sharded profile applies to any `ulysses_degree` of four or more, so eight
|
||||
ranks inherit a budget calibrated on four. Raising these caps trades a 400 for a
|
||||
CUDA OOM: the admission check passes and the request fails later inside the DiT
|
||||
forward or the VAE decode, which can abort every in-flight request on the engine
|
||||
rather than only the oversized one. Before serving with raised caps, restore the
|
||||
largest clip they admit once and confirm it completes.
|
||||
|
||||
A 362-frame 1536×2688 2x upscale was restored on eight ranks with
|
||||
`ulysses_degree=8`, VAE tiling and height sharding, using 4,300,000 for the
|
||||
per-frame cap, the memory-scaled clip cap and 2,000 for the frame cap. The long-video
|
||||
route sizes its windows from these caps, so the clip cap sets how many frames
|
||||
each window covers.
|
||||
|
||||
## Long-video restoration
|
||||
|
||||
`POST /v1/seedvr2/restore-long` accepts one uploaded 24 FPS video, a blank
|
||||
prompt, `size`, `num_frames` (up to 7,200), and optional `loop_input=true` and
|
||||
`extra_params={"color_correction_method":"wavelet"}`, which it applies to every window.
|
||||
An output frame must fit the sharded per-frame pixel cap. It returns a job ID; poll
|
||||
`GET /v1/seedvr2/restore-long/{id}` and download the completed MP4 from
|
||||
`GET /v1/seedvr2/restore-long/{id}/content`. `DELETE
|
||||
/v1/seedvr2/restore-long/{id}` asks a running job to stop; it settles as
|
||||
`cancelled` at the next window boundary. A job whose owning process is gone,
|
||||
such as after a server restart, reports `failed` on the next poll.
|
||||
|
||||
Jobs live on the API server's local disk, so the route requires
|
||||
`--api-server-count 1` and runs one job at a time. Storage is bounded by an
|
||||
upload cap and by deleting settled jobs once they age out:
|
||||
|
||||
| Variable | Default | Bounds |
|
||||
| --- | --- | --- |
|
||||
| `VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR` | system temp dir | Parent of the job directories |
|
||||
| `VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES` | 8 GiB | Largest accepted upload |
|
||||
| `VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS` | 3,600 | Age at which a settled job is deleted |
|
||||
| `VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW` | 121 | Longest model window, in frames |
|
||||
|
||||
Download a completed job before its TTL expires; the sweep runs on each new
|
||||
submission and removes the output with the job directory.
|
||||
|
||||
The service sends windows through the existing SeedVR2 endpoint, blends four
|
||||
frames at each boundary, and writes one continuous MP4 encoder. Each window is
|
||||
the longest 4n+1-frame clip that the sharded clip-pixel cap, the frame cap and
|
||||
`VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW` admit at the output size. Every window pays
|
||||
a fixed round trip through the serving stack and recomputes its overlap, and a
|
||||
window of four latent frames gets no temporal attention in the DiT, so raise
|
||||
the clip cap as far as the device allows. Windows carry source pixels, which
|
||||
the model resizes on the device, and cross the endpoint losslessly; only the
|
||||
final MP4 is compressed. The next window is restored while the previous one is
|
||||
blended and encoded. It repeats source frames and audio only when
|
||||
`loop_input=true`. The sharded serving profile above (VAE tiling, VAE height
|
||||
sharding, Ulysses degree of at least four) is required. This route keeps a
|
||||
window within the whole-clip pixel budget; it does not raise the
|
||||
`/v1/videos/sync` 257-frame or pixel limits. The long route
|
||||
uses fixed one-step, CFG=1 conditioning and requires `imageio[ffmpeg]` for
|
||||
audio muxing. The 7,200-frame 768×1344 case is still undergoing full GPU E2E.
|
||||
|
||||
## Temporal and spatial VAE tiling
|
||||
|
||||
Add `--vae-use-tiling` to bound VAE intermediate activations along time. The
|
||||
encoder processes nine frames first, then eight per chunk; the decoder processes
|
||||
two latent frames per chunk. Each causal convolution carries its past inputs
|
||||
within that encode/decode call, with temporal stride and first-frame upsampling
|
||||
alignment preserved. GroupNorm and attention still see the entire spatial frame.
|
||||
There is no overlap blend or independent restart at chunk boundaries.
|
||||
|
||||
Five-frame clips use one temporal chunk. On a single GPU, convolutions with
|
||||
more than 256 input rows additionally process tiles of 128 output rows with
|
||||
exact halos. GroupNorm and bottleneck attention retain full-frame statistics.
|
||||
This bounds convolution workspace; full-frame activations and the DiT's
|
||||
whole-clip activations still consume memory.
|
||||
|
||||
With VAE patch parallelism, each rank owns a band of latent rows and the
|
||||
corresponding encoder/decoder rows. Neighbor exchanges supply convolution halos;
|
||||
all-reduced centered statistics preserve full-frame GroupNorm. Bottleneck
|
||||
attention gathers the low-resolution frame before splitting it again. Encoder
|
||||
parameters are gathered before latent sampling, and decoded pixels are gathered
|
||||
before returning the video. Unequal bands are supported; fewer latent rows than
|
||||
ranks fall back to replicated execution. Width sharding and batch slicing are
|
||||
unsupported. FP16 reduction and convolution order can change rounding, so
|
||||
parallel outputs are numerically close rather than bitwise identical.
|
||||
|
||||
## Current integration scope
|
||||
|
||||
| Property | Contract |
|
||||
| --- | --- |
|
||||
| Weights / dtype | Released 3B DiT and VAE, FP16 |
|
||||
| Reference semantics | C0 whole clip, colour correction on by default |
|
||||
| Video input | One uploaded file, or offline TCHW RGB floats / PIL frames |
|
||||
| Timing | Constant frame rate, increasing PTS; normalize the video origin to zero |
|
||||
| Audio | First mono/stereo track, aligned by source PTS, cropped to the video interval, re-encoded as AAC |
|
||||
| Randomness | Per-request generator; preserve reference latent strides when sampling noise |
|
||||
| Sequence parallelism | Native engine SP1/2/4 correctness verified on five-frame L20 cases |
|
||||
| VAE placement | Replicated by default; optional height sharding on the window-SP group |
|
||||
| Unsupported | VFR, multichannel audio, 7B, other sampling schedules, quantization, cache acceleration, VAE width sharding / batch slicing, CPU offload, CFG/TP/PP parallelism, compiled execution, LoRA |
|
||||
|
||||
Unsupported engine modes are rejected before process hooks and worker creation.
|
||||
`ulysses_degree` selects the model-owned SP group. This branch assigns whole
|
||||
windows to ranks; the specialized head-sharded path is a dependent change.
|
||||
Ring and AllGather-KV are unsupported.
|
||||
|
||||
The practical P0 reference's five-frame batching, overlap, and CPU swapping are
|
||||
separate execution semantics. Temporal tiling and VAE patch
|
||||
parallelism reduce activation peaks but do not bound the whole-clip DiT memory.
|
||||
Large output frames may still exceed device capacity.
|
||||
|
||||
Invalid inputs return 400. A single-GPU model OOM fails that request; a
|
||||
subsequent short request succeeds. With VAE patch parallelism, restart the
|
||||
service after an OOM: another rank may still be waiting in a collective. Cancelling an in-progress video job removes it through the
|
||||
existing DELETE endpoint; in-flight GPU work may drain before memory is reusable.
|
||||
A lost SP worker fails the request and makes health return 503. Restart the
|
||||
service to restore availability; automatic rank recovery is not provided.
|
||||
|
||||
## Reproducible deployment
|
||||
|
||||
See the [RTX 5090 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/ByteDance/SeedVR2-RTX-5090.md) for the
|
||||
input/output contract, complete serving command, and media checks. The local
|
||||
tests in `tests/diffusion/models/seedvr2/test_seedvr2_e2e.py` exercise the full
|
||||
3B pipeline, all colour modes, and HTTP restoration with USP 1 and 8.
|
||||
@@ -54,6 +54,7 @@ th {
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
|
||||
| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
|
||||
| `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
|
||||
| `SeedVR2Pipeline` | [SeedVR2 3B](seedvr2.md) video restoration | `ByteDance-Seed/SeedVR2-3B` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/ByteDance/SeedVR2-RTX-5090.md) |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
|
||||
| `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
# SeedVR2 video restoration — RTX 5090
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: ByteDance
|
||||
- Model: SeedVR2 3B
|
||||
- Task: restore or upscale an uploaded video
|
||||
- Mode: native offline/HTTP pipeline, one Euler step
|
||||
- Hardware: NVIDIA GeForce RTX 5090, 32 GiB per device
|
||||
- Maintainer: 0z5a
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Run short, constant-frame-rate videos with the native SeedVR2 pipeline. This
|
||||
recipe covers the whole-clip path and the separate long-video route. Color
|
||||
correction from external applications is a separate execution semantic.
|
||||
|
||||
## Supported model contract
|
||||
|
||||
| Property | Contract |
|
||||
| --- | --- |
|
||||
| Model directory | See [checkpoint files and hashes](../../docs/models/seedvr2.md#model-directory) |
|
||||
| Input | One uploaded video; blank prompt; positive dimensions divisible by 16 |
|
||||
| Output | MP4 at requested geometry, original frame count and source frame rate |
|
||||
| Audio | First mono/stereo track, aligned by PTS; re-encoded as AAC |
|
||||
| Sampling | One Euler step, CFG=1, per-request seed |
|
||||
| Precision | 3B DiT and VAE FP16 |
|
||||
| Parallelism | Whole-window SP via `ulysses_degree`; replicated model weights |
|
||||
| Frame padding | Internal 4n+1 padding, cropped back; five returns five, six returns six |
|
||||
|
||||
## References
|
||||
|
||||
- [Canonical SeedVR](https://github.com/ByteDance-Seed/SeedVR)
|
||||
- [Model guide](../../docs/models/seedvr2.md)
|
||||
- [Integration RFC](https://github.com/vllm-project/vllm-omni/issues/7723)
|
||||
|
||||
## Hardware
|
||||
|
||||
RTX 5090 with 32 GiB per GPU, PCIe, one GPU for the default command or two/four
|
||||
for SP. Each GPU must fit a complete model copy and its activations. Host RAM
|
||||
must accommodate checkpoint staging per rank. Admission uses a padded
|
||||
five-frame 848×480 output pixel budget by default, allowing longer clips at
|
||||
smaller resolutions, plus a 257-frame decoder-work cap. The validated SP4
|
||||
configuration with VAE tiling and `vae_patch_parallel_size=4` uses a padded
|
||||
five-frame 2560×1472 budget, as does any higher matched degree. On devices
|
||||
with more memory the clip budget grows automatically; the other caps are raised
|
||||
through `VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS` and
|
||||
`VLLM_OMNI_SEEDVR2_MAX_FRAMES`, and `VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS`
|
||||
overrides the clip budget. The model guide lists
|
||||
exact bounds and which profiles have completed GPU validation.
|
||||
|
||||
## Software environment
|
||||
|
||||
Linux x86-64, Python 3.12, PyTorch 2.13.0+cu130, CUDA 13.0, vLLM 0.29.0, and
|
||||
PyAV in the existing environment. Use the vLLM-Omni branch containing this recipe.
|
||||
|
||||
## Command
|
||||
|
||||
Set `MODEL_DIR` to the authorized checkpoint directory documented in the model
|
||||
guide; it must also contain `model_index.json` selecting `SeedVR2Pipeline`.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_DIR" --omni \
|
||||
--model-class-name SeedVR2Pipeline --dtype float16 --enforce-eager \
|
||||
--num-gpus 1 --host 127.0.0.1 --port 8098
|
||||
```
|
||||
|
||||
For SP=2, expose two GPUs and replace the GPU count with:
|
||||
|
||||
```bash
|
||||
--num-gpus 2 --distributed-executor-backend mp \
|
||||
--stage-overrides '{"0":{"ulysses_degree":2}}'
|
||||
```
|
||||
|
||||
Use degree 4 and four visible GPUs for SP=4. Combine any feature-specific
|
||||
settings below into the same stage-overrides object; do not repeat that flag.
|
||||
The validated 16-frame 720p original-size profile is:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "$MODEL_DIR" --omni \
|
||||
--model-class-name SeedVR2Pipeline --dtype float16 --enforce-eager \
|
||||
--vae-use-tiling --num-gpus 4 --distributed-executor-backend mp \
|
||||
--stage-overrides '{"0":{"ulysses_degree":4,"vae_patch_parallel_size":4,"vae_parallel_mode":"spatial_shard_height"}}' \
|
||||
--host 127.0.0.1 --port 8098
|
||||
```
|
||||
|
||||
Use `size=1280x720` for a landscape 1280×720 input or `size=720x1280` for a
|
||||
portrait 720×1280 input. Both require 16 decoded source frames and retain the
|
||||
source frame rate; no upscaling is requested.
|
||||
|
||||
## Verification
|
||||
|
||||
```bash
|
||||
curl --fail-with-body http://127.0.0.1:8098/v1/videos/sync \
|
||||
-F 'prompt= ' -F 'input_references=@input.mp4;type=video/mp4' \
|
||||
-F 'size=224x128' -F 'num_inference_steps=1' \
|
||||
-F 'guidance_scale=1' -F 'seed=7723' --output restored.mp4
|
||||
python - <<'PY'
|
||||
import av
|
||||
with av.open('restored.mp4') as video:
|
||||
stream = video.streams.video[0]
|
||||
print(stream.width, stream.height, stream.average_rate,
|
||||
sum(1 for _ in video.decode(video=0)), len(video.streams.audio))
|
||||
PY
|
||||
```
|
||||
|
||||
Expect HTTP 200 and a completely decodable 224×128 video. Check the frame count,
|
||||
FPS, timestamps, and audio against the input. Repeat the same seed and compare
|
||||
decoded frames; changing the seed should change the restored video. Inspect
|
||||
matching first/middle/last frames at the same display scale. AAC is re-encoded,
|
||||
so input/output audio packet hashes need not be identical.
|
||||
|
||||
The released 3B checkpoint also passed these four-GPU SP4 HTTP cases with the
|
||||
configuration above. Each input had 16 frames at 24 FPS; the output retained
|
||||
the original size, 16 decoded frames, timestamps, and decodable audio. All four
|
||||
workers and the server exited normally after each run.
|
||||
|
||||
| 1× input → output | Earlier six-frame admission | Revised HTTP result | Request time |
|
||||
| --- | --- | --- | ---: |
|
||||
| 1280×720, 16 frames → 1280×720 | Rejected | HTTP 200; media checks passed | 3.553 s |
|
||||
| 720×1280, 16 frames → 720×1280 | Rejected | HTTP 200; media checks passed | 3.388 s |
|
||||
|
||||
These are single requests on a shared host, not performance comparisons.
|
||||
|
||||
### One long-video request
|
||||
|
||||
For a 24 FPS source, the long route can restore up to 7,200 frames at
|
||||
768×1344. Use the SP4/VAE height-sharding command above, serve with
|
||||
`--api-server-count 1`, and set a persistent
|
||||
`VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR` before starting. If the source is shorter
|
||||
than 300 seconds, `loop_input=true` repeats its frames and audio. The service
|
||||
runs the longest model windows the clip budget admits, with four-frame overlap,
|
||||
and writes one MP4;
|
||||
it does not concatenate separate video files. Download the result within
|
||||
`VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS`, after which the job directory is
|
||||
swept; `DELETE` on the job URL stops a run at the next window boundary.
|
||||
|
||||
```bash
|
||||
curl --fail-with-body http://127.0.0.1:8098/v1/seedvr2/restore-long \
|
||||
-F 'input_references=@input.mp4;type=video/mp4' \
|
||||
-F 'prompt= ' -F 'size=768x1344' -F 'num_frames=7200' \
|
||||
-F 'loop_input=true' -F 'seed=7723' -F 'extra_params={"color_correction_method":"lab"}'
|
||||
# Use the returned id to poll /v1/seedvr2/restore-long/{id};
|
||||
# download /v1/seedvr2/restore-long/{id}/content when completed.
|
||||
```
|
||||
|
||||
The input must have constant 24 FPS timestamps starting at zero. Only one job
|
||||
can run per API server, and `--api-server-count` must be one. The existing
|
||||
whole-clip `/v1/videos/sync` limits remain in force. Full 7,200-frame GPU
|
||||
validation is in progress; use this route only with a memory-checked SP4
|
||||
deployment.
|
||||
|
||||
Set `VLLM_TEST_SEEDVR2_MODEL_DIR` to the directory containing all model files:
|
||||
|
||||
```bash
|
||||
python -m pytest -o addopts='' -v tests/diffusion/models/seedvr2/test_seedvr2_e2e.py
|
||||
```
|
||||
|
||||
The tests run the complete native pipeline for every colour mode and the HTTP
|
||||
long-video route at USP 1 and 8 (eight GPUs required for the latter). The HTTP
|
||||
cases force a 13-frame window, restore 20 frames across a seam, preserve audio
|
||||
and timestamps, reject invalid `extra_params`, and cancel a longer job. They
|
||||
use the released model rather than mocking its forward pass.
|
||||
|
||||
## Supported features
|
||||
|
||||
| Feature | Status |
|
||||
| --- | --- |
|
||||
| [Window SP](../../docs/models/seedvr2.md) | Model-local regular/shifted window attention |
|
||||
| CPU offload, LoRA, compiled execution, CFG/TP/PP | Unsupported |
|
||||
| VFR or multichannel audio | Unsupported |
|
||||
| VAE temporal/spatial tiling | Available with `--vae-use-tiling` |
|
||||
| [Colour correction](../../docs/models/seedvr2.md) | On by default; `color_correction_method=lab\|wavelet\|adain\|none` |
|
||||
| Quantization | Separate opt-in feature |
|
||||
|
||||
## Measurement scope
|
||||
|
||||
Compare identical inputs, seeds, dimensions, GPU counts, and warmups. Separate
|
||||
correctness from performance: use repeated interleaved HTTP measurements before
|
||||
claiming a stable gain. Keep raw logs, tensors, and generated media outside the
|
||||
source diff; link selected visuals from the PR's test results.
|
||||
|
||||
## Temporal tiling and VAE patch parallelism
|
||||
|
||||
Use `--vae-use-tiling` for causal temporal chunks. On a single GPU this also
|
||||
bounds large convolution workspace with exact spatial halos. To shard VAE
|
||||
activations across the same two SP ranks, use:
|
||||
|
||||
```bash
|
||||
--vae-use-tiling --num-gpus 2 --distributed-executor-backend mp \
|
||||
--stage-overrides '{"0":{"ulysses_degree":2,"vae_patch_parallel_size":2,"vae_parallel_mode":"spatial_shard_height"}}'
|
||||
```
|
||||
|
||||
The VAE patch degree must equal the SP degree. Height sharding is supported;
|
||||
width sharding and batch slicing are not. Compare clips that cross a temporal
|
||||
chunk boundary as well as five/six-frame clips, and inspect frames adjacent to
|
||||
chunk boundaries. Reduced peak memory does not by itself establish lower latency.
|
||||
High-resolution clips can still exceed device capacity; validate the intended
|
||||
frame count and output size on the target GPUs.
|
||||
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 32,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 59,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 66,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 31,
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 38,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 16,
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""End-to-end native video restoration and the HTTP long-video route.
|
||||
|
||||
Set VLLM_TEST_SEEDVR2_MODEL_DIR to the directory containing the DiT, VAE,
|
||||
conditioning tensor, and model_index.json to run the complete video test.
|
||||
Each rank exits normally; the launcher does not terminate workers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from fractions import Fraction
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.local_model, pytest.mark.cuda, pytest.mark.diffusion, pytest.mark.parallel]
|
||||
MODEL_DIR_ENV = "VLLM_TEST_SEEDVR2_MODEL_DIR"
|
||||
# Two model windows with one four-frame seam, at the smallest legal frame size.
|
||||
LONG_FRAMES = 20
|
||||
LONG_SIZE = 64
|
||||
FPS = 24
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get(MODEL_DIR_ENV), reason=f"set {MODEL_DIR_ENV} to an authorized model directory")
|
||||
@pytest.mark.parametrize("method", ["lab", "wavelet", "adain", "none"])
|
||||
def test_seedvr2_native_video_e2e(method: str) -> None:
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
model_dir = Path(os.environ[MODEL_DIR_ENV]).resolve()
|
||||
for name in ("seedvr2_ema_3b_fp16.safetensors", "ema_vae_fp16.safetensors", "pos_emb.pt", "model_index.json"):
|
||||
assert (model_dir / name).is_file(), f"Missing SeedVR2 model file: {name}"
|
||||
frames = torch.rand(5, 3, 64, 112, generator=torch.Generator().manual_seed(7723))
|
||||
engine = Omni(
|
||||
model=str(model_dir),
|
||||
model_class_name="SeedVR2Pipeline",
|
||||
dtype="float16",
|
||||
enforce_eager=True,
|
||||
vae_use_tiling=True,
|
||||
)
|
||||
try:
|
||||
output = engine.generate(
|
||||
{"prompt": " ", "multi_modal_data": {"video": frames}},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=128,
|
||||
width=224,
|
||||
num_frames=5,
|
||||
fps=24,
|
||||
num_inference_steps=1,
|
||||
guidance_scale=1.0,
|
||||
seed=7723,
|
||||
output_type="np",
|
||||
extra_args={"color_correction_method": method},
|
||||
),
|
||||
use_tqdm=False,
|
||||
)
|
||||
finally:
|
||||
engine.close()
|
||||
assert len(output) == 1
|
||||
restored = np.asarray(output[0].images[0])
|
||||
assert restored.shape == (1, 5, 128, 224, 3)
|
||||
assert np.isfinite(restored).all()
|
||||
|
||||
|
||||
def _write_source(path: Path, frames: int) -> None:
|
||||
"""A constant-rate source with audio, which the route requires."""
|
||||
import imageio_ffmpeg
|
||||
|
||||
subprocess.run(
|
||||
[
|
||||
imageio_ffmpeg.get_ffmpeg_exe(),
|
||||
"-nostdin",
|
||||
"-y",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-f",
|
||||
"lavfi",
|
||||
"-i",
|
||||
f"testsrc=size={LONG_SIZE}x{LONG_SIZE}:rate={FPS}",
|
||||
"-f",
|
||||
"lavfi",
|
||||
"-i",
|
||||
"sine=frequency=440",
|
||||
# The generators are endless, so bound the video by frame count and
|
||||
# the audio by a slightly longer wall time.
|
||||
"-frames:v",
|
||||
str(frames),
|
||||
"-t",
|
||||
str((frames + 1) / FPS),
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-pix_fmt",
|
||||
"yuv420p",
|
||||
"-c:a",
|
||||
"aac",
|
||||
str(path),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def _await_status(base: str, job_id: str, timeout: float) -> dict:
|
||||
"""Poll until the job leaves the queued/running states."""
|
||||
import requests
|
||||
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
record = requests.get(f"{base}/{job_id}", timeout=30).json()
|
||||
if record["status"] not in {"queued", "running"}:
|
||||
return record
|
||||
time.sleep(2)
|
||||
raise AssertionError(f"SeedVR2 long-video job stayed {record['status']} for {timeout}s")
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get(MODEL_DIR_ENV), reason=f"set {MODEL_DIR_ENV} to an authorized model directory")
|
||||
@pytest.mark.parametrize("degree", [1, 8])
|
||||
def test_seedvr2_long_video_route_e2e(tmp_path: Path, degree: int) -> None:
|
||||
"""Restore a clip spanning two model windows, then cancel a longer job."""
|
||||
import av
|
||||
import requests
|
||||
import torch
|
||||
|
||||
if torch.accelerator.device_count() < degree:
|
||||
pytest.skip(f"SeedVR2 USP={degree} requires {degree} CUDA devices")
|
||||
|
||||
from tests.helpers.runtime import OmniServer
|
||||
from vllm_omni.diffusion.models.seedvr2.long_video import MAX_FRAMES
|
||||
|
||||
model_dir = Path(os.environ[MODEL_DIR_ENV]).resolve()
|
||||
source = tmp_path / "input.mp4"
|
||||
_write_source(source, LONG_FRAMES)
|
||||
server = OmniServer(
|
||||
model=str(model_dir),
|
||||
serve_args=[
|
||||
"--model-class-name",
|
||||
"SeedVR2Pipeline",
|
||||
"--dtype",
|
||||
"float16",
|
||||
"--enforce-eager",
|
||||
"--num-gpus",
|
||||
str(degree),
|
||||
"--vae-use-tiling",
|
||||
"--stage-overrides",
|
||||
json.dumps(
|
||||
{
|
||||
"0": {
|
||||
"ulysses_degree": degree,
|
||||
"vae_patch_parallel_size": degree,
|
||||
"vae_parallel_mode": "spatial_shard_height",
|
||||
}
|
||||
}
|
||||
),
|
||||
"--api-server-count",
|
||||
"1",
|
||||
],
|
||||
env_dict={
|
||||
"VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR": str(tmp_path / "jobs"),
|
||||
"VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW": "13",
|
||||
},
|
||||
)
|
||||
|
||||
with server:
|
||||
base = f"http://{server.host}:{server.port}/v1/seedvr2/restore-long"
|
||||
|
||||
def submit(**overrides: str) -> requests.Response:
|
||||
with source.open("rb") as upload:
|
||||
return requests.post(
|
||||
base,
|
||||
data={
|
||||
"prompt": " ",
|
||||
"size": f"{LONG_SIZE}x{LONG_SIZE}",
|
||||
"num_frames": str(LONG_FRAMES),
|
||||
"seed": "7723",
|
||||
**overrides,
|
||||
},
|
||||
files={"input_references": ("input.mp4", upload, "video/mp4")},
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
assert submit(num_frames=str(MAX_FRAMES + 1)).status_code == 400
|
||||
|
||||
# Validation must happen before a job is admitted, through model-owned
|
||||
# extra_params on both the long route and the common video endpoint.
|
||||
for invalid in ("{", "[]", '{"color_correction_method":"invalid"}'):
|
||||
assert submit(extra_params=invalid).status_code == 400
|
||||
with source.open("rb") as upload:
|
||||
invalid = requests.post(
|
||||
f"http://{server.host}:{server.port}/v1/videos/sync",
|
||||
data={
|
||||
"prompt": " ",
|
||||
"size": f"{LONG_SIZE}x{LONG_SIZE}",
|
||||
"num_inference_steps": "1",
|
||||
"guidance_scale": "1",
|
||||
"extra_params": '{"color_correction_method":"invalid"}',
|
||||
},
|
||||
files={"input_references": ("input.mp4", upload, "video/mp4")},
|
||||
timeout=120,
|
||||
)
|
||||
assert invalid.status_code == 400, invalid.text
|
||||
accepted = submit(extra_params='{"color_correction_method":"wavelet"}')
|
||||
assert accepted.status_code == 202, accepted.text
|
||||
job_id = accepted.json()["id"]
|
||||
record = _await_status(base, job_id, timeout=900)
|
||||
assert record["status"] == "completed", record["error"]
|
||||
assert record["frames"] == LONG_FRAMES
|
||||
|
||||
restored = tmp_path / "restored.mp4"
|
||||
content = requests.get(f"{base}/{job_id}/content", timeout=300)
|
||||
assert content.status_code == 200
|
||||
restored.write_bytes(content.content)
|
||||
with av.open(str(restored)) as container:
|
||||
video = container.streams.video[0]
|
||||
timestamps = [frame.pts * frame.time_base for frame in container.decode(video=0)]
|
||||
assert (video.width, video.height) == (LONG_SIZE, LONG_SIZE)
|
||||
assert video.average_rate == Fraction(FPS)
|
||||
assert timestamps == [Fraction(index, FPS) for index in range(LONG_FRAMES)]
|
||||
assert container.streams.audio, "the route must carry the source audio through"
|
||||
|
||||
# Many windows, so the cancel lands well before the job could finish.
|
||||
running = submit(num_frames="600", loop_input="true")
|
||||
assert running.status_code == 202, running.text
|
||||
cancelled_id = running.json()["id"]
|
||||
assert submit().status_code == 409, "the route runs one job per server"
|
||||
assert requests.delete(f"{base}/{cancelled_id}", timeout=30).status_code == 202
|
||||
settled = _await_status(base, cancelled_id, timeout=900)
|
||||
assert settled["status"] == "cancelled", settled
|
||||
assert settled["frames"] < 600
|
||||
assert requests.get(f"{base}/{cancelled_id}/content", timeout=30).status_code == 404
|
||||
@@ -388,8 +388,13 @@ def isolated_video_backends(tmp_path, monkeypatch):
|
||||
return store, tasks, storage
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class_name", [None, "SeedVR2Pipeline", "auto"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch):
|
||||
async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch, model_class_name):
|
||||
# Resolve imports outside the timed server-lifecycle assertions.
|
||||
from vllm_omni.diffusion.models.seedvr2 import long_video # noqa: F401
|
||||
|
||||
app = FastAPI()
|
||||
events: list[str] = []
|
||||
serve_started = asyncio.Event()
|
||||
http_shutdown = asyncio.Event()
|
||||
@@ -399,6 +404,11 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
|
||||
class FakeEngine:
|
||||
stage_configs = []
|
||||
|
||||
def get_diffusion_od_config(self):
|
||||
if model_class_name == "auto":
|
||||
return SimpleNamespace(model_class_name="SeedVR2Pipeline")
|
||||
return None
|
||||
|
||||
async def get_supported_tasks(self):
|
||||
return ("generate",)
|
||||
|
||||
@@ -435,7 +445,7 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
|
||||
events.append("init_app_state")
|
||||
|
||||
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
|
||||
monkeypatch.setattr(api_server, "build_openai_app", lambda args, supported_tasks: FastAPI())
|
||||
monkeypatch.setattr(api_server, "build_openai_app", lambda args, supported_tasks: app)
|
||||
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
|
||||
monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
|
||||
monkeypatch.setattr(api_server.openai_app_state, "_get_vllm_config", fake_get_vllm_config)
|
||||
@@ -461,9 +471,16 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
|
||||
h11_max_header_count=None,
|
||||
)
|
||||
|
||||
if model_class_name != "auto" and model_class_name is not None:
|
||||
args.model_class_name = model_class_name
|
||||
|
||||
worker_task = asyncio.create_task(api_server.omni_run_server_worker("127.0.0.1:0", sock, args))
|
||||
await asyncio.wait_for(serve_started.wait(), timeout=2)
|
||||
|
||||
route_paths = app.openapi()["paths"]
|
||||
assert ("/v1/seedvr2/restore-long" in route_paths) == (model_class_name is not None)
|
||||
if model_class_name is not None:
|
||||
assert app.state.seedvr2_long_port == args.port
|
||||
assert not engine_context_exited.is_set()
|
||||
|
||||
http_shutdown.set()
|
||||
@@ -2646,6 +2663,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"use_cfg_zero_star": True,
|
||||
"color_correction_method": "wavelet",
|
||||
}
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
@@ -2664,6 +2682,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
|
||||
assert captured.extra_args["pyramid_num_stages"] == 3
|
||||
assert captured.extra_args["pyramid_num_inference_steps_list"] == [1, 1, 1]
|
||||
assert captured.extra_args["use_cfg_zero_star"] is True
|
||||
assert captured.extra_args["color_correction_method"] == "wavelet"
|
||||
|
||||
|
||||
def test_extra_params_none_by_default(test_client, mocker: MockerFixture):
|
||||
|
||||
@@ -32,6 +32,8 @@ EXCLUDED_MODELS = [
|
||||
# dedicated MammothModa2 bridge/pipeline tests and real-checkpoint E2E.
|
||||
"MammothModa2DiTPipeline",
|
||||
"WanPipeline",
|
||||
# Uses a full checkpoint and video input; covered by its checkpoint-gated E2E test.
|
||||
"SeedVR2Pipeline",
|
||||
"WanDMDPipeline",
|
||||
"WanVACEPipeline",
|
||||
"LTX2TwoStagePipeline",
|
||||
|
||||
@@ -174,6 +174,13 @@ _MODEL_PROMOTE = (
|
||||
"VLLM_GEPARD_LOOKBACK_FRAMES",
|
||||
"VLLM_OMNI_FISH_KVCACHE_ATTN",
|
||||
"VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
|
||||
"VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS",
|
||||
"VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES",
|
||||
"VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW",
|
||||
"VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR",
|
||||
"VLLM_OMNI_SEEDVR2_MAX_FRAMES",
|
||||
"VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS",
|
||||
"VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS",
|
||||
"VLLM_OMNI_SENSENOVA_PAGED_DECODE",
|
||||
)
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ BOOGU_IMAGE_MAX_INPUT_IMAGES = 1
|
||||
|
||||
|
||||
_DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
|
||||
"SeedVR2Pipeline": DiffusionModelMetadata(supports_multimodal_inputs=True, final_output_type="video"),
|
||||
"Flux2KleinPipeline": DiffusionModelMetadata(
|
||||
supports_multimodal_inputs=True,
|
||||
max_multimodal_image_inputs=FLUX2_KLEIN_MAX_INPUT_IMAGES,
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Native SeedVR2 video restoration."""
|
||||
@@ -0,0 +1,59 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Operator limits and colour options owned by SeedVR2."""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
COLOR_CORRECTION_METHODS = ("lab", "wavelet", "adain", "none")
|
||||
DEFAULT_COLOR_CORRECTION_METHOD = "lab"
|
||||
|
||||
|
||||
def _positive_int(name: str, default: int | None) -> int | None:
|
||||
"""Read an operator-tuned limit, rejecting values that would disable it."""
|
||||
value = os.environ.get(name)
|
||||
if value is None:
|
||||
return default
|
||||
if not value.strip().isdigit() or not int(value):
|
||||
raise ValueError(f"{name} must be a positive integer, got {value!r}")
|
||||
return int(value)
|
||||
|
||||
|
||||
environment_variables: dict[str, Callable[[], Any]] = {
|
||||
# ================== SeedVR2 Restoration Env Vars ==================
|
||||
# Admission budgets for the sharded serving profile. The defaults are
|
||||
# calibrated for the smallest qualified device, so larger accelerators raise
|
||||
# them here; setting them beyond what the device holds surfaces as a CUDA
|
||||
# OOM during the forward pass. See docs/models/seedvr2.md.
|
||||
"VLLM_OMNI_SEEDVR2_MAX_FRAMES": lambda: _positive_int("VLLM_OMNI_SEEDVR2_MAX_FRAMES", 257),
|
||||
"VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS": lambda: _positive_int(
|
||||
"VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS", 2560 * 1472
|
||||
),
|
||||
# Unset, the clip budget scales with device memory from the calibrated
|
||||
# 32 GB value; see seedvr2.pipeline_seedvr2.sharded_budget.
|
||||
"VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS": lambda: _positive_int("VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS", None),
|
||||
# Bounds for the long-video restoration route.
|
||||
"VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES": lambda: _positive_int(
|
||||
"VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES", 8 * 1024**3
|
||||
),
|
||||
# Longest model window. The clip budget usually binds first; this caps it
|
||||
# at 30 latent frames, the DiT's largest temporal attention window.
|
||||
"VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW": lambda: _positive_int("VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW", 121),
|
||||
"VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS": lambda: _positive_int("VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS", 3600),
|
||||
"VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR": lambda: os.environ.get(
|
||||
"VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR", tempfile.gettempdir()
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
# lazy evaluation of environment variables
|
||||
if name in environment_variables:
|
||||
return environment_variables[name]()
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
return list(environment_variables.keys())
|
||||
@@ -0,0 +1,502 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""One asynchronous SeedVR2 request, bounded model windows, one MP4 writer."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import io
|
||||
import itertools
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from fractions import Fraction
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import av
|
||||
import imageio_ffmpeg
|
||||
import numpy as np
|
||||
import regex as re
|
||||
import requests
|
||||
from fastapi import APIRouter, FastAPI, File, Form, HTTPException, Request, UploadFile
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from vllm_omni.diffusion.models.seedvr2 import config as envs
|
||||
from vllm_omni.diffusion.models.seedvr2.config import COLOR_CORRECTION_METHODS, DEFAULT_COLOR_CORRECTION_METHOD
|
||||
from vllm_omni.diffusion.models.seedvr2.pipeline_seedvr2 import max_frames, sharded_budget
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter()
|
||||
MAX_FRAMES = 7200
|
||||
FPS = 24
|
||||
OVERLAP = 4
|
||||
# Windows cross the HTTP boundary losslessly (x264 qp 0, which is also its
|
||||
# fastest mode); only the final MP4 is compressed.
|
||||
LOSSLESS = {"preset": "ultrafast", "qp": "0"}
|
||||
UPLOAD_CHUNK = 1 << 20
|
||||
JOB_ID = re.compile(r"[0-9a-f]{32}")
|
||||
active_task: asyncio.Task[None] | None = None
|
||||
job_lock = asyncio.Lock()
|
||||
|
||||
|
||||
class JobError(Exception):
|
||||
"""An error whose message is safe to return to the client."""
|
||||
|
||||
|
||||
class JobCancelledError(JobError):
|
||||
"""The client asked the job to stop."""
|
||||
|
||||
|
||||
def _jobs_root() -> Path:
|
||||
return Path(envs.VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR) / "seedvr2-long"
|
||||
|
||||
|
||||
def _job_dir(job_id: str) -> Path:
|
||||
"""Resolve a client-supplied job id, which must never shape the path."""
|
||||
if not JOB_ID.fullmatch(job_id):
|
||||
raise HTTPException(404, "SeedVR2 long-video job not found")
|
||||
return _jobs_root() / job_id
|
||||
|
||||
|
||||
def _sweep_expired_jobs() -> None:
|
||||
"""Drop settled job directories so repeated submissions cannot fill the disk."""
|
||||
deadline = time.time() - envs.VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS
|
||||
for job in _jobs_root().glob("*"):
|
||||
status = job / "status.json"
|
||||
if not job.is_dir() or not JOB_ID.fullmatch(job.name) or not status.exists():
|
||||
continue
|
||||
record = json.loads(status.read_text())
|
||||
# A job owned by a dead process is settled too, or a restart would leak it.
|
||||
settled = record["status"] not in {"queued", "running"} or record.get("pid") != os.getpid()
|
||||
if settled and status.stat().st_mtime <= deadline:
|
||||
shutil.rmtree(job, ignore_errors=True)
|
||||
|
||||
|
||||
def _store_upload(upload: UploadFile, destination: Path) -> None:
|
||||
"""Copy the upload under a size cap so one client cannot fill the disk."""
|
||||
limit = envs.VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES
|
||||
written = 0
|
||||
with destination.open("wb") as target:
|
||||
while chunk := upload.file.read(UPLOAD_CHUNK):
|
||||
written += len(chunk)
|
||||
if written > limit:
|
||||
raise HTTPException(413, f"SeedVR2 long-video upload exceeds {limit} bytes")
|
||||
target.write(chunk)
|
||||
if not written:
|
||||
raise HTTPException(400, "SeedVR2 long-video upload is empty")
|
||||
|
||||
|
||||
def _window(width: int, height: int) -> int:
|
||||
"""Longest window the whole-clip budget admits at this output size.
|
||||
|
||||
Every window pays a fixed cost for its round trip through the serving
|
||||
stack and throws away its overlap, so the route uses the largest window the
|
||||
model accepts. The model pads clips to 4n+1 frames, so the window has that
|
||||
length too; 0 means not even one frame fits.
|
||||
"""
|
||||
frame_pixels, clip_pixels = sharded_budget()
|
||||
if width * height > frame_pixels:
|
||||
return 0
|
||||
fit = min(clip_pixels // (width * height), max_frames(), envs.VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW)
|
||||
return fit - (fit - 1) % 4
|
||||
|
||||
|
||||
def _status(job: Path, stage: str, frames: int, error: str = "") -> None:
|
||||
# The owning PID lets a poller detect a job that a server restart abandoned.
|
||||
record = {"status": stage, "frames": frames, "error": error, "pid": os.getpid()}
|
||||
pending = job / "status.json.tmp"
|
||||
pending.write_text(json.dumps(record) + "\n")
|
||||
pending.replace(job / "status.json")
|
||||
|
||||
|
||||
def _frames(source: Path, width: int, height: int, loop: bool) -> Iterator[av.VideoFrame]:
|
||||
while True:
|
||||
count = 0
|
||||
with av.open(str(source)) as container:
|
||||
video = container.streams.video[0]
|
||||
if video.average_rate != Fraction(FPS):
|
||||
raise JobError("SeedVR2 long video requires 24 FPS input")
|
||||
for frame in container.decode(video=0):
|
||||
if frame.pts is None or frame.time_base is None or frame.pts * frame.time_base != Fraction(count, FPS):
|
||||
raise JobError("SeedVR2 long video requires constant 24 FPS timestamps starting at zero")
|
||||
count += 1
|
||||
# The model resizes to the output on the device, so windows carry
|
||||
# source pixels; only a larger source shrinks here to stay in budget.
|
||||
if frame.width * frame.height > width * height:
|
||||
yield frame.reformat(width=width, height=height, format="yuv420p", interpolation="BICUBIC")
|
||||
else:
|
||||
yield frame.reformat(width=frame.width & ~1, height=frame.height & ~1, format="yuv420p")
|
||||
if count == 0:
|
||||
raise JobError("Input video has no frames")
|
||||
if not loop:
|
||||
return
|
||||
|
||||
|
||||
# Consecutive windows share their overlap frames and encoding stamps each
|
||||
# frame's pts, so two in-flight windows must not encode at the same time.
|
||||
_segment_lock = threading.Lock()
|
||||
|
||||
|
||||
def _segment(frames: list[av.VideoFrame]) -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
with _segment_lock, av.open(buffer, "w", format="matroska") as container:
|
||||
stream = container.add_stream("libx264", rate=FPS)
|
||||
stream.width, stream.height, stream.pix_fmt = frames[0].width, frames[0].height, "yuv420p"
|
||||
stream.options = LOSSLESS
|
||||
for index, frame in enumerate(frames):
|
||||
frame.pts, frame.time_base = index, Fraction(1, FPS)
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _restore(
|
||||
frames: list[av.VideoFrame],
|
||||
width: int,
|
||||
height: int,
|
||||
seed: int,
|
||||
color_correction_method: str,
|
||||
port: int,
|
||||
authorization: str,
|
||||
) -> list[np.ndarray]:
|
||||
response = requests.post(
|
||||
f"http://127.0.0.1:{port}/v1/videos/sync",
|
||||
data={
|
||||
"prompt": " ",
|
||||
"size": f"{width}x{height}",
|
||||
"num_frames": str(len(frames)),
|
||||
"num_inference_steps": "1",
|
||||
"guidance_scale": "1",
|
||||
"seed": str(seed),
|
||||
"extra_params": json.dumps(
|
||||
{"video_codec_options": LOSSLESS, "color_correction_method": color_correction_method}
|
||||
),
|
||||
},
|
||||
files={"input_references": ("window.mkv", _segment(frames), "video/x-matroska")},
|
||||
headers={"Authorization": authorization} if authorization else {},
|
||||
timeout=900,
|
||||
)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except requests.HTTPError as error:
|
||||
raise JobError("SeedVR2 window restoration failed") from error
|
||||
with av.open(io.BytesIO(response.content)) as container:
|
||||
video = container.streams.video[0]
|
||||
decoded = list(container.decode(video=0))
|
||||
if (video.width, video.height, video.average_rate, len(decoded)) != (width, height, Fraction(FPS), len(frames)):
|
||||
raise JobError("SeedVR2 window restoration returned an unexpected geometry or frame count")
|
||||
if [frame.pts * frame.time_base for frame in decoded] != [Fraction(index, FPS) for index in range(len(frames))]:
|
||||
raise JobError("SeedVR2 window restoration returned invalid timestamps")
|
||||
return [frame.to_ndarray(format="rgb24") for frame in decoded]
|
||||
|
||||
|
||||
def _run(
|
||||
job: Path,
|
||||
width: int,
|
||||
height: int,
|
||||
target: int,
|
||||
loop: bool,
|
||||
seed: int,
|
||||
color_correction_method: str,
|
||||
port: int,
|
||||
authorization: str,
|
||||
) -> None:
|
||||
window = _window(width, height)
|
||||
if window <= OVERLAP:
|
||||
raise JobError("SeedVR2 long video output exceeds the configured SeedVR2 clip budget")
|
||||
with tempfile.TemporaryDirectory(prefix="seedvr2-long-") as scratch:
|
||||
_restore_to(
|
||||
job, Path(scratch), window, width, height, target, loop, seed, color_correction_method, port, authorization
|
||||
)
|
||||
|
||||
|
||||
def _restore_to(
|
||||
job: Path,
|
||||
scratch: Path,
|
||||
window: int,
|
||||
width: int,
|
||||
height: int,
|
||||
target: int,
|
||||
loop: bool,
|
||||
seed: int,
|
||||
color_correction_method: str,
|
||||
port: int,
|
||||
authorization: str,
|
||||
) -> None:
|
||||
"""Restore into local scratch files, then publish the result into ``job``."""
|
||||
source = job / "input.mp4"
|
||||
video_path = scratch / "video.mp4"
|
||||
output = scratch / "output.mp4"
|
||||
source_frames = _frames(source, width, height, loop)
|
||||
|
||||
def windows() -> Iterator[list[av.VideoFrame]]:
|
||||
batch = list(itertools.islice(source_frames, min(window, target)))
|
||||
start = 0
|
||||
while True:
|
||||
if len(batch) != min(window, target - start):
|
||||
raise JobError("Input video has fewer frames than requested; set loop_input=true to repeat it")
|
||||
yield batch
|
||||
start += len(batch) - OVERLAP
|
||||
if start + OVERLAP >= target:
|
||||
return
|
||||
batch = batch[-OVERLAP:] + list(itertools.islice(source_frames, min(window, target - start) - OVERLAP))
|
||||
|
||||
def restore(frames: list[av.VideoFrame]) -> Future[list[np.ndarray]]:
|
||||
if (job / "cancel").exists():
|
||||
raise JobCancelledError("SeedVR2 long-video job was cancelled")
|
||||
return pool.submit(_restore, frames, width, height, seed, color_correction_method, port, authorization)
|
||||
|
||||
written = 0
|
||||
pending: list[np.ndarray] = []
|
||||
# Two windows stay in flight, so the engine always has the next one queued
|
||||
# while this thread blends and encodes, and the other prepares its upload.
|
||||
with ThreadPoolExecutor(max_workers=2) as pool, av.open(str(video_path), "w", format="mp4") as container:
|
||||
video = container.add_stream("libx264", rate=FPS)
|
||||
video.width, video.height, video.pix_fmt = width, height, "yuv420p"
|
||||
video.options = {"preset": "veryfast", "crf": "18", "tune": "zerolatency", "bf": "0"}
|
||||
|
||||
def write(array: np.ndarray) -> None:
|
||||
nonlocal written
|
||||
frame = av.VideoFrame.from_ndarray(array, format="rgb24")
|
||||
frame.pts, frame.time_base = written, Fraction(1, FPS)
|
||||
for packet in video.encode(frame):
|
||||
container.mux(packet)
|
||||
written += 1
|
||||
|
||||
batches = windows()
|
||||
in_flight = deque(restore(batch) for batch in itertools.islice(batches, 2))
|
||||
while in_flight:
|
||||
restored = in_flight.popleft().result()
|
||||
following = next(batches, None)
|
||||
if following is not None:
|
||||
in_flight.append(restore(following))
|
||||
last = not in_flight
|
||||
if not pending:
|
||||
for frame in restored if last else restored[:-OVERLAP]:
|
||||
write(frame)
|
||||
else:
|
||||
for offset in range(OVERLAP):
|
||||
newer = (offset + 1) / (OVERLAP + 1)
|
||||
write(np.rint(pending[offset] * (1 - newer) + restored[offset] * newer).astype(np.uint8))
|
||||
for frame in restored[OVERLAP:] if last else restored[OVERLAP:-OVERLAP]:
|
||||
write(frame)
|
||||
_status(job, "running", written)
|
||||
pending = restored[-OVERLAP:]
|
||||
for packet in video.encode():
|
||||
container.mux(packet)
|
||||
if written != target:
|
||||
raise JobError(f"SeedVR2 restored {written} of {target} requested frames")
|
||||
|
||||
with av.open(str(source)) as container:
|
||||
has_audio = bool(container.streams.audio)
|
||||
if has_audio:
|
||||
subprocess.run(
|
||||
[
|
||||
imageio_ffmpeg.get_ffmpeg_exe(),
|
||||
"-nostdin",
|
||||
"-y",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-stream_loop",
|
||||
"-1" if loop else "0",
|
||||
"-i",
|
||||
str(source),
|
||||
"-i",
|
||||
str(video_path),
|
||||
"-map",
|
||||
"1:v:0",
|
||||
"-map",
|
||||
"0:a:0",
|
||||
"-frames:v",
|
||||
str(target),
|
||||
"-t",
|
||||
str(target / FPS),
|
||||
"-c:v",
|
||||
"copy",
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-b:a",
|
||||
"160k",
|
||||
"-movflags",
|
||||
"+faststart",
|
||||
str(output),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
else:
|
||||
video_path.replace(output)
|
||||
with av.open(str(output)) as container:
|
||||
video = container.streams.video[0]
|
||||
decoded = 0
|
||||
for decoded, frame in enumerate(container.decode(video=0), start=1):
|
||||
if (
|
||||
frame.pts is None
|
||||
or frame.time_base is None
|
||||
or frame.pts * frame.time_base != Fraction(decoded - 1, FPS)
|
||||
):
|
||||
raise JobError("SeedVR2 output encoding produced invalid timestamps")
|
||||
if (video.width, video.height, video.average_rate, decoded) != (width, height, Fraction(FPS), target):
|
||||
raise JobError("SeedVR2 output encoding failed frame or geometry validation")
|
||||
if has_audio and not container.streams.audio:
|
||||
raise JobError("SeedVR2 output encoding lost the audio track")
|
||||
if has_audio:
|
||||
with av.open(str(output)) as container:
|
||||
audio = container.streams.audio[0]
|
||||
samples = sum(frame.samples for frame in container.decode(audio=0))
|
||||
if samples < (target / FPS - 1) * audio.rate:
|
||||
raise JobError("SeedVR2 output audio ends before the video")
|
||||
with output.open("rb") as content:
|
||||
hasher = hashlib.sha256()
|
||||
for chunk in iter(lambda: content.read(UPLOAD_CHUNK), b""):
|
||||
hasher.update(chunk)
|
||||
digest = hasher.hexdigest()
|
||||
# One sequential copy is the only write to the job directory, which may sit
|
||||
# on network storage where the encoder's and muxer's small writes are slow.
|
||||
staged = job / "output.mp4.tmp"
|
||||
shutil.move(output, staged)
|
||||
staged.replace(job / "output.mp4")
|
||||
(job / "result.json").write_text(
|
||||
json.dumps({"frames": target, "width": width, "height": height, "fps": FPS, "sha256": digest}) + "\n"
|
||||
)
|
||||
_status(job, "completed", target)
|
||||
|
||||
|
||||
def _background(
|
||||
job: Path,
|
||||
width: int,
|
||||
height: int,
|
||||
target: int,
|
||||
loop: bool,
|
||||
seed: int,
|
||||
color_correction_method: str,
|
||||
port: int,
|
||||
authorization: str,
|
||||
) -> None:
|
||||
try:
|
||||
_run(job, width, height, target, loop, seed, color_correction_method, port, authorization)
|
||||
except Exception as error:
|
||||
logger.exception("SeedVR2 long-video job %s failed", job.name)
|
||||
current = json.loads((job / "status.json").read_text())
|
||||
detail = str(error) if isinstance(error, JobError) else "SeedVR2 long-video restoration failed"
|
||||
_status(job, "cancelled" if isinstance(error, JobCancelledError) else "failed", current["frames"], detail)
|
||||
|
||||
|
||||
@router.post("/v1/seedvr2/restore-long", status_code=202)
|
||||
async def create_long_video(
|
||||
raw_request: Request,
|
||||
input_references: UploadFile = File(...),
|
||||
num_frames: int = Form(...),
|
||||
size: str = Form(...),
|
||||
loop_input: bool = Form(False),
|
||||
prompt: str = Form(" "),
|
||||
seed: int = Form(7723),
|
||||
extra_params: str | None = Form(None),
|
||||
) -> dict[str, str | int]:
|
||||
global active_task
|
||||
if raw_request.app.state.api_server_count != 1:
|
||||
raise HTTPException(409, "SeedVR2 long-video jobs require one API server")
|
||||
try:
|
||||
width, height = (int(value) for value in size.lower().split("x"))
|
||||
except ValueError as error:
|
||||
raise HTTPException(400, "size must be WIDTHxHEIGHT") from error
|
||||
if not 1 <= num_frames <= MAX_FRAMES or min(width, height) < 16 or width % 16 or height % 16:
|
||||
raise HTTPException(400, "SeedVR2 long-video frame count or dimensions are invalid")
|
||||
if prompt.strip():
|
||||
raise HTTPException(400, "SeedVR2 long video requires a blank prompt")
|
||||
if _window(width, height) <= OVERLAP:
|
||||
raise HTTPException(400, "SeedVR2 long video output exceeds the configured SeedVR2 frame or clip budget")
|
||||
if not 0 <= seed <= 2**32 - 1:
|
||||
raise HTTPException(400, "Seed must be a 32-bit unsigned integer")
|
||||
try:
|
||||
options = json.loads(extra_params) if extra_params is not None else {}
|
||||
except (TypeError, ValueError) as error:
|
||||
raise HTTPException(400, "extra_params must be a JSON object") from error
|
||||
if not isinstance(options, dict):
|
||||
raise HTTPException(400, "extra_params must be a JSON object")
|
||||
unknown = options.keys() - {"color_correction_method"}
|
||||
if unknown:
|
||||
raise HTTPException(400, f"Unsupported SeedVR2 long-video extra_params: {sorted(unknown)}")
|
||||
color_correction_method = options.get("color_correction_method", DEFAULT_COLOR_CORRECTION_METHOD)
|
||||
if color_correction_method is None:
|
||||
color_correction_method = DEFAULT_COLOR_CORRECTION_METHOD
|
||||
if color_correction_method not in COLOR_CORRECTION_METHODS:
|
||||
raise HTTPException(400, f"color_correction_method must be one of {list(COLOR_CORRECTION_METHODS)}")
|
||||
async with job_lock:
|
||||
if active_task is not None and not active_task.done():
|
||||
raise HTTPException(409, "A SeedVR2 long-video job is already running")
|
||||
await asyncio.to_thread(_sweep_expired_jobs)
|
||||
job_id = uuid4().hex
|
||||
job = _jobs_root() / job_id
|
||||
job.mkdir(parents=True)
|
||||
try:
|
||||
await asyncio.to_thread(_store_upload, input_references, job / "input.mp4")
|
||||
except HTTPException:
|
||||
shutil.rmtree(job, ignore_errors=True)
|
||||
raise
|
||||
_status(job, "queued", 0)
|
||||
active_task = asyncio.create_task(
|
||||
asyncio.to_thread(
|
||||
_background,
|
||||
job,
|
||||
width,
|
||||
height,
|
||||
num_frames,
|
||||
loop_input,
|
||||
seed,
|
||||
color_correction_method,
|
||||
raw_request.app.state.seedvr2_long_port,
|
||||
raw_request.headers.get("authorization", ""),
|
||||
)
|
||||
)
|
||||
return {"id": job_id, "status": "queued", "max_frames": MAX_FRAMES}
|
||||
|
||||
|
||||
@router.get("/v1/seedvr2/restore-long/{job_id}")
|
||||
async def get_long_video(job_id: str) -> dict[str, str | int]:
|
||||
status = _job_dir(job_id) / "status.json"
|
||||
if not status.exists():
|
||||
raise HTTPException(404, "SeedVR2 long-video job not found")
|
||||
record = json.loads(status.read_text())
|
||||
if record["status"] in {"queued", "running"} and record.get("pid") != os.getpid():
|
||||
# Only this process runs jobs, so another owner means a restart lost it.
|
||||
record["status"], record["error"] = "failed", "SeedVR2 long-video job was interrupted by a server restart"
|
||||
return record
|
||||
|
||||
|
||||
@router.get("/v1/seedvr2/restore-long/{job_id}/content")
|
||||
async def download_long_video(job_id: str) -> FileResponse:
|
||||
job = _job_dir(job_id)
|
||||
if not (job / "result.json").exists():
|
||||
raise HTTPException(404, "SeedVR2 long-video output is not ready")
|
||||
return FileResponse(job / "output.mp4", media_type="video/mp4", filename=f"{job_id}.mp4")
|
||||
|
||||
|
||||
@router.delete("/v1/seedvr2/restore-long/{job_id}", status_code=202)
|
||||
async def cancel_long_video(job_id: str) -> dict[str, str]:
|
||||
"""Ask the running job to stop; it settles at the next window boundary."""
|
||||
job = _job_dir(job_id)
|
||||
status = job / "status.json"
|
||||
if not status.exists():
|
||||
raise HTTPException(404, "SeedVR2 long-video job not found")
|
||||
record = json.loads(status.read_text())
|
||||
if record["status"] not in {"queued", "running"}:
|
||||
raise HTTPException(409, f"SeedVR2 long-video job already {record['status']}")
|
||||
(job / "cancel").touch()
|
||||
return {"id": job_id, "status": "cancelling"}
|
||||
|
||||
|
||||
def register_routes(app: FastAPI, port: int) -> None:
|
||||
"""Attach this model's optional long-video API to its serving process."""
|
||||
app.state.seedvr2_long_port = port
|
||||
app.include_router(router)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,671 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""SeedVR2 restoration pipeline: input admission, one Euler step and colour transfer."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from fractions import Fraction
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import ClassVar
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch import Tensor, nn
|
||||
from torch.nn import functional as F
|
||||
from vllm.utils.torch_utils import set_torch_threads_for_runtime
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.seedvr2 import config as envs
|
||||
from vllm_omni.diffusion.models.seedvr2.config import COLOR_CORRECTION_METHODS, DEFAULT_COLOR_CORRECTION_METHOD
|
||||
from vllm_omni.diffusion.models.seedvr2.nadit import SEEDVR2_3B_CONFIG, SeedVR2NaDiT, validate_seedvr2_parallel_config
|
||||
from vllm_omni.diffusion.models.seedvr2.vae import SeedVR2VAE
|
||||
from vllm_omni.diffusion.offloader.config import offload_enabled
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
||||
|
||||
# Bound per-frame decode bookkeeping; the padded clip-pixel budget is tighter
|
||||
# for normal video resolutions.
|
||||
MAX_FRAME_PIXELS = 848 * 480
|
||||
MAX_CLIP_PIXELS = 5 * MAX_FRAME_PIXELS
|
||||
|
||||
|
||||
def max_frames() -> int:
|
||||
"""Decoder-work frame cap, shared by every serving profile."""
|
||||
return envs.VLLM_OMNI_SEEDVR2_MAX_FRAMES
|
||||
|
||||
|
||||
# The sharded clip budget validated on the smallest qualified device (32 GB, SP4):
|
||||
# five 2560x1472 frames.
|
||||
CALIBRATED_SHARDED_CLIP_PIXELS = 5 * 2560 * 1472
|
||||
# The largest clip run through the forward at SP4 on a 268 GiB B300 (513 frames
|
||||
# at 1536x2688, 55.5 GiB peak activations). Past it the DiT's fp32 RoPE
|
||||
# intermediates and allocator fragmentation outgrow the model below, so the
|
||||
# scaled budget stops here. It also keeps one channel of a clip under 2**31
|
||||
# elements.
|
||||
VALIDATED_SHARDED_CLIP_PIXELS = 513 * 1536 * 2688
|
||||
# Peak activation bytes per rank of one sharded forward at SP4, the smallest
|
||||
# sharded degree (more ranks need less). Measured on B300 from 5 to 1,921 frames
|
||||
# at 768x1344, 2560x1472 and 1536x2688, weights excluded:
|
||||
# max(ENCODER * frame_pixels * min(frames, 9) + CLIP * clip_pixels,
|
||||
# DECODED * clip_pixels)
|
||||
# The VAE encoder's first temporal chunk (up to nine frames) peaks for short
|
||||
# clips; the whole-clip fp16 input and gathered decode peak for long ones.
|
||||
ENCODER_CHUNK_FRAMES = 9
|
||||
ENCODER_BYTES_PER_PIXEL = 1140
|
||||
CLIP_BYTES_PER_PIXEL = 8.5
|
||||
DECODED_BYTES_PER_PIXEL = 28
|
||||
WEIGHT_BYTES = 7 * 1024**3
|
||||
# Room for the CUDA context and communicator buffers, then for allocator
|
||||
# fragmentation, which reached 37% of the allocated activations near OOM.
|
||||
USABLE_MEMORY_FRACTION = 0.85
|
||||
FRAGMENTATION = 1.37
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _device_memory() -> int | None:
|
||||
"""Smallest total memory among the visible accelerators, if it can be read."""
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
try:
|
||||
count = torch.accelerator.device_count()
|
||||
return min(int(current_omni_platform.get_device_total_memory(index)) for index in range(count)) or None
|
||||
except (AttributeError, NotImplementedError, RuntimeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def scaled_clip_pixels(frame_pixels: int, device_memory: int | None) -> int:
|
||||
"""Largest padded clip the memory model admits on ``device_memory``.
|
||||
|
||||
The result never falls below the calibrated budget or rises above the
|
||||
validated one. The worst case for a clip budget is a clip of full-size
|
||||
frames, so the model is evaluated at ``frame_pixels``; smaller frames only
|
||||
need less.
|
||||
"""
|
||||
if device_memory is None:
|
||||
return CALIBRATED_SHARDED_CLIP_PIXELS
|
||||
budget = (device_memory * USABLE_MEMORY_FRACTION - WEIGHT_BYTES) / FRAGMENTATION
|
||||
chunk = ENCODER_CHUNK_FRAMES * frame_pixels
|
||||
clip = budget / (ENCODER_BYTES_PER_PIXEL + CLIP_BYTES_PER_PIXEL)
|
||||
if clip > chunk:
|
||||
clip = min(
|
||||
(budget - ENCODER_BYTES_PER_PIXEL * chunk) / CLIP_BYTES_PER_PIXEL,
|
||||
budget / DECODED_BYTES_PER_PIXEL,
|
||||
)
|
||||
return min(max(CALIBRATED_SHARDED_CLIP_PIXELS, int(clip)), VALIDATED_SHARDED_CLIP_PIXELS)
|
||||
|
||||
|
||||
def sharded_budget() -> tuple[int, int]:
|
||||
"""Per-frame and padded-clip budgets for the VAE-sharded serving profile.
|
||||
|
||||
An unset clip budget scales with the smallest visible device's memory.
|
||||
"""
|
||||
frame_pixels = envs.VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS
|
||||
clip_pixels = envs.VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS
|
||||
if clip_pixels is None:
|
||||
clip_pixels = scaled_clip_pixels(frame_pixels, _device_memory())
|
||||
return frame_pixels, clip_pixels
|
||||
|
||||
|
||||
def validate_clip_size(frame_count: int, height: int, width: int, frame_pixels: int, clip_pixels: int) -> None:
|
||||
padded_frames = frame_count + (1 - frame_count) % 4
|
||||
if frame_count > max_frames() or height * width > frame_pixels or padded_frames * height * width > clip_pixels:
|
||||
raise OmniClientError("SeedVR2 clip exceeds the configured frame or pixel budget")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SourceVideo:
|
||||
fps: float
|
||||
pts: tuple[int, ...]
|
||||
time_base: Fraction
|
||||
audio: torch.Tensor | None
|
||||
audio_sample_rate: int | None
|
||||
|
||||
|
||||
def read_video(
|
||||
path: str | Path, frame_pixels: int = MAX_FRAME_PIXELS, clip_pixels: int = MAX_CLIP_PIXELS
|
||||
) -> tuple[torch.Tensor, SourceVideo]:
|
||||
try:
|
||||
return _read_video(path, frame_pixels, clip_pixels)
|
||||
except av.FFmpegError as exc:
|
||||
raise OmniClientError(f"Cannot decode SeedVR2 input video: {exc.strerror}") from exc
|
||||
|
||||
|
||||
def _read_video(path: str | Path, frame_pixels: int, clip_pixels: int) -> tuple[torch.Tensor, SourceVideo]:
|
||||
with av.open(str(path)) as container:
|
||||
if not container.streams.video:
|
||||
raise OmniClientError("SeedVR2 input contains no video stream")
|
||||
stream = container.streams.video[0]
|
||||
rate = stream.average_rate
|
||||
if rate is None or rate <= 0:
|
||||
raise OmniClientError("SeedVR2 input must declare a positive frame rate")
|
||||
height, width = stream.codec_context.height, stream.codec_context.width
|
||||
if height > 0 and width > 0:
|
||||
validate_clip_size(max(1, stream.frames), height, width, frame_pixels, clip_pixels)
|
||||
if stream.duration is not None and stream.time_base is not None:
|
||||
if stream.duration * stream.time_base * rate > max_frames():
|
||||
raise OmniClientError("SeedVR2 input exceeds the maximum duration")
|
||||
frames = []
|
||||
pts = []
|
||||
for frame in container.decode(stream):
|
||||
validate_clip_size(len(frames) + 1, frame.height, frame.width, frame_pixels, clip_pixels)
|
||||
if frame.pts is None:
|
||||
raise OmniClientError("SeedVR2 input video is missing frame timestamps")
|
||||
frames.append(torch.from_numpy(frame.to_ndarray(format="rgb24")))
|
||||
pts.append(frame.pts)
|
||||
if not frames:
|
||||
raise OmniClientError("SeedVR2 input video contains no decoded frames")
|
||||
time_base = Fraction(stream.time_base)
|
||||
if any(b <= a for a, b in zip(pts, pts[1:])):
|
||||
raise OmniClientError("SeedVR2 input timestamps must increase")
|
||||
if any(
|
||||
abs((pts[index] - pts[0]) * time_base - Fraction(index, 1) / rate) > time_base for index in range(len(pts))
|
||||
):
|
||||
raise OmniClientError("SeedVR2 currently supports constant-frame-rate video only")
|
||||
audio = None
|
||||
sample_rate = None
|
||||
with av.open(str(path)) as container:
|
||||
if container.streams.audio:
|
||||
stream = container.streams.audio[0]
|
||||
sample_rate = stream.codec_context.sample_rate
|
||||
resampler = av.AudioResampler(format="fltp", layout=stream.codec_context.layout, rate=sample_rate)
|
||||
channels = len(stream.codec_context.layout.channels)
|
||||
if channels not in (1, 2):
|
||||
raise OmniClientError("SeedVR2 audio preservation currently supports mono or stereo")
|
||||
sample_count = round(len(frames) / rate * sample_rate)
|
||||
waveform = np.zeros((channels, sample_count), dtype=np.float32)
|
||||
|
||||
def resampled_frames() -> Iterator[av.AudioFrame]:
|
||||
for frame in container.decode(stream):
|
||||
yield from resampler.resample(frame)
|
||||
yield from resampler.resample(None)
|
||||
|
||||
for frame in resampled_frames():
|
||||
if frame.pts is None:
|
||||
raise OmniClientError("SeedVR2 audio is missing timestamps")
|
||||
offset = round((frame.pts * frame.time_base - pts[0] * time_base) * sample_rate)
|
||||
if offset >= sample_count:
|
||||
break
|
||||
chunk = frame.to_ndarray()
|
||||
start, stop = max(0, offset), min(sample_count, offset + chunk.shape[1])
|
||||
if stop > start:
|
||||
waveform[:, start:stop] = chunk[:, start - offset : stop - offset]
|
||||
audio = torch.from_numpy(waveform).unsqueeze(0)
|
||||
# Frames stay uint8 THWC: the device converts them, so the host never builds
|
||||
# (or ships to every rank) a float copy four times the size.
|
||||
return torch.stack(frames), SourceVideo(
|
||||
float(rate),
|
||||
tuple(pts),
|
||||
time_base,
|
||||
audio,
|
||||
sample_rate,
|
||||
)
|
||||
|
||||
|
||||
# Decomposition depth: the fifth band already isolates the illumination tint.
|
||||
WAVELET_LEVELS = 5
|
||||
# Restoration luminance is the reason to run the model, so chroma is matched
|
||||
# fully while luminance keeps this much of the restored value.
|
||||
DEFAULT_LUMINANCE_WEIGHT = 0.8
|
||||
|
||||
# CIE 15 lightness transfer, sRGB (IEC 61966-2-1) primaries, D65 white point.
|
||||
_CIELAB_DELTA = 6.0 / 29.0
|
||||
_D65_WHITE = (0.95047, 1.0, 1.08883)
|
||||
_SRGB_TO_XYZ = (
|
||||
(0.4124564, 0.3575761, 0.1804375),
|
||||
(0.2126729, 0.7151522, 0.0721750),
|
||||
(0.0193339, 0.1191920, 0.9503041),
|
||||
)
|
||||
_XYZ_TO_SRGB = (
|
||||
(3.2404542, -1.5371385, -0.4985314),
|
||||
(-0.9692660, 1.8760108, 0.0415560),
|
||||
(0.0556434, -0.2040259, 1.0572252),
|
||||
)
|
||||
|
||||
|
||||
def _binomial_lowpass(image: Tensor, radius: int) -> Tensor:
|
||||
"""One a-trous low-pass step with the separable 1-2-1 binomial kernel."""
|
||||
channels = image.shape[1]
|
||||
weights = image.new_tensor([1.0, 2.0, 1.0]) / 4.0
|
||||
kernel = torch.outer(weights, weights).expand(channels, 1, 3, 3)
|
||||
padded = F.pad(image, (radius,) * 4, mode="replicate")
|
||||
return F.conv2d(padded, kernel, groups=channels, dilation=radius)
|
||||
|
||||
|
||||
def _lowest_band(image: Tensor, levels: int = WAVELET_LEVELS) -> Tensor:
|
||||
"""Residual band after ``levels`` doubling-radius low-pass steps.
|
||||
|
||||
Summing the per-level high-frequency residuals telescopes to
|
||||
``image - lowest_band(image)``, so only this band has to be materialized.
|
||||
"""
|
||||
# Keep a dilated tap well inside the frame so it samples real signal rather
|
||||
# than replicated edges; this bound matches the reference decomposition.
|
||||
limit = max(1, min(image.shape[-2:]) // 8)
|
||||
for level in range(levels):
|
||||
image = _binomial_lowpass(image, min(1 << level, limit))
|
||||
return image
|
||||
|
||||
|
||||
def _match_histogram(source: Tensor, reference: Tensor) -> Tensor:
|
||||
"""Give ``source`` the value distribution of ``reference``, rank for rank.
|
||||
|
||||
Both channels hold the same pixel count, so ranks map one to one and the
|
||||
match is exact without quantile interpolation.
|
||||
"""
|
||||
order = source.flatten().argsort()
|
||||
matched = torch.empty_like(order, dtype=reference.dtype)
|
||||
matched[order] = reference.flatten().sort().values
|
||||
return matched.view_as(source)
|
||||
|
||||
|
||||
def _srgb_to_lab(rgb: Tensor) -> Tensor:
|
||||
linear = torch.where(rgb > 0.04045, ((rgb + 0.055) / 1.055).pow(2.4), rgb / 12.92)
|
||||
matrix = rgb.new_tensor(_SRGB_TO_XYZ)
|
||||
xyz = torch.einsum("ij,njhw->nihw", matrix, linear)
|
||||
xyz = xyz / xyz.new_tensor(_D65_WHITE).view(1, 3, 1, 1)
|
||||
scaled = torch.where(
|
||||
xyz > _CIELAB_DELTA**3,
|
||||
xyz.clamp_min(0.0).pow(1.0 / 3.0),
|
||||
xyz / (3.0 * _CIELAB_DELTA**2) + 4.0 / 29.0,
|
||||
)
|
||||
x, y, z = scaled[:, 0], scaled[:, 1], scaled[:, 2]
|
||||
return torch.stack([116.0 * y - 16.0, 500.0 * (x - y), 200.0 * (y - z)], dim=1)
|
||||
|
||||
|
||||
def _lab_to_srgb(lab: Tensor) -> Tensor:
|
||||
lightness, chroma_a, chroma_b = lab[:, 0], lab[:, 1], lab[:, 2]
|
||||
y = (lightness + 16.0) / 116.0
|
||||
scaled = torch.stack([y + chroma_a / 500.0, y, y - chroma_b / 200.0], dim=1)
|
||||
xyz = torch.where(
|
||||
scaled > _CIELAB_DELTA,
|
||||
scaled.pow(3.0),
|
||||
3.0 * _CIELAB_DELTA**2 * (scaled - 4.0 / 29.0),
|
||||
)
|
||||
xyz = xyz * xyz.new_tensor(_D65_WHITE).view(1, 3, 1, 1)
|
||||
linear = torch.einsum("ij,njhw->nihw", lab.new_tensor(_XYZ_TO_SRGB), xyz)
|
||||
rgb = torch.where(
|
||||
linear > 0.0031308,
|
||||
1.055 * linear.clamp_min(0.0).pow(1.0 / 2.4) - 0.055,
|
||||
12.92 * linear,
|
||||
)
|
||||
return rgb.clamp_(0.0, 1.0)
|
||||
|
||||
|
||||
def _wavelet_transfer(restored: Tensor, reference: Tensor) -> Tensor:
|
||||
return (restored - _lowest_band(restored) + _lowest_band(reference)).clamp_(0.0, 1.0)
|
||||
|
||||
|
||||
def _channel_stats(image: Tensor, eps: float) -> tuple[Tensor, Tensor]:
|
||||
dims = (2, 3)
|
||||
deviation = (image.var(dims, correction=0, keepdim=True) + eps).sqrt()
|
||||
return image.mean(dims, keepdim=True), deviation
|
||||
|
||||
|
||||
def _adain_transfer(restored: Tensor, reference: Tensor, eps: float = 1e-5) -> Tensor:
|
||||
restored_mean, restored_std = _channel_stats(restored, eps)
|
||||
reference_mean, reference_std = _channel_stats(reference, eps)
|
||||
normalized = (restored - restored_mean) / restored_std
|
||||
return (normalized * reference_std + reference_mean).clamp_(0.0, 1.0)
|
||||
|
||||
|
||||
def _lab_transfer(restored: Tensor, reference: Tensor, luminance_weight: float) -> Tensor:
|
||||
restored_lab = _srgb_to_lab(_wavelet_transfer(restored, reference))
|
||||
reference_lab = _srgb_to_lab(reference)
|
||||
lightness = restored_lab[:, 0]
|
||||
if luminance_weight < 1.0:
|
||||
matched = _match_histogram(lightness, reference_lab[:, 0])
|
||||
lightness = lightness * luminance_weight + matched * (1.0 - luminance_weight)
|
||||
corrected = torch.stack(
|
||||
[
|
||||
lightness,
|
||||
_match_histogram(restored_lab[:, 1], reference_lab[:, 1]),
|
||||
_match_histogram(restored_lab[:, 2], reference_lab[:, 2]),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
return _lab_to_srgb(corrected)
|
||||
|
||||
|
||||
def correct_video_color(
|
||||
video: Tensor,
|
||||
reference: Tensor,
|
||||
method: str = DEFAULT_COLOR_CORRECTION_METHOD,
|
||||
luminance_weight: float = DEFAULT_LUMINANCE_WEIGHT,
|
||||
) -> Tensor:
|
||||
"""Transfer ``reference`` colour onto ``video``; both are ``[B,C,T,H,W]`` in ``[0,1]``.
|
||||
|
||||
Frames are processed one at a time because histogram matching sorts every
|
||||
pixel of a channel. Inputs are never modified.
|
||||
"""
|
||||
if method not in COLOR_CORRECTION_METHODS:
|
||||
raise ValueError(f"SeedVR2 color_correction_method must be one of {list(COLOR_CORRECTION_METHODS)}: {method!r}")
|
||||
if method == "none":
|
||||
return video
|
||||
if video.shape != reference.shape:
|
||||
raise ValueError(f"SeedVR2 color correction needs matching shapes, got {video.shape} and {reference.shape}")
|
||||
if not 0.0 <= luminance_weight <= 1.0:
|
||||
raise ValueError(f"SeedVR2 luminance_weight must be within [0, 1]: {luminance_weight!r}")
|
||||
|
||||
# Colour math is scale sensitive, so it runs in fp32 regardless of the model dtype.
|
||||
frames = video.movedim(2, 0).flatten(0, 1).float()
|
||||
references = reference.movedim(2, 0).flatten(0, 1).float()
|
||||
corrected = torch.empty_like(frames)
|
||||
for index in range(frames.shape[0]):
|
||||
frame, guide = frames[index : index + 1], references[index : index + 1]
|
||||
if method == "wavelet":
|
||||
corrected[index : index + 1] = _wavelet_transfer(frame, guide)
|
||||
elif method == "adain":
|
||||
corrected[index : index + 1] = _adain_transfer(frame, guide)
|
||||
else:
|
||||
corrected[index : index + 1] = _lab_transfer(frame, guide, luminance_weight)
|
||||
corrected = corrected.unflatten(0, (video.shape[2], video.shape[0])).movedim(0, 2)
|
||||
return corrected.to(dtype=video.dtype)
|
||||
|
||||
|
||||
def validate_seedvr2_config(config: OmniDiffusionConfig) -> None:
|
||||
# Read the operator-tuned budgets here so a bad value is a startup error
|
||||
# rather than a traceback from whichever module imports them first.
|
||||
max_frames()
|
||||
sharded_budget()
|
||||
if config.dtype != torch.float16:
|
||||
raise ValueError("SeedVR2 3B requires dtype=float16")
|
||||
if not config.enforce_eager:
|
||||
raise ValueError("SeedVR2 requires enforce_eager=True; compiled execution is not supported")
|
||||
if config.cache_backend != "none" or config.quantization_config is not None:
|
||||
raise ValueError("SeedVR2 requires unquantized weights and cache_backend=none for its single Euler step")
|
||||
if config.vae_use_slicing:
|
||||
raise ValueError("SeedVR2 whole-clip VAE does not support batch slicing")
|
||||
if offload_enabled(config):
|
||||
raise ValueError("SeedVR2 native whole-clip execution does not support CPU offload")
|
||||
parallel = config.parallel_config
|
||||
validate_seedvr2_parallel_config(parallel)
|
||||
if parallel.data_parallel_size is not None and parallel.data_parallel_size > 1:
|
||||
raise ValueError("SeedVR2 does not support data_parallel_size > 1")
|
||||
if parallel.ulysses_degree > 20:
|
||||
raise ValueError("SeedVR2 3B Ulysses requires at least one of its 20 heads per rank")
|
||||
degrees = {
|
||||
"cfg_parallel_size": parallel.cfg_parallel_size,
|
||||
"tensor_parallel_size": parallel.tensor_parallel_size,
|
||||
"pipeline_parallel_size": parallel.pipeline_parallel_size,
|
||||
"text_encoder_tp_size": parallel.text_encoder_tp_size,
|
||||
"ring_degree": parallel.ring_degree,
|
||||
"allgather_degree": parallel.allgather_degree,
|
||||
}
|
||||
for name, degree in degrees.items():
|
||||
if degree != 1:
|
||||
raise ValueError(f"SeedVR2 requires {name}=1; use ulysses_degree for its model-owned SP")
|
||||
if parallel.vae_patch_parallel_size > 1 and parallel.vae_patch_parallel_size != parallel.ulysses_degree:
|
||||
raise ValueError("SeedVR2 requires vae_patch_parallel_size to match ulysses_degree")
|
||||
if parallel.vae_parallel_mode == "spatial_shard_width":
|
||||
raise ValueError("SeedVR2 VAE supports spatial_shard_height or tile mode")
|
||||
if parallel.use_hsdp or parallel.enable_expert_parallel:
|
||||
raise ValueError("SeedVR2 does not support HSDP or expert parallelism")
|
||||
if config.lora_path is not None:
|
||||
raise ValueError("SeedVR2 does not support LoRA adapters")
|
||||
|
||||
|
||||
def _check_frames(frames: torch.Tensor, height: int, width: int) -> None:
|
||||
"""Validate host frames: uint8 ``[T,H,W,3]`` or float ``[T,3,H,W]`` in ``[0,1]``."""
|
||||
rgb_axis = 3 if frames.dtype == torch.uint8 else 1
|
||||
if frames.ndim != 4 or frames.shape[rgb_axis] != 3 or min(frames.shape) < 1:
|
||||
raise OmniClientError("SeedVR2 requires nonempty RGB frames with shape [T,3,H,W]")
|
||||
if min(height, width) < 16 or height % 16 or width % 16:
|
||||
raise OmniClientError("SeedVR2 output dimensions must be positive multiples of 16")
|
||||
if frames.dtype == torch.uint8:
|
||||
return
|
||||
if not frames.is_floating_point() or not torch.isfinite(frames).all():
|
||||
raise OmniClientError("SeedVR2 frames must be finite floating-point RGB values")
|
||||
if torch.any(frames < 0) or torch.any(frames > 1):
|
||||
raise OmniClientError("SeedVR2 RGB values must be in [0,1]")
|
||||
|
||||
|
||||
def prepare_video(frames: torch.Tensor, height: int, width: int, device: torch.device) -> torch.Tensor:
|
||||
"""Resize RGB frames to the model grid on ``device`` and repeat the final frame.
|
||||
|
||||
``frames`` is uint8 ``[T,H,W,3]`` or float ``[T,3,H,W]`` in ``[0,1]``. Frames
|
||||
are converted one at a time, so the device never holds a float copy of the
|
||||
whole clip, and returns fp16 ``[1,3,T',H,W]`` in ``[-1,1]`` with ``T' = 4n+1``.
|
||||
"""
|
||||
count = frames.shape[0]
|
||||
sample = torch.empty((1, 3, count + (1 - count) % 4, height, width), device=device, dtype=torch.float16)
|
||||
for index in range(count):
|
||||
frame = frames[index].to(device, non_blocking=True)
|
||||
frame = frame.permute(2, 0, 1).float() / 255 if frame.dtype == torch.uint8 else frame.float()
|
||||
if frame.shape[-2:] != (height, width):
|
||||
frame = F.interpolate(
|
||||
frame.unsqueeze(0), size=(height, width), mode="bicubic", align_corners=False, antialias=True
|
||||
).squeeze(0)
|
||||
sample[0, :, index] = frame.clamp(0, 1) * 2 - 1
|
||||
sample[0, :, count:] = sample[0, :, count - 1 : count]
|
||||
return sample
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SeedVR2Input:
|
||||
# uint8 [T,H,W,3] for decoded video and PIL frames, float [T,3,H,W] otherwise;
|
||||
# resizing and normalization run on the device in ``prepare_video``.
|
||||
frames: torch.Tensor
|
||||
frame_count: int
|
||||
source: SourceVideo | None = None
|
||||
|
||||
|
||||
def _admission_budget(config: OmniDiffusionConfig) -> tuple[int, int]:
|
||||
"""Pick the budget for the serving profile; more ranks never admit less."""
|
||||
parallel = config.parallel_config
|
||||
sharded = parallel.ulysses_degree == parallel.vae_patch_parallel_size
|
||||
if config.vae_use_tiling and sharded and parallel.ulysses_degree >= 4:
|
||||
return sharded_budget()
|
||||
return MAX_FRAME_PIXELS, MAX_CLIP_PIXELS
|
||||
|
||||
|
||||
def prepare_request(
|
||||
request: OmniDiffusionRequest,
|
||||
*,
|
||||
frame_pixels: int = MAX_FRAME_PIXELS,
|
||||
clip_pixels: int = MAX_CLIP_PIXELS,
|
||||
) -> OmniDiffusionRequest:
|
||||
params = request.sampling_params
|
||||
if request.is_dummy_run():
|
||||
# The engine's generic warmup has no source video. Only internal warmup
|
||||
# requests may synthesize conditioning; real requests stay strict.
|
||||
request.prompt = {"prompt": "", "multi_modal_data": {"video": torch.zeros(1, 3, params.height, params.width)}}
|
||||
params.num_frames = 1
|
||||
params.num_inference_steps = 1
|
||||
params.guidance_scale = 1.0
|
||||
if params.num_inference_steps not in (None, 1) or params.guidance_scale != 1.0:
|
||||
raise OmniClientError("SeedVR2 whole-clip restoration requires one Euler step and guidance_scale=1")
|
||||
method = params.extra_args.get("color_correction_method")
|
||||
if method is not None and method not in COLOR_CORRECTION_METHODS:
|
||||
raise OmniClientError(f"SeedVR2 color_correction_method must be one of {list(COLOR_CORRECTION_METHODS)}")
|
||||
prompt = request.prompt
|
||||
if not isinstance(prompt, dict) or "multi_modal_data" not in prompt:
|
||||
raise OmniClientError("SeedVR2 requires multi_modal_data.video")
|
||||
if params.height is None or params.width is None:
|
||||
raise OmniClientError("SeedVR2 requires explicit output height and width")
|
||||
if min(params.height, params.width) < 16 or params.height % 16 or params.width % 16:
|
||||
raise OmniClientError("SeedVR2 output dimensions must be positive multiples of 16")
|
||||
validate_clip_size(1, params.height, params.width, frame_pixels, clip_pixels)
|
||||
frames = prompt["multi_modal_data"].get("video")
|
||||
source = None
|
||||
if isinstance(frames, list) and len(frames) == 1 and isinstance(frames[0], (str, Path)):
|
||||
frames = frames[0]
|
||||
if isinstance(frames, (str, Path)):
|
||||
frames, source = read_video(frames, frame_pixels, clip_pixels)
|
||||
if params.fps is not None and abs(params.fps - source.fps) > 1e-6:
|
||||
raise OmniClientError("SeedVR2 preserves source FPS; omit fps or match the input frame rate")
|
||||
if isinstance(frames, list) and frames and all(isinstance(frame, Image.Image) for frame in frames):
|
||||
for frame in frames:
|
||||
validate_clip_size(len(frames), frame.height, frame.width, frame_pixels, clip_pixels)
|
||||
frames = torch.stack([torch.from_numpy(np.array(frame.convert("RGB"))) for frame in frames])
|
||||
if not isinstance(frames, torch.Tensor):
|
||||
raise OmniClientError("SeedVR2 video must be a TCHW RGB tensor or a list of PIL frames")
|
||||
if frames.ndim == 4:
|
||||
frame_height, frame_width = frames.shape[1:3] if frames.dtype == torch.uint8 else frames.shape[2:]
|
||||
validate_clip_size(frames.shape[0], frame_height, frame_width, frame_pixels, clip_pixels)
|
||||
validate_clip_size(frames.shape[0], params.height, params.width, frame_pixels, clip_pixels)
|
||||
_check_frames(frames, params.height, params.width)
|
||||
if (prompt.get("prompt") or "").strip():
|
||||
raise OmniClientError("SeedVR2 uses fixed checkpoint conditioning; prompt text is unsupported")
|
||||
if params.num_outputs_per_prompt != 1:
|
||||
raise OmniClientError("SeedVR2 produces one restored video per request")
|
||||
request.prepared_layout = SeedVR2Input(frames, frames.shape[0], source)
|
||||
return request
|
||||
|
||||
|
||||
def get_seedvr2_pre_process_func(
|
||||
od_config: OmniDiffusionConfig,
|
||||
) -> Callable[[OmniDiffusionRequest], OmniDiffusionRequest]:
|
||||
# This factory runs in the serving process. Keep SeedVR2's CPU scheduling
|
||||
# policy local, and respect an explicit OMP_NUM_THREADS override.
|
||||
set_torch_threads_for_runtime()
|
||||
frame_pixels, clip_pixels = _admission_budget(od_config)
|
||||
return partial(prepare_request, frame_pixels=frame_pixels, clip_pixels=clip_pixels)
|
||||
|
||||
|
||||
def _seedvr2_post_process(output: dict[str, object]) -> dict[str, object]:
|
||||
payload = output["payload"]
|
||||
assert isinstance(payload, dict)
|
||||
video = payload["video"]
|
||||
assert isinstance(video, torch.Tensor) and video.dtype == torch.uint8
|
||||
# The device already produced uint8 [B,T,H,W,3], the layout encoders take.
|
||||
return {"payload": {**payload, "video": video.cpu().numpy()}, "metadata": output["metadata"]}
|
||||
|
||||
|
||||
def get_seedvr2_post_process_func(_od_config: OmniDiffusionConfig) -> Callable[[dict[str, object]], dict[str, object]]:
|
||||
return _seedvr2_post_process
|
||||
|
||||
|
||||
def finish_video(decoded: torch.Tensor, sample: torch.Tensor, frame_count: int, method: str) -> torch.Tensor:
|
||||
"""Colour-correct and quantize decoded ``[B,3,T,H,W]`` frames to uint8 ``[B,T,H,W,3]``.
|
||||
|
||||
One frame at a time: whole-clip float copies would dominate device memory on
|
||||
long clips, and a per-frame slice of one would need 64-bit indexing once the
|
||||
clip passes 2**31 pixels.
|
||||
"""
|
||||
batch_size, _, _, height, width = decoded.shape
|
||||
video = torch.empty((batch_size, frame_count, height, width, 3), device=decoded.device, dtype=torch.uint8)
|
||||
for index in range(frame_count):
|
||||
frame = slice(index, index + 1)
|
||||
# Restoration shifts global colour, so the resized input carries the
|
||||
# reference colour back onto the restored detail.
|
||||
corrected = correct_video_color(
|
||||
((decoded[:, :, frame].float() + 1) / 2).clamp(0, 1),
|
||||
((sample[:, :, frame].float() + 1) / 2).clamp(0, 1),
|
||||
method=method,
|
||||
)
|
||||
# Quantize on the device, matching the encoders' rint(clip(x) * 255), so
|
||||
# the host copy and the IPC payload are a quarter of the float size.
|
||||
video[:, index] = (corrected[:, :, 0].clamp(0, 1) * 255).round_().permute(0, 2, 3, 1)
|
||||
return video
|
||||
|
||||
|
||||
def sample_noise(condition: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
|
||||
# Reference randn_like preserves latent strides; RNG values follow storage order.
|
||||
return torch.empty_like(condition).normal_(generator=generator)
|
||||
|
||||
|
||||
class SeedVR2Pipeline(nn.Module):
|
||||
supports_request_batch = True
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_encoder_modules: ClassVar[list[str]] = []
|
||||
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
validate_seedvr2_config(od_config)
|
||||
set_torch_threads_for_runtime()
|
||||
self.device = get_local_device()
|
||||
self.od_config = od_config
|
||||
self.frame_pixels, self.clip_pixels = _admission_budget(od_config)
|
||||
self.transformer = SeedVR2NaDiT(**SEEDVR2_3B_CONFIG, use_varlen_kernel=False)
|
||||
self.vae = SeedVR2VAE()
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder=None,
|
||||
revision=od_config.revision,
|
||||
prefix=component + ".",
|
||||
fall_back_to_pt=False,
|
||||
allow_patterns_overrides=[filename],
|
||||
)
|
||||
for component, filename in (
|
||||
("transformer", "seedvr2_ema_3b_fp16.safetensors"),
|
||||
("vae", "ema_vae_fp16.safetensors"),
|
||||
)
|
||||
]
|
||||
model_path = Path(od_config.model)
|
||||
if not model_path.is_dir():
|
||||
model_path = Path(
|
||||
download_weights_from_hf_specific(od_config.model, None, ["pos_emb.pt"], revision=od_config.revision)
|
||||
)
|
||||
text = torch.load(model_path / "pos_emb.pt", map_location="cpu", weights_only=True)
|
||||
if not isinstance(text, torch.Tensor) or text.ndim != 2 or text.shape[1] != 5120 or text.shape[0] == 0:
|
||||
raise ValueError("SeedVR2 pos_emb.pt must contain a tensor of shape [L,5120]")
|
||||
self.register_buffer("text", text.to(device=self.device, dtype=od_config.dtype), persistent=False)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
targets = self.state_dict()
|
||||
loaded: set[str] = set()
|
||||
for name, weight in weights:
|
||||
if name.endswith(".rope.rope.freqs"):
|
||||
name = name.removesuffix(".rope.rope.freqs") + ".rope.freqs"
|
||||
if name not in targets or name in loaded:
|
||||
raise ValueError(f"Unexpected or duplicate SeedVR2 checkpoint key: {name}")
|
||||
if targets[name].shape != weight.shape:
|
||||
raise ValueError(f"SeedVR2 checkpoint shape mismatch for {name}")
|
||||
targets[name].copy_(weight)
|
||||
loaded.add(name)
|
||||
missing = targets.keys() - loaded
|
||||
if missing:
|
||||
raise ValueError(f"Missing SeedVR2 checkpoint keys: {sorted(missing)}")
|
||||
return loaded
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(self, batch: DiffusionRequestBatch) -> list[DiffusionOutput]:
|
||||
outputs = []
|
||||
for request in batch.requests:
|
||||
params = request.sampling_params
|
||||
if not isinstance(request.prepared_layout, SeedVR2Input):
|
||||
prepare_request(request, frame_pixels=self.frame_pixels, clip_pixels=self.clip_pixels)
|
||||
prepared = request.prepared_layout
|
||||
sample = prepare_video(prepared.frames, params.height, params.width, self.device)
|
||||
generator = params.generator
|
||||
if generator is None:
|
||||
generator = torch.Generator(device=self.device).manual_seed(params.seed)
|
||||
if not isinstance(generator, torch.Generator):
|
||||
raise ValueError("SeedVR2 accepts one generator per request")
|
||||
latent = self.vae.encode(sample).sample(generator=generator)
|
||||
condition = latent.permute(0, 2, 3, 4, 1).squeeze(0) * 0.9152
|
||||
noise = sample_noise(condition, generator)
|
||||
video = torch.cat((noise, condition, torch.ones_like(condition[..., :1])), dim=-1).reshape(-1, 33)
|
||||
shape = torch.tensor([condition.shape[:3]], device=self.device, dtype=torch.long)
|
||||
text_shape = torch.tensor([[self.text.shape[0]]], device=self.device, dtype=torch.long)
|
||||
timestep = torch.tensor([1000.0], device=self.device, dtype=torch.float16)
|
||||
runtime = self.transformer.build_runtime(
|
||||
self.transformer.token_grid_for(shape),
|
||||
text_len=self.text.shape[0],
|
||||
parallel_config=self.od_config.parallel_config,
|
||||
)
|
||||
velocity = self.transformer(video, self.text, shape, text_shape, timestep, runtime).vid_sample
|
||||
# Reference Euler returns fp32, then VAE casts to fp16 before scaling.
|
||||
restored = noise - velocity.reshape_as(noise)
|
||||
decoded = self.vae.decode((restored / 0.9152).permute(3, 0, 1, 2).unsqueeze(0))
|
||||
method = params.extra_args.get("color_correction_method") or DEFAULT_COLOR_CORRECTION_METHOD
|
||||
video = finish_video(decoded, sample, prepared.frame_count, method)
|
||||
payload: dict[str, object] = {"video": video}
|
||||
fps = prepared.source.fps if prepared.source is not None else params.fps
|
||||
metadata: dict[str, object] = {"video": {"fps": fps}} if fps is not None else {}
|
||||
if prepared.source is not None and prepared.source.audio is not None:
|
||||
payload["audio"] = prepared.source.audio
|
||||
metadata["audio"] = {"sample_rate": prepared.source.audio_sample_rate}
|
||||
outputs.append(DiffusionOutput(output={"payload": payload, "metadata": metadata}))
|
||||
return outputs
|
||||
@@ -0,0 +1,380 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""SeedVR2 causal VAE with temporal tiling and spatial height sharding."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from diffusers.models.attention_processor import Attention
|
||||
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
|
||||
from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import halo_exchange
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SpatialContext:
|
||||
group: dist.ProcessGroup
|
||||
rank: int
|
||||
units: tuple[int, ...]
|
||||
|
||||
def split(self, x: torch.Tensor) -> torch.Tensor:
|
||||
scale = x.shape[-2] // sum(self.units)
|
||||
start = sum(self.units[: self.rank]) * scale
|
||||
return x.narrow(-2, start, self.units[self.rank] * scale).contiguous()
|
||||
|
||||
def gather(self, x: torch.Tensor) -> torch.Tensor:
|
||||
scale = x.shape[-2] // self.units[self.rank]
|
||||
heights = [n * scale for n in self.units]
|
||||
padded = F.pad(x, (0, 0, 0, max(heights) - x.shape[-2])).contiguous()
|
||||
parts = [torch.empty_like(padded) for _ in heights]
|
||||
dist.all_gather(parts, padded, group=self.group)
|
||||
return torch.cat([part.narrow(-2, 0, h) for part, h in zip(parts, heights)], dim=-2)
|
||||
|
||||
def normalize(self, norm: nn.GroupNorm, x: torch.Tensor) -> torch.Tensor:
|
||||
b, c, t, h, w = x.shape
|
||||
grouped = x.to(torch.float32, copy=True).reshape(b, norm.num_groups, c // norm.num_groups, t, h, w)
|
||||
variance, mean = torch.var_mean(grouped, dim=(2, 4, 5), correction=0, keepdim=True)
|
||||
count = c // norm.num_groups * h * w
|
||||
moments = torch.stack((mean * count, torch.full_like(mean, count)))
|
||||
dist.all_reduce(moments, group=self.group)
|
||||
global_mean = moments[0] / moments[1]
|
||||
centered = (variance + (mean - global_mean).square()) * count
|
||||
dist.all_reduce(centered, group=self.group)
|
||||
normalized = grouped.sub_(global_mean).mul_(torch.rsqrt(centered / moments[1] + norm.eps)).reshape_as(x)
|
||||
normalized.mul_(norm.weight.float().view(1, -1, 1, 1, 1))
|
||||
normalized.add_(norm.bias.float().view(1, -1, 1, 1, 1))
|
||||
return normalized.to(x.dtype)
|
||||
|
||||
def convolution(self, conv: nn.Conv3d, x: torch.Tensor) -> torch.Tensor:
|
||||
halo = conv.kernel_size[1] // 2
|
||||
if halo:
|
||||
x, _, _ = halo_exchange(x, group=self.group, halo_size=halo)
|
||||
if conv.padding[1] == 0:
|
||||
# Downsampling uses right/bottom padding, not symmetric padding.
|
||||
x = x[:, :, :, halo:]
|
||||
return F.conv3d(x, conv.weight, conv.bias, conv.stride, (0, 0, conv.padding[2]), conv.dilation, conv.groups)
|
||||
|
||||
|
||||
def tiled_convolution(conv: nn.Conv3d, x: torch.Tensor, rows: int = 128) -> torch.Tensor:
|
||||
"""Tile output rows; only each tile's spatial padding/workspace is materialized."""
|
||||
shape = [(x.shape[i + 2] + 2 * conv.padding[i] - conv.kernel_size[i]) // conv.stride[i] + 1 for i in range(3)]
|
||||
output = x.new_empty((x.shape[0], conv.out_channels, *shape))
|
||||
stride, padding, kernel = conv.stride[1], conv.padding[1], conv.kernel_size[1]
|
||||
for start in range(0, shape[1], rows):
|
||||
end = min(start + rows, shape[1])
|
||||
lo, hi = start * stride - padding, (end - 1) * stride - padding + kernel
|
||||
tile = x[:, :, :, max(0, lo) : min(x.shape[-2], hi)]
|
||||
tile = F.pad(tile, (0, 0, max(0, -lo), max(0, hi - x.shape[-2])))
|
||||
output[:, :, :, start:end] = F.conv3d(
|
||||
tile, conv.weight, conv.bias, conv.stride, (0, 0, conv.padding[2]), conv.dilation, conv.groups
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
# Adapted from SeedVR, Copyright 2025 ByteDance Ltd., Apache-2.0.
|
||||
|
||||
|
||||
@dataclass
|
||||
class TemporalContext:
|
||||
"""Past convolution inputs owned by one encode or decode call."""
|
||||
|
||||
first: bool = True
|
||||
cache_temporal: bool = False
|
||||
tile_spatial: bool = False
|
||||
spatial: SpatialContext | None = None
|
||||
history: dict[nn.Module, torch.Tensor] = field(default_factory=dict)
|
||||
|
||||
|
||||
def _frame_norm_silu(norm: nn.GroupNorm, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
if context is not None and context.spatial is not None:
|
||||
return F.silu(context.spatial.normalize(norm, x))
|
||||
return F.silu(_frame_norm(norm, x))
|
||||
|
||||
|
||||
def _frame_norm(norm: nn.GroupNorm, x: torch.Tensor) -> torch.Tensor:
|
||||
frames = x.shape[2]
|
||||
x = norm(rearrange(x, "b c t h w -> (b t) c h w"))
|
||||
return rearrange(x, "(b t) c h w -> b c t h w", t=frames)
|
||||
|
||||
|
||||
class CausalConv3d(nn.Conv3d):
|
||||
"""Replicate the first frame to pad only the past temporal context."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: tuple[int, int, int] = (3, 3, 3),
|
||||
stride: tuple[int, int, int] = (1, 1, 1),
|
||||
padding: tuple[int, int, int] = (1, 1, 1),
|
||||
) -> None:
|
||||
super().__init__(in_channels, out_channels, kernel_size, stride, (0, padding[1], padding[2]))
|
||||
self.temporal_padding = 2 * padding[0]
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
if self.temporal_padding:
|
||||
if context is None or not context.cache_temporal or context.first:
|
||||
head = x[:, :, :1].expand(-1, -1, self.temporal_padding, -1, -1)
|
||||
else:
|
||||
head = context.history[self]
|
||||
x = torch.cat((head, x), dim=2)
|
||||
if context is not None and context.cache_temporal:
|
||||
# Copy the halo so it does not retain a whole chunk's storage.
|
||||
context.history[self] = x[:, :, -self.temporal_padding :].clone()
|
||||
if not context.first:
|
||||
x = x[:, :, self.stride[0] - 1 :]
|
||||
if context is not None:
|
||||
if context.spatial is not None:
|
||||
return context.spatial.convolution(self, x)
|
||||
if context.tile_spatial and x.shape[-2] > 256:
|
||||
return tiled_convolution(self, x)
|
||||
return super().forward(x)
|
||||
|
||||
|
||||
class ResnetBlock3d(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, groups: int) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = nn.GroupNorm(groups, in_channels, eps=1e-6)
|
||||
self.conv1 = CausalConv3d(in_channels, out_channels)
|
||||
self.norm2 = nn.GroupNorm(groups, out_channels, eps=1e-6)
|
||||
self.conv2 = CausalConv3d(out_channels, out_channels)
|
||||
self.conv_shortcut = (
|
||||
CausalConv3d(in_channels, out_channels, (1, 1, 1), padding=(0, 0, 0))
|
||||
if in_channels != out_channels
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
h = self.conv1(_frame_norm_silu(self.norm1, x, context), context)
|
||||
h = self.conv2(_frame_norm_silu(self.norm2, h, context), context)
|
||||
return (self.conv_shortcut(x, context) if self.conv_shortcut is not None else x) + h
|
||||
|
||||
|
||||
class Downsample3d(nn.Module):
|
||||
def __init__(self, channels: int, temporal: bool) -> None:
|
||||
super().__init__()
|
||||
self.conv = CausalConv3d(
|
||||
channels,
|
||||
channels,
|
||||
(3 if temporal else 1, 3, 3),
|
||||
(2 if temporal else 1, 2, 2),
|
||||
(1 if temporal else 0, 0, 0),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
bottom = 0 if context is not None and context.spatial is not None else 1
|
||||
return self.conv(F.pad(x, (0, 1, 0, bottom)), context)
|
||||
|
||||
|
||||
class Upsample3d(nn.Module):
|
||||
def __init__(self, channels: int, temporal: bool) -> None:
|
||||
super().__init__()
|
||||
self.temporal_ratio = 2 if temporal else 1
|
||||
self.upscale_conv = nn.Conv3d(channels, channels * 4 * self.temporal_ratio, kernel_size=1)
|
||||
self.conv = CausalConv3d(channels, channels)
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
x = rearrange(
|
||||
self.upscale_conv(x),
|
||||
"b (x y z c) t h w -> b c (t z) (h x) (w y)",
|
||||
x=2,
|
||||
y=2,
|
||||
z=self.temporal_ratio,
|
||||
)
|
||||
if self.temporal_ratio == 2 and (context is None or context.first):
|
||||
x = torch.cat((x[:, :, :1], x[:, :, 2:]), dim=2)
|
||||
return self.conv(x, context)
|
||||
|
||||
|
||||
class EncoderBlock3d(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, groups: int, index: int) -> None:
|
||||
super().__init__()
|
||||
self.resnets = nn.ModuleList(
|
||||
[ResnetBlock3d(in_channels, out_channels, groups), ResnetBlock3d(out_channels, out_channels, groups)]
|
||||
)
|
||||
self.downsamplers = nn.ModuleList([Downsample3d(out_channels, temporal=index >= 1)] if index < 3 else [])
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
for resnet in self.resnets:
|
||||
x = resnet(x, context)
|
||||
for downsampler in self.downsamplers:
|
||||
x = downsampler(x, context)
|
||||
return x
|
||||
|
||||
|
||||
class DecoderBlock3d(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, groups: int, index: int) -> None:
|
||||
super().__init__()
|
||||
self.resnets = nn.ModuleList(
|
||||
[ResnetBlock3d(in_channels, out_channels, groups)]
|
||||
+ [ResnetBlock3d(out_channels, out_channels, groups) for _ in range(2)]
|
||||
)
|
||||
self.upsamplers = nn.ModuleList([Upsample3d(out_channels, temporal=index < 2)] if index < 3 else [])
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
for resnet in self.resnets:
|
||||
x = resnet(x, context)
|
||||
for upsampler in self.upsamplers:
|
||||
x = upsampler(x, context)
|
||||
return x
|
||||
|
||||
|
||||
class MidBlock3d(nn.Module):
|
||||
def __init__(self, channels: int, groups: int) -> None:
|
||||
super().__init__()
|
||||
self.resnets = nn.ModuleList([ResnetBlock3d(channels, channels, groups) for _ in range(2)])
|
||||
self.attentions = nn.ModuleList(
|
||||
[
|
||||
Attention(
|
||||
channels,
|
||||
heads=1,
|
||||
dim_head=channels,
|
||||
eps=1e-6,
|
||||
norm_num_groups=groups,
|
||||
residual_connection=True,
|
||||
bias=True,
|
||||
upcast_softmax=True,
|
||||
_from_deprecated_attn_block=True,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
frames = x.shape[2]
|
||||
x = self.resnets[0](x, context)
|
||||
spatial = context.spatial if context is not None else None
|
||||
if spatial is not None:
|
||||
x = spatial.gather(x)
|
||||
x = self.attentions[0](rearrange(x, "b c t h w -> (b t) c h w"))
|
||||
x = rearrange(x, "(b t) c h w -> b c t h w", t=frames)
|
||||
if spatial is not None:
|
||||
x = spatial.split(x)
|
||||
return self.resnets[1](x, context)
|
||||
|
||||
|
||||
class Encoder3d(nn.Module):
|
||||
def __init__(self, channels: tuple[int, int, int, int], groups: int, latent_channels: int) -> None:
|
||||
super().__init__()
|
||||
self.conv_in = CausalConv3d(3, channels[0])
|
||||
self.down_blocks = nn.ModuleList(
|
||||
[EncoderBlock3d(channels[max(0, i - 1)], width, groups, i) for i, width in enumerate(channels)]
|
||||
)
|
||||
self.mid_block = MidBlock3d(channels[-1], groups)
|
||||
self.conv_norm_out = nn.GroupNorm(groups, channels[-1], eps=1e-6)
|
||||
self.conv_out = CausalConv3d(channels[-1], 2 * latent_channels)
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
x = self.conv_in(x, context)
|
||||
for block in self.down_blocks:
|
||||
x = block(x, context)
|
||||
x = self.mid_block(x, context)
|
||||
return self.conv_out(_frame_norm_silu(self.conv_norm_out, x, context), context)
|
||||
|
||||
|
||||
class Decoder3d(nn.Module):
|
||||
def __init__(self, channels: tuple[int, int, int, int], groups: int, latent_channels: int) -> None:
|
||||
super().__init__()
|
||||
self.conv_in = CausalConv3d(latent_channels, channels[-1])
|
||||
self.mid_block = MidBlock3d(channels[-1], groups)
|
||||
widths = tuple(reversed(channels))
|
||||
self.up_blocks = nn.ModuleList(
|
||||
[DecoderBlock3d(widths[max(0, i - 1)], width, groups, i) for i, width in enumerate(widths)]
|
||||
)
|
||||
self.conv_norm_out = nn.GroupNorm(groups, channels[0], eps=1e-6)
|
||||
self.conv_out = CausalConv3d(channels[0], 3)
|
||||
|
||||
def forward(self, x: torch.Tensor, context: TemporalContext | None = None) -> torch.Tensor:
|
||||
x = self.mid_block(self.conv_in(x, context), context)
|
||||
for block in self.up_blocks:
|
||||
x = block(x, context)
|
||||
return self.conv_out(_frame_norm_silu(self.conv_norm_out, x, context), context)
|
||||
|
||||
|
||||
class SeedVR2VAE(nn.Module, DistributedVaeMixin):
|
||||
"""Whole-clip VAE; sampling uses the caller's request-local generator."""
|
||||
|
||||
spatial_scale_factor = 8
|
||||
temporal_scale_factor = 4
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: tuple[int, int, int, int] = (128, 256, 512, 512),
|
||||
groups: int = 32,
|
||||
latent_channels: int = 16,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.use_tiling = False
|
||||
self._spatial_group: dist.ProcessGroup | None = None
|
||||
self.encoder = Encoder3d(channels, groups, latent_channels)
|
||||
self.decoder = Decoder3d(channels, groups, latent_channels)
|
||||
|
||||
def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
|
||||
if mode == "spatial_shard_width":
|
||||
raise ValueError("SeedVR2 VAE supports height sharding; select spatial_shard_height or tile")
|
||||
if parallel_size > 1:
|
||||
group = get_sp_group().device_group
|
||||
if dist.get_world_size(group) != parallel_size:
|
||||
raise ValueError("SeedVR2 VAE patch parallel size must match the SP group")
|
||||
self._spatial_group = group
|
||||
else:
|
||||
self._spatial_group = None
|
||||
|
||||
def _context(self, latent_height: int) -> TemporalContext:
|
||||
spatial = None
|
||||
if self.use_tiling and self._spatial_group is not None:
|
||||
group = self._spatial_group
|
||||
size = dist.get_world_size(group)
|
||||
if latent_height >= size:
|
||||
units = tuple((i + 1) * latent_height // size - i * latent_height // size for i in range(size))
|
||||
spatial = SpatialContext(group, dist.get_rank(group), units)
|
||||
return TemporalContext(tile_spatial=self.use_tiling, spatial=spatial)
|
||||
|
||||
def encode(self, x: torch.Tensor, *, chunk_size: int | None = None) -> DiagonalGaussianDistribution:
|
||||
if chunk_size is None and self.use_tiling:
|
||||
chunk_size = 8
|
||||
if chunk_size is not None and (chunk_size < 4 or chunk_size % 4):
|
||||
raise ValueError("SeedVR2 encode chunk_size must be a positive multiple of four")
|
||||
context = self._context(x.shape[-2] // 8)
|
||||
if context.spatial is not None:
|
||||
x = context.spatial.split(x)
|
||||
if chunk_size is None or x.shape[2] <= chunk_size + 1:
|
||||
encoded = self.encoder(x, context)
|
||||
else:
|
||||
context.cache_temporal = True
|
||||
chunks = [self.encoder(x[:, :, : chunk_size + 1], context)]
|
||||
context.first = False
|
||||
for start in range(chunk_size + 1, (x.shape[2] - 1) // 4 * 4 + 1, chunk_size):
|
||||
chunks.append(self.encoder(x[:, :, start : start + chunk_size], context))
|
||||
encoded = torch.cat(chunks, dim=2)
|
||||
if context.spatial is not None:
|
||||
encoded = context.spatial.gather(encoded)
|
||||
return DiagonalGaussianDistribution(encoded)
|
||||
|
||||
def decode(self, z: torch.Tensor, *, chunk_size: int | None = None) -> torch.Tensor:
|
||||
if chunk_size is None and self.use_tiling:
|
||||
chunk_size = 2
|
||||
if chunk_size is not None and chunk_size < 1:
|
||||
raise ValueError("SeedVR2 decode chunk_size must be positive")
|
||||
context = self._context(z.shape[-2])
|
||||
if context.spatial is not None:
|
||||
z = context.spatial.split(z)
|
||||
if chunk_size is None or z.shape[2] <= chunk_size:
|
||||
decoded = self.decoder(z, context)
|
||||
else:
|
||||
context.cache_temporal = True
|
||||
chunks = []
|
||||
for start in range(0, z.shape[2], chunk_size):
|
||||
chunks.append(self.decoder(z[:, :, start : start + chunk_size], context))
|
||||
context.first = False
|
||||
decoded = torch.cat(chunks, dim=2)
|
||||
if context.spatial is not None:
|
||||
decoded = context.spatial.gather(decoded)
|
||||
return decoded
|
||||
@@ -34,6 +34,7 @@ def resolve_native_single_file(model_class_name: str | None) -> str | None:
|
||||
|
||||
_DIFFUSION_MODELS = {
|
||||
# arch:(mod_folder, mod_relname, cls_name)
|
||||
"SeedVR2Pipeline": ("seedvr2", "pipeline_seedvr2", "SeedVR2Pipeline"),
|
||||
"QwenImagePipeline": (
|
||||
"qwen_image",
|
||||
"pipeline_qwen_image",
|
||||
@@ -693,6 +694,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"SanaVideoPipeline": "get_sana_video_post_process_func",
|
||||
"SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func",
|
||||
"Magi2Pipeline": "get_magi2_post_process_func",
|
||||
"SeedVR2Pipeline": "get_seedvr2_post_process_func",
|
||||
"OmniVoicePipeline": "get_omnivoice_post_process_func",
|
||||
"SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
|
||||
"Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
|
||||
@@ -716,6 +718,7 @@ _DIFFUSION_IR_OP_PRIORITY_FUNCS = {
|
||||
}
|
||||
|
||||
_DIFFUSION_PRE_PROCESS_FUNCS = {
|
||||
"SeedVR2Pipeline": "get_seedvr2_pre_process_func",
|
||||
# arch: pre_process_func
|
||||
# `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
|
||||
# where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
|
||||
|
||||
@@ -345,6 +345,13 @@ async def omni_run_server_worker(
|
||||
remove_route_from_app(app, "/v1/models", {"GET"}) # Remove upstream /v1/models to use omni's handler
|
||||
remove_route_from_app(app, "/health", {"GET"})
|
||||
app.include_router(router)
|
||||
get_od_config = getattr(engine_client, "get_diffusion_od_config", None)
|
||||
od_config = get_od_config() if callable(get_od_config) else None
|
||||
model_class_name = getattr(od_config, "model_class_name", None) or getattr(args, "model_class_name", None)
|
||||
if model_class_name == "SeedVR2Pipeline":
|
||||
from vllm_omni.diffusion.models.seedvr2.long_video import register_routes
|
||||
|
||||
register_routes(app, args.port)
|
||||
|
||||
# OMNI: Override upstream exception handlers with Omni-aware versions
|
||||
# that understand the multi-stage orchestrator lifecycle.
|
||||
|
||||
Reference in new issue
Block a user