Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cb46d63f07 | ||
|
|
bbbb7ab021 |
@@ -436,7 +436,7 @@ steps:
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 25m .buildkite/scripts/pr_test.sh"
|
||||
command: "timeout 15m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: LoRA Training Tests"
|
||||
env:
|
||||
- TEST_TYPE=training_lora
|
||||
|
||||
+2
-8
@@ -37,11 +37,6 @@ logs/
|
||||
official_weights/
|
||||
converted_weights/
|
||||
|
||||
# Cosmos3 local parity assets (symlinked from main worktree)
|
||||
/official_weights/
|
||||
/converted_weights/
|
||||
/cosmos-framework
|
||||
|
||||
# SSIM test outputs
|
||||
fastvideo/tests/ssim/generated_videos/
|
||||
**/.cache/**
|
||||
@@ -77,7 +72,8 @@ docs/distillation/examples/
|
||||
# Python pickle files
|
||||
*.pkl
|
||||
|
||||
# Reference videos (negations must come after the catch-all on line below)
|
||||
# Reference videos
|
||||
!fastvideo/tests/ssim/reference_videos/**/*.mp4
|
||||
|
||||
# Static images
|
||||
!docs/assets/images/**/*.png
|
||||
@@ -131,8 +127,6 @@ apps/dreamverse/web/.env.production.local
|
||||
.sisyphus/
|
||||
openspec/
|
||||
fastvideo/tests/ssim/reference_videos/**
|
||||
!fastvideo/tests/ssim/reference_videos/**/*.mp4
|
||||
!fastvideo/tests/ssim/reference_videos/**/*.png
|
||||
|
||||
# Editor logs and local Python version pins (accidentally committed)
|
||||
*.nvimlog
|
||||
|
||||
@@ -458,8 +458,6 @@ surfaces:
|
||||
guidance_scale: request.sampling.guidance_scale
|
||||
guidance_scale_2: request.sampling.guidance_scale_2
|
||||
guidance_rescale: request.sampling.guidance_rescale
|
||||
use_embedded_guidance: request.sampling.use_embedded_guidance
|
||||
true_cfg_scale: request.sampling.true_cfg_scale
|
||||
boundary_ratio: request.sampling.boundary_ratio
|
||||
sigmas: request.sampling.sigmas
|
||||
enable_teacache: request.runtime.enable_teacache
|
||||
|
||||
@@ -330,7 +330,7 @@ at FastVideo's CI — before the Dynamo-side integration even knows.
|
||||
internal; presets identify them by name on
|
||||
`PipelineSelection.preset`).
|
||||
* `fastvideo.fastvideo_args.FastVideoArgs` (legacy compat type).
|
||||
* `fastvideo.api.compat.*` private helpers
|
||||
* `fastvideo.api.translation.*` private helpers
|
||||
(`_validate_continuation_state` etc.) — the public boundary is
|
||||
`VideoGenerator` + `fastvideo.api`.
|
||||
* Any flat legacy LTX-2 kwarg (`ltx2_refine_upsampler_path`,
|
||||
|
||||
@@ -48,15 +48,17 @@ highest first:
|
||||
`request.model_fields_set` (Pydantic v2). Unset fields do not count,
|
||||
even if the Pydantic model has a schema default for them.
|
||||
2. **`ServeConfig.default_request` (operator-explicit)** — projected via
|
||||
[`explicit_request_updates()`](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/api/compat.py);
|
||||
[`explicit_request_updates()`](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/api/translation.py);
|
||||
only fields the operator actually wrote into the YAML count as
|
||||
defaults. Every other field inherits the schema default rather than
|
||||
being pinned.
|
||||
defaults (an explicit `null` counts as unset). Every other sampling
|
||||
field stays `None` — "inherit the model preset" — and other sections
|
||||
keep their schema defaults without being pinned.
|
||||
3. **Hardcoded fallback** — e.g. `fps = 24`.
|
||||
|
||||
The gate matters: both surfaces carry schema defaults. Without
|
||||
`model_fields_set` / explicit-path tracking, schema defaults would
|
||||
masquerade as intent and silently shadow the other side.
|
||||
The gate matters: the Pydantic surface carries schema defaults and the
|
||||
dataclass surface carries non-None defaults outside `sampling`. Without
|
||||
`model_fields_set` / explicit-path tracking, defaults would masquerade
|
||||
as intent and silently shadow the other side.
|
||||
|
||||
See [`video_api.py::_build_generation_kwargs`](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/entrypoints/openai/video_api.py)
|
||||
for the canonical implementation; the per-request assembly lives there,
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
# 🌊 AnyFlow Any-Step Video Distillation
|
||||
|
||||
**AnyFlow** ([paper](https://arxiv.org/abs/2605.13724), [project page](https://nvlabs.github.io/AnyFlow/), [official code](https://github.com/NVlabs/AnyFlow), [model weights](https://huggingface.co/collections/nvidia/anyflow)) is an any-step video diffusion framework built on flow maps. A single distilled checkpoint can be evaluated at NFE ∈ {1, 2, 4, 8, 16, 32} without retraining, and quality scales **monotonically** with steps — unlike consistency-based distillation, which often degrades as NFE grows.
|
||||
|
||||
The student network ``u_θ(x_t, t, r)`` predicts the *average velocity* from time ``t`` back to time ``r``, so one Euler step is
|
||||
|
||||
```
|
||||
x_r = x_t - ((t - r) / N) · u_θ(x_t, t, r)
|
||||
```
|
||||
|
||||
for any ``t > r``.
|
||||
|
||||
## 📊 Model Overview
|
||||
|
||||
NVIDIA publishes four checkpoints under [`nvidia/anyflow`](https://huggingface.co/collections/nvidia/anyflow):
|
||||
|
||||
- `nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers` — bidirectional T2V, Wan2.1 1.3B base
|
||||
- `nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers` — bidirectional T2V, Wan2.1 14B base
|
||||
- `nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers` — frame-autoregressive variant, 1.3B
|
||||
- `nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers` — frame-autoregressive variant, 14B
|
||||
|
||||
FastVideo currently supports the bidirectional T2V variants for training; the FAR variants can be loaded for inference through the diffusers integration.
|
||||
|
||||
## ⚙️ Inference
|
||||
|
||||
For inference, load the published checkpoint directly through diffusers; FastVideo's training-side ``WanModel`` config maps the HF AnyFlow ``delta_embedder`` weights onto its internal layout via ``param_names_mapping`` so the same checkpoint can be used as the ``init_from`` for the on-policy YAML below.
|
||||
|
||||
## 🧠 Algorithm
|
||||
|
||||
Training runs in two stages. Both use the dual-timestep Wan backbone — enabled by ``pipeline.dit_config.r_embedder: true`` in the YAML, which allocates a sibling ``condition_embedder.delta_embedder`` and fuses its embedding with the standard timestep embedding via either an additive or a gated mixer.
|
||||
|
||||
### Stage 1 — Pretrain (flow-map central-difference)
|
||||
|
||||
Method: ``AnyFlowPretrainMethod`` (``fastvideo/train/methods/distribution_matching/anyflow_pretrain.py``)
|
||||
|
||||
For each batch, sample ``(t, r) ∈ [0, 1]`` as ``(max, min)`` of two uniform draws, then:
|
||||
|
||||
- a ``diffusion_ratio`` fraction (default 0.5) gets ``r = t`` — recovers plain flow matching;
|
||||
- a ``consistency_ratio`` fraction (default 0.25) gets ``r = 0`` — forces consistency to clean data;
|
||||
- the remainder is free.
|
||||
|
||||
The student forward at ``(t, r)`` is trained against the central-difference target
|
||||
|
||||
```
|
||||
target = (eps - x_0) - (t - r) · dF/dt
|
||||
```
|
||||
|
||||
where ``dF/dt`` is estimated from the student's own forward at ``(t ± δ, r)`` with the sample also moved along the flow trajectory by ``v_pred · (δ / N)``. Per-timestep weighting uses ``beta08`` (``w(t) = t · sqrt(1 - t)``, renormalized). A stop-gradient scale-balance keeps the non-diffusion branches' loss magnitude aligned with the diffusion branch.
|
||||
|
||||
### Stage 2 — On-policy DMD
|
||||
|
||||
Method: ``AnyFlowMethod`` (``fastvideo/train/methods/distribution_matching/anyflow.py``)
|
||||
|
||||
Inherits ``DMD2Method``. The student is rolled out for ``student_sample_steps`` Euler-flow steps from pure noise; one randomly-chosen step is gradient-enabled (broadcast from rank 0 so every worker agrees), the rest run under ``torch.no_grad``. With ``use_mean_velocity: true`` (default) the rollout uses ``r = t_next`` at each step, matching AnyFlow's ``WanAnyFlowPipeline.training_rollout``.
|
||||
|
||||
The inherited ``_dmd_loss`` (VSD with fake-score critic) consumes the rollout output and the teacher's CFG prediction. The optional pinned ``t_list_override`` lets configs reproduce the paper's hand-tuned 4-step schedule ``[999, 937, 833, 624, 0]``.
|
||||
|
||||
## 🚀 Training Scripts
|
||||
|
||||
### Stage 1 — pretrain
|
||||
|
||||
```bash
|
||||
bash examples/train/run.sh \
|
||||
examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml
|
||||
```
|
||||
|
||||
**Key configuration** (in ``examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml``):
|
||||
|
||||
- Global batch size: 32 (8 GPUs × 4 per-GPU)
|
||||
- Learning rate: 5e-5
|
||||
- Flow shift: 5.0
|
||||
- ``diffusion_ratio`` / ``consistency_ratio``: 0.5 / 0.25
|
||||
- ``epsilon`` (finite-difference step): 5 (absolute train-timestep units)
|
||||
- ``weight_type``: ``beta08``
|
||||
- ``fuse_guidance_scale``: 3.0
|
||||
- Training steps: 6000
|
||||
|
||||
### Stage 2 — on-policy
|
||||
|
||||
```bash
|
||||
bash examples/train/run.sh \
|
||||
examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml \
|
||||
--models.student.init_from outputs/wan2.1_anyflow_pretrain/checkpoint-final
|
||||
```
|
||||
|
||||
(Or point ``models.student.init_from`` directly at ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers`` to bootstrap from the paper weights and skip Stage 1.)
|
||||
|
||||
**Key configuration**:
|
||||
|
||||
- Global batch size: 8 (8 GPUs × 1 per-GPU)
|
||||
- Learning rate: 2e-6
|
||||
- Flow shift: 5.0
|
||||
- ``student_sample_steps``: 4
|
||||
- ``t_list_override``: ``[999, 937, 833, 624, 0]``
|
||||
- ``use_mean_velocity``: ``true`` (i.e. ``r = t_next`` during rollout)
|
||||
- ``real_score_guidance_scale``: 3.0
|
||||
- ``generator_update_interval``: 5 (DMD2 alternation)
|
||||
- Training steps: 4000
|
||||
|
||||
## 🔌 Loading published AnyFlow checkpoints
|
||||
|
||||
The HF AnyFlow checkpoints expose ``condition_embedder.delta_embedder.*`` weights that FastVideo internally maps onto its ``condition_embedder.delta_embedder.mlp.*`` layout. This rename happens automatically through the regex in ``WanVideoArchConfig.param_names_mapping`` — no separate adapter is needed. The same regex is a no-op on plain Wan checkpoints (which don't contain any ``delta_embedder`` keys).
|
||||
|
||||
Set the YAML's ``pipeline.dit_config.r_embedder: true`` to allocate the ``delta_embedder`` module on the FastVideo side; when initializing from a plain Wan checkpoint the delta weights are deep-copied from ``time_embedder`` (matching AnyFlow's ``setup_flowmap_model()`` behavior).
|
||||
|
||||
## 🧭 Note on ``fuse_guidance_scale``
|
||||
|
||||
Stage 1 optionally fuses classifier-free guidance into the training target so the resulting checkpoint can be sampled at ``guidance_scale=1.0`` (no extra forward pass at inference time). The transformation is
|
||||
|
||||
```
|
||||
noise_pred ← (noise_pred - (1 - g) · noise_pred_uncond) / g
|
||||
```
|
||||
|
||||
with ``g = fuse_guidance_scale``. The negative prompt embedding comes from ``WanModel``'s ``ensure_negative_conditioning()`` — i.e. the dataset's configured ``sampling_param.negative_prompt``. Setting ``fuse_guidance_scale: 1.0`` skips the extra unconditional forward entirely.
|
||||
|
||||
The on-policy stage's ``real_score_guidance_scale`` (inherited from DMD2) follows the same parameterization conventions documented in [``dmd.md``](dmd.md#-note-on-real_score_guidance_scale).
|
||||
@@ -39,17 +39,20 @@ All you need to generate videos using multi-gpus from state-of-the-art diffusion
|
||||
|
||||
```python
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import EngineConfig, GenerationRequest, GeneratorConfig
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
num_gpus=1,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
)
|
||||
)
|
||||
|
||||
prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
|
||||
video = generator.generate_video(prompt)
|
||||
result = generator.generate(GenerationRequest(prompt=prompt))
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig, OutputConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples"
|
||||
def main():
|
||||
@@ -8,29 +9,32 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
# sampling_param.num_frames = 45
|
||||
# sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
|
||||
# Generate videos with the same simple API, regardless of GPU count
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
|
||||
# video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
|
||||
video = generator.generate(
|
||||
GenerationRequest(prompt=prompt, output=OutputConfig(output_path=OUTPUT_PATH, save_video=True)))
|
||||
|
||||
# Generate another video with a different prompt, without reloading the
|
||||
# model!
|
||||
@@ -40,7 +44,8 @@ def main():
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(prompt=prompt2, output=OutputConfig(output_path=OUTPUT_PATH, save_video=True)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,24 +1,30 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
# Point this to your local diffusers model dir (or replace with a HF model ID).
|
||||
model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_path,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
sampling_param = SamplingParam.from_pretrained(model_path)
|
||||
|
||||
# image2world example from official repo
|
||||
image_path = "assets/images/bus_terminal.jpg"
|
||||
|
||||
@@ -33,13 +39,16 @@ def main():
|
||||
"Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene."
|
||||
)
|
||||
|
||||
generator.generate_video(
|
||||
prompt,
|
||||
sampling_param=sampling_param,
|
||||
image_path=str(image_path),
|
||||
num_cond_frames=1,
|
||||
output_path="outputs_video/cosmos2_5_i2w.mp4",
|
||||
save_video=True,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=str(image_path)),
|
||||
output=OutputConfig(
|
||||
output_path="outputs_video/cosmos2_5_i2w.mp4",
|
||||
save_video=True,
|
||||
),
|
||||
extensions={"num_cond_frames": 1},
|
||||
)
|
||||
)
|
||||
|
||||
generator.shutdown()
|
||||
@@ -47,4 +56,3 @@ def main():
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
@@ -1,24 +1,29 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig, OutputConfig,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
# Point this to your local diffusers model dir (or replace with a HF model ID).
|
||||
model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_path,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Load default sampling parameters (negative_prompt, resolution, steps, etc.)
|
||||
sampling_param = SamplingParam.from_pretrained(model_path)
|
||||
|
||||
prompt = (
|
||||
"A high-definition video captures the precision of robotic welding in an industrial setting. "
|
||||
"The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. "
|
||||
@@ -34,11 +39,14 @@ def main():
|
||||
"underscoring the ongoing nature of the welding operation."
|
||||
)
|
||||
|
||||
generator.generate_video(
|
||||
prompt,
|
||||
sampling_param=sampling_param,
|
||||
output_path="outputs_video/cosmos2_5_t2w.mp4",
|
||||
save_video=True,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
output=OutputConfig(
|
||||
output_path="outputs_video/cosmos2_5_t2w.mp4",
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
generator.shutdown()
|
||||
@@ -46,6 +54,3 @@ def main():
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,23 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig,
|
||||
OffloadConfig, OutputConfig,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
# Point this to your local diffusers model dir (or replace with a HF model ID).
|
||||
model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_path,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
)
|
||||
|
||||
sampling_param = SamplingParam.from_pretrained(model_path)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
# video2world example from official repo
|
||||
video_path = "assets/videos/robot_pouring.mp4"
|
||||
@@ -36,18 +42,19 @@ def main():
|
||||
"The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper."
|
||||
)
|
||||
|
||||
generator.generate_video(
|
||||
prompt,
|
||||
sampling_param=sampling_param,
|
||||
video_path=str(video_path),
|
||||
num_cond_frames=1,
|
||||
output_path="outputs_video/cosmos2_5_v2w.mp4",
|
||||
save_video=True,
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(video_path=str(video_path)),
|
||||
output=OutputConfig(
|
||||
output_path="outputs_video/cosmos2_5_v2w.mp4",
|
||||
save_video=True,
|
||||
),
|
||||
extensions={"num_cond_frames": 1},
|
||||
))
|
||||
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
@@ -1,81 +0,0 @@
|
||||
import os
|
||||
import time
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
# NVIDIA Cosmos3-Nano omni world model — image-to-video (I2V) path through
|
||||
# FastVideo's native Cosmos3 pipeline. The input image conditions latent frame 0
|
||||
# (kept clean during denoising); the rest of the clip is generated to follow it.
|
||||
# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g.
|
||||
# ``official_weights/cosmos3``) to skip the Hugging Face download.
|
||||
OUTPUT_PATH = "video_samples_cosmos3_i2v"
|
||||
|
||||
|
||||
def main():
|
||||
model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
|
||||
image_path = os.environ.get("COSMOS3_IMAGE_PATH", "assets/images/cyclist.jpg")
|
||||
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
dit=False,
|
||||
vae=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
load_start_time = time.perf_counter()
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
load_time = time.perf_counter() - load_start_time
|
||||
|
||||
prompt = (
|
||||
"A mountain biker rides forward along the sunlit forest trail, wheels "
|
||||
"kicking up dust as trees and dappled light sweep past, smooth cinematic "
|
||||
"tracking shot from behind."
|
||||
)
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
sampling=SamplingConfig(
|
||||
# cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps;
|
||||
# overridable via env for quick smoke runs.
|
||||
num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
|
||||
height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
|
||||
width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
|
||||
num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
|
||||
guidance_scale=6.0,
|
||||
fps=24,
|
||||
seed=1024,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
result = generator.generate(request)
|
||||
gen_time = time.perf_counter() - start_time
|
||||
|
||||
print(f"Time taken to load model: {load_time} seconds")
|
||||
print(f"Time taken to generate video: {gen_time} seconds")
|
||||
print(f"Output written to: {result.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,77 +0,0 @@
|
||||
import os
|
||||
import time
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
# NVIDIA Cosmos3-Nano omni world model — this example exercises the
|
||||
# text-to-video (T2V) path through FastVideo's native Cosmos3 pipeline.
|
||||
# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g.
|
||||
# ``official_weights/cosmos3``) to skip the Hugging Face download.
|
||||
OUTPUT_PATH = "video_samples_cosmos3"
|
||||
|
||||
|
||||
def main():
|
||||
model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
|
||||
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
dit=False,
|
||||
vae=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
load_start_time = time.perf_counter()
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
load_time = time.perf_counter() - load_start_time
|
||||
|
||||
prompt = (
|
||||
"A golden retriever puppy runs across a sunlit meadow toward the camera, "
|
||||
"ears flopping and wildflowers swaying in the breeze. Shallow depth of "
|
||||
"field, warm afternoon light, smooth cinematic tracking shot."
|
||||
)
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(
|
||||
# cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps;
|
||||
# overridable via env for quick smoke runs.
|
||||
num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
|
||||
height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
|
||||
width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
|
||||
num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
|
||||
guidance_scale=6.0,
|
||||
fps=24,
|
||||
seed=1024,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
result = generator.generate(request)
|
||||
gen_time = time.perf_counter() - start_time
|
||||
|
||||
print(f"Time taken to load model: {load_time} seconds")
|
||||
print(f"Time taken to generate video: {gen_time} seconds")
|
||||
print(f"Output written to: {result.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,78 +0,0 @@
|
||||
import os
|
||||
import time
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
# NVIDIA Cosmos3-Nano omni world model — text-to-image (T2I) path through
|
||||
# FastVideo's native Cosmos3 pipeline. T2I is the single-frame case
|
||||
# (num_frames=1); the canonical Cosmos3 T2I resolution is 960x960 (the model's
|
||||
# "720" bucket, UniPC flow_shift=10.0). Point COSMOS3_MODEL_PATH at a local
|
||||
# diffusers checkpoint (e.g. ``official_weights/cosmos3``) to skip the download.
|
||||
OUTPUT_PATH = "video_samples_cosmos3_t2i"
|
||||
|
||||
|
||||
def main():
|
||||
model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
|
||||
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
dit=False,
|
||||
vae=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
load_start_time = time.perf_counter()
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
load_time = time.perf_counter() - load_start_time
|
||||
|
||||
prompt = (
|
||||
"A photograph of a red panda sitting on a mossy log in a misty bamboo "
|
||||
"forest, soft golden morning light filtering through the leaves, shallow "
|
||||
"depth of field, crisp fur detail, serene atmosphere."
|
||||
)
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(
|
||||
# T2I is single-frame; canonical Cosmos3 T2I is 960x960. Overridable
|
||||
# via env for quick smoke runs.
|
||||
num_frames=1,
|
||||
height=int(os.environ.get("COSMOS3_HEIGHT", "960")),
|
||||
width=int(os.environ.get("COSMOS3_WIDTH", "960")),
|
||||
num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
|
||||
guidance_scale=6.0,
|
||||
fps=24,
|
||||
seed=1024,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
result = generator.generate(request)
|
||||
gen_time = time.perf_counter() - start_time
|
||||
|
||||
print(f"Time taken to load model: {load_time} seconds")
|
||||
print(f"Time taken to generate image: {gen_time} seconds")
|
||||
print(f"Output written to: {result.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,67 +0,0 @@
|
||||
import os
|
||||
import time
|
||||
|
||||
# t2vs (text -> video + sound). The Cosmos3 denoise stage generates a joint
|
||||
# [vision | sound] latent and AVAE-decodes the sound to a waveform muxed into the
|
||||
# mp4. The joint-sound path is gated on COSMOS3_T2VS (set here for the example).
|
||||
os.environ.setdefault("COSMOS3_T2VS", "1")
|
||||
|
||||
from fastvideo import VideoGenerator # noqa: E402
|
||||
from fastvideo.api import ( # noqa: E402
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_cosmos3_t2vs"
|
||||
|
||||
|
||||
def main():
|
||||
model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
|
||||
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(text_encoder=True, pin_cpu_memory=True, dit=False, vae=False),
|
||||
),
|
||||
)
|
||||
|
||||
load_start = time.perf_counter()
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
load_time = time.perf_counter() - load_start
|
||||
|
||||
prompt = (
|
||||
"Ocean waves crash against a rocky shore at sunset, white foam spraying "
|
||||
"into the air as seagulls wheel overhead. Golden light, cinematic wide "
|
||||
"shot, the rhythmic roar of the surf."
|
||||
)
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(
|
||||
num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
|
||||
height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
|
||||
width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
|
||||
num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
|
||||
guidance_scale=6.0,
|
||||
fps=24,
|
||||
seed=1024,
|
||||
),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True, return_frames=False),
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
result = generator.generate(request)
|
||||
gen_time = time.perf_counter() - start
|
||||
|
||||
print(f"Time taken to load model: {load_time} seconds")
|
||||
print(f"Time taken to generate video+sound: {gen_time} seconds")
|
||||
print(f"Output written to: {result.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -2,7 +2,9 @@ import os
|
||||
import time
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, PipelineSelection,
|
||||
SamplingConfig)
|
||||
|
||||
OUTPUT_PATH = "video_samples_dmd2"
|
||||
def main():
|
||||
@@ -10,30 +12,36 @@ def main():
|
||||
|
||||
load_start_time = time.perf_counter()
|
||||
model_name = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
# Adjust these offload parameters if you have < 32GB of VRAM
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
VSA_sparsity=0.8,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
# Adjust these offload parameters if you have < 32GB of VRAM
|
||||
offload=OffloadConfig(
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
dit=False,
|
||||
vae=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(experimental={"VSA_sparsity": 0.8}),
|
||||
))
|
||||
load_end_time = time.perf_counter()
|
||||
load_time = load_end_time - load_start_time
|
||||
|
||||
|
||||
sampling_param = SamplingParam.from_pretrained(model_name)
|
||||
sampling_param.num_frames = 81
|
||||
|
||||
prompt = (
|
||||
"A neon-lit alley in futuristic Tokyo during a heavy rainstorm at night. The puddles reflect glowing signs in kanji, advertising ramen, karaoke, and VR arcades. A woman in a translucent raincoat walks briskly with an LED umbrella. Steam rises from a street food cart, and a cat darts across the screen. Raindrops are visible on the camera lens, creating a cinematic bokeh effect."
|
||||
)
|
||||
start_time = time.perf_counter()
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_frames=81),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
end_time = time.perf_counter()
|
||||
gen_time = end_time - start_time
|
||||
|
||||
@@ -46,7 +54,12 @@ def main():
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
start_time = time.perf_counter()
|
||||
video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, num_frames=81)
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
sampling=SamplingConfig(num_frames=81),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
end_time = time.perf_counter()
|
||||
gen_time2 = end_time - start_time
|
||||
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import os
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (ComponentConfig, EngineConfig, GenerationRequest,
|
||||
GeneratorConfig, InputConfig, OffloadConfig,
|
||||
OutputConfig, PipelineSelection, SamplingConfig)
|
||||
|
||||
|
||||
OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world")
|
||||
@@ -16,16 +19,22 @@ def _env_float(name: str, default: float) -> float:
|
||||
|
||||
def main():
|
||||
model_name = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
override_pipeline_cls_name="DreamXWorldPipeline",
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(override_pipeline_cls_name="DreamXWorldPipeline"), ),
|
||||
))
|
||||
|
||||
prompt = os.getenv(
|
||||
"DREAMX_WORLD_PROMPT",
|
||||
@@ -37,25 +46,31 @@ def main():
|
||||
"https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG",
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"output_path": OUTPUT_PATH,
|
||||
"save_video": os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
|
||||
"height": _env_int("DREAMX_WORLD_HEIGHT", 480),
|
||||
"width": _env_int("DREAMX_WORLD_WIDTH", 832),
|
||||
"num_frames": _env_int("DREAMX_WORLD_NUM_FRAMES", 161),
|
||||
"num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30),
|
||||
"guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0),
|
||||
"action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
|
||||
"action_speed_list": [
|
||||
float(value)
|
||||
for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")
|
||||
],
|
||||
}
|
||||
if image_path:
|
||||
kwargs["image_path"] = image_path
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path or None),
|
||||
sampling=SamplingConfig(
|
||||
height=_env_int("DREAMX_WORLD_HEIGHT", 480),
|
||||
width=_env_int("DREAMX_WORLD_WIDTH", 832),
|
||||
num_frames=_env_int("DREAMX_WORLD_NUM_FRAMES", 161),
|
||||
num_inference_steps=_env_int("DREAMX_WORLD_STEPS", 30),
|
||||
guidance_scale=_env_float("DREAMX_WORLD_GUIDANCE", 5.0),
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
|
||||
),
|
||||
extensions={
|
||||
"action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
|
||||
"action_speed_list": [
|
||||
float(value)
|
||||
for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
generator.generate_video(prompt, **kwargs)
|
||||
generator.generate(request)
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
@@ -1,140 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import os
|
||||
import re
|
||||
|
||||
DEFAULT_PROMPTS = [
|
||||
"a photo of a cat",
|
||||
(
|
||||
"a cinematic photo of a red panda wearing a tiny backpack, standing on a "
|
||||
"rainy neon-lit street at night, shallow depth of field, sharp focus, "
|
||||
"35mm, bokeh"
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _safe_filename(text: str, max_len: int = 100) -> str:
|
||||
"""Make a stable, filesystem-friendly filename base."""
|
||||
s = text[:max_len].strip()
|
||||
s = s.replace(os.sep, "_")
|
||||
if os.altsep:
|
||||
s = s.replace(os.altsep, "_")
|
||||
s = re.sub(r"\s+", " ", s)
|
||||
s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
|
||||
s = s.strip(" .")
|
||||
return s or "prompt"
|
||||
|
||||
|
||||
def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
|
||||
"""Delete prior outputs so reruns do not get _1, _2 suffixes."""
|
||||
if not os.path.isdir(out_dir):
|
||||
return
|
||||
|
||||
pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
|
||||
for fn in os.listdir(out_dir):
|
||||
if pattern.match(fn):
|
||||
with contextlib.suppress(FileNotFoundError):
|
||||
os.remove(os.path.join(out_dir, fn))
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(
|
||||
description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--model-path",
|
||||
default="official_weights/FLUX.1-dev",
|
||||
help="Local Diffusers checkpoint dir or HF repo id.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--out-dir",
|
||||
"--outdir",
|
||||
default="outputs/flux_dev/samples",
|
||||
help="Directory for saved PNG outputs.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--prompt",
|
||||
action="append",
|
||||
default=None,
|
||||
help="Prompt. Repeat for multiple images.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--backend",
|
||||
default=None,
|
||||
help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA).",
|
||||
)
|
||||
p.add_argument("--seed", type=int, default=42, help="Base seed; each prompt uses seed + index.")
|
||||
p.add_argument("--height", type=int, default=1024, help="Output height.")
|
||||
p.add_argument("--width", type=int, default=1024, help="Output width.")
|
||||
p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
|
||||
p.add_argument("--guidance", type=float, default=3.5, help="Guidance scale.")
|
||||
p.add_argument("--num-gpus", type=int, default=1, help="GPU count.")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
prompts: list[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
|
||||
|
||||
if args.backend:
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
init_kwargs = {
|
||||
"num_gpus": args.num_gpus,
|
||||
"workload_type": "t2i",
|
||||
"sp_size": 1,
|
||||
"tp_size": 1,
|
||||
"dit_cpu_offload": False,
|
||||
"dit_layerwise_offload": False,
|
||||
"text_encoder_cpu_offload": False,
|
||||
"vae_cpu_offload": False,
|
||||
"image_encoder_cpu_offload": False,
|
||||
"pin_cpu_memory": False,
|
||||
"use_fsdp_inference": False,
|
||||
}
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_path=args.model_path,
|
||||
**init_kwargs,
|
||||
)
|
||||
try:
|
||||
for i, prompt in enumerate(prompts):
|
||||
seed = args.seed + i
|
||||
filename_base = (
|
||||
f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
|
||||
)
|
||||
_remove_existing_outputs(args.out_dir, filename_base)
|
||||
output_path = os.path.join(args.out_dir, f"{filename_base}.png")
|
||||
print(f"[flux] prompt_idx={i} seed={seed} output_path={output_path}")
|
||||
|
||||
generation_kwargs = {
|
||||
"output_path": output_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": 1,
|
||||
"fps": 1,
|
||||
"num_inference_steps": args.steps,
|
||||
"guidance_scale": args.guidance,
|
||||
"use_embedded_guidance": True,
|
||||
"true_cfg_scale": 1.0,
|
||||
"seed": seed,
|
||||
"save_video": True,
|
||||
}
|
||||
|
||||
generator.generate_video(prompt, **generation_kwargs)
|
||||
|
||||
print(f"[flux] done. outputs written to: {args.out_dir}")
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -26,6 +26,15 @@ import os
|
||||
import torch
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.models.camera import create_camera_trajectory
|
||||
|
||||
# Model configuration (use GAMECRAFT_MODEL_PATH for local weights)
|
||||
@@ -55,14 +64,20 @@ OUTPUT_PATH = "video_samples_gamecraft"
|
||||
def main():
|
||||
# Initialize generator
|
||||
# FastVideo will automatically download weights from HuggingFace
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
MODEL_PATH,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=MODEL_PATH,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Video parameters
|
||||
@@ -96,23 +111,27 @@ def main():
|
||||
prompt = DEFAULT_I2V_PROMPT if is_i2v else DEFAULT_PROMPTS["temple"]
|
||||
print(f"Mode: {'I2V' if is_i2v else 'T2V'}, prompt: {prompt[:60]}...")
|
||||
|
||||
gen_kw = dict(
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt="",
|
||||
camera_states=camera_states,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=50,
|
||||
guidance_scale=6.0,
|
||||
seed=42,
|
||||
fps=24,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
sampling=SamplingConfig(
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=50,
|
||||
guidance_scale=6.0,
|
||||
seed=42,
|
||||
fps=24,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
extensions={"camera_states": camera_states},
|
||||
)
|
||||
if is_i2v:
|
||||
gen_kw["image_path"] = image_path
|
||||
generator.generate_video(**gen_kw)
|
||||
request.inputs = InputConfig(image_path=image_path)
|
||||
generator.generate(request)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -22,6 +22,10 @@ Requirements:
|
||||
import argparse
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig,
|
||||
OffloadConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
@@ -74,33 +78,47 @@ def main():
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
args = parser.parse_args()
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
args.model_path,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
video = generator.generate_video(
|
||||
args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
image_path=args.image_path,
|
||||
trajectory_type=args.trajectory,
|
||||
movement_distance=args.movement_distance,
|
||||
camera_rotation=args.camera_rotation,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
fps=24,
|
||||
seed=args.seed,
|
||||
output_path=args.output_path,
|
||||
save_video=True,
|
||||
)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
inputs=InputConfig(
|
||||
image_path=args.image_path,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
fps=24,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=args.output_path,
|
||||
save_video=True,
|
||||
),
|
||||
extensions={
|
||||
"trajectory_type": args.trajectory,
|
||||
"movement_distance": args.movement_distance,
|
||||
"camera_rotation": args.camera_rotation,
|
||||
},
|
||||
))
|
||||
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
@@ -1,107 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Run GLM-Image text-to-image generation through FastVideo.
|
||||
|
||||
User story:
|
||||
"I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
|
||||
text-to-image generation command, saved as a PNG."
|
||||
"""
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default="zai-org/GLM-Image",
|
||||
help="HF id or local diffusers-format GLM-Image weights directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="image_output/landscape.png",
|
||||
help="Output PNG path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default=("A beautiful landscape photography with rolling hills, "
|
||||
"a winding river, and a vibrant sunset in the background. "
|
||||
"Warm golden light, photorealistic style."),
|
||||
help="Text prompt.",
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--steps", type=int, default=50)
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.5)
|
||||
parser.add_argument("--seed", type=int, default=1024)
|
||||
parser.add_argument("--num-gpus", type=int, default=1)
|
||||
parser.add_argument("--tp-size", type=int, default=None)
|
||||
parser.add_argument("--sp-size", type=int, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
output = Path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
|
||||
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
|
||||
|
||||
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
|
||||
# pipeline class come from the model's registered defaults — don't override.
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
trust_remote_code=True,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
|
||||
),
|
||||
pipeline=PipelineSelection(workload_type="t2i"),
|
||||
)
|
||||
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
try:
|
||||
request = GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=1,
|
||||
fps=1,
|
||||
num_inference_steps=args.steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output.parent),
|
||||
save_video=False,
|
||||
return_frames=True,
|
||||
),
|
||||
)
|
||||
result = generator.generate(request)
|
||||
if isinstance(result, list):
|
||||
result = result[0]
|
||||
|
||||
frames = result.frames
|
||||
if frames is not None and len(frames):
|
||||
Image.fromarray(frames[0]).save(output)
|
||||
print(f"Saved image to {output}")
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,6 +1,13 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
import json
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
OUTPUT_PATH = "video_samples_hy15"
|
||||
def main():
|
||||
@@ -8,17 +15,21 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
|
||||
generator = VideoGenerator.from_config(GeneratorConfig(
|
||||
model_path="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder_cpu_offload=False,
|
||||
)
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
@@ -26,7 +37,12 @@ def main():
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(num_frames=81, fps=16),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
prompt2 = (
|
||||
"A majestic lion strides across the golden savanna, its powerful frame "
|
||||
@@ -35,8 +51,13 @@ def main():
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
|
||||
video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt2,
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(num_frames=81, fps=16),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -1,6 +1,12 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
)
|
||||
import json
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
OUTPUT_PATH = "video_samples_hy15_1080p"
|
||||
def main():
|
||||
@@ -8,17 +14,23 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p
|
||||
# or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p
|
||||
# or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
prompt = (
|
||||
@@ -27,7 +39,13 @@ def main():
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt="",
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
prompt2 = (
|
||||
"A majestic lion strides across the golden savanna, its powerful frame "
|
||||
@@ -36,7 +54,13 @@ def main():
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
|
||||
video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
negative_prompt="",
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig,
|
||||
SamplingConfig)
|
||||
from fastvideo.models.dits.hyworld.resolution_utils import get_resolution_from_image
|
||||
|
||||
# Default prompt from HY-WorldPlay run.sh
|
||||
@@ -31,33 +33,45 @@ def main():
|
||||
|
||||
# Initialize generator
|
||||
print("\nInitializing VideoGenerator for HYWorld...")
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
image_encoder_cpu_offload=True,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
image_encoder=True,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
# Generate video
|
||||
# The pose string is automatically converted to camera matrices by the pipeline
|
||||
print("\nGenerating video...")
|
||||
generator.generate_video(
|
||||
prompt=args.prompt,
|
||||
image_path=args.image,
|
||||
pose=args.pose, # Camera trajectory control
|
||||
output_path=args.output_path,
|
||||
save_video=True,
|
||||
negative_prompt="",
|
||||
num_frames=args.num_frames,
|
||||
fps=24,
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
seed=args.seed,
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt="",
|
||||
inputs=InputConfig(
|
||||
image_path=args.image,
|
||||
pose=args.pose, # Camera trajectory control
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
num_frames=args.num_frames,
|
||||
fps=24,
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=args.output_path,
|
||||
save_video=True,
|
||||
),
|
||||
))
|
||||
|
||||
print(f"\nVideo saved to: {args.output_path}")
|
||||
|
||||
|
||||
@@ -1,36 +1,43 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
InputConfig, OffloadConfig, OutputConfig,
|
||||
SamplingConfig)
|
||||
|
||||
OUTPUT_PATH = "video_samples_kandinsky5_i2v"
|
||||
|
||||
IMAGE_PATH = "assets/girl.png"
|
||||
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers",
|
||||
# "kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers",
|
||||
# "kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
prompt = (
|
||||
"A woman stands up and walks away"
|
||||
)
|
||||
_ = generator.generate_video(
|
||||
prompt,
|
||||
image_path=IMAGE_PATH,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
height=1024,
|
||||
width=1024,
|
||||
num_frames=121,
|
||||
)
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=IMAGE_PATH),
|
||||
sampling=SamplingConfig(height=1024, width=1024, num_frames=121),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,28 +1,41 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, SamplingConfig)
|
||||
|
||||
OUTPUT_PATH = "video_samples_kandinsky5_t2v"
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers"
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers"
|
||||
# "kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers"
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
_ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True,height=512, width=768, num_frames=121)
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(height=512, width=768, num_frames=121),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
prompt2 = (
|
||||
"A majestic lion strides across the golden savanna, its powerful frame "
|
||||
@@ -30,8 +43,13 @@ def main():
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
_ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=512, width=768, num_frames=121)
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
sampling=SamplingConfig(height=512, width=768, num_frames=121),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -1,24 +1,31 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
OUTPUT_PATH = "video_samples_lingbotworld"
|
||||
def main():
|
||||
# FastVideo will automatically use the optimal default arguments for the
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LingBot-World-Base-Cam-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LingBot-World-Base-Cam-Diffusers",
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
num_frames = 81
|
||||
@@ -33,15 +40,23 @@ def main():
|
||||
spatial_scale=8,
|
||||
)
|
||||
|
||||
generator.generate_video(
|
||||
prompt,
|
||||
image_path=image_path,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
num_frames=num_frames,
|
||||
height=480,
|
||||
width=832,
|
||||
c2ws_plucker_emb=c2ws_plucker_emb,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(
|
||||
image_path=image_path,
|
||||
c2ws_plucker_emb=c2ws_plucker_emb,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
num_frames=num_frames,
|
||||
height=480,
|
||||
width=832,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -19,6 +19,10 @@ import glob
|
||||
import os
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
InputConfig, OffloadConfig, OutputConfig, PipelineSelection, SamplingConfig,
|
||||
)
|
||||
|
||||
# Common prompts and settings matching the shell script examples
|
||||
PROMPT = (
|
||||
@@ -45,41 +49,50 @@ SEED = 42
|
||||
def basic_generation():
|
||||
"""
|
||||
Run basic LongCat I2V generation (50 steps at 480p).
|
||||
|
||||
|
||||
This uses the full 50-step denoising process for highest quality.
|
||||
"""
|
||||
print("=" * 60)
|
||||
print("LongCat I2V: Basic Generation (50 steps, 480p)")
|
||||
print("=" * 60)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-I2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=False,
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-I2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(experimental={"enable_bsa": False}),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
output_path = "outputs_video/longcat_i2v_basic"
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
image_path=IMAGE_PATH,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=480, # Square
|
||||
num_frames=93,
|
||||
num_inference_steps=50,
|
||||
fps=15,
|
||||
guidance_scale=4.0,
|
||||
seed=SEED,
|
||||
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
inputs=InputConfig(image_path=IMAGE_PATH),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=480, # Square
|
||||
num_frames=93,
|
||||
num_inference_steps=50,
|
||||
fps=15,
|
||||
guidance_scale=4.0,
|
||||
seed=SEED,
|
||||
),
|
||||
output=OutputConfig(output_path=output_path, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"\nBasic generation complete! Video saved to: {output_path}")
|
||||
generator.shutdown()
|
||||
|
||||
@@ -87,55 +100,70 @@ def basic_generation():
|
||||
def distill_refine_generation():
|
||||
"""
|
||||
Run LongCat I2V with distill+refine pipeline (16 steps + refinement to 768p).
|
||||
|
||||
|
||||
This uses the distilled LoRA for fast 480p generation (16 steps),
|
||||
then refines to 768p using the refinement LoRA with BSA enabled.
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("LongCat I2V: Distill + Refine Pipeline")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
# Stage 1: Distilled generation (16 steps at 480p)
|
||||
print("\n[Stage 1] Distilled generation (16 steps, 480p)")
|
||||
print("-" * 40)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-I2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=False,
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
|
||||
lora_nickname="distilled",
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-I2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
|
||||
),
|
||||
experimental={
|
||||
"enable_bsa": False,
|
||||
"lora_nickname": "distilled",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
distill_output_path = "outputs_video/longcat_i2v_distill"
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
image_path=IMAGE_PATH,
|
||||
output_path=distill_output_path,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=480, # Square
|
||||
num_frames=93,
|
||||
num_inference_steps=16,
|
||||
fps=15,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
inputs=InputConfig(image_path=IMAGE_PATH),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=480, # Square
|
||||
num_frames=93,
|
||||
num_inference_steps=16,
|
||||
fps=15,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
),
|
||||
output=OutputConfig(output_path=distill_output_path, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"Distilled generation complete! Video saved to: {distill_output_path}")
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
# Stage 2: Refinement (480p -> 768p)
|
||||
print("\n[Stage 2] Refinement (480p -> 768p with BSA)")
|
||||
print("-" * 40)
|
||||
|
||||
|
||||
# Find the actual saved video file from stage 1
|
||||
video_files = glob.glob(os.path.join(distill_output_path, "*.mp4"))
|
||||
if not video_files:
|
||||
@@ -143,46 +171,63 @@ def distill_refine_generation():
|
||||
# Use the most recently created video file
|
||||
distill_video_path = max(video_files, key=os.path.getmtime)
|
||||
print(f"Using stage 1 video: {distill_video_path}")
|
||||
|
||||
|
||||
# Create a new generator with refinement LoRA and BSA enabled
|
||||
# Note: Refinement uses the T2V model (not I2V) since it's upscaling the generated video
|
||||
# For BSA [4, 4, 8]: latent must be divisible by 8
|
||||
# 768x768: latent 48x48, 48%8=0 ✓
|
||||
refine_generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=True,
|
||||
bsa_sparsity=0.875,
|
||||
bsa_chunk_q=[4, 4, 4],
|
||||
bsa_chunk_k=[4, 4, 4],
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
|
||||
lora_nickname="refinement",
|
||||
refine_generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
|
||||
),
|
||||
experimental={
|
||||
"enable_bsa": True,
|
||||
"bsa_sparsity": 0.875,
|
||||
"bsa_chunk_q": [4, 4, 4],
|
||||
"bsa_chunk_k": [4, 4, 4],
|
||||
"lora_nickname": "refinement",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
refine_output_path = "outputs_video/longcat_i2v_refine_720p"
|
||||
|
||||
refine_generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output_path=refine_output_path,
|
||||
save_video=True,
|
||||
refine_from=distill_video_path,
|
||||
t_thresh=0.5,
|
||||
spatial_refine_only=False,
|
||||
num_cond_frames=0,
|
||||
height=720,
|
||||
width=720,
|
||||
num_inference_steps=50,
|
||||
fps=30,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
|
||||
refine_generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
inputs=InputConfig(refine_from=distill_video_path),
|
||||
sampling=SamplingConfig(
|
||||
height=720,
|
||||
width=720,
|
||||
num_inference_steps=50,
|
||||
fps=30,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
),
|
||||
output=OutputConfig(output_path=refine_output_path, save_video=True),
|
||||
extensions={
|
||||
"t_thresh": 0.5,
|
||||
"spatial_refine_only": False,
|
||||
"num_cond_frames": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"Refinement complete! Video saved to: {refine_output_path}")
|
||||
refine_generator.shutdown()
|
||||
|
||||
@@ -192,13 +237,13 @@ def main():
|
||||
print("\n" + "=" * 60)
|
||||
print("LongCat Image-to-Video Example")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
|
||||
# Run basic generation
|
||||
basic_generation()
|
||||
|
||||
|
||||
# Run distill+refine pipeline
|
||||
distill_refine_generation()
|
||||
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All generations complete!")
|
||||
print("=" * 60)
|
||||
|
||||
@@ -13,6 +13,10 @@ import glob
|
||||
import os
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
InputConfig, OffloadConfig, OutputConfig, PipelineSelection, SamplingConfig,
|
||||
)
|
||||
|
||||
# Common prompts and settings matching the shell script examples
|
||||
PROMPT = (
|
||||
@@ -38,40 +42,54 @@ SEED = 42
|
||||
def basic_generation():
|
||||
"""
|
||||
Run basic LongCat T2V generation (50 steps at 480p).
|
||||
|
||||
|
||||
This uses the full 50-step denoising process for highest quality.
|
||||
"""
|
||||
print("=" * 60)
|
||||
print("LongCat T2V: Basic Generation (50 steps, 480p)")
|
||||
print("=" * 60)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=False,
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
experimental={"enable_bsa": False},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
output_path = "outputs_video/longcat_t2v_basic"
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=50,
|
||||
fps=15,
|
||||
guidance_scale=4.0,
|
||||
seed=SEED,
|
||||
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=50,
|
||||
fps=15,
|
||||
guidance_scale=4.0,
|
||||
seed=SEED,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"\nBasic generation complete! Video saved to: {output_path}")
|
||||
generator.shutdown()
|
||||
|
||||
@@ -79,54 +97,72 @@ def basic_generation():
|
||||
def distill_refine_generation():
|
||||
"""
|
||||
Run LongCat T2V with distill+refine pipeline (16 steps + refinement to 720p).
|
||||
|
||||
|
||||
This uses the distilled LoRA for fast 480p generation (16 steps),
|
||||
then refines to 720p using the refinement LoRA with BSA enabled.
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("LongCat T2V: Distill + Refine Pipeline")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
# Stage 1: Distilled generation (16 steps at 480p)
|
||||
print("\n[Stage 1] Distilled generation (16 steps, 480p)")
|
||||
print("-" * 40)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=False,
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
|
||||
lora_nickname="distilled",
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
|
||||
),
|
||||
experimental={
|
||||
"enable_bsa": False,
|
||||
"lora_nickname": "distilled",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
distill_output_path = "outputs_video/longcat_t2v_distill"
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output_path=distill_output_path,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=16,
|
||||
fps=15,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path=distill_output_path,
|
||||
save_video=True,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=16,
|
||||
fps=15,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"Distilled generation complete! Video saved to: {distill_output_path}")
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
# Stage 2: Refinement (480p -> 720p)
|
||||
print("\n[Stage 2] Refinement (480p -> 720p with BSA)")
|
||||
print("-" * 40)
|
||||
|
||||
|
||||
# Find the actual saved video file from stage 1
|
||||
video_files = glob.glob(os.path.join(distill_output_path, "*.mp4"))
|
||||
if not video_files:
|
||||
@@ -134,43 +170,65 @@ def distill_refine_generation():
|
||||
# Use the most recently created video file
|
||||
distill_video_path = max(video_files, key=os.path.getmtime)
|
||||
print(f"Using stage 1 video: {distill_video_path}")
|
||||
|
||||
|
||||
# Create a new generator with refinement LoRA and BSA enabled
|
||||
refine_generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=True,
|
||||
bsa_sparsity=0.875,
|
||||
bsa_chunk_q=[4, 4, 8],
|
||||
bsa_chunk_k=[4, 4, 8],
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
|
||||
lora_nickname="refinement",
|
||||
refine_generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
|
||||
),
|
||||
experimental={
|
||||
"enable_bsa": True,
|
||||
"bsa_sparsity": 0.875,
|
||||
"bsa_chunk_q": [4, 4, 8],
|
||||
"bsa_chunk_k": [4, 4, 8],
|
||||
"lora_nickname": "refinement",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
refine_output_path = "outputs_video/longcat_t2v_refine_720p"
|
||||
|
||||
refine_generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output_path=refine_output_path,
|
||||
save_video=True,
|
||||
refine_from=distill_video_path,
|
||||
t_thresh=0.5,
|
||||
spatial_refine_only=False,
|
||||
num_cond_frames=0,
|
||||
height=720,
|
||||
width=1280,
|
||||
num_inference_steps=50,
|
||||
fps=30,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
|
||||
refine_generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path=refine_output_path,
|
||||
save_video=True,
|
||||
),
|
||||
inputs=InputConfig(
|
||||
refine_from=distill_video_path,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
height=720,
|
||||
width=1280,
|
||||
num_inference_steps=50,
|
||||
fps=30,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
),
|
||||
extensions={
|
||||
"t_thresh": 0.5,
|
||||
"spatial_refine_only": False,
|
||||
"num_cond_frames": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"Refinement complete! Video saved to: {refine_output_path}")
|
||||
refine_generator.shutdown()
|
||||
|
||||
@@ -180,13 +238,13 @@ def main():
|
||||
print("\n" + "=" * 60)
|
||||
print("LongCat Text-to-Video Example")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
|
||||
# Run basic generation
|
||||
basic_generation()
|
||||
|
||||
|
||||
# Run distill+refine pipeline
|
||||
distill_refine_generation()
|
||||
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All generations complete!")
|
||||
print("=" * 60)
|
||||
@@ -194,5 +252,3 @@ def main():
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
|
||||
@@ -19,6 +19,10 @@ import glob
|
||||
import os
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig,
|
||||
PipelineSelection, SamplingConfig,
|
||||
)
|
||||
|
||||
# Common prompts and settings matching the shell script examples
|
||||
PROMPT = (
|
||||
@@ -63,35 +67,49 @@ def basic_generation():
|
||||
"Please provide a valid video path."
|
||||
)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-VC-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-VC-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
experimental={"enable_bsa": False},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
output_path = "outputs_video/longcat_vc_basic"
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
video_path=VIDEO_PATH,
|
||||
num_cond_frames=NUM_COND_FRAMES,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=50,
|
||||
fps=15,
|
||||
guidance_scale=4.0,
|
||||
seed=SEED,
|
||||
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
inputs=InputConfig(video_path=VIDEO_PATH),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=50,
|
||||
fps=15,
|
||||
guidance_scale=4.0,
|
||||
seed=SEED,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
),
|
||||
extensions={"num_cond_frames": NUM_COND_FRAMES},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"\nBasic generation complete! Video saved to: {output_path}")
|
||||
generator.shutdown()
|
||||
|
||||
@@ -118,37 +136,55 @@ def distill_refine_generation():
|
||||
print("\n[Stage 1] Distilled generation (16 steps, 480p)")
|
||||
print("-" * 40)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-VC-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=False,
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
|
||||
lora_nickname="distilled",
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-VC-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
|
||||
),
|
||||
experimental={
|
||||
"enable_bsa": False,
|
||||
"lora_nickname": "distilled",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
distill_output_path = "outputs_video/longcat_vc_distill"
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
video_path=VIDEO_PATH,
|
||||
num_cond_frames=NUM_COND_FRAMES,
|
||||
output_path=distill_output_path,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=16,
|
||||
fps=15,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
inputs=InputConfig(video_path=VIDEO_PATH),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=93,
|
||||
num_inference_steps=16,
|
||||
fps=15,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=distill_output_path,
|
||||
save_video=True,
|
||||
),
|
||||
extensions={"num_cond_frames": NUM_COND_FRAMES},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"Distilled generation complete! Video saved to: {distill_output_path}")
|
||||
generator.shutdown()
|
||||
|
||||
@@ -166,41 +202,61 @@ def distill_refine_generation():
|
||||
|
||||
# Create a new generator with refinement LoRA and BSA enabled
|
||||
# Note: Refinement uses the T2V model (not VC) since it's upscaling the generated video
|
||||
refine_generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=False,
|
||||
enable_bsa=True,
|
||||
bsa_sparsity=0.875,
|
||||
bsa_chunk_q=[4, 4, 8],
|
||||
bsa_chunk_k=[4, 4, 8],
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
|
||||
lora_nickname="refinement",
|
||||
refine_generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LongCat-Video-T2V-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
|
||||
),
|
||||
experimental={
|
||||
"enable_bsa": True,
|
||||
"bsa_sparsity": 0.875,
|
||||
"bsa_chunk_q": [4, 4, 8],
|
||||
"bsa_chunk_k": [4, 4, 8],
|
||||
"lora_nickname": "refinement",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
refine_output_path = "outputs_video/longcat_vc_refine_720p"
|
||||
|
||||
refine_generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
output_path=refine_output_path,
|
||||
save_video=True,
|
||||
refine_from=distill_video_path,
|
||||
t_thresh=0.5,
|
||||
spatial_refine_only=False,
|
||||
num_cond_frames=0, # For refinement, no conditioning frames
|
||||
height=720,
|
||||
width=1280,
|
||||
num_inference_steps=50,
|
||||
fps=30,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
|
||||
refine_generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
negative_prompt=NEGATIVE_PROMPT,
|
||||
inputs=InputConfig(refine_from=distill_video_path),
|
||||
sampling=SamplingConfig(
|
||||
height=720,
|
||||
width=1280,
|
||||
num_inference_steps=50,
|
||||
fps=30,
|
||||
guidance_scale=1.0,
|
||||
seed=SEED,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=refine_output_path,
|
||||
save_video=True,
|
||||
),
|
||||
extensions={
|
||||
"t_thresh": 0.5,
|
||||
"spatial_refine_only": False,
|
||||
"num_cond_frames": 0, # For refinement, no conditioning frames
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
print(f"Refinement complete! Video saved to: {refine_output_path}")
|
||||
refine_generator.shutdown()
|
||||
|
||||
|
||||
@@ -1,4 +1,11 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
|
||||
PROMPT = (
|
||||
@@ -18,22 +25,32 @@ PROMPT = (
|
||||
|
||||
def main() -> None:
|
||||
# Uses FastVideo default sampling settings for LTX2 base.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Davids048/LTX2-Base-Diffusers",
|
||||
num_gpus=1,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Davids048/LTX2-Base-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.1.mp4"
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
num_frames=121,
|
||||
height=1088,
|
||||
width=1920,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
num_frames=121,
|
||||
height=1088,
|
||||
width=1920,
|
||||
),
|
||||
)
|
||||
)
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -49,6 +49,11 @@ from pathlib import Path
|
||||
import torch._inductor.config as _inductor
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig, ComponentConfig, EngineConfig, GenerationRequest,
|
||||
GeneratorConfig, OffloadConfig, OutputConfig, PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.utils import maybe_download_model
|
||||
|
||||
@@ -86,9 +91,9 @@ PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
|
||||
|
||||
# Per-stage timing helpers --------------------------------------------------
|
||||
|
||||
def _print_stage_breakdown(result: dict, label: str) -> float | None:
|
||||
def _print_stage_breakdown(result, label: str) -> float | None:
|
||||
"""Print stage execution times and return the sum, or None if missing."""
|
||||
logging_info = result.get("logging_info")
|
||||
logging_info = result.logging_info
|
||||
stages = getattr(logging_info, "stages", None) if logging_info else None
|
||||
if not stages:
|
||||
print(f" [{label}] stage breakdown unavailable")
|
||||
@@ -104,11 +109,11 @@ def _print_stage_breakdown(result: dict, label: str) -> float | None:
|
||||
|
||||
|
||||
def _collect_stage_times(
|
||||
result: dict,
|
||||
result,
|
||||
stage_times: dict[str, list[float]],
|
||||
stage_order: OrderedDict[str, None],
|
||||
) -> None:
|
||||
logging_info = result.get("logging_info")
|
||||
logging_info = result.logging_info
|
||||
stages = getattr(logging_info, "stages", None) if logging_info else None
|
||||
if not stages:
|
||||
return
|
||||
@@ -169,34 +174,54 @@ def main() -> None:
|
||||
pipeline_config = PipelineConfig.from_pretrained(model_root)
|
||||
pipeline_config.dit_config.quant_config = None
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_root,
|
||||
num_gpus=1,
|
||||
# LTX-2.3 distilled uses the two-stage refine pipeline; the refine
|
||||
# LoRA is intentionally empty for the distilled student.
|
||||
ltx2_refine_enabled=True,
|
||||
ltx2_refine_upsampler_path=str(refine_upsampler_path),
|
||||
ltx2_refine_lora_path="",
|
||||
ltx2_refine_num_inference_steps=3,
|
||||
ltx2_refine_guidance_scale=1.0,
|
||||
ltx2_refine_add_noise=True,
|
||||
pipeline_config=pipeline_config,
|
||||
enable_torch_compile=True,
|
||||
enable_torch_compile_text_encoder=True,
|
||||
# Compile the VAE codec submodules (encoder / decoder) too. The
|
||||
# `LTX2CausalVideoAutoencoder` declares `_compile_conditions` so
|
||||
# `_compile_with_conditions` targets just those submodules and
|
||||
# leaves the surrounding tiling control flow eager — needed for
|
||||
# fullgraph + dynamic=False to succeed. VAE eager decode is
|
||||
# ~1.0s; compiling it brings the stage to ~0.3s.
|
||||
enable_torch_compile_vae=True,
|
||||
torch_compile_kwargs=torch_compile_kwargs,
|
||||
torch_compile_kwargs_vae=torch_compile_kwargs,
|
||||
# Keep everything resident — no CPU offload for serving-style runs.
|
||||
dit_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
ltx2_vae_tiling=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_root,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
compile=CompileConfig(
|
||||
enabled=True,
|
||||
text_encoder_enabled=True,
|
||||
# Compile the VAE codec submodules (encoder / decoder)
|
||||
# too. The `LTX2CausalVideoAutoencoder` declares
|
||||
# `_compile_conditions` so `_compile_with_conditions`
|
||||
# targets just those submodules and leaves the
|
||||
# surrounding tiling control flow eager — needed for
|
||||
# fullgraph + dynamic=False to succeed. VAE eager decode
|
||||
# is ~1.0s; compiling it brings the stage to ~0.3s.
|
||||
vae_enabled=True,
|
||||
backend=torch_compile_kwargs["backend"],
|
||||
fullgraph=torch_compile_kwargs["fullgraph"],
|
||||
mode=torch_compile_kwargs["mode"],
|
||||
dynamic=torch_compile_kwargs["dynamic"],
|
||||
vae_kwargs=torch_compile_kwargs,
|
||||
),
|
||||
# Keep everything resident — no CPU offload for serving runs.
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
text_encoder=False,
|
||||
vae=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
vae_tiling=False,
|
||||
# LTX-2.3 distilled uses the two-stage refine pipeline; the
|
||||
# refine LoRA is intentionally empty for the distilled
|
||||
# student.
|
||||
components=ComponentConfig(
|
||||
upsampler_weights=str(refine_upsampler_path),
|
||||
),
|
||||
preset_overrides={
|
||||
"refine": {
|
||||
"enabled": True,
|
||||
"num_inference_steps": 3,
|
||||
"guidance_scale": 1.0,
|
||||
"add_noise": True,
|
||||
}
|
||||
},
|
||||
experimental={"pipeline_config": pipeline_config},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
common_kwargs = dict(
|
||||
@@ -206,12 +231,15 @@ def main() -> None:
|
||||
height=1280, width=832, # portrait runway aspect
|
||||
num_frames=121, fps=24, # ~5s clip
|
||||
num_inference_steps=8, # distilled denoise steps
|
||||
# i2v: anchor the input image at frame 0 with full strength.
|
||||
# `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already
|
||||
# JPEG conditioning image.
|
||||
)
|
||||
|
||||
# i2v: anchor the input image at frame 0 with full strength.
|
||||
# `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already
|
||||
# JPEG conditioning image. These are model-specific knobs routed through
|
||||
# the request extensions escape hatch.
|
||||
common_extensions = dict(
|
||||
ltx2_images=[(I2V_IMAGE, 0, 1.0)],
|
||||
ltx2_image_crf=0.0,
|
||||
save_video=True,
|
||||
)
|
||||
|
||||
warmup_runs = 2
|
||||
@@ -227,10 +255,25 @@ def main() -> None:
|
||||
for w in range(warmup_runs):
|
||||
t0 = time.perf_counter()
|
||||
print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
|
||||
generator.generate_video(
|
||||
output_path=str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4"),
|
||||
seed=7,
|
||||
**common_kwargs,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=common_kwargs["prompt"],
|
||||
negative_prompt=common_kwargs["negative_prompt"],
|
||||
sampling=SamplingConfig(
|
||||
guidance_scale=common_kwargs["guidance_scale"],
|
||||
height=common_kwargs["height"],
|
||||
width=common_kwargs["width"],
|
||||
num_frames=common_kwargs["num_frames"],
|
||||
fps=common_kwargs["fps"],
|
||||
num_inference_steps=common_kwargs["num_inference_steps"],
|
||||
seed=7,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4"),
|
||||
save_video=True,
|
||||
),
|
||||
extensions=common_extensions,
|
||||
)
|
||||
)
|
||||
dt = time.perf_counter() - t0
|
||||
warmup_secs.append(dt)
|
||||
@@ -245,19 +288,31 @@ def main() -> None:
|
||||
out_path = OUTPUT_DIR / f"output_ltx2_3_distilled_i2v_run_{m + 1}.mp4"
|
||||
print(f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}")
|
||||
t0 = time.perf_counter()
|
||||
result = generator.generate_video(
|
||||
output_path=str(out_path),
|
||||
seed=2002 + m,
|
||||
**common_kwargs,
|
||||
result = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=common_kwargs["prompt"],
|
||||
negative_prompt=common_kwargs["negative_prompt"],
|
||||
sampling=SamplingConfig(
|
||||
guidance_scale=common_kwargs["guidance_scale"],
|
||||
height=common_kwargs["height"],
|
||||
width=common_kwargs["width"],
|
||||
num_frames=common_kwargs["num_frames"],
|
||||
fps=common_kwargs["fps"],
|
||||
num_inference_steps=common_kwargs["num_inference_steps"],
|
||||
seed=2002 + m,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(out_path),
|
||||
save_video=True,
|
||||
),
|
||||
extensions=common_extensions,
|
||||
)
|
||||
)
|
||||
wall = time.perf_counter() - t0
|
||||
e2e = (
|
||||
result.get("e2e_latency")
|
||||
if isinstance(result, dict) else None
|
||||
) or wall
|
||||
e2e = (result.extra.get("e2e_latency") if result is not None else None) or wall
|
||||
measured_secs.append(e2e)
|
||||
print(f"[measured {m + 1}/{measured_runs}] e2e={e2e:.2f}s wall={wall:.2f}s")
|
||||
if isinstance(result, dict):
|
||||
if result is not None:
|
||||
_print_stage_breakdown(result, f"measured {m + 1}")
|
||||
_collect_stage_times(result, stage_times, stage_order)
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import EngineConfig, GenerationRequest, GeneratorConfig, OutputConfig
|
||||
|
||||
PROMPT = (
|
||||
"A warm sunny backyard. The camera starts in a tight cinematic close-up "
|
||||
@@ -17,16 +18,19 @@ import os
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/LTX2-Distilled-Diffusers",
|
||||
num_gpus=4,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/LTX2-Distilled-Diffusers",
|
||||
engine=EngineConfig(num_gpus=4),
|
||||
)
|
||||
)
|
||||
|
||||
output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4"
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(output_path=output_path, save_video=True),
|
||||
)
|
||||
)
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
@@ -8,6 +8,11 @@ from pathlib import Path
|
||||
import torch
|
||||
import torch._inductor.config
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig, ComponentConfig, EngineConfig, GenerationRequest,
|
||||
GenerationResult, GeneratorConfig, OffloadConfig, OutputConfig,
|
||||
PipelineSelection, SamplingConfig,
|
||||
)
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
|
||||
from fastvideo.utils import maybe_download_model
|
||||
@@ -45,11 +50,11 @@ def load_validation_entries(path: Path) -> list[dict]:
|
||||
|
||||
|
||||
def print_stage_breakdown(
|
||||
result: dict,
|
||||
result: GenerationResult,
|
||||
run_idx: int,
|
||||
num_runs: int,
|
||||
) -> float | None:
|
||||
logging_info = result.get("logging_info")
|
||||
logging_info = result.logging_info
|
||||
if logging_info is None:
|
||||
print(f"[{run_idx}/{num_runs}] Stage breakdown unavailable: no logging_info")
|
||||
return None
|
||||
@@ -70,9 +75,9 @@ def print_stage_breakdown(
|
||||
|
||||
|
||||
def extract_sr_forward_latency(
|
||||
result: dict,
|
||||
result: GenerationResult,
|
||||
) -> tuple[float | None, list[tuple[str, float]], list[str]]:
|
||||
logging_info = result.get("logging_info")
|
||||
logging_info = result.logging_info
|
||||
if logging_info is None:
|
||||
return None, [], []
|
||||
|
||||
@@ -106,11 +111,11 @@ def extract_sr_forward_latency(
|
||||
|
||||
|
||||
def collect_stage_times(
|
||||
result: dict,
|
||||
result: GenerationResult,
|
||||
stage_times: dict[str, list[float]],
|
||||
stage_order: OrderedDict[str, None],
|
||||
) -> None:
|
||||
logging_info = result.get("logging_info")
|
||||
logging_info = result.logging_info
|
||||
if logging_info is None:
|
||||
return
|
||||
stages = getattr(logging_info, "stages", None)
|
||||
@@ -202,26 +207,45 @@ def main() -> None:
|
||||
"dynamic": False,
|
||||
}
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_root,
|
||||
num_gpus=1,
|
||||
ltx2_refine_enabled=True,
|
||||
ltx2_refine_upsampler_path=str(refine_upsampler_path),
|
||||
refine_lora_path="", # keep refine LoRA disabled in this repo's typed adapter
|
||||
ltx2_refine_lora_path="", # keep refine LoRA disabled for distilled model
|
||||
ltx2_refine_num_inference_steps=2,
|
||||
ltx2_refine_guidance_scale=1.0,
|
||||
ltx2_refine_add_noise=True,
|
||||
pipeline_config=pipeline_config,
|
||||
enable_torch_compile=True,
|
||||
enable_torch_compile_text_encoder=True,
|
||||
enable_torch_compile_vae=True,
|
||||
torch_compile_kwargs=torch_compile_kwargs,
|
||||
torch_compile_kwargs_vae=torch_compile_kwargs,
|
||||
dit_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
ltx2_vae_tiling=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_root,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
text_encoder=False,
|
||||
vae=False,
|
||||
),
|
||||
compile=CompileConfig(
|
||||
enabled=True,
|
||||
text_encoder_enabled=True,
|
||||
vae_enabled=True,
|
||||
backend="inductor",
|
||||
fullgraph=True,
|
||||
dynamic=False,
|
||||
vae_kwargs=torch_compile_kwargs,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
vae_tiling=False,
|
||||
components=ComponentConfig(
|
||||
upsampler_weights=str(refine_upsampler_path),
|
||||
),
|
||||
preset_overrides={
|
||||
"refine": {
|
||||
"enabled": True,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 1.0,
|
||||
"add_noise": True,
|
||||
}
|
||||
},
|
||||
experimental={
|
||||
"refine_lora_path": "", # keep refine LoRA disabled in this repo's typed adapter
|
||||
"pipeline_config": pipeline_config,
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
run_times: list[float] = []
|
||||
@@ -243,25 +267,31 @@ def main() -> None:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start = time.perf_counter()
|
||||
result = generator.generate_video(
|
||||
prompt=prompt,
|
||||
output_path=str(output_path),
|
||||
fps=24,
|
||||
seed=10,
|
||||
save_video=True,
|
||||
guidance_scale=1.0,
|
||||
height=benchmark_entry.get("height", 1088),
|
||||
width=benchmark_entry.get("width", 1920),
|
||||
num_frames=121,
|
||||
num_inference_steps=5,
|
||||
# image_path="examples/inference/basic/prompt1.png",
|
||||
# ltx2_image_crf=0.0
|
||||
result = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(
|
||||
fps=24,
|
||||
seed=10,
|
||||
guidance_scale=1.0,
|
||||
height=benchmark_entry.get("height", 1088),
|
||||
width=benchmark_entry.get("width", 1920),
|
||||
num_frames=121,
|
||||
num_inference_steps=5,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output_path),
|
||||
save_video=True,
|
||||
),
|
||||
# inputs=InputConfig(image_path="examples/inference/basic/prompt1.png"),
|
||||
# extensions={"ltx2_image_crf": 0.0},
|
||||
)
|
||||
)
|
||||
if os.environ.get("FASTVIDEO_STAGE_LOGGING") == "0":
|
||||
torch.cuda.synchronize()
|
||||
|
||||
elapsed = result.get("generation_time") if isinstance(result, dict) else None
|
||||
e2e_elapsed = result.get("e2e_latency") if isinstance(result, dict) else None
|
||||
elapsed = result.generation_time if isinstance(result, GenerationResult) else None
|
||||
e2e_elapsed = result.extra.get("e2e_latency") if isinstance(result, GenerationResult) else None
|
||||
if elapsed is None:
|
||||
elapsed = time.perf_counter() - start
|
||||
if e2e_elapsed is None:
|
||||
@@ -272,7 +302,7 @@ def main() -> None:
|
||||
print(f"[{i + 1}/{num_runs}] Generation time: {elapsed:.2f}s")
|
||||
print(f"[{i + 1}/{num_runs}] End-to-end latency: {e2e_elapsed:.2f}s")
|
||||
|
||||
if isinstance(result, dict):
|
||||
if isinstance(result, GenerationResult):
|
||||
stage_sum = print_stage_breakdown(result, i + 1, num_runs)
|
||||
if stage_sum is not None:
|
||||
non_stage_overhead = e2e_elapsed - stage_sum
|
||||
|
||||
@@ -1,18 +1,27 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig,
|
||||
OffloadConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_lucy_edit"
|
||||
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"decart-ai/Lucy-Edit-Dev",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="decart-ai/Lucy-Edit-Dev",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
prompt = ("Change the apron and blouse to a classic clown costume: satin "
|
||||
"polka-dot jumpsuit in bright primary colors, ruffled white collar, "
|
||||
@@ -20,18 +29,20 @@ def main():
|
||||
"foam nose; soft window light from left, eye-level medium shot.")
|
||||
video_path = "https://d2drjpuinn46lb.cloudfront.net/painter_original_edit.mp4"
|
||||
|
||||
generator.generate_video(
|
||||
prompt,
|
||||
negative_prompt="",
|
||||
video_path=video_path,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=81,
|
||||
fps=24,
|
||||
guidance_scale=5.0,
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt="",
|
||||
inputs=InputConfig(video_path=video_path),
|
||||
sampling=SamplingConfig(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=81,
|
||||
fps=24,
|
||||
guidance_scale=5.0,
|
||||
),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig,
|
||||
SamplingConfig)
|
||||
from fastvideo.models.dits.matrixgame2.utils import create_action_presets
|
||||
|
||||
import torch
|
||||
@@ -38,35 +40,48 @@ def main():
|
||||
# attempt to identify the optimal arguments.
|
||||
config = VARIANT_CONFIG[MODEL_VARIANT]
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
config["model_path"],
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=config["model_path"],
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
num_frames = 597
|
||||
actions = create_action_presets(num_frames, keyboard_dim=config["keyboard_dim"])
|
||||
grid_sizes = torch.tensor([150, 44, 80])
|
||||
|
||||
generator.generate_video(
|
||||
prompt="",
|
||||
image_path=config["image_url"],
|
||||
mouse_cond=actions["mouse"].unsqueeze(0),
|
||||
keyboard_cond=actions["keyboard"].unsqueeze(0),
|
||||
grid_sizes=grid_sizes,
|
||||
num_frames=num_frames,
|
||||
height=352,
|
||||
width=640,
|
||||
num_inference_steps=50,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt="",
|
||||
inputs=InputConfig(
|
||||
image_path=config["image_url"],
|
||||
mouse_cond=actions["mouse"].unsqueeze(0),
|
||||
keyboard_cond=actions["keyboard"].unsqueeze(0),
|
||||
grid_sizes=grid_sizes,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
num_frames=num_frames,
|
||||
height=352,
|
||||
width=640,
|
||||
num_inference_steps=50,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator
|
||||
from fastvideo.models.dits.matrixgame2.utils import get_current_action_async, expand_action_to_frames
|
||||
from fastvideo.api import EngineConfig, GeneratorConfig, OffloadConfig
|
||||
|
||||
import torch
|
||||
import asyncio
|
||||
@@ -42,17 +43,23 @@ async def main():
|
||||
# attempt to identify the optimal arguments.
|
||||
config = VARIANT_CONFIG[MODEL_VARIANT]
|
||||
|
||||
generator = StreamingVideoGenerator.from_pretrained(
|
||||
config["model_path"],
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = StreamingVideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=config["model_path"],
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
max_blocks = 50
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
MODEL_PATH = "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers"
|
||||
IMAGE_URL = "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-3/demo_images/001/image.png"
|
||||
@@ -7,28 +10,38 @@ OUTPUT_PATH = "video_samples_matrixgame3"
|
||||
|
||||
|
||||
def main():
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
MODEL_PATH,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=MODEL_PATH,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
image_path=IMAGE_URL,
|
||||
height=720,
|
||||
width=1280,
|
||||
num_frames=57,
|
||||
num_inference_steps=3,
|
||||
guidance_scale=1.0,
|
||||
seed=42,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
inputs=InputConfig(image_path=IMAGE_URL),
|
||||
sampling=SamplingConfig(
|
||||
height=720,
|
||||
width=1280,
|
||||
num_frames=57,
|
||||
num_inference_steps=3,
|
||||
guidance_scale=1.0,
|
||||
seed=42,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,40 +1,56 @@
|
||||
from fastvideo import VideoGenerator, PipelineConfig
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
def main():
|
||||
config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
config.text_encoder_precisions = ["fp16"]
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
pipeline_config=config,
|
||||
use_fsdp_inference=False, # Disable FSDP for MPS
|
||||
dit_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
disable_autocast=False,
|
||||
num_gpus=1,
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # Disable FSDP for MPS
|
||||
disable_autocast=False,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
experimental={"pipeline_config": config},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Create sampling parameters with reduced number of frames
|
||||
sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
sampling_param.num_frames = 25 # Reduce from default 81 to 25 frames bc we have to use the SDPA attn backend for mps
|
||||
sampling_param.height = 256
|
||||
sampling_param.width = 256
|
||||
# Reduce from default 81 to 25 frames bc we have to use the SDPA attn backend for mps
|
||||
sampling = SamplingConfig(
|
||||
num_frames=25,
|
||||
height=256,
|
||||
width=256,
|
||||
)
|
||||
|
||||
prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
|
||||
|
||||
video = generator.generate_video(prompt, sampling_param=sampling_param)
|
||||
|
||||
video = generator.generate(GenerationRequest(prompt=prompt, sampling=sampling))
|
||||
|
||||
prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame "
|
||||
"glistening under the warm afternoon sun. The tall grass ripples gently in "
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
|
||||
video2 = generator.generate_video(prompt2, sampling_param=sampling_param)
|
||||
|
||||
video2 = generator.generate(GenerationRequest(prompt=prompt2, sampling=sampling))
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig,
|
||||
OutputConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples"
|
||||
def main():
|
||||
@@ -8,17 +10,23 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=2,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
distributed_executor_backend="ray",
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=2,
|
||||
use_fsdp_inference=True,
|
||||
execution_backend="ray",
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Generate videos with the same simple API, regardless of GPU count
|
||||
@@ -27,7 +35,8 @@ def main():
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
|
||||
video = generator.generate(
|
||||
GenerationRequest(prompt=prompt, output=OutputConfig(output_path=OUTPUT_PATH, save_video=True)))
|
||||
|
||||
# Generate another video with a different prompt, without reloading the
|
||||
# model!
|
||||
@@ -37,7 +46,8 @@ def main():
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(prompt=prompt2, output=OutputConfig(output_path=OUTPUT_PATH, save_video=True)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -85,24 +85,35 @@ def main() -> None:
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig, OutputConfig,
|
||||
ParallelismConfig, PipelineSelection, SamplingConfig,
|
||||
)
|
||||
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
init_kwargs = {
|
||||
"num_gpus": args.num_gpus,
|
||||
"workload_type": "t2i",
|
||||
"sp_size": 1,
|
||||
"tp_size": 1,
|
||||
"dit_cpu_offload": False,
|
||||
"dit_layerwise_offload": False,
|
||||
"text_encoder_cpu_offload": False,
|
||||
"vae_cpu_offload": False,
|
||||
"image_encoder_cpu_offload": False,
|
||||
"pin_cpu_memory": False,
|
||||
"use_fsdp_inference": False,
|
||||
}
|
||||
|
||||
generator = VideoGenerator.from_pretrained(model_path=args.model_path, **init_kwargs)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=False,
|
||||
parallelism=ParallelismConfig(
|
||||
sp_size=1,
|
||||
tp_size=1,
|
||||
),
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=False,
|
||||
text_encoder=False,
|
||||
vae=False,
|
||||
image_encoder=False,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(workload_type="t2i"),
|
||||
)
|
||||
)
|
||||
try:
|
||||
for i, prompt in enumerate(prompts):
|
||||
seed = args.seed + i
|
||||
@@ -113,20 +124,25 @@ def main() -> None:
|
||||
output_path = os.path.join(args.out_dir, f"{filename_base}.png")
|
||||
print(f"[sd35] prompt_idx={i} seed={seed} output_path={output_path}")
|
||||
|
||||
generation_kwargs = {
|
||||
"output_path": output_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": 1,
|
||||
"fps": 1,
|
||||
"num_inference_steps": args.steps,
|
||||
"guidance_scale": args.guidance,
|
||||
"seed": seed,
|
||||
"negative_prompt": args.negative,
|
||||
"save_video": True,
|
||||
}
|
||||
|
||||
generator.generate_video(prompt, **generation_kwargs)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=args.negative,
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=1,
|
||||
fps=1,
|
||||
num_inference_steps=args.steps,
|
||||
guidance_scale=args.guidance,
|
||||
seed=seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
print(f"[sd35] done. outputs written to: {args.out_dir}")
|
||||
finally:
|
||||
|
||||
@@ -1,6 +1,14 @@
|
||||
import os
|
||||
import time
|
||||
from fastvideo import VideoGenerator, SamplingParam
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_causal"
|
||||
def main():
|
||||
@@ -9,23 +17,33 @@ def main():
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=model_name,
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
text_encoder_cpu_offload=False,
|
||||
dit_cpu_offload=False,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
text_encoder=False,
|
||||
dit=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
sampling_param = SamplingParam.from_pretrained(model_name)
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
|
||||
request = GenerationRequest(
|
||||
prompt=prompt,
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
video = generator.generate(request)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,8 +1,17 @@
|
||||
# NOTE: This is still a work in progress, and the checkpoints are not released yet.
|
||||
|
||||
from fastvideo import VideoGenerator, SamplingParam
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
import json
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_i2v"
|
||||
def main():
|
||||
@@ -10,26 +19,37 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
dit_precision="fp32",
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
experimental={
|
||||
"dit_precision": "fp32",
|
||||
"dmd_denoising_steps": [1000, 850, 700, 550, 350, 275, 200, 125],
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
sampling_param = SamplingParam.from_pretrained("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers")
|
||||
sampling_param.num_frames = 81
|
||||
sampling_param.width = 832
|
||||
sampling_param.height = 480
|
||||
sampling_param.seed = 1000
|
||||
sampling = SamplingConfig(
|
||||
num_frames=81,
|
||||
width=832,
|
||||
height=480,
|
||||
seed=1000,
|
||||
)
|
||||
|
||||
with open("assets/prompts/mixkit_i2v.jsonl", "r") as f:
|
||||
prompt_image_pairs = json.load(f)
|
||||
@@ -37,7 +57,14 @@ def main():
|
||||
for prompt_image_pair in prompt_image_pairs:
|
||||
prompt = prompt_image_pair["prompt"]
|
||||
image_path = prompt_image_pair["image_path"]
|
||||
_ = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
sampling=sampling,
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
# NOTE: This is still a work in progress, and the checkpoints are not released yet.
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, PipelineSelection, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_t2v"
|
||||
def main():
|
||||
@@ -10,34 +12,49 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/",
|
||||
init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/",
|
||||
num_frame_per_block=7,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
transformer_weights="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/",
|
||||
transformer_2_weights="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/",
|
||||
),
|
||||
experimental={
|
||||
"dmd_denoising_steps": [1000, 850, 700, 550, 350, 275, 200, 125],
|
||||
"num_frame_per_block": 7,
|
||||
},
|
||||
),
|
||||
# image_encoder_cpu_offload=False,
|
||||
)
|
||||
)
|
||||
|
||||
# sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
# sampling_param.num_frames = 45
|
||||
# sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
|
||||
# Generate videos with the same simple API, regardless of GPU count
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
_ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, num_frames=81)
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_frames=81),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -51,25 +51,31 @@ Prerequisites:
|
||||
uv pip install k_diffusion einops_exts alias_free_torch torchsde
|
||||
"""
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OutputConfig)
|
||||
|
||||
PROMPT = "Lo-fi hip hop instrumental with vinyl crackle and gentle piano."
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/stable-audio-open-1.0-Diffusers",
|
||||
num_gpus=1,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/stable-audio-open-1.0-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
))
|
||||
output_path = "outputs_audio/stable_audio_basic/output_stable_audio.wav"
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
# 6-second clip; the model max is ~47.5s.
|
||||
audio_end_in_s=6.0,
|
||||
# The registered preset gives 100 steps + CFG=7.0 by default;
|
||||
# override num_inference_steps / guidance_scale here for QA.
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
),
|
||||
# 6-second clip; the model max is ~47.5s.
|
||||
extensions={"audio_end_in_s": 6.0},
|
||||
# The registered preset gives 100 steps + CFG=7.0 by default;
|
||||
# override num_inference_steps / guidance_scale here for QA.
|
||||
))
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
|
||||
@@ -48,6 +48,12 @@ Picking `init_audio_strength` (0.0 to 1.0):
|
||||
Prerequisites: same as `basic_stable_audio.py`.
|
||||
"""
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
)
|
||||
|
||||
PROMPT = "Change the piano to a cello playing the same notes"
|
||||
# Path to any audio-bearing file (wav, mp3, mp4, m4a, flac, ...).
|
||||
@@ -58,18 +64,24 @@ INIT_AUDIO_STRENGTH = 0.6
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/stable-audio-open-1.0-Diffusers",
|
||||
num_gpus=1,
|
||||
)
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path="outputs_audio/stable_audio_a2a/output_a2a.wav",
|
||||
save_video=True,
|
||||
audio_end_in_s=6.0,
|
||||
init_audio=INIT_AUDIO_PATH,
|
||||
init_audio_strength=INIT_AUDIO_STRENGTH,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/stable-audio-open-1.0-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
))
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path="outputs_audio/stable_audio_a2a/output_a2a.wav",
|
||||
save_video=True,
|
||||
),
|
||||
extensions={
|
||||
"audio_end_in_s": 6.0,
|
||||
"init_audio": INIT_AUDIO_PATH,
|
||||
"init_audio_strength": INIT_AUDIO_STRENGTH,
|
||||
},
|
||||
))
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
|
||||
@@ -48,6 +48,9 @@ Prerequisites: same as `basic_stable_audio.py`.
|
||||
import os
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GeneratorConfig, GenerationRequest, OutputConfig,
|
||||
)
|
||||
|
||||
PROMPT = "Steady lo-fi hip hop drum loop with vinyl crackle."
|
||||
# Required: path to the reference audio file (wav, mp3, mp4, m4a, flac,
|
||||
@@ -64,19 +67,25 @@ def main() -> None:
|
||||
f"REFERENCE_AUDIO_PATH={REFERENCE_AUDIO_PATH!r} does not exist. "
|
||||
"Edit this script to point at a real audio file (wav/mp3/mp4/"
|
||||
"m4a/flac) before running.")
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/stable-audio-open-1.0-Diffusers",
|
||||
num_gpus=1,
|
||||
)
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path="outputs_audio/stable_audio_inpaint/output_inpaint.wav",
|
||||
save_video=True,
|
||||
audio_end_in_s=TOTAL_SECONDS,
|
||||
inpaint_audio=REFERENCE_AUDIO_PATH,
|
||||
# Tuple form: keep first KEEP_SECONDS, regenerate the rest.
|
||||
inpaint_mask=(KEEP_SECONDS, TOTAL_SECONDS),
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/stable-audio-open-1.0-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
))
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(
|
||||
output_path="outputs_audio/stable_audio_inpaint/output_inpaint.wav",
|
||||
save_video=True,
|
||||
),
|
||||
extensions={
|
||||
"audio_end_in_s": TOTAL_SECONDS,
|
||||
"inpaint_audio": REFERENCE_AUDIO_PATH,
|
||||
# Tuple form: keep first KEEP_SECONDS, regenerate the rest.
|
||||
"inpaint_mask": (KEEP_SECONDS, TOTAL_SECONDS),
|
||||
},
|
||||
))
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
|
||||
@@ -28,24 +28,27 @@ Prerequisites: same as `basic_stable_audio.py`. The converted repo is
|
||||
public so no gated-access flow is required.
|
||||
"""
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OutputConfig)
|
||||
|
||||
PROMPT = "Lo-fi hip hop instrumental with vinyl crackle and gentle piano."
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"FastVideo/stable-audio-open-small-Diffusers",
|
||||
num_gpus=1,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="FastVideo/stable-audio-open-small-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
))
|
||||
output_path = "outputs_audio/stable_audio_small/output_stable_audio_small.wav"
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
# Small variant trains on a ~11.9s window — keep `audio_end_in_s`
|
||||
# at or below that.
|
||||
audio_end_in_s=6.0,
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(output_path=output_path, save_video=True),
|
||||
# Small variant trains on a ~11.9s window — keep `audio_end_in_s`
|
||||
# at or below that.
|
||||
extensions={"audio_end_in_s": 6.0},
|
||||
))
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
|
||||
@@ -4,6 +4,9 @@ import os
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_turbodiffusion"
|
||||
|
||||
@@ -11,14 +14,17 @@ OUTPUT_PATH = "video_samples_turbodiffusion"
|
||||
def main() -> None:
|
||||
# TurboDiffusion: 1-4 step video generation using RCM scheduler + SLA attention
|
||||
# FastVideo will automatically use TurboDiffusionPipeline when specified
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"loayrashid/TurboWan2.1-T2V-1.3B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
|
||||
# set to false if using RTX 4090
|
||||
# pin_cpu_memory=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="loayrashid/TurboWan2.1-T2V-1.3B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
),
|
||||
# set to false if using RTX 4090
|
||||
# pin_cpu_memory=False,
|
||||
)
|
||||
)
|
||||
|
||||
# Generate videos with the same simple API, regardless of GPU count
|
||||
@@ -28,11 +34,17 @@ def main() -> None:
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
seed=42,
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
seed=42,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Generate another video with a different prompt, without reloading the model!
|
||||
@@ -43,11 +55,17 @@ def main() -> None:
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic."
|
||||
)
|
||||
video2 = generator.generate_video(
|
||||
prompt2,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
seed=42,
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
sampling=SamplingConfig(
|
||||
seed=42,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -4,6 +4,9 @@ import os
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_turbodiffusion_14B"
|
||||
|
||||
@@ -11,10 +14,12 @@ OUTPUT_PATH = "video_samples_turbodiffusion_14B"
|
||||
def main() -> None:
|
||||
# TurboDiffusion 14B: 1-4 step video generation using RCM scheduler + SLA attention
|
||||
# FastVideo will automatically use TurboDiffusionPipeline when specified
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"loayrashid/TurboWan2.1-T2V-14B-Diffusers",
|
||||
# 14B model needs more GPUs
|
||||
num_gpus=2,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="loayrashid/TurboWan2.1-T2V-14B-Diffusers",
|
||||
# 14B model needs more GPUs
|
||||
engine=EngineConfig(num_gpus=2),
|
||||
)
|
||||
)
|
||||
|
||||
prompt = (
|
||||
@@ -22,11 +27,12 @@ def main() -> None:
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
seed=42,
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
sampling=SamplingConfig(seed=42),
|
||||
)
|
||||
)
|
||||
|
||||
# Generate another video with a different prompt, without reloading the model!
|
||||
@@ -37,11 +43,12 @@ def main() -> None:
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic."
|
||||
)
|
||||
video2 = generator.generate_video(
|
||||
prompt2,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
seed=42,
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
sampling=SamplingConfig(seed=42),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -4,6 +4,10 @@ import os
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig,
|
||||
OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
# Use local model path
|
||||
MODEL_PATH = "loayrashid/TurboWan2.2-I2V-A14B-Diffusers"
|
||||
@@ -12,9 +16,11 @@ OUTPUT_PATH = "video_samples_turbodiffusion_i2v"
|
||||
|
||||
def main() -> None:
|
||||
# TurboDiffusion I2V: 1-4 step image-to-video generation
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
MODEL_PATH,
|
||||
num_gpus=2,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=MODEL_PATH,
|
||||
engine=EngineConfig(num_gpus=2),
|
||||
)
|
||||
)
|
||||
|
||||
# Example prompt and image for I2V
|
||||
@@ -24,12 +30,13 @@ def main() -> None:
|
||||
image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
|
||||
|
||||
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
image_path=image_path,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
seed=42,
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
sampling=SamplingConfig(seed=42),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig,
|
||||
OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_wan2_2_14B_t2v"
|
||||
def main():
|
||||
@@ -8,30 +10,37 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=2,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=2,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
# sampling_param.num_frames = 45
|
||||
# sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
|
||||
# Generate videos with the same simple API, regardless of GPU count
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
||||
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
|
||||
)
|
||||
_ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81)
|
||||
# video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(height=720, width=1280, num_frames=81),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
# Generate another video with a different prompt, without reloading the
|
||||
# model!
|
||||
@@ -41,8 +50,14 @@ def main():
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
_ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81)
|
||||
_ = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
sampling=SamplingConfig(height=720, width=1280, num_frames=81),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -1,6 +1,12 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_wan2_1_Fun"
|
||||
OUTPUT_NAME = "wan2.1_test"
|
||||
@@ -9,18 +15,24 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
|
||||
# "alibaba-pai/Wan2.2-Fun-A14B-Control",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
|
||||
# "alibaba-pai/Wan2.2-Fun-A14B-Control",
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
prompt = "一位年轻女性穿着一件粉色的连衣裙,裙子上有白色的装饰和粉色的纽扣。她的头发是紫色的,头上戴着一个红色的大蝴蝶结,显得非常可爱和精致。她还戴着一个红色的领结,整体造型充满了少女感和活力。她的表情温柔,双手轻轻交叉放在身前,姿态优雅。背景是简单的灰色,没有任何多余的装饰,使得人物更加突出。她的妆容清淡自然,突显了她的清新气质。整体画面给人一种甜美、梦幻的感觉,仿佛置身于童话世界中。"
|
||||
@@ -30,7 +42,14 @@ def main():
|
||||
image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png"
|
||||
control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4"
|
||||
|
||||
video = generator.generate_video(prompt, negative_prompt=negative_prompt, image_path=image_path, video_path=control_video_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
inputs=InputConfig(image_path=image_path, video_path=control_video_path),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
# from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig,
|
||||
OffloadConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_wan2_2_14B_i2v"
|
||||
def main():
|
||||
@@ -8,23 +10,36 @@ def main():
|
||||
# model.
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True, # DiT need to be offloaded for MoE
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True, # DiT need to be offloaded for MoE
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
|
||||
pin_cpu_memory=True,
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
|
||||
image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
|
||||
|
||||
video = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, height=832, width=480, num_frames=81)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
sampling=SamplingConfig(height=832, width=480, num_frames=81),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig, OutputConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"
|
||||
def main():
|
||||
@@ -7,22 +10,34 @@ def main():
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder_cpu_offload=False,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_name,
|
||||
engine=EngineConfig(
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False, # set to True if GPU is out of memory
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
# image_encoder=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# I2V is triggered just by passing in an image_path argument
|
||||
prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
|
||||
image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
# Generate another video with a different prompt, without reloading the
|
||||
# model!
|
||||
@@ -34,8 +49,13 @@ def main():
|
||||
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
|
||||
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
|
||||
"cinematic.")
|
||||
video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
|
||||
video2 = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt2,
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Run GLM-Image image-to-image (edit) generation through FastVideo.
|
||||
|
||||
User story:
|
||||
"I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
|
||||
want a minimal edit command (text + image -> edited image), saved as a PNG."
|
||||
|
||||
GLM-Image is a single unified pipeline: passing a condition image switches it
|
||||
from text-to-image to the edit path (the condition enters the DiT via a KV-cache
|
||||
write pass), so the generator config is identical to `basic_glm_image.py` — the
|
||||
`inputs.pil_image` on the request is what selects the edit mode.
|
||||
"""
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default="zai-org/GLM-Image",
|
||||
help="HF id or local diffusers-format GLM-Image weights directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image",
|
||||
default="assets/images/couple.jpg",
|
||||
help="Condition image to edit.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="image_output/edited.png",
|
||||
help="Output PNG path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="Change the background to a snowy mountain landscape at golden hour.",
|
||||
help="Edit instruction.",
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--steps", type=int, default=50)
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.5)
|
||||
parser.add_argument("--seed", type=int, default=1024)
|
||||
parser.add_argument("--num-gpus", type=int, default=1)
|
||||
parser.add_argument("--tp-size", type=int, default=None)
|
||||
parser.add_argument("--sp-size", type=int, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
output = Path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
condition = Image.open(args.image).convert("RGB")
|
||||
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
|
||||
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
|
||||
|
||||
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
|
||||
# pipeline class come from the model's registered defaults — don't override.
|
||||
# The pipeline is registered as t2i; passing inputs.pil_image below switches
|
||||
# it to the edit path.
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
trust_remote_code=True,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
|
||||
),
|
||||
pipeline=PipelineSelection(workload_type="t2i"),
|
||||
)
|
||||
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
try:
|
||||
request = GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
inputs=InputConfig(pil_image=condition),
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=1,
|
||||
fps=1,
|
||||
num_inference_steps=args.steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output.parent),
|
||||
save_video=False,
|
||||
return_frames=True,
|
||||
),
|
||||
)
|
||||
result = generator.generate(request)
|
||||
if isinstance(result, list):
|
||||
result = result[0]
|
||||
|
||||
frames = result.frames
|
||||
if frames is not None and len(frames):
|
||||
Image.fromarray(frames[0]).save(output)
|
||||
print(f"Saved image to {output}")
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -50,23 +50,33 @@ N_DUP = 4 # how many times to duplicate the video for the gen/ref corpora
|
||||
def generate_one_ltx2_video() -> str:
|
||||
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OutputConfig, SamplingConfig)
|
||||
|
||||
Path(OUTPUT_PATH).parent.mkdir(parents=True, exist_ok=True)
|
||||
# Davids048/LTX2-Base-Diffusers is the audio-capable LTX-2 checkpoint
|
||||
# (the Distilled variant ships without the audio VAE, so its mp4
|
||||
# audio track is silence/noise — unusable for audio.* metrics).
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Davids048/LTX2-Base-Diffusers",
|
||||
num_gpus=1,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Davids048/LTX2-Base-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
)
|
||||
)
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
num_frames=121, # ~5s @ 24 fps — long enough for audio.desync (Synchformer ≥14 segments)
|
||||
height=480,
|
||||
width=832,
|
||||
fps=24,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
sampling=SamplingConfig(
|
||||
num_frames=121, # ~5s @ 24 fps — long enough for audio.desync (Synchformer ≥14 segments)
|
||||
height=480,
|
||||
width=832,
|
||||
fps=24,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
generator.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
@@ -21,6 +21,10 @@ Install: ``uv pip install -e .[eval-audio]`` covers both metrics here
|
||||
import torch
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.eval import create_evaluator
|
||||
|
||||
PROMPT = (
|
||||
@@ -39,20 +43,26 @@ METRICS = [
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Davids048/LTX2-Base-Diffusers",
|
||||
num_gpus=1,
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Davids048/LTX2-Base-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
))
|
||||
|
||||
output_path = "outputs_video/ltx2_audio_eval/output.mp4"
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
num_frames=121,
|
||||
height=1088,
|
||||
width=1920,
|
||||
)
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
sampling=SamplingConfig(
|
||||
num_frames=121,
|
||||
height=1088,
|
||||
width=1920,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
),
|
||||
))
|
||||
generator.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
@@ -22,6 +22,10 @@ sharing, or run on a smaller-resolution generation.
|
||||
import torch
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.eval import Evaluator
|
||||
from fastvideo.eval.io import build_eval_kwargs
|
||||
|
||||
@@ -58,19 +62,20 @@ METRICS = [
|
||||
|
||||
def main() -> None:
|
||||
# ----- generation (matches examples/inference/basic/basic_ltx2.py) -----
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Davids048/LTX2-Base-Diffusers",
|
||||
num_gpus=1,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Davids048/LTX2-Base-Diffusers",
|
||||
engine=EngineConfig(num_gpus=1),
|
||||
)
|
||||
)
|
||||
|
||||
output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.1.mp4"
|
||||
generator.generate_video(
|
||||
prompt=PROMPT,
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
num_frames=121,
|
||||
height=1088,
|
||||
width=1920,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
output=OutputConfig(output_path=output_path, save_video=True),
|
||||
sampling=SamplingConfig(num_frames=121, height=1088, width=1920),
|
||||
)
|
||||
)
|
||||
generator.shutdown()
|
||||
# Free residual CUDA memory the generator left behind so the
|
||||
|
||||
@@ -45,6 +45,9 @@ def _generate_videos(rows: list[dict], videos_dir: Path,
|
||||
model: str, num_gpus: int,
|
||||
num_frames: int, height: int, width: int) -> None:
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
videos_dir.mkdir(parents=True, exist_ok=True)
|
||||
todo = [(row, videos_dir / _expected_filename(row)) for row in rows]
|
||||
@@ -55,13 +58,16 @@ def _generate_videos(rows: list[dict], videos_dir: Path,
|
||||
|
||||
print(f"[gen] {len(todo)}/{len(rows)} scenarios to render with {model} "
|
||||
f"({num_frames}x{height}x{width})...")
|
||||
gen = VideoGenerator.from_pretrained(model, num_gpus=num_gpus)
|
||||
gen = VideoGenerator.from_config(GeneratorConfig(
|
||||
model_path=model, engine=EngineConfig(num_gpus=num_gpus),
|
||||
))
|
||||
try:
|
||||
for row, out_path in todo:
|
||||
gen.generate_video(
|
||||
prompt=row["prompt"], output_path=str(out_path), save_video=True,
|
||||
num_frames=num_frames, height=height, width=width,
|
||||
)
|
||||
gen.generate(GenerationRequest(
|
||||
prompt=row["prompt"],
|
||||
sampling=SamplingConfig(num_frames=num_frames, height=height, width=width),
|
||||
output=OutputConfig(output_path=str(out_path), save_video=True),
|
||||
))
|
||||
finally:
|
||||
gen.shutdown()
|
||||
|
||||
|
||||
@@ -43,6 +43,8 @@ def _generate_videos(prompts: list[str], videos_dir: Path,
|
||||
model: str, num_gpus: int,
|
||||
num_frames: int, height: int, width: int) -> None:
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OutputConfig, SamplingConfig)
|
||||
|
||||
videos_dir.mkdir(parents=True, exist_ok=True)
|
||||
todo = [(p, videos_dir / f"{_slugify(p)}.mp4") for p in prompts]
|
||||
@@ -53,13 +55,15 @@ def _generate_videos(prompts: list[str], videos_dir: Path,
|
||||
|
||||
print(f"[gen] {len(todo)}/{len(prompts)} prompts to render with {model} "
|
||||
f"({num_frames}x{height}x{width})...")
|
||||
gen = VideoGenerator.from_pretrained(model, num_gpus=num_gpus)
|
||||
gen = VideoGenerator.from_config(GeneratorConfig(
|
||||
model_path=model, engine=EngineConfig(num_gpus=num_gpus)))
|
||||
try:
|
||||
for prompt, out_path in todo:
|
||||
gen.generate_video(
|
||||
prompt=prompt, output_path=str(out_path), save_video=True,
|
||||
num_frames=num_frames, height=height, width=width,
|
||||
)
|
||||
gen.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_frames=num_frames, height=height, width=width),
|
||||
output=OutputConfig(output_path=str(out_path), save_video=True),
|
||||
))
|
||||
finally:
|
||||
gen.shutdown()
|
||||
|
||||
|
||||
@@ -33,6 +33,13 @@ import json
|
||||
from pathlib import Path
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.eval import create_evaluator
|
||||
from fastvideo.eval.io import load_video
|
||||
|
||||
@@ -99,16 +106,27 @@ def generate(args: argparse.Namespace) -> Path:
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"[gen] loading {args.model} ({args.num_gpus} GPU)...")
|
||||
generator = VideoGenerator.from_pretrained(args.model, num_gpus=args.num_gpus)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model,
|
||||
engine=EngineConfig(num_gpus=args.num_gpus),
|
||||
)
|
||||
)
|
||||
try:
|
||||
print(f"[gen] generating to {out}...")
|
||||
generator.generate_video(
|
||||
prompt=args.prompt,
|
||||
output_path=str(out),
|
||||
save_video=True,
|
||||
num_frames=args.num_frames,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
sampling=SamplingConfig(
|
||||
num_frames=args.num_frames,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(out),
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
@@ -33,7 +33,7 @@ This demo initializes a `VideoGenerator` with the minimum required arguments for
|
||||
|
||||
The core functionality is in the `generate_video` function, which:
|
||||
1. Processes user inputs
|
||||
2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate_video()`)
|
||||
2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate(GenerationRequest(...))`)
|
||||
|
||||
## Gradio Interface
|
||||
|
||||
|
||||
@@ -5,7 +5,13 @@ import time
|
||||
|
||||
import gradio as gr
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
SamplingParam,
|
||||
)
|
||||
from copy import deepcopy
|
||||
|
||||
|
||||
@@ -129,9 +135,22 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators
|
||||
output_dir = "outputs/"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
start_time = time.time()
|
||||
result = generator.generate_video(prompt=prompt, sampling_param=params, save_video=True, return_frames=False)
|
||||
result = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=params.negative_prompt,
|
||||
sampling=SamplingConfig(
|
||||
seed=int(params.seed),
|
||||
guidance_scale=params.guidance_scale,
|
||||
num_frames=int(params.num_frames),
|
||||
height=int(params.height),
|
||||
width=int(params.width),
|
||||
),
|
||||
output=OutputConfig(save_video=True, return_frames=False),
|
||||
)
|
||||
)
|
||||
inference_time = time.time() - start_time
|
||||
logging_info = result.get("logging_info", None)
|
||||
logging_info = result.logging_info
|
||||
if logging_info:
|
||||
stage_names = logging_info.get_execution_order()
|
||||
stage_execution_times = [
|
||||
@@ -550,7 +569,7 @@ def main():
|
||||
for model_path in model_paths:
|
||||
print(f"Loading model: {model_path}")
|
||||
setup_model_environment(model_path)
|
||||
generators[model_path] = VideoGenerator.from_pretrained(model_path)
|
||||
generators[model_path] = VideoGenerator.from_config(GeneratorConfig(model_path=model_path))
|
||||
default_params[model_path] = SamplingParam.from_pretrained(model_path)
|
||||
demo = create_gradio_interface(default_params, generators)
|
||||
print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
|
||||
|
||||
@@ -55,10 +55,11 @@ demo can actually boot:
|
||||
`fastvideo/fastvideo_args.py` currently wires only `ltx2_vae_tiling`.
|
||||
The backing stages (`ltx2_refine.py`, `ltx2_i2v_conditioning.py`) are
|
||||
also missing from `fastvideo/pipelines/stages/`.
|
||||
3. **`fastvideo.configs.sample.base.SamplingParam`** — the import path used
|
||||
by this demo. Upstream moved sampling params to
|
||||
`fastvideo.api.sampling_param`. A re-export shim at the old path, or an
|
||||
import update here once the other two prereqs land, will resolve it.
|
||||
3. **`SamplingParam`** — now imported from `fastvideo.api` (the public
|
||||
re-export of `fastvideo.api.sampling_param`); the old
|
||||
`fastvideo.configs.sample.base` path was removed upstream. `SamplingParam`
|
||||
here only sources model-default slider values — generation itself runs
|
||||
through the typed `GenerationRequest` / `generator.generate(...)` path.
|
||||
|
||||
## Environment variables
|
||||
|
||||
|
||||
@@ -4,8 +4,16 @@ from pathlib import Path
|
||||
|
||||
import gradio as gr
|
||||
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
ComponentConfig,
|
||||
EngineConfig,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
PipelineSelection,
|
||||
SamplingParam,
|
||||
)
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.configs.sample.base import SamplingParam
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo.layers.quantization.fp4_config import FP4Config
|
||||
from fastvideo.utils import maybe_download_model
|
||||
@@ -48,28 +56,44 @@ def main():
|
||||
refine_upsampler_path = resolve_refine_upsampler_path(resolved_model_path)
|
||||
print(f"Using refine upsampler: {refine_upsampler_path}")
|
||||
|
||||
generators[model_path] = VideoGenerator.from_pretrained(
|
||||
str(resolved_model_path),
|
||||
num_gpus=1,
|
||||
ltx2_refine_enabled=True,
|
||||
ltx2_refine_upsampler_path=str(refine_upsampler_path),
|
||||
ltx2_refine_lora_path="", # disable refine LoRA for distilled model
|
||||
ltx2_refine_num_inference_steps=2,
|
||||
ltx2_refine_guidance_scale=1.0,
|
||||
ltx2_refine_add_noise=True,
|
||||
pipeline_config=pipeline_config,
|
||||
enable_torch_compile=True,
|
||||
enable_torch_compile_text_encoder=True,
|
||||
torch_compile_kwargs={
|
||||
"backend": "inductor",
|
||||
"fullgraph": True,
|
||||
"mode": "max-autotune-no-cudagraphs",
|
||||
"dynamic": False,
|
||||
},
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
ltx2_vae_tiling=False,
|
||||
generators[model_path] = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=str(resolved_model_path),
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=False,
|
||||
),
|
||||
compile=CompileConfig(
|
||||
enabled=True,
|
||||
text_encoder_enabled=True,
|
||||
backend="inductor",
|
||||
fullgraph=True,
|
||||
mode="max-autotune-no-cudagraphs",
|
||||
dynamic=False,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
upsampler_weights=str(refine_upsampler_path),
|
||||
# Empty refine LoRA path (distilled needs none) -> omit.
|
||||
),
|
||||
vae_tiling=False,
|
||||
preset_overrides={
|
||||
"refine": {
|
||||
"enabled": True,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 1.0,
|
||||
"add_noise": True,
|
||||
},
|
||||
},
|
||||
# PipelineConfig object (with FP4 quant wired on above) has
|
||||
# no first-class typed field; route via experimental.
|
||||
experimental={"pipeline_config": pipeline_config},
|
||||
),
|
||||
)
|
||||
)
|
||||
default_params[model_path] = apply_ltx2_defaults(
|
||||
SamplingParam.from_pretrained(str(resolved_model_path))
|
||||
|
||||
@@ -4,7 +4,7 @@ from pathlib import Path
|
||||
import torch
|
||||
import torch._inductor.config
|
||||
|
||||
from fastvideo.configs.sample.base import SamplingParam
|
||||
from fastvideo.api import SamplingParam
|
||||
|
||||
LOCAL_DEMO_DIR = Path(__file__).resolve().parent
|
||||
CLASSIFIER_DIR = Path(
|
||||
|
||||
@@ -5,8 +5,14 @@ from copy import deepcopy
|
||||
|
||||
import gradio as gr
|
||||
|
||||
from fastvideo.configs.sample.base import SamplingParam
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
GenerationRequest,
|
||||
InputConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
SamplingParam,
|
||||
)
|
||||
|
||||
from .config import (
|
||||
DEFAULT_FPS,
|
||||
@@ -69,40 +75,38 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators
|
||||
output_path = str(OUTPUT_DIR / video_filename)
|
||||
params.output_path = output_path
|
||||
start_time = time.perf_counter()
|
||||
result = generator.generate_video(
|
||||
prompt=prompt,
|
||||
output_path=output_path,
|
||||
fps=DEFAULT_FPS,
|
||||
seed=int(params.seed),
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
guidance_scale=float(params.guidance_scale),
|
||||
height=int(params.height),
|
||||
width=int(params.width),
|
||||
num_frames=int(params.num_frames),
|
||||
num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS,
|
||||
negative_prompt=params.negative_prompt,
|
||||
image_path=params.image_path,
|
||||
ltx2_image_crf=0.0
|
||||
result = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=params.negative_prompt,
|
||||
inputs=InputConfig(image_path=params.image_path),
|
||||
sampling=SamplingConfig(
|
||||
seed=int(params.seed),
|
||||
fps=DEFAULT_FPS,
|
||||
guidance_scale=float(params.guidance_scale),
|
||||
height=int(params.height),
|
||||
width=int(params.width),
|
||||
num_frames=int(params.num_frames),
|
||||
num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=output_path,
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
# LTX-2 i2v knob without a first-class typed field yet.
|
||||
extensions={"ltx2_image_crf": 0.0},
|
||||
)
|
||||
)
|
||||
wall_time = time.perf_counter() - start_time
|
||||
generation_time = (
|
||||
result.get("generation_time")
|
||||
if isinstance(result, dict) else None
|
||||
)
|
||||
e2e_latency = (
|
||||
result.get("e2e_latency")
|
||||
if isinstance(result, dict) else None
|
||||
)
|
||||
generation_time = result.generation_time
|
||||
e2e_latency = result.extra.get("e2e_latency")
|
||||
if generation_time is None:
|
||||
generation_time = wall_time
|
||||
if e2e_latency is None:
|
||||
e2e_latency = wall_time
|
||||
resolved_output_path = (
|
||||
result.get("output_path", output_path)
|
||||
if isinstance(result, dict) else output_path
|
||||
)
|
||||
logging_info = result.get("logging_info", None) if isinstance(result, dict) else None
|
||||
resolved_output_path = result.video_path or output_path
|
||||
logging_info = result.logging_info
|
||||
if logging_info:
|
||||
stage_names = logging_info.get_execution_order()
|
||||
stage_execution_times = [
|
||||
|
||||
@@ -9,6 +9,7 @@ import uvicorn
|
||||
from fastapi import FastAPI, Request, HTTPException
|
||||
from fastapi.responses import HTMLResponse, FileResponse
|
||||
|
||||
from fastvideo.api import EngineConfig, GeneratorConfig, OffloadConfig
|
||||
from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator
|
||||
from fastvideo.models.dits.matrixgame2.utils import expand_action_to_frames
|
||||
|
||||
@@ -572,14 +573,20 @@ def main():
|
||||
|
||||
print(f"Loading model: {model_path}")
|
||||
setup_model_environment(model_path)
|
||||
generator = StreamingVideoGenerator.from_pretrained(
|
||||
model_path,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True,
|
||||
generator = StreamingVideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
generators = {model_path: generator}
|
||||
|
||||
@@ -3,7 +3,6 @@ import os
|
||||
import torch
|
||||
import base64
|
||||
import io
|
||||
from copy import deepcopy
|
||||
from typing import Dict, Any, Optional, List
|
||||
import signal
|
||||
import sys
|
||||
@@ -20,6 +19,17 @@ import imageio
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
from prometheus_client import Counter, Histogram, generate_latest
|
||||
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
InputConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
NUM_GPUS = 16
|
||||
DEFAULT_FPS = 16
|
||||
SEED_RANGE_MAX = 1_000_000
|
||||
@@ -136,10 +146,10 @@ def setup_model_environment(model_path: str) -> None:
|
||||
|
||||
|
||||
def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, List[str], List[float]]:
|
||||
frames = result if isinstance(result, list) else result.get("frames", [])
|
||||
generation_time = result.get("generation_time", 0.0) if isinstance(result, dict) else 0.0
|
||||
|
||||
logging_info = result.get("logging_info", None)
|
||||
frames = result.frames or []
|
||||
generation_time = result.generation_time or 0.0
|
||||
|
||||
logging_info = result.logging_info
|
||||
if logging_info:
|
||||
stage_names = logging_info.get_execution_order()
|
||||
stage_execution_times = [
|
||||
@@ -153,24 +163,29 @@ def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, Lis
|
||||
return frames, generation_time, stage_names, stage_execution_times
|
||||
|
||||
|
||||
def prepare_sampling_params(video_request: VideoGenerationRequest, default_params: Any) -> Any:
|
||||
params = deepcopy(default_params)
|
||||
params.prompt = video_request.prompt
|
||||
|
||||
if video_request.use_negative_prompt:
|
||||
params.negative_prompt = video_request.negative_prompt
|
||||
def prepare_generation_request(video_request: VideoGenerationRequest, image_path: Optional[str] = None) -> Any:
|
||||
seed = (video_request.seed if not video_request.randomize_seed
|
||||
else torch.randint(0, SEED_RANGE_MAX, (1,)).item())
|
||||
|
||||
params.seed = (video_request.seed if not video_request.randomize_seed
|
||||
else torch.randint(0, SEED_RANGE_MAX, (1,)).item())
|
||||
params.randomize_seed = video_request.randomize_seed
|
||||
params.guidance_scale = video_request.guidance_scale
|
||||
params.num_frames = video_request.num_frames
|
||||
params.height = video_request.height
|
||||
params.width = video_request.width
|
||||
params.save_video = False
|
||||
params.return_frames = True
|
||||
|
||||
return params
|
||||
# "" explicitly clears the model preset's negative prompt (None would
|
||||
# inherit it, changing this demo's long-standing behavior).
|
||||
negative_prompt = video_request.negative_prompt if video_request.use_negative_prompt else ""
|
||||
|
||||
request = GenerationRequest(
|
||||
prompt=video_request.prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
sampling=SamplingConfig(
|
||||
seed=seed,
|
||||
guidance_scale=video_request.guidance_scale,
|
||||
num_frames=video_request.num_frames,
|
||||
height=video_request.height,
|
||||
width=video_request.width,
|
||||
),
|
||||
output=OutputConfig(save_video=False, return_frames=True),
|
||||
)
|
||||
|
||||
return request, seed
|
||||
|
||||
|
||||
class BaseModelDeployment:
|
||||
@@ -185,31 +200,34 @@ class BaseModelDeployment:
|
||||
|
||||
def _initialize_generator(self, config: Dict[str, Any]) -> None:
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
print(f"Initializing model: {self.model_path}")
|
||||
self.generator = VideoGenerator.from_pretrained(
|
||||
model_path=self.model_path,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
|
||||
dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # TODO: hardocde for I2V
|
||||
dit_precision="fp32", # TODO: hardocde for I2V
|
||||
dit_cpu_offload=config["dit_cpu_offload"],
|
||||
vae_cpu_offload=config["vae_cpu_offload"],
|
||||
VSA_sparsity=config["VSA_sparsity"],
|
||||
enable_stage_verification=False,
|
||||
self.generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=self.model_path,
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
enable_stage_verification=False,
|
||||
offload=OffloadConfig(
|
||||
text_encoder=config["text_encoder_cpu_offload"],
|
||||
dit=config["dit_cpu_offload"],
|
||||
vae=config["vae_cpu_offload"],
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
# I2V knobs without first-class typed fields yet.
|
||||
experimental={
|
||||
"dmd_denoising_steps": [1000, 850, 700, 550, 350, 275, 200, 125],
|
||||
"dit_precision": "fp32",
|
||||
"VSA_sparsity": config["VSA_sparsity"],
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
self.default_params = SamplingParam.from_pretrained(self.model_path)
|
||||
self.default_params.seed = 1000
|
||||
self.default_params.num_frames = 73
|
||||
self.default_params.width = 832
|
||||
self.default_params.height = 480
|
||||
|
||||
def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
|
||||
total_start_time = time.time()
|
||||
|
||||
params = prepare_sampling_params(video_request, self.default_params)
|
||||
|
||||
# Save image if provided (for I2V)
|
||||
image_path = None
|
||||
@@ -218,19 +236,15 @@ class BaseModelDeployment:
|
||||
if image_path is None:
|
||||
return VideoGenerationResponse(
|
||||
video_data=None,
|
||||
seed=params.seed,
|
||||
seed=video_request.seed,
|
||||
success=False,
|
||||
error_message="Failed to save input image",
|
||||
)
|
||||
|
||||
request, seed = prepare_generation_request(video_request, image_path)
|
||||
|
||||
inference_start_time = time.time()
|
||||
result = self.generator.generate_video(
|
||||
prompt=video_request.prompt,
|
||||
sampling_param=params,
|
||||
image_path=image_path,
|
||||
save_video=False,
|
||||
return_frames=True,
|
||||
)
|
||||
result = self.generator.generate(request)
|
||||
inference_time = time.time() - inference_start_time
|
||||
|
||||
frames, generation_time, stage_names, stage_execution_times = process_generation_result(result)
|
||||
@@ -250,7 +264,7 @@ class BaseModelDeployment:
|
||||
|
||||
return VideoGenerationResponse(
|
||||
video_data=video_data,
|
||||
seed=params.seed,
|
||||
seed=seed,
|
||||
success=True,
|
||||
generation_time=generation_time,
|
||||
inference_time=inference_time,
|
||||
|
||||
@@ -1,18 +1,31 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (
|
||||
ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, PipelineSelection, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "./lora_out"
|
||||
def main():
|
||||
# Initialize VideoGenerator with the Wan model
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
num_gpus=1,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
lora_path="benjamin-paine/steamboat-willie-1.3b",
|
||||
lora_nickname="steamboat"
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="benjamin-paine/steamboat-willie-1.3b",
|
||||
),
|
||||
experimental={"lora_nickname": "steamboat"},
|
||||
),
|
||||
)
|
||||
)
|
||||
kwargs = {
|
||||
"height": 480,
|
||||
@@ -26,25 +39,32 @@ def main():
|
||||
prompt = "steamboat willie style, golden era animation, close-up of a short fluffy monster kneeling beside a melting red candle. the mood is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image."
|
||||
negative_prompt = "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"
|
||||
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
# sampling_param=sampling_param,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
negative_prompt=negative_prompt,
|
||||
**kwargs
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling=SamplingConfig(**kwargs),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
|
||||
prompt = "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco], 1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds, live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid."
|
||||
negative_prompt = "bad quality video,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
negative_prompt=negative_prompt,
|
||||
**kwargs
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling=SamplingConfig(**kwargs),
|
||||
output=OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -2,45 +2,60 @@
|
||||
Inference using a LoRA checkpoint from FastVideo trainer.
|
||||
"""
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.api import (ComponentConfig, EngineConfig, GenerationRequest,
|
||||
GeneratorConfig, OffloadConfig, OutputConfig,
|
||||
PipelineSelection, SamplingConfig)
|
||||
|
||||
OUTPUT_PATH = "./lora_out"
|
||||
def main():
|
||||
# Initialize VideoGenerator with the Wan model
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
num_gpus=1,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-160/transformer",
|
||||
lora_nickname="crush_smol"
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-160/transformer",
|
||||
),
|
||||
experimental={"lora_nickname": "crush_smol"},
|
||||
),
|
||||
))
|
||||
generator.unmerge_lora_weights()
|
||||
kwargs = {
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 77,
|
||||
"guidance_scale": 6.0,
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
}
|
||||
sampling = SamplingConfig(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=77,
|
||||
guidance_scale=6.0,
|
||||
num_inference_steps=50,
|
||||
seed=42,
|
||||
)
|
||||
output = OutputConfig(
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
)
|
||||
# Generate video with LoRA style
|
||||
prompt = "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press."
|
||||
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
**kwargs
|
||||
)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=sampling,
|
||||
output=output,
|
||||
))
|
||||
prompt = "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press."
|
||||
video = generator.generate_video(
|
||||
prompt,
|
||||
output_path=OUTPUT_PATH,
|
||||
save_video=True,
|
||||
**kwargs
|
||||
)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=sampling,
|
||||
output=output,
|
||||
))
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
@@ -37,6 +37,17 @@ import imageio
|
||||
import torch
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
ComponentConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
|
||||
|
||||
@@ -148,35 +159,49 @@ def build_generator(args: argparse.Namespace) -> VideoGenerator:
|
||||
|
||||
compile_enabled = not args.no_compile
|
||||
|
||||
extra_kwargs = {}
|
||||
# ``pipeline_config`` is a PipelineConfig object (not a string path) and
|
||||
# ``output_type`` has no first-class typed field, so both are routed through
|
||||
# the pipeline experimental escape hatch.
|
||||
experimental = {"pipeline_config": pipeline_config}
|
||||
|
||||
components = ComponentConfig()
|
||||
if args.distilled_model:
|
||||
weights_path = resolve_distilled_weights(args.distilled_model)
|
||||
print(f"Using distilled weights: {args.distilled_model} -> {weights_path}")
|
||||
extra_kwargs["init_weights_from_safetensors"] = weights_path
|
||||
components.transformer_weights = weights_path
|
||||
|
||||
if args.taehv:
|
||||
# Skip the in-pipeline VAE decode entirely: the pipeline returns raw
|
||||
# latents, the Wan VAE is offloaded to CPU (and not compiled) since we
|
||||
# decode with TAEHV in this script instead.
|
||||
extra_kwargs["output_type"] = "latent"
|
||||
experimental["output_type"] = "latent"
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_id,
|
||||
pipeline_config=pipeline_config,
|
||||
num_gpus=args.num_gpus,
|
||||
# Keep everything resident on the GPU -- no offloading, except the
|
||||
# unused Wan VAE when TAEHV handles decoding.
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
vae_cpu_offload=args.taehv,
|
||||
text_encoder_cpu_offload=False,
|
||||
pin_cpu_memory=False,
|
||||
enable_torch_compile=compile_enabled,
|
||||
enable_torch_compile_text_encoder=compile_enabled,
|
||||
enable_torch_compile_vae=compile_enabled and not args.taehv,
|
||||
**extra_kwargs,
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=model_id,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
# Keep everything resident on the GPU -- no offloading, except the
|
||||
# unused Wan VAE when TAEHV handles decoding.
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=False,
|
||||
vae=args.taehv,
|
||||
text_encoder=False,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
compile=CompileConfig(
|
||||
enabled=compile_enabled,
|
||||
text_encoder_enabled=compile_enabled,
|
||||
vae_enabled=compile_enabled and not args.taehv,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=components,
|
||||
experimental=experimental,
|
||||
),
|
||||
)
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
return generator
|
||||
|
||||
|
||||
@@ -237,11 +262,11 @@ def main() -> None:
|
||||
# runs below measure steady-state latency only.
|
||||
with silence_request_log():
|
||||
for _ in range(args.warmups):
|
||||
warm = generator.generate(request={
|
||||
"prompt": PROMPT,
|
||||
"sampling": {"num_inference_steps": 2, "guidance_scale": args.guidance_scale},
|
||||
"output": {"save_video": False, "return_frames": args.taehv},
|
||||
})
|
||||
warm = generator.generate(GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
sampling=SamplingConfig(num_inference_steps=2, guidance_scale=args.guidance_scale),
|
||||
output=OutputConfig(save_video=False, return_frames=args.taehv),
|
||||
))
|
||||
if args.taehv:
|
||||
taehv.decode(warm.samples)
|
||||
|
||||
@@ -257,18 +282,18 @@ def main() -> None:
|
||||
frames = None
|
||||
with silence_request_log():
|
||||
for i in range(args.benchmark_runs):
|
||||
result = generator.generate(request={
|
||||
"prompt": PROMPT,
|
||||
"sampling": {
|
||||
"num_inference_steps": args.infer_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
},
|
||||
"output": {
|
||||
"save_video": False,
|
||||
"return_frames": args.taehv,
|
||||
"output_path": output_path,
|
||||
},
|
||||
})
|
||||
result = generator.generate(GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
sampling=SamplingConfig(
|
||||
num_inference_steps=args.infer_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
),
|
||||
output=OutputConfig(
|
||||
save_video=False,
|
||||
return_frames=args.taehv,
|
||||
output_path=output_path,
|
||||
),
|
||||
))
|
||||
denoise_elapsed = result.generation_time
|
||||
denoise_times.append(denoise_elapsed)
|
||||
|
||||
|
||||
@@ -2,31 +2,41 @@ import os
|
||||
import time
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig,
|
||||
OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
def main():
|
||||
# set the attention backend
|
||||
# set the attention backend
|
||||
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
|
||||
|
||||
start_time = time.perf_counter()
|
||||
gen = VideoGenerator.from_pretrained(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
num_gpus=1,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
)
|
||||
gen = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
vae=False,
|
||||
text_encoder=True,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
),
|
||||
),
|
||||
))
|
||||
load_time = time.perf_counter() - start_time
|
||||
print(f"Model loading time: {load_time:.2f} seconds")
|
||||
|
||||
gen_start_time = time.perf_counter()
|
||||
|
||||
gen.generate_video(
|
||||
prompt=
|
||||
"Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
|
||||
seed=1024,
|
||||
output_path="example_outputs/")
|
||||
|
||||
gen.generate(
|
||||
GenerationRequest(
|
||||
prompt=
|
||||
"Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
|
||||
sampling=SamplingConfig(seed=1024),
|
||||
output=OutputConfig(output_path="example_outputs/")))
|
||||
|
||||
generation_time = time.perf_counter() - gen_start_time
|
||||
print(f"Video generation time: {generation_time:.2f} seconds")
|
||||
|
||||
|
||||
@@ -18,6 +18,10 @@ import os
|
||||
import time
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "video_samples"
|
||||
|
||||
@@ -39,17 +43,24 @@ def main():
|
||||
mode += "_compile"
|
||||
print(f"Mode: {mode.upper()}")
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
args.model,
|
||||
num_gpus=args.num_gpus,
|
||||
nvfp4_fa4=args.nvfp4_fa4,
|
||||
use_fsdp_inference=not args.nvfp4_fa4,
|
||||
dit_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
enable_torch_compile=args.compile,
|
||||
)
|
||||
if args.nvfp4_fa4:
|
||||
os.environ["FASTVIDEO_NVFP4_FA4"] = "1"
|
||||
os.environ.setdefault("CUTE_DSL_ENABLE_TVM_FFI", "1")
|
||||
|
||||
generator = VideoGenerator.from_config(GeneratorConfig(
|
||||
model_path=args.model,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=not args.nvfp4_fa4,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=False,
|
||||
vae=True,
|
||||
text_encoder=True,
|
||||
),
|
||||
compile=CompileConfig(enabled=args.compile),
|
||||
),
|
||||
))
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
@@ -59,16 +70,22 @@ def main():
|
||||
|
||||
n_warmup = 2 if args.compile else 1
|
||||
for i in range(n_warmup):
|
||||
generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 2},
|
||||
"output": {"save_video": False}})
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=2),
|
||||
output=OutputConfig(save_video=False),
|
||||
))
|
||||
|
||||
os.makedirs(OUTPUT_PATH, exist_ok=True)
|
||||
start = time.time()
|
||||
generator.generate(request={
|
||||
"prompt": prompt,
|
||||
"sampling": {"num_inference_steps": args.infer_steps},
|
||||
"output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")},
|
||||
})
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=args.infer_steps),
|
||||
output=OutputConfig(
|
||||
save_video=True,
|
||||
output_path=os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4"),
|
||||
),
|
||||
))
|
||||
elapsed = time.time() - start
|
||||
print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
|
||||
f"({args.infer_steps / elapsed:.2f} it/s)")
|
||||
|
||||
@@ -69,6 +69,16 @@ def main():
|
||||
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.layers.quantization import get_quantization_config
|
||||
|
||||
mode = "bf16" if args.bf16 else f"fp8_{args.granularity}"
|
||||
@@ -79,25 +89,33 @@ def main():
|
||||
|
||||
taehv_model = load_taehv(args.taehv_checkpoint) if use_taehv else None
|
||||
|
||||
# transformer_quant needs a QuantizationConfig *instance* — the bare string
|
||||
# is not resolved on the from_pretrained kwarg path.
|
||||
extra = {} if args.bf16 else {
|
||||
"transformer_quant": get_quantization_config("FP8")(granularity=args.granularity)
|
||||
}
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
args.model,
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
vae_cpu_offload=use_taehv,
|
||||
text_encoder_cpu_offload=False,
|
||||
pin_cpu_memory=False,
|
||||
enable_torch_compile=not args.no_compile,
|
||||
enable_torch_compile_vae=not args.no_compile and not use_taehv,
|
||||
output_type="latent" if use_taehv else "pil",
|
||||
**extra,
|
||||
)
|
||||
# ``output_type`` and ``transformer_quant`` have no first-class typed
|
||||
# fields yet, so they ride the pipeline.experimental escape hatch (same
|
||||
# place the legacy from_pretrained shim routed them). The typed
|
||||
# QuantizationConfig only accepts a quant-name string, so it can't carry
|
||||
# FP8's ``granularity`` arg — pass the resolved config instance instead.
|
||||
experimental = {"output_type": "latent" if use_taehv else "pil"}
|
||||
if not args.bf16:
|
||||
experimental["transformer_quant"] = get_quantization_config("FP8")(granularity=args.granularity)
|
||||
generator = VideoGenerator.from_config(GeneratorConfig(
|
||||
model_path=args.model,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=False,
|
||||
vae=use_taehv,
|
||||
text_encoder=False,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
compile=CompileConfig(
|
||||
enabled=not args.no_compile,
|
||||
vae_enabled=not args.no_compile and not use_taehv,
|
||||
),
|
||||
),
|
||||
pipeline=PipelineSelection(experimental=experimental),
|
||||
))
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
@@ -107,28 +125,34 @@ def main():
|
||||
|
||||
n_warmup = 1 if not args.no_compile else 0
|
||||
for _ in range(n_warmup):
|
||||
generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 3, "guidance_scale": 1.0},
|
||||
"output": {"save_video": False}})
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=3, guidance_scale=1.0),
|
||||
output=OutputConfig(save_video=False),
|
||||
))
|
||||
|
||||
os.makedirs(OUTPUT_PATH, exist_ok=True)
|
||||
start = time.time()
|
||||
if use_taehv:
|
||||
result = generator.generate(request={
|
||||
"prompt": prompt,
|
||||
"sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
|
||||
"output": {"save_video": False},
|
||||
})
|
||||
result = generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=args.infer_steps, guidance_scale=1.0),
|
||||
output=OutputConfig(save_video=False),
|
||||
))
|
||||
import imageio
|
||||
frames = decode_with_taehv(taehv_model, result.samples)
|
||||
video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
|
||||
imageio.mimsave(video_path, frames, fps=16, format="mp4")
|
||||
print(f"Saved TAEHV-decoded video to: {video_path}")
|
||||
else:
|
||||
generator.generate(request={
|
||||
"prompt": prompt,
|
||||
"sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
|
||||
"output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")},
|
||||
})
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=args.infer_steps, guidance_scale=1.0),
|
||||
output=OutputConfig(
|
||||
save_video=True,
|
||||
output_path=os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4"),
|
||||
),
|
||||
))
|
||||
elapsed = time.time() - start
|
||||
print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
|
||||
f"({args.infer_steps / elapsed:.2f} it/s)")
|
||||
|
||||
@@ -45,26 +45,29 @@ def main():
|
||||
|
||||
# Import after the env var so the platform picks up the selection.
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.layers.quantization import get_quantization_config
|
||||
from fastvideo.api import (
|
||||
CompileConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, QuantizationConfig, SamplingConfig,
|
||||
)
|
||||
|
||||
mode = "bf16" if args.bf16 else args.quant_method
|
||||
if args.compile:
|
||||
mode += "_compile"
|
||||
print(f"Mode: {mode.upper()}")
|
||||
|
||||
# transformer_quant needs a QuantizationConfig *instance* — the bare string
|
||||
# is not resolved on the from_pretrained kwarg path.
|
||||
extra = {} if args.bf16 else {"transformer_quant": get_quantization_config(args.quant_method)()}
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
args.model,
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=args.bf16,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
enable_torch_compile=args.compile,
|
||||
**extra,
|
||||
)
|
||||
# transformer_quant takes the config name (string); the typed path resolves
|
||||
# the QuantizationConfig class.
|
||||
quantization = None if args.bf16 else QuantizationConfig(transformer_quant=args.quant_method)
|
||||
generator = VideoGenerator.from_config(GeneratorConfig(
|
||||
model_path=args.model,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=args.bf16,
|
||||
offload=OffloadConfig(dit=False, vae=True, text_encoder=True),
|
||||
compile=CompileConfig(enabled=args.compile),
|
||||
quantization=quantization,
|
||||
),
|
||||
))
|
||||
|
||||
prompt = (
|
||||
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
||||
@@ -74,16 +77,20 @@ def main():
|
||||
|
||||
n_warmup = 2 if args.compile else 1
|
||||
for _ in range(n_warmup):
|
||||
generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 2},
|
||||
"output": {"save_video": False}})
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=2),
|
||||
output=OutputConfig(save_video=False),
|
||||
))
|
||||
|
||||
os.makedirs(OUTPUT_PATH, exist_ok=True)
|
||||
start = time.time()
|
||||
generator.generate(request={
|
||||
"prompt": prompt,
|
||||
"sampling": {"num_inference_steps": args.infer_steps},
|
||||
"output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")},
|
||||
})
|
||||
generator.generate(GenerationRequest(
|
||||
prompt=prompt,
|
||||
sampling=SamplingConfig(num_inference_steps=args.infer_steps),
|
||||
output=OutputConfig(save_video=True,
|
||||
output_path=os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")),
|
||||
))
|
||||
elapsed = time.time() - start
|
||||
print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
|
||||
f"({args.infer_steps / elapsed:.2f} it/s)")
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (ComponentConfig, EngineConfig, GenerationRequest,
|
||||
GeneratorConfig, InputConfig, OffloadConfig,
|
||||
OutputConfig, PipelineSelection, QuantizationConfig)
|
||||
import argparse
|
||||
|
||||
OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"
|
||||
@@ -10,28 +13,39 @@ def main(text_encoder_path: str):
|
||||
# If a local path is provided, FastVideo will make a best effort
|
||||
# attempt to identify the optimal arguments.
|
||||
model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
# FastVideo will automatically handle distributed setup
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
dit_cpu_offload=True,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
# AbsMaxFP8 is the quantization method used by ComfyUI;
|
||||
# check fastvideo/layers/quantization/* for more quantization methods
|
||||
override_text_encoder_quant="AbsMaxFP8",
|
||||
# for Wan 2.2, this is the path to "umt5_xxl_fp8_e4m3fn_scaled.safetensors"
|
||||
override_text_encoder_safetensors=text_encoder_path,
|
||||
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=model_name,
|
||||
# FastVideo will automatically handle distributed setup
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=True,
|
||||
offload=OffloadConfig(
|
||||
dit=True,
|
||||
vae=False,
|
||||
text_encoder=False,
|
||||
pin_cpu_memory=
|
||||
True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
|
||||
),
|
||||
# AbsMaxFP8 is the quantization method used by ComfyUI;
|
||||
# check fastvideo/layers/quantization/* for more quantization methods
|
||||
quantization=QuantizationConfig(text_encoder_quant="AbsMaxFP8"),
|
||||
),
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
# for Wan 2.2, this is the path to "umt5_xxl_fp8_e4m3fn_scaled.safetensors"
|
||||
text_encoder_weights=text_encoder_path, ), ),
|
||||
))
|
||||
|
||||
# I2V is triggered just by passing in an image_path argument
|
||||
prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
|
||||
image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
|
||||
video = generator.generate_video(
|
||||
prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path
|
||||
)
|
||||
video = generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=prompt,
|
||||
inputs=InputConfig(image_path=image_path),
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, save_video=True),
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -22,6 +22,14 @@ import os
|
||||
import time
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
PROMPT = (
|
||||
"A high-definition video of a robotic arm welding a metal structure, "
|
||||
@@ -42,24 +50,27 @@ def main() -> None:
|
||||
|
||||
os.makedirs("video_samples", exist_ok=True)
|
||||
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
args.model,
|
||||
num_gpus=args.num_gpus,
|
||||
enable_torch_compile=args.compile,
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
compile=CompileConfig(enabled=args.compile),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def _run(tag: str) -> float:
|
||||
save = tag == "measured"
|
||||
# Modern typed-request API (generate_video is deprecated). Same
|
||||
# prompt/seed/shapes both runs so the compiled graph is reused.
|
||||
request: dict = {
|
||||
"prompt": PROMPT,
|
||||
"sampling": {"seed": 1024},
|
||||
"output": {"save_video": save},
|
||||
}
|
||||
# Same prompt/seed/shapes both runs so the compiled graph is reused.
|
||||
output = OutputConfig(save_video=save)
|
||||
if save:
|
||||
request["output"]["output_path"] = (
|
||||
f"video_samples/torch_compile_{tag}.mp4")
|
||||
output.output_path = f"video_samples/torch_compile_{tag}.mp4"
|
||||
request = GenerationRequest(
|
||||
prompt=PROMPT,
|
||||
sampling=SamplingConfig(seed=1024),
|
||||
output=output,
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
generator.generate(request)
|
||||
return time.perf_counter() - t0
|
||||
|
||||
@@ -1,105 +0,0 @@
|
||||
# AnyFlow on-policy DMD — Wan 2.1 T2V 1.3B.
|
||||
#
|
||||
# Stage 2 of the AnyFlow two-stage recipe. Continues from the pretrain
|
||||
# checkpoint; refines the student via DMD2 with a multi-step Euler-flow
|
||||
# rollout from pure noise. Teacher provides the real score, critic
|
||||
# learns the fake score; both inherited from DMD2Method.
|
||||
#
|
||||
# Replace <PATH_TO_PRETRAIN_CKPT> with the output of the pretrain stage,
|
||||
# or with the NVIDIA-released checkpoint
|
||||
# nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers to bootstrap directly from
|
||||
# the paper weights (the delta_embedder rename is handled by the
|
||||
# param_names_mapping in WanVideoArchConfig).
|
||||
|
||||
models:
|
||||
student:
|
||||
_target_: fastvideo.train.models.wan.WanModel
|
||||
init_from: <PATH_TO_PRETRAIN_CKPT>
|
||||
trainable: true
|
||||
teacher:
|
||||
_target_: fastvideo.train.models.wan.WanModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
|
||||
trainable: false
|
||||
disable_custom_init_weights: true
|
||||
critic:
|
||||
_target_: fastvideo.train.models.wan.WanModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: true
|
||||
disable_custom_init_weights: true
|
||||
|
||||
method:
|
||||
_target_: fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod
|
||||
rollout_mode: simulate
|
||||
generator_update_interval: 5
|
||||
real_score_guidance_scale: 3.0
|
||||
dmd_denoising_steps: [999, 937, 833, 624]
|
||||
warp_denoising_step: false
|
||||
|
||||
# AnyFlow rollout knobs.
|
||||
student_sample_steps: 4
|
||||
use_mean_velocity: true
|
||||
t_list_override: [999.0, 937.0, 833.0, 624.0, 0.0]
|
||||
dmd_score_r_value: 0.0 # DMD scoring conditioning is at r=0 (consistency target).
|
||||
|
||||
# Critic optimizer (DMD2 inherited).
|
||||
fake_score_learning_rate: 8.0e-6
|
||||
fake_score_betas: [0.0, 0.999]
|
||||
fake_score_lr_scheduler: constant
|
||||
|
||||
attn_kind: vsa
|
||||
|
||||
training:
|
||||
distributed:
|
||||
num_gpus: 8
|
||||
sp_size: 1
|
||||
tp_size: 1
|
||||
hsdp_replicate_dim: 1
|
||||
hsdp_shard_dim: 8
|
||||
|
||||
data:
|
||||
data_path: data/preprocessed
|
||||
dataloader_num_workers: 4
|
||||
train_batch_size: 1
|
||||
training_cfg_rate: 0.0
|
||||
seed: 1000
|
||||
num_latent_t: 21
|
||||
num_height: 480
|
||||
num_width: 832
|
||||
num_frames: 81
|
||||
|
||||
optimizer:
|
||||
learning_rate: 2.0e-6
|
||||
betas: [0.0, 0.999]
|
||||
weight_decay: 0.01
|
||||
lr_scheduler: constant
|
||||
lr_warmup_steps: 0
|
||||
|
||||
loop:
|
||||
max_train_steps: 4000
|
||||
gradient_accumulation_steps: 1
|
||||
|
||||
checkpoint:
|
||||
output_dir: outputs/wan2.1_anyflow_onpolicy
|
||||
training_state_checkpointing_steps: 500
|
||||
checkpoints_total_limit: 3
|
||||
resume_from_checkpoint: latest
|
||||
|
||||
tracker:
|
||||
project_name: anyflow-wan
|
||||
run_name: wan2.1_t2v_anyflow_onpolicy
|
||||
|
||||
model:
|
||||
enable_gradient_checkpointing_type: full
|
||||
|
||||
callbacks:
|
||||
grad_clip:
|
||||
_target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
|
||||
max_grad_norm: 1.0
|
||||
|
||||
pipeline:
|
||||
flow_shift: 5.0
|
||||
dit_config:
|
||||
r_embedder: true
|
||||
r_embedder_fusion: gated
|
||||
r_embedder_gate_value: 0.25
|
||||
r_embedder_deltatime_type: r
|
||||
@@ -1,83 +0,0 @@
|
||||
# AnyFlow pretrain (flow-map central-difference) — Wan 2.1 T2V 1.3B.
|
||||
#
|
||||
# Stage 1 of the AnyFlow two-stage recipe. Trains the dual-timestep
|
||||
# u_θ(x_t, t, r) on the central-difference target so the same checkpoint
|
||||
# can be sampled at arbitrary NFE in the on-policy stage.
|
||||
#
|
||||
# Initialize from base Wan 2.1 T2V 1.3B. No teacher or critic at this
|
||||
# stage; AnyFlowPretrainMethod owns a single student + one optimizer.
|
||||
|
||||
models:
|
||||
student:
|
||||
_target_: fastvideo.train.models.wan.WanModel
|
||||
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
trainable: true
|
||||
|
||||
method:
|
||||
_target_: fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod
|
||||
diffusion_ratio: 0.5
|
||||
consistency_ratio: 0.25
|
||||
epsilon: 5 # finite-difference step in absolute train-timestep units
|
||||
weight_type: beta08 # per-timestep loss weight = t * sqrt(1 - t), renormalized
|
||||
fuse_guidance_scale: 3.0
|
||||
# shift is taken from pipeline.flow_shift below.
|
||||
|
||||
training:
|
||||
distributed:
|
||||
num_gpus: 8
|
||||
sp_size: 1
|
||||
tp_size: 1
|
||||
hsdp_replicate_dim: 1
|
||||
hsdp_shard_dim: 8
|
||||
|
||||
data:
|
||||
data_path: data/preprocessed
|
||||
dataloader_num_workers: 4
|
||||
train_batch_size: 4
|
||||
training_cfg_rate: 0.0
|
||||
seed: 1000
|
||||
num_latent_t: 21
|
||||
num_height: 480
|
||||
num_width: 832
|
||||
num_frames: 81
|
||||
|
||||
optimizer:
|
||||
learning_rate: 5.0e-5
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
lr_scheduler: constant
|
||||
lr_warmup_steps: 0
|
||||
|
||||
loop:
|
||||
max_train_steps: 6000
|
||||
gradient_accumulation_steps: 1
|
||||
|
||||
checkpoint:
|
||||
output_dir: outputs/wan2.1_anyflow_pretrain
|
||||
training_state_checkpointing_steps: 500
|
||||
checkpoints_total_limit: 3
|
||||
resume_from_checkpoint: latest
|
||||
|
||||
tracker:
|
||||
project_name: anyflow-wan
|
||||
run_name: wan2.1_t2v_anyflow_pretrain
|
||||
|
||||
model:
|
||||
enable_gradient_checkpointing_type: full
|
||||
|
||||
callbacks:
|
||||
grad_clip:
|
||||
_target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
|
||||
max_grad_norm: 1.0
|
||||
|
||||
pipeline:
|
||||
flow_shift: 5.0
|
||||
dit_config:
|
||||
# Enable AnyFlow dual-timestep conditioning. The student loads from
|
||||
# base Wan 2.1 — its checkpoint has no delta_embedder weights, so they
|
||||
# get initialized identically to time_embedder via deep-copy in
|
||||
# WanTimeTextImageEmbedding.__init__.
|
||||
r_embedder: true
|
||||
r_embedder_fusion: gated
|
||||
r_embedder_gate_value: 0.25
|
||||
r_embedder_deltatime_type: r
|
||||
@@ -1,27 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxSamplingParam(SamplingParam):
|
||||
|
||||
prompt: str | None = "a photo of a cat"
|
||||
negative_prompt: str = ""
|
||||
|
||||
num_videos_per_prompt: int = 1
|
||||
seed: int = 0
|
||||
|
||||
num_frames: int = 1
|
||||
height: int = 1024
|
||||
width: int = 1024
|
||||
fps: int = 1
|
||||
|
||||
num_inference_steps: int = 28
|
||||
guidance_scale: float = 3.5
|
||||
use_embedded_guidance: bool = True
|
||||
true_cfg_scale: float = 1.0
|
||||
+22
-3
@@ -37,15 +37,34 @@ def parse_config(config_type: type[T], raw: Mapping[str, Any] | T) -> T:
|
||||
if not isinstance(raw, Mapping):
|
||||
raise ConfigValidationError("", f"expected mapping for {config_type.__name__}")
|
||||
parsed = _SchemaParser().parse_dataclass(config_type, raw, "")
|
||||
# None is the schema-wide sentinel for "not specified, inherit the
|
||||
# model preset", so an explicit YAML/JSON null must not bind as an
|
||||
# explicit path — otherwise it would stomp preset values with None.
|
||||
if config_type is GenerationRequest:
|
||||
return bind_generation_request_raw(parsed, raw)
|
||||
return bind_generation_request_raw(parsed, drop_none_leaves(raw))
|
||||
if config_type is RunConfig:
|
||||
return bind_run_config_raw(parsed, raw)
|
||||
return bind_run_config_raw(parsed, drop_none_leaves(raw))
|
||||
if config_type is ServeConfig:
|
||||
return bind_serve_config_raw(parsed, raw)
|
||||
return bind_serve_config_raw(parsed, drop_none_leaves(raw))
|
||||
return parsed
|
||||
|
||||
|
||||
def drop_none_leaves(raw: Any) -> Any:
|
||||
"""Prune None leaves — and dicts emptied by the pruning — from a raw
|
||||
config mapping, so they are never recorded as explicit paths."""
|
||||
if not isinstance(raw, dict):
|
||||
return raw
|
||||
pruned: dict[str, Any] = {}
|
||||
for key, value in raw.items():
|
||||
if value is None:
|
||||
continue
|
||||
value = drop_none_leaves(value)
|
||||
if isinstance(value, dict) and not value:
|
||||
continue
|
||||
pruned[key] = value
|
||||
return pruned
|
||||
|
||||
|
||||
def config_to_dict(config: Any) -> Any:
|
||||
"""Serialize a typed config object into plain Python containers."""
|
||||
if dataclasses.is_dataclass(config) and not isinstance(config, type):
|
||||
|
||||
@@ -90,10 +90,6 @@ class SamplingParam:
|
||||
num_inference_steps_sr: int = 50
|
||||
guidance_scale: float = 1.0
|
||||
guidance_scale_2: float | None = None
|
||||
# Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
|
||||
use_embedded_guidance: bool = False
|
||||
# Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
|
||||
true_cfg_scale: float = 1.0
|
||||
guidance_rescale: float = 0.0
|
||||
boundary_ratio: float | None = None
|
||||
sigmas: list[float] | None = None
|
||||
@@ -329,18 +325,6 @@ class SamplingParam:
|
||||
default=SamplingParam.guidance_rescale,
|
||||
help="Guidance rescale factor",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-embedded-guidance",
|
||||
action="store_true",
|
||||
default=SamplingParam.use_embedded_guidance,
|
||||
help="Use embedded guidance scale (FLUX-style) instead of classic CFG",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--true-cfg-scale",
|
||||
type=float,
|
||||
default=SamplingParam.true_cfg_scale,
|
||||
help="True CFG scale for FLUX when > 1 (requires negative prompt encoding)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--boundary-ratio",
|
||||
type=float,
|
||||
|
||||
+27
-17
@@ -136,21 +136,28 @@ class InputConfig:
|
||||
|
||||
@dataclass
|
||||
class SamplingConfig:
|
||||
num_videos_per_prompt: int = 1
|
||||
seed: int = 1024
|
||||
num_frames: int = 125
|
||||
height: int = 720
|
||||
width: int = 1280
|
||||
height_sr: int = 1072
|
||||
width_sr: int = 1920
|
||||
fps: int = 24
|
||||
num_inference_steps: int = 50
|
||||
num_inference_steps_sr: int = 50
|
||||
guidance_scale: float = 1.0
|
||||
"""Sampling knobs for a generation request.
|
||||
|
||||
``None`` means "not specified": the value is inherited from the
|
||||
model's preset (``SamplingParam.from_pretrained``) at generate time.
|
||||
Set a field to override the preset. There is deliberately no way to
|
||||
express "explicitly the schema default" — pass the concrete value
|
||||
you want instead.
|
||||
"""
|
||||
num_videos_per_prompt: int | None = None
|
||||
seed: int | None = None
|
||||
num_frames: int | None = None
|
||||
height: int | None = None
|
||||
width: int | None = None
|
||||
height_sr: int | None = None
|
||||
width_sr: int | None = None
|
||||
fps: int | None = None
|
||||
num_inference_steps: int | None = None
|
||||
num_inference_steps_sr: int | None = None
|
||||
guidance_scale: float | None = None
|
||||
guidance_scale_2: float | None = None
|
||||
guidance_rescale: float = 0.0
|
||||
guidance_rescale: float | None = None
|
||||
true_cfg_scale: float | None = None
|
||||
use_embedded_guidance: bool | None = None
|
||||
boundary_ratio: float | None = None
|
||||
sigmas: list[float] | None = None
|
||||
|
||||
@@ -195,6 +202,8 @@ class GenerationPlan:
|
||||
class GenerationRequest:
|
||||
prompt: str | list[str] | None = None
|
||||
negative_prompt: str | None = None
|
||||
"""``None`` inherits the model preset's negative prompt; pass ``""``
|
||||
to explicitly clear it."""
|
||||
inputs: InputConfig = field(default_factory=InputConfig)
|
||||
sampling: SamplingConfig = field(default_factory=SamplingConfig)
|
||||
runtime: RequestRuntimeConfig = field(default_factory=RequestRuntimeConfig)
|
||||
@@ -262,15 +271,16 @@ class ServeConfig:
|
||||
the incoming body as the operator-pinned baseline.
|
||||
|
||||
Important nuance: only fields the operator **explicitly wrote** in the
|
||||
serve YAML/JSON count as defaults. Although the in-memory object is
|
||||
fully populated (schema defaults fill every unset field), the merge
|
||||
walks ``_fastvideo_explicit_paths`` — populated during parse — so
|
||||
serve YAML/JSON count as defaults. Unset sampling fields stay ``None``
|
||||
("inherit the model preset"), other sections keep their schema
|
||||
defaults in memory — but the merge walks ``_fastvideo_explicit_paths``
|
||||
(populated during parse; explicit ``null`` is treated as unset), so
|
||||
unset fields are *not* forced onto requests. Per-request precedence:
|
||||
|
||||
body (client-explicit) > default_request (operator-explicit)
|
||||
> hardcoded fallback (e.g. ``fps=24``)
|
||||
|
||||
See :func:`fastvideo.api.compat.explicit_request_updates` for the
|
||||
See :func:`fastvideo.api.translation.explicit_request_updates` for the
|
||||
projection and ``entrypoints/openai/video_api.py::_build_generation_kwargs``
|
||||
for the merge.
|
||||
"""
|
||||
|
||||
@@ -1,4 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Translation layer between the public typed API and the legacy internals.
|
||||
|
||||
This module (formerly ``fastvideo/api/compat.py``) is the single bridge that
|
||||
converts typed public config objects (:class:`GeneratorConfig`,
|
||||
:class:`GenerationRequest`) into the internal :class:`FastVideoArgs` /
|
||||
:class:`SamplingParam` the runtime still consumes, plus the inbound
|
||||
normalization helpers used by every public entrypoint.
|
||||
|
||||
It is *internal* and not part of the stable public surface. The legacy
|
||||
forward-translation helpers (``legacy_from_pretrained_to_config``,
|
||||
``legacy_generate_call_to_request`` and friends) exist only to support the
|
||||
deprecated ``VideoGenerator.from_pretrained(model, **kwargs)`` /
|
||||
``generate_video(...)`` entry points. The reverse-translation
|
||||
(``generator_config_to_fastvideo_args``, ``request_to_sampling_param``) is
|
||||
slated for elimination once ``FastVideoArgs`` becomes a view over
|
||||
``GeneratorConfig`` and ``ForwardBatch`` reads ``GenerationRequest`` directly
|
||||
(see ``.agents/memory/dreamverse-integration/pr-roadmap.md``, PRs 15-17).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
@@ -8,7 +26,12 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.api.overrides import apply_overrides, normalize_overrides
|
||||
from fastvideo.api.parser import config_to_dict, load_raw_config, parse_config
|
||||
from fastvideo.api.parser import (
|
||||
config_to_dict,
|
||||
drop_none_leaves,
|
||||
load_raw_config,
|
||||
parse_config,
|
||||
)
|
||||
from fastvideo.api.request_metadata import (
|
||||
EXPLICIT_PATHS_ATTR,
|
||||
bind_generation_request_raw,
|
||||
@@ -314,8 +337,12 @@ def normalize_generation_request(request: GenerationRequest | Mapping[str, Any],
|
||||
|
||||
if not hasattr(normalized, EXPLICIT_PATHS_ATTR):
|
||||
# Request wasn't bound through the parser (e.g. constructed
|
||||
# directly). Treat every currently-set field as explicit.
|
||||
bind_generation_request_raw(normalized, _serialize_generation_request(normalized))
|
||||
# directly). Bind only non-None leaves as explicit: None is the
|
||||
# schema-wide sentinel for "not specified, inherit the model
|
||||
# preset", so treating it as explicit would stomp preset values
|
||||
# with schema defaults. (parse_config applies the same pruning
|
||||
# to YAML/JSON nulls.)
|
||||
bind_generation_request_raw(normalized, drop_none_leaves(_serialize_generation_request(normalized)))
|
||||
return normalized
|
||||
|
||||
|
||||
@@ -364,15 +391,13 @@ def request_to_sampling_param(
|
||||
updates = explicit_request_updates(request)
|
||||
|
||||
for key, value in updates.items():
|
||||
if key == "return_state":
|
||||
# Already translated to return_continuation_state above.
|
||||
continue
|
||||
if hasattr(sampling_param, key):
|
||||
setattr(sampling_param, key, deepcopy(value))
|
||||
elif key in _REQUEST_PIPELINE_OVERRIDE_FIELDS:
|
||||
continue
|
||||
elif value == _SCHEMA_DEFAULT_UPDATES.get(key, _MISSING):
|
||||
# Schema-default field that isn't on SamplingParam; tolerated
|
||||
# because direct GenerationRequest(...) construction has no
|
||||
# way to distinguish "user set" from "schema default".
|
||||
continue
|
||||
else:
|
||||
raise ValueError(f"Request field {key!r} is not supported by sampling params for {model_path}")
|
||||
|
||||
@@ -476,11 +501,13 @@ def explicit_request_updates(request: GenerationRequest) -> dict[str, Any]:
|
||||
This is what makes ``ServeConfig.default_request`` work as an
|
||||
operator-pinned baseline rather than a full override: a YAML with just
|
||||
``sampling.seed: 42`` yields ``{"seed": 42}``, not the full sampling
|
||||
config with its 15 schema defaults.
|
||||
config. (Sampling fields default to ``None`` = "inherit the model
|
||||
preset", so for directly-constructed requests only non-None leaves
|
||||
are bound as explicit.)
|
||||
|
||||
Precondition: the request must carry ``_fastvideo_explicit_paths`` —
|
||||
populated by :func:`fastvideo.api.parser.parse_config` or
|
||||
:func:`fastvideo.api.compat.normalize_generation_request`. Calling on
|
||||
:func:`fastvideo.api.translation.normalize_generation_request`. Calling on
|
||||
a raw ``GenerationRequest()`` asserts.
|
||||
"""
|
||||
assert hasattr(request,
|
||||
@@ -579,8 +606,6 @@ def _serialize_generation_request(request: GenerationRequest) -> dict[str, Any]:
|
||||
return deepcopy(config_to_dict(request))
|
||||
|
||||
|
||||
_SCHEMA_DEFAULT_UPDATES = _extract_request_updates(config_to_dict(GenerationRequest()))
|
||||
|
||||
_KNOWN_CONTINUATION_KINDS: set[str] = set()
|
||||
|
||||
|
||||
@@ -1,15 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastvideo.attention.utils.flash_attn_default import (
|
||||
fa_version,
|
||||
flash_attn_func_compilable,
|
||||
)
|
||||
|
||||
from fastvideo import envs
|
||||
from fastvideo.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
@@ -19,6 +17,119 @@ from fastvideo.attention.backends.abstract import (
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1, mirroring the
|
||||
# kernel package's FASTVIDEO_VSA_CUTEDSL: its CuTeDSL kernels JIT-compile per
|
||||
# shape family and can fail at runtime on some arch/shape combinations, so it
|
||||
# is never auto-selected just because it is installed. Below sm90 a capability
|
||||
# gate in flash_attn_cute routes to FA2 the calls FA4 cannot serve there:
|
||||
# grad-enabled (its backward asserts sm90+) and GQA (pack_gqa fails CuTeDSL
|
||||
# JIT, observed on sm_89).
|
||||
if envs.FASTVIDEO_FA4:
|
||||
try:
|
||||
from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
|
||||
except ImportError as e:
|
||||
raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
|
||||
"fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
|
||||
"or unset FASTVIDEO_FA4.") from e
|
||||
fa_version = "4"
|
||||
else:
|
||||
try:
|
||||
from flash_attn_interface import flash_attn_func as flash_attn_3_func
|
||||
|
||||
# flash_attn 3 no longer have a different API, see following commit:
|
||||
# https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
|
||||
flash_attn_func = flash_attn_3_func
|
||||
fa_version = "3"
|
||||
except ImportError:
|
||||
from flash_attn import flash_attn_func as flash_attn_2_func
|
||||
flash_attn_func = flash_attn_2_func
|
||||
fa_version = "2"
|
||||
try:
|
||||
if importlib.util.find_spec("flash_attn.cute") is not None:
|
||||
logger.info("flash_attn.cute (FA4) is installed but not enabled; "
|
||||
"set FASTVIDEO_FA4=1 to use it for inference.")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# torch.compile traceability: the FA4/cute path (fa_version=="4") is
|
||||
# already a registered torch.library custom op, so dynamo treats it as a
|
||||
# graph node. The external FA2/FA3 `flash_attn_func` is NOT — dynamo
|
||||
# breaks the graph at the call site (observed: wanvideo.py self-attn,
|
||||
# once per layer every step), which fragments the compiled region and
|
||||
# blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom
|
||||
# op (mirrors the FP4 `flash_attn_cute` template) so it becomes an
|
||||
# opaque-but-traceable node. The kernel still runs eager inside the op
|
||||
# (correct — flash-attn must run eager); only dynamo's treatment of the
|
||||
# boundary changes, so numerics are unchanged (SSIM-gate to confirm).
|
||||
if fa_version in ("2", "3"):
|
||||
_fa_default = flash_attn_func
|
||||
|
||||
# Scope: this op covers exactly the q/k/v + softmax_scale + causal
|
||||
# call shape used by FlashAttentionImpl.forward's default branch
|
||||
# (see `flash_attn_func_compilable(...)` call site below). The
|
||||
# masked/no-pad and varlen / cross-attn paths use different
|
||||
# entry points (`flash_attn_no_pad`, `flash_attn_varlen_*`) which
|
||||
# are intentionally out of scope for this PR — wrapping them is a
|
||||
# natural follow-up. The wrapper's signature is the contract: any
|
||||
# extra kwarg (dropout_p, window_size, alibi_slopes, deterministic,
|
||||
# return_attn_probs, ...) raises TypeError at the call site, so
|
||||
# silent loss of kwargs is not a failure mode.
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_default_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_default_forward(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> torch.Tensor:
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_default_forward")
|
||||
def _flash_attn_default_forward_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> torch.Tensor:
|
||||
del softmax_scale, causal
|
||||
# FA2/FA3 default path: [batch, seqlen_q, nheads, head_dim_v],
|
||||
# same dtype/device as q (head dim taken from v).
|
||||
return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
|
||||
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
# Autograd carve-out. The custom op above registers a forward + fake
|
||||
# kernel but NO backward (register_autograd), so it is opaque to
|
||||
# autograd. Inference runs under no_grad / inference_mode and routes
|
||||
# through the traceable custom op — that is the torch.compile win, and
|
||||
# the only path this PR claims. Training backprops through attention,
|
||||
# so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
|
||||
# (itself an autograd.Function, so backward is correct) at the cost of a
|
||||
# dynamo graph break on the training path — i.e. pre-PR behavior, no
|
||||
# regression. Full autograd parity for the custom op (mirroring the FP4
|
||||
# cute template) is a tracked follow-up.
|
||||
if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
|
||||
elif fa_version == "4":
|
||||
# FA4 path: `flash_attn_func` (from `flash_attn_cute`) goes through a
|
||||
# registered torch.library custom op (with an FA4 backward on sm90+;
|
||||
# grad-enabled and GQA calls below sm90 route to FA2), so a passthrough
|
||||
# is enough — no extra registration needed.
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
else:
|
||||
# Defensive: the probe above only ever sets fa_version to "2", "3",
|
||||
# or "4"; an unexpected value means an import/probe regression and
|
||||
# we want a loud error at import, not a silent NameError later.
|
||||
raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
|
||||
f"'2', '3', or '4' from the import probe above.")
|
||||
|
||||
logger.info("Using FlashAttention-%s backend", fa_version)
|
||||
|
||||
# FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
|
||||
@@ -198,17 +309,9 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
attn_metadata: FlashAttnMetadata,
|
||||
):
|
||||
if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
|
||||
# Route through the *_compilable wrappers so dynamo sees one
|
||||
# traceable node for each masked entry point (the unpad/pad
|
||||
# bookkeeping runs eager inside the custom op). On FA2 these
|
||||
# wrappers go through ops with full register_autograd, so
|
||||
# training also backprops through the op (no graph break on
|
||||
# the training path); on FA3/FA4 they carve out to the
|
||||
# autograd.Function for grad-enabled calls — see
|
||||
# fastvideo/attention/utils/flash_attn_no_pad.py.
|
||||
from fastvideo.attention.utils.flash_attn_no_pad import (
|
||||
flash_attn_no_pad_compilable as flash_attn_no_pad,
|
||||
flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
|
||||
flash_attn_no_pad,
|
||||
flash_attn_varlen_qk_no_pad,
|
||||
)
|
||||
|
||||
attn_mask = attn_metadata.attn_mask
|
||||
|
||||
@@ -1,251 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""torch.compile-traceable wrapper for the FA2/FA3/FA4 default attention path.
|
||||
|
||||
The FA4/cute path (`fa_version == "4"`) is already a registered
|
||||
`torch.library.custom_op` in `fastvideo.attention.utils.flash_attn_cute`, so
|
||||
dynamo treats it as a graph node. The external FA2/FA3 ``flash_attn_func`` is
|
||||
NOT — dynamo breaks the graph at the call site (observed: wanvideo.py
|
||||
self-attn, once per layer every step), which fragments the compiled region
|
||||
and blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom op
|
||||
(mirrors the FP4 `flash_attn_cute` template) so it becomes an
|
||||
opaque-but-traceable node. The kernel still runs eager inside the op
|
||||
(correct — flash-attn must run eager); only dynamo's treatment of the
|
||||
boundary changes, so numerics are unchanged (SSIM-gated).
|
||||
|
||||
Autograd: FA2 has full ``register_autograd`` parity — the custom op's
|
||||
backward calls flash_attn's ``_flash_attn_backward`` directly, so training
|
||||
backprops *through* the op (no graph break on the training path either).
|
||||
FA3 currently keeps the no-backward + carve-out pattern from PR #1373
|
||||
because FA3's private backward signature wants validation on a real Hopper
|
||||
box (gated on Kuan-Hao's Modal FA3 setup PR). Once that lands the FA3 path
|
||||
can mirror FA2.
|
||||
|
||||
Lives in `attention/utils/` (sibling of `flash_attn_cute.py` and
|
||||
`flash_attn_no_pad.py`) so it can be imported by any backend that wants the
|
||||
traceable FA default call without pulling in backend dispatch logic. The
|
||||
backend (`attention/backends/flash_attn.py`) just imports
|
||||
`flash_attn_func_compilable` and `fa_version` from here.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo import envs
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Pick the same backend the rest of FastVideo picked for `flash_attn_func`
|
||||
# (FA4/cute → FA3 → FA2). Mirror the precedence used in
|
||||
# `attention/utils/flash_attn_no_pad.py` so the two probes always agree.
|
||||
#
|
||||
# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1: its CuTeDSL
|
||||
# kernels JIT-compile per shape family and can fail at runtime on some
|
||||
# arch/shape combinations, so it is never auto-selected just because it is
|
||||
# installed.
|
||||
if envs.FASTVIDEO_FA4:
|
||||
try:
|
||||
from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
|
||||
except ImportError as e:
|
||||
raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
|
||||
"fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
|
||||
"or unset FASTVIDEO_FA4.") from e
|
||||
fa_version = "4"
|
||||
else:
|
||||
try:
|
||||
from flash_attn_interface import flash_attn_func as flash_attn_3_func
|
||||
|
||||
# flash_attn 3 no longer has a different API, see following commit:
|
||||
# https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
|
||||
flash_attn_func = flash_attn_3_func
|
||||
fa_version = "3"
|
||||
except ImportError:
|
||||
from flash_attn import flash_attn_func as flash_attn_2_func
|
||||
flash_attn_func = flash_attn_2_func
|
||||
fa_version = "2"
|
||||
try:
|
||||
if importlib.util.find_spec("flash_attn.cute") is not None:
|
||||
logger.info("flash_attn.cute (FA4) is installed but not enabled; "
|
||||
"set FASTVIDEO_FA4=1 to use it for inference.")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if fa_version == "2":
|
||||
# Scope: this op covers exactly the q/k/v + softmax_scale + causal call
|
||||
# shape used by FlashAttentionImpl.forward's default branch (see
|
||||
# `flash_attn_func_compilable(...)` call site in
|
||||
# `attention/backends/flash_attn.py`). The masked/no-pad and varlen /
|
||||
# cross-attn paths use different entry points
|
||||
# (`flash_attn_no_pad`, `flash_attn_varlen_*`) which live in
|
||||
# `attention/utils/flash_attn_no_pad.py`. The wrapper's signature is the
|
||||
# contract: any extra kwarg (dropout_p, window_size, alibi_slopes,
|
||||
# deterministic, return_attn_probs, ...) raises TypeError at the call
|
||||
# site, so silent loss of kwargs is not a failure mode.
|
||||
from flash_attn.flash_attn_interface import _flash_attn_backward as _fa2_backward
|
||||
_fa_default = flash_attn_func
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_default_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_default_forward(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# `return_attn_probs=True` asks FA2 to also return softmax_lse +
|
||||
# S_dmask. We need softmax_lse to feed the backward; S_dmask is the
|
||||
# dropout mask (always None here since dropout_p is fixed at 0).
|
||||
out, softmax_lse, _ = _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal, return_attn_probs=True)
|
||||
return out, softmax_lse
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_default_forward")
|
||||
def _flash_attn_default_forward_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
del softmax_scale, causal
|
||||
# FA2 default path: out = [batch, seqlen_q, nheads, head_dim_v],
|
||||
# softmax_lse = [batch, nheads, seqlen_q], fp32 regardless of q dtype.
|
||||
b, sq, hq = q.shape[0], q.shape[1], q.shape[2]
|
||||
out = q.new_empty(b, sq, hq, v.shape[-1])
|
||||
lse = q.new_empty(b, hq, sq, dtype=torch.float32)
|
||||
return out, lse
|
||||
|
||||
def _flash_attn_default_setup_context(ctx, inputs, output):
|
||||
q, k, v, softmax_scale, causal = inputs
|
||||
out, lse = output
|
||||
ctx.save_for_backward(q, k, v, out, lse)
|
||||
# `lse` is an auxiliary output we save to feed FA2's backward; nobody
|
||||
# should differentiate through it. Mark it non-differentiable so
|
||||
# autograd errors loudly if a caller wires it into a loss, rather
|
||||
# than silently producing zero/None grads through the `del grad_lse`
|
||||
# in our backward.
|
||||
ctx.mark_non_differentiable(lse)
|
||||
# FA2's *forward* substitutes `1 / sqrt(head_dim)` for `softmax_scale=None`
|
||||
# internally; FA2's *backward* (`_flash_attn_backward`) demands a concrete
|
||||
# float in its C++ schema and rejects None at the binding boundary. Resolve
|
||||
# the default here so the value saved on ctx (and passed to backward) is
|
||||
# always a real float — matches what FA2's own autograd.Function does.
|
||||
if softmax_scale is None:
|
||||
softmax_scale = q.shape[-1]**-0.5
|
||||
ctx.softmax_scale = softmax_scale
|
||||
ctx.causal = causal
|
||||
|
||||
def _flash_attn_default_backward(ctx, grad_out, grad_lse):
|
||||
# We only differentiate `out`; softmax_lse is saved-for-backward, not
|
||||
# a real differentiable output. (Mirrors the FP4 cute template.)
|
||||
del grad_lse
|
||||
q, k, v, out, lse = ctx.saved_tensors
|
||||
dq = torch.empty_like(q)
|
||||
dk = torch.empty_like(k)
|
||||
dv = torch.empty_like(v)
|
||||
# FA2's `_flash_attn_backward` writes into dq/dk/dv in place. The
|
||||
# extra kwargs (window_size_*, softcap, alibi_slopes, deterministic,
|
||||
# rng_state) are pinned to the same defaults the forward wrapper
|
||||
# uses — flash-attn==2.8.1 (the version FastVideo pins) requires
|
||||
# all of them explicitly. `rng_state=None` is correct for our
|
||||
# `dropout_p=0` configuration.
|
||||
_fa2_backward(
|
||||
grad_out,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
out,
|
||||
lse,
|
||||
dq,
|
||||
dk,
|
||||
dv,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=ctx.softmax_scale,
|
||||
causal=ctx.causal,
|
||||
window_size_left=-1,
|
||||
window_size_right=-1,
|
||||
softcap=0.0,
|
||||
alibi_slopes=None,
|
||||
deterministic=False,
|
||||
rng_state=None,
|
||||
)
|
||||
return dq, dk, dv, None, None
|
||||
|
||||
torch.library.register_autograd(
|
||||
"fastvideo::_flash_attn_default_forward",
|
||||
_flash_attn_default_backward,
|
||||
setup_context=_flash_attn_default_setup_context,
|
||||
)
|
||||
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
# Backward is registered: autograd flows through the op (training
|
||||
# path is also traceable; no carve-out needed). Public API matches
|
||||
# `flash_attn_func` — returns just `out`; we drop the saved-for-
|
||||
# backward `lse` here so callers see the original single-tensor
|
||||
# contract.
|
||||
out, _ = torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
|
||||
return out
|
||||
elif fa_version == "3":
|
||||
# FA3 path: same forward+fake custom op as the original PR #1373, with
|
||||
# the autograd carve-out kept. The full backward (mirroring the FA2 leg
|
||||
# above) wants a Hopper box for grad-check validation, which we don't
|
||||
# have until Kuan-Hao's Modal FA3 setup PR lands. Until then this keeps
|
||||
# inference traceable + training correct (via the original
|
||||
# autograd.Function path + a pre-PR-style graph break on training).
|
||||
_fa_default = flash_attn_func
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_default_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_default_forward(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> torch.Tensor:
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_default_forward")
|
||||
def _flash_attn_default_forward_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> torch.Tensor:
|
||||
del softmax_scale, causal
|
||||
return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
|
||||
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
# Autograd carve-out. The custom op above registers a forward + fake
|
||||
# kernel but NO backward (register_autograd), so it is opaque to
|
||||
# autograd. Inference runs under no_grad / inference_mode and routes
|
||||
# through the traceable custom op — that is the torch.compile win, and
|
||||
# the only path this PR claims. Training backprops through attention,
|
||||
# so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
|
||||
# (itself an autograd.Function, so backward is correct) at the cost of a
|
||||
# dynamo graph break on the training path — i.e. pre-PR behavior, no
|
||||
# regression. Full autograd parity for the custom op (mirroring the FP4
|
||||
# cute template) is a tracked follow-up.
|
||||
if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
|
||||
elif fa_version == "4":
|
||||
# FA4 path: `flash_attn_func` is already a torch.library custom op
|
||||
# (registered in `fastvideo.attention.utils.flash_attn_cute`), so a
|
||||
# passthrough is enough — no extra registration needed.
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
else:
|
||||
# Defensive: the probe above only ever sets fa_version to "2", "3",
|
||||
# or "4"; an unexpected value means an import/probe regression and
|
||||
# we want a loud error at import, not a silent NameError later.
|
||||
raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
|
||||
f"'2', '3', or '4' from the import probe above.")
|
||||
@@ -24,7 +24,7 @@ from flash_attn.bert_padding import pad_input, unpad_input
|
||||
from fastvideo import envs
|
||||
|
||||
|
||||
def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
|
||||
def _resolve_flash_attn_varlen_func() -> Any:
|
||||
if envs.FASTVIDEO_FA4:
|
||||
# FA4 cute is explicit opt-in (see fastvideo/attention/backends/
|
||||
# flash_attn.py); with FASTVIDEO_FA4=1 an unimportable FA4 build must
|
||||
@@ -39,28 +39,20 @@ def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
|
||||
"fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
|
||||
"or unset FASTVIDEO_FA4.") from e
|
||||
|
||||
return flash_attn_varlen_func_cute, "4"
|
||||
return flash_attn_varlen_func_cute
|
||||
try:
|
||||
from flash_attn_interface import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_interface, )
|
||||
|
||||
return flash_attn_varlen_func_interface, "3"
|
||||
return flash_attn_varlen_func_interface
|
||||
except ImportError:
|
||||
from flash_attn import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_flash, )
|
||||
|
||||
return flash_attn_varlen_func_flash, "2"
|
||||
return flash_attn_varlen_func_flash
|
||||
|
||||
|
||||
flash_attn_varlen_func_impl, _FA_VARLEN_VERSION = _resolve_flash_attn_varlen_func()
|
||||
|
||||
# FA2-only: the private varlen backward we register against the custom ops
|
||||
# below. FA3 / FA4 have different private signatures and validation paths
|
||||
# (Hopper / Blackwell boxes) — those legs keep the autograd carve-out
|
||||
# pattern from PR #1373 until their setup PRs land.
|
||||
if _FA_VARLEN_VERSION == "2":
|
||||
from flash_attn.flash_attn_interface import (
|
||||
_flash_attn_varlen_backward as _fa2_varlen_backward, )
|
||||
flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
|
||||
|
||||
|
||||
def flash_attn_no_pad(
|
||||
@@ -199,473 +191,3 @@ def flash_attn_varlen_qk_no_pad(
|
||||
h=nheads,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# torch.compile traceability + register_autograd parity for the masked /
|
||||
# varlen attention paths.
|
||||
#
|
||||
# Wraps the two entry points `FlashAttentionImpl.forward` calls
|
||||
# (`flash_attn_no_pad`, `flash_attn_varlen_qk_no_pad`) as
|
||||
# `torch.library.custom_op`s so dynamo sees one traceable node — the
|
||||
# internal unpad / pad bookkeeping (data-dependent `nnz` shapes) runs
|
||||
# eager inside the op, and the op's outputs are the statically-shaped
|
||||
# padded tensors. This mirrors the FA2 default-path wrapper in
|
||||
# `fastvideo/attention/backends/flash_attn.py`.
|
||||
#
|
||||
# Autograd: on FA2 we register a real backward (`register_autograd`)
|
||||
# that calls FA2's `_flash_attn_varlen_backward` on the unpadded form
|
||||
# — re-unpadding the saved padded tensors using the saved mask. The
|
||||
# `softmax_lse` from the varlen forward is naturally unpadded
|
||||
# (`[nheads, total_q]`); we pad it to `[batch, nheads, seqlen]` on
|
||||
# the way out (statically shaped) and re-unpad in backward. So
|
||||
# training backprops *through* the op (no graph break on the training
|
||||
# path either).
|
||||
#
|
||||
# FA3 / FA4 keep the autograd carve-out pattern from PR #1373: the
|
||||
# custom op has forward + fake only, and `*_compilable` falls back to
|
||||
# the original autograd.Function for grad-enabled calls. Those legs
|
||||
# are gated on Hopper-class / Blackwell-class boxes for backward
|
||||
# validation and ship as separate follow-ups.
|
||||
|
||||
if _FA_VARLEN_VERSION == "2":
|
||||
# ---------- masked self-attention: flash_attn_no_pad (FA2) ----------
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_no_pad_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_no_pad_forward(
|
||||
qkv: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
b, s, _three, h, d = qkv.shape
|
||||
x = rearrange(qkv, "b s three h d -> b s (three h d)")
|
||||
x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
|
||||
x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
|
||||
out_unpad, lse_unpad, _ = flash_attn_varlen_qkvpacked_func(x_unpad,
|
||||
cu_seqlens,
|
||||
max_s,
|
||||
dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
deterministic=deterministic,
|
||||
return_attn_probs=True)
|
||||
# Pad out: [nnz, h, d] -> [b, s, h, d]
|
||||
out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), indices, b, s),
|
||||
"b s (h d) -> b s h d",
|
||||
h=h)
|
||||
# Pad lse: FA2 varlen returns [nheads, total_q]. Transpose to [total_q,
|
||||
# nheads], pad to [b, s, nheads], permute to [b, nheads, s] — statically
|
||||
# shaped so register_fake matches.
|
||||
lse_padded = pad_input(lse_unpad.t().contiguous(), indices, b, s).permute(0, 2, 1).contiguous()
|
||||
return out_padded, lse_padded
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
|
||||
def _flash_attn_no_pad_forward_fake(
|
||||
qkv: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
|
||||
b, s, _three, h, d = qkv.shape
|
||||
out = qkv.new_empty(b, s, h, d)
|
||||
lse = qkv.new_empty(b, h, s, dtype=torch.float32)
|
||||
return out, lse
|
||||
|
||||
def _flash_attn_no_pad_setup_context(ctx, inputs, output):
|
||||
qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic = inputs
|
||||
out, lse = output
|
||||
ctx.save_for_backward(qkv, out, lse, key_padding_mask)
|
||||
# Auxiliary output, not differentiable — see default-path note.
|
||||
ctx.mark_non_differentiable(lse)
|
||||
# FA2's varlen backward requires a concrete float for softmax_scale.
|
||||
if softmax_scale is None:
|
||||
softmax_scale = qkv.shape[-1]**-0.5 # head_dim from qkv's last dim
|
||||
ctx.softmax_scale = softmax_scale
|
||||
ctx.causal = causal
|
||||
ctx.dropout_p = dropout_p
|
||||
ctx.deterministic = deterministic
|
||||
|
||||
def _flash_attn_no_pad_backward(ctx, grad_out, grad_lse):
|
||||
# lse is saved-for-backward, not differentiated.
|
||||
del grad_lse
|
||||
qkv, out_padded, lse_padded, key_padding_mask = ctx.saved_tensors
|
||||
b, s, _three, h, d = qkv.shape
|
||||
|
||||
# One `unpad_input` call (on qkv) gives us indices + cu_seqlens + max_s;
|
||||
# reuse those for out / dout / lse below via direct indexing instead
|
||||
# of redundant `unpad_input` calls (each of which would re-run
|
||||
# `nonzero` + `cumsum` + a `.max().item()` GPU→CPU sync).
|
||||
x = rearrange(qkv, "b s three h d -> b s (three h d)")
|
||||
x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
|
||||
x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
|
||||
q_unpad, k_unpad, v_unpad = (t.contiguous() for t in x_unpad.unbind(dim=1))
|
||||
|
||||
# Direct-index variants reuse `indices` (computed above).
|
||||
out_unpad = out_padded.flatten(0, 1)[indices].view(-1, h, d).contiguous()
|
||||
dout_unpad = grad_out.flatten(0, 1)[indices].view(-1, h, d).contiguous()
|
||||
# lse_padded [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
|
||||
lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[indices].t().contiguous()
|
||||
|
||||
dq_unpad = torch.empty_like(q_unpad)
|
||||
dk_unpad = torch.empty_like(k_unpad)
|
||||
dv_unpad = torch.empty_like(v_unpad)
|
||||
_fa2_varlen_backward(
|
||||
dout_unpad,
|
||||
q_unpad,
|
||||
k_unpad,
|
||||
v_unpad,
|
||||
out_unpad,
|
||||
lse_unpad,
|
||||
dq_unpad,
|
||||
dk_unpad,
|
||||
dv_unpad,
|
||||
cu_seqlens_q=cu_seqlens,
|
||||
cu_seqlens_k=cu_seqlens,
|
||||
max_seqlen_q=max_s,
|
||||
max_seqlen_k=max_s,
|
||||
dropout_p=ctx.dropout_p,
|
||||
softmax_scale=ctx.softmax_scale,
|
||||
causal=ctx.causal,
|
||||
window_size_left=-1,
|
||||
window_size_right=-1,
|
||||
softcap=0.0,
|
||||
alibi_slopes=None,
|
||||
deterministic=ctx.deterministic,
|
||||
rng_state=None,
|
||||
)
|
||||
|
||||
# Re-pad each grad and stack into dqkv.
|
||||
def _repad(dt_unpad: torch.Tensor) -> torch.Tensor:
|
||||
padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, b, s)
|
||||
return rearrange(padded, "b s (h d) -> b s h d", h=h)
|
||||
|
||||
dqkv = torch.stack([_repad(dq_unpad), _repad(dk_unpad), _repad(dv_unpad)], dim=2)
|
||||
# 6 inputs total: qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic.
|
||||
return dqkv, None, None, None, None, None
|
||||
|
||||
torch.library.register_autograd(
|
||||
"fastvideo::_flash_attn_no_pad_forward",
|
||||
_flash_attn_no_pad_backward,
|
||||
setup_context=_flash_attn_no_pad_setup_context,
|
||||
)
|
||||
|
||||
# ---------- cross-attention: flash_attn_varlen_qk_no_pad (FA2) ----------
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_varlen_qk_no_pad_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_varlen_qk_no_pad_forward(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
query_padding_mask: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
b, sq, h, d = query.shape
|
||||
q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
|
||||
query_padding_mask)
|
||||
k_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
|
||||
key_padding_mask)
|
||||
v_unpad, _, _, _, _ = unpad_input(rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
|
||||
q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
|
||||
k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
|
||||
v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
|
||||
out_unpad, lse_unpad, _ = flash_attn_varlen_func_impl(q_unpad,
|
||||
k_unpad,
|
||||
v_unpad,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
deterministic=deterministic,
|
||||
return_attn_probs=True)
|
||||
# Pad out: [nnz_q, h, d] -> [b, sq, h, d]
|
||||
out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), q_indices, b, sq),
|
||||
"b s (h d) -> b s h d",
|
||||
h=h)
|
||||
# Pad lse: [h, nnz_q] -> [b, h, sq]
|
||||
lse_padded = pad_input(lse_unpad.t().contiguous(), q_indices, b, sq).permute(0, 2, 1).contiguous()
|
||||
return out_padded, lse_padded
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
|
||||
def _flash_attn_varlen_qk_no_pad_forward_fake(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
query_padding_mask: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
del key, query_padding_mask, key_padding_mask
|
||||
del causal, dropout_p, softmax_scale, deterministic
|
||||
b, sq, h, _ = query.shape
|
||||
# `out`'s head_dim comes from value (d_v), matching the real forward's
|
||||
# out_padded ([b, sq, h, d_v]); it can differ from query's d_q.
|
||||
out = query.new_empty(b, sq, h, value.shape[-1])
|
||||
lse = query.new_empty(b, h, sq, dtype=torch.float32)
|
||||
return out, lse
|
||||
|
||||
def _flash_attn_varlen_qk_no_pad_setup_context(ctx, inputs, output):
|
||||
(query, key, value, query_padding_mask, key_padding_mask, causal, dropout_p, softmax_scale,
|
||||
deterministic) = inputs
|
||||
out, lse = output
|
||||
ctx.save_for_backward(query, key, value, out, lse, query_padding_mask, key_padding_mask)
|
||||
# Auxiliary output, not differentiable — see default-path note.
|
||||
ctx.mark_non_differentiable(lse)
|
||||
if softmax_scale is None:
|
||||
softmax_scale = query.shape[-1]**-0.5
|
||||
ctx.softmax_scale = softmax_scale
|
||||
ctx.causal = causal
|
||||
ctx.dropout_p = dropout_p
|
||||
ctx.deterministic = deterministic
|
||||
|
||||
def _flash_attn_varlen_qk_no_pad_backward(ctx, grad_out, grad_lse):
|
||||
del grad_lse
|
||||
(query, key, value, out_padded, lse_padded, query_padding_mask, key_padding_mask) = ctx.saved_tensors
|
||||
b, sq, h, d = query.shape
|
||||
sk = key.shape[1]
|
||||
|
||||
# One `unpad_input` call per distinct mask; reuse the returned
|
||||
# indices via direct indexing for everything else that shares
|
||||
# the same mask (v with k_mask; out/dout/lse with q_mask; the
|
||||
# final repad of dk/dv also reuses k_indices). Avoids ~4
|
||||
# redundant `unpad_input` calls + their GPU→CPU `.max().item()`
|
||||
# syncs.
|
||||
q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
|
||||
query_padding_mask)
|
||||
k_unpad, k_indices, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
|
||||
key_padding_mask)
|
||||
q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
|
||||
k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
|
||||
v_unpad = value.flatten(0, 1)[k_indices].view(-1, h, d).contiguous()
|
||||
|
||||
# out / dout / lse follow q's shape, so index with q_indices.
|
||||
out_unpad = out_padded.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
|
||||
dout_unpad = grad_out.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
|
||||
lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[q_indices].t().contiguous()
|
||||
|
||||
dq_unpad = torch.empty_like(q_unpad)
|
||||
dk_unpad = torch.empty_like(k_unpad)
|
||||
dv_unpad = torch.empty_like(v_unpad)
|
||||
_fa2_varlen_backward(
|
||||
dout_unpad,
|
||||
q_unpad,
|
||||
k_unpad,
|
||||
v_unpad,
|
||||
out_unpad,
|
||||
lse_unpad,
|
||||
dq_unpad,
|
||||
dk_unpad,
|
||||
dv_unpad,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
dropout_p=ctx.dropout_p,
|
||||
softmax_scale=ctx.softmax_scale,
|
||||
causal=ctx.causal,
|
||||
window_size_left=-1,
|
||||
window_size_right=-1,
|
||||
softcap=0.0,
|
||||
alibi_slopes=None,
|
||||
deterministic=ctx.deterministic,
|
||||
rng_state=None,
|
||||
)
|
||||
|
||||
# k_indices is already available from the unpad_input above —
|
||||
# no need to recompute it for the dk/dv repad.
|
||||
def _repad(dt_unpad: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
|
||||
padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, batch, seqlen)
|
||||
return rearrange(padded, "b s (h d) -> b s h d", h=h)
|
||||
|
||||
dq_padded = _repad(dq_unpad, q_indices, b, sq)
|
||||
dk_padded = _repad(dk_unpad, k_indices, b, sk)
|
||||
dv_padded = _repad(dv_unpad, k_indices, b, sk)
|
||||
# 9 inputs total: query, key, value, q_mask, k_mask, causal, dropout_p,
|
||||
# softmax_scale, deterministic.
|
||||
return dq_padded, dk_padded, dv_padded, None, None, None, None, None, None
|
||||
|
||||
torch.library.register_autograd(
|
||||
"fastvideo::_flash_attn_varlen_qk_no_pad_forward",
|
||||
_flash_attn_varlen_qk_no_pad_backward,
|
||||
setup_context=_flash_attn_varlen_qk_no_pad_setup_context,
|
||||
)
|
||||
|
||||
# ---------- public dispatchers (FA2: autograd flows through the op) -----
|
||||
|
||||
|
||||
def flash_attn_no_pad_compilable(qkv,
|
||||
key_padding_mask,
|
||||
causal=False,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
deterministic=False):
|
||||
"""dynamo-traceable wrapper around ``flash_attn_no_pad`` (registered op,
|
||||
full register_autograd on FA2 — both inference and training go through
|
||||
the op, no graph break on either)."""
|
||||
out, _ = torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
|
||||
deterministic)
|
||||
return out
|
||||
|
||||
def flash_attn_varlen_qk_no_pad_compilable(query,
|
||||
key,
|
||||
value,
|
||||
query_padding_mask,
|
||||
key_padding_mask,
|
||||
causal=False,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
deterministic=False):
|
||||
"""dynamo-traceable wrapper around ``flash_attn_varlen_qk_no_pad`` (registered
|
||||
op, full register_autograd on FA2)."""
|
||||
out, _ = torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
|
||||
key_padding_mask, causal, dropout_p,
|
||||
softmax_scale, deterministic)
|
||||
return out
|
||||
|
||||
else:
|
||||
# ---------- FA3 / FA4: carve-out (forward+fake only, no real backward) ---
|
||||
# Same pattern as the parked varlen-extension and the FA3 default leg in
|
||||
# `fastvideo/attention/backends/flash_attn.py`. Real backward for these
|
||||
# versions is a follow-up gated on Hopper / Blackwell box validation.
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_no_pad_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_no_pad_forward(
|
||||
qkv: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> torch.Tensor:
|
||||
return flash_attn_no_pad( # type: ignore[no-untyped-call]
|
||||
qkv,
|
||||
key_padding_mask,
|
||||
causal=causal,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
deterministic=deterministic)
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
|
||||
def _flash_attn_no_pad_forward_fake(
|
||||
qkv: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> torch.Tensor:
|
||||
del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
|
||||
b, s, _three, h, d = qkv.shape
|
||||
return qkv.new_empty(b, s, h, d)
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_varlen_qk_no_pad_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_varlen_qk_no_pad_forward(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
query_padding_mask: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> torch.Tensor:
|
||||
return flash_attn_varlen_qk_no_pad( # type: ignore[no-untyped-call]
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
query_padding_mask=query_padding_mask,
|
||||
key_padding_mask=key_padding_mask,
|
||||
causal=causal,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
deterministic=deterministic)
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
|
||||
def _flash_attn_varlen_qk_no_pad_forward_fake(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
query_padding_mask: torch.Tensor,
|
||||
key_padding_mask: torch.Tensor,
|
||||
causal: bool,
|
||||
dropout_p: float,
|
||||
softmax_scale: float | None,
|
||||
deterministic: bool,
|
||||
) -> torch.Tensor:
|
||||
del key, query_padding_mask, key_padding_mask
|
||||
del causal, dropout_p, softmax_scale, deterministic
|
||||
b, sq, h, _ = query.shape
|
||||
# `out`'s head_dim comes from value (d_v), matching the real forward's
|
||||
# output ([b, sq, h, d_v]); it can differ from query's d_q.
|
||||
return query.new_empty(b, sq, h, value.shape[-1])
|
||||
|
||||
def flash_attn_no_pad_compilable(qkv,
|
||||
key_padding_mask,
|
||||
causal=False,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
deterministic=False):
|
||||
if torch.is_grad_enabled() and qkv.requires_grad:
|
||||
return flash_attn_no_pad(qkv,
|
||||
key_padding_mask,
|
||||
causal=causal,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
deterministic=deterministic)
|
||||
return torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
|
||||
deterministic)
|
||||
|
||||
def flash_attn_varlen_qk_no_pad_compilable(query,
|
||||
key,
|
||||
value,
|
||||
query_padding_mask,
|
||||
key_padding_mask,
|
||||
causal=False,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
deterministic=False):
|
||||
if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
|
||||
return flash_attn_varlen_qk_no_pad(query,
|
||||
key,
|
||||
value,
|
||||
query_padding_mask=query_padding_mask,
|
||||
key_padding_mask=key_padding_mask,
|
||||
causal=causal,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
deterministic=deterministic)
|
||||
return torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
|
||||
key_padding_mask, causal, dropout_p,
|
||||
softmax_scale, deterministic)
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
|
||||
from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
|
||||
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
|
||||
from fastvideo.configs.models.dits.flux import FluxDiTConfig
|
||||
from fastvideo.configs.models.dits.flux_2 import Flux2Config
|
||||
from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
|
||||
from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
|
||||
from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
|
||||
from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
|
||||
@@ -17,7 +15,6 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
|
||||
|
||||
__all__ = [
|
||||
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
|
||||
"DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
|
||||
"LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
|
||||
"StableAudioConfig", "GlmImageDiTConfig"
|
||||
"DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
|
||||
"HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
|
||||
]
|
||||
|
||||
@@ -1,119 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Cosmos3 VFM Transformer FastVideo dataclass configs.
|
||||
|
||||
Architecture is 1:1 with the published ``nvidia/Cosmos3-Nano`` checkpoint
|
||||
(``transformer/config.json``; class ``Cosmos3OmniTransformer`` / framework
|
||||
``Cosmos3VFMNetwork``). Field values match that config so the FastVideo native
|
||||
DiT builds a parameter tree matching the checkpoint's state-dict surface
|
||||
(814 tensors / 44 patterns, validated 2026-06-06).
|
||||
|
||||
Reference of record: ``cosmos-framework`` (NVIDIA). The checkpoint is a single
|
||||
``layers`` ModuleList of dual-pathway (understanding/text + generation/vision)
|
||||
decoder blocks; per layer: ``self_attn`` with und (``to_{q,k,v}``/``to_out``)
|
||||
and gen (``add_{q,k,v}_proj``/``to_add_out``) projections + QK-norms, plus
|
||||
``mlp`` (und) and ``mlp_moe_gen`` (gen), and four RMSNorms. Top level adds
|
||||
``embed_tokens``/``norm``/``norm_moe_gen``/``lm_head``/``proj_in``/``proj_out``/
|
||||
``time_embedder`` and dormant ``action_*``/``audio_*`` heads. The checkpoint
|
||||
remap lives in ``scripts/checkpoint_conversion/cosmos3_convert.py``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
|
||||
|
||||
def is_cosmos3_transformer_block(name: str, module) -> bool:
|
||||
"""FSDP shard boundary: the dual-pathway decoder blocks ``layers.{i}``."""
|
||||
del module
|
||||
parts = name.split(".")
|
||||
return "layers" in parts and parts[-1].isdigit()
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cosmos3ArchConfig(DiTArchConfig):
|
||||
"""Architecture config for the Cosmos3 omni DiT (Cosmos3-Nano).
|
||||
|
||||
1:1 with ``transformer/config.json``. The action/sound heads ship in the
|
||||
checkpoint, so they are constructed for strict-load parity even though the
|
||||
PR1 video path (T2V/I2V/T2I) leaves them dormant.
|
||||
"""
|
||||
|
||||
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_cosmos3_transformer_block])
|
||||
|
||||
# Conversion is owned by scripts/checkpoint_conversion/cosmos3_convert.py;
|
||||
# the native module tree is the source of truth for parameter names.
|
||||
param_names_mapping: dict = field(default_factory=dict)
|
||||
|
||||
# ---- Backbone (Qwen3-VL-text) ----
|
||||
hidden_size: int = 4096
|
||||
num_hidden_layers: int = 36
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 8 # GQA (4 query groups)
|
||||
head_dim: int = 128
|
||||
intermediate_size: int = 12288
|
||||
hidden_act: str = "silu"
|
||||
vocab_size: int = 151936
|
||||
rms_norm_eps: float = 1e-6
|
||||
attention_bias: bool = False
|
||||
qk_norm_for_diffusion: bool = True
|
||||
qk_norm_for_text: bool = True
|
||||
use_moe: bool = True # dual-pathway weights; sparse routing unused
|
||||
joint_attn_implementation: str = "two_way"
|
||||
freeze_und: bool = False
|
||||
|
||||
# ---- Position embedding (unified 3D MRoPE) ----
|
||||
position_embedding_type: str = "unified_3d_mrope"
|
||||
rope_theta: float = 5_000_000.0
|
||||
max_position_embeddings: int = 262144
|
||||
mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20])
|
||||
mrope_interleaved: bool = True
|
||||
unified_3d_mrope_reset_spatial_ids: bool = True
|
||||
temporal_modality_margin: int = 15000 # unified_3d_mrope_temporal_modality_margin
|
||||
|
||||
# ---- VAE / patch geometry ----
|
||||
latent_patch_size: int = 2
|
||||
latent_channel: int = 48
|
||||
patch_latent_dim: int = 192 # latent_patch_size**2 * latent_channel
|
||||
|
||||
# ---- Diffusion conditioning ----
|
||||
timestep_scale: float = 0.001
|
||||
|
||||
# ---- Temporal / FPS modulation ----
|
||||
base_fps: float = 24.0
|
||||
temporal_compression_factor: int = 4
|
||||
enable_fps_modulation: bool = True
|
||||
video_temporal_causal: bool = False
|
||||
|
||||
# ---- Action generation head (dormant in PR1 video path) ----
|
||||
action_gen: bool = True
|
||||
action_dim: int = 64
|
||||
max_action_dim: int = 64
|
||||
num_embodiment_domains: int = 32
|
||||
|
||||
# ---- Sound generation head (dormant in PR1 video path) ----
|
||||
sound_gen: bool = True
|
||||
sound_dim: int = 64
|
||||
sound_latent_fps: float = 25.0
|
||||
temporal_compression_factor_sound: int = 1
|
||||
|
||||
# ---- BaseDiT bookkeeping ----
|
||||
in_channels: int = 48
|
||||
out_channels: int = 48
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
# Video DiT contract: latent channels == VAE z_dim.
|
||||
self.num_channels_latents = self.latent_channel
|
||||
if not self.out_channels:
|
||||
self.out_channels = self.in_channels
|
||||
# Derived: patchify packs latent_patch_size**2 spatial patches * channels.
|
||||
self.patch_latent_dim = self.latent_patch_size**2 * self.latent_channel
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cosmos3VideoConfig(DiTConfig):
|
||||
"""Pipeline-level Cosmos3 DiT config (T2V / I2V / T2I share this surface)."""
|
||||
|
||||
arch_config: DiTArchConfig = field(default_factory=Cosmos3ArchConfig)
|
||||
prefix: str = "Cosmos3"
|
||||
@@ -1,27 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxTransformer2DArchConfig(DiTArchConfig):
|
||||
|
||||
patch_size: int = 1
|
||||
in_channels: int = 64
|
||||
out_channels: int | None = None
|
||||
num_layers: int = 19
|
||||
num_single_layers: int = 38
|
||||
attention_head_dim: int = 128
|
||||
num_attention_heads: int = 24
|
||||
joint_attention_dim: int = 4096
|
||||
pooled_projection_dim: int = 768
|
||||
guidance_embeds: bool = True
|
||||
axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxDiTConfig(DiTConfig):
|
||||
arch_config: DiTArchConfig = field(default_factory=FluxTransformer2DArchConfig)
|
||||
prefix: str = "flux"
|
||||
@@ -1,61 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
|
||||
|
||||
def is_blocks(n: str, m) -> bool:
|
||||
return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageDiTArchConfig(DiTArchConfig):
|
||||
|
||||
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
|
||||
|
||||
hidden_size: int = 4096
|
||||
num_attention_heads: int = 32
|
||||
attention_head_dim: int = 128
|
||||
in_channels: int = 16
|
||||
out_channels: int = 16
|
||||
num_layers: int = 30
|
||||
|
||||
text_embed_dim: int = 1472
|
||||
time_embed_dim: int = 512
|
||||
condition_dim: int = 256
|
||||
|
||||
prior_vq_quantizer_codebook_size: int = 16384
|
||||
|
||||
patch_size: int = 2
|
||||
|
||||
max_height: int = 2048
|
||||
max_width: int = 2048
|
||||
|
||||
qk_norm: str = "layer_norm"
|
||||
eps: float = 1e-5
|
||||
|
||||
exclude_lora_layers: list[str] = field(
|
||||
default_factory=lambda: ["image_projector", "glyph_projector", "prior_token_embedding"])
|
||||
|
||||
param_names_mapping: dict = field(
|
||||
default_factory=lambda: {
|
||||
r"^glyph_projector\.net\.0\.proj\.(.*)$": r"glyph_projector.fc_in.\1",
|
||||
r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1",
|
||||
r"^prior_projector\.net\.0\.proj\.(.*)$": r"prior_projector.fc_in.\1",
|
||||
r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1",
|
||||
r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
|
||||
r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
|
||||
})
|
||||
|
||||
reverse_param_names_mapping: dict = field(default_factory=dict)
|
||||
lora_param_names_mapping: dict = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.num_channels_latents = self.out_channels
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageDiTConfig(DiTConfig):
|
||||
arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig)
|
||||
prefix: str = "GlmImage"
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
|
||||
@@ -20,11 +19,6 @@ class WanVideoArchConfig(DiTArchConfig):
|
||||
r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
|
||||
r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
|
||||
r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
|
||||
# AnyFlow dual-timestep checkpoints expose delta_embedder weights with the
|
||||
# same internal layout as time_embedder. The regex is harmless on plain
|
||||
# Wan checkpoints (no delta_embedder keys to match).
|
||||
r"^condition_embedder\.delta_embedder\.linear_1\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_in.\1",
|
||||
r"^condition_embedder\.delta_embedder\.linear_2\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_out.\1",
|
||||
r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
|
||||
r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
|
||||
r"condition_embedder.image_embedder.ff.fc_in.\1",
|
||||
@@ -92,14 +86,6 @@ class WanVideoArchConfig(DiTArchConfig):
|
||||
# "relativistic" keeps long rollouts in-distribution; a no-op unless sink_size > 0 and local_attn_size > 0.
|
||||
rope_cache_policy: str = "absolute"
|
||||
|
||||
# AnyFlow dual-timestep conditioning. Defaults preserve bit-identity with
|
||||
# the legacy single-timestep forward (no delta_embedder allocated, no
|
||||
# extra computation on the embedder forward path).
|
||||
r_embedder: bool = False
|
||||
r_embedder_fusion: Literal["additive", "gated"] = "additive"
|
||||
r_embedder_gate_value: float = 0.25
|
||||
r_embedder_deltatime_type: Literal["r", "t-r"] = "r"
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.out_channels = self.out_channels or self.in_channels
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
|
||||
from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
|
||||
from fastvideo.configs.models.vaes.cosmos3vae import Cosmos3VAEConfig
|
||||
from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
|
||||
from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
|
||||
from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
|
||||
from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
|
||||
from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
|
||||
from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
|
||||
@@ -17,12 +15,10 @@ __all__ = [
|
||||
"WanVAEConfig",
|
||||
"CosmosVAEConfig",
|
||||
"Cosmos25VAEConfig",
|
||||
"Cosmos3VAEConfig",
|
||||
"Gen3CVAEConfig",
|
||||
"Hunyuan15VAEConfig",
|
||||
"LTX2VAEConfig",
|
||||
"OobleckVAEArchConfig",
|
||||
"OobleckVAEConfig",
|
||||
"Flux2VAEConfig",
|
||||
"GlmImageVAEConfig",
|
||||
]
|
||||
|
||||
@@ -1,277 +0,0 @@
|
||||
"""Cosmos3 (Wan2.2-TI2V-5B) VAE config and checkpoint-key mapping.
|
||||
|
||||
The Cosmos3 checkpoint VAE is literally ``Wan-AI/Wan2.2-TI2V-5B-Diffusers``
|
||||
(diffusers ``AutoencoderKLWan``), so this config locks the Wan2.2 geometry:
|
||||
residual down/up blocks, ``patch_size=2``, ``z_dim=48``, ``base_dim=160``,
|
||||
``decoder_base_dim=256``, and ``scale_factor_spatial=16``. The 48-dim
|
||||
``latents_mean``/``latents_std`` are taken verbatim from the Cosmos3
|
||||
checkpoint's ``vae/config.json`` (identical to the canonical Wan2.2-TI2V-5B
|
||||
statistics).
|
||||
|
||||
Mirrors the :class:`Cosmos25VAEArchConfig` pattern. ``param_names_mapping`` /
|
||||
``map_official_key`` translate the *official* Wan2.2 VAE state-dict keys
|
||||
(nested-residual naming, e.g. ``encoder.downsamples.{b}.downsamples.{j}`` and
|
||||
``decoder.upsamples.{b}.upsamples.{j}``) into FastVideo's ``AutoencoderKLWan``
|
||||
key space. The standard diffusers checkpoint already ships native FastVideo
|
||||
keys, so these helpers exist for parity tooling and official ``.pth`` loading.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig, WanVAEConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cosmos3VAEArchConfig(WanVAEArchConfig):
|
||||
# Wan2.2-TI2V-5B geometry (differs from the Wan2.1 WanVAEArchConfig
|
||||
# defaults: residual blocks, patch_size=2, z_dim=48, base_dim=160,
|
||||
# decoder_base_dim=256, scale_factor_spatial=16, 12 patch channels).
|
||||
_name_or_path: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
base_dim: int = 160
|
||||
decoder_base_dim: int | None = 256
|
||||
z_dim: int = 48
|
||||
dim_mult: tuple[int, ...] = (1, 2, 4, 4)
|
||||
num_res_blocks: int = 2
|
||||
attn_scales: tuple[float, ...] = ()
|
||||
temperal_downsample: tuple[bool, ...] = (False, True, True)
|
||||
dropout: float = 0.0
|
||||
is_residual: bool = True
|
||||
in_channels: int = 12
|
||||
out_channels: int = 12
|
||||
patch_size: int | None = 2
|
||||
scale_factor_temporal: int = 4
|
||||
scale_factor_spatial: int = 16
|
||||
clip_output: bool = False
|
||||
|
||||
# 48-dim statistics copied verbatim from the Cosmos3 checkpoint
|
||||
# (official_weights/cosmos3/vae/config.json).
|
||||
latents_mean: tuple[float, ...] = (
|
||||
-0.2289,
|
||||
-0.0052,
|
||||
-0.1323,
|
||||
-0.2339,
|
||||
-0.2799,
|
||||
0.0174,
|
||||
0.1838,
|
||||
0.1557,
|
||||
-0.1382,
|
||||
0.0542,
|
||||
0.2813,
|
||||
0.0891,
|
||||
0.157,
|
||||
-0.0098,
|
||||
0.0375,
|
||||
-0.1825,
|
||||
-0.2246,
|
||||
-0.1207,
|
||||
-0.0698,
|
||||
0.5109,
|
||||
0.2665,
|
||||
-0.2108,
|
||||
-0.2158,
|
||||
0.2502,
|
||||
-0.2055,
|
||||
-0.0322,
|
||||
0.1109,
|
||||
0.1567,
|
||||
-0.0729,
|
||||
0.0899,
|
||||
-0.2799,
|
||||
-0.123,
|
||||
-0.0313,
|
||||
-0.1649,
|
||||
0.0117,
|
||||
0.0723,
|
||||
-0.2839,
|
||||
-0.2083,
|
||||
-0.052,
|
||||
0.3748,
|
||||
0.0152,
|
||||
0.1957,
|
||||
0.1433,
|
||||
-0.2944,
|
||||
0.3573,
|
||||
-0.0548,
|
||||
-0.1681,
|
||||
-0.0667,
|
||||
)
|
||||
latents_std: tuple[float, ...] = (
|
||||
0.4765,
|
||||
1.0364,
|
||||
0.4514,
|
||||
1.1677,
|
||||
0.5313,
|
||||
0.499,
|
||||
0.4818,
|
||||
0.5013,
|
||||
0.8158,
|
||||
1.0344,
|
||||
0.5894,
|
||||
1.0901,
|
||||
0.6885,
|
||||
0.6165,
|
||||
0.8454,
|
||||
0.4978,
|
||||
0.5759,
|
||||
0.3523,
|
||||
0.7135,
|
||||
0.6804,
|
||||
0.5833,
|
||||
1.4146,
|
||||
0.8986,
|
||||
0.5659,
|
||||
0.7069,
|
||||
0.5338,
|
||||
0.4889,
|
||||
0.4917,
|
||||
0.4069,
|
||||
0.4999,
|
||||
0.6866,
|
||||
0.4093,
|
||||
0.5709,
|
||||
0.6065,
|
||||
0.6415,
|
||||
0.4944,
|
||||
0.5726,
|
||||
1.2042,
|
||||
0.5458,
|
||||
1.6887,
|
||||
0.3971,
|
||||
1.06,
|
||||
0.3943,
|
||||
0.5537,
|
||||
0.5444,
|
||||
0.4089,
|
||||
0.7468,
|
||||
0.7744,
|
||||
)
|
||||
|
||||
# Simple 1:1 renames. The nested-residual block remapping (encoder
|
||||
# downsamples / decoder upsamples / middle / head) is handled by
|
||||
# ``map_official_key()``.
|
||||
param_names_mapping: dict[str, str] = field(
|
||||
default_factory=lambda: {
|
||||
r"^conv1\.(.*)$": r"quant_conv.\1",
|
||||
r"^conv2\.(.*)$": r"post_quant_conv.\1",
|
||||
r"^encoder\.conv1\.(.*)$": r"encoder.conv_in.\1",
|
||||
r"^decoder\.conv1\.(.*)$": r"decoder.conv_in.\1",
|
||||
r"^encoder\.head\.0\.gamma$": r"encoder.norm_out.gamma",
|
||||
r"^encoder\.head\.2\.(.*)$": r"encoder.conv_out.\1",
|
||||
r"^decoder\.head\.0\.gamma$": r"decoder.norm_out.gamma",
|
||||
r"^decoder\.head\.2\.(.*)$": r"decoder.conv_out.\1",
|
||||
})
|
||||
|
||||
@staticmethod
|
||||
def map_official_key(key: str) -> str | None:
|
||||
"""Map a single official Wan2.2 VAE key into FastVideo key space.
|
||||
|
||||
Handles the residual (Wan2.2) module layout where each down/up block
|
||||
is a nested ``Sequential`` (``downsamples.{b}.downsamples.{j}`` /
|
||||
``upsamples.{b}.upsamples.{j}``) rather than the flat Wan2.1 indexing.
|
||||
Returns ``None`` for keys with no FastVideo counterpart.
|
||||
"""
|
||||
|
||||
def map_residual_subkey(prefix: str, sub: str) -> str | None:
|
||||
if re.match(r"^residual\.0\.gamma$", sub):
|
||||
return f"{prefix}.norm1.gamma"
|
||||
m = re.match(r"^residual\.2\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.conv1.{m.group(1)}"
|
||||
if re.match(r"^residual\.3\.gamma$", sub):
|
||||
return f"{prefix}.norm2.gamma"
|
||||
m = re.match(r"^residual\.6\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.conv2.{m.group(1)}"
|
||||
m = re.match(r"^shortcut\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.conv_shortcut.{m.group(1)}"
|
||||
return None
|
||||
|
||||
def map_attn_subkey(prefix: str, sub: str) -> str | None:
|
||||
if re.match(r"^norm\.gamma$", sub):
|
||||
return f"{prefix}.norm.gamma"
|
||||
m = re.match(r"^to_qkv\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.to_qkv.{m.group(1)}"
|
||||
m = re.match(r"^proj\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.proj.{m.group(1)}"
|
||||
return None
|
||||
|
||||
def map_resample_subkey(prefix: str, sub: str) -> str | None:
|
||||
m = re.match(r"^resample\.1\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.resample.1.{m.group(1)}"
|
||||
m = re.match(r"^time_conv\.(weight|bias)$", sub)
|
||||
if m:
|
||||
return f"{prefix}.time_conv.{m.group(1)}"
|
||||
return None
|
||||
|
||||
m = re.match(r"^conv1\.(weight|bias)$", key)
|
||||
if m:
|
||||
return f"quant_conv.{m.group(1)}"
|
||||
m = re.match(r"^conv2\.(weight|bias)$", key)
|
||||
if m:
|
||||
return f"post_quant_conv.{m.group(1)}"
|
||||
m = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key)
|
||||
if m:
|
||||
return f"{m.group(1)}.conv_in.{m.group(2)}"
|
||||
m = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key)
|
||||
if m:
|
||||
return f"{m.group(1)}.norm_out.gamma"
|
||||
m = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key)
|
||||
if m:
|
||||
return f"{m.group(1)}.conv_out.{m.group(2)}"
|
||||
|
||||
m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
|
||||
if m:
|
||||
return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0", m.group(2))
|
||||
m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
|
||||
if m:
|
||||
return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0", m.group(2))
|
||||
m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
|
||||
if m:
|
||||
return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1", m.group(2))
|
||||
|
||||
# Encoder: downsamples.{block}.downsamples.{j}.* (nested residual layout)
|
||||
m = re.match(r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.(.*)$", key)
|
||||
if m:
|
||||
block_i, res_i, sub = int(m.group(1)), int(m.group(2)), m.group(3)
|
||||
if sub.startswith("resample.") or sub.startswith("time_conv."):
|
||||
return map_resample_subkey(f"encoder.down_blocks.{block_i}.downsampler", sub)
|
||||
return map_residual_subkey(f"encoder.down_blocks.{block_i}.resnets.{res_i}", sub)
|
||||
|
||||
# Decoder: upsamples.{block}.upsamples.{j}.* (nested residual layout)
|
||||
m = re.match(r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.(.*)$", key)
|
||||
if m:
|
||||
block_i, res_i, sub = int(m.group(1)), int(m.group(2)), m.group(3)
|
||||
if sub.startswith("resample.") or sub.startswith("time_conv."):
|
||||
return map_resample_subkey(f"decoder.up_blocks.{block_i}.upsampler", sub)
|
||||
return map_residual_subkey(f"decoder.up_blocks.{block_i}.resnets.{res_i}", sub)
|
||||
|
||||
return None
|
||||
|
||||
# ``__post_init__`` (scaling_factor / shift_factor / compression ratios) is
|
||||
# inherited unchanged from ``WanVAEArchConfig``.
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cosmos3VAEConfig(WanVAEConfig):
|
||||
"""Cosmos3 VAE config (reuses FastVideo's Wan2.2 ``AutoencoderKLWan``).
|
||||
|
||||
Subclasses :class:`WanVAEConfig` so the model reads the same runtime flags
|
||||
(``use_feature_cache``, ``use_light_vae``, tiling) and only swaps in the
|
||||
Cosmos3 = Wan2.2 ``arch_config``.
|
||||
"""
|
||||
|
||||
arch_config: Cosmos3VAEArchConfig = field(default_factory=Cosmos3VAEArchConfig)
|
||||
|
||||
use_feature_cache: bool = True
|
||||
use_tiling: bool = False
|
||||
use_temporal_tiling: bool = False
|
||||
use_parallel_tiling: bool = False
|
||||
|
||||
# ``__post_init__`` (blend_num_frames) is inherited from ``WanVAEConfig``.
|
||||
@@ -1,95 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.vaes.autoencoder_kl import (AutoencoderKLArchConfig, AutoencoderKLVAEConfig)
|
||||
|
||||
_GLM_IMAGE_LATENTS_MEAN: tuple[float, ...] = (
|
||||
-0.2080078125,
|
||||
1.875,
|
||||
-0.470703125,
|
||||
-1.265625,
|
||||
-1.421875,
|
||||
0.77734375,
|
||||
-0.3671875,
|
||||
-0.9453125,
|
||||
0.318359375,
|
||||
0.7734375,
|
||||
-0.1884765625,
|
||||
-0.022216796875,
|
||||
-0.220703125,
|
||||
-1.59375,
|
||||
-0.81640625,
|
||||
-0.255859375,
|
||||
)
|
||||
_GLM_IMAGE_LATENTS_STD: tuple[float, ...] = (
|
||||
3.0625,
|
||||
2.203125,
|
||||
2.265625,
|
||||
4.84375,
|
||||
2.5,
|
||||
3.9375,
|
||||
2.203125,
|
||||
3.03125,
|
||||
2.1875,
|
||||
2.046875,
|
||||
2.71875,
|
||||
2.390625,
|
||||
2.390625,
|
||||
2.453125,
|
||||
2.25,
|
||||
2.15625,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageVAEArchConfig(AutoencoderKLArchConfig):
|
||||
act_fn: str = "silu"
|
||||
block_out_channels: tuple[int, ...] = (128, 512, 1024, 1024)
|
||||
down_block_types: tuple[str, ...] = (
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
)
|
||||
up_block_types: tuple[str, ...] = (
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
)
|
||||
force_upcast: bool = True
|
||||
in_channels: int = 3
|
||||
out_channels: int = 3
|
||||
latent_channels: int = 16
|
||||
latents_mean: tuple[float, ...] = _GLM_IMAGE_LATENTS_MEAN
|
||||
latents_std: tuple[float, ...] = _GLM_IMAGE_LATENTS_STD
|
||||
layers_per_block: int = 3
|
||||
mid_block_add_attention: bool = False
|
||||
norm_num_groups: int = 32
|
||||
sample_size: int = 1024
|
||||
scaling_factor: float = 0.18215
|
||||
shift_factor: float | None = None
|
||||
use_quant_conv: bool = False
|
||||
use_post_quant_conv: bool = False
|
||||
|
||||
temporal_compression_ratio: int = 1
|
||||
spatial_compression_ratio: int = 8
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageVAEConfig(AutoencoderKLVAEConfig):
|
||||
arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig)
|
||||
|
||||
use_tiling: bool = True
|
||||
use_temporal_tiling: bool = False
|
||||
use_parallel_tiling: bool = False
|
||||
|
||||
tile_sample_min_height: int = 512
|
||||
tile_sample_min_width: int = 512
|
||||
tile_sample_stride_height: int = 384
|
||||
tile_sample_stride_width: int = 384
|
||||
|
||||
load_encoder: bool = True
|
||||
load_decoder: bool = True
|
||||
@@ -1,69 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Cosmos3 pipeline configuration.
|
||||
|
||||
Reference of record: the official ``cosmos-framework`` / ``nvidia/Cosmos3-Nano``
|
||||
checkpoint (``model_index.json``). Cosmos3 is structurally different from
|
||||
Cosmos 2.5:
|
||||
|
||||
- Dual-pathway (UND + GEN) DiT lives entirely inside ``Cosmos3VFMTransformer``
|
||||
(``Cosmos3VideoConfig``).
|
||||
- No separate text encoder — the Qwen3-VL-text backbone is inside the DiT, so
|
||||
``text_encoder_configs`` is the empty tuple. The Qwen2 tokenizer is loaded as
|
||||
the ``text_tokenizer`` checkpoint module by the component loader.
|
||||
- VAE is Wan2.2 ``AutoencoderKLWan`` (z_dim=48, scale_factor_spatial=16),
|
||||
configured by ``Cosmos3VAEConfig`` (the checkpoint's exact latents_mean/std).
|
||||
- Scheduler is FastVideo-native ``UniPCMultistepScheduler`` configured for
|
||||
pure flow matching (flow_prediction, use_flow_sigmas), equivalent to the
|
||||
framework's ``FlowUniPCMultistepScheduler``. The checkpoint's diffusers-style
|
||||
scheduler config (karras/sigma_min/max) is coerced to the flow setup in
|
||||
``Cosmos3OmniDiffusersPipeline.initialize_pipeline``.
|
||||
- T2I default ``flow_shift`` is 3.0 (set per-request by ``_set_flow_shift``);
|
||||
T2V/I2V use the engine-init default of 1.0 baked into this config.
|
||||
"""
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
|
||||
from fastvideo.configs.models.dits.cosmos3 import (Cosmos3ArchConfig, Cosmos3VideoConfig)
|
||||
from fastvideo.configs.models.encoders import BaseEncoderOutput
|
||||
from fastvideo.configs.models.vaes import Cosmos3VAEConfig # Wan2.2 AutoencoderKLWan
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cosmos3Config(PipelineConfig):
|
||||
"""Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I).
|
||||
|
||||
Wires the framework-parity-verified Cosmos3 components: the native
|
||||
``Cosmos3VideoConfig`` DiT, the Wan2.2 ``Cosmos3VAEConfig`` VAE, the Qwen2
|
||||
tokenizer (loaded as ``text_tokenizer``), and the UniPC scheduler.
|
||||
"""
|
||||
|
||||
dit_config: DiTConfig = field(default_factory=lambda: Cosmos3VideoConfig(arch_config=Cosmos3ArchConfig()))
|
||||
|
||||
vae_config: VAEConfig = field(default_factory=Cosmos3VAEConfig)
|
||||
|
||||
# No separate text encoder: the Qwen3-VL-text backbone lives inside the DiT
|
||||
# and the pipeline tokenizes in Cosmos3DenoisingStage, so all three
|
||||
# text-encoder lists are empty (the generic text-encode stage is not used).
|
||||
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple)
|
||||
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=tuple)
|
||||
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=tuple)
|
||||
|
||||
dit_precision: str = "bf16"
|
||||
vae_precision: str = "bf16"
|
||||
text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple)
|
||||
|
||||
embedded_cfg_scale: float = 0.0
|
||||
# T2V/I2V engine-init flow_shift (framework text2video/image2video default);
|
||||
# T2I overrides to 3.0 per request via Cosmos3DenoisingStage._set_flow_shift.
|
||||
flow_shift: float = 10.0
|
||||
|
||||
vae_tiling: bool = False
|
||||
vae_sp: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
self.vae_config.load_encoder = True
|
||||
self.vae_config.load_decoder = True
|
||||
@@ -1,74 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models import EncoderConfig
|
||||
from fastvideo.configs.models.dits.flux import FluxDiTConfig
|
||||
from fastvideo.configs.models.encoders import (
|
||||
BaseEncoderOutput,
|
||||
CLIPTextConfig,
|
||||
T5LargeConfig,
|
||||
)
|
||||
from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
|
||||
|
||||
|
||||
def _flux_clip_pooled_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
|
||||
"""CLIP branch for FLUX: Diffusers uses pooled prompt embeddings only."""
|
||||
if outputs.pooler_output is None:
|
||||
raise RuntimeError(
|
||||
"FLUX CLIP conditioning requires pooler_output. Ensure the CLIP text encoder returns pooled features.")
|
||||
return outputs.pooler_output
|
||||
|
||||
|
||||
def _flux_t5_sequence_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
|
||||
if outputs.last_hidden_state is None:
|
||||
raise RuntimeError("FLUX T5 conditioning requires last_hidden_state.")
|
||||
return outputs.last_hidden_state
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxPipelineConfig(PipelineConfig):
|
||||
"""Pipeline layout for Diffusers FLUX.1-dev (CLIP + T5 + packed DiT + FlowMatch)."""
|
||||
|
||||
scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
|
||||
transformer_arch: str = "FluxTransformer2DModel"
|
||||
vae_arch: str = "AutoencoderKL"
|
||||
text_encoder_archs: tuple[str, ...] = ("CLIPTextModel", "T5EncoderModel")
|
||||
tokenizer_archs: tuple[str, ...] = ("CLIPTokenizer", "T5TokenizerFast")
|
||||
|
||||
dit_config: FluxDiTConfig = field(default_factory=FluxDiTConfig)
|
||||
vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
|
||||
|
||||
embedded_cfg_scale: float = 3.5
|
||||
flow_shift: float | None = None
|
||||
|
||||
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (CLIPTextConfig(), T5LargeConfig()))
|
||||
preprocess_text_funcs: tuple[Callable[[str], str],
|
||||
...] = field(default_factory=lambda: (preprocess_text, preprocess_text))
|
||||
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
|
||||
default_factory=lambda: (_flux_clip_pooled_postprocess, _flux_t5_sequence_postprocess))
|
||||
|
||||
dit_precision: str = "bf16"
|
||||
vae_precision: str = "fp32"
|
||||
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", "bf16"))
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
te_cfgs = list(self.text_encoder_configs)
|
||||
if len(te_cfgs) >= 1:
|
||||
te_cfgs[0].tokenizer_kwargs.setdefault("padding", "max_length")
|
||||
te_cfgs[0].tokenizer_kwargs.setdefault("max_length", 77)
|
||||
te_cfgs[0].tokenizer_kwargs.setdefault("truncation", True)
|
||||
te_cfgs[0].tokenizer_kwargs.setdefault("return_tensors", "pt")
|
||||
if len(te_cfgs) >= 2:
|
||||
cap = 512
|
||||
te_cfgs[1].tokenizer_kwargs["max_length"] = min(int(te_cfgs[1].tokenizer_kwargs.get("max_length", cap)),
|
||||
cap)
|
||||
te_cfgs[1].tokenizer_kwargs.setdefault("padding", "max_length")
|
||||
te_cfgs[1].tokenizer_kwargs.setdefault("truncation", True)
|
||||
te_cfgs[1].tokenizer_kwargs.setdefault("return_tensors", "pt")
|
||||
@@ -1,46 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
|
||||
from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
|
||||
from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
|
||||
from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
|
||||
|
||||
def glm_image_t5_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
|
||||
mask: torch.Tensor = outputs.attention_mask
|
||||
hidden_state: torch.Tensor = outputs.last_hidden_state
|
||||
seq_lens = mask.gt(0).sum(dim=1).long()
|
||||
|
||||
assert torch.isnan(hidden_state).sum() == 0, "T5 hidden states contain NaN"
|
||||
|
||||
max_len = 512
|
||||
prompt_embeds = [u[:min(v, max_len)] for u, v in zip(hidden_state, seq_lens, strict=True)]
|
||||
prompt_embeds_tensor: torch.Tensor = torch.stack(
|
||||
[torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
|
||||
|
||||
return prompt_embeds_tensor
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageConfig(PipelineConfig):
|
||||
|
||||
dit_config: DiTConfig = field(default_factory=GlmImageDiTConfig)
|
||||
dit_precision: str = "bf16"
|
||||
|
||||
vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig)
|
||||
vae_precision: str = "fp32"
|
||||
vae_tiling: bool = True
|
||||
vae_sp: bool = False
|
||||
|
||||
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
|
||||
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
|
||||
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
|
||||
...] = field(default_factory=lambda: (glm_image_t5_postprocess, ))
|
||||
|
||||
flow_shift: float | None = 1.0
|
||||
embedded_cfg_scale: float = 7.5
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user