nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[misc]: remove DreamX World, HunyuanGameCraft, GLM-Image, LongCat, TurboDiffusion, and Z-Image

Browse Source 
These ports were marked pending removal in #1820. Each has a newer,
maintained counterpart in the repo, and their model-specific hooks
leaked into shared stages, request fields, CLI args, and loaders.

Remove each model's code, configs, presets, examples, converters,
tests, cookbook pages, CI planner rules, and registry entries, along
with the shared-code hooks only they used:

- camera kwargs (GameCraft camera_states, DreamX y_camera) in the
  shared denoising loop
- LongCat refine branches in latent preparation, the decoding crop,
  the --bsa-* CLI args, its train model, and third_party/longcat_video
- Z-Image cfg_normalization/cfg_truncation and GameCraft/DreamX
  camera/action fields in SamplingParam, ForwardBatch, and the schema
- the GLM-Image vision-language loader and the TurboDiffusion rCM
  scheduler and SLA auto-select in the train entrypoint
- the Wan definition groups that only preserved DreamX detector order

Removed env vars and their legacy aliases move to DEPRECATED_VARIABLES
per the env-var policy. Matrix Game 2.0 stays: the new trainer's world
model scenarios depend on it. The SLA and BSA attention backends stay.
The TurboDiffusion int8 kernels in fastvideo-kernel are left for a
separate kernel PR.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
SolitaryThinker
2026-10-06 22:40:54 +00:00
co-authored by Claude Opus 5.5
parent 02a027c49f
commit d692d1cc3b
236 changed files with 243 additions and 29701 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/skills/add-model-02-parity/SKILL.md
+1 -1
View File
@@ -68,7 +68,7 @@ Copy `templates/component_parity_test.py` and fill every `TODO` marker. The
template is distilled from:
- `tests/local_tests/transformers/test_ltx2.py`
- `tests/local_tests/transformers/test_gamecraft_parity.py`
- `tests/local_tests/gen3c/test_gen3c.py`
- `tests/local_tests/encoders/test_ltx2_gemma_parity.py`
- `tests/local_tests/vaes/test_oobleck_vae_parity.py`
- `tests/local_tests/sd35/test_sd35_component_parity.py`
.agents/skills/add-model-02-parity/templates/component_parity_test.py
+1 -1
View File
@@ -87,7 +87,7 @@ def _load_official_model(device: torch.device, dtype: torch.dtype) -> torch.nn.M
# TODO: import official class/factory and load real weights strictly.
# Examples in-tree:
# - LTX2: SingleGPUModelBuilder(...).build(device=device, dtype=dtype)
# - GameCraft: torch.load(...)["module"] -> official_model.load_state_dict(...)
# - GEN3C: torch.load(...)["state_dict"] -> official_model.load_state_dict(...)
# - Oobleck: create_model_from_config(config) + ckpt state_dict
OfficialClass = _import_or_skip(OFFICIAL_MODULE, OFFICIAL_CLASS)
model = OfficialClass() # TODO: pass official config kwargs.
.agents/skills/add-model-03-port-dit/SKILL.md
+2 -2
View File
@@ -52,8 +52,8 @@ scaling constants, dtype casts, state-dict names, and every output head.
- Loader path: `TransformerLoader` reads `transformer/config.json`, calls
`dit_config.update_model_arch(config)`, resolves `_class_name` through
`ModelRegistry`, and constructs the class with `config` and `hf_config`.
- Reference examples: `stable_audio.py`, `wanvideo.py`, `sd3.py`, `longcat.py`,
and `ltx2.py`.
- Reference examples: `stable_audio.py`, `wanvideo.py`, `sd3.py`, and
`ltx2.py`.
- Layer guidance: `fastvideo/layers/AGENTS.md`.
## Implementation Rules
.agents/skills/add-model-04-port-vae/SKILL.md
+1 -1
View File
@@ -51,7 +51,7 @@ posterior behavior, encode/decode output objects, tiling flags, and cropping.
- Loader path: VAE loaders resolve `_class_name` through `ModelRegistry` and
load converted component weights from the VAE subdir.
- Reference examples: `oobleck.py`, `autoencoder_kl.py`, `wanvae.py`,
`ltx2vae.py`, and `gamecraftvae.py`.
`ltx2vae.py`, and `hunyuanvae.py`.
- Layer guidance: `fastvideo/layers/AGENTS.md`.
## Implementation Rules
.agents/skills/add-model-07-conversion/SKILL.md
+6 -5
View File
@@ -43,11 +43,12 @@ from `../add-model/contracts/conversion_request.md`.
- `scripts/checkpoint_conversion/stable_audio_to_diffusers.py`: monolithic
`model.safetensors` split into transformer/VAE/conditioner, plus copied
passthrough subfolders. Use this shape for single-checkpoint official repos.
- `scripts/checkpoint_conversion/convert_gamecraft_full.py`: separate official
sources for transformer, VAE, text encoders, tokenizers, scheduler, and root
`model_index.json`.
- `scripts/checkpoint_conversion/longcat_to_fastvideo.py`: fused QKV/KV split,
renamed native transformer weights, and copied existing Diffusers components.
- `scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py`: separate
official sources for transformer, VAE, encoders, and vocoder, assembled under
a root `model_index.json`.
- `scripts/checkpoint_conversion/convert_flux2_klein.py`: fused QKV split,
renamed native transformer weights, and copied passthrough text encoder,
tokenizer, and scheduler components.
- `scripts/checkpoint_conversion/pt_to_safetensors.py`: simple `.pt` extraction
helper for nested checkpoint dictionaries.
.agents/skills/add-model-09-pipeline/SKILL.md
+1 -1
View File
@@ -247,7 +247,7 @@ setup gap, not a pass.
- `fastvideo/configs/pipelines/stable_audio.py` and
`fastvideo/pipelines/basic/stable_audio/presets.py` for config/preset shape.
- `fastvideo/registry.py` for `register_configs(...)` and preset registration.
- `tests/local_tests/pipelines/test_gamecraft_pipeline_parity.py` for latent
- `tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py` for latent
parity structure.
- `tests/local_tests/pipelines/test_stable_audio_pipeline_parity.py` for audio
parity structure.
.agents/skills/add-model/SKILL.md
+1 -1
View File
@@ -413,7 +413,7 @@ matching `*secret*`.
- `fastvideo/pipelines/basic/wan/` for standard T2V/I2V/DMD/Causal variants.
- `fastvideo/pipelines/basic/ltx2/` for non-standard stages and audio/video
patterns.
- `tests/local_tests/pipelines/test_gamecraft_pipeline_parity.py` for pipeline
- `tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py` for pipeline
parity shape.
- `tests/local_tests/transformers/test_ltx2.py`,
`tests/local_tests/vaes/test_ltx2_vae.py`, and
.agents/skills/seed-ssim-references/SKILL.md
+2 -2
View File
@@ -105,8 +105,8 @@ Detect artefact type by inspecting the file's imports / helper call:
- **pixel** (`.mp4`) — file imports
`run_text_to_video_similarity_test` / `run_image_to_video_similarity_test`
from `fastvideo.tests.ssim.inference_similarity_utils`, OR uses the
legacy custom-inline helper pattern (see `test_gamecraft`,
`test_longcat`, etc.). Default to pixel when both heuristics fail.
legacy custom-inline helper pattern (see `test_gen3c`). Default to pixel
when both heuristics fail.
Record `ARTEFACT_TYPE ∈ {pixel, latent}` for use in step 4. Steps 2, 3, 5,
and 6 are artefact-type-agnostic — `_iter_reference_files`,
.github/scripts/plan_merge_ci.py
-30
View File
@@ -130,11 +130,6 @@ class FamilyCoverage:
FAMILY_COVERAGE = (
FamilyCoverage(
re.compile(r"(^|[/_.-])dreamx(_world)?([/_.-]|$)"),
("test_dreamx.py", ),
("test_dreamx_world_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"),
("test_flux2_klein.py", ),
@@ -145,21 +140,11 @@ FAMILY_COVERAGE = (
("test_flux.py", ),
("test_flux_t2i_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])(hunyuan)?gamecraft([/_.-]|$)"),
("test_gamecraft.py", ),
("test_gamecraft_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])gen3c([/_.-]|$)"),
("test_gen3c.py", ),
("test_gen3c_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])glm[_-]?image([/_.-]|$)"),
("test_glm_image.py", ),
("test_glm_image_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])kandinsky[_-]?5([/_.-]|$)"),
("test_kandinsky5.py", ),
@@ -170,11 +155,6 @@ FAMILY_COVERAGE = (
("test_lingbot.py", ),
("test_lingbot_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])longcat([/_.-]|$)"),
("test_longcat.py", ),
("test_longcat_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])ltx[_-]?2([/_.-]|$)"),
("test_ltx2.py", ),
@@ -205,11 +185,6 @@ FAMILY_COVERAGE = (
("test_stable_audio.py", ),
("test_stable_audio_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])turbo(diffusion)?([/_.-]|$)"),
(),
("test_turbodiffusion_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])wan(video|vae)?([/_.-]|$)"),
("test_wan_t2v.py", "test_wan_vae.py", "test_wan_causal.py", "test_wan_denoising.py"),
@@ -219,11 +194,6 @@ FAMILY_COVERAGE = (
"test_wan_t2v_similarity.py",
),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])z[_-]?image([/_.-]|$)"),
("test_zimage.py", ),
("test_zimage_similarity.py", ),
),
)
.gitignore
-1
View File
@@ -36,7 +36,6 @@ env
*.log
weights/
logs/
/Z-Image/
official_weights/
converted_weights/
docs/assets/cookbook-recipes.json
+1 -127
View File
@@ -1,5 +1,5 @@
{
"version": 14,
"version": 15,
"recipes": [
{
"id": "fastwan21-t2v",
@@ -161,53 +161,6 @@
"Native MLX FastMetal T2V on 36 GB+ unified memory. Same --fast, --fast-spatial, and --refine flags as the 1.3B script."
]
},
{
"id": "turbodiffusion-wan21-1-3b-t2v",
"family": "turbodiffusion",
"stage": "inference",
"task": "Text to video",
"label": "TurboWan2.1 1.3B",
"summary": "A TurboDiffusion-accelerated Wan2.1 1.3B text-to-video run from its maintained single-GPU example.",
"model": "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers",
"source": "examples/inference/basic/basic_turbodiffusion.py",
"command": "python examples/inference/basic/basic_turbodiffusion.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "MP4 videos under video_samples_turbodiffusion/",
"related": ["turbodiffusion-wan21-14b-t2v", "turbowan22-i2v"]
},
{
"id": "turbodiffusion-wan21-14b-t2v",
"family": "turbodiffusion",
"stage": "inference",
"task": "Text to video",
"label": "TurboWan2.1 14B",
"summary": "TurboDiffusion acceleration applied to the 14B Wan2.1 text-to-video checkpoint; the checked-in source is configured for two GPUs.",
"model": "loayrashid/TurboWan2.1-T2V-14B-Diffusers",
"source": "examples/inference/basic/basic_turbodiffusion_14b.py",
"command": "python examples/inference/basic/basic_turbodiffusion_14b.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 2, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "MP4 videos under video_samples_turbodiffusion_14B/",
"related": ["turbodiffusion-wan21-1-3b-t2v", "turbowan22-i2v"]
},
{
"id": "turbowan22-i2v",
"family": "turbodiffusion",
"stage": "inference",
"task": "Image to video",
"label": "TurboWan2.2 A14B",
"summary": "A one-to-four-step image-to-video path using TurboDiffusion and the SLA attention backend from its maintained example.",
"model": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers",
"source": "examples/inference/basic/basic_turbodiffusion_i2v.py",
"command": "python examples/inference/basic/basic_turbodiffusion_i2v.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 2, "evidence": "source-configured"},
"evidence": "Source-backed",
"related": ["turbodiffusion-wan21-14b-t2v"]
},
{
"id": "ltx2-distilled-t2v",
"family": "ltx2",
@@ -444,53 +397,6 @@
"expected_artifact": "PNG images under outputs/flux_dev/samples/",
"limitations": ["FLUX.1 is loadable by ID but registers no model_family in fastvideo/registry.py; it is grouped under FLUX for documentation only."]
},
{
"id": "glm-image-t2i",
"family": "glm_image",
"stage": "inference",
"task": "Text to image",
"label": "GLM-Image",
"summary": "GLM-Image text-to-image generation from its maintained example.",
"model": "zai-org/GLM-Image",
"source": "examples/inference/basic/basic_glm_image.py",
"command": "python examples/inference/basic/basic_glm_image.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "PNG image at image_output/landscape.png",
"related": ["glm-image-edit"]
},
{
"id": "glm-image-edit",
"family": "glm_image",
"stage": "inference",
"task": "Image editing",
"label": "GLM-Image editing",
"summary": "Edit an input image with an instruction prompt using GLM-Image, from its maintained editing example.",
"model": "zai-org/GLM-Image",
"source": "examples/inference/basic/edit_glm_image.py",
"command": "python examples/inference/basic/edit_glm_image.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "PNG image at image_output/edited.png (input: assets/images/couple.jpg)",
"related": ["glm-image-t2i"]
},
{
"id": "zimage-turbo-t2i",
"family": "zimage",
"stage": "inference",
"task": "Text to image",
"label": "Z-Image Turbo",
"summary": "Z-Image Turbo text-to-image on a single GPU from its maintained example.",
"model": "Tongyi-MAI/Z-Image-Turbo",
"source": "examples/inference/basic/basic_zimage.py",
"command": "python examples/inference/basic/basic_zimage.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "PNG image at outputs/zimage/zimage_turbo.png"
},
{
"id": "sd35-medium-t2i",
"family": "sd35",
@@ -882,38 +788,6 @@
],
"limitations": ["Supply a compatible FastH3 adapter. The script infers dense or VSA attention from the adapter payload unless you override it."]
},
{
"id": "longcat-t2v",
"family": "longcat",
"stage": "inference",
"task": "Text to video",
"label": "LongCat Video T2V",
"summary": "LongCat Video text-to-video at 480p (50 steps), with distilled and 720p refinement passes included in the same maintained script.",
"model": "FastVideo/LongCat-Video-T2V-Diffusers",
"source": "examples/inference/basic/basic_longcat_t2v.py",
"command": "python examples/inference/basic/basic_longcat_t2v.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "MP4 videos under outputs_video/longcat_t2v_basic/, longcat_t2v_distill/, and longcat_t2v_refine_720p/",
"related": ["longcat-i2v"]
},
{
"id": "longcat-i2v",
"family": "longcat",
"stage": "inference",
"task": "Image to video",
"label": "LongCat Video I2V",
"summary": "LongCat Video image-to-video with optional distilled and refinement passes, from its maintained example.",
"model": "FastVideo/LongCat-Video-I2V-Diffusers",
"source": "examples/inference/basic/basic_longcat_i2v.py",
"command": "python examples/inference/basic/basic_longcat_i2v.py",
"gpu_types": ["NVIDIA"],
"hardware": {"gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "MP4 videos under outputs_video/longcat_i2v_basic/ and longcat_i2v_distill/",
"related": ["longcat-t2v"]
},
{
"id": "stable-audio-open-t2a",
"family": "stable_audio",
docs/assets/logos/SOURCES.md
+4 -7
View File
@@ -9,18 +9,15 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo.
| --- | --- | --- |
| `wan-ai.webp` | [Official Wan-AI Hugging Face organization avatar](https://huggingface.co/Wan-AI) | Wan family card and page header |
| `ltx.webp` | [Official Lightricks Hugging Face organization avatar](https://huggingface.co/Lightricks) | LTX family card and page header |
| `tencent-hunyuan.webp` | [Official Tencent Hunyuan Hugging Face organization avatar](https://huggingface.co/Tencent-Hunyuan) | Hunyuan and GameCraft cards, Hunyuan page header |
| `tencent-hunyuan.webp` | [Official Tencent Hunyuan Hugging Face organization avatar](https://huggingface.co/Tencent-Hunyuan) | Hunyuan family card and page header |
| `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header |
| `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers |
| `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header |
| `minimax.webp` | [Official MiniMax Hugging Face organization avatar](https://huggingface.co/MiniMaxAI) | MiniMax H3 family card and page header |
| `tongyi.webp` | [Official Tongyi MAI Hugging Face organization avatar](https://huggingface.co/Tongyi-MAI) | Z-Image family card and page header |
| `zai.webp` | [Official Z.ai Hugging Face organization avatar](https://huggingface.co/zai-org) | GLM-Image family card and page header |
| `stabilityai.webp` | [Official Stability AI Hugging Face organization avatar](https://huggingface.co/stabilityai) | Stable Diffusion and Stable Audio cards and page headers |
| `meituan-longcat.webp` | [Official Meituan LongCat Hugging Face organization avatar](https://huggingface.co/meituan-longcat) | LongCat family card and page header |
| `fastvideo.webp` | [Official FastVideo Hugging Face organization avatar](https://huggingface.co/FastVideo) | Matrix Game and MMAudio cards (converted weights published by this org), DreamX card |
| `fastvideo.webp` | [Official FastVideo Hugging Face organization avatar](https://huggingface.co/FastVideo) | Matrix Game and MMAudio cards (converted weights published by this org) |
Typographic tiles (no vendored image): TurboDiffusion ("Turbo"), HY-World
("HY"), LingBot ("LB"), MMAudio page header ("MMA"). These publishers have no
Typographic tiles (no vendored image): HY-World ("HY"), LingBot ("LB"),
MMAudio page header ("MMA"). These publishers have no
single official mark appropriate for reuse in the catalog; add a licensed
asset here if one becomes available.
docs/assets/logos/meituan-longcat.webp
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 5.1 KiB

docs/assets/logos/tongyi.webp
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 6.0 KiB

docs/assets/logos/zai.webp
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 2.4 KiB

docs/contributing/env_vars.md
+25 -17
View File
@@ -289,35 +289,43 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
| `FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY` | bool | `0` | test | Run the MLX Wan2.2 5B real-weights parity test on hosts with little memory. Deprecated names: `FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY`. |
| `FASTVIDEO_TEST_WAN22_5B_ROOT` | str | unset | test | Local Wan2.2 5B checkpoint for the MLX real-weights parity test. Deprecated names: `FASTVIDEO_WAN22_5B_ROOT`. |
| `FASTVIDEO_TEST_GRADNORM_UPDATE` | bool | `0` | test | Gradient-norm regression tests update their references. Deprecated names: `FASTVIDEO_GRADNORM_UPDATE`. |
| `FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Cam-Diffusers` | test | Model for the DreamX-World camera SSIM test. Deprecated names: `DREAMX_WORLD_SSIM_MODEL_PATH`. |
| `FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Diffusers` | test | Model for the DreamX-World autoregressive SSIM test. Deprecated names: `DREAMX_WORLD_AR_SSIM_MODEL_PATH`. |
| `FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR` | str | `black-forest-labs/FLUX.1-dev` | test | Model for the Flux text-to-image SSIM test. Deprecated names: `FLUX_T2I_MODEL_DIR`. |
| `FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH` | str | unset | test | Local Flux transformer for the Flux transformer test. Deprecated names: `FLUX_TRANSFORMER_PATH`. |
| `FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH` | str | `FastVideo/HunyuanGameCraft-Diffusers` | test | Model for the HunyuanGameCraft SSIM test. Deprecated names: `GAMECRAFT_MODEL_PATH`. |
| `FASTVIDEO_TEST_GEN3C_MODEL_PATH` | str | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | test | Model for the GEN3C SSIM test. Deprecated names: `GEN3C_MODEL_PATH`. |
| `FASTVIDEO_TEST_GEN3C_IMAGE_PATH` | str | unset | test | Input image for the GEN3C SSIM test. Deprecated names: `GEN3C_TEST_IMAGE_PATH`. |
| `FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR` | str | unset | test | Local official GLM-Image weights for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_LOCAL_WEIGHTS_DIR`. |
| `FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR` | str | unset | test | Model for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_MODEL_DIR`. |
| `FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS` | int | `1` | test | GPUs for the Kandinsky5 nightly end-to-end overfit test. Deprecated names: `KANDINSKY5_E2E_NUM_GPUS`. |
| `FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE` | bool | `0` | test | The Kandinsky5 nightly end-to-end test writes a missing reference video. Deprecated names: `KANDINSKY5_E2E_WRITE_REFERENCE`. |
| `FASTVIDEO_TEST_LONGCAT_MODEL_ROOT` | str | unset | test | Local LongCat-Video checkpoint for the golden-gate test. Deprecated names: `LONGCAT_MODEL_ROOT`. |
| `FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR` | str | unset | test | Local directory of MiniMax-H3 golden-gate tensors. Deprecated names: `MINIMAX_H3_GATE_GOLDEN_DIR`. |
| `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER` | int | `0` | test | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`. |
| `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT` | str | unset | test | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`. |
| `FASTVIDEO_TEST_SD35_MODEL_DIR` | str | `stabilityai/stable-diffusion-3.5-medium` | test | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`. |
| `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR` | str | unset | test | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`. |
| `FASTVIDEO_TEST_ZIMAGE_MODEL_DIR` | str | `Tongyi-MAI/Z-Image-Turbo` | test | Model for the Z-Image SSIM test. Deprecated names: `ZIMAGE_MODEL_DIR`. |
| `FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION` | str | `f332072aa78be7aecdf3ee76d5c247082da564a6` | test | Hugging Face revision of the Z-Image model for its SSIM test. Deprecated names: `ZIMAGE_MODEL_REVISION`. |
Variables that FastVideo no longer reads; setting one logs a warning:
| Deprecated variable | Reason |
| ----------------------------------------- | ---------------- |
| `FASTVIDEO_TARGET_DEVICE` | no code reads it |
| `FASTVIDEO_USE_PRECOMPILED` | no code reads it |
| `FASTVIDEO_RINGBUFFER_WARNING_INTERVAL` | no code reads it |
| `FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S` | no code reads it |
| `FASTVIDEO_SERVER_DEV_MODE` | no code reads it |
| `FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE` | no code reads it |
| `FASTVIDEO_TRACE_FUNCTION` | no code reads it |
| Deprecated variable | Reason |
| ------------------------------------------------ | ---------------------------- |
| `FASTVIDEO_TARGET_DEVICE` | no code reads it |
| `FASTVIDEO_USE_PRECOMPILED` | no code reads it |
| `FASTVIDEO_RINGBUFFER_WARNING_INTERVAL` | no code reads it |
| `FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S` | no code reads it |
| `FASTVIDEO_SERVER_DEV_MODE` | no code reads it |
| `FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE` | no code reads it |
| `FASTVIDEO_TRACE_FUNCTION` | no code reads it |
| `FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH` | DreamX World was removed |
| `DREAMX_WORLD_SSIM_MODEL_PATH` | DreamX World was removed |
| `FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH` | DreamX World was removed |
| `DREAMX_WORLD_AR_SSIM_MODEL_PATH` | DreamX World was removed |
| `FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH` | HunyuanGameCraft was removed |
| `GAMECRAFT_MODEL_PATH` | HunyuanGameCraft was removed |
| `FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR` | GLM-Image was removed |
| `GLM_IMAGE_LOCAL_WEIGHTS_DIR` | GLM-Image was removed |
| `FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR` | GLM-Image was removed |
| `GLM_IMAGE_MODEL_DIR` | GLM-Image was removed |
| `FASTVIDEO_TEST_LONGCAT_MODEL_ROOT` | LongCat was removed |
| `LONGCAT_MODEL_ROOT` | LongCat was removed |
| `FASTVIDEO_TEST_ZIMAGE_MODEL_DIR` | Z-Image was removed |
| `ZIMAGE_MODEL_DIR` | Z-Image was removed |
| `FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION` | Z-Image was removed |
| `ZIMAGE_MODEL_REVISION` | Z-Image was removed |
<!-- END GENERATED ENV TABLE -->
docs/cookbook/cosmos.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Cosmos recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="cosmos" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="cosmos" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/flux.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# FLUX recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="flux" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="flux" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/glm-image.md
-140
View File
@@ -1,140 +0,0 @@
---
hide:
- toc
---
# GLM-Image recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="glm_image" data-recipes="../../assets/cookbook-recipes.json?v=14">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
<span class="cookbook-family-header__logo">
<img class="off-glb" src="../../assets/logos/zai.webp" alt="Z.ai" width="112" height="112">
</span>
<div>
<p class="cookbook-eyebrow">Maintained family · Inference</p>
<h2>GLM-Image inference recipes</h2>
<p>GLM-Image from Z.ai supports both text-to-image generation and instruction-based image editing, each with a maintained example.</p>
</div>
</div>
<div class="cookbook-lifecycle" aria-label="Lifecycle stages">
<span class="cookbook-lifecycle__stage cookbook-lifecycle__stage--active">Inference <small>live</small></span>
<span class="cookbook-lifecycle__stage">Distillation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Fine-tuning <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Training <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Evaluation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Optimization <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Deployment <small>planned</small></span>
</div>
</header>
<nav class="cookbook-jumpnav" aria-label="Recipe page sections">
<a href="#recipe-builder">Builder</a>
<a href="#cookbook-setup">Setup</a>
<a href="#cookbook-troubleshooting">Troubleshooting</a>
<a href="#cookbook-evidence">Evidence</a>
</nav>
<section class="cookbook-builder" id="recipe-builder" aria-labelledby="builder-heading">
<div class="cookbook-builder__intro">
<h2 id="builder-heading">Pick a recipe and runtime</h2>
<p>Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.</p>
</div>
<div class="cookbook-builder__layout">
<div class="cookbook-controls">
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Recipe</strong>
<span>Task and checkpoint</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--models" data-cookbook-model-options role="group" aria-label="Recipe">
<button type="button" disabled>Loading GLM-Image recipes...</button>
</div>
</div>
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Runtime</strong>
<span>Maintained paths only</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--hardware" data-cookbook-hardware-options role="group" aria-label="Runtime">
<button type="button" disabled>Loading runtimes...</button>
</div>
</div>
<p class="cookbook-selection-description" data-cookbook-description>Loading recipe details...</p>
<p class="cookbook-hardware-note">Exact device and memory details appear only when a recorded run supports them.</p>
<div class="cookbook-hardware-state" data-cookbook-hardware-state role="status" aria-live="polite">
Reading recipe evidence...
</div>
</div>
<article class="cookbook-result">
<div class="cookbook-result__header">
<h3 data-cookbook-label>Loading...</h3>
<div class="cookbook-result__badges">
<span class="cookbook-badge">Maintained</span>
<span class="cookbook-badge" data-cookbook-evidence>Source-backed</span>
<span class="cookbook-badge cookbook-badge--neutral" data-cookbook-hardware-badge>Source config</span>
</div>
</div>
<dl class="cookbook-result__facts">
<div><dt>Model</dt><dd data-cookbook-model>Loading...</dd></div>
<div><dt>Workload</dt><dd data-cookbook-task>Loading...</dd></div>
<div><dt>Source configuration</dt><dd data-cookbook-gpus>Loading...</dd></div>
<div><dt>Expected output</dt><dd data-cookbook-artifact>Loading...</dd></div>
</dl>
<div class="cookbook-command">
<div class="cookbook-command__bar">
<span>Terminal</span>
</div>
<pre><code class="language-bash" data-cookbook-command>Loading...</code></pre>
</div>
<div class="cookbook-result__footer">
<a data-cookbook-source href="../../inference/examples/basic/">Open example source</a>
<a data-cookbook-model-link href="https://huggingface.co/zai-org">View model card</a>
</div>
<p class="cookbook-picker__status" role="status" aria-live="polite" data-cookbook-status></p>
</article>
</div>
<noscript>
<div class="cookbook-noscript">
JavaScript is needed for the guided selector. You can still browse the
<a href="../../inference/examples/examples_inference_index/">maintained inference examples</a>.
</div>
</noscript>
</section>
</div>
<details class="cookbook-collapsible" id="cookbook-setup">
<summary>Setup</summary>
<div class="cookbook-collapsible__body">
<p>The generated commands expect a local clone:</p>
<pre><code>git clone https://github.com/hao-ai-lab/FastVideo.git
cd FastVideo</code></pre>
<p>Use <a href="../../inference/configuration/">Configuration</a> for supported Python and CLI settings, <a href="../../inference/optimizations/">Optimizations</a> for attention and memory tradeoffs, and the <a href="../../inference/support_matrix/">support matrix</a> for the supported model and optimization surface.</p>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-troubleshooting">
<summary>Troubleshooting</summary>
<div class="cookbook-collapsible__body">
<ul>
<li>The editing example reads <code>assets/images/couple.jpg</code> from the repository root, so run it from a repo checkout rather than an arbitrary working directory.</li>
<li>Output paths default under <code>image_output/</code>; pass <code>--output</code> to change them.</li>
</ul>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-evidence">
<summary>Evidence status</summary>
<div class="cookbook-collapsible__body">
<p>All recipes on this page are <strong>Source-backed</strong>: their commands, model IDs, and flags were validated against the checked-in FastVideo sources listed above (static validation). No runtime GPU validation is recorded for these recipes, so GPU model fit, memory use, throughput, and runtime duration are <strong>Unknown</strong> and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.</p>
</div>
</details>
docs/cookbook/hunyuan.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Hunyuan recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="hunyuan" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="hunyuan" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/index.md
-115
View File
@@ -163,50 +163,6 @@ hide:
</ul>
</span>
</a>
<a class="cookbook-family-tile cookbook-family-tile--ready" href="./longcat/" aria-label="Open LongCat recipes">
<span class="cookbook-family-tile__visual" data-evervault>
<span class="cookbook-evervault" aria-hidden="true">
<span class="cookbook-evervault__gradient"></span>
<span class="cookbook-evervault__noise" data-cookbook-pattern></span>
</span>
<span class="cookbook-family-tile__logo-wrap">
<img class="off-glb" src="../assets/logos/meituan-longcat.webp" alt="" width="132" height="132" loading="lazy">
</span>
</span>
<span class="cookbook-family-tile__footer">
<span class="cookbook-family-tile__footer-top">
<span><strong>LongCat</strong><small>T2V, I2V, optional refine</small></span>
<span class="cookbook-count">2 recipes</span>
</span>
<ul class="cookbook-mode-row">
<li>T2V</li>
<li>I2V</li>
</ul>
</span>
</a>
<a class="cookbook-family-tile cookbook-family-tile--ready" href="./turbodiffusion/" aria-label="Open TurboDiffusion recipes">
<span class="cookbook-family-tile__visual" data-evervault>
<span class="cookbook-evervault" aria-hidden="true">
<span class="cookbook-evervault__gradient"></span>
<span class="cookbook-evervault__noise" data-cookbook-pattern></span>
</span>
<span class="cookbook-family-tile__logo-wrap">
<span class="cookbook-family-tile__monogram" aria-hidden="true">Turbo</span>
</span>
</span>
<span class="cookbook-family-tile__footer">
<span class="cookbook-family-tile__footer-top">
<span><strong>TurboDiffusion</strong><small>Accelerated Wan profiles</small></span>
<span class="cookbook-count">3 recipes</span>
</span>
<ul class="cookbook-mode-row">
<li>T2V</li>
<li>I2V</li>
</ul>
</span>
</a>
</div>
</section>
@@ -236,49 +192,6 @@ hide:
</span>
</a>
<a class="cookbook-family-tile cookbook-family-tile--ready" href="./glm-image/" aria-label="Open GLM-Image recipes">
<span class="cookbook-family-tile__visual" data-evervault>
<span class="cookbook-evervault" aria-hidden="true">
<span class="cookbook-evervault__gradient"></span>
<span class="cookbook-evervault__noise" data-cookbook-pattern></span>
</span>
<span class="cookbook-family-tile__logo-wrap">
<img class="off-glb" src="../assets/logos/zai.webp" alt="" width="132" height="132" loading="lazy">
</span>
</span>
<span class="cookbook-family-tile__footer">
<span class="cookbook-family-tile__footer-top">
<span><strong>GLM-Image</strong><small>Generate and edit</small></span>
<span class="cookbook-count">2 recipes</span>
</span>
<ul class="cookbook-mode-row">
<li>T2I</li>
<li>Edit</li>
</ul>
</span>
</a>
<a class="cookbook-family-tile cookbook-family-tile--ready" href="./z-image/" aria-label="Open Z-Image recipes">
<span class="cookbook-family-tile__visual" data-evervault>
<span class="cookbook-evervault" aria-hidden="true">
<span class="cookbook-evervault__gradient"></span>
<span class="cookbook-evervault__noise" data-cookbook-pattern></span>
</span>
<span class="cookbook-family-tile__logo-wrap">
<img class="off-glb" src="../assets/logos/tongyi.webp" alt="" width="132" height="132" loading="lazy">
</span>
</span>
<span class="cookbook-family-tile__footer">
<span class="cookbook-family-tile__footer-top">
<span><strong>Z-Image</strong><small>Turbo text to image</small></span>
<span class="cookbook-count">1 recipe</span>
</span>
<ul class="cookbook-mode-row">
<li>T2I</li>
</ul>
</span>
</a>
<a class="cookbook-family-tile cookbook-family-tile--ready" href="./stable-diffusion/" aria-label="Open Stable Diffusion recipes">
<span class="cookbook-family-tile__visual" data-evervault>
<span class="cookbook-evervault" aria-hidden="true">
@@ -407,20 +320,6 @@ hide:
<p>These families already have runnable examples. The cookbook page is not ready, so the cards are not links.</p>
</div>
<div class="cookbook-family-grid">
<article class="cookbook-family-tile cookbook-family-tile--coming" aria-label="GameCraft cookbook page planned; runnable examples exist">
<span class="cookbook-family-tile__visual">
<span class="cookbook-family-tile__logo-wrap">
<img class="off-glb" src="../assets/logos/tencent-hunyuan.webp" alt="" width="132" height="132" loading="lazy">
</span>
</span>
<span class="cookbook-family-tile__footer">
<span class="cookbook-family-tile__footer-top">
<span><strong>GameCraft</strong><small>Game world generation</small></span>
<span class="cookbook-count">Page planned</span>
</span>
</span>
</article>
<article class="cookbook-family-tile cookbook-family-tile--coming" aria-label="GEN3C cookbook page planned; runnable examples exist">
<span class="cookbook-family-tile__visual">
<span class="cookbook-family-tile__logo-wrap">
@@ -449,20 +348,6 @@ hide:
</span>
</article>
<article class="cookbook-family-tile cookbook-family-tile--coming" aria-label="DreamX cookbook page planned; runnable examples exist">
<span class="cookbook-family-tile__visual">
<span class="cookbook-family-tile__logo-wrap">
<img class="off-glb" src="../assets/logos/fastvideo.webp" alt="" width="132" height="132" loading="lazy">
</span>
</span>
<span class="cookbook-family-tile__footer">
<span class="cookbook-family-tile__footer-top">
<span><strong>DreamX</strong><small>World generation</small></span>
<span class="cookbook-count">Page planned</span>
</span>
</span>
</article>
<article class="cookbook-family-tile cookbook-family-tile--coming" aria-label="LingBot cookbook page planned; runnable examples exist">
<span class="cookbook-family-tile__visual">
<span class="cookbook-family-tile__logo-wrap">
docs/cookbook/kandinsky5.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Kandinsky 5 recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="kandinsky5" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="kandinsky5" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/kandinsky6.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Kandinsky 6 recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="kandinsky6" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="kandinsky6" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/longcat.md
-140
View File
@@ -1,140 +0,0 @@
---
hide:
- toc
---
# LongCat recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="longcat" data-recipes="../../assets/cookbook-recipes.json?v=14">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
<span class="cookbook-family-header__logo">
<img class="off-glb" src="../../assets/logos/meituan-longcat.webp" alt="Meituan LongCat" width="112" height="112">
</span>
<div>
<p class="cookbook-eyebrow">Maintained family · Inference</p>
<h2>LongCat inference recipes</h2>
<p>LongCat Video from Meituan covers text-to-video and image-to-video, and its maintained examples chain optional distilled and 720p refinement passes.</p>
</div>
</div>
<div class="cookbook-lifecycle" aria-label="Lifecycle stages">
<span class="cookbook-lifecycle__stage cookbook-lifecycle__stage--active">Inference <small>live</small></span>
<span class="cookbook-lifecycle__stage">Distillation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Fine-tuning <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Training <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Evaluation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Optimization <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Deployment <small>planned</small></span>
</div>
</header>
<nav class="cookbook-jumpnav" aria-label="Recipe page sections">
<a href="#recipe-builder">Builder</a>
<a href="#cookbook-setup">Setup</a>
<a href="#cookbook-troubleshooting">Troubleshooting</a>
<a href="#cookbook-evidence">Evidence</a>
</nav>
<section class="cookbook-builder" id="recipe-builder" aria-labelledby="builder-heading">
<div class="cookbook-builder__intro">
<h2 id="builder-heading">Pick a recipe and runtime</h2>
<p>Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.</p>
</div>
<div class="cookbook-builder__layout">
<div class="cookbook-controls">
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Recipe</strong>
<span>Task and checkpoint</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--models" data-cookbook-model-options role="group" aria-label="Recipe">
<button type="button" disabled>Loading LongCat recipes...</button>
</div>
</div>
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Runtime</strong>
<span>Maintained paths only</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--hardware" data-cookbook-hardware-options role="group" aria-label="Runtime">
<button type="button" disabled>Loading runtimes...</button>
</div>
</div>
<p class="cookbook-selection-description" data-cookbook-description>Loading recipe details...</p>
<p class="cookbook-hardware-note">Exact device and memory details appear only when a recorded run supports them.</p>
<div class="cookbook-hardware-state" data-cookbook-hardware-state role="status" aria-live="polite">
Reading recipe evidence...
</div>
</div>
<article class="cookbook-result">
<div class="cookbook-result__header">
<h3 data-cookbook-label>Loading...</h3>
<div class="cookbook-result__badges">
<span class="cookbook-badge">Maintained</span>
<span class="cookbook-badge" data-cookbook-evidence>Source-backed</span>
<span class="cookbook-badge cookbook-badge--neutral" data-cookbook-hardware-badge>Source config</span>
</div>
</div>
<dl class="cookbook-result__facts">
<div><dt>Model</dt><dd data-cookbook-model>Loading...</dd></div>
<div><dt>Workload</dt><dd data-cookbook-task>Loading...</dd></div>
<div><dt>Source configuration</dt><dd data-cookbook-gpus>Loading...</dd></div>
<div><dt>Expected output</dt><dd data-cookbook-artifact>Loading...</dd></div>
</dl>
<div class="cookbook-command">
<div class="cookbook-command__bar">
<span>Terminal</span>
</div>
<pre><code class="language-bash" data-cookbook-command>Loading...</code></pre>
</div>
<div class="cookbook-result__footer">
<a data-cookbook-source href="../../inference/examples/basic/">Open example source</a>
<a data-cookbook-model-link href="https://huggingface.co/meituan-longcat">View model card</a>
</div>
<p class="cookbook-picker__status" role="status" aria-live="polite" data-cookbook-status></p>
</article>
</div>
<noscript>
<div class="cookbook-noscript">
JavaScript is needed for the guided selector. You can still browse the
<a href="../../inference/examples/examples_inference_index/">maintained inference examples</a>.
</div>
</noscript>
</section>
</div>
<details class="cookbook-collapsible" id="cookbook-setup">
<summary>Setup</summary>
<div class="cookbook-collapsible__body">
<p>The generated commands expect a local clone:</p>
<pre><code>git clone https://github.com/hao-ai-lab/FastVideo.git
cd FastVideo</code></pre>
<p>Use <a href="../../inference/configuration/">Configuration</a> for supported Python and CLI settings, <a href="../../inference/optimizations/">Optimizations</a> for attention and memory tradeoffs, and the <a href="../../inference/support_matrix/">support matrix</a> for the supported model and optimization surface.</p>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-troubleshooting">
<summary>Troubleshooting</summary>
<div class="cookbook-collapsible__body">
<ul>
<li>Each LongCat script runs multiple passes (basic, distilled, refine); total runtime scales accordingly, and every pass prints its own output directory.</li>
<li>Out of memory: the sources already enable VAE and text-encoder CPU offload; further options are covered in <a href="../../inference/offloading/">Offloading</a>.</li>
</ul>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-evidence">
<summary>Evidence status</summary>
<div class="cookbook-collapsible__body">
<p>All recipes on this page are <strong>Source-backed</strong>: their commands, model IDs, and flags were validated against the checked-in FastVideo sources listed above (static validation). No runtime GPU validation is recorded for these recipes, so GPU model fit, memory use, throughput, and runtime duration are <strong>Unknown</strong> and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.</p>
</div>
</details>
docs/cookbook/ltx.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# LTX recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="ltx2" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="ltx2" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/matrix-game.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Matrix Game recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="matrixgame" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="matrixgame" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/minimax-h3.md
+1 -1
View File
@@ -18,7 +18,7 @@ eight-forward T2VA command](../getting_started/installation/mlx.md#pruned-eight-
It reads `fastvideo_inference.json` for the trained schedule. The command
uses native 832x480 resolution and all requested frames.
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="minimax_h3" data-default-recipe="fasth3-preview-cuda" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="minimax_h3" data-default-recipe="fasth3-preview-cuda" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/mmaudio.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# MMAudio recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="mmaudio" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="mmaudio" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/stable-audio.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Stable Audio recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="stable_audio" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="stable_audio" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/stable-diffusion.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Stable Diffusion recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="sd35" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="sd35" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/turbodiffusion.md
-140
View File
@@ -1,140 +0,0 @@
---
hide:
- toc
---
# TurboDiffusion recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="turbodiffusion" data-recipes="../../assets/cookbook-recipes.json?v=14">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
<span class="cookbook-family-header__logo">
<span class="cookbook-family-tile__monogram" aria-hidden="true">Turbo</span>
</span>
<div>
<p class="cookbook-eyebrow">Maintained family · Inference</p>
<h2>TurboDiffusion inference recipes</h2>
<p>TurboDiffusion profiles accelerate Wan checkpoints with step-distilled sampling and the SLA attention backend. These recipes follow the registry's `turbodiffusion` model family.</p>
</div>
</div>
<div class="cookbook-lifecycle" aria-label="Lifecycle stages">
<span class="cookbook-lifecycle__stage cookbook-lifecycle__stage--active">Inference <small>live</small></span>
<span class="cookbook-lifecycle__stage">Distillation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Fine-tuning <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Training <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Evaluation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Optimization <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Deployment <small>planned</small></span>
</div>
</header>
<nav class="cookbook-jumpnav" aria-label="Recipe page sections">
<a href="#recipe-builder">Builder</a>
<a href="#cookbook-setup">Setup</a>
<a href="#cookbook-troubleshooting">Troubleshooting</a>
<a href="#cookbook-evidence">Evidence</a>
</nav>
<section class="cookbook-builder" id="recipe-builder" aria-labelledby="builder-heading">
<div class="cookbook-builder__intro">
<h2 id="builder-heading">Pick a recipe and runtime</h2>
<p>Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.</p>
</div>
<div class="cookbook-builder__layout">
<div class="cookbook-controls">
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Recipe</strong>
<span>Task and checkpoint</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--models" data-cookbook-model-options role="group" aria-label="Recipe">
<button type="button" disabled>Loading TurboDiffusion recipes...</button>
</div>
</div>
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Runtime</strong>
<span>Maintained paths only</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--hardware" data-cookbook-hardware-options role="group" aria-label="Runtime">
<button type="button" disabled>Loading runtimes...</button>
</div>
</div>
<p class="cookbook-selection-description" data-cookbook-description>Loading recipe details...</p>
<p class="cookbook-hardware-note">Exact device and memory details appear only when a recorded run supports them.</p>
<div class="cookbook-hardware-state" data-cookbook-hardware-state role="status" aria-live="polite">
Reading recipe evidence...
</div>
</div>
<article class="cookbook-result">
<div class="cookbook-result__header">
<h3 data-cookbook-label>Loading...</h3>
<div class="cookbook-result__badges">
<span class="cookbook-badge">Maintained</span>
<span class="cookbook-badge" data-cookbook-evidence>Source-backed</span>
<span class="cookbook-badge cookbook-badge--neutral" data-cookbook-hardware-badge>Source config</span>
</div>
</div>
<dl class="cookbook-result__facts">
<div><dt>Model</dt><dd data-cookbook-model>Loading...</dd></div>
<div><dt>Workload</dt><dd data-cookbook-task>Loading...</dd></div>
<div><dt>Source configuration</dt><dd data-cookbook-gpus>Loading...</dd></div>
<div><dt>Expected output</dt><dd data-cookbook-artifact>Loading...</dd></div>
</dl>
<div class="cookbook-command">
<div class="cookbook-command__bar">
<span>Terminal</span>
</div>
<pre><code class="language-bash" data-cookbook-command>Loading...</code></pre>
</div>
<div class="cookbook-result__footer">
<a data-cookbook-source href="../../inference/examples/basic/">Open example source</a>
<a data-cookbook-model-link href="https://huggingface.co/loayrashid">View model card</a>
</div>
<p class="cookbook-picker__status" role="status" aria-live="polite" data-cookbook-status></p>
</article>
</div>
<noscript>
<div class="cookbook-noscript">
JavaScript is needed for the guided selector. You can still browse the
<a href="../../inference/examples/examples_inference_index/">maintained inference examples</a>.
</div>
</noscript>
</section>
</div>
<details class="cookbook-collapsible" id="cookbook-setup">
<summary>Setup</summary>
<div class="cookbook-collapsible__body">
<p>The generated commands expect a local clone:</p>
<pre><code>git clone https://github.com/hao-ai-lab/FastVideo.git
cd FastVideo</code></pre>
<p>Use <a href="../../inference/configuration/">Configuration</a> for supported Python and CLI settings, <a href="../../inference/optimizations/">Optimizations</a> for attention and memory tradeoffs, and the <a href="../../inference/support_matrix/">support matrix</a> for the supported model and optimization surface.</p>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-troubleshooting">
<summary>Troubleshooting</summary>
<div class="cookbook-collapsible__body">
<ul>
<li>TurboDiffusion paths load community-published <code>loayrashid/TurboWan*</code> checkpoints; availability is governed by those repos.</li>
<li>The SLA attention backend used by the I2V recipe is selected inside the example source; do not combine it with another <code>FASTVIDEO_ATTENTION_BACKEND</code> override in the same shell.</li>
</ul>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-evidence">
<summary>Evidence status</summary>
<div class="cookbook-collapsible__body">
<p>All recipes on this page are <strong>Source-backed</strong>: their commands, model IDs, and flags were validated against the checked-in FastVideo sources listed above (static validation). No runtime GPU validation is recorded for these recipes, so GPU model fit, memory use, throughput, and runtime duration are <strong>Unknown</strong> and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.</p>
</div>
</details>
docs/cookbook/wan.md
+1 -1
View File
@@ -5,7 +5,7 @@ hide:
# Wan recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="wan" data-recipes="../../assets/cookbook-recipes.json?v=14">
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="wan" data-recipes="../../assets/cookbook-recipes.json?v=15">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
docs/cookbook/z-image.md
-140
View File
@@ -1,140 +0,0 @@
---
hide:
- toc
---
# Z-Image recipes
<div class="cookbook-shell cookbook-family-page" data-cookbook data-family="zimage" data-recipes="../../assets/cookbook-recipes.json?v=14">
<header class="cookbook-family-header">
<a class="cookbook-back-link" href="../"><span aria-hidden="true">←</span> All model families</a>
<div class="cookbook-family-header__body">
<span class="cookbook-family-header__logo">
<img class="off-glb" src="../../assets/logos/tongyi.webp" alt="Tongyi MAI" width="112" height="112">
</span>
<div>
<p class="cookbook-eyebrow">Maintained family · Inference</p>
<h2>Z-Image inference recipes</h2>
<p>Z-Image Turbo from Tongyi MAI is a fast text-to-image model. The maintained example runs it on a single GPU.</p>
</div>
</div>
<div class="cookbook-lifecycle" aria-label="Lifecycle stages">
<span class="cookbook-lifecycle__stage cookbook-lifecycle__stage--active">Inference <small>live</small></span>
<span class="cookbook-lifecycle__stage">Distillation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Fine-tuning <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Training <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Evaluation <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Optimization <small>planned</small></span>
<span class="cookbook-lifecycle__stage">Deployment <small>planned</small></span>
</div>
</header>
<nav class="cookbook-jumpnav" aria-label="Recipe page sections">
<a href="#recipe-builder">Builder</a>
<a href="#cookbook-setup">Setup</a>
<a href="#cookbook-troubleshooting">Troubleshooting</a>
<a href="#cookbook-evidence">Evidence</a>
</nav>
<section class="cookbook-builder" id="recipe-builder" aria-labelledby="builder-heading">
<div class="cookbook-builder__intro">
<h2 id="builder-heading">Pick a recipe and runtime</h2>
<p>Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.</p>
</div>
<div class="cookbook-builder__layout">
<div class="cookbook-controls">
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Recipe</strong>
<span>Task and checkpoint</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--models" data-cookbook-model-options role="group" aria-label="Recipe">
<button type="button" disabled>Loading Z-Image recipes...</button>
</div>
</div>
<div class="cookbook-selection-row">
<div class="cookbook-selection-row__label">
<strong>Runtime</strong>
<span>Maintained paths only</span>
</div>
<div class="cookbook-option-grid cookbook-option-grid--hardware" data-cookbook-hardware-options role="group" aria-label="Runtime">
<button type="button" disabled>Loading runtimes...</button>
</div>
</div>
<p class="cookbook-selection-description" data-cookbook-description>Loading recipe details...</p>
<p class="cookbook-hardware-note">Exact device and memory details appear only when a recorded run supports them.</p>
<div class="cookbook-hardware-state" data-cookbook-hardware-state role="status" aria-live="polite">
Reading recipe evidence...
</div>
</div>
<article class="cookbook-result">
<div class="cookbook-result__header">
<h3 data-cookbook-label>Loading...</h3>
<div class="cookbook-result__badges">
<span class="cookbook-badge">Maintained</span>
<span class="cookbook-badge" data-cookbook-evidence>Source-backed</span>
<span class="cookbook-badge cookbook-badge--neutral" data-cookbook-hardware-badge>Source config</span>
</div>
</div>
<dl class="cookbook-result__facts">
<div><dt>Model</dt><dd data-cookbook-model>Loading...</dd></div>
<div><dt>Workload</dt><dd data-cookbook-task>Loading...</dd></div>
<div><dt>Source configuration</dt><dd data-cookbook-gpus>Loading...</dd></div>
<div><dt>Expected output</dt><dd data-cookbook-artifact>Loading...</dd></div>
</dl>
<div class="cookbook-command">
<div class="cookbook-command__bar">
<span>Terminal</span>
</div>
<pre><code class="language-bash" data-cookbook-command>Loading...</code></pre>
</div>
<div class="cookbook-result__footer">
<a data-cookbook-source href="../../inference/examples/basic/">Open example source</a>
<a data-cookbook-model-link href="https://huggingface.co/Tongyi-MAI">View model card</a>
</div>
<p class="cookbook-picker__status" role="status" aria-live="polite" data-cookbook-status></p>
</article>
</div>
<noscript>
<div class="cookbook-noscript">
JavaScript is needed for the guided selector. You can still browse the
<a href="../../inference/examples/examples_inference_index/">maintained inference examples</a>.
</div>
</noscript>
</section>
</div>
<details class="cookbook-collapsible" id="cookbook-setup">
<summary>Setup</summary>
<div class="cookbook-collapsible__body">
<p>The generated commands expect a local clone:</p>
<pre><code>git clone https://github.com/hao-ai-lab/FastVideo.git
cd FastVideo</code></pre>
<p>Use <a href="../../inference/configuration/">Configuration</a> for supported Python and CLI settings, <a href="../../inference/optimizations/">Optimizations</a> for attention and memory tradeoffs, and the <a href="../../inference/support_matrix/">support matrix</a> for the supported model and optimization surface.</p>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-troubleshooting">
<summary>Troubleshooting</summary>
<div class="cookbook-collapsible__body">
<ul>
<li>Output defaults to <code>outputs/zimage/zimage_turbo.png</code>; pass <code>--output</code> to redirect.</li>
<li>Gated or missing checkpoints: run <code>huggingface-cli login</code> and confirm you accepted the model's license on Hugging Face.</li>
</ul>
</div>
</details>
<details class="cookbook-collapsible" id="cookbook-evidence">
<summary>Evidence status</summary>
<div class="cookbook-collapsible__body">
<p>All recipes on this page are <strong>Source-backed</strong>: their commands, model IDs, and flags were validated against the checked-in FastVideo sources listed above (static validation). No runtime GPU validation is recorded for these recipes, so GPU model fit, memory use, throughput, and runtime duration are <strong>Unknown</strong> and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.</p>
</div>
</details>
docs/design/inference_schema_parity_inventory.yaml
-53
View File
@@ -205,9 +205,6 @@ surfaces:
sources:
- fastvideo.configs.pipelines.gen3c.Gen3CConfig
- fastvideo.configs.pipelines.gen3c.Gen3CInferenceConfig
color_correction_strength:
sources:
- fastvideo.configs.pipelines.dreamx_world.DreamXWorld5BARPipelineConfig
default_camera_rotation:
sources:
- fastvideo.configs.pipelines.gen3c.Gen3CConfig
@@ -282,11 +279,6 @@ surfaces:
sources:
- fastvideo.configs.pipelines.lingbotworld.LingBotWorldI2V480PConfig
- fastvideo.configs.pipelines.lingbotworld.Wan2_2_I2V_A14B_Config
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionI2VConfig
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionI2V_A14B_Config
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionT2VConfig
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionT2V_14B_Config
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionT2V_1_3B_Config
- fastvideo.configs.pipelines.wan.FastWan2_1_T2V_480P_Config
- fastvideo.configs.pipelines.wan.FastWan2_2_TI2V_5B_Config
- fastvideo.configs.pipelines.matrixgame2.MatrixGame2BaseI2V480PConfig
@@ -305,11 +297,6 @@ surfaces:
sources:
- fastvideo.configs.pipelines.lingbotworld.LingBotWorldI2V480PConfig
- fastvideo.configs.pipelines.lingbotworld.Wan2_2_I2V_A14B_Config
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionI2VConfig
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionI2V_A14B_Config
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionT2VConfig
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionT2V_14B_Config
- fastvideo.configs.pipelines.turbodiffusion.TurboDiffusionT2V_1_3B_Config
- fastvideo.configs.pipelines.wan.FastWan2_1_T2V_480P_Config
- fastvideo.configs.pipelines.wan.FastWan2_2_TI2V_5B_Config
- fastvideo.configs.pipelines.matrixgame2.MatrixGame2BaseI2V480PConfig
@@ -324,48 +311,21 @@ surfaces:
- fastvideo.configs.pipelines.wan.WanI2V720PConfig
- fastvideo.configs.pipelines.wan.WanT2V480PConfig
- fastvideo.configs.pipelines.wan.WanT2V720PConfig
bsa_cdf_threshold:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
bsa_chunk_k:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
bsa_chunk_q:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
bsa_params:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
bsa_sparsity:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
enable_bsa:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
enable_kv_cache:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
enhance_hf:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
offload_kv_cache:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
t_thresh:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
use_distill:
sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
scheduler_arch:
sources:
- fastvideo.configs.pipelines.sd35.SD35Config
- fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
text_encoder_archs:
sources:
- fastvideo.configs.pipelines.sd35.SD35Config
- fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
tokenizer_archs:
sources:
- fastvideo.configs.pipelines.sd35.SD35Config
- fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
transformer_arch:
sources:
- fastvideo.configs.pipelines.sd35.SD35Config
- fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
vae_arch:
sources:
- fastvideo.configs.pipelines.sd35.SD35Config
- fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
expand_timesteps:
sources:
- fastvideo.configs.pipelines.wan.FastWan2_2_TI2V_5B_Config
@@ -480,8 +440,6 @@ surfaces:
image_conditioning: "MagiHuman preset variant marker for reference-image conditioning."
pdd_step_indices: "MiniMax-H3 PDD fine-grid partition; resolve_checkpoint_settings sets it from fastvideo_inference.json."
ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset."
scheduler_sigma_min: "Z-Image scheduler parity invariant; not part of the public typed inference API."
scheduler_use_reference_discrete_timesteps: "Z-Image scheduler parity invariant; not part of the public typed inference API."
sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection."
text_offset: "MagiHuman internal data-proxy text alignment offset."
vae_stride: "MagiHuman internal VAE/data-proxy stride setting."
@@ -507,8 +465,6 @@ surfaces:
pose: request.inputs.pose
c2ws_plucker_emb: request.inputs.c2ws_plucker_emb
action_path: request.inputs.action_path
refine_from: request.inputs.refine_from
stage1_video: request.inputs.stage1_video
prompt: request.prompt
negative_prompt: request.negative_prompt
prompt_path: request.inputs.prompt_path
@@ -528,8 +484,6 @@ surfaces:
guidance_scale: request.sampling.guidance_scale
batch_cfg: request.sampling.batch_cfg
guidance_scale_2: request.sampling.guidance_scale_2
cfg_normalization: request.sampling.cfg_normalization
cfg_truncation: request.sampling.cfg_truncation
guidance_rescale: request.sampling.guidance_rescale
use_embedded_guidance: request.sampling.use_embedded_guidance
true_cfg_scale: request.sampling.true_cfg_scale
@@ -544,19 +498,12 @@ surfaces:
return_continuation_state: request.output.return_state
preset_owned:
t_thresh: request.stage_overrides.refine.t_thresh
spatial_refine_only: request.stage_overrides.refine.spatial_refine_only
num_cond_frames: request.stage_overrides.refine.num_cond_frames
trajectory_type: request.extensions.gen3c.trajectory_type
movement_distance: request.extensions.gen3c.movement_distance
camera_rotation: request.extensions.gen3c.camera_rotation
prompt_attention_mask: request.extensions.hyworld.prompt_attention_mask
negative_attention_mask: request.extensions.hyworld.negative_attention_mask
camera_states: request.extensions.hunyuangamecraft.camera_states
camera_trajectory: request.extensions.hunyuangamecraft.camera_trajectory
action_list: request.extensions.hunyuangamecraft.action_list
action_speed_list: request.extensions.hunyuangamecraft.action_speed_list
gt_latents: request.extensions.hunyuangamecraft.gt_latents
conditioning_mask: request.extensions.hunyuangamecraft.conditioning_mask
ltx2_cfg_scale_video: request.extensions.ltx2.cfg_scale_video
ltx2_cfg_scale_audio: request.extensions.ltx2.cfg_scale_audio
ltx2_modality_scale_video: request.extensions.ltx2.modality_scale_video
docs/generate_examples.py
-4
View File
@@ -38,18 +38,14 @@ COOKBOOK_SOURCE_ROOTS = (
# (e.g. black-forest-labs/FLUX.1-dev) and are grouped for documentation only.
COOKBOOK_FAMILIES = {
"wan",
"turbodiffusion",
"ltx2",
"hunyuan",
"cosmos",
"kandinsky5",
"kandinsky6",
"flux",
"glm_image",
"zimage",
"sd35",
"minimax_h3",
"longcat",
"stable_audio",
"mmaudio",
"matrixgame",
docs/inference/architecture.md
+1 -3
View File
@@ -232,8 +232,7 @@ Dataclass carrying all pipeline state between stages. Key field groups:
- **Scheduler**: `timesteps`, `num_inference_steps`, `guidance_scale`,
`sigmas`.
- **Task-specific**: `mouse_cond`/`keyboard_cond` (Matrix-Game 2.0), `pose`
(HYWorld), `camera_states` (GameCraft), `c2ws_plucker_emb`
(LingBotWorld).
(HYWorld), `c2ws_plucker_emb` (LingBotWorld).
- **Output**: `output: Tensor | None`.
- **Logging**: `logging_info: PipelineLoggingInfo`.
@@ -256,7 +255,6 @@ Standard stages (typical execution order):
| `DecodingStage` | `stages/decoding.py` | Decodes latents to video via VAE |
Specialized variants: `CausalDenoisingStage`, `LTX2DenoisingStage`,
`LongCatDenoisingStage`, `GameCraftDenoisingStage`,
`HYWorldDenoisingStage`, `MatrixGame2CausalDenoisingStage`,
`SRDenoisingStage`, `LTX2AudioDecodingStage`, `SD35ConditioningStage`,
`LTX2TextEncodingStage`, `LTX2LatentPreparationStage`.
docs/inference/offloading.md
+2 -2
View File
@@ -168,8 +168,8 @@ and logs a warning if the flag is set.
Deferral is opt-in per pipeline. Releasing a component and loading it again is
only safe when nothing outside the loader has changed it, and two common habits
break that without raising: mutating a component after load, as LongCat does
when it enables block-sparse attention, and reading a component's attributes
break that without raising: mutating a component after load, such as enabling
a sparse-attention mode in `initialize_pipeline`, and reading a component's attributes
while stages are built, as the shared denoising stage does to pick an attention
backend. A pipeline therefore lists the components it has checked in
`_lazy_module_names`, which is empty in the base class. MiniMax-H3 opts in. On
docs/inference/support_matrix.md
-22
View File
@@ -31,14 +31,10 @@ column links a runnable script in `examples/inference/basic/` where one exists.
| cosmos | `nvidia/Cosmos-Predict2-2B-Video2World` | T2V | — |
| cosmos25 | `KyleShao/Cosmos-Predict2.5-2B-Diffusers` | T2V | [basic_cosmos2_5_t2w.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_cosmos2_5_t2w.py) |
| cosmos25 | `nvidia/Cosmos-Predict2.5-14B` | T2V | [basic_cosmos2_5_t2w.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_cosmos2_5_t2w.py) |
| dreamx_world | `FastVideo/DreamX-World-5B-Cam-Diffusers` | I2V | [basic_dreamx_world.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dreamx_world.py) |
| dreamx_world | `FastVideo/DreamX-World-5B-Diffusers` | I2V | [basic_dreamx_world.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dreamx_world.py) |
| flux | `black-forest-labs/FLUX.1-dev` | T2I | [basic_flux_dev.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_flux_dev.py) |
| flux2 | `black-forest-labs/FLUX.2-klein-4B`<br>`black-forest-labs/FLUX.2-klein-9B` | T2I | [basic_flux2_klein.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_flux2_klein.py) |
| flux2 | `black-forest-labs/FLUX.2-dev` | T2I | [basic_flux2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_flux2.py) |
| gamecraft | `FastVideo/HunyuanGameCraft-Diffusers` | I2V | [basic_gamecraft.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_gamecraft.py) |
| gen3c | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | T2V | [basic_gen3c.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_gen3c.py) |
| glm_image | `zai-org/GLM-Image` | T2I | [basic_glm_image.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_glm_image.py) |
| hunyuan | `hunyuanvideo-community/HunyuanVideo` | T2V | — |
| hunyuan | `FastVideo/FastHunyuan-diffusers` | T2V | — |
| hunyuan15 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | T2V | [basic_hy15.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hy15.py) |
@@ -61,9 +57,6 @@ column links a runnable script in `examples/inference/basic/` where one exists.
| lingbot_video | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) |
| lingbotworld | `FastVideo/LingBot-World-Base-Cam-Diffusers` | I2V | [basic_lingbotworld_base_cam.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld_base_cam.py) |
| lingbotworld2 | `robbyant/lingbot-world-v2-14b-causal-fast` | I2V | [basic_lingbotworld2_causal_fast.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld2_causal_fast.py) |
| longcat | `FastVideo/LongCat-Video-T2V-Diffusers` | T2V | [basic_longcat_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_t2v.py) |
| longcat | `FastVideo/LongCat-Video-I2V-Diffusers` | I2V | [basic_longcat_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_i2v.py) |
| longcat | `FastVideo/LongCat-Video-VC-Diffusers` | — | [basic_longcat_vc.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_vc.py) |
| ltx2 | `FastVideo/LTX2-Distilled-Diffusers`<br>`FastVideo/LTX2.3-Distilled-Diffusers`<br>`FastVideo/LTX-2.3-Distilled-Diffusers` | T2V | [basic_ltx2_distilled.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_distilled.py) |
| ltx2 | `Lightricks/LTX-2.3`<br>`FastVideo/LTX2.3-base`<br>`FastVideo/LTX2.3-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
| ltx2 | `Lightricks/LTX-2`<br>`FastVideo/LTX2-base`<br>`FastVideo/LTX2-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
@@ -74,9 +67,6 @@ column links a runnable script in `examples/inference/basic/` where one exists.
| sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) |
| stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) |
| stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) |
| turbodiffusion | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | T2V | [basic_turbodiffusion.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_turbodiffusion.py) |
| turbodiffusion | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | T2V | [basic_turbodiffusion_14b.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_turbodiffusion_14b.py) |
| turbodiffusion | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | I2V | [basic_turbodiffusion_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_turbodiffusion_i2v.py) |
| wan | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | T2V | [basic.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic.py) |
| wan | `Wan-AI/Wan2.1-T2V-14B-Diffusers`<br>`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` | T2V | — |
| wan | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | I2V | — |
@@ -92,7 +82,6 @@ column links a runnable script in `examples/inference/basic/` where one exists.
| wan | `wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers` | T2V | [basic_self_forcing_causal.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal.py) |
| wan | `rand0nmr/SFWan2.2-T2V-A14B-Diffusers` | T2V | [basic_self_forcing_causal_wan2_2_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py) |
| wan | `FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers` | I2V | [basic_self_forcing_causal_wan2_2_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py) |
| zimage | `Tongyi-MAI/Z-Image-Turbo` | T2I | [basic_zimage.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_zimage.py) |
**Note (stable_audio)**: the Stable Audio Open pipelines generate audio
(`StableAudioT2AConfig` / `StableAudioOpenSmallConfig`); they are registered
@@ -168,8 +157,6 @@ optimizations: absence means **untested**, not incompatible.
| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
| DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
@@ -179,10 +166,6 @@ optimizations: absence means **untested**, not incompatible.
| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
| TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| LongCat T2V 13.6B | `FastVideo/LongCat-Video-T2V-Diffusers` | 480P<br>720P | ❌ | ❌ | ❌ | ⭕ | ✅ |
| Matrix Game 2.0 Base Distilled | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| Matrix Game 2.0 GTA Distilled | `FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| Matrix Game 2.0 TempleRun Distilled | `FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -251,11 +234,6 @@ listed under [Special requirements](#special-requirements).
https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete
- STA currently requires Hopper GPUs (H100s).
### TurboWan2.1 (TurboDiffusion)
- Uses TurboDiffusionPipeline with RCM scheduler for 1-4 step generation
- Requires SLA attention backend: `export FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN`
- Uses `guidance_scale=1.0` (no classifier-free guidance)
### Matrix Game 2.0
- Image-to-video game world models with keyboard/mouse control input
- Three variants available: Base (universal), GTA, and TempleRun
examples/inference/basic/basic_dreamx_world.py
-61
View File
@@ -1,61 +0,0 @@
import os
from fastvideo import VideoGenerator
OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world")
def _env_int(name: str, default: int) -> int:
return int(os.getenv(name, str(default)))
def _env_float(name: str, default: float) -> float:
return float(os.getenv(name, str(default)))
def main():
model_name = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
generator = VideoGenerator.from_pretrained(
model_name,
num_gpus=1,
use_fsdp_inference=False,
dit_cpu_offload=False,
vae_cpu_offload=True,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
override_pipeline_cls_name="DreamXWorldPipeline",
)
prompt = os.getenv(
"DREAMX_WORLD_PROMPT",
"A cinematic first-person drive through a futuristic coastal city at "
"sunrise, reflective glass towers, clean streets, soft volumetric light.",
)
image_path = os.getenv(
"DREAMX_WORLD_IMAGE_PATH",
"https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG",
)
kwargs = {
"output_path": OUTPUT_PATH,
"save_video": os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
"height": _env_int("DREAMX_WORLD_HEIGHT", 480),
"width": _env_int("DREAMX_WORLD_WIDTH", 832),
"num_frames": _env_int("DREAMX_WORLD_NUM_FRAMES", 161),
"num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30),
"guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0),
"action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
"action_speed_list":
[float(value) for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")],
}
if image_path:
kwargs["image_path"] = image_path
try:
generator.generate_video(prompt, **kwargs)
finally:
generator.shutdown()
if __name__ == "__main__":
main()
examples/inference/basic/basic_gamecraft.py
-118
View File
@@ -1,118 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
Basic inference script for HunyuanGameCraft video generation.
HunyuanGameCraft generates game-like videos with camera/action control.
It takes an optional image input and generates video with camera motion
based on simple action commands (forward, left, right, backward, rotations).
Available actions:
- forward (w): Move camera forward
- backward (s): Move camera backward
- left (a): Move camera left (strafe)
- right (d): Move camera right (strafe)
- left_rot: Rotate camera left (pan)
- right_rot: Rotate camera right (pan)
- up_rot: Rotate camera up (tilt)
- down_rot: Rotate camera down (tilt)
T2V vs I2V:
- Default: I2V (uses a default reference image). Set GAMECRAFT_I2V_IMAGE to a
URL or path to use a different image.
- T2V only (no reference image): run with GAMECRAFT_I2V_IMAGE= (empty).
"""
import os
import torch
from fastvideo import VideoGenerator
from fastvideo.models.camera import create_camera_trajectory
# Model configuration (use GAMECRAFT_MODEL_PATH for local weights)
MODEL_PATH = os.environ.get("GAMECRAFT_MODEL_PATH", "FastVideo/HunyuanGameCraft-Diffusers")
# Default prompts for demo
DEFAULT_PROMPTS = {
"village":
"A charming medieval village with cobblestone streets, thatched-roof houses, and vibrant flower gardens under a bright blue sky.",
"temple":
"A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.",
"forest":
"A lush green forest with tall trees, dappled sunlight filtering through the leaves, and a winding dirt path.",
"beach": "A tropical beach with crystal clear turquoise water, white sand, and palm trees swaying in the breeze.",
}
# I2V: default reference image (URL). Can override with a local path.
DEFAULT_I2V_IMAGE_URL = ("https://huggingface.co/datasets/huggingface/documentation-images/"
"resolve/main/diffusers/astronaut.jpg")
DEFAULT_I2V_PROMPT = ("An astronaut hatching from an egg, on the surface of the moon, "
"the darkness and depth of space realised in the background.")
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,
)
# Video parameters
height = 704
width = 1280
num_frames = 33
action = "forward"
action_speed = 0.2
# Create camera trajectory (Plücker coordinates)
camera_states = create_camera_trajectory(
action=action,
height=height,
width=width,
num_frames=num_frames,
action_speed=action_speed,
dtype=torch.bfloat16,
)
print(f"Camera states shape: {camera_states.shape}")
# I2V vs T2V: unset GAMECRAFT_I2V_IMAGE -> I2V (default image). Set to "" -> T2V.
env_image = os.environ.get("GAMECRAFT_I2V_IMAGE")
if env_image is None:
image_path = DEFAULT_I2V_IMAGE_URL # default: I2V
elif env_image.strip() == "":
image_path = None # T2V
else:
image_path = env_image.strip() # I2V with given URL/path
is_i2v = image_path is not None
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(
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,
)
if is_i2v:
gen_kw["image_path"] = image_path
generator.generate_video(**gen_kw)
if __name__ == "__main__":
main()
examples/inference/basic/basic_glm_image.py
-107
View File
@@ -1,107 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Run GLM-Image text-to-image generation through FastVideo.
User story:
"I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
text-to-image generation command, saved as a PNG."
"""
import argparse
from pathlib import Path
from PIL import Image
from fastvideo import VideoGenerator
from fastvideo.api import (
EngineConfig,
GenerationRequest,
GeneratorConfig,
OutputConfig,
ParallelismConfig,
PipelineSelection,
SamplingConfig,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
parser.add_argument(
"--model-path",
default="zai-org/GLM-Image",
help="HF id or local diffusers-format GLM-Image weights directory.",
)
parser.add_argument(
"--output",
default="image_output/landscape.png",
help="Output PNG path.",
)
parser.add_argument(
"--prompt",
default=("A beautiful landscape photography with rolling hills, "
"a winding river, and a vibrant sunset in the background. "
"Warm golden light, photorealistic style."),
help="Text prompt.",
)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--steps", type=int, default=50)
parser.add_argument("--guidance-scale", type=float, default=1.5)
parser.add_argument("--seed", type=int, default=1024)
parser.add_argument("--num-gpus", type=int, default=1)
parser.add_argument("--tp-size", type=int, default=None)
parser.add_argument("--sp-size", type=int, default=None)
return parser.parse_args()
def main() -> None:
args = parse_args()
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
# pipeline class come from the model's registered defaults — don't override.
generator_config = GeneratorConfig(
model_path=args.model_path,
trust_remote_code=True,
engine=EngineConfig(
num_gpus=args.num_gpus,
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
),
pipeline=PipelineSelection(workload_type="t2i"),
)
generator = VideoGenerator.from_config(generator_config)
try:
request = GenerationRequest(
prompt=args.prompt,
sampling=SamplingConfig(
height=args.height,
width=args.width,
num_frames=1,
fps=1,
num_inference_steps=args.steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
),
output=OutputConfig(
output_path=str(output.parent),
save_video=False,
return_frames=True,
),
)
result = generator.generate(request)
if isinstance(result, list):
result = result[0]
frames = result.frames
if frames is not None and len(frames):
Image.fromarray(frames[0]).save(output)
print(f"Saved image to {output}")
finally:
generator.shutdown()
if __name__ == "__main__":
main()
examples/inference/basic/basic_longcat_i2v.py
-204
View File
@@ -1,204 +0,0 @@
"""
LongCat Image-to-Video (I2V) Example Script
This script demonstrates LongCat I2V inference using the FastVideo Python API.
LongCat I2V takes an input image and generates a video from it.
It runs both basic generation (50 steps) and distill+refine generation
(16 steps distill + 50 steps refinement to 720p with BSA).
Usage:
python examples/inference/basic/basic_longcat_i2v.py
Note:
Refinement uses 768x768 dimensions where latent (48x48) is divisible by 8,
compatible with BSA chunks [4, 4, 8].
"""
import glob
import os
from fastvideo import VideoGenerator
# Common prompts and settings matching the shell script examples
PROMPT = ("A woman sits at a wooden table by the window in a cozy café. She reaches out "
"with her right hand, picks up the white coffee cup from the saucer, and gently "
"brings it to her lips to take a sip. After drinking, she places the cup back on "
"the table and looks out the window, enjoying the peaceful atmosphere.")
NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, "
"paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
"residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, "
"deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, "
"three legs, many people in the background, walking backwards")
# Input image path
IMAGE_PATH = "assets/girl.png"
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,
)
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,
)
print(f"\nBasic generation complete! Video saved to: {output_path}")
generator.shutdown()
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",
)
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,
)
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:
raise FileNotFoundError(f"No video file found in {distill_output_path}")
# 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_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,
)
print(f"Refinement complete! Video saved to: {refine_output_path}")
refine_generator.shutdown()
def main():
"""Run both basic and distill+refine generation pipelines."""
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)
if __name__ == "__main__":
main()
examples/inference/basic/basic_longcat_t2v.py
-192
View File
@@ -1,192 +0,0 @@
"""
LongCat Text-to-Video (T2V) Example Script
This script demonstrates LongCat T2V inference using the FastVideo Python API.
It runs both basic generation (50 steps) and distill+refine generation
(16 steps distill + 50 steps refinement to 720p).
Usage:
python examples/inference/basic/basic_longcat_t2v.py
"""
import glob
import os
from fastvideo import VideoGenerator
# Common prompts and settings matching the shell script examples
PROMPT = ("In a realistic photography style, a white boy around seven or eight years old "
"sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. "
"He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a "
"medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, "
"who eagerly licks it with its tongue. The sun is shining brightly, and the background "
"features a green lawn and several tall trees, creating a warm and loving scene.")
NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, "
"paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
"residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, "
"deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, "
"three legs, many people in the background, walking backwards")
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,
)
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,
)
print(f"\nBasic generation complete! Video saved to: {output_path}")
generator.shutdown()
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",
)
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,
)
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:
raise FileNotFoundError(f"No video file found in {distill_output_path}")
# 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_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,
)
print(f"Refinement complete! Video saved to: {refine_output_path}")
refine_generator.shutdown()
def main():
"""Run both basic and distill+refine generation pipelines."""
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)
if __name__ == "__main__":
main()
examples/inference/basic/basic_longcat_vc.py
-218
View File
@@ -1,218 +0,0 @@
"""
LongCat Video Continuation (VC) Example Script
This script demonstrates LongCat VC inference using the FastVideo Python API.
LongCat VC takes an input video and generates a continuation of it.
It runs both basic generation (50 steps) and distill+refine generation
(16 steps distill + 50 steps refinement to 720p).
Usage:
python examples/inference/basic/basic_longcat_vc.py
Prerequisites:
- Ensure the input video exists at assets/motorcycle.mp4
(or provide your own video)
"""
import glob
import os
from fastvideo import VideoGenerator
# Common prompts and settings matching the shell script examples
PROMPT = ("A person rides a motorcycle along a long, straight road that stretches between "
"a body of water and a forested hillside. The rider steadily accelerates, keeping "
"the motorcycle centered between the guardrails, while the scenery passes by on "
"both sides. The video captures the journey from the rider's perspective, emphasizing "
"the sense of motion and adventure.")
NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, "
"paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
"residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, "
"deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, "
"three legs, many people in the background, walking backwards")
# Input video path
VIDEO_PATH = "assets/motorcycle.mp4"
# Number of conditioning frames from the input video
NUM_COND_FRAMES = 13
SEED = 42
def basic_generation():
"""
Run basic LongCat VC generation (50 steps at 480p).
This uses the full 50-step denoising process for highest quality.
"""
print("=" * 60)
print("LongCat VC: Basic Generation (50 steps, 480p)")
print("=" * 60)
# Check if video exists
if not os.path.exists(VIDEO_PATH):
raise FileNotFoundError(f"Video not found at {VIDEO_PATH}. "
"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,
)
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,
)
print(f"\nBasic generation complete! Video saved to: {output_path}")
generator.shutdown()
def distill_refine_generation():
"""
Run LongCat VC 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 VC: Distill + Refine Pipeline")
print("=" * 60)
# Check if video exists
if not os.path.exists(VIDEO_PATH):
raise FileNotFoundError(f"Video not found at {VIDEO_PATH}. "
"Please provide a valid video path.")
# 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-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",
)
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,
)
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:
raise FileNotFoundError(f"No video file found in {distill_output_path}")
# 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 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_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,
)
print(f"Refinement complete! Video saved to: {refine_output_path}")
refine_generator.shutdown()
def main():
"""Run both basic and distill+refine generation pipelines."""
print("\n" + "=" * 60)
print("LongCat Video Continuation 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)
if __name__ == "__main__":
main()
examples/inference/basic/basic_turbodiffusion.py
-51
View File
@@ -1,51 +0,0 @@
import os
# Set SLA attention backend BEFORE fastvideo imports
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
from fastvideo import VideoGenerator
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,
)
# Generate videos with the same simple API, regardless of GPU count
# TurboDiffusion defaults: guidance_scale=1.0 and num_inference_steps=4 (from 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,
seed=42,
)
# Generate another video with a different prompt, without reloading the model!
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,
output_path=OUTPUT_PATH,
save_video=True,
seed=42,
)
if __name__ == "__main__":
main()
examples/inference/basic/basic_turbodiffusion_14b.py
-45
View File
@@ -1,45 +0,0 @@
import os
# Set SLA attention backend BEFORE fastvideo imports
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
from fastvideo import VideoGenerator
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,
)
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,
seed=42,
)
# Generate another video with a different prompt, without reloading the model!
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,
output_path=OUTPUT_PATH,
save_video=True,
seed=42,
)
if __name__ == "__main__":
main()
examples/inference/basic/basic_turbodiffusion_i2v.py
-38
View File
@@ -1,38 +0,0 @@
import os
# Set SLA attention backend BEFORE fastvideo imports
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
from fastvideo import VideoGenerator
# Use local model path
MODEL_PATH = "loayrashid/TurboWan2.2-I2V-A14B-Diffusers"
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,
)
# Example prompt and image for I2V
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."
)
# Use an example image path
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,
)
if __name__ == "__main__":
main()
examples/inference/basic/basic_zimage.py
-94
View File
@@ -1,94 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Run Z-Image-Turbo text-to-image generation through FastVideo.
User story:
"I want the official Z-Image-Turbo defaults and a deterministic PNG from
a local or Hugging Face checkpoint."
"""
import argparse
from pathlib import Path
from fastvideo import VideoGenerator
from fastvideo.api import (
EngineConfig,
GenerationRequest,
GeneratorConfig,
OutputConfig,
ParallelismConfig,
PipelineSelection,
SamplingConfig,
)
DEFAULT_PROMPT = (
"Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. "
"Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. "
"Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, "
"silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.")
DEFAULT_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run Z-Image-Turbo text-to-image generation.")
parser.add_argument("--model-path", default="Tongyi-MAI/Z-Image-Turbo")
parser.add_argument("--revision", default=DEFAULT_REVISION)
parser.add_argument("--output", default="outputs/zimage/zimage_turbo.png")
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--negative-prompt", default="")
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--steps", type=int, default=8)
parser.add_argument("--guidance-scale", type=float, default=0.0)
parser.add_argument("--max-sequence-length", type=int, default=512)
parser.add_argument("--cfg-normalization", action=argparse.BooleanOptionalAction, default=False)
parser.add_argument("--cfg-truncation", type=float, default=1.0)
parser.add_argument("--seed", type=int, default=42)
return parser.parse_args()
def main() -> None:
args = parse_args()
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
generator = VideoGenerator.from_config(
GeneratorConfig(
model_path=args.model_path,
revision=args.revision,
engine=EngineConfig(
num_gpus=1,
parallelism=ParallelismConfig(tp_size=1, sp_size=1),
use_fsdp_inference=False,
),
# The model registry selects the native zimage_turbo preset.
pipeline=PipelineSelection(workload_type="t2i"),
))
try:
generator.generate(
GenerationRequest(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
sampling=SamplingConfig(
height=args.height,
width=args.width,
num_frames=1,
fps=1,
num_inference_steps=args.steps,
guidance_scale=args.guidance_scale,
max_sequence_length=args.max_sequence_length,
cfg_normalization=args.cfg_normalization,
cfg_truncation=args.cfg_truncation,
seed=args.seed,
),
output=OutputConfig(
output_path=str(output),
save_video=True,
return_frames=False,
),
))
finally:
generator.shutdown()
if __name__ == "__main__":
main()
examples/inference/basic/edit_glm_image.py
-120
View File
@@ -1,120 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Run GLM-Image image-to-image (edit) generation through FastVideo.
User story:
"I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
want a minimal edit command (text + image -> edited image), saved as a PNG."
GLM-Image is a single unified pipeline: passing a condition image switches it
from text-to-image to the edit path (the condition enters the DiT via a KV-cache
write pass), so the generator config is identical to `basic_glm_image.py` — the
`inputs.pil_image` on the request is what selects the edit mode.
"""
import argparse
from pathlib import Path
from PIL import Image
from fastvideo import VideoGenerator
from fastvideo.api import (
EngineConfig,
GenerationRequest,
GeneratorConfig,
InputConfig,
OutputConfig,
ParallelismConfig,
PipelineSelection,
SamplingConfig,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
parser.add_argument(
"--model-path",
default="zai-org/GLM-Image",
help="HF id or local diffusers-format GLM-Image weights directory.",
)
parser.add_argument(
"--image",
default="assets/images/couple.jpg",
help="Condition image to edit.",
)
parser.add_argument(
"--output",
default="image_output/edited.png",
help="Output PNG path.",
)
parser.add_argument(
"--prompt",
default="Change the background to a snowy mountain landscape at golden hour.",
help="Edit instruction.",
)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--steps", type=int, default=50)
parser.add_argument("--guidance-scale", type=float, default=1.5)
parser.add_argument("--seed", type=int, default=1024)
parser.add_argument("--num-gpus", type=int, default=1)
parser.add_argument("--tp-size", type=int, default=None)
parser.add_argument("--sp-size", type=int, default=None)
return parser.parse_args()
def main() -> None:
args = parse_args()
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
condition = Image.open(args.image).convert("RGB")
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
# pipeline class come from the model's registered defaults — don't override.
# The pipeline is registered as t2i; passing inputs.pil_image below switches
# it to the edit path.
generator_config = GeneratorConfig(
model_path=args.model_path,
trust_remote_code=True,
engine=EngineConfig(
num_gpus=args.num_gpus,
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
),
pipeline=PipelineSelection(workload_type="t2i"),
)
generator = VideoGenerator.from_config(generator_config)
try:
request = GenerationRequest(
prompt=args.prompt,
inputs=InputConfig(pil_image=condition),
sampling=SamplingConfig(
height=args.height,
width=args.width,
num_frames=1,
fps=1,
num_inference_steps=args.steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
),
output=OutputConfig(
output_path=str(output.parent),
save_video=False,
return_frames=True,
),
)
result = generator.generate(request)
if isinstance(result, list):
result = result[0]
frames = result.frames
if frames is not None and len(frames):
Image.fromarray(frames[0]).save(output)
print(f"Saved image to {output}")
finally:
generator.shutdown()
if __name__ == "__main__":
main()
examples/train/configs/fine_tuning/longcat/t2v.yaml
-71
View File
@@ -1,71 +0,0 @@
# LongCat-Video T2V 13.6B bidirectional finetune.
models:
student:
_target_: fastvideo.train.models.longcat.LongCatModel
init_from: FastVideo/LongCat-Video-T2V-Diffusers
trainable: true
method:
_target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
training:
model_path: FastVideo/LongCat-Video-T2V-Diffusers
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path: data/LongCat-Syn
dataloader_num_workers: 4
train_batch_size: 1
training_cfg_rate: 0.0
seed: 1000
num_latent_t: 20
num_height: 480
num_width: 848
num_frames: 77
optimizer:
learning_rate: 1.0e-6
betas: [0.9, 0.999]
weight_decay: 0.01
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 4000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/longcat_finetune
training_state_checkpointing_steps: 1000
checkpoints_total_limit: 3
tracker:
project_name: fastvideo
run_name: longcat_finetune
model:
enable_gradient_checkpointing_type: full
callbacks:
grad_clip:
_target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
max_grad_norm: 1.0
validation:
_target_: fastvideo.train.callbacks.validation.ValidationCallback
pipeline_target: fastvideo.pipelines.basic.longcat.longcat_pipeline.LongCatPipeline
dataset_file: data/validation_prompts.json
every_steps: 100
sampling_steps: [50]
guidance_scale: 5.0
pipeline:
# Match the released LongCat scheduler config. flow_shift=0.0 collapses
# FlowMatch training timesteps to zero in FastVideo's scheduler.
flow_shift: 12.0
examples/train/configs/fine_tuning/wan/turbo_t2v_lora_sla.yaml
-89
View File
@@ -1,89 +0,0 @@
# TurboWan 2.1 T2V 1.3B LoRA finetune with SLA attention.
#
# Notes:
# - TurboDiffusion checkpoints expect the SLA attention backend. The training
# entrypoint now auto-selects `SLA_ATTN` for model IDs containing
# `turbodiffusion` or `turbowan`, unless you override the environment
# variable manually.
# - This debug config reuses the existing crush-smol T2V parquet recipe. If you
# preprocess data with the official TurboWan recipe, update the temporal and
# spatial fields to match it.
models:
student:
_target_: fastvideo.train.models.wan.WanModel
init_from: loayrashid/TurboWan2.1-T2V-1.3B-Diffusers
trainable: true
lora:
enable: true
rank: 16
alpha: 32
target_modules:
- to_q
- to_k
- to_v
- to_out
method:
_target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
training:
dit_precision: bf16
distributed:
num_gpus: 1
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 1
data:
data_path: data/crush-smol_processed_t2v/combined_parquet_dataset
dataloader_num_workers: 2
train_batch_size: 1
training_cfg_rate: 0.0
seed: 42
num_latent_t: 20
num_height: 480
num_width: 832
num_frames: 77
optimizer:
learning_rate: 1.0e-5
betas: [0.9, 0.999]
weight_decay: 0.01
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 10
gradient_accumulation_steps: 5
checkpoint:
output_dir: outputs/turbo_wan_t2v_lora_sla
training_state_checkpointing_steps: 5
checkpoints_total_limit: 2
resume_from_checkpoint: null
tracker:
project_name: fastvideo
run_name: turbo_wan_t2v_lora_sla
model:
enable_gradient_checkpointing_type: full
callbacks:
grad_clip:
_target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
max_grad_norm: 1.0
validation:
_target_: fastvideo.train.callbacks.validation.ValidationCallback
pipeline_target: fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_pipeline.TurboDiffusionPipeline
dataset_file: examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json
every_steps: 5
sampling_steps: [4]
guidance_scale: 1.0
pipeline:
flow_shift: 3
fastvideo/api/sampling_param.py
+1 -45
View File
@@ -46,28 +46,13 @@ class SamplingParam:
prompt_attention_mask: list = field(default_factory=list)
negative_attention_mask: list = field(default_factory=list)
# Camera/action control inputs (GameCraft)
camera_states: Any | None = None # Plücker coordinates [B, T_video, 6, H, W]
camera_trajectory: str | None = None
action_list: list[str] | None = None
action_speed_list: list[float] | None = None
gt_latents: Any | None = None # Ground truth latents [B, 16, T, H, W]
conditioning_mask: Any | None = None # Mask [B, 1, T, H, W]
# Camera control inputs (LingBotWorld and LingBotWorld2)
c2ws_plucker_emb: Any | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
action_path: str | None = None # Directory containing poses.npy and intrinsics.npy
# Refine inputs (LongCat 480p->720p upscaling)
# Path-based refine (load stage1 video from disk, e.g. MP4)
refine_from: str | None = None # Path to stage1 video (480p output from distill)
# Refine / conditioning-frame inputs
t_thresh: float = 0.5 # Threshold for timestep scheduling in refinement
spatial_refine_only: bool = False # If True, only spatial (no temporal doubling)
num_cond_frames: int = 0 # Number of conditioning frames
# In-memory refine input (for two-stage pipeline where stage1 frames are already in memory)
# This mirrors LongCat's demo where a list of frames (e.g. np.ndarray or PIL.Image)
# is passed directly to the refinement pipeline instead of reloading from disk.
stage1_video: Any | None = None
# Text inputs
prompt: str | list[str] | None = None
@@ -95,11 +80,6 @@ class SamplingParam:
guidance_scale: float = 1.0
batch_cfg: bool = False
guidance_scale_2: float | None = None
# Z-Image CFG controls. ``cfg_normalization=True`` caps the guided
# prediction norm at the positive-prediction norm; ``cfg_truncation``
# disables CFG above the normalized-noise threshold.
cfg_normalization: bool = False
cfg_truncation: float | None = 1.0
# 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).
@@ -333,18 +313,6 @@ class SamplingParam:
default=SamplingParam.guidance_scale,
help="Classifier-free guidance scale",
)
parser.add_argument(
"--cfg-normalization",
action=StoreBoolean,
default=SamplingParam.cfg_normalization,
help="Cap Z-Image CFG prediction norm to the positive-prediction norm",
)
parser.add_argument(
"--cfg-truncation",
type=float,
default=SamplingParam.cfg_truncation,
help="Disable Z-Image CFG above this normalized-noise threshold",
)
parser.add_argument(
"--batch-cfg",
action=StoreBoolean,
@@ -405,24 +373,12 @@ class SamplingParam:
default=SamplingParam.video_path,
help="Path to input video for video-to-video generation",
)
parser.add_argument(
"--refine-from",
type=str,
default=SamplingParam.refine_from,
help="Path to stage1 video for refinement (LongCat 480p->720p)",
)
parser.add_argument(
"--t-thresh",
type=float,
default=SamplingParam.t_thresh,
help="Threshold for timestep scheduling in refinement (default: 0.5)",
)
parser.add_argument(
"--spatial-refine-only",
action=StoreBoolean,
default=SamplingParam.spatial_refine_only,
help="Only perform spatial super-resolution (no temporal doubling)",
)
parser.add_argument(
"--num-cond-frames",
type=int,
fastvideo/api/schema.py
-4
View File
@@ -149,8 +149,6 @@ class InputConfig:
grid_sizes: Any | None = None
c2ws_plucker_emb: Any | None = None
action_path: str | None = None
refine_from: str | None = None
stage1_video: Any | None = None
@dataclass
@@ -169,8 +167,6 @@ class SamplingConfig:
guidance_scale: float = 1.0
batch_cfg: bool = False
guidance_scale_2: float | None = None
cfg_normalization: bool = False
cfg_truncation: float | None = 1.0
guidance_rescale: float = 0.0
true_cfg_scale: float | None = None
use_embedded_guidance: bool | None = None
fastvideo/attention/utils/flash_attn_cute.py
+1 -1
View File
@@ -67,7 +67,7 @@ def _use_fa2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:
# capability gate, not a runtime fallback):
# * the backward asserts sm90+ (L40S/sm_89 dies on its arch check);
# * GQA fails CuTeDSL JIT in pack_gqa ("ValueError: Operation creation
# failed", observed on sm_89 with HunyuanGameCraft/LTX2).
# failed", observed on sm_89 with LTX2).
if q.shape[-2] != k.shape[-2]:
return True
return torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v))
fastvideo/configs/models/dits/__init__.py
+4 -10
View File
@@ -1,20 +1,15 @@
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
from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
from fastvideo.configs.models.dits.mmaudio import MMAudioArchConfig, MMAudioTransformerConfig
from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
from fastvideo.models.wan.config import WanVideoConfig
from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
from fastvideo.configs.models.dits.hyworld import HYWorldConfig
from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig, Kandinsky6VideoAudioConfig
@@ -23,10 +18,9 @@ from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastV
from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig
__all__ = [
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
"DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
"LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig",
"HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
"FluxDiTConfig", "Flux2Config", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig",
"Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig",
"StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig",
"MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig"
"StableAudioConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config",
"MMAudioArchConfig", "MMAudioTransformerConfig"
]
fastvideo/configs/models/dits/dreamx_world.py
-71
View File
@@ -1,71 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.models.wan.config import WanVideoArchConfig
@dataclass
class DreamXWorldArchConfig(WanVideoArchConfig):
"""DreamX-World DiT config with camera PRoPE control fields."""
add_control_adapter: bool = True
cam_method: str | None = "prope"
attn_compress: int = 1
cam_self_attn_layers: tuple[int, ...] | None = None
@dataclass
class DreamXWorldConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=DreamXWorldArchConfig)
prefix: str = "Wan"
@dataclass
class DreamXWorldARArchConfig(DreamXWorldArchConfig):
"""DreamX-World-5B autoregressive causal DiT config."""
model_type: str = "ti2v"
patch_size: tuple[int, int, int] = (1, 2, 2)
text_len: int = 512
text_dim: int = 4096
freq_dim: int = 256
attn_compress: int = 4
cam_self_attn_layers: tuple[int, ...] | None = tuple(range(30))
local_attn_size: int = 12
sink_size: int = 3
num_frames_per_block: int = 3
rope_cache_policy: str = "block_relativistic"
# The official AR checkpoint (AMAP-ML/DreamX-World ``model.safetensors``)
# already uses FastVideo's native key names and the converter copies the
# tensors verbatim, so every rule is an identity. The rules enumerate the
# full state-dict surface of ``DreamXWorldARTransformer3DModel`` (norm1 /
# norm2 / head.norm are affine-free and have no parameters).
param_names_mapping: dict = field(
default_factory=lambda: {
r"^patch_embedding\.(.*)$": r"patch_embedding.\1",
r"^text_embedding\.([02])\.(.*)$": r"text_embedding.\1.\2",
r"^time_embedding\.([02])\.(.*)$": r"time_embedding.\1.\2",
r"^time_projection\.1\.(.*)$": r"time_projection.1.\1",
r"^blocks\.(\d+)\.self_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.self_attn.\2.\3",
r"^blocks\.(\d+)\.self_attn\.norm_(q|k)\.weight$": r"blocks.\1.self_attn.norm_\2.weight",
r"^blocks\.(\d+)\.cross_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.cross_attn.\2.\3",
r"^blocks\.(\d+)\.cross_attn\.norm_(q|k)\.weight$": r"blocks.\1.cross_attn.norm_\2.weight",
r"^blocks\.(\d+)\.cam_self_attn\.(q_proj|k_proj|v_proj|out_proj)\.(.*)$": r"blocks.\1.cam_self_attn.\2.\3",
r"^blocks\.(\d+)\.cam_self_attn\.norm_(q|k)\.weight$": r"blocks.\1.cam_self_attn.norm_\2.weight",
r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm3.\2",
r"^blocks\.(\d+)\.ffn\.([02])\.(.*)$": r"blocks.\1.ffn.\2.\3",
r"^blocks\.(\d+)\.modulation$": r"blocks.\1.modulation",
r"^head\.head\.(.*)$": r"head.head.\1",
r"^head\.modulation$": r"head.modulation",
})
reverse_param_names_mapping: dict = field(default_factory=dict)
lora_param_names_mapping: dict = field(default_factory=dict)
@dataclass
class DreamXWorldARConfig(DreamXWorldConfig):
arch_config: DiTArchConfig = field(default_factory=DreamXWorldARArchConfig)
prefix: str = "Wan"
fastvideo/configs/models/dits/glm_image.py
-61
View File
@@ -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"
fastvideo/configs/models/dits/hunyuangamecraft.py
-143
View File
@@ -1,143 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
Configuration for HunyuanGameCraft transformer model.
HunyuanGameCraft extends HunyuanVideo with:
1. CameraNet for camera/action conditioning
2. 33 input channels (16 latent + 16 gt_latent + 1 mask)
3. Mask-based conditioning for autoregressive generation
"""
from dataclasses import dataclass, field
import torch
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
def is_double_block(n: str, m) -> bool:
return "double" in n and str.isdigit(n.split(".")[-1])
def is_single_block(n: str, m) -> bool:
return "single" in n and str.isdigit(n.split(".")[-1])
def is_refiner_block(n: str, m) -> bool:
return "refiner" in n and str.isdigit(n.split(".")[-1])
def is_txt_in(n: str, m) -> bool:
return n.split(".")[-1] == "txt_in"
def is_camera_net(n: str, m) -> bool:
return "camera_net" in n
@dataclass
class HunyuanGameCraftArchConfig(DiTArchConfig):
"""Architecture config for HunyuanGameCraft transformer."""
# Version field for compatibility with saved config.json
_fastvideo_version: str = "0.1.0"
# Camera net flag (for config.json compatibility)
camera_net: bool = True
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_double_block, is_single_block, is_refiner_block, is_camera_net])
_compile_conditions: list = field(default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
# Parameter names mapping from official checkpoint to FastVideo naming
# GameCraft weights are already close to FastVideo format with minor adjustments
param_names_mapping: dict = field(
default_factory=lambda: {
# MLP naming: fc1 -> fc_in, fc2 -> fc_out
r"^(.*)\.img_mlp\.fc1\.(.*)$": r"\1.img_mlp.fc_in.\2",
r"^(.*)\.img_mlp\.fc2\.(.*)$": r"\1.img_mlp.fc_out.\2",
r"^(.*)\.txt_mlp\.fc1\.(.*)$": r"\1.txt_mlp.fc_in.\2",
r"^(.*)\.txt_mlp\.fc2\.(.*)$": r"\1.txt_mlp.fc_out.\2",
# Single block MLP naming
r"^single_blocks\.(\d+)\.mlp\.fc1\.(.*)$": r"single_blocks.\1.mlp.fc_in.\2",
r"^single_blocks\.(\d+)\.mlp\.fc2\.(.*)$": r"single_blocks.\1.mlp.fc_out.\2",
# Token refiner naming
r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.(.*)$": r"txt_in.refiner_blocks.\1.\2",
# Vector in naming
r"^vector_in\.in_layer\.(.*)$": r"vector_in.fc_in.\1",
r"^vector_in\.out_layer\.(.*)$": r"vector_in.fc_out.\1",
# Time embedder naming
r"^time_in\.mlp\.0\.(.*)$": r"time_in.mlp.fc_in.\1",
r"^time_in\.mlp\.2\.(.*)$": r"time_in.mlp.fc_out.\1",
# Guidance embedder naming (if present)
r"^guidance_in\.mlp\.0\.(.*)$": r"guidance_in.mlp.fc_in.\1",
r"^guidance_in\.mlp\.2\.(.*)$": r"guidance_in.mlp.fc_out.\1",
# Final layer adaLN modulation
r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"final_layer.adaLN_modulation.linear.\1",
# Refiner block MLP naming
r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc1\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc2\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
# Camera net weights are already correctly named
})
# Reverse mapping for saving checkpoints
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
# Model architecture parameters
# patch_size can be int or tuple - if tuple, it's [T, H, W]
patch_size: int | tuple[int, int, int] = 2
patch_size_t: int = 1
in_channels: int = 33 # 16 latent + 16 gt_latent + 1 mask
out_channels: int = 16
num_attention_heads: int = 24
attention_head_dim: int = 128
mlp_ratio: float = 4.0
num_layers: int = 20 # Double stream blocks
num_single_layers: int = 40 # Single stream blocks
num_refiner_layers: int = 2
rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
guidance_embeds: bool = False # GameCraft doesn't use guidance
dtype: torch.dtype | None = None
text_embed_dim: int = 4096 # LLaMA-3 hidden size
pooled_projection_dim: int = 768 # CLIP pooled output dim
rope_theta: int = 256
qk_norm: str = "rms_norm"
# Camera net parameters
camera_in_channels: int = 6 # Plücker coordinates
camera_downscale_coef: int = 8
camera_out_channels: int = 16
# Layers to exclude from LoRA
exclude_lora_layers: list[str] = field(
default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in", "camera_net"])
def __post_init__(self):
super().__post_init__()
self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
self.num_channels_latents: int = 16 # Output is 16 channels
# Convert patch_size list to tuple if needed (from JSON deserialization)
if isinstance(self.patch_size, list):
self.patch_size = tuple(self.patch_size)
# Convert rope_axes_dim list to tuple if needed
if isinstance(self.rope_axes_dim, list):
self.rope_axes_dim = tuple(self.rope_axes_dim)
@dataclass
class HunyuanGameCraftConfig(DiTConfig):
"""Full config for HunyuanGameCraft model."""
arch_config: DiTArchConfig = field(default_factory=HunyuanGameCraftArchConfig)
prefix: str = "HunyuanGameCraft"
fastvideo/configs/models/dits/longcat.py
-133
View File
@@ -1,133 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
LongCat Video DiT configuration for native FastVideo implementation.
"""
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.platforms import AttentionBackendEnum
def is_longcat_blocks(n: str, m) -> bool:
"""FSDP shard condition for LongCat transformer blocks."""
return "blocks" in n and str.isdigit(n.split(".")[-1])
@dataclass
class LongCatVideoArchConfig(DiTArchConfig):
"""Architecture configuration for native LongCat Video DiT."""
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_longcat_blocks])
# Enable torch.compile for transformer blocks (major speedup!)
_compile_conditions: list = field(default_factory=lambda: [is_longcat_blocks])
# Parameter name mapping for weight conversion
param_names_mapping: dict = field(
default_factory=lambda: {
# Embedders
r"^x_embedder\.(.*)$": r"patch_embed.\1",
r"^t_embedder\.mlp\.0\.(.*)$": r"time_embedder.linear_1.\1",
r"^t_embedder\.mlp\.2\.(.*)$": r"time_embedder.linear_2.\1",
r"^y_embedder\.y_proj\.0\.(.*)$": r"caption_embedder.linear_1.\1",
r"^y_embedder\.y_proj\.2\.(.*)$": r"caption_embedder.linear_2.\1",
# Transformer blocks - AdaLN modulation
r"^blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$": r"blocks.\1.adaln_linear_1.\2",
# Transformer blocks - Normalization
r"^blocks\.(\d+)\.mod_norm_attn\.(.*)$": r"blocks.\1.norm_attn.\2",
r"^blocks\.(\d+)\.mod_norm_ffn\.(.*)$": r"blocks.\1.norm_ffn.\2",
r"^blocks\.(\d+)\.pre_crs_attn_norm\.(.*)$": r"blocks.\1.norm_cross.\2",
# Self-attention: QKV fused -> separate (will need splitting in converter)
# Original has attn.qkv.weight -> need to split into to_q, to_k, to_v
r"^blocks\.(\d+)\.attn\.qkv\.(.*)$": r"blocks.\1.self_attn.qkv_fused.\2", # Marker for splitting
r"^blocks\.(\d+)\.attn\.proj\.(.*)$": r"blocks.\1.self_attn.to_out.\2",
r"^blocks\.(\d+)\.attn\.q_norm\.(.*)$": r"blocks.\1.self_attn.q_norm.\2",
r"^blocks\.(\d+)\.attn\.k_norm\.(.*)$": r"blocks.\1.self_attn.k_norm.\2",
# Cross-attention
r"^blocks\.(\d+)\.cross_attn\.q_linear\.(.*)$": r"blocks.\1.cross_attn.to_q.\2",
r"^blocks\.(\d+)\.cross_attn\.kv_linear\.(.*)$":
r"blocks.\1.cross_attn.kv_fused.\2", # Marker for splitting
r"^blocks\.(\d+)\.cross_attn\.proj\.(.*)$": r"blocks.\1.cross_attn.to_out.\2",
r"^blocks\.(\d+)\.cross_attn\.q_norm\.(.*)$": r"blocks.\1.cross_attn.q_norm.\2",
r"^blocks\.(\d+)\.cross_attn\.k_norm\.(.*)$": r"blocks.\1.cross_attn.k_norm.\2",
# FFN (SwiGLU)
r"^blocks\.(\d+)\.ffn\.w1\.(.*)$": r"blocks.\1.ffn.w1.\2", # gate
r"^blocks\.(\d+)\.ffn\.w2\.(.*)$": r"blocks.\1.ffn.w2.\2", # down
r"^blocks\.(\d+)\.ffn\.w3\.(.*)$": r"blocks.\1.ffn.w3.\2", # up
# Final layer
r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"final_layer.adaln_linear.\1",
r"^final_layer\.norm_final\.(.*)$": r"final_layer.norm.\1",
r"^final_layer\.linear\.(.*)$": r"final_layer.proj.\1",
})
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
# LoRA parameter name mapping
lora_param_names_mapping: dict = field(default_factory=lambda: {})
# Model architecture parameters
hidden_size: int = 4096
depth: int = 48 # Number of transformer blocks
num_attention_heads: int = 32
attention_head_dim: int = 128 # hidden_size / num_attention_heads
in_channels: int = 16 # Latent space channels
out_channels: int = 16
num_channels_latents: int = 16
# Patch embedding
patch_size: tuple[int, int, int] = (1, 2, 2) # [T, H, W] - no temporal compression
# Text/caption embedding
caption_channels: int = 4096 # UMT5 d_model
# Timestep embedding
adaln_tembed_dim: int = 512
frequency_embedding_size: int = 256
# FFN
mlp_ratio: int = 4
# Attention backend support
_supported_attention_backends: tuple = field(default_factory=lambda: (
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
))
# Text padding behavior
text_tokens_zero_pad: bool = True
# Block Sparse Attention (BSA)
enable_bsa: bool = False
bsa_params: dict | None = field(default_factory=lambda: {
"sparsity": 0.9375,
"cdf_threshold": None,
"chunk_3d_shape_q": [4, 4, 4],
"chunk_3d_shape_k": [4, 4, 4],
})
# LoRA exclusions
exclude_lora_layers: list[str] = field(default_factory=lambda: [])
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
# Ensure attention_head_dim matches
self.attention_head_dim = self.hidden_size // self.num_attention_heads
@dataclass
class LongCatVideoConfig(DiTConfig):
"""Main configuration for LongCat Video DiT."""
arch_config: DiTArchConfig = field(default_factory=LongCatVideoArchConfig)
prefix: str = "longcat"
fastvideo/configs/models/dits/zimage.py
-60
View File
@@ -1,60 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.platforms import AttentionBackendEnum
def is_zimage_block(name: str, module) -> bool:
parts = name.split(".")
return len(parts) >= 2 and parts[-2] in {"noise_refiner", "context_refiner", "layers"} and parts[-1].isdigit()
@dataclass
class ZImageDiTArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_block])
_supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.TORCH_SDPA, )
all_patch_size: tuple[int, ...] = (2, )
all_f_patch_size: tuple[int, ...] = (1, )
in_channels: int = 16
dim: int = 3840
n_layers: int = 30
n_refiner_layers: int = 2
n_heads: int = 30
n_kv_heads: int = 30
norm_eps: float = 1e-5
qk_norm: bool = True
cap_feat_dim: int = 2560
rope_theta: float = 256.0
t_scale: float = 1000.0
axes_dims: tuple[int, ...] = (32, 48, 48)
axes_lens: tuple[int, ...] = (1536, 512, 512)
adaln_embed_dim: int = 256
frequency_embedding_size: int = 256
timestep_mid_size: int = 1024
max_period: int = 10000
seq_multi_of: int = 32
def __post_init__(self) -> None:
super().__post_init__()
if len(self.all_patch_size) != len(self.all_f_patch_size):
raise ValueError("all_patch_size and all_f_patch_size must have equal length")
if self.dim % self.n_heads:
raise ValueError("dim must be divisible by n_heads")
if self.dim // self.n_heads != sum(self.axes_dims):
raise ValueError("attention head dimension must equal sum(axes_dims)")
if len(self.axes_dims) != len(self.axes_lens) or any(dim % 2 for dim in self.axes_dims):
raise ValueError("RoPE axes require matching lengths and even dimensions")
self.hidden_size = self.dim
self.num_attention_heads = self.n_heads
self.num_channels_latents = self.in_channels
self.out_channels = self.in_channels
@dataclass
class ZImageDiTConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=ZImageDiTArchConfig)
prefix: str = "ZImage"
fastvideo/configs/models/encoders/base.py
+1 -1
View File
@@ -44,7 +44,7 @@ class TextEncoderArchConfig(EncoderArchConfig):
def __post_init__(self) -> None:
# update_model_arch re-runs __post_init__ after pipeline configs may
# have customized tokenizer_kwargs (e.g. kandinsky5/gen3c/longcat set
# have customized tokenizer_kwargs (e.g. kandinsky5/gen3c set
# "padding"); rebuilding the dict here would silently wipe those
# customizations, so only fill in defaults for keys not already set.
defaults = {
fastvideo/configs/models/vaes/__init__.py
-4
View File
@@ -1,8 +1,6 @@
from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEArchConfig, Kandinsky6SRVAEConfig
@@ -17,7 +15,6 @@ from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
from fastvideo.models.wan.vae_config import WanVAEConfig
__all__ = [
"GameCraftVAEConfig",
"HunyuanVAEConfig",
"WanVAEConfig",
"CosmosVAEConfig",
@@ -34,5 +31,4 @@ __all__ = [
"OobleckVAEArchConfig",
"OobleckVAEConfig",
"Flux2VAEConfig",
"GlmImageVAEConfig",
]
fastvideo/configs/models/vaes/gamecraftvae.py
-50
View File
@@ -1,50 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
GameCraft VAE config - matches official config.json from Hunyuan-GameCraft-1.0.
"""
from dataclasses import dataclass, field
from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class GameCraftVAEArchConfig(VAEArchConfig):
"""Architecture config matching official AutoencoderKLCausal3D config.json."""
in_channels: int = 3
out_channels: int = 3
latent_channels: int = 16
down_block_types: tuple[str, ...] = (
"DownEncoderBlockCausal3D",
"DownEncoderBlockCausal3D",
"DownEncoderBlockCausal3D",
"DownEncoderBlockCausal3D",
)
up_block_types: tuple[str, ...] = (
"UpDecoderBlockCausal3D",
"UpDecoderBlockCausal3D",
"UpDecoderBlockCausal3D",
"UpDecoderBlockCausal3D",
)
block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
layers_per_block: int = 2
act_fn: str = "silu"
norm_num_groups: int = 32
scaling_factor: float = 0.476986
spatial_compression_ratio: int = 8
temporal_compression_ratio: int = 4
time_compression_ratio: int = 4 # alias for DecoderCausal3D
mid_block_add_attention: bool = True
mid_block_causal_attn: bool = True
sample_size: int = 256 # from config.json
sample_tsize: int = 64 # from config.json
def __post_init__(self):
self.spatial_compression_ratio = 2**(len(self.block_out_channels) - 1)
@dataclass
class GameCraftVAEConfig(VAEConfig):
"""Full config for GameCraft VAE."""
arch_config: VAEArchConfig = field(default_factory=GameCraftVAEArchConfig)
fastvideo/configs/models/vaes/glm_image.py
-95
View File
@@ -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
fastvideo/configs/pipelines/__init__.py
+3 -6
View File
@@ -1,10 +1,8 @@
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.configs.pipelines.cosmos import CosmosConfig
from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
from fastvideo.configs.pipelines.hyworld import HYWorldConfig
from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig
from fastvideo.configs.pipelines.kandinsky6 import Kandinsky6TI2VAConfig
@@ -20,10 +18,9 @@ from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcing
WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig)
__all__ = [
"HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
"Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
"SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
"DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
"HunyuanConfig", "FastHunyuanConfig", "PipelineConfig", "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig",
"WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig",
"LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
"Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig", "Kandinsky6SRPipelineConfig",
"LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig",
"MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name"
fastvideo/configs/pipelines/dreamx_world.py
-127
View File
@@ -1,127 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World-5B-Cam FastVideo model configuration helpers."""
from collections.abc import Callable
from dataclasses import dataclass, field
import torch
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldARArchConfig, DreamXWorldARConfig,
DreamXWorldArchConfig, DreamXWorldConfig)
from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
from fastvideo.configs.models.encoders.t5 import T5ArchConfig
from fastvideo.configs.models.vaes import WanVAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.models.wan.pipeline_config import LucyEditDevConfig, t5_postprocess_text
def make_dreamx_world_5b_cam_dit_config() -> DreamXWorldConfig:
"""Return the DreamX-World DiT config matching DreamX-World-5B-Cam."""
return DreamXWorldConfig(arch_config=DreamXWorldArchConfig(
num_attention_heads=24,
attention_head_dim=128,
in_channels=48,
out_channels=48,
ffn_dim=14336,
num_layers=30,
cross_attn_norm=True,
qk_norm="rms_norm_across_heads",
add_control_adapter=True,
cam_method="prope",
attn_compress=1,
cam_self_attn_layers=None,
))
def make_dreamx_world_5b_ar_dit_config() -> DreamXWorldARConfig:
"""Return the DreamX-World-5B autoregressive causal DiT config."""
return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig(
model_type="ti2v",
num_attention_heads=24,
attention_head_dim=128,
in_channels=48,
out_channels=48,
ffn_dim=14336,
num_layers=30,
cross_attn_norm=True,
qk_norm=True,
add_control_adapter=True,
cam_method="prope",
attn_compress=4,
cam_self_attn_layers=tuple(range(30)),
local_attn_size=12,
sink_size=3,
num_frames_per_block=3,
))
def make_dreamx_world_5b_cam_vae_config() -> WanVAEConfig:
"""Return the Wan2.2 48-channel VAE config used by DreamX-World-5B-Cam."""
return LucyEditDevConfig().vae_config
def make_dreamx_world_5b_cam_text_encoder_config() -> T5Config:
"""Return the UMT5-XXL text encoder config used by DreamX-World-5B-Cam."""
return T5Config(
arch_config=T5ArchConfig(
vocab_size=256384,
d_model=4096,
d_kv=64,
d_ff=10240,
num_layers=24,
num_decoder_layers=None,
num_heads=64,
relative_attention_num_buckets=32,
dropout_rate=0.0,
text_len=512,
feed_forward_proj="gelu",
is_encoder_decoder=False,
),
prefix="umt5",
)
@dataclass
class DreamXWorld5BCamPipelineConfig(PipelineConfig):
"""Pipeline config for the first-scope DreamX-World-5B-Cam mode."""
dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_cam_dit_config)
vae_config: VAEConfig = field(default_factory=make_dreamx_world_5b_cam_vae_config)
text_encoder_configs: tuple[EncoderConfig,
...] = field(default_factory=lambda: (make_dreamx_world_5b_cam_text_encoder_config(), ))
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
...] = field(default_factory=lambda: (t5_postprocess_text, ))
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
flow_shift: float | None = 3.0
ti2v_task: bool = True
expand_timesteps: bool = True
vae_tiling: bool = False
vae_sp: bool = False
vae_precision: str = "fp32"
vae_decode_precision: str | None = "bf16"
dit_precision: str = "bf16"
def __post_init__(self) -> None:
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
self.dit_config.expand_timesteps = self.expand_timesteps
@dataclass
class DreamXWorld5BARPipelineConfig(DreamXWorld5BCamPipelineConfig):
"""Pipeline config for DreamX-World-5B autoregressive forcing."""
dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_ar_dit_config)
flow_shift: float | None = 5.0
ti2v_task: bool = True
is_causal: bool = True
dmd_denoising_steps: tuple[int, ...] = (1000, 750, 500, 250)
warp_denoising_step: bool = True
context_noise: float = 0.1
num_frames_per_block: int = 3
color_correction_strength: float = 1.0
def __post_init__(self) -> None:
super().__post_init__()
self.dit_config.expand_timesteps = True
fastvideo/configs/pipelines/glm_image.py
-46
View File
@@ -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
fastvideo/configs/pipelines/hunyuangamecraft.py
-69
View File
@@ -1,69 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
Pipeline configuration for HunyuanGameCraft.
HunyuanGameCraft extends HunyuanVideo with:
1. CameraNet for camera/action conditioning (Plücker coordinates)
2. Mask-based conditioning for autoregressive generation
3. 33 input channels (16 latent + 16 gt_latent + 1 mask)
Text encoders are the same as HunyuanVideo:
- LLaVA-LLaMA-3-8B for primary text encoding (4096 dim)
- CLIP ViT-L/14 for secondary pooled embeddings (768 dim)
"""
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 import HunyuanGameCraftConfig
from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPTextConfig, LlamaConfig)
from fastvideo.configs.models.vaes import GameCraftVAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.configs.pipelines.hunyuan import (clip_postprocess_text, clip_preprocess_text, llama_postprocess_text,
llama_preprocess_text)
@dataclass
class HunyuanGameCraftPipelineConfig(PipelineConfig):
"""Configuration for HunyuanGameCraft pipeline.
Inherits text encoding from HunyuanVideo but uses:
- GameCraft DiT with CameraNet
- Same VAE (HunyuanVAE)
- Same text encoders (LLaMA + CLIP)
"""
# DiT config - uses GameCraft config (33 input channels)
dit_config: DiTConfig = field(default_factory=HunyuanGameCraftConfig)
# VAE config - GameCraft VAE (mid_block_causal_attn=True, etc.)
vae_config: VAEConfig = field(default_factory=GameCraftVAEConfig)
# Denoising parameters
# Official GameCraft does NOT use embedded guidance (passes guidance=None)
# It uses standard CFG with guidance_scale=6.0 instead
embedded_cfg_scale = None
flow_shift: int = 5 # Official GameCraft uses flow_shift=5.0
# Text encoding stage - same as HunyuanVideo
# Uses LLaMA-3-8B (via LLaVA) + CLIP
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
preprocess_text_funcs: tuple[Callable[[str], str],
...] = field(default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
...] = field(default_factory=lambda: (llama_postprocess_text, clip_postprocess_text))
# Precision for each component
dit_precision: str = "bf16"
vae_precision: str = "fp16"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp16", "fp16"))
def __post_init__(self):
# VAE only needs decoder for inference
self.vae_config.load_encoder = False
self.vae_config.load_decoder = True
if self.text_encoder_configs:
# Hunyuan postprocess consumes intermediate LLaMA hidden states.
self.text_encoder_configs[0].arch_config.output_hidden_states = True
fastvideo/configs/pipelines/longcat.py
-362
View File
@@ -1,362 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from dataclasses import dataclass, field
import html
import ftfy
import regex as re
import torch
from fastvideo.configs.models import DiTConfig, VAEConfig
from fastvideo.configs.models.dits.base import DiTArchConfig
from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
from fastvideo.configs.models.encoders.base import TextEncoderArchConfig
from fastvideo.configs.models.encoders.t5 import T5ArchConfig
from fastvideo.configs.models.vaes import WanVAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig
@dataclass
class LongCatT5ArchConfig(T5ArchConfig):
"""T5 arch that pads tokenizer output to ``max_length``.
LongCat's denoising stage concatenates positive and negative
attention masks along the batch dimension for CFG, which requires
uniform seq length. The shared :class:`T5ArchConfig` dropped the
``"padding": "max_length"`` tokenizer kwarg so other DiTs could run
with variable-length masks; LongCat still needs the uniform
contract.
"""
def __post_init__(self) -> None:
super().__post_init__()
self.tokenizer_kwargs["padding"] = "max_length"
@dataclass
class LongCatT5Config(T5Config):
arch_config: TextEncoderArchConfig = field(default_factory=LongCatT5ArchConfig)
@dataclass
class LongCatDiTArchConfig(DiTArchConfig):
"""Extended DiTArchConfig with LongCat-specific fields."""
# LongCat-specific architecture parameters
adaln_tembed_dim: int = 512
caption_channels: int = 4096
depth: int = 48
enable_bsa: bool = False
enable_flashattn3: bool = False
enable_flashattn2: bool = True
enable_xformers: bool = False
frequency_embedding_size: int = 256
in_channels: int = 16
mlp_ratio: int = 4
num_heads: int = 32
out_channels: int = 16
text_tokens_zero_pad: bool = True
patch_size: list[int] = field(default_factory=lambda: [1, 2, 2])
cp_split_hw: list[int] | None = None
bsa_params: dict | None = None
def longcat_preprocess_text(prompt: str) -> str:
"""Clean and preprocess text like original LongCat implementation.
This function applies the same text cleaning pipeline as the original
LongCat-Video implementation to ensure identical tokenization results.
Steps:
1. basic_clean: Fix unicode issues and unescape HTML entities
2. whitespace_clean: Normalize whitespace to single spaces
Args:
prompt: Raw input text prompt
Returns:
Cleaned and normalized text prompt
"""
# basic_clean: fix unicode and HTML entities
text = ftfy.fix_text(prompt)
text = html.unescape(html.unescape(text))
text = text.strip()
# whitespace_clean: normalize whitespace
text = re.sub(r"\s+", " ", text)
text = text.strip()
return text
def umt5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
"""
Postprocess UMT5/T5 encoder outputs to fixed length 512 embeddings.
"""
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
prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
prompt_embeds_tensor: torch.Tensor = torch.stack(
[torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
return prompt_embeds_tensor
@dataclass
class LongCatT2V480PConfig(PipelineConfig):
"""Configuration for LongCat pipeline (480p).
Components expected by loaders:
- tokenizer: AutoTokenizer
- text_encoder: UMT5EncoderModel
- transformer: LongCatTransformer3DModel
- vae: AutoencoderKLWan (Wan VAE, 4x8 compression)
- scheduler: FlowMatchEulerDiscreteScheduler
"""
dit_config: DiTConfig = field(default_factory=lambda: DiTConfig(arch_config=LongCatDiTArchConfig()))
# VAE config: Wan VAE with encoder+decoder enabled
vae_config: VAEConfig = field(default_factory=WanVAEConfig)
vae_tiling: bool = False
vae_sp: bool = False
# Precision defaults
dit_precision: str = "bf16"
vae_precision: str = "bf16"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
# UMT5 uses T5-like config; postprocess pads to 512. LongCatT5Config
# restores ``padding="max_length"`` for the CFG concat contract.
text_encoder_configs: tuple[T5Config, ...] = field(default_factory=lambda: (LongCatT5Config(), ))
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (longcat_preprocess_text, ))
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
...] = field(default_factory=lambda: (umt5_postprocess_text, ))
# LongCat-specific runtime toggles (consumed by pipeline/stages)
enable_kv_cache: bool = True
offload_kv_cache: bool = False
enable_bsa: bool = False
use_distill: bool = False
enhance_hf: bool = False
# Optional BSA parameter dict (kept for backward/phase-1 compatibility).
# `LongCatPipeline.initialize_pipeline()` uses this as a base and then applies
# CLI overrides (bsa_sparsity / bsa_chunk_{q,k} / bsa_cdf_threshold).
bsa_params: dict | None = None
# BSA runtime overrides (preferred over bsa_params if provided via CLI)
bsa_sparsity: float | None = None
bsa_cdf_threshold: float | None = None
bsa_chunk_q: list[int] | None = None
bsa_chunk_k: list[int] | None = None
t_thresh: float | None = None # refine stage default controlled by sampling args
# LongCat does not need flow_shift
flow_shift: float | None = None
dmd_denoising_steps: list[int] | None = None
def __post_init__(self):
# LongCat inference requires vae encoder and decoder
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
@dataclass
class LongCatT2V704PConfig(LongCatT2V480PConfig):
"""Configuration for LongCat pipeline (704p) with BSA enabled by default.
Uses the same resolution and BSA parameters as original LongCat refinement stage.
BSA parameters configured in transformer config.json with chunk_3d_shape=[4,4,4]:
- Input: 704×1280×96
- VAE (8x): 88×160×96
- Patch [1,2,2]: 44×80×96
- chunk [4,4,4]: 96%4=0, 44%4=0, 80%4=0 ✅
This configuration matches the original LongCat refinement stage parameters.
"""
# Enable BSA by default for 704p
enable_bsa: bool = True
ASPECT_RATIO_627 = {
'0.26': ([320, 1216], 1),
'0.31': ([352, 1120], 1),
'0.38': ([384, 1024], 1),
'0.43': ([416, 960], 1),
'0.52': ([448, 864], 1),
'0.58': ([480, 832], 1),
'0.67': ([512, 768], 1),
'0.74': ([544, 736], 1),
'0.86': ([576, 672], 1),
'0.95': ([608, 640], 1),
'1.05': ([640, 608], 1),
'1.17': ([672, 576], 1),
'1.29': ([704, 544], 1),
'1.35': ([736, 544], 1),
'1.50': ([768, 512], 1),
'1.67': ([800, 480], 1),
'1.73': ([832, 480], 1),
'2.00': ([896, 448], 1),
'2.31': ([960, 416], 1),
'2.58': ([992, 384], 1),
'2.75': ([1056, 384], 1),
'3.09': ([1088, 352], 1),
'3.70': ([1184, 320], 1),
'3.80': ([1216, 320], 1),
'3.90': ([1248, 320], 1),
'4.00': ([1280, 320], 1)
}
ASPECT_RATIO_627_F64 = {
'0.26': ([320, 1216], 1),
'0.38': ([384, 1024], 1),
'0.50': ([448, 896], 1),
'0.67': ([512, 768], 1),
'0.82': ([576, 704], 1),
'1.00': ([640, 640], 1),
'1.22': ([704, 576], 1),
'1.50': ([768, 512], 1),
'1.86': ([832, 448], 1),
'2.00': ([896, 448], 1),
'2.50': ([960, 384], 1),
'2.83': ([1088, 384], 1),
'3.60': ([1152, 320], 1),
'3.80': ([1216, 320], 1),
'4.00': ([1280, 320], 1)
}
ASPECT_RATIO_627_F128 = {
'0.25': ([256, 1024], 1),
'0.38': ([384, 1024], 1),
'0.43': ([384, 896], 1),
'0.57': ([512, 896], 1),
'0.67': ([512, 768], 1),
'1.00': ([640, 640], 1),
'1.50': ([768, 512], 1),
'1.75': ([896, 512], 1),
'2.33': ([896, 384], 1),
'2.67': ([1024, 384], 1),
'4.00': ([1024, 256], 1),
}
ASPECT_RATIO_627_F256 = {
'0.25': ([256, 1024], 1),
'0.33': ([256, 768], 1),
'0.50': ([256, 512], 1),
'0.67': ([512, 768], 1),
'1.00': ([512, 512], 1),
'1.50': ([768, 512], 1),
'2.00': ([512, 256], 1),
'3.00': ([768, 256], 1),
'4.00': ([1024, 256], 1),
}
ASPECT_RATIO_960 = {
'0.25': ([480, 1920], 1),
'0.29': ([512, 1792], 1),
'0.32': ([544, 1696], 1),
'0.36': ([576, 1600], 1),
'0.40': ([608, 1504], 1),
'0.49': ([672, 1376], 1),
'0.54': ([704, 1312], 1),
'0.59': ([736, 1248], 1),
'0.69': ([800, 1152], 1),
'0.74': ([832, 1120], 1),
'0.82': ([864, 1056], 1),
'0.88': ([896, 1024], 1),
'0.94': ([928, 992], 1),
'1.00': ([960, 960], 1),
'1.07': ([992, 928], 1),
'1.14': ([1024, 896], 1),
'1.22': ([1056, 864], 1),
'1.31': ([1088, 832], 1),
'1.35': ([1120, 832], 1),
'1.44': ([1152, 800], 1),
'1.70': ([1248, 736], 1),
'2.00': ([1344, 672], 1),
'2.05': ([1376, 672], 1),
'2.47': ([1504, 608], 1),
'2.53': ([1536, 608], 1),
'2.83': ([1632, 576], 1),
'3.06': ([1664, 544], 1),
'3.12': ([1696, 544], 1),
'3.62': ([1856, 512], 1),
'3.93': ([1888, 480], 1),
'4.00': ([1920, 480], 1)
}
ASPECT_RATIO_960_F64 = {
'0.22': ([448, 2048], 1),
'0.29': ([512, 1792], 1),
'0.36': ([576, 1600], 1),
'0.45': ([640, 1408], 1),
'0.55': ([704, 1280], 1),
'0.63': ([768, 1216], 1),
'0.76': ([832, 1088], 1),
'0.88': ([896, 1024], 1),
'1.00': ([960, 960], 1),
'1.14': ([1024, 896], 1),
'1.31': ([1088, 832], 1),
'1.50': ([1152, 768], 1),
'1.58': ([1216, 768], 1),
'1.82': ([1280, 704], 1),
'1.91': ([1344, 704], 1),
'2.20': ([1408, 640], 1),
'2.30': ([1472, 640], 1),
'2.67': ([1536, 576], 1),
'2.89': ([1664, 576], 1),
'3.62': ([1856, 512], 1),
'3.75': ([1920, 512], 1)
}
ASPECT_RATIO_960_F128 = {
'0.20': ([384, 1920], 1),
'0.27': ([512, 1920], 1),
'0.33': ([512, 1536], 1),
'0.42': ([640, 1536], 1),
'0.50': ([640, 1280], 1),
'0.60': ([768, 1280], 1),
'0.67': ([768, 1152], 1),
'0.78': ([896, 1152], 1),
'1.00': ([1024, 1024], 1),
'1.29': ([1152, 896], 1),
'1.50': ([1152, 768], 1),
'1.67': ([1280, 768], 1),
'2.00': ([1280, 640], 1),
'2.40': ([1536, 640], 1),
'3.00': ([1536, 512], 1),
'3.75': ([1920, 512], 1),
'5.00': ([1920, 384], 1),
}
ASPECT_RATIO_960_F256 = {
'0.33': ([512, 1536], 1),
'0.60': ([768, 1280], 1),
'1.00': ([1024, 1024], 1),
'1.67': ([1280, 768], 1),
'3.00': ([1536, 512], 1),
}
def get_bucket_config(resolution, scale_factor_spatial):
if resolution == '480p':
if scale_factor_spatial == 16 or scale_factor_spatial == 32:
return ASPECT_RATIO_627
elif scale_factor_spatial == 64:
return ASPECT_RATIO_627_F64
elif scale_factor_spatial == 128:
return ASPECT_RATIO_627_F128
elif scale_factor_spatial == 256:
return ASPECT_RATIO_627_F256
elif resolution == '720p':
if scale_factor_spatial == 16 or scale_factor_spatial == 32:
return ASPECT_RATIO_960
elif scale_factor_spatial == 64:
return ASPECT_RATIO_960_F64
elif scale_factor_spatial == 128:
return ASPECT_RATIO_960_F128
elif scale_factor_spatial == 256:
return ASPECT_RATIO_960_F256
raise ValueError(f"Unsupported resolution '{resolution}' or scale_factor_spatial '{scale_factor_spatial}'")
fastvideo/configs/pipelines/turbodiffusion.py
-121
View File
@@ -1,121 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
TurboDiffusion pipeline configurations.
TurboDiffusion uses RCM (recurrent Consistency Model) scheduler with
SLA (Sparse-Linear Attention) for fast 1-4 step video generation.
"""
from dataclasses import dataclass, field
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
from fastvideo.configs.models.dits import WanVideoConfig
from fastvideo.configs.models.encoders import CLIPVisionConfig
from fastvideo.configs.models.vaes import WanVAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.models.wan.pipeline_config import t5_postprocess_text, T5Config, BaseEncoderOutput
import torch
from collections.abc import Callable
@dataclass
class TurboDiffusionT2VConfig(PipelineConfig):
"""Base configuration for TurboDiffusion T2V pipeline.
Uses RCM scheduler with sigma_max=80 for 1-4 step generation.
No boundary_ratio (single model, no switching).
"""
# DiT
dit_config: DiTConfig = field(default_factory=WanVideoConfig)
# VAE
vae_config: VAEConfig = field(default_factory=WanVAEConfig)
vae_tiling: bool = False
vae_sp: bool = False
# Denoising stage
flow_shift: float | None = 3.0
# No boundary_ratio for T2V (single model)
boundary_ratio: float | None = None
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
...] = field(default_factory=lambda: (t5_postprocess_text, ))
# Precision for each component
precision: str = "bf16"
vae_precision: str = "fp32"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
# self-forcing params
warp_denoising_step: bool = True
def __post_init__(self):
self.vae_config.load_encoder = False
self.vae_config.load_decoder = True
# Ensure no boundary_ratio is set in dit_config
self.dit_config.boundary_ratio = None
@dataclass
class TurboDiffusionT2V_1_3B_Config(TurboDiffusionT2VConfig):
"""Configuration for TurboDiffusion T2V 1.3B model."""
pass
@dataclass
class TurboDiffusionT2V_14B_Config(TurboDiffusionT2VConfig):
"""Configuration for TurboDiffusion T2V 14B model.
Uses same config as 1.3B but with higher flow_shift for 14B model.
"""
flow_shift: float | None = 5.0
@dataclass
class TurboDiffusionI2VConfig(PipelineConfig):
"""Base configuration for TurboDiffusion I2V pipeline.
Uses RCM scheduler with sigma_max=200 for 1-4 step generation.
Uses boundary_ratio=0.9 for high-noise to low-noise model switching.
"""
# DiT
dit_config: DiTConfig = field(default_factory=WanVideoConfig)
# VAE
vae_config: VAEConfig = field(default_factory=WanVAEConfig)
vae_tiling: bool = False
vae_sp: bool = False
# Denoising stage
flow_shift: float | None = 5.0
boundary_ratio: float | None = 0.9
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
...] = field(default_factory=lambda: (t5_postprocess_text, ))
# Image encoder for I2V
image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig)
image_encoder_precision: str = "fp32"
# Precision for each component
precision: str = "bf16"
vae_precision: str = "fp32"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
# self-forcing params
warp_denoising_step: bool = True
def __post_init__(self):
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
self.dit_config.boundary_ratio = self.boundary_ratio
@dataclass
class TurboDiffusionI2V_A14B_Config(TurboDiffusionI2VConfig):
"""Configuration for TurboDiffusion I2V A14B model."""
pass
fastvideo/configs/pipelines/zimage.py
-52
View File
@@ -1,52 +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.zimage import ZImageDiTConfig
from fastvideo.configs.models.encoders import BaseEncoderOutput
from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
def _zimage_text_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
if outputs.hidden_states is None:
raise RuntimeError("Z-Image requires Qwen3 hidden states")
return outputs.hidden_states[-2]
@dataclass
class ZImagePipelineConfig(PipelineConfig):
"""Configuration for the native Z-Image text-to-image pipeline."""
scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
transformer_arch: str = "ZImageTransformer2DModel"
vae_arch: str = "AutoencoderKL"
text_encoder_archs: tuple[str, ...] = ("Qwen3Model", )
tokenizer_archs: tuple[str, ...] = ("Qwen2Tokenizer", )
dit_config: ZImageDiTConfig = field(default_factory=ZImageDiTConfig)
vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (Qwen3TextConfig(chat_template_enable_thinking=True), ))
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (preprocess_text, ))
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
...] = field(default_factory=lambda: (_zimage_text_postprocess, ))
dit_precision: str = "bf16"
vae_precision: str = "fp32"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
vae_tiling: bool = False
vae_sp: bool = False
embedded_cfg_scale: float = 0.0
flow_shift: float | None = 3.0
scheduler_step_in_fp32: bool = True
scheduler_sigma_min: float = 0.0
scheduler_use_reference_discrete_timesteps: bool = True
fastvideo/envs.py
+16 -32
View File
@@ -719,14 +719,6 @@ FASTVIDEO_TEST_GRADNORM_UPDATE = EnvBool(False,
category="test",
doc="Gradient-norm regression tests update their references.",
deprecated_names=("FASTVIDEO_GRADNORM_UPDATE", ))
FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH = EnvStr("FastVideo/DreamX-World-5B-Cam-Diffusers",
category="test",
doc="Model for the DreamX-World camera SSIM test.",
deprecated_names=("DREAMX_WORLD_SSIM_MODEL_PATH", ))
FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH = EnvStr("FastVideo/DreamX-World-5B-Diffusers",
category="test",
doc="Model for the DreamX-World autoregressive SSIM test.",
deprecated_names=("DREAMX_WORLD_AR_SSIM_MODEL_PATH", ))
FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR = EnvStr("black-forest-labs/FLUX.1-dev",
category="test",
doc="Model for the Flux text-to-image SSIM test.",
@@ -735,10 +727,6 @@ FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH = EnvStr(None,
category="test",
doc="Local Flux transformer for the Flux transformer test.",
deprecated_names=("FLUX_TRANSFORMER_PATH", ))
FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH = EnvStr("FastVideo/HunyuanGameCraft-Diffusers",
category="test",
doc="Model for the HunyuanGameCraft SSIM test.",
deprecated_names=("GAMECRAFT_MODEL_PATH", ))
FASTVIDEO_TEST_GEN3C_MODEL_PATH = EnvStr("FastVideo/GEN3C-Cosmos-7B-Diffusers",
category="test",
doc="Model for the GEN3C SSIM test.",
@@ -747,14 +735,6 @@ FASTVIDEO_TEST_GEN3C_IMAGE_PATH = EnvStr(None,
category="test",
doc="Input image for the GEN3C SSIM test.",
deprecated_names=("GEN3C_TEST_IMAGE_PATH", ))
FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR = EnvStr(None,
category="test",
doc="Local official GLM-Image weights for the GLM-Image SSIM test.",
deprecated_names=("GLM_IMAGE_LOCAL_WEIGHTS_DIR", ))
FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR = EnvStr(None,
category="test",
doc="Model for the GLM-Image SSIM test.",
deprecated_names=("GLM_IMAGE_MODEL_DIR", ))
FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS = EnvInt(1,
category="test",
doc="GPUs for the Kandinsky5 nightly end-to-end overfit test.",
@@ -764,10 +744,6 @@ FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE = EnvBool(
category="test",
doc="The Kandinsky5 nightly end-to-end test writes a missing reference video.",
deprecated_names=("KANDINSKY5_E2E_WRITE_REFERENCE", ))
FASTVIDEO_TEST_LONGCAT_MODEL_ROOT = EnvStr(None,
category="test",
doc="Local LongCat-Video checkpoint for the golden-gate test.",
deprecated_names=("LONGCAT_MODEL_ROOT", ))
FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR = EnvStr(None,
category="test",
doc="Local directory of MiniMax-H3 golden-gate tensors.",
@@ -788,14 +764,6 @@ FASTVIDEO_TEST_TAEH3_REFERENCE_DIR = EnvStr(None,
category="test",
doc="Upstream taehv checkout for the MLX TAEH3 parity test.",
deprecated_names=("TAEH3_REFERENCE_DIR", ))
FASTVIDEO_TEST_ZIMAGE_MODEL_DIR = EnvStr("Tongyi-MAI/Z-Image-Turbo",
category="test",
doc="Model for the Z-Image SSIM test.",
deprecated_names=("ZIMAGE_MODEL_DIR", ))
FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION = EnvStr("f332072aa78be7aecdf3ee76d5c247082da564a6",
category="test",
doc="Hugging Face revision of the Z-Image model for its SSIM test.",
deprecated_names=("ZIMAGE_MODEL_REVISION", ))
# Variables that FastVideo no longer reads. Setting one logs a warning; delete
# the entries in the next minor release.
@@ -807,6 +775,22 @@ DEPRECATED_VARIABLES = {
"FASTVIDEO_SERVER_DEV_MODE": "no code reads it",
"FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE": "no code reads it",
"FASTVIDEO_TRACE_FUNCTION": "no code reads it",
"FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH": "DreamX World was removed",
"DREAMX_WORLD_SSIM_MODEL_PATH": "DreamX World was removed",
"FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH": "DreamX World was removed",
"DREAMX_WORLD_AR_SSIM_MODEL_PATH": "DreamX World was removed",
"FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH": "HunyuanGameCraft was removed",
"GAMECRAFT_MODEL_PATH": "HunyuanGameCraft was removed",
"FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR": "GLM-Image was removed",
"GLM_IMAGE_LOCAL_WEIGHTS_DIR": "GLM-Image was removed",
"FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR": "GLM-Image was removed",
"GLM_IMAGE_MODEL_DIR": "GLM-Image was removed",
"FASTVIDEO_TEST_LONGCAT_MODEL_ROOT": "LongCat was removed",
"LONGCAT_MODEL_ROOT": "LongCat was removed",
"FASTVIDEO_TEST_ZIMAGE_MODEL_DIR": "Z-Image was removed",
"ZIMAGE_MODEL_DIR": "Z-Image was removed",
"FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION": "Z-Image was removed",
"ZIMAGE_MODEL_REVISION": "Z-Image was removed",
}
fastvideo/fastvideo_args.py
+2 -33
View File
@@ -655,37 +655,6 @@ class FastVideoArgs:
help="Optional list of module name substrings to restrict LoRA injection (e.g. q_proj k_proj v_proj).",
)
# BSA runtime control (LongCat)
parser.add_argument(
"--enable-bsa",
action=StoreBoolean,
help="Enable Block Sparse Attention (BSA) at runtime (overrides config).",
)
parser.add_argument(
"--bsa-sparsity",
type=float,
help="BSA sparsity (e.g., 0.9375).",
)
parser.add_argument(
"--bsa-cdf-threshold",
type=float,
help="BSA CDF threshold (optional).",
)
parser.add_argument(
"--bsa-chunk-q",
nargs=3,
type=int,
metavar=("T", "H", "W"),
help="BSA chunk_3d_shape_q as three ints, e.g., 4 4 4.",
)
parser.add_argument(
"--bsa-chunk-k",
nargs=3,
type=int,
metavar=("T", "H", "W"),
help="BSA chunk_3d_shape_k as three ints, e.g., 4 4 4.",
)
parser.add_argument(
"--enable-torch-compile",
action=StoreBoolean,
@@ -936,8 +905,8 @@ class FastVideoArgs:
kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs)
kwargs['preprocess_config'] = PreprocessConfig.from_kwargs(kwargs)
# Filter to only FastVideoArgs dataclass fields — pipeline-specific CLI
# args (e.g. enable_bsa, bsa_sparsity) live in PipelineConfig and must
# not be forwarded to the FastVideoArgs constructor.
# args live in PipelineConfig and must not be forwarded to the
# FastVideoArgs constructor.
valid_fields = {f.name for f in dataclasses.fields(cls)}
return cls(**{k: v for k, v in kwargs.items() if k in valid_fields})
fastvideo/layers/rotary_embedding.py
+3 -3
View File
@@ -117,7 +117,7 @@ def _apply_rotary_emb(
positional embeddings.
The function auto-detects whether cos/sin are full or half head_size:
- If cos/sin have head_size: use rotate_half style (for HunyuanVideo/GameCraft)
- If cos/sin have head_size: use rotate_half style (for HunyuanVideo)
- If cos/sin have head_size // 2: use Neox/GPT-J style
"""
head_size = x.shape[-1]
@@ -125,7 +125,7 @@ def _apply_rotary_emb(
# Check if cos/sin are full head_dim (rotate_half style) or half (traditional style)
if rope_dim == head_size:
# Full head_dim - use rotate_half style (HunyuanVideo, GameCraft)
# Full head_dim - use rotate_half style (HunyuanVideo)
# x * cos + rotate_half(x) * sin
cos = cos.unsqueeze(-2) # [num_tokens, 1, head_size]
sin = sin.unsqueeze(-2) # [num_tokens, 1, head_size]
@@ -335,7 +335,7 @@ def get_1d_rotary_pos_embed(
freqs_sin = freqs.sin() # [S, D/2]
if use_real:
# For rotate_half style RoPE (used by HunyuanVideo, GameCraft),
# For rotate_half style RoPE (used by HunyuanVideo),
# we need to expand cos/sin to full head_dim using repeat_interleave.
# The rotate_half operation works on consecutive PAIRS: (x0,x1), (x2,x3)...
# so cos/sin must be interleaved: [c0,c0,c1,c1,...] to match the pairing.
fastvideo/models/camera/__init__.py
-6
View File
@@ -1,6 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Camera trajectory utilities for video generation models."""
from fastvideo.models.camera.trajectory import create_camera_trajectory
__all__ = ["create_camera_trajectory"]
fastvideo/models/camera/trajectory.py
-387
View File
@@ -1,387 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
Camera trajectory generation for HunyuanGameCraft.
Generates Plücker coordinate embeddings from simple action commands
(forward, backward, left, right, rotations) for camera conditioning.
This is a self-contained implementation that does not depend on the
official Hunyuan-GameCraft-1.0 repository.
"""
import math
import numpy as np
import torch
from packaging import version as pver
# Action name -> motion type mapping
ACTION_DICT = {
"w": "forward",
"a": "left",
"d": "right",
"s": "backward",
"forward": "forward",
"backward": "backward",
"left": "left",
"right": "right",
"left_rot": "left_rot",
"right_rot": "right_rot",
"up_rot": "up_rot",
"down_rot": "down_rot",
}
def _custom_meshgrid(*args):
"""Torch meshgrid with consistent indexing."""
if pver.parse(torch.__version__) < pver.parse("1.10"):
return torch.meshgrid(*args)
else:
return torch.meshgrid(*args, indexing="ij")
def _generate_motion_segment(
current_pose: dict,
motion_type: str,
value: float,
duration: int = 30,
) -> tuple[list, list, dict]:
"""
Generate camera motion segment.
Args:
current_pose: Dict with 'position' (xyz) and 'rotation' (pitch, yaw, roll).
motion_type: One of 'forward', 'backward', 'left', 'right',
'left_rot', 'right_rot', 'up_rot', 'down_rot'.
value: Translation (meters) or rotation (degrees).
duration: Number of frames.
Returns:
positions: List of position arrays.
rotations: List of rotation arrays.
current_pose: Updated pose dict.
"""
positions = []
rotations = []
if motion_type in ["forward", "backward"]:
yaw_rad = np.radians(current_pose["rotation"][1])
pitch_rad = np.radians(current_pose["rotation"][0])
forward_vec = np.array([
-math.sin(yaw_rad) * math.cos(pitch_rad),
math.sin(pitch_rad),
-math.cos(yaw_rad) * math.cos(pitch_rad),
])
direction = 1 if motion_type == "forward" else -1
total_move = forward_vec * value * direction
step = total_move / duration
for i in range(1, duration + 1):
new_pos = current_pose["position"] + step * i
positions.append(new_pos.copy())
rotations.append(current_pose["rotation"].copy())
current_pose["position"] = positions[-1]
elif motion_type in ["left", "right"]:
yaw_rad = np.radians(current_pose["rotation"][1])
right_vec = np.array([math.cos(yaw_rad), 0, -math.sin(yaw_rad)])
direction = -1 if motion_type == "right" else 1
total_move = right_vec * value * direction
step = total_move / duration
for i in range(1, duration + 1):
new_pos = current_pose["position"] + step * i
positions.append(new_pos.copy())
rotations.append(current_pose["rotation"].copy())
current_pose["position"] = positions[-1]
elif motion_type.endswith("rot"):
axis = motion_type.split("_")[0]
total_rotation = np.zeros(3)
if axis == "left":
total_rotation[0] = value
elif axis == "right":
total_rotation[0] = -value
elif axis == "up":
total_rotation[2] = -value
elif axis == "down":
total_rotation[2] = value
step = total_rotation / duration
for i in range(1, duration + 1):
positions.append(current_pose["position"].copy())
new_rot = current_pose["rotation"] + step * i
rotations.append(new_rot.copy())
current_pose["rotation"] = rotations[-1]
return positions, rotations, current_pose
def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
"""Convert Euler angles (pitch, yaw, roll in degrees) to quaternion."""
pitch, yaw, roll = np.radians(angles)
cy = math.cos(yaw * 0.5)
sy = math.sin(yaw * 0.5)
cp = math.cos(pitch * 0.5)
sp = math.sin(pitch * 0.5)
cr = math.cos(roll * 0.5)
sr = math.sin(roll * 0.5)
qw = cy * cp * cr + sy * sp * sr
qx = cy * cp * sr - sy * sp * cr
qy = sy * cp * sr + cy * sp * cr
qz = sy * cp * cr - cy * sp * sr
return [qw, qx, qy, qz]
def _quaternion_to_rotation_matrix(q: list[float]) -> np.ndarray:
"""Convert quaternion to 3x3 rotation matrix."""
qw, qx, qy, qz = q
return np.array([
[1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
[2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
[2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
])
def _action_to_pose_list(action_id: str, value: float = 0.2, duration: int = 33) -> list[str]:
"""
Convert an action ID to a list of pose strings.
Args:
action_id: Action identifier (e.g., 'w', 'forward', 'left_rot').
value: Motion magnitude (translation in meters, rotation in degrees).
duration: Number of frames.
Returns:
List of pose strings in the official GameCraft format.
"""
all_positions = []
all_rotations = []
current_pose = {
"position": np.array([0.0, 0.0, 0.0]),
"rotation": np.array([0.0, 0.0, 0.0]),
}
intrinsic = [0.50505, 0.8979, 0.5, 0.5]
motion_type = ACTION_DICT.get(action_id, action_id)
positions, rotations, current_pose = _generate_motion_segment(current_pose, motion_type, value, duration)
all_positions.extend(positions)
all_rotations.extend(rotations)
pose_list = []
# First frame: identity pose
row = [0] + intrinsic + [0, 0] + [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]
first_row = " ".join(map(str, row))
pose_list.append(first_row)
for i, (pos, rot) in enumerate(zip(all_positions, all_rotations)):
quat = _euler_to_quaternion(rot)
R = _quaternion_to_rotation_matrix(quat)
extrinsic = np.hstack([R, pos.reshape(3, 1)])
row = [i] + intrinsic + [0, 0] + extrinsic.flatten().tolist()
pose_list.append(" ".join(map(str, row)))
return pose_list
class _Camera:
"""Camera parameters from a pose string."""
def __init__(self, entry: list[float]):
fx, fy, cx, cy = entry[1:5]
self.fx = fx
self.fy = fy
self.cx = cx
self.cy = cy
w2c_mat = np.array(entry[7:]).reshape(3, 4)
w2c_mat_4x4 = np.eye(4)
w2c_mat_4x4[:3, :] = w2c_mat
self.w2c_mat = w2c_mat_4x4
self.c2w_mat = np.linalg.inv(w2c_mat_4x4)
def _get_relative_pose(cam_params: list[_Camera]) -> np.ndarray:
"""Convert camera parameters to relative poses."""
abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1],
])
abs2rel = target_cam_c2w @ abs_w2cs[0]
ret_poses = [target_cam_c2w] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
for pose in ret_poses:
pose[:3, -1:] *= 10
ret_poses = np.array(ret_poses, dtype=np.float32)
return ret_poses
def _get_c2w(w2cs: list[np.ndarray], transform_matrix: np.ndarray) -> np.ndarray:
"""Convert w2c matrices to c2w with relative transform."""
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1],
])
abs2rel = target_cam_c2w @ w2cs[0]
ret_poses = [target_cam_c2w] + [abs2rel @ np.linalg.inv(w2c) for w2c in w2cs[1:]]
for pose in ret_poses:
pose[:3, -1:] *= 2
ret_poses = [transform_matrix @ x for x in ret_poses]
return np.array(ret_poses, dtype=np.float32)
def _ray_condition(
K: torch.Tensor,
c2w: torch.Tensor,
H: int,
W: int,
device: str | torch.device = "cpu",
flip_flag: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Compute Plücker coordinates from camera intrinsics and extrinsics.
Args:
K: Intrinsics [B, V, 4] (fx, fy, cx, cy).
c2w: Camera-to-world matrices [B, V, 4, 4].
H: Image height.
W: Image width.
device: Torch device.
flip_flag: Optional flip flags [V].
Returns:
Plücker coordinates [B, V, H, W, 6].
"""
B, V = K.shape[:2]
j, i = _custom_meshgrid(
torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),
torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype),
)
i = i.reshape([1, 1, H * W]).expand([B, V, H * W]) + 0.5
j = j.reshape([1, 1, H * W]).expand([B, V, H * W]) + 0.5
n_flip = torch.sum(flip_flag).item() if flip_flag is not None else 0
if n_flip > 0:
j_flip, i_flip = _custom_meshgrid(
torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),
torch.linspace(W - 1, 0, W, device=device, dtype=c2w.dtype),
)
i_flip = i_flip.reshape([1, 1, H * W]).expand(B, 1, H * W) + 0.5
j_flip = j_flip.reshape([1, 1, H * W]).expand(B, 1, H * W) + 0.5
i[:, flip_flag, ...] = i_flip
j[:, flip_flag, ...] = j_flip
fx, fy, cx, cy = K.chunk(4, dim=-1)
zs = torch.ones_like(i)
xs = (i - cx) / fx * zs
ys = (j - cy) / fy * zs
zs = zs.expand_as(ys)
directions = torch.stack((xs, ys, zs), dim=-1)
directions = directions / directions.norm(dim=-1, keepdim=True)
rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2)
rays_o = c2w[..., :3, 3]
rays_o = rays_o[:, :, None].expand_as(rays_d)
rays_dxo = torch.linalg.cross(rays_o, rays_d)
plucker = torch.cat([rays_dxo, rays_d], dim=-1)
plucker = plucker.reshape(B, c2w.shape[1], H, W, 6)
return plucker
def create_camera_trajectory(
action: str,
height: int,
width: int,
num_frames: int,
action_speed: float = 0.2,
device: torch.device | str = "cpu",
dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
"""
Create Plücker coordinate embeddings from an action command.
Args:
action: One of 'forward', 'backward', 'left', 'right',
'left_rot', 'right_rot', 'up_rot', 'down_rot'
(or shorthand 'w', 'a', 's', 'd').
height: Video height in pixels.
width: Video width in pixels.
num_frames: Number of video frames.
action_speed: Speed of motion (default 0.2).
device: Torch device for output tensor.
dtype: Torch dtype for output tensor.
Returns:
camera_states: [1, num_frames, 6, height, width] Plücker embeddings.
"""
# Generate pose list from action
poses = _action_to_pose_list(action, value=action_speed, duration=num_frames)
# Parse poses
poses_parsed = [pose.split(" ") for pose in poses]
start_idx = 0
sample_id = [start_idx + i for i in range(num_frames)]
poses_parsed = [poses_parsed[i] for i in sample_id]
# Convert to w2c matrices
w2cs = [np.asarray([float(p) for p in pose[7:]]).reshape(3, 4) for pose in poses_parsed]
transform_matrix = np.asarray([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]]).reshape(4, 4)
last_row = np.zeros((1, 4))
last_row[0, -1] = 1.0
w2cs = [np.concatenate((w2c, last_row), axis=0) for w2c in w2cs]
c2ws = _get_c2w(w2cs, transform_matrix)
# Parse camera parameters
cam_params = [[float(x) for x in pose] for pose in poses_parsed]
assert len(cam_params) == num_frames
cam_params = [_Camera(cam_param) for cam_param in cam_params]
# Compute scaled intrinsics
monst3r_w = cam_params[0].cx * 2
monst3r_h = cam_params[0].cy * 2
ratio_w, ratio_h = width / monst3r_w, height / monst3r_h
intrinsics = np.asarray(
[[
cam_param.fx * ratio_w,
cam_param.fy * ratio_h,
cam_param.cx * ratio_w,
cam_param.cy * ratio_h,
] for cam_param in cam_params],
dtype=np.float32,
)
intrinsics = torch.as_tensor(intrinsics)[None] # [1, n_frame, 4]
# Get relative poses
c2w_poses = _get_relative_pose(cam_params)
c2w = torch.as_tensor(c2w_poses)[None] # [1, n_frame, 4, 4]
# Compute Plücker embeddings
flip_flag = torch.zeros(num_frames, dtype=torch.bool, device="cpu")
plucker_embedding = _ray_condition(intrinsics, c2w, height, width, device="cpu", flip_flag=flip_flag)
# [1, n_frame, H, W, 6] -> [1, n_frame, 6, H, W]
plucker_embedding = plucker_embedding[0].permute(0, 3, 1, 2).contiguous()
# Add batch dim and convert to target dtype/device
# Shape: [n_frame, 6, H, W] -> [1, n_frame, 6, H, W]
camera_states = plucker_embedding.unsqueeze(0).to(device=device, dtype=dtype)
return camera_states
fastvideo/models/dits/dreamx_world.py
-452
View File
@@ -1,452 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
import math
from typing import Any
import torch
import torch.nn as nn
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig
from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad,
sequence_model_parallel_shard)
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.layers.layernorm import RMSNorm
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
from fastvideo.models.wan.transformer import (LayerNormScaleShift, PatchEmbed, WanTimeTextImageEmbedding,
WanTransformer3DModel, WanTransformerBlock)
from fastvideo.platforms import AttentionBackendEnum, current_platform
from fastvideo.attention import LocalAttention
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.layers.quantization import QuantizationConfig
from fastvideo.models.dits.base import BaseDiT
def _dreamx_invert_se3(transforms: torch.Tensor) -> torch.Tensor:
assert transforms.shape[-2:] == (4, 4)
rot_inv = transforms[..., :3, :3].transpose(-1, -2)
out = torch.zeros_like(transforms)
out[..., :3, :3] = rot_inv
out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, transforms[..., :3, 3])
out[..., 3, 3] = 1.0
return out.to(dtype=transforms.dtype)
def _dreamx_lift_k(intrinsics: torch.Tensor) -> torch.Tensor:
assert intrinsics.shape[-2:] == (3, 3)
out = torch.zeros(intrinsics.shape[:-2] + (4, 4), device=intrinsics.device, dtype=intrinsics.dtype)
out[..., :3, :3] = intrinsics
out[..., 3, 3] = 1.0
return out
def _dreamx_invert_k(intrinsics: torch.Tensor) -> torch.Tensor:
assert intrinsics.shape[-2:] == (3, 3)
out = torch.zeros_like(intrinsics)
out[..., 0, 0] = 1.0 / intrinsics[..., 0, 0]
out[..., 1, 1] = 1.0 / intrinsics[..., 1, 1]
out[..., 0, 2] = -intrinsics[..., 0, 2] / intrinsics[..., 0, 0]
out[..., 1, 2] = -intrinsics[..., 1, 2] / intrinsics[..., 1, 1]
out[..., 2, 2] = 1.0
return out.to(dtype=intrinsics.dtype)
def _dreamx_apply_tiled_projmat(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
batch, num_heads, seq_len, feat_dim = feats.shape
proj_dim = matrix.shape[-1]
assert feat_dim % proj_dim == 0
if matrix.shape[1] == seq_len:
feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim, proj_dim)
out = torch.einsum("btij,bntpj->bntpi", matrix, feats)
return out.reshape(batch, num_heads, seq_len, feat_dim)
cameras = matrix.shape[1]
assert seq_len > cameras and seq_len % cameras == 0
feats = feats.reshape(batch, num_heads, cameras, -1, feat_dim // proj_dim, proj_dim)
out = torch.einsum("bcij,bncpkj->bncpki", matrix, feats)
return out.reshape(batch, num_heads, seq_len, feat_dim)
def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, viewmats: torch.Tensor,
intrinsics: torch.Tensor):
batch, num_heads, seq_len, head_dim = q.shape
cameras = viewmats.shape[1]
assert q.shape == k.shape == v.shape
assert viewmats.shape == (batch, cameras, 4, 4)
assert intrinsics.shape == (batch, cameras, 3, 3)
assert head_dim % 4 == 0
intrinsics_norm = torch.zeros_like(intrinsics)
intrinsics_norm[..., 0, 0] = intrinsics[..., 0, 0]
intrinsics_norm[..., 1, 1] = intrinsics[..., 1, 1]
intrinsics_norm[..., 2, 2] = 1.0
proj = torch.einsum("...ij,...jk->...ik", _dreamx_lift_k(intrinsics_norm), viewmats)
proj_t = proj.transpose(-1, -2).to(dtype=viewmats.dtype)
proj_inv = torch.einsum(
"...ij,...jk->...ik",
_dreamx_invert_se3(viewmats),
_dreamx_lift_k(_dreamx_invert_k(intrinsics_norm)),
).to(dtype=viewmats.dtype)
q = _dreamx_apply_tiled_projmat(q, proj_t)
k = _dreamx_apply_tiled_projmat(k, proj_inv)
v = _dreamx_apply_tiled_projmat(v, proj_inv)
return q, k, v, proj
class DreamXPropeSelfAttention(nn.Module):
"""DreamX-World parallel PRoPE camera self-attention branch."""
def __init__(self,
dim: int,
attn_dim: int,
num_heads: int,
qk_norm: str | bool = True,
eps: float = 1e-6,
quant_config: QuantizationConfig | None = None,
prefix: str = ""):
super().__init__()
assert attn_dim % num_heads == 0
self.attn_dim = attn_dim
self.num_heads = num_heads
self.head_dim = attn_dim // num_heads
self.qk_norm = qk_norm
self.q_proj = ReplicatedLinear(dim, attn_dim, quant_config=quant_config, prefix=f"{prefix}.q_proj")
self.k_proj = ReplicatedLinear(dim, attn_dim, quant_config=quant_config, prefix=f"{prefix}.k_proj")
self.v_proj = ReplicatedLinear(dim, attn_dim, quant_config=quant_config, prefix=f"{prefix}.v_proj")
self.out_proj = ReplicatedLinear(attn_dim, dim, quant_config=quant_config, prefix=f"{prefix}.out_proj")
if qk_norm == "rms_norm":
self.norm_q = RMSNorm(self.head_dim, eps=eps)
self.norm_k = RMSNorm(self.head_dim, eps=eps)
elif qk_norm in (True, "rms_norm_across_heads"):
self.norm_q = RMSNorm(attn_dim, eps=eps)
self.norm_k = RMSNorm(attn_dim, eps=eps)
elif qk_norm is False:
self.norm_q = nn.Identity()
self.norm_k = nn.Identity()
else:
raise ValueError(f"Unsupported qk_norm for DreamX PRoPE: {qk_norm}")
nn.init.zeros_(self.out_proj.weight)
if self.out_proj.bias is not None:
nn.init.zeros_(self.out_proj.bias)
self.attn = LocalAttention(num_heads=num_heads,
head_size=self.head_dim,
dropout_rate=0,
softmax_scale=None,
causal=False,
supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA))
def forward(self, hidden_states: torch.Tensor, y_camera: dict[str, torch.Tensor]) -> torch.Tensor:
if get_sp_world_size() > 1:
# The transformer shards the sequence before the block loop and
# this branch uses LocalAttention (no all-to-all): under
# sequence parallelism each rank would attend only within its
# own shard — silently wrong output. Fail loudly until this
# path is ported to DistributedAttention and validated.
raise NotImplementedError("DreamXPropeSelfAttention does not support sequence "
"parallelism yet (LocalAttention on a sharded sequence "
"corrupts output). Run with sp_size=1.")
batch_size, seq_len, _ = hidden_states.shape
query, _ = self.q_proj(hidden_states)
key, _ = self.k_proj(hidden_states)
value, _ = self.v_proj(hidden_states)
if self.qk_norm == "rms_norm":
query = query.view(batch_size, seq_len, self.num_heads, self.head_dim)
key = key.view(batch_size, seq_len, self.num_heads, self.head_dim)
query = self.norm_q(query)
key = self.norm_k(key)
else:
query = self.norm_q(query).view(batch_size, seq_len, self.num_heads, self.head_dim)
key = self.norm_k(key).view(batch_size, seq_len, self.num_heads, self.head_dim)
value = value.view(batch_size, seq_len, self.num_heads, self.head_dim)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
query, key, value, output_projection = _dreamx_prope_qkv(
query,
key,
value,
viewmats=y_camera["viewmats"],
intrinsics=y_camera["K"],
)
out = self.attn(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
out = _dreamx_apply_tiled_projmat(out.transpose(1, 2), output_projection).transpose(1, 2)
out = out.flatten(2)
out, _ = self.out_proj(out)
return out
class DreamXWorldTransformerBlock(WanTransformerBlock):
def __init__(self,
dim: int,
ffn_dim: int,
num_heads: int,
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
added_kv_proj_dim: int | None = None,
supported_attention_backends: tuple[AttentionBackendEnum, ...]
| None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
add_control_adapter: bool = True,
cam_method: str | None = "prope",
attn_compress: int = 1,
cam_self_attn_layers: tuple[int, ...] | None = None,
layer_idx: int | None = None):
super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim,
supported_attention_backends, quant_config, prefix)
self.cam_self_attn = None
add_cam_attn = add_control_adapter and cam_method == "prope"
if add_cam_attn and cam_self_attn_layers is not None:
add_cam_attn = layer_idx in cam_self_attn_layers
if add_cam_attn:
if num_heads % attn_compress != 0 or dim % attn_compress != 0:
raise ValueError("DreamX attn_compress must divide dim and num_heads")
self.cam_self_attn = DreamXPropeSelfAttention(dim,
dim // attn_compress,
num_heads // attn_compress,
qk_norm=qk_norm,
eps=eps,
quant_config=quant_config,
prefix=f"{prefix}.cam_self_attn")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
original_seq_len: int,
y_camera: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
if hidden_states.dim() == 4:
hidden_states = hidden_states.squeeze(1)
orig_dtype = hidden_states.dtype
if temb.dim() == 4:
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
shift_msa = shift_msa.squeeze(2)
scale_msa = scale_msa.squeeze(2)
gate_msa = gate_msa.squeeze(2)
c_shift_msa = c_shift_msa.squeeze(2)
c_scale_msa = c_scale_msa.squeeze(2)
c_gate_msa = c_gate_msa.squeeze(2)
else:
e = self.scale_shift_table + temb.float()
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1)
assert shift_msa.dtype == torch.float32
norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype)
query, _ = self.to_q(norm_hidden_states)
key, _ = self.to_k(norm_hidden_states)
value, _ = self.to_v(norm_hidden_states)
if self.norm_q is not None:
query = self.norm_q(query)
if self.norm_k is not None:
key = self.norm_k(key)
query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
attn_output, _ = self.attn1(
query,
key,
value,
original_seq_len,
freqs_cis=freqs_cis,
)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)
if self.cam_self_attn is not None and y_camera is not None:
attn_output = attn_output + self.cam_self_attn(norm_hidden_states, y_camera)
null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
null_shift, null_scale)
norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)
norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa,
c_scale_msa)
norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
ff_output = self.ffn(norm_hidden_states)
hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
hidden_states = hidden_states.to(orig_dtype)
return hidden_states
class DreamXWorldTransformer3DModel(WanTransformer3DModel):
_fsdp_shard_conditions = DreamXWorldConfig()._fsdp_shard_conditions
_compile_conditions = DreamXWorldConfig()._compile_conditions
_supported_attention_backends = DreamXWorldConfig()._supported_attention_backends
param_names_mapping = DreamXWorldConfig().param_names_mapping
reverse_param_names_mapping = DreamXWorldConfig().reverse_param_names_mapping
lora_param_names_mapping = DreamXWorldConfig().lora_param_names_mapping
def __init__(self, config: DreamXWorldConfig, hf_config: dict[str, Any]) -> None:
BaseDiT.__init__(self, config=config, hf_config=hf_config)
self.quant_config = config.quant_config
inner_dim = config.num_attention_heads * config.attention_head_dim
self.hidden_size = config.hidden_size
self.num_attention_heads = config.num_attention_heads
self.in_channels = config.in_channels
self.out_channels = config.out_channels
self.num_channels_latents = config.num_channels_latents
self.patch_size = config.patch_size
self.text_len = config.text_len
assert config.num_attention_heads % get_sp_world_size(
) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
embed_dim=inner_dim,
patch_size=config.patch_size,
flatten=False)
self.condition_embedder = WanTimeTextImageEmbedding(
dim=inner_dim,
time_freq_dim=config.freq_dim,
text_embed_dim=config.text_dim,
image_embed_dim=config.image_dim,
)
self.blocks = nn.ModuleList([
DreamXWorldTransformerBlock(inner_dim,
config.ffn_dim,
config.num_attention_heads,
config.qk_norm,
config.cross_attn_norm,
config.eps,
config.added_kv_proj_dim,
self._supported_attention_backends,
quant_config=config.quant_config,
prefix=f"{config.prefix}.blocks.{i}",
add_control_adapter=config.add_control_adapter,
cam_method=config.cam_method,
attn_compress=config.attn_compress,
cam_self_attn_layers=config.cam_self_attn_layers,
layer_idx=i) for i in range(config.num_layers)
])
self.norm_out = LayerNormScaleShift(inner_dim,
norm_type="layer",
eps=config.eps,
elementwise_affine=False,
dtype=torch.float32,
compute_dtype=torch.float32)
self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size))
self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
self.gradient_checkpointing = False
self.__post_init__()
def forward(self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | list[torch.Tensor],
timestep: torch.LongTensor,
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
| None = None,
guidance=None,
y_camera: dict[str, torch.Tensor] | None = None,
**kwargs) -> torch.Tensor:
orig_dtype = hidden_states.dtype
if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
encoder_hidden_states = encoder_hidden_states[0]
if isinstance(encoder_hidden_states_image, list):
encoder_hidden_states_image = (encoder_hidden_states_image[0]
if len(encoder_hidden_states_image) > 0 else None)
batch_size, _, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.patch_size
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p_h
post_patch_width = width // p_w
d = self.hidden_size // self.num_attention_heads
rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width),
self.hidden_size,
self.num_attention_heads,
rope_dim_list,
dtype=torch.float32 if current_platform.is_mps() else torch.float64,
rope_theta=10000)
freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float())
hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
if timestep.dim() == 2:
ts_seq_len = timestep.shape[1]
timestep = timestep.flatten()
else:
ts_seq_len = None
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
if ts_seq_len is not None:
timestep_proj = timestep_proj.unflatten(2, (6, -1))
else:
timestep_proj = timestep_proj.unflatten(1, (6, -1))
if encoder_hidden_states_image is not None:
if encoder_hidden_states is not None:
encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
else:
encoder_hidden_states = encoder_hidden_states_image
if current_platform.is_mps() or current_platform.is_npu():
encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
assert encoder_hidden_states.dtype == orig_dtype
if torch.is_grad_enabled() and self.gradient_checkpointing:
for block in self.blocks:
hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, original_seq_len, y_camera)
else:
for block in self.blocks:
hidden_states = block(hidden_states,
encoder_hidden_states,
timestep_proj,
freqs_cis,
original_seq_len,
y_camera=y_camera)
if temb.dim() == 3:
shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
shift = shift.squeeze(2)
scale = scale.squeeze(2)
else:
shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
hidden_states = self.norm_out(hidden_states, shift, scale)
hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
hidden_states = self.proj_out(hidden_states)
hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
p_t, p_h, p_w, -1)
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
return output
EntryClass = DreamXWorldTransformer3DModel
fastvideo/models/dits/dreamx_world_ar.py
-982
View File
@@ -1,982 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World autoregressive causal DiT.
Adapted from DreamX-World's Apache-2.0
``wan/modules/causal_camera_model_2_2_prope_infinity.py``. The implementation is
kept native to FastVideo: no production import from DreamX, Diffusers, or
Transformers is required.
"""
from typing import Any
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.models.dits.base import BaseDiT
from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat, _dreamx_prope_qkv)
def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
# Deliberately raw SDPA rather than fastvideo.attention.LocalAttention:
# (1) LocalAttention dispatches through the attention-backend registry, so
# FLASH_ATTN could be selected and its kernel is not bit-identical to
# torch SDPA — the AR KV-cache rollout must stay numerically frozen;
# (2) LocalAttention requires an active ForwardContext, which direct
# transformer invocations (parity tests) do not set;
# (3) the sibling causal model keeps raw SDPA in the same KV-cache window
# path (matrixgame2/causal_model.py).
# Sequence-parallel gap: this model never shards the sequence; run with
# sp_size=1 (see fastvideo/layers/AGENTS.md on documenting raw SDPA).
q_bhld = q.transpose(1, 2)
k_bhld = k.transpose(1, 2)
v_bhld = v.transpose(1, 2)
out = F.scaled_dot_product_attention(q_bhld, k_bhld, v_bhld, dropout_p=0.0)
return out.transpose(1, 2)
def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, viewmats: torch.Tensor, Ks: torch.Tensor):
q, k, v, output_projection = _dreamx_prope_qkv(q, k, v, viewmats, Ks)
def apply_fn_o(x: torch.Tensor) -> torch.Tensor:
return _dreamx_apply_tiled_projmat(x, output_projection)
return q, k, v, apply_fn_o
def sinusoidal_embedding_1d(dim, position):
assert dim % 2 == 0
half = dim // 2
position = position.type(torch.float64)
sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
def rope_params(max_seq_len, dim, theta=10000):
assert dim % 2 == 0
freqs = torch.outer(torch.arange(max_seq_len), 1.0 / torch.pow(theta,
torch.arange(0, dim, 2).to(torch.float64).div(dim)))
return torch.polar(torch.ones_like(freqs), freqs)
class WanRMSNorm(nn.Module):
"""Kept private instead of fastvideo.layers.layernorm.RMSNorm.
The official DreamX-World ``model_2_2.py`` computes the RMS statistics in
the *input* dtype — the upstream code has the fp32 upcast explicitly
commented out (``# return self._norm(x.float())...``). FastVideo's RMSNorm
always normalizes in fp32, which is not bit-identical under bf16, so the
verbatim implementation stays.
"""
def __init__(self, dim, eps=1e-5):
super().__init__()
self.dim = dim
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
return self._norm(x).type_as(x) * self.weight
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
class WanLayerNorm(nn.LayerNorm):
"""Kept private instead of fastvideo.layers.layernorm.FP32LayerNorm.
The official DreamX-World ``model_2_2.py`` normalizes in the *input* dtype
(no ``x.float()`` upcast, unlike Wan2.1). FP32LayerNorm casts input and
affine params to fp32, which is not bit-identical under bf16, so the
verbatim implementation stays.
"""
def __init__(self, dim, eps=1e-6, elementwise_affine=False):
super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
def forward(self, x):
return super().forward(x).type_as(x)
class WanCrossAttention(nn.Module):
def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.q = ReplicatedLinear(dim, dim)
self.k = ReplicatedLinear(dim, dim)
self.v = ReplicatedLinear(dim, dim)
self.o = ReplicatedLinear(dim, dim)
self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
def _kv(self, context, b, n, d):
k, _ = self.k(context)
k = self.norm_k(k).view(b, -1, n, d)
v, _ = self.v(context)
v = v.view(b, -1, n, d)
return k, v
def forward(self, x, context, context_lens, crossattn_cache=None):
b, n, d = x.size(0), self.num_heads, self.head_dim
q, _ = self.q(x)
q = self.norm_q(q).view(b, -1, n, d)
if crossattn_cache is not None:
if not crossattn_cache["is_init"]:
crossattn_cache["is_init"] = True
k, v = self._kv(context, b, n, d)
crossattn_cache["k"] = k
crossattn_cache["v"] = v
else:
k = crossattn_cache["k"]
v = crossattn_cache["v"]
else:
k, v = self._kv(context, b, n, d)
x = attention(q, k, v)
x = x.flatten(2)
out, _ = self.o(x)
return out
def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_indices=None):
"""
Apply Block-Relativistic RoPE to input tensor.
Adapted from Infinity-RoPE (https://arxiv.org/abs/2511.20649).
Args:
x: Input tensor [B, L, num_heads, head_dim]
grid_sizes: Tensor [B, 3] containing (F, H, W)
freqs: RoPE frequencies
start_frame: Starting frame index for sequential RoPE
relative_frame_indices: Optional tensor [F] specifying explicit frame indices
for Block-Relativistic RoPE. Overrides start_frame if provided.
"""
n, c = x.size(2), x.size(3) // 2
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
output = []
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
seq_len = f * h * w
x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2))
if relative_frame_indices is not None:
frame_indices = relative_frame_indices.long()
freqs_temporal = freqs[0][frame_indices].view(f, 1, 1, -1).expand(f, h, w, -1)
else:
freqs_temporal = freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1)
freqs_i = torch.cat([
freqs_temporal, freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), freqs[2][:w].view(1, 1, w, -1).expand(
f, h, w, -1)
],
dim=-1).reshape(seq_len, 1, -1)
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
x_i = torch.cat([x_i, x[i, seq_len:]])
output.append(x_i)
return torch.stack(output).type_as(x)
class CausalWanSelfAttention(nn.Module):
"""Self-attention with KV cache and Block-Relativistic RoPE for causal inference."""
def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1, qk_norm=True, eps=1e-6):
assert dim % num_heads == 0
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.local_attn_size = local_attn_size
self.sink_size = sink_size
self.qk_norm = qk_norm
self.eps = eps
self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
self.q = ReplicatedLinear(dim, dim)
self.k = ReplicatedLinear(dim, dim)
self.v = ReplicatedLinear(dim, dim)
self.o = ReplicatedLinear(dim, dim)
self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
def forward(self,
x,
seq_lens,
grid_sizes,
freqs,
kv_cache,
current_start=0,
cache_start=None,
sink_recache_after_switch=False):
"""
Args:
x: Shape [B, L, C]
seq_lens: Shape [B]
grid_sizes: Shape [B, 3] containing (F, H, W)
freqs: RoPE frequencies [1024, head_dim / 2]
kv_cache: Dict with 'k', 'v', 'global_end_index', 'local_end_index'
current_start: Current position in the global token sequence
"""
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
if cache_start is None:
cache_start = current_start
q, _ = self.q(x)
q = self.norm_q(q).view(b, s, n, d)
k, _ = self.k(x)
k = self.norm_k(k).view(b, s, n, d)
v, _ = self.v(x)
v = v.view(b, s, n, d)
frame_seqlen = math.prod(grid_sizes[0][1:]).item()
num_new_frames = grid_sizes[0][0].item()
current_end = current_start + q.shape[1]
sink_tokens = self.sink_size * frame_seqlen
kv_cache_size = kv_cache["k"].shape[1]
num_new_tokens = q.shape[1]
cache_update_info = None
is_recompute = current_end <= kv_cache["global_end_index"].item() and current_start > 0
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
# === ROLLING MODE: cache full, evict oldest non-sink tokens ===
num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
local_end_index = kv_cache["local_end_index"].item() + current_end - \
kv_cache["global_end_index"].item() - num_evicted_tokens
local_start_index = local_end_index - num_new_tokens
temp_k = kv_cache["k"].detach().clone()
temp_v = kv_cache["v"].detach().clone()
temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
roped_offset = max(0, write_start_index - local_start_index)
write_len = max(0, local_end_index - write_start_index)
if write_len > 0:
temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
# Block-Relativistic RoPE: query uses window-relative indices
query_relative_indices = torch.arange(self.local_attn_size - num_new_frames,
self.local_attn_size,
device=q.device)
roped_query = block_relativistic_rope(q, grid_sizes, freqs,
relative_frame_indices=query_relative_indices).type_as(v)
# Block-Relativistic RoPE: cached K uses position-in-window indices
num_cache_frames = local_end_index // frame_seqlen
cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
cache_grid_sizes = grid_sizes.clone()
cache_grid_sizes[0, 0] = num_cache_frames
roped_temp_k = block_relativistic_rope(temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen,
n, d).flatten(1, 2),
cache_grid_sizes,
freqs,
relative_frame_indices=cache_relative_indices).type_as(v)
cache_update_info = {
"action": "roll_and_insert",
"sink_tokens": sink_tokens,
"num_rolled_tokens": num_rolled_tokens,
"num_evicted_tokens": num_evicted_tokens,
"local_start_index": local_start_index,
"local_end_index": local_end_index,
"write_start_index": write_start_index,
"write_end_index": local_end_index,
"new_k": k[:, roped_offset:roped_offset + write_len],
"new_v": v[:, roped_offset:roped_offset + write_len],
"current_end": current_end,
"is_recompute": is_recompute
}
else:
# === DIRECT INSERT MODE: cache not yet full ===
local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
local_start_index = local_end_index - num_new_tokens
temp_k = kv_cache["k"].detach().clone()
temp_v = kv_cache["v"].detach().clone()
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
if sink_recache_after_switch:
write_start_index = local_start_index
roped_offset = max(0, write_start_index - local_start_index)
write_len = max(0, local_end_index - write_start_index)
if write_len > 0:
temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
# RoPE with relative indices (growing sequentially before cache fills)
current_frame_in_window = local_start_index // frame_seqlen
query_relative_indices = torch.arange(current_frame_in_window,
current_frame_in_window + num_new_frames,
device=q.device)
roped_query = block_relativistic_rope(q, grid_sizes, freqs,
relative_frame_indices=query_relative_indices).type_as(v)
num_cache_frames = local_end_index // frame_seqlen
cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
cache_grid_sizes = grid_sizes.clone()
cache_grid_sizes[0, 0] = num_cache_frames
roped_temp_k = block_relativistic_rope(temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen,
n, d).flatten(1, 2),
cache_grid_sizes,
freqs,
relative_frame_indices=cache_relative_indices).type_as(v)
cache_update_info = {
"action": "direct_insert",
"local_start_index": local_start_index,
"local_end_index": local_end_index,
"write_start_index": write_start_index,
"write_end_index": local_end_index,
"new_k": k[:, roped_offset:roped_offset + write_len],
"new_v": v[:, roped_offset:roped_offset + write_len],
"current_end": current_end,
"is_recompute": is_recompute
}
# Attention: sink tokens + local window
if sink_tokens > 0:
local_budget = self.max_attention_size - sink_tokens
k_sink = roped_temp_k[:, :sink_tokens]
v_sink = temp_v[:, :sink_tokens]
if local_budget > 0:
local_start_for_window = max(sink_tokens, local_end_index - local_budget)
k_local = roped_temp_k[:, local_start_for_window:local_end_index]
v_local = temp_v[:, local_start_for_window:local_end_index]
k_cat = torch.cat([k_sink, k_local], dim=1)
v_cat = torch.cat([v_sink, v_local], dim=1)
else:
k_cat = k_sink
v_cat = v_sink
x = attention(roped_query, k_cat, v_cat)
else:
window_start = max(0, local_end_index - self.max_attention_size)
x = attention(roped_query, roped_temp_k[:, window_start:local_end_index],
temp_v[:, window_start:local_end_index])
x = x.flatten(2)
x, _ = self.o(x)
return x, (current_end, local_end_index, cache_update_info)
class CausalPropeSelfAttention(nn.Module):
"""PRoPE self-attention with optional KV cache for camera-controlled inference."""
def __init__(self,
dim,
attn_dim,
num_heads,
window_size=(-1, -1),
local_attn_size=-1,
sink_size=0,
qk_norm=True,
eps=1e-6):
assert dim % num_heads == 0
assert attn_dim % num_heads == 0
super().__init__()
self.dim = dim
self.attn_dim = attn_dim
self.num_heads = num_heads
self.head_dim = attn_dim // num_heads
self.local_attn_size = local_attn_size
self.sink_size = sink_size
self.qk_norm = qk_norm
self.eps = eps
self.window_size = window_size
self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
self.q_proj = ReplicatedLinear(dim, attn_dim)
self.k_proj = ReplicatedLinear(dim, attn_dim)
self.v_proj = ReplicatedLinear(dim, attn_dim)
self.out_proj = ReplicatedLinear(attn_dim, dim)
self.norm_q = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
self.norm_k = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
nn.init.zeros_(self.out_proj.weight)
nn.init.zeros_(self.out_proj.bias)
def forward(self,
x,
cam_viewmats,
cam_K,
seq_lens,
grid_sizes,
freqs,
kv_cache=None,
current_start=0,
cache_start=None,
sink_recache_after_switch=False,
cache_update_policy="commit_detached"):
"""
Args:
x: Shape [B, L, C]
cam_viewmats: Camera view matrices
cam_K: Camera intrinsics
kv_cache: Optional KV cache dict. When None, runs full attention over current chunk.
"""
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
if cache_start is None:
cache_start = current_start
q, _ = self.q_proj(x)
q = self.norm_q(q).view(b, s, n, d)
k, _ = self.k_proj(x)
k = self.norm_k(k).view(b, s, n, d)
v, _ = self.v_proj(x)
v = v.view(b, s, n, d)
# Apply PRoPE (Positional Rotary Position Embedding from camera parameters)
q_t, k_t, v_t, apply_fn_o = prope_qkv(q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
viewmats=cam_viewmats,
Ks=cam_K)
proped_q = q_t.transpose(1, 2)
proped_k = k_t.transpose(1, 2)
proped_v = v_t.transpose(1, 2)
if kv_cache is None:
# No cache: full attention over current chunk
x_out = attention(proped_q, proped_k, proped_v)
else:
# KV cache mode with rolling cache support
frame_seqlen = math.prod(grid_sizes[0][1:]).item()
num_new_tokens = s
current_end = current_start + num_new_tokens
sink_tokens = self.sink_size * frame_seqlen
kv_cache_size = kv_cache["k"].shape[1]
is_recompute = (current_end <= kv_cache["global_end_index"].item()) and (current_start > 0)
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
# === ROLLING MODE ===
num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
local_end_index = kv_cache["local_end_index"].item() + current_end - \
kv_cache["global_end_index"].item() - num_evicted_tokens
local_start_index = local_end_index - num_new_tokens
if cache_update_policy != "none":
with torch.no_grad():
kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
roped_offset = max(0, write_start_index - local_start_index)
write_len = max(0, local_end_index - write_start_index)
if write_len > 0:
with torch.no_grad():
kv_cache["k"][:,
write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset +
write_len].detach()
kv_cache["v"][:,
write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset +
write_len].detach()
else:
# === DIRECT INSERT MODE ===
local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
local_start_index = local_end_index - num_new_tokens
if cache_update_policy != "none":
write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
if sink_recache_after_switch:
write_start_index = local_start_index
roped_offset = max(0, write_start_index - local_start_index)
write_len = max(0, local_end_index - write_start_index)
if write_len > 0:
with torch.no_grad():
kv_cache["k"][:,
write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset +
write_len].detach()
kv_cache["v"][:,
write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset +
write_len].detach()
# Attention: sink tokens + local window
if sink_tokens > 0:
local_budget = self.max_attention_size - sink_tokens
k_sink = kv_cache["k"][:, :sink_tokens].detach()
v_sink = kv_cache["v"][:, :sink_tokens].detach()
if local_budget > 0:
local_start_for_window = max(sink_tokens, local_end_index - local_budget)
k_local = kv_cache["k"][:, local_start_for_window:local_end_index].detach()
v_local = kv_cache["v"][:, local_start_for_window:local_end_index].detach()
k_cat = torch.cat([k_sink, k_local], dim=1)
v_cat = torch.cat([v_sink, v_local], dim=1)
else:
k_cat = k_sink
v_cat = v_sink
x_out = attention(proped_q, k_cat, v_cat)
else:
window_start = max(0, local_end_index - self.max_attention_size)
x_out = attention(proped_q, kv_cache["k"][:, window_start:local_end_index].detach(),
kv_cache["v"][:, window_start:local_end_index].detach())
if not is_recompute and cache_update_policy != "none":
kv_cache["global_end_index"].fill_(current_end)
kv_cache["local_end_index"].fill_(local_end_index)
# Apply inverse PRoPE
x = apply_fn_o(x_out.transpose(1, 2)).transpose(1, 2)
x = x.flatten(2)
x, _ = self.out_proj(x)
return x
class CausalWanAttentionBlock(nn.Module):
def __init__(self,
dim,
ffn_dim,
num_heads,
local_attn_size=-1,
sink_size=0,
qk_norm=True,
cross_attn_norm=False,
eps=1e-6,
**kwargs):
super().__init__()
self.dim = dim
self.ffn_dim = ffn_dim
self.num_heads = num_heads
self.local_attn_size = local_attn_size
self.qk_norm = qk_norm
self.cross_attn_norm = cross_attn_norm
self.eps = eps
self.add_control_adapter = kwargs.get('add_control_adapter', False)
self.cam_method = kwargs.get('cam_method')
self.attn_compress = kwargs.get('attn_compress', 1)
self.layer_idx = kwargs.get('layer_idx')
cam_self_attn_layers = kwargs.get('cam_self_attn_layers')
# layers
self.norm1 = WanLayerNorm(dim, eps)
self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
self.norm2 = WanLayerNorm(dim, eps)
# nn.Linear (not ReplicatedLinear) on purpose: the official checkpoint
# stores these as positional Sequential keys (ffn.0 / ffn.2) that the
# copy-only converter and the strict-load tests require verbatim, and
# ReplicatedLinear's (out, bias) tuple return cannot compose inside
# nn.Sequential without renaming the state-dict surface.
self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), nn.Linear(ffn_dim, dim))
# PRoPE self-attention branch for camera control
add_cam_attn = self.add_control_adapter and self.cam_method == 'prope'
if add_cam_attn and cam_self_attn_layers is not None:
add_cam_attn = self.layer_idx in cam_self_attn_layers
if add_cam_attn:
self.cam_self_attn = CausalPropeSelfAttention(dim,
dim // self.attn_compress,
num_heads,
local_attn_size=local_attn_size,
sink_size=sink_size,
qk_norm=qk_norm,
eps=eps)
# modulation
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
def forward(self,
x,
e,
seq_lens,
grid_sizes,
freqs,
context,
context_lens,
kv_cache,
crossattn_cache=None,
current_start=0,
cache_start=None,
cam_viewmats=None,
cam_K=None,
sink_recache_after_switch=False,
cache_update_policy="commit_detached"):
num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
# self-attention
attn_input = (self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(
1, 2)
y, cache_update_info = self.self_attn(attn_input, seq_lens, grid_sizes, freqs, kv_cache, current_start,
cache_start, sink_recache_after_switch)
# PRoPE camera attention (parallel branch)
if hasattr(self, 'cam_self_attn') and cam_viewmats is not None and cam_K is not None:
prope_kv_cache = None
if kv_cache is not None and "prope_k" in kv_cache:
prope_kv_cache = {
"k": kv_cache["prope_k"],
"v": kv_cache["prope_v"],
"global_end_index": kv_cache["prope_global_end_index"],
"local_end_index": kv_cache["prope_local_end_index"],
}
y = y + self.cam_self_attn(attn_input,
cam_viewmats,
cam_K,
seq_lens,
grid_sizes,
freqs,
kv_cache=prope_kv_cache,
current_start=current_start,
cache_start=cache_start,
cache_update_policy=cache_update_policy)
x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)
# cross-attention & FFN
x = x + self.cross_attn(self.norm3(x), context, context_lens, crossattn_cache=crossattn_cache)
y = self.ffn(
(self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[4]) + e[3]).flatten(1, 2))
x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2)
return x, cache_update_info
class CausalHead(nn.Module):
def __init__(self, dim, out_dim, patch_size, eps=1e-6):
super().__init__()
self.dim = dim
self.out_dim = out_dim
self.patch_size = patch_size
self.eps = eps
out_dim = math.prod(patch_size) * out_dim
self.norm = WanLayerNorm(dim, eps)
self.head = ReplicatedLinear(dim, out_dim)
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
def forward(self, x, e):
num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
x, _ = self.head(self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0])
return x
class DreamXWorldARTransformer3DModel(BaseDiT):
"""DreamX-World-5B autoregressive causal transformer."""
_fsdp_shard_conditions = DreamXWorldARConfig()._fsdp_shard_conditions
_compile_conditions = DreamXWorldARConfig()._compile_conditions
_supported_attention_backends = DreamXWorldARConfig()._supported_attention_backends
param_names_mapping = DreamXWorldARConfig().param_names_mapping
reverse_param_names_mapping = DreamXWorldARConfig().reverse_param_names_mapping
lora_param_names_mapping = DreamXWorldARConfig().lora_param_names_mapping
_no_split_modules = ["CausalWanAttentionBlock"]
def __init__(self, config: DreamXWorldARConfig, hf_config: dict[str, Any]) -> None:
super().__init__(config=config, hf_config=hf_config)
model_type = config.model_type
patch_size = config.patch_size
text_len = config.text_len
in_dim = config.in_channels
dim = config.hidden_size
ffn_dim = config.ffn_dim
freq_dim = config.freq_dim
text_dim = config.text_dim
out_dim = config.out_channels
num_heads = config.num_attention_heads
num_layers = config.num_layers
local_attn_size = config.local_attn_size
sink_size = config.sink_size
qk_norm = bool(config.qk_norm)
cross_attn_norm = config.cross_attn_norm
eps = config.eps
add_control_adapter = config.add_control_adapter
cam_method = config.cam_method
attn_compress = config.attn_compress
cam_self_attn_layers = config.cam_self_attn_layers
assert model_type in ['t2v', 'i2v', 'ti2v']
self.model_type = model_type
self.patch_size = patch_size
self.text_len = text_len
self.in_dim = in_dim
self.dim = dim
self.ffn_dim = ffn_dim
self.freq_dim = freq_dim
self.text_dim = text_dim
self.out_dim = out_dim
self.num_heads = num_heads
self.num_layers = num_layers
self.local_attn_size = local_attn_size
self.qk_norm = qk_norm
self.cross_attn_norm = cross_attn_norm
self.eps = eps
# embeddings — nn.Linear inside nn.Sequential on purpose: the official
# checkpoint keys are positional (text_embedding.0/.2, time_embedding.0/.2,
# time_projection.1) and must load verbatim (see ffn comment above).
self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), nn.Linear(dim, dim))
self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
# transformer blocks
self.blocks = nn.ModuleList([
CausalWanAttentionBlock(dim,
ffn_dim,
num_heads,
local_attn_size,
sink_size,
qk_norm,
cross_attn_norm,
eps,
add_control_adapter=add_control_adapter,
cam_method=cam_method,
attn_compress=attn_compress,
layer_idx=layer_idx,
cam_self_attn_layers=cam_self_attn_layers) for layer_idx in range(num_layers)
])
for layer_idx, block in enumerate(self.blocks):
block.self_attn.layer_idx = layer_idx
block.self_attn.num_layers = self.num_layers
# head
self.head = CausalHead(dim, out_dim, patch_size, eps)
# RoPE frequencies
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
d = dim // num_heads
self.freqs = torch.cat(
[rope_params(1024, d - 4 * (d // 6)),
rope_params(1024, 2 * (d // 6)),
rope_params(1024, 2 * (d // 6))],
dim=1)
self.num_attention_heads = num_heads
self.attention_head_dim = dim // num_heads
self.hidden_size = dim
self.in_channels = in_dim
self.out_channels = out_dim
self.num_channels_latents = out_dim
self.init_weights()
self.num_frame_per_block = config.arch_config.num_frames_per_block
self.__post_init__()
def forward(self,
x=None,
t=None,
context=None,
seq_len=None,
y=None,
y_camera=None,
kv_cache=None,
crossattn_cache=None,
current_start=0,
cache_start=0,
cache_update_policy="commit_detached",
hidden_states=None,
encoder_hidden_states=None,
timestep=None,
**kwargs):
"""
Causal inference with KV caching.
See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772).
Args:
x: List of input video tensors [C_in, F, H, W]
t: Timestep tensor [B, L]
context: List of text embeddings [L, C]
seq_len: Maximum sequence length for positional encoding
y: Optional conditional video inputs (I2V mode)
y_camera: Camera parameters dict {'viewmats': ..., 'K': ...}
kv_cache: List of KV cache dicts per transformer block
crossattn_cache: List of cross-attention cache dicts
current_start: Current position in global token sequence
cache_start: Cache start position
cache_update_policy: Cache update strategy ('commit_detached' or 'none')
Returns:
Stacked output tensors [B, C_out, F, H/8, W/8]
"""
if x is None and hidden_states is not None:
x = [sample for sample in hidden_states]
if t is None and timestep is not None:
t = timestep
if context is None and encoder_hidden_states is not None:
if isinstance(encoder_hidden_states, torch.Tensor):
context = [sample for sample in encoder_hidden_states]
else:
context = encoder_hidden_states
if seq_len is None:
if torch.is_tensor(t):
seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel())
elif x is not None:
sample = x[0]
seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (
sample.shape[3] // self.patch_size[2])
if x is None or t is None or context is None or seq_len is None:
raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases")
device = self.patch_embedding.weight.device
if self.freqs.is_meta or self.freqs.device != device:
d = self.dim // self.num_heads
self.freqs = torch.cat([
rope_params(1024, d - 4 * (d // 6)),
rope_params(1024, 2 * (d // 6)),
rope_params(1024, 2 * (d // 6)),
],
dim=1).to(device)
if y is not None:
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]
# patch embedding
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
x = [u.flatten(2).transpose(1, 2) for u in x]
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
assert seq_lens.max() <= seq_len
x = torch.cat(x)
# time embedding
e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
e0 = self.time_projection(e).unflatten(1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
# text embedding
context_lens = None
context = self.text_embedding(
torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]))
# camera parameters
if y_camera is not None and isinstance(y_camera, dict):
cam_viewmats = y_camera['viewmats']
cam_K = y_camera['K']
else:
cam_viewmats = None
cam_K = None
block_kwargs = dict(
e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=self.freqs,
context=context,
context_lens=context_lens,
cam_viewmats=cam_viewmats,
cam_K=cam_K,
cache_update_policy=cache_update_policy,
)
cache_update_infos = []
for block_index, block in enumerate(self.blocks):
block_kwargs.update({
"kv_cache": kv_cache[block_index] if kv_cache is not None else None,
"crossattn_cache": crossattn_cache[block_index] if crossattn_cache is not None else None,
"current_start": current_start,
"cache_start": cache_start,
})
x, block_cache_update_info = block(x, **block_kwargs)
if kv_cache is not None:
cache_update_infos.append((block_index, block_cache_update_info))
# Apply deferred cache updates
if kv_cache is not None and cache_update_infos and cache_update_policy != "none":
self._apply_cache_updates(kv_cache, cache_update_infos)
# head & unpatchify
x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
x = self.unpatchify(x, grid_sizes)
return torch.stack(x)
def _apply_cache_updates(self, kv_cache, cache_update_infos):
"""Apply deferred cache updates collected from all transformer blocks.
For Block-Relativistic RoPE, this stores un-roped K values in the cache.
RoPE is applied dynamically during attention based on each token's current
relative position in the sliding window.
"""
with torch.no_grad():
for block_index, (current_end, local_end_index, update_info) in cache_update_infos:
if update_info is not None:
cache = kv_cache[block_index]
if update_info["action"] == "roll_and_insert":
sink_tokens = update_info["sink_tokens"]
num_rolled_tokens = update_info["num_rolled_tokens"]
num_evicted_tokens = update_info["num_evicted_tokens"]
write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
new_k = update_info["new_k"].detach()
new_v = update_info["new_v"].detach()
cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
if write_end_index > write_start_index and new_k.shape[1] == (write_end_index -
write_start_index):
cache["k"][:, write_start_index:write_end_index] = new_k
cache["v"][:, write_start_index:write_end_index] = new_v
elif update_info["action"] == "direct_insert":
write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
new_k = update_info["new_k"].detach()
new_v = update_info["new_v"].detach()
if write_end_index > write_start_index and new_k.shape[1] == (write_end_index -
write_start_index):
cache["k"][:, write_start_index:write_end_index] = new_k
cache["v"][:, write_start_index:write_end_index] = new_v
is_recompute = False if update_info is None else update_info.get("is_recompute", False)
if not is_recompute:
kv_cache[block_index]["global_end_index"].fill_(current_end)
kv_cache[block_index]["local_end_index"].fill_(local_end_index)
def unpatchify(self, x, grid_sizes):
"""Reconstruct video tensors from patch embeddings."""
c = self.out_dim
out = []
for u, v in zip(x, grid_sizes.tolist(), strict=True):
u = u[:math.prod(v)].view(*v, *self.patch_size, c)
u = torch.einsum('fhwpqrc->cfphqwr', u)
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
out.append(u)
return out
def init_weights(self):
"""Initialize model parameters using Xavier initialization."""
for m in self.modules():
if isinstance(m, (nn.Linear, ReplicatedLinear)):
nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
for m in self.text_embedding.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, std=.02)
for m in self.time_embedding.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, std=.02)
nn.init.zeros_(self.head.head.weight)
EntryClass = DreamXWorldARTransformer3DModel
fastvideo/models/dits/glm_image.py
-702
View File
@@ -1,702 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastvideo.layers.mlp import MLP
from fastvideo.attention import LocalAttention
from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.layers.rotary_embedding import _apply_rotary_emb
from fastvideo.layers.visual_embedding import Timesteps
from fastvideo.models.dits.base import BaseDiT
from fastvideo.platforms import AttentionBackendEnum
class GlmImageLayerKVCache:
def __init__(self):
self.k_cache = None
self.v_cache = None
self.mode: Optional[str] = None
def store(self, k: torch.Tensor, v: torch.Tensor):
# Append along seq (dim=1).
if self.k_cache is None:
self.k_cache = k
self.v_cache = v
else:
self.k_cache = torch.cat([self.k_cache, k], dim=1)
self.v_cache = torch.cat([self.v_cache, v], dim=1)
def get(self):
return self.k_cache, self.v_cache
def clear(self):
self.k_cache = None
self.v_cache = None
self.mode = None
class GlmImageKVCache:
def __init__(self, num_layers: int):
self.num_layers = num_layers
self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
return self.caches[layer_idx]
def set_mode(self, mode: Optional[str]):
if mode is not None and mode not in ["write", "read", "skip"]:
raise ValueError(f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'")
for cache in self.caches:
cache.mode = mode
def clear(self):
for cache in self.caches:
cache.clear()
# =============================================================================
# Timestep and Text Projection
# =============================================================================
class GlmImageTimestepEmbedding(nn.Module):
def __init__(
self,
in_channels: int,
time_embed_dim: int,
act_fn: str = "silu",
out_dim: int = None,
):
super().__init__()
if out_dim is None:
out_dim = time_embed_dim
self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True)
if act_fn == "silu":
self.act = nn.SiLU()
elif act_fn == "gelu":
self.act = nn.GELU(approximate="tanh")
else:
self.act = nn.SiLU()
self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True)
def forward(self, sample: torch.Tensor) -> torch.Tensor:
sample, _ = self.linear_1(sample)
sample = self.act(sample)
sample, _ = self.linear_2(sample)
return sample
class GlmImageTextProjection(nn.Module):
def __init__(
self,
in_features: int,
hidden_size: int,
out_features: int = None,
act_fn: str = "silu",
):
super().__init__()
if out_features is None:
out_features = hidden_size
self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True)
if act_fn == "silu":
self.act_1 = nn.SiLU()
elif act_fn == "gelu_tanh":
self.act_1 = nn.GELU(approximate="tanh")
else:
self.act_1 = nn.SiLU()
self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True)
def forward(self, caption: torch.Tensor) -> torch.Tensor:
hidden_states, _ = self.linear_1(caption)
hidden_states = self.act_1(hidden_states)
hidden_states, _ = self.linear_2(hidden_states)
return hidden_states
class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
def __init__(
self,
embedding_dim: int,
condition_dim: int,
pooled_projection_dim: int,
timesteps_dim: int = 256,
):
super().__init__()
self.time_proj = Timesteps(num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
self.condition_proj = Timesteps(num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = GlmImageTimestepEmbedding(in_channels=timesteps_dim, time_embed_dim=embedding_dim)
self.condition_embedder = GlmImageTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
def forward(
self,
timestep: torch.Tensor,
target_size: torch.Tensor,
crop_coords: torch.Tensor,
hidden_dtype: torch.dtype,
) -> torch.Tensor:
timesteps_proj = self.time_proj(timestep)
crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(crop_coords.size(0), -1)
target_size_proj = self.condition_proj(target_size.flatten()).view(target_size.size(0), -1)
condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))
condition_emb = self.condition_embedder(condition_proj.to(dtype=hidden_dtype))
conditioning = timesteps_emb + condition_emb
return conditioning
# =============================================================================
# Image Projector
# =============================================================================
class GlmImageImageProjector(nn.Module):
def __init__(
self,
in_channels: int = 16,
hidden_size: int = 2560,
patch_size: int = 2,
):
super().__init__()
self.patch_size = patch_size
self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
batch_size, channel, height, width = hidden_states.shape
post_patch_height = height // self.patch_size
post_patch_width = width // self.patch_size
hidden_states = hidden_states.reshape(
batch_size,
channel,
post_patch_height,
self.patch_size,
post_patch_width,
self.patch_size,
)
hidden_states = (hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2))
hidden_states = self.proj(hidden_states)
return hidden_states
# =============================================================================
# AdaLayerNorm
# =============================================================================
class GlmImageAdaLayerNormZero(nn.Module):
def __init__(self, embedding_dim: int, dim: int) -> None:
super().__init__()
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
) -> Tuple[torch.Tensor, ...]:
dtype = hidden_states.dtype
norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
emb, _ = self.linear(temb)
(
shift_msa,
c_shift_msa,
scale_msa,
c_scale_msa,
gate_msa,
c_gate_msa,
shift_mlp,
c_shift_mlp,
scale_mlp,
c_scale_mlp,
gate_mlp,
c_gate_mlp,
) = emb.chunk(12, dim=1)
hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1)
return (
hidden_states,
gate_msa,
shift_mlp,
scale_mlp,
gate_mlp,
encoder_hidden_states,
c_gate_msa,
c_shift_mlp,
c_scale_mlp,
c_gate_mlp,
)
# =============================================================================
# Attention
# =============================================================================
class GlmImageAttention(nn.Module):
def __init__(
self,
query_dim: int,
heads: int,
dim_head: int,
out_dim: int,
bias: bool = True,
qk_norm: str = "layer_norm",
elementwise_affine: bool = False,
eps: float = 1e-5,
supported_attention_backends: set[AttentionBackendEnum] | None = None,
prefix: str = "",
):
super().__init__()
self.heads = out_dim // dim_head if out_dim is not None else heads
self.num_kv_heads = self.heads
self.dim_head = dim_head
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
self.inner_kv_dim = self.inner_dim
self.out_dim = out_dim if out_dim is not None else query_dim
self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
self.to_out = nn.ModuleList([ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)])
if qk_norm is None:
self.norm_q = None
self.norm_k = None
elif qk_norm == "layer_norm":
self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
else:
raise ValueError(f"unknown qk_norm: {qk_norm}")
self.attn = LocalAttention(
num_heads=self.heads,
head_size=dim_head,
num_kv_heads=self.heads,
softmax_scale=None,
causal=False,
supported_attention_backends=supported_attention_backends,
)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
kv_cache: Optional[GlmImageLayerKVCache] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
dtype = encoder_hidden_states.dtype
batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
batch_size, image_seq_length, embed_dim = hidden_states.shape
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
# 1. QKV projections
query, _ = self.to_q(hidden_states)
key, _ = self.to_k(hidden_states)
value, _ = self.to_v(hidden_states)
query = query.unflatten(2, (self.heads, -1))
key = key.unflatten(2, (self.heads, -1))
value = value.unflatten(2, (self.heads, -1))
# 2. QK normalization
if self.norm_q is not None:
query = self.norm_q(query).to(dtype=dtype)
if self.norm_k is not None:
key = self.norm_k(key).to(dtype=dtype)
# 3. Rotational positional embeddings applied to latent stream
if image_rotary_emb is not None:
cos, sin = image_rotary_emb
query[:, text_seq_length:, :, :] = _apply_rotary_emb(query[:, text_seq_length:, :, :],
cos,
sin,
is_neox_style=True)
key[:, text_seq_length:, :, :] = _apply_rotary_emb(key[:, text_seq_length:, :, :],
cos,
sin,
is_neox_style=True)
# 4. KV Cache handling
if kv_cache is not None:
if kv_cache.mode == "write":
kv_cache.store(key, value)
elif kv_cache.mode == "read":
# Prepend cached condition k/v along seq (dim=1).
k_cache, v_cache = kv_cache.get()
key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key
value = (torch.cat([v_cache, value], dim=1) if v_cache is not None else value)
elif kv_cache.mode == "skip":
pass
hidden_states = self.attn(query, key, value)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
# 6. Output projection
hidden_states, _ = self.to_out[0](hidden_states)
encoder_hidden_states, hidden_states = hidden_states.split(
[text_seq_length, hidden_states.size(1) - text_seq_length], dim=1)
return hidden_states, encoder_hidden_states
# =============================================================================
# Transformer Block
# =============================================================================
class GlmImageTransformerBlock(nn.Module):
def __init__(
self,
dim: int = 2560,
num_attention_heads: int = 64,
attention_head_dim: int = 40,
time_embed_dim: int = 512,
supported_attention_backends: set[AttentionBackendEnum] | None = None,
prefix: str = "",
) -> None:
super().__init__()
# 1. Attention
self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
self.attn1 = GlmImageAttention(
query_dim=dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
out_dim=dim,
bias=True,
qk_norm="layer_norm",
elementwise_affine=False,
eps=1e-5,
supported_attention_backends=supported_attention_backends,
prefix=f"{prefix}.attn1",
)
# 2. Feedforward with fused ScaleResidualLayerNorm
self.norm2 = ScaleResidualLayerNormScaleShift(dim, norm_type="layer", eps=1e-5, elementwise_affine=False)
self.norm2_context = ScaleResidualLayerNormScaleShift(dim,
norm_type="layer",
eps=1e-5,
elementwise_affine=False)
self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: Optional[torch.Tensor] = None,
image_rotary_emb: Optional[Union[
Tuple[torch.Tensor, torch.Tensor],
List[Tuple[torch.Tensor, torch.Tensor]],
]] = None,
attention_kwargs: Optional[Dict[str, Any]] = None,
kv_cache: Optional[GlmImageLayerKVCache] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# 1. Timestep conditioning
(
norm_hidden_states,
gate_msa,
shift_mlp,
scale_mlp,
gate_mlp,
norm_encoder_hidden_states,
c_gate_msa,
c_shift_mlp,
c_scale_mlp,
c_gate_mlp,
) = self.norm1(hidden_states, encoder_hidden_states, temb)
# 2. Attention
if attention_kwargs is None:
attention_kwargs = {}
attn_hidden_states, attn_encoder_hidden_states = self.attn1(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
image_rotary_emb=image_rotary_emb,
kv_cache=kv_cache,
**attention_kwargs,
)
# 3. Feedforward (fused residual + norm + scale/shift)
norm_hidden_states, hidden_states = self.norm2(
hidden_states,
attn_hidden_states,
gate_msa.unsqueeze(1),
shift_mlp.unsqueeze(1),
scale_mlp.unsqueeze(1),
)
norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context(
encoder_hidden_states,
attn_encoder_hidden_states,
c_gate_msa.unsqueeze(1),
c_shift_mlp.unsqueeze(1),
c_scale_mlp.unsqueeze(1),
)
ff_output = self.ff(norm_hidden_states)
ff_output_context = self.ff(norm_encoder_hidden_states)
hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
encoder_hidden_states = (encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1))
return hidden_states, encoder_hidden_states
# =============================================================================
# Rotary Positional Embedding
# =============================================================================
class GlmImageRotaryPosEmbed(nn.Module):
def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
super().__init__()
self.dim = dim
self.patch_size = patch_size
self.theta = theta
self._cache_key: tuple | None = None
self._cache_value: tuple[torch.Tensor, torch.Tensor] | None = None
def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
_, _, raw_h, raw_w = hidden_states.shape
height = raw_h // self.patch_size
width = raw_w // self.patch_size
device = hidden_states.device
cache_key = (height, width, device.type, device.index if device.index is not None else -1)
if self._cache_key == cache_key and self._cache_value is not None:
return self._cache_value
dim_h, dim_w = self.dim // 2, self.dim // 2
h_inv_freq = 1.0 / (self.theta
**(torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[:(dim_h // 2)] / dim_h))
w_inv_freq = 1.0 / (self.theta
**(torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[:(dim_w // 2)] / dim_w))
h_seq = torch.arange(height, device=device)
w_seq = torch.arange(width, device=device)
freqs_h = torch.outer(h_seq, h_inv_freq).unsqueeze(1).expand(height, width, -1)
freqs_w = torch.outer(w_seq, w_inv_freq).unsqueeze(0).expand(height, width, -1)
freqs = torch.cat([freqs_h, freqs_w], dim=-1).reshape(height * width, -1)
result = (freqs.cos(), freqs.sin())
self._cache_key = cache_key
self._cache_value = result
return result
# =============================================================================
# Final AdaLayerNorm
# =============================================================================
class GlmImageAdaLayerNormContinuous(nn.Module):
def __init__(
self,
embedding_dim: int,
conditioning_embedding_dim: int,
elementwise_affine: bool = True,
eps: float = 1e-5,
bias: bool = True,
norm_type: str = "layer_norm",
):
super().__init__()
self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
if norm_type == "layer_norm":
self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
elif norm_type == "rms_norm":
self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine)
else:
raise ValueError(f"unknown norm_type {norm_type}")
def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
emb = self.linear(conditioning_embedding.to(x.dtype))
scale, shift = torch.chunk(emb, 2, dim=1)
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
return x
# =============================================================================
# Main Model
# =============================================================================
class GlmImageTransformer2DModel(BaseDiT):
_fsdp_shard_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions
_compile_conditions = GlmImageDiTConfig().arch_config._compile_conditions
_supported_attention_backends = (
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
)
param_names_mapping = GlmImageDiTConfig().arch_config.param_names_mapping
reverse_param_names_mapping = {}
lora_param_names_mapping = {}
def __init__(
self,
config: GlmImageDiTConfig,
hf_config: dict[str, Any],
):
super().__init__(config=config, hf_config=hf_config)
arch_config = config.arch_config
self.in_channels = arch_config.in_channels
self.out_channels = arch_config.out_channels
self.patch_size = arch_config.patch_size
self.num_layers = arch_config.num_layers
self.attention_head_dim = arch_config.attention_head_dim
self.num_attention_heads = arch_config.num_attention_heads
self.text_embed_dim = arch_config.text_embed_dim
self.time_embed_dim = arch_config.time_embed_dim
# GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
# Each of these are sincos embeddings of shape 2 * condition_dim
pooled_projection_dim = 2 * 2 * arch_config.condition_dim
inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
self.hidden_size = inner_dim
self.num_channels_latents = arch_config.out_channels
# 1. RoPE
self.rotary_emb = GlmImageRotaryPosEmbed(arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0)
# 2. Patch & Text-timestep embedding
self.image_projector = GlmImageImageProjector(arch_config.in_channels, inner_dim, arch_config.patch_size)
self.glyph_projector = MLP(
input_dim=arch_config.text_embed_dim,
mlp_hidden_dim=inner_dim,
output_dim=inner_dim,
act_type="gelu",
)
self.prior_token_embedding = nn.Embedding(arch_config.prior_vq_quantizer_codebook_size, inner_dim)
self.prior_projector = MLP(
input_dim=inner_dim,
mlp_hidden_dim=inner_dim,
output_dim=inner_dim,
act_type="silu",
)
self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
embedding_dim=arch_config.time_embed_dim,
condition_dim=arch_config.condition_dim,
pooled_projection_dim=pooled_projection_dim,
timesteps_dim=arch_config.time_embed_dim,
)
# 3. Transformer blocks
self.transformer_blocks = nn.ModuleList([
GlmImageTransformerBlock(
inner_dim,
arch_config.num_attention_heads,
arch_config.attention_head_dim,
arch_config.time_embed_dim,
supported_attention_backends=self._supported_attention_backends,
prefix=f"transformer_blocks.{i}",
) for i in range(arch_config.num_layers)
])
# 4. Output projection
self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, arch_config.time_embed_dim, elementwise_affine=False)
self.proj_out = nn.Linear(
inner_dim,
arch_config.patch_size * arch_config.patch_size * arch_config.out_channels,
bias=True,
)
self.gradient_checkpointing = False
self.__post_init__()
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
prior_token_id: torch.Tensor,
prior_token_drop: torch.Tensor,
timestep: torch.LongTensor,
target_size: torch.Tensor,
crop_coords: torch.Tensor,
attention_kwargs: Optional[Dict[str, Any]] = None,
kv_caches: Optional[GlmImageKVCache] = None,
kv_caches_mode: Optional[str] = None,
freqs_cis: Optional[Union[
Tuple[torch.Tensor, torch.Tensor],
List[Tuple[torch.Tensor, torch.Tensor]],
]] = None,
guidance: torch.Tensor = None,
**kwargs,
) -> torch.Tensor:
if kv_caches is not None:
kv_caches.set_mode(kv_caches_mode)
batch_size, num_channels, height, width = hidden_states.shape
if isinstance(encoder_hidden_states, list):
encoder_hidden_states = encoder_hidden_states[0]
# 1. RoPE
image_rotary_emb = freqs_cis
if image_rotary_emb is None:
image_rotary_emb = self.rotary_emb(hidden_states)
# 2. Patch & Timestep embeddings
p = self.patch_size
post_patch_height = height // p
post_patch_width = width // p
hidden_states = self.image_projector(hidden_states)
encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
prior_embedding = self.prior_token_embedding(prior_token_id)
# Zero dropped priors by multiply: boolean indexing + .any() syncs each step.
keep = (~prior_token_drop).to(device=prior_embedding.device, dtype=prior_embedding.dtype)
while keep.dim() < prior_embedding.dim():
keep = keep.unsqueeze(-1)
prior_embedding = prior_embedding * keep
prior_hidden_states = self.prior_projector(prior_embedding)
hidden_states = hidden_states + prior_hidden_states
temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype)
temb = F.silu(temb)
# 3. Transformer blocks
for idx, block in enumerate(self.transformer_blocks):
hidden_states, encoder_hidden_states = block(
hidden_states,
encoder_hidden_states,
temb,
image_rotary_emb,
attention_kwargs,
kv_cache=kv_caches[idx] if kv_caches is not None else None,
)
# 4. Output norm & projection
hidden_states = self.norm_out(hidden_states, temb)
hidden_states = self.proj_out(hidden_states)
# 5. Unpatchify
hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p)
output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
return output.float()
EntryClass = GlmImageTransformer2DModel
fastvideo/models/dits/hunyuangamecraft.py
-348
View File
@@ -1,348 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
HunyuanGameCraft Transformer model for FastVideo.
Ported from official Hunyuan-GameCraft-1.0 implementation.
"""
from typing import Any, List, Optional
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from fastvideo.configs.models.dits.hunyuangamecraft import (
HunyuanGameCraftArchConfig,
HunyuanGameCraftConfig,
)
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.layers.mlp import MLP
from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
from fastvideo.layers.visual_embedding import ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify
from fastvideo.models.dits.base import BaseDiT
from fastvideo.models.dits.hunyuanvideo import (
MMDoubleStreamBlock,
MMSingleStreamBlock,
SingleTokenRefiner,
)
class GameCraftFinalLayer(nn.Module):
"""
GameCraft-specific FinalLayer with correct shift/scale order.
The official GameCraft implementation uses shift, scale order (not scale, shift).
This differs from the HunyuanVideo FinalLayer.
"""
def __init__(self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "") -> None:
super().__init__()
self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype)
output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels
self.linear = ReplicatedLinear(hidden_size,
output_dim,
bias=True,
params_dtype=dtype,
prefix=f"{prefix}.linear")
self.adaLN_modulation = ModulateProjection(hidden_size,
factor=2,
act_layer="silu",
dtype=dtype,
prefix=f"{prefix}.adaLN_modulation")
def forward(self, x, c):
# GameCraft uses shift, scale order (verified against official implementation)
shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
x, _ = self.linear(x)
return x
class CameraNet(nn.Module):
"""
Camera state encoding network - ported from official GameCraft.
Processes camera parameters (Plücker coordinates) into feature embeddings.
"""
def __init__(
self,
in_channels: int = 6,
downscale_coef: int = 8,
out_channels: int = 16,
patch_size: List[int] = [1, 2, 2],
hidden_size: int = 3072,
dtype: Optional[torch.dtype] = None,
prefix: str = "",
):
super().__init__()
_ = prefix # Unused
start_channels = in_channels * (downscale_coef**2)
input_channels = [start_channels, start_channels // 2, start_channels // 4]
self.input_channels = input_channels
self.unshuffle = nn.PixelUnshuffle(downscale_coef)
self.encode_first = nn.Sequential(
nn.Conv2d(input_channels[0], input_channels[1], kernel_size=1, stride=1, padding=0),
nn.GroupNorm(2, input_channels[1]),
nn.ReLU(),
)
self._initialize_weights(self.encode_first)
self.encode_second = nn.Sequential(
nn.Conv2d(input_channels[1], input_channels[2], kernel_size=1, stride=1, padding=0),
nn.GroupNorm(2, input_channels[2]),
nn.ReLU(),
)
self._initialize_weights(self.encode_second)
self.final_proj = nn.Conv2d(input_channels[2], out_channels, kernel_size=1)
self._zeros_init_linear(self.final_proj)
self.scale = nn.Parameter(torch.ones(1))
self.camera_in = PatchEmbed(
patch_size=patch_size,
in_chans=out_channels,
embed_dim=hidden_size,
)
def _zeros_init_linear(self, linear):
if hasattr(linear, "weight"):
nn.init.zeros_(linear.weight)
if hasattr(linear, "bias") and linear.bias is not None:
nn.init.zeros_(linear.bias)
def _initialize_weights(self, block):
for m in block:
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.in_channels
nn.init.normal_(m.weight, mean=0.0, std=np.sqrt(2.0 / n))
if m.bias is not None:
nn.init.zeros_(m.bias)
def compress_time(self, x: torch.Tensor, num_frames: int) -> torch.Tensor:
x = rearrange(x, '(b f) c h w -> b f c h w', f=num_frames)
batch_size, frames, channels, height, width = x.shape
x = rearrange(x, 'b f c h w -> (b h w) c f')
if x.shape[-1] == 66 or x.shape[-1] == 34:
x_len = x.shape[-1]
x_clip1 = x[..., :x_len // 2]
x_clip1_first = x_clip1[..., 0].unsqueeze(-1)
x_clip1_rest = F.avg_pool1d(x_clip1[..., 1:], kernel_size=2, stride=2)
x_clip2 = x[..., x_len // 2:]
x_clip2_first = x_clip2[..., 0].unsqueeze(-1)
x_clip2_rest = F.avg_pool1d(x_clip2[..., 1:], kernel_size=2, stride=2)
x = torch.cat([x_clip1_first, x_clip1_rest, x_clip2_first, x_clip2_rest], dim=-1)
elif x.shape[-1] % 2 == 1:
x_first = x[..., 0]
x_rest = x[..., 1:]
if x_rest.shape[-1] > 0:
x_rest = F.avg_pool1d(x_rest, kernel_size=2, stride=2)
x = torch.cat([x_first[..., None], x_rest], dim=-1)
else:
x = F.avg_pool1d(x, kernel_size=2, stride=2)
x = rearrange(x, '(b h w) c f -> (b f) c h w', b=batch_size, h=height, w=width)
return x
def forward(self, camera_states: torch.Tensor) -> torch.Tensor:
batch_size, num_frames, channels, height, width = camera_states.shape
camera_states = rearrange(camera_states, 'b f c h w -> (b f) c h w')
camera_states = self.unshuffle(camera_states)
camera_states = self.encode_first(camera_states)
camera_states = self.compress_time(camera_states, num_frames=num_frames)
num_frames = camera_states.shape[0] // batch_size
camera_states = self.encode_second(camera_states)
camera_states = self.compress_time(camera_states, num_frames=num_frames)
camera_states = self.final_proj(camera_states)
camera_states = rearrange(camera_states, "(b f) c h w -> b c f h w", b=batch_size)
camera_states = self.camera_in(camera_states)
return camera_states * self.scale
class HunyuanGameCraftTransformer3DModel(BaseDiT):
"""
HunyuanGameCraft Transformer - ported from official implementation.
"""
_fsdp_shard_conditions = HunyuanGameCraftArchConfig()._fsdp_shard_conditions
_compile_conditions = HunyuanGameCraftArchConfig()._compile_conditions
_supported_attention_backends = HunyuanGameCraftArchConfig()._supported_attention_backends
param_names_mapping = HunyuanGameCraftConfig().param_names_mapping
reverse_param_names_mapping = HunyuanGameCraftConfig().reverse_param_names_mapping
def __init__(self, config: HunyuanGameCraftConfig, hf_config: dict[str, Any]):
super().__init__(config=config, hf_config=hf_config)
arch = config.arch_config
if isinstance(arch.patch_size, (list, tuple)):
self.patch_size = list(arch.patch_size)
else:
self.patch_size = [arch.patch_size_t, arch.patch_size, arch.patch_size]
self.in_channels = arch.in_channels
self.out_channels = arch.out_channels
self.unpatchify_channels = self.out_channels
self.num_channels_latents = self.out_channels # Alias for latent_preparation stage
self.hidden_size = arch.hidden_size
self.num_heads = arch.num_attention_heads
self.num_attention_heads = arch.num_attention_heads # Alias for compatibility
self.guidance_embeds = arch.guidance_embeds
self.rope_dim_list = list(arch.rope_axes_dim)
self.rope_theta = arch.rope_theta
self.text_states_dim = arch.text_embed_dim
self.text_states_dim_2 = arch.pooled_projection_dim
self.dtype = arch.dtype
pe_dim = self.hidden_size // self.num_heads
if sum(self.rope_dim_list) != pe_dim:
raise ValueError(f"rope_axes_dim sum {sum(self.rope_dim_list)} != {pe_dim}")
factory_kwargs = {'dtype': self.dtype}
self.img_in = PatchEmbed(
patch_size=self.patch_size,
in_chans=self.in_channels,
embed_dim=self.hidden_size,
**factory_kwargs,
)
self.txt_in = SingleTokenRefiner(
self.text_states_dim,
self.hidden_size,
self.num_heads,
depth=arch.num_refiner_layers,
**factory_kwargs,
)
self.time_in = TimestepEmbedder(self.hidden_size, **factory_kwargs)
self.vector_in = MLP(
self.text_states_dim_2,
self.hidden_size,
self.hidden_size,
act_type="silu",
**factory_kwargs,
)
self.guidance_in = (TimestepEmbedder(self.hidden_size, **factory_kwargs) if self.guidance_embeds else None)
self.double_blocks = nn.ModuleList([
MMDoubleStreamBlock(
hidden_size=self.hidden_size,
num_attention_heads=self.num_heads,
mlp_ratio=arch.mlp_ratio,
supported_attention_backends=self._supported_attention_backends,
**factory_kwargs,
) for _ in range(arch.num_layers)
])
self.single_blocks = nn.ModuleList([
MMSingleStreamBlock(
hidden_size=self.hidden_size,
num_attention_heads=self.num_heads,
mlp_ratio=arch.mlp_ratio,
supported_attention_backends=self._supported_attention_backends,
**factory_kwargs,
) for _ in range(arch.num_single_layers)
])
self.final_layer = GameCraftFinalLayer(
self.hidden_size,
self.patch_size,
self.out_channels,
**factory_kwargs,
)
self.camera_net = CameraNet(
in_channels=arch.camera_in_channels,
out_channels=16,
downscale_coef=arch.camera_downscale_coef,
patch_size=self.patch_size,
hidden_size=self.hidden_size,
)
def forward(
self,
x: torch.Tensor,
encoder_hidden_states: List[torch.Tensor],
timestep: torch.Tensor,
camera_states: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[List[torch.Tensor]] = None,
guidance: Optional[torch.Tensor] = None,
return_dict: bool = False,
) -> torch.Tensor:
img = x
_, _, ot, oh, ow = x.shape
tt = ot // self.patch_size[0]
th = oh // self.patch_size[1]
tw = ow // self.patch_size[2]
text_states = encoder_hidden_states[0]
text_states_2 = encoder_hidden_states[1] if len(encoder_hidden_states) > 1 else None
text_mask = encoder_attention_mask[0] if encoder_attention_mask else None
vec = self.time_in(timestep)
if text_states_2 is not None:
vec = vec + self.vector_in(text_states_2)
if self.guidance_in is not None and guidance is not None:
vec = vec + self.guidance_in(guidance)
img = self.img_in(img)
if camera_states is not None:
latent_len = ot
if latent_len == 18:
camera_latents = torch.cat(
[self.camera_net(torch.zeros_like(camera_states)),
self.camera_net(camera_states)], dim=1)
elif latent_len == 9:
camera_latents = self.camera_net(camera_states)
elif latent_len == 10:
camera_latents = torch.cat(
[self.camera_net(torch.zeros_like(camera_states[:, 0:4, :, :, :])),
self.camera_net(camera_states)],
dim=1)
else:
camera_latents = self.camera_net(camera_states)
img = img + camera_latents
txt = self.txt_in(text_states, timestep)
txt_seq_len = txt.shape[1]
freqs_cos, freqs_sin = get_rotary_pos_embed(
(tt, th, tw),
self.hidden_size,
self.num_heads,
self.rope_dim_list,
self.rope_theta,
)
freqs_cos = freqs_cos.to(device=img.device, dtype=img.dtype)
freqs_sin = freqs_sin.to(device=img.device, dtype=img.dtype)
freqs_cis = (freqs_cos, freqs_sin)
for block in self.double_blocks:
img, txt = block(img, txt, vec, freqs_cis)
x = torch.cat([img, txt], dim=1)
for block in self.single_blocks:
x = block(x, vec, txt_seq_len, freqs_cis)
img = x[:, :-txt_seq_len, ...]
img = self.final_layer(img, vec)
img = unpatchify(img, tt, th, tw, self.patch_size, self.out_channels)
return img
fastvideo/models/dits/longcat.py
-1141
View File
File diff suppressed because it is too large Load Diff
fastvideo/models/dits/zimage.py
-563
View File
@@ -1,563 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""FastVideo-native Z-Image transformer.
Z-Image attends over padded variable-length image/text streams and requires a
key-padding mask. FastVideo's distributed attention wrappers do not expose that
mask contract yet, so this implementation uses torch SDPA and is SP=1 only.
"""
from __future__ import annotations
import math
from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.models.dits.base import BaseDiT
from fastvideo.platforms import AttentionBackendEnum
def _linear(layer: ReplicatedLinear, x: torch.Tensor) -> torch.Tensor:
return layer(x)[0]
def _prepare_attention_mask(attention_mask: torch.Tensor | None, dtype: torch.dtype) -> torch.Tensor | None:
if attention_mask is None:
return None
if attention_mask.ndim == 2:
attention_mask = attention_mask[:, None, None, :]
if attention_mask.dtype == torch.bool:
additive_mask = torch.zeros_like(attention_mask, dtype=dtype)
additive_mask.masked_fill_(~attention_mask, float("-inf"))
return additive_mask
return attention_mask
class TimestepEmbedder(nn.Module):
def __init__(
self,
out_size: int,
mid_size: int,
frequency_embedding_size: int,
max_period: int,
) -> None:
super().__init__()
self.mlp = nn.ModuleList([
ReplicatedLinear(frequency_embedding_size, mid_size, bias=True),
nn.SiLU(),
ReplicatedLinear(mid_size, out_size, bias=True),
])
self.frequency_embedding_size = frequency_embedding_size
self.max_period = max_period
@staticmethod
def timestep_embedding(t: torch.Tensor, dim: int, max_period: int) -> torch.Tensor:
with torch.amp.autocast("cuda", enabled=False):
half = dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
args = t[:, None].float() * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2:
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
return embedding
def forward(self, t: torch.Tensor) -> torch.Tensor:
t_freq = self.timestep_embedding(t, self.frequency_embedding_size, self.max_period)
weight_dtype = self.mlp[0].weight.dtype
if weight_dtype.is_floating_point:
t_freq = t_freq.to(weight_dtype)
return _linear(self.mlp[2], self.mlp[1](_linear(self.mlp[0], t_freq)))
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-5) -> None:
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return output * self.weight
class FeedForward(nn.Module):
def __init__(self, dim: int, hidden_dim: int) -> None:
super().__init__()
self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False)
self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False)
self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _linear(self.w2, F.silu(_linear(self.w1, x)) * _linear(self.w3, x))
def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
with torch.amp.autocast("cuda", enabled=False):
x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
x_out = torch.view_as_real(x * freqs_cis.unsqueeze(2)).flatten(3)
return x_out.type_as(x_in)
class ZImageAttention(nn.Module):
def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool = True, eps: float = 1e-5) -> None:
super().__init__()
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.head_dim = dim // n_heads
self.to_q = ReplicatedLinear(dim, n_heads * self.head_dim, bias=False)
self.to_k = ReplicatedLinear(dim, n_kv_heads * self.head_dim, bias=False)
self.to_v = ReplicatedLinear(dim, n_kv_heads * self.head_dim, bias=False)
self.to_out = nn.ModuleList([ReplicatedLinear(n_heads * self.head_dim, dim, bias=False)])
self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else None
self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else None
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
freqs_cis: torch.Tensor | None = None,
) -> torch.Tensor:
query = _linear(self.to_q, hidden_states).unflatten(-1, (self.n_heads, -1))
key = _linear(self.to_k, hidden_states).unflatten(-1, (self.n_kv_heads, -1))
value = _linear(self.to_v, hidden_states).unflatten(-1, (self.n_kv_heads, -1))
if self.norm_q is not None:
query = self.norm_q(query)
if self.norm_k is not None:
key = self.norm_k(key)
if freqs_cis is not None:
query = apply_rotary_emb(query, freqs_cis)
key = apply_rotary_emb(key, freqs_cis)
mask = _prepare_attention_mask(attention_mask, query.dtype)
hidden_states = F.scaled_dot_product_attention(
query.transpose(1, 2),
key.transpose(1, 2),
value.transpose(1, 2),
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
).transpose(1, 2).contiguous()
return _linear(self.to_out[0], hidden_states.flatten(2, 3).to(query.dtype))
class ZImageTransformerBlock(nn.Module):
def __init__(
self,
layer_id: int,
dim: int,
n_heads: int,
n_kv_heads: int,
norm_eps: float,
qk_norm: bool,
adaln_embed_dim: int,
modulation: bool = True,
) -> None:
super().__init__()
self.dim = dim
self.head_dim = dim // n_heads
self.layer_id = layer_id
self.modulation = modulation
self.attention = ZImageAttention(dim, n_heads, n_kv_heads, qk_norm, norm_eps)
self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
if modulation:
self.adaLN_modulation = nn.ModuleList([ReplicatedLinear(min(dim, adaln_embed_dim), 4 * dim, bias=True)])
def forward(
self,
x: torch.Tensor,
attn_mask: torch.Tensor,
freqs_cis: torch.Tensor,
adaln_input: torch.Tensor | None = None,
) -> torch.Tensor:
if self.modulation:
assert adaln_input is not None
scale_msa, gate_msa, scale_mlp, gate_mlp = _linear(self.adaLN_modulation[0],
adaln_input).unsqueeze(1).chunk(4, dim=2)
gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
attn_out = self.attention(
self.attention_norm1(x) * scale_msa,
attention_mask=attn_mask,
freqs_cis=freqs_cis,
)
x = x + gate_msa * self.attention_norm2(attn_out)
x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
else:
attn_out = self.attention(
self.attention_norm1(x),
attention_mask=attn_mask,
freqs_cis=freqs_cis,
)
x = x + self.attention_norm2(attn_out)
x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
return x
class FinalLayer(nn.Module):
def __init__(self, hidden_size: int, out_channels: int, adaln_embed_dim: int) -> None:
super().__init__()
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = ReplicatedLinear(hidden_size, out_channels, bias=True)
self.adaLN_modulation = nn.ModuleList([
nn.SiLU(),
ReplicatedLinear(min(hidden_size, adaln_embed_dim), hidden_size, bias=True),
])
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
scale = 1.0 + _linear(self.adaLN_modulation[1], self.adaLN_modulation[0](c))
return _linear(self.linear, self.norm_final(x) * scale.unsqueeze(1))
class RopeEmbedder:
def __init__(self, theta: float, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...]) -> None:
if len(axes_dims) != len(axes_lens):
raise ValueError("RoPE axes require matching dimensions and lengths")
self.theta = theta
self.axes_dims = axes_dims
self.axes_lens = axes_lens
self.freqs_cis: list[torch.Tensor] | None = None
@staticmethod
def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float) -> list[torch.Tensor]:
with torch.device("cpu"):
freqs_cis = []
for axis_dim, axis_end in zip(dim, end):
freqs = 1.0 / (theta**(torch.arange(0, axis_dim, 2, dtype=torch.float64) / axis_dim))
timestep = torch.arange(axis_end, dtype=torch.float64)
angles = torch.outer(timestep, freqs).float()
freqs_cis.append(torch.polar(torch.ones_like(angles), angles).to(torch.complex64))
return freqs_cis
def __call__(self, ids: torch.Tensor) -> torch.Tensor:
if ids.ndim != 2 or ids.shape[-1] != len(self.axes_dims):
raise ValueError("RoPE ids must have shape [sequence, number_of_axes]")
if self.freqs_cis is None:
self.freqs_cis = [
freqs.to(ids.device) for freqs in self.precompute_freqs_cis(self.axes_dims, self.axes_lens, self.theta)
]
elif self.freqs_cis[0].device != ids.device:
self.freqs_cis = [freqs.to(ids.device) for freqs in self.freqs_cis]
return torch.cat([self.freqs_cis[i][ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
class ZImageTransformer2DModel(BaseDiT):
_default_config = ZImageDiTConfig()
_fsdp_shard_conditions = _default_config.arch_config._fsdp_shard_conditions
_compile_conditions = _default_config.arch_config._compile_conditions
_supported_attention_backends = (AttentionBackendEnum.TORCH_SDPA, )
param_names_mapping = _default_config.arch_config.param_names_mapping
reverse_param_names_mapping = _default_config.arch_config.reverse_param_names_mapping
def __init__(self, config: ZImageDiTConfig, hf_config: dict[str, Any]) -> None:
super().__init__(config=config, hf_config=hf_config)
arch = config.arch_config
self.in_channels = arch.in_channels
self.out_channels = arch.in_channels
self.all_patch_size = tuple(arch.all_patch_size)
self.all_f_patch_size = tuple(arch.all_f_patch_size)
self.dim = arch.dim
self.n_heads = arch.n_heads
self.rope_theta = arch.rope_theta
self.t_scale = arch.t_scale
self.seq_multi_of = arch.seq_multi_of
self.hidden_size = arch.dim
self.num_attention_heads = arch.n_heads
self.num_channels_latents = arch.in_channels
self.all_x_embedder = nn.ModuleDict({
f"{patch_size}-{f_patch_size}":
ReplicatedLinear(f_patch_size * patch_size * patch_size * arch.in_channels, arch.dim, bias=True)
for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size)
})
self.all_final_layer = nn.ModuleDict({
f"{patch_size}-{f_patch_size}":
FinalLayer(
arch.dim,
patch_size * patch_size * f_patch_size * self.out_channels,
arch.adaln_embed_dim,
)
for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size)
})
block_kwargs = {
"dim": arch.dim,
"n_heads": arch.n_heads,
"n_kv_heads": arch.n_kv_heads,
"norm_eps": arch.norm_eps,
"qk_norm": arch.qk_norm,
"adaln_embed_dim": arch.adaln_embed_dim,
}
self.noise_refiner = nn.ModuleList([
ZImageTransformerBlock(1000 + layer_id, modulation=True, **block_kwargs)
for layer_id in range(arch.n_refiner_layers)
])
self.context_refiner = nn.ModuleList([
ZImageTransformerBlock(layer_id, modulation=False, **block_kwargs)
for layer_id in range(arch.n_refiner_layers)
])
self.t_embedder = TimestepEmbedder(
min(arch.dim, arch.adaln_embed_dim),
mid_size=arch.timestep_mid_size,
frequency_embedding_size=arch.frequency_embedding_size,
max_period=arch.max_period,
)
self.cap_embedder = nn.ModuleList([
RMSNorm(arch.cap_feat_dim, eps=arch.norm_eps),
ReplicatedLinear(arch.cap_feat_dim, arch.dim, bias=True),
])
self.x_pad_token = nn.Parameter(torch.empty((1, arch.dim)))
self.cap_pad_token = nn.Parameter(torch.empty((1, arch.dim)))
self.layers = nn.ModuleList(
[ZImageTransformerBlock(layer_id, modulation=True, **block_kwargs) for layer_id in range(arch.n_layers)])
self.axes_dims = tuple(arch.axes_dims)
self.axes_lens = tuple(arch.axes_lens)
self.rope_embedder = RopeEmbedder(arch.rope_theta, self.axes_dims, self.axes_lens)
self.__post_init__()
def unpatchify(
self,
x: list[torch.Tensor],
size: list[tuple[int, int, int]],
patch_size: int,
f_patch_size: int,
) -> list[torch.Tensor]:
patch_height = patch_width = patch_size
patch_frames = f_patch_size
if len(x) != len(size):
raise ValueError("output batch and original sizes must have equal length")
for i, (frames, height, width) in enumerate(size):
original_length = (frames // patch_frames) * (height // patch_height) * (width // patch_width)
x[i] = (x[i][:original_length].view(
frames // patch_frames,
height // patch_height,
width // patch_width,
patch_frames,
patch_height,
patch_width,
self.out_channels,
).permute(6, 0, 3, 1, 4, 2, 5).reshape(self.out_channels, frames, height, width))
return x
@staticmethod
def create_coordinate_grid(
size: tuple[int, int, int],
start: tuple[int, int, int] | None = None,
device: torch.device | None = None,
) -> torch.Tensor:
start = start or (0, ) * len(size)
axes = [
torch.arange(axis_start, axis_start + span, dtype=torch.int32, device=device)
for axis_start, span in zip(start, size)
]
return torch.stack(torch.meshgrid(axes, indexing="ij"), dim=-1)
def patchify_and_embed(
self,
all_image: list[torch.Tensor],
all_cap_feats: list[torch.Tensor],
patch_size: int,
f_patch_size: int,
) -> tuple[
list[torch.Tensor],
list[torch.Tensor],
list[tuple[int, int, int]],
list[torch.Tensor],
list[torch.Tensor],
list[torch.Tensor],
list[torch.Tensor],
]:
patch_height = patch_width = patch_size
patch_frames = f_patch_size
device = all_image[0].device
image_out = []
image_sizes = []
image_pos_ids = []
image_pad_masks = []
cap_pos_ids = []
cap_pad_masks = []
cap_feats_out = []
for image, cap_feat in zip(all_image, all_cap_feats):
cap_length = len(cap_feat)
cap_padding = (-cap_length) % self.seq_multi_of
cap_pos_ids.append(
self.create_coordinate_grid(
(cap_length + cap_padding, 1, 1),
start=(1, 0, 0),
device=device,
).flatten(0, 2))
cap_pad_masks.append(
torch.cat([
torch.zeros(cap_length, dtype=torch.bool, device=device),
torch.ones(cap_padding, dtype=torch.bool, device=device),
]) if cap_padding else torch.zeros(cap_length, dtype=torch.bool, device=device))
cap_feats_out.append(
torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding, 1)]) if cap_padding else cap_feat)
channels, frames, height, width = image.size()
image_sizes.append((frames, height, width))
frame_tokens, height_tokens, width_tokens = (
frames // patch_frames,
height // patch_height,
width // patch_width,
)
image = image.view(
channels,
frame_tokens,
patch_frames,
height_tokens,
patch_height,
width_tokens,
patch_width,
)
image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
frame_tokens * height_tokens * width_tokens,
patch_frames * patch_height * patch_width * channels,
)
image_length = len(image)
image_padding = (-image_length) % self.seq_multi_of
original_pos_ids = self.create_coordinate_grid(
(frame_tokens, height_tokens, width_tokens),
start=(cap_length + cap_padding + 1, 0, 0),
device=device,
).flatten(0, 2)
if image_padding:
padding_pos_ids = self.create_coordinate_grid((1, 1, 1),
device=device).flatten(0, 2).repeat(image_padding, 1)
image_pos_ids.append(torch.cat([original_pos_ids, padding_pos_ids]))
else:
image_pos_ids.append(original_pos_ids)
image_pad_masks.append(
torch.cat([
torch.zeros(image_length, dtype=torch.bool, device=device),
torch.ones(image_padding, dtype=torch.bool, device=device),
]) if image_padding else torch.zeros(image_length, dtype=torch.bool, device=device))
image_out.append(torch.cat([image, image[-1:].repeat(image_padding, 1)]) if image_padding else image)
return (
image_out,
cap_feats_out,
image_sizes,
image_pos_ids,
cap_pos_ids,
image_pad_masks,
cap_pad_masks,
)
@staticmethod
def _attention_mask(lengths: list[int], device: torch.device) -> torch.Tensor:
mask = torch.zeros((len(lengths), max(lengths)), dtype=torch.bool, device=device)
for i, length in enumerate(lengths):
mask[i, :length] = True
return mask
def forward(
self,
hidden_states: torch.Tensor | list[torch.Tensor],
encoder_hidden_states: torch.Tensor | list[torch.Tensor],
timestep: torch.LongTensor,
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
guidance=None,
patch_size: int = 2,
f_patch_size: int = 1,
**kwargs,
) -> tuple[list[torch.Tensor], dict]:
del encoder_hidden_states_image, guidance, kwargs
if model_parallel_is_initialized() and get_sp_world_size() != 1:
raise NotImplementedError("Z-Image masked SDPA does not support sequence parallelism; run with sp_size=1")
if patch_size not in self.all_patch_size or f_patch_size not in self.all_f_patch_size:
raise ValueError(f"unsupported patch sizes: spatial={patch_size}, temporal={f_patch_size}")
if isinstance(hidden_states, torch.Tensor):
hidden_states = list(hidden_states.unbind(0))
if isinstance(encoder_hidden_states, torch.Tensor):
encoder_hidden_states = list(encoder_hidden_states.unbind(0))
device = hidden_states[0].device
timestep_embedding = self.t_embedder(timestep * self.t_scale)
(
hidden_states,
encoder_hidden_states,
image_sizes,
image_pos_ids,
cap_pos_ids,
image_inner_pad_masks,
cap_inner_pad_masks,
) = self.patchify_and_embed(hidden_states, encoder_hidden_states, patch_size, f_patch_size)
image_lengths = [len(item) for item in hidden_states]
if not all(length % self.seq_multi_of == 0 for length in image_lengths):
raise ValueError("padded image sequence lengths must be aligned")
hidden_states = torch.cat(hidden_states)
hidden_states = _linear(self.all_x_embedder[f"{patch_size}-{f_patch_size}"], hidden_states)
adaln_input = timestep_embedding.type_as(hidden_states)
hidden_states[torch.cat(image_inner_pad_masks)] = self.x_pad_token
hidden_states = list(hidden_states.split(image_lengths))
image_freqs_cis = list(
self.rope_embedder(torch.cat(image_pos_ids)).split([len(item) for item in image_pos_ids]))
hidden_states = pad_sequence(hidden_states, batch_first=True, padding_value=0.0)
image_freqs_cis = pad_sequence(image_freqs_cis, batch_first=True, padding_value=0.0)
image_freqs_cis = image_freqs_cis[:, :hidden_states.shape[1]]
image_attn_mask = self._attention_mask(image_lengths, device)
for layer in self.noise_refiner:
hidden_states = layer(hidden_states, image_attn_mask, image_freqs_cis, adaln_input)
cap_lengths = [len(item) for item in encoder_hidden_states]
if not all(length % self.seq_multi_of == 0 for length in cap_lengths):
raise ValueError("padded caption sequence lengths must be aligned")
encoder_hidden_states = torch.cat(encoder_hidden_states)
encoder_hidden_states = _linear(self.cap_embedder[1], self.cap_embedder[0](encoder_hidden_states))
encoder_hidden_states[torch.cat(cap_inner_pad_masks)] = self.cap_pad_token
encoder_hidden_states = list(encoder_hidden_states.split(cap_lengths))
cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids)).split([len(item) for item in cap_pos_ids]))
encoder_hidden_states = pad_sequence(encoder_hidden_states, batch_first=True, padding_value=0.0)
cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0)
cap_freqs_cis = cap_freqs_cis[:, :encoder_hidden_states.shape[1]]
cap_attn_mask = self._attention_mask(cap_lengths, device)
for layer in self.context_refiner:
encoder_hidden_states = layer(encoder_hidden_states, cap_attn_mask, cap_freqs_cis)
unified = []
unified_freqs_cis = []
for i, (image_length, cap_length) in enumerate(zip(image_lengths, cap_lengths)):
unified.append(torch.cat([hidden_states[i][:image_length], encoder_hidden_states[i][:cap_length]]))
unified_freqs_cis.append(torch.cat([image_freqs_cis[i][:image_length], cap_freqs_cis[i][:cap_length]]))
unified_lengths = [image_length + cap_length for image_length, cap_length in zip(image_lengths, cap_lengths)]
unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0)
unified_attn_mask = self._attention_mask(unified_lengths, device)
for layer in self.layers:
unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input)
unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
outputs = self.unpatchify(list(unified.unbind(0)), image_sizes, patch_size, f_patch_size)
return outputs, {}
EntryClass = ZImageTransformer2DModel
fastvideo/models/encoders/glm_image_ar_loader.py
-63
View File
@@ -1,63 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Sole HF-import boundary for GLM-Image's AR encoder (lazy-wrapper exception E001)."""
from __future__ import annotations
from typing import Any
import torch
import torch.nn as nn
from fastvideo.logger import init_logger
logger = init_logger(__name__)
class GlmImageARLoader(nn.Module):
def __init__(self,
model_path: str,
processor_path: str | None = None,
*,
torch_dtype: torch.dtype = torch.bfloat16,
trust_remote_code: bool = True) -> None:
super().__init__()
from transformers import (AutoProcessor, GlmImageForConditionalGeneration)
logger.info("Loading GLM-Image AR encoder from %s", model_path)
self._model = GlmImageForConditionalGeneration.from_pretrained(
model_path,
torch_dtype=torch_dtype,
trust_remote_code=trust_remote_code,
)
if processor_path is not None:
logger.info("Loading GLM-Image processor from %s", processor_path)
self.processor = AutoProcessor.from_pretrained(processor_path, trust_remote_code=trust_remote_code)
else:
self.processor = None
@torch.no_grad()
def generate(self, *args: Any, **kwargs: Any) -> torch.Tensor:
return self._model.generate(*args, **kwargs)
@torch.no_grad()
def get_image_features(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor) -> Any:
return self._model.get_image_features(pixel_values, image_grid_thw)
@torch.no_grad()
def get_image_tokens(self, image_embeds: torch.Tensor, image_grid_thw: torch.Tensor) -> torch.Tensor:
return self._model.get_image_tokens(image_embeds, image_grid_thw)
@property
def config(self): # type: ignore[no-untyped-def]
return self._model.config
@property
def generation_config(self): # type: ignore[no-untyped-def]
return self._model.generation_config
def to(self, *args, **kwargs): # type: ignore[override]
self._model = self._model.to(*args, **kwargs)
return super().to(*args, **kwargs)
def eval(self): # type: ignore[override]
self._model = self._model.eval()
return super().eval()
fastvideo/models/loader/component_loader.py
-18
View File
@@ -110,7 +110,6 @@ class ComponentLoader(ABC):
"image_encoder": (ImageEncoderLoader, "transformers"),
"image_encoder_2": (ImageEncoderLoader, "transformers"),
"image_encoder_3": (ImageEncoderLoader, "transformers"),
"vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
"processor": (ProcessorLoader, "transformers"),
"upsampler": (UpsamplerLoader, "diffusers"),
"upsampler_2": (UpsamplerLoader, "diffusers"),
@@ -594,23 +593,6 @@ class ImageEncoderLoader(TextEncoderLoader):
)
class VisionLanguageEncoderLoader(ComponentLoader):
"""Loader for vision-language autoregressive encoders."""
def load(self, model_path: str, fastvideo_args: FastVideoArgs):
from fastvideo.distributed.parallel_state import get_local_torch_device
from fastvideo.models.encoders.glm_image_ar_loader import (GlmImageARLoader)
logger.info("Loading vision-language encoder from %s", model_path)
target_device = get_local_torch_device()
loader = GlmImageARLoader(
model_path,
torch_dtype=torch.bfloat16,
trust_remote_code=fastvideo_args.trust_remote_code,
).to(target_device).eval()
return loader
class ProcessorLoader(ComponentLoader):
"""Loader for HF processors that pair with vision-language encoders."""
fastvideo/models/registry.py
-15
View File
@@ -25,18 +25,12 @@ logger = init_logger(__name__)
_TEXT_TO_VIDEO_DIT_MODELS = {
"MMAudioTransformer": ("dits", "mmaudio", "MMAudioTransformer"),
"HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
"HunyuanGameCraftTransformer3DModel": ("dits", "hunyuangamecraft", "HunyuanGameCraftTransformer3DModel"),
"HunyuanVideo15Transformer3DModel": ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
"HYWorldTransformer3DModel": ("dits", "hyworld", "HYWorldTransformer3DModel"),
"WanTransformer3DModel": ("wan", "transformer", "WanTransformer3DModel"),
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
"DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
"CausalWanTransformer3DModel": ("wan", "causal_transformer", "CausalWanTransformer3DModel"),
"CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
"Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
"LongCatVideoTransformer3DModel":
("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1)
"LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2)
"LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
"SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"),
"LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
@@ -54,8 +48,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
_IMAGE_TO_VIDEO_DIT_MODELS = {
# "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
"WanTransformer3DModel": ("wan", "transformer", "WanTransformer3DModel"),
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
"DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
"CausalWanTransformer3DModel": ("wan", "causal_transformer", "CausalWanTransformer3DModel"),
"LingBotWorld2CausalFastTransformer3DModel": (
"dits",
@@ -72,8 +64,6 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
# Text-to-image DiT models (2D image generation)
_TEXT_TO_IMAGE_DIT_MODELS = {
"GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"),
"ZImageTransformer2DModel": ("dits", "zimage", "ZImageTransformer2DModel"),
}
_TEXT_ENCODER_MODELS = {
@@ -88,9 +78,6 @@ _TEXT_ENCODER_MODELS = {
"Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"),
"Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"),
"Qwen2_5_VLForConditionalGeneration": ("encoders", "reason1", "Reason1TextEncoder"),
# Z-Image-Turbo's text_encoder/config.json declares architecture
# "Qwen3Model"; route it to the shared Qwen3 encoder (added for Flux2 Klein).
"Qwen3Model": ("encoders", "qwen3", "Qwen3ForCausalLM"),
"LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
"Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
"Mistral3ForConditionalGeneration": ("encoders", "mistral3", "Mistral3ForConditionalGeneration"),
@@ -111,7 +98,6 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
_VAE_MODELS = {
"AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
"AutoencoderKLCausal3D": ("vaes", "gamecraftvae", "GameCraftVAE"),
"AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"),
"AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
"AutoencoderKLWan": ("wan", "vae", "AutoencoderKLWan"),
@@ -142,7 +128,6 @@ _SCHEDULERS = {
"FlowUniPCMultistepScheduler": ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"),
"SelfForcingFlowMatchScheduler":
("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"),
"RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"),
}
_UPSAMPLERS = {
fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py
+2 -2
View File
@@ -96,7 +96,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
sigma_data (`float`, *optional*):
The sigma data value for scaling.
use_reference_discrete_timesteps (`bool`, defaults to False):
Some reference schedulers (e.g. Z-Image) construct the timestep
Some reference schedulers (e.g. MMAudio) construct the timestep
schedule by linspacing `num_inference_steps + 1` points from
`t_max` to `t_min` and dropping the terminal point. Default
(`False`) preserves the original `np.linspace(t_max, t_min,
@@ -340,7 +340,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
t_max = self._sigma_to_t(self.sigma_max)
t_min = self._sigma_to_t(self.sigma_min)
if self.config.use_reference_discrete_timesteps:
# Some reference schedulers (for example Z-Image) build a
# Some reference schedulers (for example MMAudio) build a
# float64 num_steps+1 linspace and drop the terminal point.
timesteps_array = np.linspace(
t_max,
fastvideo/models/schedulers/scheduling_rcm.py
-323
View File
@@ -1,323 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# rCM (recurrent Consistency Model) Scheduler for TurboDiffusion support
#
# This scheduler implements the rCM sampling method from TurboDiffusion,
# enabling 1-4 step video generation with distilled checkpoints.
#
# Reference:
# TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times
# https://arxiv.org/pdf/2512.16093
import math
from dataclasses import dataclass
from typing import Any
import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput
from fastvideo.logger import init_logger
from fastvideo.models.schedulers.base import BaseScheduler
logger = init_logger(__name__)
@dataclass
class RCMSchedulerOutput(BaseOutput):
"""
Output class for the RCM scheduler's `step` function output.
Args:
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, ...)`):
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used
as next model input in the denoising loop.
"""
prev_sample: torch.FloatTensor
class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
"""
rCM (recurrent Consistency Model) scheduler for TurboDiffusion.
This scheduler implements the rCM sampling method which enables 1-4 step
video generation using distilled checkpoints. It uses:
1. TrigFlow → RectifiedFlow timestep conversion
2. SDE sampling formula: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise
Args:
num_train_timesteps (`int`, defaults to 1000):
The number of diffusion steps used to train the model.
sigma_max (`float`, defaults to 80.0):
The initial sigma value for rCM sampling. Controls the noise level
at the start of sampling.
mid_timesteps (`list[float]`, *optional*):
Custom intermediate timesteps. If None, uses optimized defaults
[1.5, 1.4, 1.0] for better visual quality.
"""
_compatibles: list[Any] = []
order = 1
@register_to_config
def __init__(
self,
num_train_timesteps: int = 1000,
sigma_max: float = 80.0,
mid_timesteps: list[float] | None = None,
):
# Default mid timesteps optimized for visual quality
if mid_timesteps is None:
mid_timesteps = [1.5, 1.4, 1.0]
self._mid_timesteps = mid_timesteps
self.num_train_timesteps = num_train_timesteps
self.sigma_max = sigma_max
# Initialize with default timesteps (will be set properly via set_timesteps)
self.timesteps = torch.tensor([1.0, 0.0], dtype=torch.float64)
self.sigmas = self.timesteps.clone()
self._step_index: int | None = None
self._begin_index: int | None = None
BaseScheduler.__init__(self)
@property
def step_index(self) -> int | None:
"""
The index counter for current timestep. Increases by 1 after each scheduler step.
"""
return self._step_index
@property
def begin_index(self) -> int | None:
"""
The index for the first timestep.
"""
return self._begin_index
@property
def init_noise_sigma(self) -> float:
"""
Initial noise sigma for scaling latents.
In rCM, initial noise is scaled by the first sigma value:
x_0 = noise * sigmas[0]
This property is used by LatentPreparationStage to scale initial latents.
"""
return float(self.sigmas[0])
def set_begin_index(self, begin_index: int = 0) -> None:
"""
Sets the begin index for the scheduler.
Args:
begin_index (`int`):
The begin index for the scheduler.
"""
self._begin_index = begin_index
def set_shift(self, shift: float) -> None:
"""
rCM doesn't use shift parameter, but required by BaseScheduler.
"""
pass
def set_timesteps(
self,
num_inference_steps: int,
device: str | torch.device | None = None,
sigma_max: float | None = None,
) -> None:
"""
Sets the discrete timesteps used for the rCM sampling process.
The timesteps are computed using TrigFlow → RectifiedFlow conversion:
1. Start with atan(sigma_max) and intermediate values
2. Convert via: t = sin(t) / (cos(t) + sin(t))
Args:
num_inference_steps (`int`):
The number of diffusion steps (1-4 for rCM).
device (`str` or `torch.device`, *optional*):
The device to move timesteps to.
sigma_max (`float`, *optional*):
Override the initial sigma value.
"""
if num_inference_steps < 1 or num_inference_steps > 4:
logger.warning("rCM is optimized for 1-4 steps, got %d steps. "
"Performance may be suboptimal.", num_inference_steps)
self.num_inference_steps = num_inference_steps
if sigma_max is not None:
self.sigma_max = sigma_max
# Build timestep schedule
mid_t = self._mid_timesteps[:num_inference_steps - 1]
# TrigFlow timesteps: [atan(sigma_max), mid_t..., 0]
t_steps = torch.tensor(
[math.atan(self.sigma_max), *mid_t, 0],
dtype=torch.float64,
device=device,
)
# Convert TrigFlow → RectifiedFlow: t = sin(t) / (cos(t) + sin(t))
t_steps = torch.sin(t_steps) / (torch.cos(t_steps) + torch.sin(t_steps))
# Store raw sigmas for use in step() formula
self.sigmas = t_steps.clone()
# Scale timesteps by 1000 for model input (as per TurboDiffusion)
self.timesteps = t_steps * 1000
self._step_index = None
self._begin_index = None
logger.debug("rCM timesteps (scaled): %s", self.timesteps.tolist())
logger.debug("rCM sigmas (raw): %s", self.sigmas.tolist())
def _init_step_index(self, timestep: torch.FloatTensor | None = None) -> None:
"""Initialize step index at the beginning of sampling."""
if self._begin_index is None:
self._step_index = 0
else:
self._step_index = self._begin_index
def scale_model_input(
self,
sample: torch.Tensor,
timestep: int | None = None,
) -> torch.Tensor:
"""
rCM doesn't scale model input, returns sample as-is.
"""
return sample
def scale_noise(
self,
sample: torch.FloatTensor,
timestep: torch.FloatTensor | None = None,
noise: torch.FloatTensor | None = None,
) -> torch.FloatTensor:
"""
Scale initial noise for rCM sampling.
In rCM, initial noise is scaled by the first timestep (raw sigma):
x_0 = noise * t_steps[0]
Args:
sample: Not used (for API compatibility)
timestep: Not used (for API compatibility)
noise: The noise tensor to scale
Returns:
Scaled noise tensor ready for sampling
"""
if noise is None:
raise ValueError("noise must be provided for rCM scale_noise")
# Use raw sigma (not scaled timestep) for initial noise scaling
t_initial = self.sigmas[0]
return noise.to(torch.float64) * t_initial
def step(
self,
model_output: torch.FloatTensor,
timestep: int | torch.Tensor,
sample: torch.FloatTensor,
generator: torch.Generator | None = None,
return_dict: bool = True,
) -> RCMSchedulerOutput | tuple[torch.FloatTensor, ...]:
"""
Predict the sample from the previous timestep using rCM update rule.
The rCM update formula is:
x_{t+1} = (1 - t_next) * (x_t - t_cur * v_pred) + t_next * noise
Args:
model_output (`torch.FloatTensor`):
The velocity prediction from the model (v_pred).
timestep (`int` or `torch.Tensor`):
The current timestep index (not the actual timestep value).
For rCM, this should be the index into self.timesteps.
sample (`torch.FloatTensor`):
Current sample x_t.
generator (`torch.Generator`, *optional*):
Random number generator for noise.
return_dict (`bool`):
Whether to return RCMSchedulerOutput or tuple.
Returns:
`RCMSchedulerOutput` or `tuple`:
The denoised sample for the next step.
"""
if self._step_index is None:
self._init_step_index()
assert self._step_index is not None
# Get current and next sigma values (raw, unscaled) for rCM formula
# Note: self.timesteps is scaled by 1000 for model input,
# but we need raw values for the step formula
t_cur = self.sigmas[self._step_index]
# On the final step, t_next should be 0 (fully denoised)
if self._step_index + 1 < len(self.sigmas):
t_next = self.sigmas[self._step_index + 1]
else:
t_next = torch.tensor(0.0, device=sample.device, dtype=torch.float64)
# Ensure we're working in float64 for precision
sample = sample.to(torch.float64)
model_output = model_output.to(torch.float64)
# rCM update: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise
x_denoised = sample - t_cur * model_output
# Generate noise for SDE sampling
if isinstance(generator, list):
generator = generator[0]
noise = torch.randn(
sample.shape,
dtype=torch.float32,
device="cpu",
generator=generator,
).to(sample.device).to(torch.float64)
prev_sample = (1 - t_next) * x_denoised + t_next * noise
# Increment step counter
self._step_index += 1
# Cast back to model output dtype
prev_sample = prev_sample.to(model_output.dtype)
if not return_dict:
return (prev_sample, )
return RCMSchedulerOutput(prev_sample=prev_sample)
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.IntTensor,
) -> torch.Tensor:
"""
Add noise to samples (forward diffusion process).
Not typically used for rCM inference, but provided for API compatibility.
"""
raise NotImplementedError("add_noise is not implemented for RCMScheduler. "
"Use scale_noise for initializing the sampling process.")
def __len__(self) -> int:
return self.config.num_train_timesteps
# Entry point for model registry
EntryClass = RCMScheduler
fastvideo/models/vaes/gamecraftvae.py
-450
View File
@@ -1,450 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
GameCraft VAE - ported from official Hunyuan-GameCraft-1.0/hymm_sp/vae/.
Matches the official AutoencoderKLCausal3D structure exactly for weight loading.
"""
from dataclasses import dataclass
from typing import Optional, Tuple
import numpy as np
import torch
import torch.nn as nn
from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
from fastvideo.models.vaes.common import DiagonalGaussianDistribution
from fastvideo.models.vaes.gamecraftvae_blocks import (
CausalConv3d,
DownEncoderBlockCausal3D,
UpDecoderBlockCausal3D,
UNetMidBlockCausal3D,
)
@dataclass
class AutoencoderKLOutput:
"""Matches official AutoencoderKLOutput interface."""
latent_dist: DiagonalGaussianDistribution
@dataclass
class DecoderOutput:
"""Matches official DecoderOutput interface."""
sample: torch.Tensor
class EncoderCausal3D(nn.Module):
"""Encoder - ported from official vae.py. Structure matches for weight loading."""
def __init__(
self,
in_channels: int = 3,
out_channels: int = 16,
down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ),
block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
double_z: bool = True,
mid_block_add_attention: bool = True,
time_compression_ratio: int = 4,
spatial_compression_ratio: int = 8,
disable_causal: bool = False,
mid_block_causal_attn: bool = False,
):
super().__init__()
self.layers_per_block = layers_per_block
self.conv_in = CausalConv3d(in_channels,
block_out_channels[0],
kernel_size=3,
stride=1,
disable_causal=disable_causal)
self.down_blocks = nn.ModuleList([])
output_channel = block_out_channels[0]
for i, _ in enumerate(down_block_types):
input_channel = output_channel
output_channel = block_out_channels[i]
is_final_block = i == len(block_out_channels) - 1
num_spatial = int(np.log2(spatial_compression_ratio))
num_time = int(np.log2(time_compression_ratio))
if time_compression_ratio == 4:
add_spatial = bool(i < num_spatial)
add_time = bool(i >= (len(block_out_channels) - 1 - num_time) and not is_final_block)
elif time_compression_ratio == 8:
add_spatial = bool(i < num_spatial)
add_time = bool(i < num_time)
else:
raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}")
downsample_stride_HW = (2, 2) if add_spatial else (1, 1)
downsample_stride_T = (2, ) if add_time else (1, )
downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
down_block = DownEncoderBlockCausal3D(
in_channels=input_channel,
out_channels=output_channel,
num_layers=layers_per_block,
resnet_eps=1e-6,
resnet_act_fn=act_fn,
resnet_groups=norm_num_groups,
add_downsample=bool(add_spatial or add_time),
downsample_stride=downsample_stride,
downsample_padding=0,
disable_causal=disable_causal,
)
self.down_blocks.append(down_block)
self.mid_block = UNetMidBlockCausal3D(
in_channels=block_out_channels[-1],
temb_channels=None,
num_layers=1,
resnet_eps=1e-6,
resnet_act_fn=act_fn,
resnet_groups=norm_num_groups,
add_attention=mid_block_add_attention,
attention_head_dim=block_out_channels[-1],
disable_causal=disable_causal,
causal_attention=mid_block_causal_attn,
)
self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
self.conv_act = nn.SiLU()
conv_out_channels = 2 * out_channels if double_z else out_channels
self.conv_out = CausalConv3d(
block_out_channels[-1],
conv_out_channels,
kernel_size=3,
disable_causal=disable_causal,
)
def forward(self, sample: torch.Tensor) -> torch.Tensor:
sample = self.conv_in(sample)
for down_block in self.down_blocks:
sample = down_block(sample, scale=1.0)
sample = self.mid_block(sample, temb=None)
sample = self.conv_norm_out(sample)
sample = self.conv_act(sample)
sample = self.conv_out(sample)
return sample
class DecoderCausal3D(nn.Module):
"""Decoder - ported from official vae.py. Structure matches for weight loading."""
def __init__(
self,
in_channels: int = 16,
out_channels: int = 3,
up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ),
block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
mid_block_add_attention: bool = True,
time_compression_ratio: int = 4,
spatial_compression_ratio: int = 8,
disable_causal: bool = False,
mid_block_causal_attn: bool = False,
):
super().__init__()
self.layers_per_block = layers_per_block
self.conv_in = CausalConv3d(
in_channels,
block_out_channels[-1],
kernel_size=3,
stride=1,
disable_causal=disable_causal,
)
self.mid_block = UNetMidBlockCausal3D(
in_channels=block_out_channels[-1],
temb_channels=None,
num_layers=1,
resnet_eps=1e-6,
resnet_act_fn=act_fn,
resnet_groups=norm_num_groups,
add_attention=mid_block_add_attention,
attention_head_dim=block_out_channels[-1],
disable_causal=disable_causal,
causal_attention=mid_block_causal_attn,
)
self.up_blocks = nn.ModuleList([])
reversed_channels = list(reversed(block_out_channels))
output_channel = reversed_channels[0]
for i, _ in enumerate(up_block_types):
prev_output_channel = output_channel
output_channel = reversed_channels[i]
is_final_block = i == len(block_out_channels) - 1
num_spatial = int(np.log2(spatial_compression_ratio))
num_time = int(np.log2(time_compression_ratio))
if time_compression_ratio == 4:
add_spatial = bool(i < num_spatial)
add_time = bool(i >= len(block_out_channels) - 1 - num_time and not is_final_block)
elif time_compression_ratio == 8:
add_spatial = bool(i >= len(block_out_channels) - num_spatial)
add_time = bool(i >= len(block_out_channels) - num_time)
else:
raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}")
upsample_HW = (2, 2) if add_spatial else (1, 1)
upsample_T = (2, ) if add_time else (1, )
upsample_scale_factor = tuple(upsample_T + upsample_HW)
up_block = UpDecoderBlockCausal3D(
in_channels=prev_output_channel,
out_channels=output_channel,
num_layers=layers_per_block + 1,
resnet_eps=1e-6,
resnet_act_fn=act_fn,
resnet_groups=norm_num_groups,
add_upsample=bool(add_spatial or add_time),
upsample_scale_factor=upsample_scale_factor,
temb_channels=None,
disable_causal=disable_causal,
)
self.up_blocks.append(up_block)
output_channel = output_channel
self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
self.conv_act = nn.SiLU()
self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3, disable_causal=disable_causal)
def forward(
self,
sample: torch.Tensor,
latent_embeds: Optional[torch.Tensor] = None,
) -> torch.Tensor:
sample = self.conv_in(sample)
sample = self.mid_block(sample, temb=latent_embeds)
for up_block in self.up_blocks:
sample = up_block(sample, temb=latent_embeds, scale=1.0)
sample = self.conv_norm_out(sample)
sample = self.conv_act(sample)
sample = self.conv_out(sample)
return sample
class GameCraftVAE(nn.Module):
"""
GameCraft VAE - ported from official AutoencoderKLCausal3D.
Structure matches exactly for loading official weights.
"""
def __init__(self, config: GameCraftVAEConfig):
super().__init__()
self.config = config
arch = config.arch_config
time_ratio = getattr(arch, "time_compression_ratio", arch.temporal_compression_ratio)
self.encoder = EncoderCausal3D(
in_channels=arch.in_channels,
out_channels=arch.latent_channels,
down_block_types=tuple(arch.down_block_types),
block_out_channels=tuple(arch.block_out_channels),
layers_per_block=arch.layers_per_block,
norm_num_groups=arch.norm_num_groups,
act_fn=arch.act_fn,
double_z=True,
time_compression_ratio=time_ratio,
spatial_compression_ratio=arch.spatial_compression_ratio,
disable_causal=getattr(arch, "disable_causal_conv", False),
mid_block_add_attention=arch.mid_block_add_attention,
mid_block_causal_attn=getattr(arch, "mid_block_causal_attn", False),
)
self.decoder = DecoderCausal3D(
in_channels=arch.latent_channels,
out_channels=arch.out_channels,
up_block_types=tuple(arch.up_block_types),
block_out_channels=tuple(arch.block_out_channels),
layers_per_block=arch.layers_per_block,
norm_num_groups=arch.norm_num_groups,
act_fn=arch.act_fn,
time_compression_ratio=time_ratio,
spatial_compression_ratio=arch.spatial_compression_ratio,
disable_causal=getattr(arch, "disable_causal_conv", False),
mid_block_add_attention=arch.mid_block_add_attention,
mid_block_causal_attn=getattr(arch, "mid_block_causal_attn", False),
)
self.quant_conv = nn.Conv3d(2 * arch.latent_channels, 2 * arch.latent_channels, kernel_size=1)
self.post_quant_conv = nn.Conv3d(arch.latent_channels, arch.latent_channels, kernel_size=1)
# Scaling factor for latent normalization (required for decoding stage)
self.scaling_factor = arch.scaling_factor
# Tiling support - matches official GameCraft VAE settings
self._tiling_enabled = False
self.tile_overlap_factor = 0.25
# Temporal tiling params (for >64 output frames)
self.tile_sample_min_tsize = 64 # Minimum sample temporal size (video frames)
self.tile_latent_min_tsize = 16 # = 64 // 4 (time_compression_ratio)
# Spatial tiling params - use small tiles to reduce memory
self.tile_sample_min_size = 256 # Minimum spatial tile size in pixel space
self.tile_latent_min_size = 32 # = 256 // 8 (spatial_compression_ratio)
def encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
"""Encode to latent distribution."""
h = self.encoder(x)
moments = self.quant_conv(h)
posterior = DiagonalGaussianDistribution(moments)
return AutoencoderKLOutput(latent_dist=posterior)
def decode(self, z: torch.Tensor) -> torch.Tensor:
"""Decode from latents.
Args:
z: Latent tensor [B, C, T, H, W]
Returns:
Decoded tensor [B, C, T_out, H_out, W_out]
"""
# Use tiled decode for memory efficiency when enabled
if self._tiling_enabled:
# Check if temporal tiling needed (>64 output frames)
if z.shape[2] > self.tile_latent_min_tsize:
return self._temporal_tiled_decode(z)
# Check if spatial tiling needed (large H or W)
if z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size:
return self._spatial_tiled_decode(z)
z = self.post_quant_conv(z)
dec = self.decoder(z, latent_embeds=None)
return dec
def _temporal_tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
"""Decode latents in temporal tiles with overlapping and blending.
Based on official GameCraft temporal_tiled_decode implementation.
Only used when T > tile_latent_min_tsize (16).
"""
B, C, T, H, W = z.shape
# Use the pre-configured tiling parameters
overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))
blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)
t_limit = self.tile_sample_min_tsize - blend_extent
row = []
for i in range(0, T, overlap_size):
tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :]
tile = self.post_quant_conv(tile)
decoded = self.decoder(tile, latent_embeds=None)
if i > 0:
decoded = decoded[:, :, 1:, :, :] # Skip first frame for non-first tiles
row.append(decoded)
# Blend overlapping regions
result_row = []
for i, tile in enumerate(row):
if i > 0:
tile = self._blend_t(row[i - 1], tile, blend_extent)
result_row.append(tile[:, :, :t_limit, :, :])
else:
result_row.append(tile[:, :, :t_limit + 1, :, :])
return torch.cat(result_row, dim=2)
def _spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
"""Decode latents in spatial tiles with overlapping and blending.
Based on official GameCraft spatial_tiled_decode implementation.
"""
overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
row_limit = self.tile_sample_min_size - blend_extent
# Split z into overlapping tiles and decode them separately
rows = []
for i in range(0, z.shape[-2], overlap_size):
row = []
for j in range(0, z.shape[-1], overlap_size):
tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size]
tile = self.post_quant_conv(tile)
decoded = self.decoder(tile, latent_embeds=None)
row.append(decoded)
rows.append(row)
# Blend overlapping regions
result_rows = []
for i, row in enumerate(rows):
result_row = []
for j, tile in enumerate(row):
# Blend with above tile and left tile
if i > 0:
tile = self._blend_v(rows[i - 1][j], tile, blend_extent)
if j > 0:
tile = self._blend_h(row[j - 1], tile, blend_extent)
result_row.append(tile[:, :, :, :row_limit, :row_limit])
result_rows.append(torch.cat(result_row, dim=-1))
return torch.cat(result_rows, dim=-2)
def _blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
"""Blend two tensors along temporal dimension."""
blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
if blend_extent == 0:
return b
a_region = a[..., -blend_extent:, :, :]
b_region = b[..., :blend_extent, :, :]
weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent
weights = weights.view(1, 1, blend_extent, 1, 1)
blended = a_region * (1 - weights) + b_region * weights
b[..., :blend_extent, :, :] = blended
return b
def _blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
"""Blend two tensors along vertical (height) dimension."""
blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
if blend_extent == 0:
return b
a_region = a[..., -blend_extent:, :]
b_region = b[..., :blend_extent, :]
weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent
weights = weights.view(1, 1, 1, blend_extent, 1)
blended = a_region * (1 - weights) + b_region * weights
b[..., :blend_extent, :] = blended
return b
def _blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
"""Blend two tensors along horizontal (width) dimension."""
blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
if blend_extent == 0:
return b
a_region = a[..., -blend_extent:]
b_region = b[..., :blend_extent]
weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent
weights = weights.view(1, 1, 1, 1, blend_extent)
blended = a_region * (1 - weights) + b_region * weights
b[..., :blend_extent] = blended
return b
def enable_tiling(self) -> None:
"""Enable tiling for large inputs."""
self._tiling_enabled = True
def disable_tiling(self) -> None:
"""Disable tiling."""
self._tiling_enabled = False
EntryClass = GameCraftVAE
fastvideo/models/vaes/gamecraftvae_blocks.py
-420
View File
@@ -1,420 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""
GameCraft VAE building blocks - ported from official Hunyuan-GameCraft-1.0/hymm_sp/vae/unet_causal_3d_blocks.py.
Matches the official structure exactly for weight loading.
"""
from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: Optional[int] = None):
seq_len = n_frame * n_hw
mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
for i in range(seq_len):
i_frame = i // n_hw
mask[i, :(i_frame + 1) * n_hw] = 0
if batch_size is not None:
mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
return mask
class CausalConv3d(nn.Module):
"""Causal 3D convolution - matches official structure (has .conv)."""
def __init__(
self,
chan_in: int,
chan_out: int,
kernel_size: Union[int, Tuple[int, int, int]] = 3,
stride: Union[int, Tuple[int, int, int]] = 1,
dilation: Union[int, Tuple[int, int, int]] = 1,
pad_mode: str = "replicate",
disable_causal: bool = False,
**kwargs,
):
super().__init__()
self.pad_mode = pad_mode
k = kernel_size if isinstance(kernel_size, int) else kernel_size[0]
if disable_causal:
padding = (k // 2, k // 2, k // 2, k // 2, k // 2, k // 2)
else:
padding = (k // 2, k // 2, k // 2, k // 2, k - 1, 0)
self.time_causal_padding = padding
self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
return self.conv(x)
class DownsampleCausal3D(nn.Module):
"""Causal 3D downsampling - matches official (has .conv)."""
def __init__(
self,
channels: int,
out_channels: Optional[int] = None,
padding: int = 1,
stride: Union[int, Tuple[int, int, int]] = 2,
kernel_size: int = 3,
bias: bool = True,
disable_causal: bool = False,
):
super().__init__()
self.out_channels = out_channels or channels
self.conv = CausalConv3d(
channels,
self.out_channels,
kernel_size=kernel_size,
stride=stride,
padding=0,
disable_causal=disable_causal,
bias=bias,
)
def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
return self.conv(hidden_states)
class UpsampleCausal3D(nn.Module):
"""Causal 3D upsampling - matches official (has .conv when use_conv=True)."""
def __init__(
self,
channels: int,
out_channels: Optional[int] = None,
kernel_size: int = 3,
upsample_factor: Tuple[int, int, int] = (2, 2, 2),
disable_causal: bool = False,
bias: bool = True,
):
super().__init__()
self.out_channels = out_channels or channels
self.upsample_factor = upsample_factor
self.disable_causal = disable_causal
self.conv = CausalConv3d(
channels,
self.out_channels,
kernel_size=kernel_size,
stride=1,
disable_causal=disable_causal,
bias=bias,
)
def forward(
self,
hidden_states: torch.Tensor,
output_size: Optional[int] = None,
scale: float = 1.0,
) -> torch.Tensor:
B, C, T, H, W = hidden_states.shape
dtype = hidden_states.dtype
if dtype == torch.bfloat16:
hidden_states = hidden_states.to(torch.float32)
if not self.disable_causal and T > 1:
first_h, other_h = hidden_states.split((1, T - 1), dim=2)
other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest")
first_h = F.interpolate(first_h.squeeze(2), scale_factor=self.upsample_factor[1:],
mode="nearest").unsqueeze(2)
hidden_states = torch.cat((first_h, other_h), dim=2)
else:
hidden_states = F.interpolate(hidden_states, scale_factor=self.upsample_factor, mode="nearest")
if dtype == torch.bfloat16:
hidden_states = hidden_states.to(dtype)
return self.conv(hidden_states)
class GameCraftVAEAttention(nn.Module):
"""Attention block matching official diffusers Attention structure (group_norm, to_q, to_k, to_v, to_out)."""
def __init__(
self,
in_channels: int,
heads: int,
dim_head: int,
eps: float = 1e-6,
norm_num_groups: Optional[int] = 32,
bias: bool = True,
):
super().__init__()
self.heads = heads
self.dim_head = dim_head
inner_dim = heads * dim_head
self.group_norm = nn.GroupNorm(norm_num_groups or in_channels, in_channels, eps=eps)
self.to_q = nn.Linear(in_channels, inner_dim, bias=bias)
self.to_k = nn.Linear(in_channels, inner_dim, bias=bias)
self.to_v = nn.Linear(in_channels, inner_dim, bias=bias)
self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, bias=bias))
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
residual = hidden_states
batch_size, seq_len, _ = hidden_states.shape
hidden_states = self.group_norm(hidden_states.permute(0, 2, 1)).permute(0, 2, 1)
q = self.to_q(hidden_states)
k = self.to_k(hidden_states)
v = self.to_v(hidden_states)
q = q.view(batch_size, seq_len, self.heads, self.dim_head).transpose(1, 2)
k = k.view(batch_size, seq_len, self.heads, self.dim_head).transpose(1, 2)
v = v.view(batch_size, seq_len, self.heads, self.dim_head).transpose(1, 2)
scale = self.dim_head**-0.5
attn = torch.matmul(q, k.transpose(-2, -1)) * scale
if attention_mask is not None:
attn = attn + attention_mask
# Official uses upcast_softmax=True: compute softmax in fp32 for numerical stability
attn = F.softmax(attn.float(), dim=-1).to(q.dtype)
hidden_states = torch.matmul(attn, v)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, seq_len, -1)
hidden_states = self.to_out(hidden_states) + residual
return hidden_states
class ResnetBlockCausal3D(nn.Module):
"""ResNet block - matches official structure (conv1.conv, conv2.conv, norm1, norm2)."""
def __init__(
self,
in_channels: int,
out_channels: Optional[int] = None,
temb_channels: Optional[int] = None,
eps: float = 1e-6,
groups: int = 32,
dropout: float = 0.0,
non_linearity: str = "swish",
disable_causal: bool = False,
):
super().__init__()
out_channels = out_channels or in_channels
self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps)
self.conv1 = CausalConv3d(in_channels, out_channels, 3, 1, disable_causal=disable_causal)
self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps)
self.conv2 = CausalConv3d(out_channels, out_channels, 3, 1, disable_causal=disable_causal)
self.dropout = nn.Dropout(dropout)
self.conv_shortcut = (CausalConv3d(in_channels, out_channels, 1, 1, disable_causal=disable_causal)
if in_channels != out_channels else None)
self.nonlinearity = getattr(F, non_linearity, F.silu)
def forward(
self,
x: torch.Tensor,
temb: Optional[torch.Tensor] = None,
scale: float = 1.0,
) -> torch.Tensor:
h = self.norm1(x)
h = self.nonlinearity(h)
h = self.conv1(h)
h = self.norm2(h)
h = self.nonlinearity(h)
h = self.dropout(h)
h = self.conv2(h)
if self.conv_shortcut is not None:
x = self.conv_shortcut(x)
return (x + h) / 1.0
class UNetMidBlockCausal3D(nn.Module):
"""Mid block with resnets and optional attention - matches official structure."""
def __init__(
self,
in_channels: int,
temb_channels: Optional[int] = None,
num_layers: int = 1,
resnet_eps: float = 1e-6,
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
add_attention: bool = True,
attention_head_dim: int = 1,
disable_causal: bool = False,
causal_attention: bool = False,
):
super().__init__()
self.add_attention = add_attention
self.causal_attention = causal_attention
self.resnets = nn.ModuleList()
self.attentions = nn.ModuleList()
self.resnets.append(
ResnetBlockCausal3D(
in_channels=in_channels,
out_channels=in_channels,
temb_channels=temb_channels,
eps=resnet_eps,
groups=resnet_groups,
dropout=0.0,
non_linearity=resnet_act_fn,
disable_causal=disable_causal,
))
for _ in range(num_layers):
if add_attention:
self.attentions.append(
GameCraftVAEAttention(
in_channels=in_channels,
heads=in_channels // attention_head_dim,
dim_head=attention_head_dim,
eps=resnet_eps,
norm_num_groups=resnet_groups,
bias=True,
))
else:
self.attentions.append(None)
self.resnets.append(
ResnetBlockCausal3D(
in_channels=in_channels,
out_channels=in_channels,
temb_channels=temb_channels,
eps=resnet_eps,
groups=resnet_groups,
dropout=0.0,
non_linearity=resnet_act_fn,
disable_causal=disable_causal,
))
def forward(
self,
hidden_states: torch.Tensor,
temb: Optional[torch.Tensor] = None,
) -> torch.Tensor:
hidden_states = self.resnets[0](hidden_states, temb)
for attn, resnet in zip(self.attentions, self.resnets[1:]):
if attn is not None:
B, C, T, H, W = hidden_states.shape
hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c")
if self.causal_attention:
mask = prepare_causal_attention_mask(T,
H * W,
hidden_states.dtype,
hidden_states.device,
batch_size=B)
else:
mask = None
hidden_states = attn(hidden_states, attention_mask=mask)
hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W)
hidden_states = resnet(hidden_states, temb)
return hidden_states
class DownEncoderBlockCausal3D(nn.Module):
"""Encoder down block - matches official structure."""
def __init__(
self,
in_channels: int,
out_channels: int,
num_layers: int = 2,
resnet_eps: float = 1e-6,
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
add_downsample: bool = True,
downsample_stride: Union[int, Tuple[int, int, int]] = 2,
downsample_padding: int = 0,
disable_causal: bool = False,
):
super().__init__()
self.resnets = nn.ModuleList()
for i in range(num_layers):
inc = in_channels if i == 0 else out_channels
self.resnets.append(
ResnetBlockCausal3D(
in_channels=inc,
out_channels=out_channels,
temb_channels=None,
eps=resnet_eps,
groups=resnet_groups,
dropout=0.0,
non_linearity=resnet_act_fn,
disable_causal=disable_causal,
))
self.downsamplers = None
if add_downsample:
self.downsamplers = nn.ModuleList([
DownsampleCausal3D(
out_channels,
out_channels=out_channels,
padding=downsample_padding,
stride=downsample_stride,
disable_causal=disable_causal,
)
])
def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
for resnet in self.resnets:
hidden_states = resnet(hidden_states, temb=None, scale=scale)
if self.downsamplers is not None:
for ds in self.downsamplers:
hidden_states = ds(hidden_states, scale)
return hidden_states
class UpDecoderBlockCausal3D(nn.Module):
"""Decoder up block - matches official structure."""
def __init__(
self,
in_channels: int,
out_channels: int,
num_layers: int = 3,
resnet_eps: float = 1e-6,
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
add_upsample: bool = True,
upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2),
temb_channels: Optional[int] = None,
disable_causal: bool = False,
):
super().__init__()
self.resnets = nn.ModuleList()
for i in range(num_layers):
inc = in_channels if i == 0 else out_channels
self.resnets.append(
ResnetBlockCausal3D(
in_channels=inc,
out_channels=out_channels,
temb_channels=temb_channels,
eps=resnet_eps,
groups=resnet_groups,
dropout=0.0,
non_linearity=resnet_act_fn,
disable_causal=disable_causal,
))
self.upsamplers = None
if add_upsample:
self.upsamplers = nn.ModuleList([
UpsampleCausal3D(
out_channels,
out_channels=out_channels,
upsample_factor=upsample_scale_factor,
disable_causal=disable_causal,
)
])
def forward(
self,
hidden_states: torch.Tensor,
temb: Optional[torch.Tensor] = None,
scale: float = 1.0,
) -> torch.Tensor:
for resnet in self.resnets:
hidden_states = resnet(hidden_states, temb=temb, scale=scale)
if self.upsamplers is not None:
for us in self.upsamplers:
hidden_states = us(hidden_states)
return hidden_states
fastvideo/models/wan/AGENTS.md
+2 -2
View File
@@ -15,8 +15,8 @@ explicit compatibility exports.
- Keep `__init__.py` lightweight: no eager model or pipeline imports.
- Keep `definition.py` data-only. Config and preset names reference their
existing owners; do not copy defaults into the catalog. Its registration
groups preserve first-match detector ordering around DreamX. Checkpoint
existing owners; do not copy defaults into the catalog. Declaration order
is first-match detector order. Checkpoint
manifests and explicit pipeline overrides still choose the executable
pipeline; a definition must not silently pin that choice.
- Keep `configs.pipelines.wan` as explicit compatibility aliases, including
fastvideo/models/wan/definition.py
+114 -124
View File
@@ -34,130 +34,120 @@ class WanModelDefinition:
and not any(token in value for token in self.exclude))
# The two groups preserve the registry's existing first-match ordering around
# DreamX. Keeping this boundary avoids changing local/custom-checkpoint
# resolution when both a Wan and DreamX detector match. Within each group,
# declaration order is registration order.
WAN_MODEL_DEFINITION_GROUPS = (
(
WanModelDefinition(
pipeline_config="WanT2V480PConfig",
preset="wan_t2v_1_3b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", ),
workload_types=("t2v", ),
match_any=("wanpipeline", ),
),
WanModelDefinition(
pipeline_config="WanT2V720PConfig",
preset="wan_t2v_14b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-T2V-14B-Diffusers", "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers"),
workload_types=("t2v", ),
),
WanModelDefinition(
pipeline_config="WanI2V480PConfig",
preset="wan_i2v_14b_480p",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", ),
workload_types=("i2v", ),
match_any=("wanimagetovideo", ),
),
WanModelDefinition(
pipeline_config="WanI2V720PConfig",
preset="wan_i2v_14b_720p",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", ),
workload_types=("i2v", ),
),
WanModelDefinition(
pipeline_config="WanI2V480PConfig",
preset="wan_fun_1_3b_inp",
sampling="unipc",
hf_model_paths=("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", ),
workload_types=("i2v", ),
),
WanModelDefinition(
pipeline_config="WANV2VConfig",
preset="wan_fun_1_3b_control",
sampling="unipc",
hf_model_paths=("IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers", ),
workload_types=(),
),
WanModelDefinition(
pipeline_config="FastWan2_1_T2V_480P_Config",
preset="fast_wan_t2v_480p",
sampling="dmd",
hf_model_paths=("FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers"),
workload_types=("t2v", ),
match_any=("wandmdpipeline", ),
),
WanModelDefinition(
pipeline_config="Wan2_2_TI2V_5B_Config",
preset="wan_2_2_ti2v_5b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.2-TI2V-5B-Diffusers", ),
workload_types=("t2v", "i2v"),
),
# Declaration order is registration order; model detectors are first-match.
WAN_MODEL_DEFINITIONS = (
WanModelDefinition(
pipeline_config="WanT2V480PConfig",
preset="wan_t2v_1_3b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", ),
workload_types=("t2v", ),
match_any=("wanpipeline", ),
),
(
WanModelDefinition(
pipeline_config="FastWan2_2_TI2V_5B_Config",
preset="fast_wan_2_2_ti2v_5b",
sampling="dmd",
hf_model_paths=("FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
"FastVideo/FastWan2.2-TI2V-5B-Diffusers"),
workload_types=("t2v", "i2v"),
),
WanModelDefinition(
pipeline_config="LucyEditDevConfig",
preset="lucy_edit_dev",
sampling="unipc",
hf_model_paths=("decart-ai/Lucy-Edit-Dev", "decart-ai/Lucy-Edit-1.1-Dev"),
workload_types=(),
match_any=("lucy-edit", ),
),
WanModelDefinition(
pipeline_config="Wan2_2_T2V_A14B_Config",
preset="wan_2_2_t2v_a14b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.2-T2V-A14B-Diffusers", ),
workload_types=("t2v", ),
),
WanModelDefinition(
pipeline_config="Wan2_2_I2V_A14B_Config",
preset="wan_2_2_i2v_a14b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.2-I2V-A14B-Diffusers", ),
workload_types=("i2v", ),
),
WanModelDefinition(
pipeline_config="SelfForcingWanT2V480PConfig",
preset="sf_wan_t2v_1_3b",
sampling="causal_dmd",
hf_model_paths=("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", ),
workload_types=("t2v", ),
match_any=("wancausaldmdpipeline", ),
),
WanModelDefinition(
pipeline_config="SelfForcingWan2_2_T2V480PConfig",
preset="sf_wan_2_2_t2v_a14b",
sampling="causal_dmd",
hf_model_paths=("rand0nmr/SFWan2.2-T2V-A14B-Diffusers", ),
workload_types=("t2v", ),
match_any=("sfwan2.2", "sfwan2_2"),
exclude=("i2v", ),
),
WanModelDefinition(
pipeline_config="SelfForcingWan2_2_T2V480PConfig",
preset="sf_wan_2_2_i2v_a14b",
sampling="causal_dmd",
hf_model_paths=("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", ),
workload_types=("i2v", ),
match_any=("sfwan2.2", "sfwan2_2"),
match_all=("i2v", ),
),
WanModelDefinition(
pipeline_config="WanT2V720PConfig",
preset="wan_t2v_14b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-T2V-14B-Diffusers", "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers"),
workload_types=("t2v", ),
),
WanModelDefinition(
pipeline_config="WanI2V480PConfig",
preset="wan_i2v_14b_480p",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", ),
workload_types=("i2v", ),
match_any=("wanimagetovideo", ),
),
WanModelDefinition(
pipeline_config="WanI2V720PConfig",
preset="wan_i2v_14b_720p",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", ),
workload_types=("i2v", ),
),
WanModelDefinition(
pipeline_config="WanI2V480PConfig",
preset="wan_fun_1_3b_inp",
sampling="unipc",
hf_model_paths=("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", ),
workload_types=("i2v", ),
),
WanModelDefinition(
pipeline_config="WANV2VConfig",
preset="wan_fun_1_3b_control",
sampling="unipc",
hf_model_paths=("IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers", ),
workload_types=(),
),
WanModelDefinition(
pipeline_config="FastWan2_1_T2V_480P_Config",
preset="fast_wan_t2v_480p",
sampling="dmd",
hf_model_paths=("FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers"),
workload_types=("t2v", ),
match_any=("wandmdpipeline", ),
),
WanModelDefinition(
pipeline_config="Wan2_2_TI2V_5B_Config",
preset="wan_2_2_ti2v_5b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.2-TI2V-5B-Diffusers", ),
workload_types=("t2v", "i2v"),
),
WanModelDefinition(
pipeline_config="FastWan2_2_TI2V_5B_Config",
preset="fast_wan_2_2_ti2v_5b",
sampling="dmd",
hf_model_paths=("FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", "FastVideo/FastWan2.2-TI2V-5B-Diffusers"),
workload_types=("t2v", "i2v"),
),
WanModelDefinition(
pipeline_config="LucyEditDevConfig",
preset="lucy_edit_dev",
sampling="unipc",
hf_model_paths=("decart-ai/Lucy-Edit-Dev", "decart-ai/Lucy-Edit-1.1-Dev"),
workload_types=(),
match_any=("lucy-edit", ),
),
WanModelDefinition(
pipeline_config="Wan2_2_T2V_A14B_Config",
preset="wan_2_2_t2v_a14b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.2-T2V-A14B-Diffusers", ),
workload_types=("t2v", ),
),
WanModelDefinition(
pipeline_config="Wan2_2_I2V_A14B_Config",
preset="wan_2_2_i2v_a14b",
sampling="unipc",
hf_model_paths=("Wan-AI/Wan2.2-I2V-A14B-Diffusers", ),
workload_types=("i2v", ),
),
WanModelDefinition(
pipeline_config="SelfForcingWanT2V480PConfig",
preset="sf_wan_t2v_1_3b",
sampling="causal_dmd",
hf_model_paths=("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", ),
workload_types=("t2v", ),
match_any=("wancausaldmdpipeline", ),
),
WanModelDefinition(
pipeline_config="SelfForcingWan2_2_T2V480PConfig",
preset="sf_wan_2_2_t2v_a14b",
sampling="causal_dmd",
hf_model_paths=("rand0nmr/SFWan2.2-T2V-A14B-Diffusers", ),
workload_types=("t2v", ),
match_any=("sfwan2.2", "sfwan2_2"),
exclude=("i2v", ),
),
WanModelDefinition(
pipeline_config="SelfForcingWan2_2_T2V480PConfig",
preset="sf_wan_2_2_i2v_a14b",
sampling="causal_dmd",
hf_model_paths=("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", ),
workload_types=("i2v", ),
match_any=("sfwan2.2", "sfwan2_2"),
match_all=("i2v", ),
),
)
WAN_MODEL_DEFINITIONS = tuple(definition for group in WAN_MODEL_DEFINITION_GROUPS for definition in group)
fastvideo/pipelines/AGENTS.md
+4 -6
View File
@@ -20,15 +20,13 @@ pipelines/
│ ├── conditioning.py # CFG / negative prompt fan-out
│ ├── denoising.py # Standard diffusion loop
│ ├── sd35_conditioning.py # Per-model overrides (named by family)
│ ├── longcat_*.py # LongCat I2V/V2V/refine variants
│ ├── gen3c_stages.py # Gen3C-specific stages
│ ├── gamecraft_denoising.py # GameCraft-specific
│ └── matrixgame2_denoising.py # Matrix-Game 2.0-specific
├── basic/ # Per-model end-to-end pipelines
│ ├── cosmos/, dreamx_world/, flux/, flux_2/, gamecraft/, gen3c/, glm_image/
│ ├── cosmos/, flux/, flux_2/, gen3c/
│ ├── hunyuan/, hunyuan15/, hyworld/, kandinsky5/, lingbot_video/, lingbotworld/
│ ├── lingbotworld2/, longcat/, ltx2/, magi_human/, matrixgame2/, matrixgame3/
│ ├── minimax_h3/, mmaudio/, sd35/, stable_audio/, turbodiffusion/, wan/, zimage/
│ ├── lingbotworld2/, ltx2/, magi_human/, matrixgame2/, matrixgame3/
│ ├── minimax_h3/, mmaudio/, sd35/, stable_audio/, wan/
│ └── <model>/{<model>_pipeline.py, presets.py, __init__.py}
├── preprocess/ # Data preprocessing pipelines (ltx2, wan, matrixgame2)
└── training/ # Training-time pipeline glue
@@ -63,7 +61,7 @@ Reuse `stages/text_encoding.py` if your model takes text → embeddings via a st
- The model needs a **different ForwardBatch shape** (extra inputs, different output keys).
- The denoising loop has structural differences (causal, refine-then-denoise, multi-stream).
When forking, keep the file name model-prefixed (`longcat_*`, `gamecraft_*`) so the registry stays grep-able.
When forking, keep the file name model-prefixed (`gen3c_*`, `matrixgame2_*`) so the registry stays grep-able.
Wan's family-specific stages live in `basic/wan/stages/`; see
`basic/wan/AGENTS.md`. Shared dense scheduling/CFG stays in `stages/denoising.py`.
fastvideo/pipelines/basic/dreamx_world/__init__.py
-28
View File
@@ -1,28 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from fastvideo.configs.pipelines.dreamx_world import (
DreamXWorld5BARPipelineConfig,
DreamXWorld5BCamPipelineConfig,
make_dreamx_world_5b_ar_dit_config,
make_dreamx_world_5b_cam_dit_config,
make_dreamx_world_5b_cam_text_encoder_config,
make_dreamx_world_5b_cam_vae_config,
)
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
from fastvideo.pipelines.basic.dreamx_world.stages import (
DREAMX_Y_CAMERA_KEY,
DreamXWorldCameraConditioningStage,
)
__all__ = [
"DREAMX_Y_CAMERA_KEY",
"DreamXWorld5BARPipelineConfig",
"DreamXWorld5BCamPipelineConfig",
"DreamXWorldCameraConditioningStage",
"DreamXWorldARPipeline",
"DreamXWorldPipeline",
"make_dreamx_world_5b_ar_dit_config",
"make_dreamx_world_5b_cam_dit_config",
"make_dreamx_world_5b_cam_text_encoder_config",
"make_dreamx_world_5b_cam_vae_config",
]
fastvideo/pipelines/basic/dreamx_world/ar_denoising.py
-219
View File
@@ -1,219 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World autoregressive causal denoising stage."""
from __future__ import annotations
from typing import Any
import torch
from tqdm.auto import tqdm
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.pipelines.basic.dreamx_world.stages import DREAMX_Y_CAMERA_KEY
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.denoising import DenoisingStage
class DreamXWorldARCausalDenoisingStage(DenoisingStage):
"""Official DreamX AR-forcing denoising loop with KV cache."""
_AR_NOISE_SEED_OFFSET = 1_000_003
def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
self.num_transformer_blocks = len(self.transformer.blocks)
self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
assert batch.latents is not None, "latents must be prepared before DreamX AR denoising"
assert batch.prompt_embeds, "prompt embeds must be prepared before DreamX AR denoising"
latents = batch.latents
device = latents.device
target_dtype = torch.bfloat16
autocast_enabled = device.type == "cuda" and not fastvideo_args.disable_autocast
frame_seq_length = (latents.shape[-2] // self.transformer.patch_size[1]) * (latents.shape[-1] //
self.transformer.patch_size[2])
timesteps = torch.tensor(
tuple(getattr(fastvideo_args.pipeline_config, "dmd_denoising_steps", (1000, 750, 500, 250))),
dtype=torch.long,
).cpu()
if getattr(fastvideo_args.pipeline_config, "warp_denoising_step", True):
self.scheduler.set_timesteps(1000)
scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
timesteps = scheduler_timesteps[1000 - timesteps]
timesteps = timesteps.to(device)
if latents.shape[2] % self.num_frame_per_block != 0:
raise ValueError("DreamX AR latent frames must be divisible by num_frame_per_block")
y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY, batch.extra.get("y_camera"))
if isinstance(y_camera, dict):
y_camera = {
k: v.to(device=device, dtype=target_dtype) if torch.is_tensor(v) else v
for k, v in y_camera.items()
}
if batch.image_latent is not None and batch.image_latent.shape[1] == latents.shape[1]:
latents[:, :, :batch.image_latent.shape[2]] = batch.image_latent.to(device=device, dtype=latents.dtype)
kv_cache = self._initialize_kv_cache(latents.shape[0], target_dtype, device, frame_seq_length)
crossattn_cache = self._initialize_crossattn_cache(latents.shape[0], target_dtype, device)
prompt = batch.prompt_embeds[0]
if torch.is_tensor(prompt):
prompt = prompt.to(device=device, dtype=target_dtype)
context = [sample for sample in prompt]
else:
context = prompt
num_blocks = latents.shape[2] // self.num_frame_per_block
start = 0
first_frame_mask = torch.ones_like(latents)
first_frame_mask[:, :, 0] = 0
base_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
noise_generator = self._make_noise_generator(base_generator, device)
with tqdm(total=num_blocks * len(timesteps), desc="DreamX AR denoising", leave=False) as progress:
for _ in range(num_blocks):
current_num_frames = self.num_frame_per_block
block_latents = latents[:, :, start:start + current_num_frames]
noisy_input = block_latents.clone()
mask_block = first_frame_mask[:, :, start:start + current_num_frames]
camera_block = self._slice_camera(y_camera, start, current_num_frames)
for idx, current_timestep in enumerate(timesteps):
timestep = torch.full(
(latents.shape[0], current_num_frames * frame_seq_length),
int(current_timestep.item()),
device=device,
dtype=torch.long,
)
if start == 0:
timestep[:, :frame_seq_length] = 0
with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
denoised = self.transformer(
hidden_states=block_latents.to(target_dtype),
encoder_hidden_states=torch.stack(context).to(target_dtype),
timestep=timestep,
y_camera=camera_block,
kv_cache=kv_cache,
crossattn_cache=crossattn_cache,
current_start=start * frame_seq_length,
)
denoised = denoised.to(latents.dtype)
if idx < len(timesteps) - 1:
next_timestep = torch.full((latents.shape[0], current_num_frames),
int(timesteps[idx + 1].item()),
device=device,
dtype=torch.long)
noise_kwargs = {"device": device, "dtype": denoised.dtype}
if noise_generator is not None:
noise_kwargs["generator"] = noise_generator
noise = torch.randn(denoised.permute(0, 2, 1, 3, 4).shape, **noise_kwargs)
block_btchw = self.scheduler.add_noise(
denoised.permute(0, 2, 1, 3, 4).flatten(0, 1),
noise.flatten(0, 1),
next_timestep.flatten(),
).unflatten(0, (latents.shape[0], current_num_frames))
block_latents = block_btchw.permute(0, 2, 1, 3, 4)
block_latents = block_latents * mask_block + noisy_input * (1 - mask_block)
else:
block_latents = denoised * mask_block + noisy_input * (1 - mask_block)
progress.update()
latents[:, :, start:start + current_num_frames] = block_latents
self._update_context_cache(block_latents, context, camera_block, kv_cache, crossattn_cache, start,
frame_seq_length, target_dtype, autocast_enabled,
float(getattr(fastvideo_args.pipeline_config, "context_noise", 0.1)))
start += current_num_frames
batch.latents = latents
return batch
def _make_noise_generator(self, generator: torch.Generator | None, device: torch.device) -> torch.Generator | None:
if generator is None:
return None
if getattr(generator, "device", None) == device:
return generator
seed = int(generator.initial_seed()) + self._AR_NOISE_SEED_OFFSET
return torch.Generator(device=device).manual_seed(seed)
@staticmethod
def _context_noise_timestep(context_noise: float) -> int:
if 0.0 < context_noise <= 1.0:
return int(context_noise * 1000)
return int(context_noise)
def _slice_camera(self, y_camera: Any, start: int, num_frames: int):
if not isinstance(y_camera, dict):
return y_camera
return {
"viewmats": y_camera["viewmats"][:, start:start + num_frames],
"K": y_camera["K"][:, start:start + num_frames],
}
def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device,
frame_seq_length: int) -> list[dict[str, Any]]:
size = self.local_attn_size * frame_seq_length if self.local_attn_size != -1 else 18480
heads = self.transformer.num_attention_heads
head_dim = self.transformer.attention_head_dim
cam_self_attn = next(
(getattr(block, "cam_self_attn", None)
for block in self.transformer.blocks if getattr(block, "cam_self_attn", None) is not None),
None,
)
caches = []
for _ in range(self.num_transformer_blocks):
cache = {
"k": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
"v": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
"global_end_index": torch.tensor([0], dtype=torch.long, device=device),
"local_end_index": torch.tensor([0], dtype=torch.long, device=device),
}
if cam_self_attn is not None:
cam_heads = int(cam_self_attn.num_heads)
cam_head_dim = int(cam_self_attn.head_dim)
cache.update({
"prope_k":
torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
"prope_v":
torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
"prope_global_end_index":
torch.tensor([0], dtype=torch.long, device=device),
"prope_local_end_index":
torch.tensor([0], dtype=torch.long, device=device),
})
caches.append(cache)
return caches
def _initialize_crossattn_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
heads = self.transformer.num_attention_heads
head_dim = self.transformer.attention_head_dim
return [{
"k": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
"v": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
"is_init": False,
} for _ in range(self.num_transformer_blocks)]
def _update_context_cache(self, block_latents: torch.Tensor, context: Any, camera_block: Any,
kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], start: int,
frame_seq_length: int, target_dtype: torch.dtype, autocast_enabled: bool,
context_noise: float) -> None:
timestep = torch.full(
(block_latents.shape[0], block_latents.shape[2] * frame_seq_length),
self._context_noise_timestep(context_noise),
device=block_latents.device,
dtype=torch.long,
)
with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
self.transformer(
hidden_states=block_latents.to(target_dtype),
encoder_hidden_states=torch.stack(context).to(target_dtype),
timestep=timestep,
y_camera=camera_block,
kv_cache=kv_cache,
crossattn_cache=crossattn_cache,
current_start=start * frame_seq_length,
)
fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py
-228
View File
@@ -1,228 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import math
from dataclasses import dataclass
import numpy as np
import torch
from scipy.interpolate import interp1d
from scipy.spatial.transform import Rotation, Slerp
_ACTION_TO_MOTION = {
"w": "forward",
"a": "left",
"d": "right",
"s": "backward",
"j": "left_rot",
"l": "right_rot",
"i": "up_rot",
"k": "down_rot",
}
_TRANSLATION_BASE_UNIT = 1.0
_ROTATION_BASE_UNIT = 10.0
_INTRINSIC_ROW = [0.8, 0.5, 0.5, 0.5]
@dataclass
class DreamXCamera:
fx: float
fy: float
cx: float
cy: float
w2c_mat: np.ndarray
@property
def c2w_mat(self) -> np.ndarray:
return np.linalg.inv(self.w2c_mat)
@classmethod
def from_pose_row(cls, row: list[float]) -> DreamXCamera:
w2c_mat = np.eye(4, dtype=np.float64)
w2c_mat[:3, :] = np.asarray(row[7:], dtype=np.float64).reshape(3, 4)
return cls(
fx=float(row[1]),
fy=float(row[2]),
cx=float(row[3]),
cy=float(row[4]),
w2c_mat=w2c_mat,
)
def _translation_step(motion_type: str, current_pose: dict[str, np.ndarray], value: float, duration: int) -> np.ndarray:
if motion_type in ("forward", "backward"):
yaw = np.radians(current_pose["rotation"][1])
pitch = np.radians(current_pose["rotation"][0])
forward = np.array([-math.sin(yaw) * math.cos(pitch), math.sin(pitch), math.cos(yaw) * math.cos(pitch)])
direction = 1 if motion_type == "forward" else -1
return forward * value * direction / duration
if motion_type in ("left", "right"):
yaw = np.radians(current_pose["rotation"][1])
right = np.array([math.cos(yaw), 0.0, math.sin(yaw)])
direction = -1 if motion_type == "left" else 1
return right * value * direction / duration
return np.zeros(3)
def _rotation_step(motion_type: str, value: float, duration: int) -> np.ndarray:
if not motion_type.endswith("rot"):
return np.zeros(3)
axis = motion_type.split("_")[0]
rotation = np.zeros(3)
if axis == "left":
rotation[1] = value
elif axis == "right":
rotation[1] = -value
elif axis == "up":
rotation[0] = -value
elif axis == "down":
rotation[0] = value
return rotation / duration
def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
pitch, yaw, roll = np.radians(angles)
cy = math.cos(yaw * 0.5)
sy = math.sin(yaw * 0.5)
cp = math.cos(pitch * 0.5)
sp = math.sin(pitch * 0.5)
cr = math.cos(roll * 0.5)
sr = math.sin(roll * 0.5)
return [
cy * cp * cr + sy * sp * sr,
cy * sp * cr + sy * cp * sr,
sy * cp * cr - cy * sp * sr,
cy * cp * sr - sy * sp * cr,
]
def _quaternion_to_rotation_matrix(quaternion: list[float]) -> np.ndarray:
qw, qx, qy, qz = quaternion
return np.array([
[1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
[2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
[2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
])
def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float], duration: int) -> list[list[float]]:
if len(action_seq) != len(action_speed_list):
raise ValueError("action_seq and action_speed_list must have the same length")
positions: list[np.ndarray] = []
rotations: list[np.ndarray] = []
current_pose = {
"position": np.array([0.0, 0.0, 0.0]),
"rotation": np.array([0.0, 0.0, 0.0]),
}
for action_id, speed in zip(action_seq, action_speed_list, strict=True):
motion_types = [_ACTION_TO_MOTION[key] for key in list(action_id)]
translation_step = np.zeros(3)
rotation_step = np.zeros(3)
for motion_type in motion_types:
translation_step += _translation_step(motion_type, current_pose,
float(speed) * _TRANSLATION_BASE_UNIT, duration)
rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
segment_positions = []
segment_rotations = []
for index in range(1, duration + 1):
segment_positions.append(current_pose["position"] + translation_step * index)
segment_rotations.append(current_pose["rotation"] + rotation_step * index)
current_pose["position"] = segment_positions[-1].copy()
current_pose["rotation"] = segment_rotations[-1].copy()
positions.extend(segment_positions)
rotations.extend(segment_rotations)
rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] +
[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
for index, (position, rotation) in enumerate(zip(positions, rotations, strict=False)):
rotation_matrix = _quaternion_to_rotation_matrix(_euler_to_quaternion(rotation))
translation = -rotation_matrix @ position
extrinsic = np.hstack([rotation_matrix, translation.reshape(3, 1)])
rows.append([float(index)] + _INTRINSIC_ROW + [0.0, 0.0] + extrinsic.flatten().tolist())
return rows
def _interpolate_camera_poses(
cameras: list[DreamXCamera],
src_indices: np.ndarray,
tgt_indices: np.ndarray,
) -> list[DreamXCamera]:
if len(cameras) <= 1:
return [cameras[0]] * len(tgt_indices) if cameras else []
src_rot_mat = np.array([camera.w2c_mat[:3, :3] for camera in cameras])
src_trans_vec = np.array([camera.w2c_mat[:3, 3] for camera in cameras])
dets = np.linalg.det(src_rot_mat)
flip_handedness = dets.size > 0 and np.median(dets) < 0.0
if flip_handedness:
flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
src_rot_mat = src_rot_mat @ flip_mat
trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False,
fill_value="extrapolate")(tgt_indices)
quats = Rotation.from_matrix(src_rot_mat).as_quat().copy()
for index in range(1, len(quats)):
if np.dot(quats[index], quats[index - 1]) < 0:
quats[index] = -quats[index]
rot = Slerp(src_indices, Rotation.from_quat(quats))(tgt_indices).as_matrix()
if flip_handedness:
rot = rot @ flip_mat
ref = cameras[0]
result = []
for index in range(len(tgt_indices)):
w2c_mat = np.eye(4, dtype=np.float64)
w2c_mat[:3, :] = np.hstack([rot[index], trans[index].reshape(3, 1)])
result.append(DreamXCamera(ref.fx, ref.fy, ref.cx, ref.cy, w2c_mat))
return result
def _relative_c2w_poses(cameras: list[DreamXCamera]) -> np.ndarray:
abs_w2cs = [camera.w2c_mat for camera in cameras]
abs_c2ws = [camera.c2w_mat for camera in cameras]
target_cam_c2w = np.eye(4, dtype=np.float64)
abs2rel = target_cam_c2w @ abs_w2cs[0]
poses = [target_cam_c2w] + [abs2rel @ c2w for c2w in abs_c2ws[1:]]
return np.asarray(poses, dtype=np.float32)
def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
rotation_inv = transforms[..., :3, :3].transpose(-1, -2)
output = torch.zeros_like(transforms)
output[..., :3, :3] = rotation_inv
output[..., :3, 3] = -torch.einsum("...ij,...j->...i", rotation_inv, transforms[..., :3, 3])
output[..., 3, 3] = 1.0
return output
def build_dreamx_camera_condition(
action_seq: list[str],
action_speed_list: list[float],
*,
num_frames: int,
height: int,
width: int,
dtype: torch.dtype = torch.float32,
device: torch.device | str = "cpu",
) -> dict[str, torch.Tensor]:
del height, width # DreamX-World-5B-Cam uses fixed normalized intrinsics.
duration = math.ceil(num_frames / len(action_seq))
rows = _pose_rows_from_actions(action_seq, action_speed_list, duration)[:num_frames]
cameras = [DreamXCamera.from_pose_row(row) for row in rows]
latent_frame_count = 1 + (len(cameras) - 1) // 4
src_indices = np.arange(len(cameras), dtype=np.float64)
tgt_indices = np.linspace(0, len(cameras) - 1, latent_frame_count)
cameras = _interpolate_camera_poses(cameras, src_indices, tgt_indices)
c2ws = torch.as_tensor(_relative_c2w_poses(cameras), dtype=dtype, device=device)
viewmats = _invert_se3(c2ws)
intrinsics = torch.zeros((latent_frame_count, 3, 3), dtype=dtype, device=device)
intrinsics[:, 0, 0] = 969.6969696969696 / (960.0 * 2)
intrinsics[:, 1, 1] = 969.6969696969696 / (540.0 * 2)
intrinsics[:, 2, 2] = 1.0
return {"viewmats": viewmats, "K": intrinsics}
fastvideo/pipelines/basic/dreamx_world/config.py
-20
View File
@@ -1,20 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Compatibility exports for DreamX-World pipeline configs."""
from fastvideo.configs.pipelines.dreamx_world import (
DreamXWorld5BARPipelineConfig,
DreamXWorld5BCamPipelineConfig,
make_dreamx_world_5b_ar_dit_config,
make_dreamx_world_5b_cam_dit_config,
make_dreamx_world_5b_cam_text_encoder_config,
make_dreamx_world_5b_cam_vae_config,
)
__all__ = [
"DreamXWorld5BARPipelineConfig",
"DreamXWorld5BCamPipelineConfig",
"make_dreamx_world_5b_ar_dit_config",
"make_dreamx_world_5b_cam_dit_config",
"make_dreamx_world_5b_cam_text_encoder_config",
"make_dreamx_world_5b_cam_vae_config",
]
fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py
-67
View File
@@ -1,67 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World-5B autoregressive pipeline entrypoint."""
from fastvideo.api.sampling_param import SamplingParam
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
from fastvideo.pipelines.basic.dreamx_world.stages import (
DreamXWorldCameraConditioningStage,
DreamXWorldImageVAEEncodingStage,
)
from fastvideo.pipelines.stages import (
ConditioningStage,
DecodingStage,
InputValidationStage,
LatentPreparationStage,
TextEncodingStage,
TimestepPreparationStage,
)
logger = init_logger(__name__)
class DreamXWorldARPipeline(LoRAPipeline, ComposedPipelineBase):
"""DreamX-World-5B autoregressive causal camera pipeline."""
_required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
pipeline_config_cls = DreamXWorld5BARPipelineConfig
sampling_params_cls = SamplingParam
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
self.modules["scheduler"].set_timesteps(1000)
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
self.add_stage(stage_name="prompt_encoding_stage",
stage=TextEncodingStage(
text_encoders=[self.get_module("text_encoder")],
tokenizers=[self.get_module("tokenizer")],
))
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
self.add_stage(stage_name="timestep_preparation_stage",
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
self.add_stage(stage_name="latent_preparation_stage",
stage=LatentPreparationStage(
scheduler=self.get_module("scheduler"),
transformer=self.get_module("transformer", None),
))
self.add_stage(stage_name="image_latent_preparation_stage",
stage=DreamXWorldImageVAEEncodingStage(vae=self.get_module("vae")))
self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
self.add_stage(stage_name="denoising_stage",
stage=DreamXWorldARCausalDenoisingStage(
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler"),
vae=self.get_module("vae"),
pipeline=self,
))
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
logger.info("DreamXWorldARPipeline initialized with autoregressive causal denoising")
EntryClass = DreamXWorldARPipeline
fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py
-81
View File
@@ -1,81 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World video pipeline entrypoint."""
from fastvideo.api.sampling_param import SamplingParam
from fastvideo.pipelines.basic.wan.stages.conditioning import WanFirstFrameEncodingStage
from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraConditioningStage
from fastvideo.pipelines.stages import (
ConditioningStage,
DecodingStage,
InputValidationStage,
LatentPreparationStage,
TextEncodingStage,
TimestepPreparationStage,
)
logger = init_logger(__name__)
class DreamXWorldPipeline(LoRAPipeline, ComposedPipelineBase):
"""DreamX-World-5B-Cam pipeline with native FastVideo camera conditioning."""
_required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
pipeline_config_cls = DreamXWorld5BCamPipelineConfig
sampling_params_cls = SamplingParam
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
self.add_stage(
stage_name="prompt_encoding_stage",
stage=TextEncodingStage(
text_encoders=[self.get_module("text_encoder")],
tokenizers=[self.get_module("tokenizer")],
),
)
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
self.add_stage(
stage_name="timestep_preparation_stage",
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
)
self.add_stage(
stage_name="latent_preparation_stage",
stage=LatentPreparationStage(
scheduler=self.get_module("scheduler"),
transformer=self.get_module("transformer", None),
),
)
self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
self.add_stage(stage_name="first_frame_encoding_stage",
stage=WanFirstFrameEncodingStage(vae=self.get_module("vae")))
self.add_stage(
stage_name="denoising_stage",
stage=WanDenoisingStage(
transformer=self.get_module("transformer"),
transformer_2=self.get_module("transformer_2", None),
scheduler=self.get_module("scheduler"),
pipeline=self,
),
)
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
logger.info("DreamXWorldPipeline initialized with native camera conditioning")
EntryClass = DreamXWorldPipeline
fastvideo/pipelines/basic/dreamx_world/presets.py
-58
View File
@@ -1,58 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World model family pipeline presets."""
from fastvideo.api.presets import InferencePreset, PresetStageSpec
_NEGATIVE_PROMPT_CN = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，"
"静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，"
"多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，"
"形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，"
"背景人很多，倒着走")
_DENOISE_STAGE = PresetStageSpec(
name="denoise",
kind="denoising",
description="DreamX-World camera-conditioned denoising pass",
allowed_overrides=frozenset({
"num_inference_steps",
"guidance_scale",
}),
)
DREAMX_WORLD_5B_CAM = InferencePreset(
name="dreamx_world_5b_cam",
version=1,
model_family="dreamx_world",
description="DreamX-World 5B camera-control video generation",
workload_type="i2v",
stage_schemas=(_DENOISE_STAGE, ),
defaults={
"height": 480,
"width": 832,
"num_frames": 161,
"fps": 16,
"guidance_scale": 5.0,
"num_inference_steps": 30,
"negative_prompt": _NEGATIVE_PROMPT_CN,
},
)
DREAMX_WORLD_5B_AR = InferencePreset(
name="dreamx_world_5b_ar",
version=1,
model_family="dreamx_world",
description="DreamX-World 5B autoregressive camera-control generation",
workload_type="i2v",
stage_schemas=(_DENOISE_STAGE, ),
defaults={
"height": 704,
"width": 1280,
"num_frames": 1005,
"fps": 16,
"guidance_scale": 1.0,
"num_inference_steps": 4,
"negative_prompt": _NEGATIVE_PROMPT_CN,
},
)
ALL_PRESETS = (DREAMX_WORLD_5B_CAM, DREAMX_WORLD_5B_AR)
fastvideo/pipelines/basic/dreamx_world/stages.py
-156
View File
@@ -1,156 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World pipeline stages."""
from __future__ import annotations
import torch
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import VerificationResult
from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (
build_dreamx_camera_condition, )
DREAMX_Y_CAMERA_KEY = "dreamx_y_camera"
logger = init_logger(__name__)
class DreamXWorldCameraConditioningStage(PipelineStage):
"""Build PRoPE camera conditioning for DreamX-World-5B-Cam."""
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
del fastvideo_args
if DREAMX_Y_CAMERA_KEY in batch.extra:
return batch
action_seq = batch.extra.get("dreamx_action_seq", batch.action_list)
action_speed_list = batch.extra.get("dreamx_action_speed_list", batch.action_speed_list)
if action_seq is None:
action_seq = ["w"]
if action_speed_list is None:
action_speed_list = [4]
if isinstance(action_seq, str):
action_seq = [action_seq]
if isinstance(action_speed_list, int | float):
action_speed_list = [action_speed_list]
if len(action_speed_list) == 1 and len(action_seq) > 1:
action_speed_list = list(action_speed_list) * len(action_seq)
action_speed_list = [float(speed) for speed in action_speed_list]
height = int(batch.height) if batch.height is not None else 704
width = int(batch.width) if batch.width is not None else 1280
num_frames = int(batch.num_frames)
dtype = batch.latents.dtype if torch.is_tensor(batch.latents) else torch.float32
device = batch.latents.device if torch.is_tensor(batch.latents) else "cpu"
y_camera = build_dreamx_camera_condition(
list(action_seq),
action_speed_list,
num_frames=num_frames,
height=height,
width=width,
dtype=dtype,
device=device,
)
batch.extra[DREAMX_Y_CAMERA_KEY] = {key: value.unsqueeze(0) for key, value in y_camera.items()}
return batch
def verify_output(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> VerificationResult:
del fastvideo_args
result = VerificationResult()
y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY)
result.add_check("dreamx_y_camera", y_camera, lambda value: isinstance(value, dict))
if isinstance(y_camera, dict):
result.add_check("dreamx_y_camera.viewmats", y_camera.get("viewmats"), torch.is_tensor)
result.add_check("dreamx_y_camera.K", y_camera.get("K"), torch.is_tensor)
return result
class DreamXWorldImageVAEEncodingStage(PipelineStage):
"""Encode the conditioning image into the first-frame latent.
Official AR-forcing flow (AMAP-ML/DreamX-World inference_ar_forcing.py):
the input image is resized, normalized to [-1, 1], VAE-encoded
deterministically, and written into frame 0 of the noise — the causal
denoiser then treats frame 0 as clean context. This stage produces
``batch.image_latent`` ([B, C, 1, H_lat, W_lat]); the injection into
the latents happens in DreamXWorldARCausalDenoisingStage.
"""
def __init__(self, vae) -> None:
super().__init__()
self.vae = vae
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
if batch.pil_image is None:
# No conditioning image: the causal denoiser falls back to
# running from pure noise (frame 0 uninitialized). Warn loudly —
# this pipeline is registered I2V and the official flow always
# forces from a frame.
logger.warning("DreamXWorldARPipeline called without an input image; "
"first-frame context will be noise (T2V-style). Pass an "
"image for the official AR-forcing behavior.")
return batch
from fastvideo.platforms import get_local_torch_device
from fastvideo.utils import PRECISION_TO_TYPE
device = get_local_torch_device()
image = batch.pil_image
if not isinstance(image, torch.Tensor):
import numpy as np
import PIL.Image
assert isinstance(image, PIL.Image.Image)
width = batch.width if isinstance(batch.width, int) else batch.width[0]
height = batch.height if isinstance(batch.height, int) else batch.height[0]
image = image.convert("RGB").resize((width, height), PIL.Image.Resampling.LANCZOS)
arr = torch.from_numpy(np.asarray(image)).float().permute(2, 0, 1) / 255.0
image = (arr - 0.5) / 0.5 # official Normalize([0.5], [0.5])
image = image.unsqueeze(0) # [1, C, H, W]
if image.dim() == 4:
image = image.unsqueeze(2) # [B, C, 1, H, W]
elif image.dim() == 5:
image = image[:, :, :1]
image = image.to(device=device, dtype=torch.float32)
vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
self.vae = self.vae.to(device)
with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
if not vae_autocast_enabled:
image = image.to(vae_dtype)
encoder_output = self.vae.encode(image)
# Official encode_to_latent is deterministic ((mean - mu) / sigma per
# channel); the posterior mean + shift/scale is the FastVideo
# equivalent of that normalization.
latent = encoder_output.mean
if getattr(self.vae, "shift_factor", None) is not None:
shift = self.vae.shift_factor
latent = latent - (shift.to(latent.device, latent.dtype) if isinstance(shift, torch.Tensor) else shift)
scale = self.vae.scaling_factor
latent = latent * (scale.to(latent.device, latent.dtype) if isinstance(scale, torch.Tensor) else scale)
batch.image_latent = latent
return batch
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
result = VerificationResult()
return result

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