Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 SolitaryThinkerandClaude Opus 5.5
|
b1bbc9e6d0 | ||
|
SolitaryThinkerandClaude Opus 5.5
|
d692d1cc3b |
@@ -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`
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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`,
|
||||
|
||||
@@ -12,6 +12,7 @@ from __future__ import annotations
|
||||
import argparse
|
||||
import fnmatch
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TextIO
|
||||
@@ -130,11 +131,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 +141,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 +156,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 +186,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 +195,6 @@ FAMILY_COVERAGE = (
|
||||
"test_wan_t2v_similarity.py",
|
||||
),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])z[_-]?image([/_.-]|$)"),
|
||||
("test_zimage.py", ),
|
||||
("test_zimage_similarity.py", ),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -319,16 +290,22 @@ def _select_output_coverage(plan: MergePlan, path: str) -> None:
|
||||
plan.add_ssim(SSIM_SMOKE_TESTS, reason=f"shared output SSIM smoke coverage: {path}")
|
||||
|
||||
|
||||
def classify_paths(paths: list[str]) -> MergePlan:
|
||||
def _normalize_path(raw_path: str) -> str:
|
||||
path = raw_path.strip()
|
||||
while path.startswith("./"):
|
||||
path = path[2:]
|
||||
return path
|
||||
|
||||
|
||||
def classify_paths(paths: list[str], removed_paths: Iterable[str] = ()) -> MergePlan:
|
||||
"""Plan merge lanes for ``paths``.
|
||||
|
||||
``removed_paths`` lists changed paths that no longer exist at the PR head
|
||||
(deleted files and rename sources); removed golden/SSIM tests are not run.
|
||||
"""
|
||||
plan = MergePlan()
|
||||
normalized_paths: list[str] = []
|
||||
for raw_path in paths:
|
||||
path = raw_path.strip()
|
||||
while path.startswith("./"):
|
||||
path = path[2:]
|
||||
if path:
|
||||
normalized_paths.append(path)
|
||||
normalized_paths = sorted(set(normalized_paths))
|
||||
normalized_paths = sorted({path for path in map(_normalize_path, paths) if path})
|
||||
removed = {path for path in map(_normalize_path, removed_paths) if path}
|
||||
if not normalized_paths:
|
||||
plan.require_all("changed-file list was empty; failing closed")
|
||||
return plan
|
||||
@@ -368,7 +345,10 @@ def classify_paths(paths: list[str]) -> MergePlan:
|
||||
if path.startswith("fastvideo/tests/golden_gate/"):
|
||||
name = Path(path).name
|
||||
if name.startswith("test_") and name.endswith(".py"):
|
||||
plan.add_golden((name, ), reason=f"changed golden test: {path}")
|
||||
if path in removed:
|
||||
plan.reasons.append(f"removed golden test has nothing to run: {path}")
|
||||
else:
|
||||
plan.add_golden((name, ), reason=f"changed golden test: {path}")
|
||||
elif name in {"AGENTS.md", "README.md"}:
|
||||
plan.reasons.append(f"golden documentation only: {path}")
|
||||
else:
|
||||
@@ -379,7 +359,10 @@ def classify_paths(paths: list[str]) -> MergePlan:
|
||||
if path.startswith("fastvideo/tests/ssim/"):
|
||||
name = Path(path).name
|
||||
if name.startswith("test_") and name.endswith(".py"):
|
||||
plan.add_ssim((name, ), reason=f"changed SSIM test: {path}")
|
||||
if path in removed:
|
||||
plan.reasons.append(f"removed SSIM test has nothing to run: {path}")
|
||||
else:
|
||||
plan.add_ssim((name, ), reason=f"changed SSIM test: {path}")
|
||||
elif path.endswith((".py", ".json", ".pt", ".png", ".mp4")):
|
||||
plan.ssim_all = True
|
||||
plan.add_lanes("ssim", reason=f"shared SSIM harness/reference: {path}")
|
||||
@@ -555,6 +538,7 @@ def _write_summary(output: TextIO, plan: MergePlan) -> None:
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--paths-file", type=Path, required=True)
|
||||
parser.add_argument("--removed-paths-file", type=Path)
|
||||
parser.add_argument("--github-output", type=Path)
|
||||
parser.add_argument("--summary-file", type=Path)
|
||||
return parser.parse_args()
|
||||
@@ -563,7 +547,9 @@ def parse_args() -> argparse.Namespace:
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
paths = args.paths_file.read_text(encoding="utf-8").splitlines()
|
||||
plan = classify_paths(paths)
|
||||
removed_paths = (args.removed_paths_file.read_text(encoding="utf-8").splitlines()
|
||||
if args.removed_paths_file else [])
|
||||
plan = classify_paths(paths, removed_paths)
|
||||
print(f"MERGE_TEST_PLAN={plan.encoded_lanes()}")
|
||||
print(f"MERGE_GOLDEN_TESTS={plan.encoded_golden_tests()}")
|
||||
print(f"MERGE_SSIM_TESTS={plan.encoded_ssim_tests()}")
|
||||
|
||||
@@ -76,6 +76,8 @@ jobs:
|
||||
set -euo pipefail
|
||||
changed_json="$RUNNER_TEMP/merge-changed-files.json"
|
||||
changed_paths="$RUNNER_TEMP/merge-changed-paths.txt"
|
||||
removed_paths="$RUNNER_TEMP/merge-removed-paths.txt"
|
||||
: > "$removed_paths"
|
||||
if gh api --paginate --slurp \
|
||||
"repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}/files?per_page=100" \
|
||||
> "$changed_json"; then
|
||||
@@ -83,6 +85,10 @@ jobs:
|
||||
if [ "$observed" = "$EXPECTED_CHANGED_FILES" ]; then
|
||||
jq -r '.[][] | .filename, (.previous_filename // empty)' "$changed_json" \
|
||||
| sort -u > "$changed_paths"
|
||||
# Paths absent from the PR head, so the planner never selects a deleted test.
|
||||
jq -r '.[][] | if .status == "removed" then .filename
|
||||
elif .status == "renamed" then (.previous_filename // empty) else empty end' \
|
||||
"$changed_json" | sort -u > "$removed_paths"
|
||||
else
|
||||
echo "::warning::Changed-file API returned $observed of $EXPECTED_CHANGED_FILES paths; selecting all merge lanes."
|
||||
echo '__FASTVIDEO_CI_PLAN_ALL__' > "$changed_paths"
|
||||
@@ -98,6 +104,7 @@ jobs:
|
||||
run: |
|
||||
python3 .github/scripts/plan_merge_ci.py \
|
||||
--paths-file "$RUNNER_TEMP/merge-changed-paths.txt" \
|
||||
--removed-paths-file "$RUNNER_TEMP/merge-removed-paths.txt" \
|
||||
--github-output "$GITHUB_OUTPUT" \
|
||||
--summary-file "$GITHUB_STEP_SUMMARY"
|
||||
|
||||
|
||||
@@ -36,7 +36,6 @@ env
|
||||
*.log
|
||||
weights/
|
||||
logs/
|
||||
/Z-Image/
|
||||
official_weights/
|
||||
converted_weights/
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 5.1 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.0 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 2.4 KiB |
@@ -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 -->
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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>
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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>
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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">
|
||||
|
||||
@@ -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>
|
||||
@@ -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">
|
||||
|
||||
@@ -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>
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -163,31 +152,25 @@ optimizations: absence means **untested**, not incompatible.
|
||||
}
|
||||
</style>
|
||||
|
||||
| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn (Legacy Branch) | Sage Attn | VSA | BSA |
|
||||
|------------|---------------------|-------------|----------|-------------------|-----------|-----|-----|
|
||||
| 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 | ❌ | ❌ | ✅ | ⭕ | ⭕ |
|
||||
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| 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 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |
|
||||
| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn (Legacy Branch) | Sage Attn | VSA |
|
||||
|------------|---------------------|-------------|----------|-------------------|-----------|-----|
|
||||
| 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 | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| 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 | ❌ | ❌ | ✅ | ⭕ |
|
||||
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ |
|
||||
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ |
|
||||
| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ |
|
||||
| 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 | ✅ | ✅ | ✅ | ⭕ |
|
||||
| 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 | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ |
|
||||
|
||||
## Apple Silicon native runtime
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -1,107 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Run GLM-Image text-to-image generation through FastVideo.
|
||||
|
||||
User story:
|
||||
"I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
|
||||
text-to-image generation command, saved as a PNG."
|
||||
"""
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default="zai-org/GLM-Image",
|
||||
help="HF id or local diffusers-format GLM-Image weights directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="image_output/landscape.png",
|
||||
help="Output PNG path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default=("A beautiful landscape photography with rolling hills, "
|
||||
"a winding river, and a vibrant sunset in the background. "
|
||||
"Warm golden light, photorealistic style."),
|
||||
help="Text prompt.",
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--steps", type=int, default=50)
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.5)
|
||||
parser.add_argument("--seed", type=int, default=1024)
|
||||
parser.add_argument("--num-gpus", type=int, default=1)
|
||||
parser.add_argument("--tp-size", type=int, default=None)
|
||||
parser.add_argument("--sp-size", type=int, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
output = Path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
|
||||
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
|
||||
|
||||
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
|
||||
# pipeline class come from the model's registered defaults — don't override.
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
trust_remote_code=True,
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
|
||||
),
|
||||
pipeline=PipelineSelection(workload_type="t2i"),
|
||||
)
|
||||
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
try:
|
||||
request = GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=1,
|
||||
fps=1,
|
||||
num_inference_steps=args.steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output.parent),
|
||||
save_video=False,
|
||||
return_frames=True,
|
||||
),
|
||||
)
|
||||
result = generator.generate(request)
|
||||
if isinstance(result, list):
|
||||
result = result[0]
|
||||
|
||||
frames = result.frames
|
||||
if frames is not None and len(frames):
|
||||
Image.fromarray(frames[0]).save(output)
|
||||
print(f"Saved image to {output}")
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
|
||||
@@ -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"
|
||||
@@ -1,61 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
|
||||
|
||||
def is_blocks(n: str, m) -> bool:
|
||||
return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageDiTArchConfig(DiTArchConfig):
|
||||
|
||||
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
|
||||
|
||||
hidden_size: int = 4096
|
||||
num_attention_heads: int = 32
|
||||
attention_head_dim: int = 128
|
||||
in_channels: int = 16
|
||||
out_channels: int = 16
|
||||
num_layers: int = 30
|
||||
|
||||
text_embed_dim: int = 1472
|
||||
time_embed_dim: int = 512
|
||||
condition_dim: int = 256
|
||||
|
||||
prior_vq_quantizer_codebook_size: int = 16384
|
||||
|
||||
patch_size: int = 2
|
||||
|
||||
max_height: int = 2048
|
||||
max_width: int = 2048
|
||||
|
||||
qk_norm: str = "layer_norm"
|
||||
eps: float = 1e-5
|
||||
|
||||
exclude_lora_layers: list[str] = field(
|
||||
default_factory=lambda: ["image_projector", "glyph_projector", "prior_token_embedding"])
|
||||
|
||||
param_names_mapping: dict = field(
|
||||
default_factory=lambda: {
|
||||
r"^glyph_projector\.net\.0\.proj\.(.*)$": r"glyph_projector.fc_in.\1",
|
||||
r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1",
|
||||
r"^prior_projector\.net\.0\.proj\.(.*)$": r"prior_projector.fc_in.\1",
|
||||
r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1",
|
||||
r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
|
||||
r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
|
||||
})
|
||||
|
||||
reverse_param_names_mapping: dict = field(default_factory=dict)
|
||||
lora_param_names_mapping: dict = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.num_channels_latents = self.out_channels
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageDiTConfig(DiTConfig):
|
||||
arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig)
|
||||
prefix: str = "GlmImage"
|
||||
@@ -1,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"
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
@@ -1,95 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.vaes.autoencoder_kl import (AutoencoderKLArchConfig, AutoencoderKLVAEConfig)
|
||||
|
||||
_GLM_IMAGE_LATENTS_MEAN: tuple[float, ...] = (
|
||||
-0.2080078125,
|
||||
1.875,
|
||||
-0.470703125,
|
||||
-1.265625,
|
||||
-1.421875,
|
||||
0.77734375,
|
||||
-0.3671875,
|
||||
-0.9453125,
|
||||
0.318359375,
|
||||
0.7734375,
|
||||
-0.1884765625,
|
||||
-0.022216796875,
|
||||
-0.220703125,
|
||||
-1.59375,
|
||||
-0.81640625,
|
||||
-0.255859375,
|
||||
)
|
||||
_GLM_IMAGE_LATENTS_STD: tuple[float, ...] = (
|
||||
3.0625,
|
||||
2.203125,
|
||||
2.265625,
|
||||
4.84375,
|
||||
2.5,
|
||||
3.9375,
|
||||
2.203125,
|
||||
3.03125,
|
||||
2.1875,
|
||||
2.046875,
|
||||
2.71875,
|
||||
2.390625,
|
||||
2.390625,
|
||||
2.453125,
|
||||
2.25,
|
||||
2.15625,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageVAEArchConfig(AutoencoderKLArchConfig):
|
||||
act_fn: str = "silu"
|
||||
block_out_channels: tuple[int, ...] = (128, 512, 1024, 1024)
|
||||
down_block_types: tuple[str, ...] = (
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
)
|
||||
up_block_types: tuple[str, ...] = (
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
)
|
||||
force_upcast: bool = True
|
||||
in_channels: int = 3
|
||||
out_channels: int = 3
|
||||
latent_channels: int = 16
|
||||
latents_mean: tuple[float, ...] = _GLM_IMAGE_LATENTS_MEAN
|
||||
latents_std: tuple[float, ...] = _GLM_IMAGE_LATENTS_STD
|
||||
layers_per_block: int = 3
|
||||
mid_block_add_attention: bool = False
|
||||
norm_num_groups: int = 32
|
||||
sample_size: int = 1024
|
||||
scaling_factor: float = 0.18215
|
||||
shift_factor: float | None = None
|
||||
use_quant_conv: bool = False
|
||||
use_post_quant_conv: bool = False
|
||||
|
||||
temporal_compression_ratio: int = 1
|
||||
spatial_compression_ratio: int = 8
|
||||
|
||||
|
||||
@dataclass
|
||||
class GlmImageVAEConfig(AutoencoderKLVAEConfig):
|
||||
arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig)
|
||||
|
||||
use_tiling: bool = True
|
||||
use_temporal_tiling: bool = False
|
||||
use_parallel_tiling: bool = False
|
||||
|
||||
tile_sample_min_height: int = 512
|
||||
tile_sample_min_width: int = 512
|
||||
tile_sample_stride_height: int = 384
|
||||
tile_sample_stride_width: int = 384
|
||||
|
||||
load_encoder: bool = True
|
||||
load_decoder: bool = True
|
||||
@@ -1,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"
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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}'")
|
||||
@@ -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
|
||||
@@ -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
|
||||
+16
-32
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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})
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -21,7 +21,6 @@ models/
|
||||
├── schedulers/ # FlowMatch / EulerDiscrete / DPM custom schedulers
|
||||
├── upsamplers/ # Hunyuan15 super-resolution
|
||||
├── audio/ # Audio-VAE/decoder modules (LTX-2 audio, Stable Audio)
|
||||
├── camera/ # Camera-conditioning modules (Gen3C)
|
||||
└── loader/ # component_loader.py, fsdp_load.py, weight_utils.py
|
||||
```
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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()
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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",
|
||||
@@ -70,12 +62,6 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
|
||||
"MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
|
||||
}
|
||||
|
||||
# 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 = {
|
||||
"MMAudioDFNCLIPTextEncoder": ("encoders", "mmaudio_clip", "MMAudioDFNCLIPTextEncoder"),
|
||||
"CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
|
||||
@@ -88,9 +74,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,14 +94,12 @@ _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"),
|
||||
"LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"),
|
||||
"AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"),
|
||||
"AutoencoderKLGen3CTokenizer": ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
|
||||
"AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"),
|
||||
"CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
|
||||
"AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"),
|
||||
# `stable-audio-open-1.0/vae/config.json` ships `_class_name="AutoencoderOobleck"`
|
||||
@@ -142,7 +123,6 @@ _SCHEDULERS = {
|
||||
"FlowUniPCMultistepScheduler": ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"),
|
||||
"SelfForcingFlowMatchScheduler":
|
||||
("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"),
|
||||
"RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"),
|
||||
}
|
||||
|
||||
_UPSAMPLERS = {
|
||||
@@ -154,7 +134,6 @@ _UPSAMPLERS = {
|
||||
_LEGACY_FAST_VIDEO_MODELS = {
|
||||
**_TEXT_TO_VIDEO_DIT_MODELS,
|
||||
**_IMAGE_TO_VIDEO_DIT_MODELS,
|
||||
**_TEXT_TO_IMAGE_DIT_MODELS,
|
||||
**_TEXT_ENCODER_MODELS,
|
||||
**_IMAGE_ENCODER_MODELS,
|
||||
**_VAE_MODELS,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
+114
-124
@@ -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)
|
||||
|
||||
@@ -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/
|
||||
│ ├── hunyuan/, hunyuan15/, hyworld/, kandinsky5/, lingbot_video/, lingbotworld/
|
||||
│ ├── lingbotworld2/, longcat/, ltx2/, magi_human/, matrixgame2/, matrixgame3/
|
||||
│ ├── minimax_h3/, mmaudio/, sd35/, stable_audio/, turbodiffusion/, wan/, zimage/
|
||||
│ ├── cosmos/, flux/, flux_2/, gen3c/, hunyuan/, hunyuan15/, hyworld/
|
||||
│ ├── kandinsky5/, kandinsky6/, kandinsky6_sr/, lingbot_video/, lingbotworld/
|
||||
│ ├── 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`.
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from fastvideo.configs.pipelines.dreamx_world import (
|
||||
DreamXWorld5BARPipelineConfig,
|
||||
DreamXWorld5BCamPipelineConfig,
|
||||
make_dreamx_world_5b_ar_dit_config,
|
||||
make_dreamx_world_5b_cam_dit_config,
|
||||
make_dreamx_world_5b_cam_text_encoder_config,
|
||||
make_dreamx_world_5b_cam_vae_config,
|
||||
)
|
||||
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
|
||||
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import (
|
||||
DREAMX_Y_CAMERA_KEY,
|
||||
DreamXWorldCameraConditioningStage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DREAMX_Y_CAMERA_KEY",
|
||||
"DreamXWorld5BARPipelineConfig",
|
||||
"DreamXWorld5BCamPipelineConfig",
|
||||
"DreamXWorldCameraConditioningStage",
|
||||
"DreamXWorldARPipeline",
|
||||
"DreamXWorldPipeline",
|
||||
"make_dreamx_world_5b_ar_dit_config",
|
||||
"make_dreamx_world_5b_cam_dit_config",
|
||||
"make_dreamx_world_5b_cam_text_encoder_config",
|
||||
"make_dreamx_world_5b_cam_vae_config",
|
||||
]
|
||||
@@ -1,219 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World autoregressive causal denoising stage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import DREAMX_Y_CAMERA_KEY
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.denoising import DenoisingStage
|
||||
|
||||
|
||||
class DreamXWorldARCausalDenoisingStage(DenoisingStage):
|
||||
"""Official DreamX AR-forcing denoising loop with KV cache."""
|
||||
|
||||
_AR_NOISE_SEED_OFFSET = 1_000_003
|
||||
|
||||
def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
|
||||
super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
|
||||
self.num_transformer_blocks = len(self.transformer.blocks)
|
||||
self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
|
||||
self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
assert batch.latents is not None, "latents must be prepared before DreamX AR denoising"
|
||||
assert batch.prompt_embeds, "prompt embeds must be prepared before DreamX AR denoising"
|
||||
latents = batch.latents
|
||||
device = latents.device
|
||||
target_dtype = torch.bfloat16
|
||||
autocast_enabled = device.type == "cuda" and not fastvideo_args.disable_autocast
|
||||
|
||||
frame_seq_length = (latents.shape[-2] // self.transformer.patch_size[1]) * (latents.shape[-1] //
|
||||
self.transformer.patch_size[2])
|
||||
timesteps = torch.tensor(
|
||||
tuple(getattr(fastvideo_args.pipeline_config, "dmd_denoising_steps", (1000, 750, 500, 250))),
|
||||
dtype=torch.long,
|
||||
).cpu()
|
||||
if getattr(fastvideo_args.pipeline_config, "warp_denoising_step", True):
|
||||
self.scheduler.set_timesteps(1000)
|
||||
scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
|
||||
timesteps = scheduler_timesteps[1000 - timesteps]
|
||||
timesteps = timesteps.to(device)
|
||||
|
||||
if latents.shape[2] % self.num_frame_per_block != 0:
|
||||
raise ValueError("DreamX AR latent frames must be divisible by num_frame_per_block")
|
||||
|
||||
y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY, batch.extra.get("y_camera"))
|
||||
if isinstance(y_camera, dict):
|
||||
y_camera = {
|
||||
k: v.to(device=device, dtype=target_dtype) if torch.is_tensor(v) else v
|
||||
for k, v in y_camera.items()
|
||||
}
|
||||
|
||||
if batch.image_latent is not None and batch.image_latent.shape[1] == latents.shape[1]:
|
||||
latents[:, :, :batch.image_latent.shape[2]] = batch.image_latent.to(device=device, dtype=latents.dtype)
|
||||
|
||||
kv_cache = self._initialize_kv_cache(latents.shape[0], target_dtype, device, frame_seq_length)
|
||||
crossattn_cache = self._initialize_crossattn_cache(latents.shape[0], target_dtype, device)
|
||||
prompt = batch.prompt_embeds[0]
|
||||
if torch.is_tensor(prompt):
|
||||
prompt = prompt.to(device=device, dtype=target_dtype)
|
||||
context = [sample for sample in prompt]
|
||||
else:
|
||||
context = prompt
|
||||
|
||||
num_blocks = latents.shape[2] // self.num_frame_per_block
|
||||
start = 0
|
||||
first_frame_mask = torch.ones_like(latents)
|
||||
first_frame_mask[:, :, 0] = 0
|
||||
base_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
|
||||
noise_generator = self._make_noise_generator(base_generator, device)
|
||||
|
||||
with tqdm(total=num_blocks * len(timesteps), desc="DreamX AR denoising", leave=False) as progress:
|
||||
for _ in range(num_blocks):
|
||||
current_num_frames = self.num_frame_per_block
|
||||
block_latents = latents[:, :, start:start + current_num_frames]
|
||||
noisy_input = block_latents.clone()
|
||||
mask_block = first_frame_mask[:, :, start:start + current_num_frames]
|
||||
camera_block = self._slice_camera(y_camera, start, current_num_frames)
|
||||
|
||||
for idx, current_timestep in enumerate(timesteps):
|
||||
timestep = torch.full(
|
||||
(latents.shape[0], current_num_frames * frame_seq_length),
|
||||
int(current_timestep.item()),
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
if start == 0:
|
||||
timestep[:, :frame_seq_length] = 0
|
||||
with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
|
||||
denoised = self.transformer(
|
||||
hidden_states=block_latents.to(target_dtype),
|
||||
encoder_hidden_states=torch.stack(context).to(target_dtype),
|
||||
timestep=timestep,
|
||||
y_camera=camera_block,
|
||||
kv_cache=kv_cache,
|
||||
crossattn_cache=crossattn_cache,
|
||||
current_start=start * frame_seq_length,
|
||||
)
|
||||
denoised = denoised.to(latents.dtype)
|
||||
if idx < len(timesteps) - 1:
|
||||
next_timestep = torch.full((latents.shape[0], current_num_frames),
|
||||
int(timesteps[idx + 1].item()),
|
||||
device=device,
|
||||
dtype=torch.long)
|
||||
noise_kwargs = {"device": device, "dtype": denoised.dtype}
|
||||
if noise_generator is not None:
|
||||
noise_kwargs["generator"] = noise_generator
|
||||
noise = torch.randn(denoised.permute(0, 2, 1, 3, 4).shape, **noise_kwargs)
|
||||
block_btchw = self.scheduler.add_noise(
|
||||
denoised.permute(0, 2, 1, 3, 4).flatten(0, 1),
|
||||
noise.flatten(0, 1),
|
||||
next_timestep.flatten(),
|
||||
).unflatten(0, (latents.shape[0], current_num_frames))
|
||||
block_latents = block_btchw.permute(0, 2, 1, 3, 4)
|
||||
block_latents = block_latents * mask_block + noisy_input * (1 - mask_block)
|
||||
else:
|
||||
block_latents = denoised * mask_block + noisy_input * (1 - mask_block)
|
||||
progress.update()
|
||||
|
||||
latents[:, :, start:start + current_num_frames] = block_latents
|
||||
self._update_context_cache(block_latents, context, camera_block, kv_cache, crossattn_cache, start,
|
||||
frame_seq_length, target_dtype, autocast_enabled,
|
||||
float(getattr(fastvideo_args.pipeline_config, "context_noise", 0.1)))
|
||||
start += current_num_frames
|
||||
|
||||
batch.latents = latents
|
||||
return batch
|
||||
|
||||
def _make_noise_generator(self, generator: torch.Generator | None, device: torch.device) -> torch.Generator | None:
|
||||
if generator is None:
|
||||
return None
|
||||
if getattr(generator, "device", None) == device:
|
||||
return generator
|
||||
seed = int(generator.initial_seed()) + self._AR_NOISE_SEED_OFFSET
|
||||
return torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
@staticmethod
|
||||
def _context_noise_timestep(context_noise: float) -> int:
|
||||
if 0.0 < context_noise <= 1.0:
|
||||
return int(context_noise * 1000)
|
||||
return int(context_noise)
|
||||
|
||||
def _slice_camera(self, y_camera: Any, start: int, num_frames: int):
|
||||
if not isinstance(y_camera, dict):
|
||||
return y_camera
|
||||
return {
|
||||
"viewmats": y_camera["viewmats"][:, start:start + num_frames],
|
||||
"K": y_camera["K"][:, start:start + num_frames],
|
||||
}
|
||||
|
||||
def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device,
|
||||
frame_seq_length: int) -> list[dict[str, Any]]:
|
||||
size = self.local_attn_size * frame_seq_length if self.local_attn_size != -1 else 18480
|
||||
heads = self.transformer.num_attention_heads
|
||||
head_dim = self.transformer.attention_head_dim
|
||||
cam_self_attn = next(
|
||||
(getattr(block, "cam_self_attn", None)
|
||||
for block in self.transformer.blocks if getattr(block, "cam_self_attn", None) is not None),
|
||||
None,
|
||||
)
|
||||
|
||||
caches = []
|
||||
for _ in range(self.num_transformer_blocks):
|
||||
cache = {
|
||||
"k": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
|
||||
"v": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
|
||||
"global_end_index": torch.tensor([0], dtype=torch.long, device=device),
|
||||
"local_end_index": torch.tensor([0], dtype=torch.long, device=device),
|
||||
}
|
||||
if cam_self_attn is not None:
|
||||
cam_heads = int(cam_self_attn.num_heads)
|
||||
cam_head_dim = int(cam_self_attn.head_dim)
|
||||
cache.update({
|
||||
"prope_k":
|
||||
torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
|
||||
"prope_v":
|
||||
torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
|
||||
"prope_global_end_index":
|
||||
torch.tensor([0], dtype=torch.long, device=device),
|
||||
"prope_local_end_index":
|
||||
torch.tensor([0], dtype=torch.long, device=device),
|
||||
})
|
||||
caches.append(cache)
|
||||
return caches
|
||||
|
||||
def _initialize_crossattn_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
|
||||
heads = self.transformer.num_attention_heads
|
||||
head_dim = self.transformer.attention_head_dim
|
||||
return [{
|
||||
"k": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
|
||||
"v": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
|
||||
"is_init": False,
|
||||
} for _ in range(self.num_transformer_blocks)]
|
||||
|
||||
def _update_context_cache(self, block_latents: torch.Tensor, context: Any, camera_block: Any,
|
||||
kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], start: int,
|
||||
frame_seq_length: int, target_dtype: torch.dtype, autocast_enabled: bool,
|
||||
context_noise: float) -> None:
|
||||
timestep = torch.full(
|
||||
(block_latents.shape[0], block_latents.shape[2] * frame_seq_length),
|
||||
self._context_noise_timestep(context_noise),
|
||||
device=block_latents.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
|
||||
self.transformer(
|
||||
hidden_states=block_latents.to(target_dtype),
|
||||
encoder_hidden_states=torch.stack(context).to(target_dtype),
|
||||
timestep=timestep,
|
||||
y_camera=camera_block,
|
||||
kv_cache=kv_cache,
|
||||
crossattn_cache=crossattn_cache,
|
||||
current_start=start * frame_seq_length,
|
||||
)
|
||||
@@ -1,228 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.interpolate import interp1d
|
||||
from scipy.spatial.transform import Rotation, Slerp
|
||||
|
||||
_ACTION_TO_MOTION = {
|
||||
"w": "forward",
|
||||
"a": "left",
|
||||
"d": "right",
|
||||
"s": "backward",
|
||||
"j": "left_rot",
|
||||
"l": "right_rot",
|
||||
"i": "up_rot",
|
||||
"k": "down_rot",
|
||||
}
|
||||
_TRANSLATION_BASE_UNIT = 1.0
|
||||
_ROTATION_BASE_UNIT = 10.0
|
||||
_INTRINSIC_ROW = [0.8, 0.5, 0.5, 0.5]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DreamXCamera:
|
||||
fx: float
|
||||
fy: float
|
||||
cx: float
|
||||
cy: float
|
||||
w2c_mat: np.ndarray
|
||||
|
||||
@property
|
||||
def c2w_mat(self) -> np.ndarray:
|
||||
return np.linalg.inv(self.w2c_mat)
|
||||
|
||||
@classmethod
|
||||
def from_pose_row(cls, row: list[float]) -> DreamXCamera:
|
||||
w2c_mat = np.eye(4, dtype=np.float64)
|
||||
w2c_mat[:3, :] = np.asarray(row[7:], dtype=np.float64).reshape(3, 4)
|
||||
return cls(
|
||||
fx=float(row[1]),
|
||||
fy=float(row[2]),
|
||||
cx=float(row[3]),
|
||||
cy=float(row[4]),
|
||||
w2c_mat=w2c_mat,
|
||||
)
|
||||
|
||||
|
||||
def _translation_step(motion_type: str, current_pose: dict[str, np.ndarray], value: float, duration: int) -> np.ndarray:
|
||||
if motion_type in ("forward", "backward"):
|
||||
yaw = np.radians(current_pose["rotation"][1])
|
||||
pitch = np.radians(current_pose["rotation"][0])
|
||||
forward = np.array([-math.sin(yaw) * math.cos(pitch), math.sin(pitch), math.cos(yaw) * math.cos(pitch)])
|
||||
direction = 1 if motion_type == "forward" else -1
|
||||
return forward * value * direction / duration
|
||||
if motion_type in ("left", "right"):
|
||||
yaw = np.radians(current_pose["rotation"][1])
|
||||
right = np.array([math.cos(yaw), 0.0, math.sin(yaw)])
|
||||
direction = -1 if motion_type == "left" else 1
|
||||
return right * value * direction / duration
|
||||
return np.zeros(3)
|
||||
|
||||
|
||||
def _rotation_step(motion_type: str, value: float, duration: int) -> np.ndarray:
|
||||
if not motion_type.endswith("rot"):
|
||||
return np.zeros(3)
|
||||
axis = motion_type.split("_")[0]
|
||||
rotation = np.zeros(3)
|
||||
if axis == "left":
|
||||
rotation[1] = value
|
||||
elif axis == "right":
|
||||
rotation[1] = -value
|
||||
elif axis == "up":
|
||||
rotation[0] = -value
|
||||
elif axis == "down":
|
||||
rotation[0] = value
|
||||
return rotation / duration
|
||||
|
||||
|
||||
def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
|
||||
pitch, yaw, roll = np.radians(angles)
|
||||
cy = math.cos(yaw * 0.5)
|
||||
sy = math.sin(yaw * 0.5)
|
||||
cp = math.cos(pitch * 0.5)
|
||||
sp = math.sin(pitch * 0.5)
|
||||
cr = math.cos(roll * 0.5)
|
||||
sr = math.sin(roll * 0.5)
|
||||
return [
|
||||
cy * cp * cr + sy * sp * sr,
|
||||
cy * sp * cr + sy * cp * sr,
|
||||
sy * cp * cr - cy * sp * sr,
|
||||
cy * cp * sr - sy * sp * cr,
|
||||
]
|
||||
|
||||
|
||||
def _quaternion_to_rotation_matrix(quaternion: list[float]) -> np.ndarray:
|
||||
qw, qx, qy, qz = quaternion
|
||||
return np.array([
|
||||
[1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
|
||||
[2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
|
||||
[2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
|
||||
])
|
||||
|
||||
|
||||
def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float], duration: int) -> list[list[float]]:
|
||||
if len(action_seq) != len(action_speed_list):
|
||||
raise ValueError("action_seq and action_speed_list must have the same length")
|
||||
|
||||
positions: list[np.ndarray] = []
|
||||
rotations: list[np.ndarray] = []
|
||||
current_pose = {
|
||||
"position": np.array([0.0, 0.0, 0.0]),
|
||||
"rotation": np.array([0.0, 0.0, 0.0]),
|
||||
}
|
||||
|
||||
for action_id, speed in zip(action_seq, action_speed_list, strict=True):
|
||||
motion_types = [_ACTION_TO_MOTION[key] for key in list(action_id)]
|
||||
translation_step = np.zeros(3)
|
||||
rotation_step = np.zeros(3)
|
||||
for motion_type in motion_types:
|
||||
translation_step += _translation_step(motion_type, current_pose,
|
||||
float(speed) * _TRANSLATION_BASE_UNIT, duration)
|
||||
rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
|
||||
|
||||
segment_positions = []
|
||||
segment_rotations = []
|
||||
for index in range(1, duration + 1):
|
||||
segment_positions.append(current_pose["position"] + translation_step * index)
|
||||
segment_rotations.append(current_pose["rotation"] + rotation_step * index)
|
||||
current_pose["position"] = segment_positions[-1].copy()
|
||||
current_pose["rotation"] = segment_rotations[-1].copy()
|
||||
positions.extend(segment_positions)
|
||||
rotations.extend(segment_rotations)
|
||||
|
||||
rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] +
|
||||
[1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
|
||||
for index, (position, rotation) in enumerate(zip(positions, rotations, strict=False)):
|
||||
rotation_matrix = _quaternion_to_rotation_matrix(_euler_to_quaternion(rotation))
|
||||
translation = -rotation_matrix @ position
|
||||
extrinsic = np.hstack([rotation_matrix, translation.reshape(3, 1)])
|
||||
rows.append([float(index)] + _INTRINSIC_ROW + [0.0, 0.0] + extrinsic.flatten().tolist())
|
||||
return rows
|
||||
|
||||
|
||||
def _interpolate_camera_poses(
|
||||
cameras: list[DreamXCamera],
|
||||
src_indices: np.ndarray,
|
||||
tgt_indices: np.ndarray,
|
||||
) -> list[DreamXCamera]:
|
||||
if len(cameras) <= 1:
|
||||
return [cameras[0]] * len(tgt_indices) if cameras else []
|
||||
src_rot_mat = np.array([camera.w2c_mat[:3, :3] for camera in cameras])
|
||||
src_trans_vec = np.array([camera.w2c_mat[:3, 3] for camera in cameras])
|
||||
|
||||
dets = np.linalg.det(src_rot_mat)
|
||||
flip_handedness = dets.size > 0 and np.median(dets) < 0.0
|
||||
if flip_handedness:
|
||||
flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
|
||||
src_rot_mat = src_rot_mat @ flip_mat
|
||||
|
||||
trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False,
|
||||
fill_value="extrapolate")(tgt_indices)
|
||||
quats = Rotation.from_matrix(src_rot_mat).as_quat().copy()
|
||||
for index in range(1, len(quats)):
|
||||
if np.dot(quats[index], quats[index - 1]) < 0:
|
||||
quats[index] = -quats[index]
|
||||
rot = Slerp(src_indices, Rotation.from_quat(quats))(tgt_indices).as_matrix()
|
||||
if flip_handedness:
|
||||
rot = rot @ flip_mat
|
||||
|
||||
ref = cameras[0]
|
||||
result = []
|
||||
for index in range(len(tgt_indices)):
|
||||
w2c_mat = np.eye(4, dtype=np.float64)
|
||||
w2c_mat[:3, :] = np.hstack([rot[index], trans[index].reshape(3, 1)])
|
||||
result.append(DreamXCamera(ref.fx, ref.fy, ref.cx, ref.cy, w2c_mat))
|
||||
return result
|
||||
|
||||
|
||||
def _relative_c2w_poses(cameras: list[DreamXCamera]) -> np.ndarray:
|
||||
abs_w2cs = [camera.w2c_mat for camera in cameras]
|
||||
abs_c2ws = [camera.c2w_mat for camera in cameras]
|
||||
target_cam_c2w = np.eye(4, dtype=np.float64)
|
||||
abs2rel = target_cam_c2w @ abs_w2cs[0]
|
||||
poses = [target_cam_c2w] + [abs2rel @ c2w for c2w in abs_c2ws[1:]]
|
||||
return np.asarray(poses, dtype=np.float32)
|
||||
|
||||
|
||||
def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
|
||||
rotation_inv = transforms[..., :3, :3].transpose(-1, -2)
|
||||
output = torch.zeros_like(transforms)
|
||||
output[..., :3, :3] = rotation_inv
|
||||
output[..., :3, 3] = -torch.einsum("...ij,...j->...i", rotation_inv, transforms[..., :3, 3])
|
||||
output[..., 3, 3] = 1.0
|
||||
return output
|
||||
|
||||
|
||||
def build_dreamx_camera_condition(
|
||||
action_seq: list[str],
|
||||
action_speed_list: list[float],
|
||||
*,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
device: torch.device | str = "cpu",
|
||||
) -> dict[str, torch.Tensor]:
|
||||
del height, width # DreamX-World-5B-Cam uses fixed normalized intrinsics.
|
||||
duration = math.ceil(num_frames / len(action_seq))
|
||||
rows = _pose_rows_from_actions(action_seq, action_speed_list, duration)[:num_frames]
|
||||
cameras = [DreamXCamera.from_pose_row(row) for row in rows]
|
||||
|
||||
latent_frame_count = 1 + (len(cameras) - 1) // 4
|
||||
src_indices = np.arange(len(cameras), dtype=np.float64)
|
||||
tgt_indices = np.linspace(0, len(cameras) - 1, latent_frame_count)
|
||||
cameras = _interpolate_camera_poses(cameras, src_indices, tgt_indices)
|
||||
|
||||
c2ws = torch.as_tensor(_relative_c2w_poses(cameras), dtype=dtype, device=device)
|
||||
viewmats = _invert_se3(c2ws)
|
||||
|
||||
intrinsics = torch.zeros((latent_frame_count, 3, 3), dtype=dtype, device=device)
|
||||
intrinsics[:, 0, 0] = 969.6969696969696 / (960.0 * 2)
|
||||
intrinsics[:, 1, 1] = 969.6969696969696 / (540.0 * 2)
|
||||
intrinsics[:, 2, 2] = 1.0
|
||||
return {"viewmats": viewmats, "K": intrinsics}
|
||||
@@ -1,20 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Compatibility exports for DreamX-World pipeline configs."""
|
||||
|
||||
from fastvideo.configs.pipelines.dreamx_world import (
|
||||
DreamXWorld5BARPipelineConfig,
|
||||
DreamXWorld5BCamPipelineConfig,
|
||||
make_dreamx_world_5b_ar_dit_config,
|
||||
make_dreamx_world_5b_cam_dit_config,
|
||||
make_dreamx_world_5b_cam_text_encoder_config,
|
||||
make_dreamx_world_5b_cam_vae_config,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DreamXWorld5BARPipelineConfig",
|
||||
"DreamXWorld5BCamPipelineConfig",
|
||||
"make_dreamx_world_5b_ar_dit_config",
|
||||
"make_dreamx_world_5b_cam_dit_config",
|
||||
"make_dreamx_world_5b_cam_text_encoder_config",
|
||||
"make_dreamx_world_5b_cam_vae_config",
|
||||
]
|
||||
@@ -1,67 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""DreamX-World-5B autoregressive pipeline entrypoint."""
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
||||
from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
|
||||
from fastvideo.pipelines.basic.dreamx_world.stages import (
|
||||
DreamXWorldCameraConditioningStage,
|
||||
DreamXWorldImageVAEEncodingStage,
|
||||
)
|
||||
from fastvideo.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DreamXWorldARPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
"""DreamX-World-5B autoregressive causal camera pipeline."""
|
||||
|
||||
_required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
|
||||
pipeline_config_cls = DreamXWorld5BARPipelineConfig
|
||||
sampling_params_cls = SamplingParam
|
||||
|
||||
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
|
||||
self.modules["scheduler"].set_timesteps(1000)
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
|
||||
self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
|
||||
self.add_stage(stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
))
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
self.add_stage(stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
|
||||
self.add_stage(stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
))
|
||||
self.add_stage(stage_name="image_latent_preparation_stage",
|
||||
stage=DreamXWorldImageVAEEncodingStage(vae=self.get_module("vae")))
|
||||
self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
|
||||
self.add_stage(stage_name="denoising_stage",
|
||||
stage=DreamXWorldARCausalDenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
vae=self.get_module("vae"),
|
||||
pipeline=self,
|
||||
))
|
||||
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
|
||||
logger.info("DreamXWorldARPipeline initialized with autoregressive causal denoising")
|
||||
|
||||
|
||||
EntryClass = DreamXWorldARPipeline
|
||||
@@ -1,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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user