diff --git a/.agents/skills/add-model-02-parity/SKILL.md b/.agents/skills/add-model-02-parity/SKILL.md index 5a96f34a5..a11341759 100644 --- a/.agents/skills/add-model-02-parity/SKILL.md +++ b/.agents/skills/add-model-02-parity/SKILL.md @@ -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` diff --git a/.agents/skills/add-model-02-parity/templates/component_parity_test.py b/.agents/skills/add-model-02-parity/templates/component_parity_test.py index 8d24e292a..c1137f298 100644 --- a/.agents/skills/add-model-02-parity/templates/component_parity_test.py +++ b/.agents/skills/add-model-02-parity/templates/component_parity_test.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. diff --git a/.agents/skills/add-model-03-port-dit/SKILL.md b/.agents/skills/add-model-03-port-dit/SKILL.md index 028f6b346..4d185bda6 100644 --- a/.agents/skills/add-model-03-port-dit/SKILL.md +++ b/.agents/skills/add-model-03-port-dit/SKILL.md @@ -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 diff --git a/.agents/skills/add-model-04-port-vae/SKILL.md b/.agents/skills/add-model-04-port-vae/SKILL.md index 16c50bc3f..14e8047f9 100644 --- a/.agents/skills/add-model-04-port-vae/SKILL.md +++ b/.agents/skills/add-model-04-port-vae/SKILL.md @@ -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 diff --git a/.agents/skills/add-model-07-conversion/SKILL.md b/.agents/skills/add-model-07-conversion/SKILL.md index 727b80f06..73d8f201b 100644 --- a/.agents/skills/add-model-07-conversion/SKILL.md +++ b/.agents/skills/add-model-07-conversion/SKILL.md @@ -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. diff --git a/.agents/skills/add-model-09-pipeline/SKILL.md b/.agents/skills/add-model-09-pipeline/SKILL.md index e91257b0a..1f3e55557 100644 --- a/.agents/skills/add-model-09-pipeline/SKILL.md +++ b/.agents/skills/add-model-09-pipeline/SKILL.md @@ -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. diff --git a/.agents/skills/add-model/SKILL.md b/.agents/skills/add-model/SKILL.md index d7f83de23..674967140 100644 --- a/.agents/skills/add-model/SKILL.md +++ b/.agents/skills/add-model/SKILL.md @@ -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 diff --git a/.agents/skills/seed-ssim-references/SKILL.md b/.agents/skills/seed-ssim-references/SKILL.md index ca2266ae1..851494605 100644 --- a/.agents/skills/seed-ssim-references/SKILL.md +++ b/.agents/skills/seed-ssim-references/SKILL.md @@ -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`, diff --git a/.github/scripts/plan_merge_ci.py b/.github/scripts/plan_merge_ci.py index db79723ce..e00c4199d 100644 --- a/.github/scripts/plan_merge_ci.py +++ b/.github/scripts/plan_merge_ci.py @@ -130,11 +130,6 @@ class FamilyCoverage: FAMILY_COVERAGE = ( - FamilyCoverage( - re.compile(r"(^|[/_.-])dreamx(_world)?([/_.-]|$)"), - ("test_dreamx.py", ), - ("test_dreamx_world_similarity.py", ), - ), FamilyCoverage( re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"), ("test_flux2_klein.py", ), @@ -145,21 +140,11 @@ FAMILY_COVERAGE = ( ("test_flux.py", ), ("test_flux_t2i_similarity.py", ), ), - FamilyCoverage( - re.compile(r"(^|[/_.-])(hunyuan)?gamecraft([/_.-]|$)"), - ("test_gamecraft.py", ), - ("test_gamecraft_similarity.py", ), - ), FamilyCoverage( re.compile(r"(^|[/_.-])gen3c([/_.-]|$)"), ("test_gen3c.py", ), ("test_gen3c_similarity.py", ), ), - FamilyCoverage( - re.compile(r"(^|[/_.-])glm[_-]?image([/_.-]|$)"), - ("test_glm_image.py", ), - ("test_glm_image_similarity.py", ), - ), FamilyCoverage( re.compile(r"(^|[/_.-])kandinsky[_-]?5([/_.-]|$)"), ("test_kandinsky5.py", ), @@ -170,11 +155,6 @@ FAMILY_COVERAGE = ( ("test_lingbot.py", ), ("test_lingbot_similarity.py", ), ), - FamilyCoverage( - re.compile(r"(^|[/_.-])longcat([/_.-]|$)"), - ("test_longcat.py", ), - ("test_longcat_similarity.py", ), - ), FamilyCoverage( re.compile(r"(^|[/_.-])ltx[_-]?2([/_.-]|$)"), ("test_ltx2.py", ), @@ -205,11 +185,6 @@ FAMILY_COVERAGE = ( ("test_stable_audio.py", ), ("test_stable_audio_similarity.py", ), ), - FamilyCoverage( - re.compile(r"(^|[/_.-])turbo(diffusion)?([/_.-]|$)"), - (), - ("test_turbodiffusion_similarity.py", ), - ), FamilyCoverage( re.compile(r"(^|[/_.-])wan(video|vae)?([/_.-]|$)"), ("test_wan_t2v.py", "test_wan_vae.py", "test_wan_causal.py", "test_wan_denoising.py"), @@ -219,11 +194,6 @@ FAMILY_COVERAGE = ( "test_wan_t2v_similarity.py", ), ), - FamilyCoverage( - re.compile(r"(^|[/_.-])z[_-]?image([/_.-]|$)"), - ("test_zimage.py", ), - ("test_zimage_similarity.py", ), - ), ) diff --git a/.gitignore b/.gitignore index b65ab4377..8922ac204 100644 --- a/.gitignore +++ b/.gitignore @@ -36,7 +36,6 @@ env *.log weights/ logs/ -/Z-Image/ official_weights/ converted_weights/ diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 63a0cdbf5..0c50bca20 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -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", diff --git a/docs/assets/logos/SOURCES.md b/docs/assets/logos/SOURCES.md index 21247b395..17b3907f4 100644 --- a/docs/assets/logos/SOURCES.md +++ b/docs/assets/logos/SOURCES.md @@ -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. diff --git a/docs/assets/logos/meituan-longcat.webp b/docs/assets/logos/meituan-longcat.webp deleted file mode 100644 index f12949441..000000000 Binary files a/docs/assets/logos/meituan-longcat.webp and /dev/null differ diff --git a/docs/assets/logos/tongyi.webp b/docs/assets/logos/tongyi.webp deleted file mode 100644 index efd767aa1..000000000 Binary files a/docs/assets/logos/tongyi.webp and /dev/null differ diff --git a/docs/assets/logos/zai.webp b/docs/assets/logos/zai.webp deleted file mode 100644 index 93ceda24a..000000000 Binary files a/docs/assets/logos/zai.webp and /dev/null differ diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 3b85bfe18..6043d0b02 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -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 | diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index 4d9c8c0fb..89ad17313 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
+
All model families
diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index 160c56ca5..bcb59a2fe 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
+
All model families
diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md deleted file mode 100644 index 93f8b5cf2..000000000 --- a/docs/cookbook/glm-image.md +++ /dev/null @@ -1,140 +0,0 @@ ---- -hide: -- toc ---- - -# GLM-Image recipes - -
-
- All model families -
- -
-

Maintained family · Inference

-

GLM-Image inference recipes

-

GLM-Image from Z.ai supports both text-to-image generation and instruction-based image editing, each with a maintained example.

-
-
-
- Inference live - Distillation planned - Fine-tuning planned - Training planned - Evaluation planned - Optimization planned - Deployment planned -
-
- - -
-
-

Pick a recipe and runtime

-

Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.

-
- -
-
-
-
- Recipe - Task and checkpoint -
-
- -
-
- -
-
- Runtime - Maintained paths only -
-
- -
-
- -

Loading recipe details...

-

Exact device and memory details appear only when a recorded run supports them.

- -
- Reading recipe evidence... -
-
- -
-
-

Loading...

-
- Maintained - Source-backed - Source config -
-
- -
-
Model
Loading...
-
Workload
Loading...
-
Source configuration
Loading...
-
Expected output
Loading...
-
- -
-
- Terminal -
- 
Loading...
-
- -
- Open example source - View model card -
-

-
-
- - -
-
- -
- Setup -
-

The generated commands expect a local clone:

- 
git clone https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo
-

Use Configuration for supported Python and CLI settings, Optimizations for attention and memory tradeoffs, and the support matrix for the supported model and optimization surface.

-
-
- -
- Troubleshooting -
-
    -
  • The editing example reads assets/images/couple.jpg from the repository root, so run it from a repo checkout rather than an arbitrary working directory.
    • -
  • Output paths default under image_output/; pass --output to change them.
    • -
-
-
- -
- Evidence status -
-

All recipes on this page are Source-backed: 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 Unknown and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.

-
-
diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index dede8af94..52c11cd56 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
+
All model families
diff --git a/docs/cookbook/index.md b/docs/cookbook/index.md index 6b4b43497..22e735066 100644 --- a/docs/cookbook/index.md +++ b/docs/cookbook/index.md @@ -163,50 +163,6 @@ hide: - - - - - - - - - - - LongCatT2V, I2V, optional refine - 2 recipes - -
    -
  • T2V
    • -
  • I2V
    • -
- - - - - - - - - - - - - TurboDiffusionAccelerated Wan profiles - 3 recipes - -
    -
  • T2V
    • -
  • I2V
    • -
- -
@@ -236,49 +192,6 @@ hide: - - - - - - - - - - GLM-ImageGenerate and edit - 2 recipes - -
    -
  • T2I
    • -
  • Edit
    • -
- - - - - - - - - - - - - Z-ImageTurbo text to image - 1 recipe - -
    -
  • T2I
    • -
- - -
-
- - - - - - - - GameCraftGame world generation - Page planned - - -
-
@@ -449,20 +348,6 @@ hide:
-
- - - - - - - - DreamXWorld generation - Page planned - - -
-
diff --git a/docs/cookbook/kandinsky5.md b/docs/cookbook/kandinsky5.md index d899dd977..2389f1bd6 100644 --- a/docs/cookbook/kandinsky5.md +++ b/docs/cookbook/kandinsky5.md @@ -5,7 +5,7 @@ hide: # Kandinsky 5 recipes -
+
All model families
diff --git a/docs/cookbook/kandinsky6.md b/docs/cookbook/kandinsky6.md index fc9f32dff..4e4c0aaf4 100644 --- a/docs/cookbook/kandinsky6.md +++ b/docs/cookbook/kandinsky6.md @@ -5,7 +5,7 @@ hide: # Kandinsky 6 recipes -
+
All model families
diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md deleted file mode 100644 index 28f303deb..000000000 --- a/docs/cookbook/longcat.md +++ /dev/null @@ -1,140 +0,0 @@ ---- -hide: -- toc ---- - -# LongCat recipes - -
-
- All model families -
- -
-

Maintained family · Inference

-

LongCat inference recipes

-

LongCat Video from Meituan covers text-to-video and image-to-video, and its maintained examples chain optional distilled and 720p refinement passes.

-
-
-
- Inference live - Distillation planned - Fine-tuning planned - Training planned - Evaluation planned - Optimization planned - Deployment planned -
-
- - -
-
-

Pick a recipe and runtime

-

Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.

-
- -
-
-
-
- Recipe - Task and checkpoint -
-
- -
-
- -
-
- Runtime - Maintained paths only -
-
- -
-
- -

Loading recipe details...

-

Exact device and memory details appear only when a recorded run supports them.

- -
- Reading recipe evidence... -
-
- -
-
-

Loading...

-
- Maintained - Source-backed - Source config -
-
- -
-
Model
Loading...
-
Workload
Loading...
-
Source configuration
Loading...
-
Expected output
Loading...
-
- -
-
- Terminal -
- 
Loading...
-
- -
- Open example source - View model card -
-

-
-
- - -
-
- -
- Setup -
-

The generated commands expect a local clone:

- 
git clone https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo
-

Use Configuration for supported Python and CLI settings, Optimizations for attention and memory tradeoffs, and the support matrix for the supported model and optimization surface.

-
-
- -
- Troubleshooting -
-
    -
  • Each LongCat script runs multiple passes (basic, distilled, refine); total runtime scales accordingly, and every pass prints its own output directory.
    • -
  • Out of memory: the sources already enable VAE and text-encoder CPU offload; further options are covered in Offloading.
    • -
-
-
- -
- Evidence status -
-

All recipes on this page are Source-backed: 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 Unknown and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.

-
-
diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index f751661c9..d8878be7b 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
+
All model families
diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index b6d4e16cd..10f0eb4ca 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
+
All model families
diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 0ccb0b3bb..265a508d6 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -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. -
+
All model families
diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index 544241d44..8cd584cfa 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
+
All model families
diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index 5fa15b299..64f52ce32 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
+
All model families
diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index 845a30594..9bd005ea0 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
+
All model families
diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md deleted file mode 100644 index df3dac64f..000000000 --- a/docs/cookbook/turbodiffusion.md +++ /dev/null @@ -1,140 +0,0 @@ ---- -hide: -- toc ---- - -# TurboDiffusion recipes - -
-
- All model families -
- -
-

Maintained family · Inference

-

TurboDiffusion inference recipes

-

TurboDiffusion profiles accelerate Wan checkpoints with step-distilled sampling and the SLA attention backend. These recipes follow the registry's `turbodiffusion` model family.

-
-
-
- Inference live - Distillation planned - Fine-tuning planned - Training planned - Evaluation planned - Optimization planned - Deployment planned -
-
- - -
-
-

Pick a recipe and runtime

-

Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.

-
- -
-
-
-
- Recipe - Task and checkpoint -
-
- -
-
- -
-
- Runtime - Maintained paths only -
-
- -
-
- -

Loading recipe details...

-

Exact device and memory details appear only when a recorded run supports them.

- -
- Reading recipe evidence... -
-
- -
-
-

Loading...

-
- Maintained - Source-backed - Source config -
-
- -
-
Model
Loading...
-
Workload
Loading...
-
Source configuration
Loading...
-
Expected output
Loading...
-
- -
-
- Terminal -
- 
Loading...
-
- -
- Open example source - View model card -
-

-
-
- - -
-
- -
- Setup -
-

The generated commands expect a local clone:

- 
git clone https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo
-

Use Configuration for supported Python and CLI settings, Optimizations for attention and memory tradeoffs, and the support matrix for the supported model and optimization surface.

-
-
- -
- Troubleshooting -
-
    -
  • TurboDiffusion paths load community-published loayrashid/TurboWan* checkpoints; availability is governed by those repos.
    • -
  • The SLA attention backend used by the I2V recipe is selected inside the example source; do not combine it with another FASTVIDEO_ATTENTION_BACKEND override in the same shell.
    • -
-
-
- -
- Evidence status -
-

All recipes on this page are Source-backed: 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 Unknown and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.

-
-
diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index 7139d02fb..d23766ab3 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
+
All model families
diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md deleted file mode 100644 index 4db21f991..000000000 --- a/docs/cookbook/z-image.md +++ /dev/null @@ -1,140 +0,0 @@ ---- -hide: -- toc ---- - -# Z-Image recipes - -
-
- All model families -
- -
-

Maintained family · Inference

-

Z-Image inference recipes

-

Z-Image Turbo from Tongyi MAI is a fast text-to-image model. The maintained example runs it on a single GPU.

-
-
-
- Inference live - Distillation planned - Fine-tuning planned - Training planned - Evaluation planned - Optimization planned - Deployment planned -
-
- - -
-
-

Pick a recipe and runtime

-

Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.

-
- -
-
-
-
- Recipe - Task and checkpoint -
-
- -
-
- -
-
- Runtime - Maintained paths only -
-
- -
-
- -

Loading recipe details...

-

Exact device and memory details appear only when a recorded run supports them.

- -
- Reading recipe evidence... -
-
- -
-
-

Loading...

-
- Maintained - Source-backed - Source config -
-
- -
-
Model
Loading...
-
Workload
Loading...
-
Source configuration
Loading...
-
Expected output
Loading...
-
- -
-
- Terminal -
- 
Loading...
-
- -
- Open example source - View model card -
-

-
-
- - -
-
- -
- Setup -
-

The generated commands expect a local clone:

- 
git clone https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo
-

Use Configuration for supported Python and CLI settings, Optimizations for attention and memory tradeoffs, and the support matrix for the supported model and optimization surface.

-
-
- -
- Troubleshooting -
-
    -
  • Output defaults to outputs/zimage/zimage_turbo.png; pass --output to redirect.
    • -
  • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
    • -
-
-
- -
- Evidence status -
-

All recipes on this page are Source-backed: 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 Unknown and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.

-
-
diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index d45543555..7f8f01816 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -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 diff --git a/docs/generate_examples.py b/docs/generate_examples.py index bc2d9dd63..21c49f372 100644 --- a/docs/generate_examples.py +++ b/docs/generate_examples.py @@ -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", diff --git a/docs/inference/architecture.md b/docs/inference/architecture.md index c86038750..766f69ce7 100644 --- a/docs/inference/architecture.md +++ b/docs/inference/architecture.md @@ -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`. diff --git a/docs/inference/offloading.md b/docs/inference/offloading.md index 904581271..ca27b88d2 100644 --- a/docs/inference/offloading.md +++ b/docs/inference/offloading.md @@ -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 diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index c89be5ce8..6a713a01c 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -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`
`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`
`FastVideo/LTX2.3-Distilled-Diffusers`
`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`
`FastVideo/LTX2.3-base`
`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`
`FastVideo/LTX2-base`
`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`
`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` | T2V | — | | wan | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | I2V | — | @@ -92,7 +82,6 @@ column links a runnable script in `examples/inference/basic/` where one exists. | wan | `wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers` | T2V | [basic_self_forcing_causal.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal.py) | | wan | `rand0nmr/SFWan2.2-T2V-A14B-Diffusers` | T2V | [basic_self_forcing_causal_wan2_2_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py) | | wan | `FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers` | I2V | [basic_self_forcing_causal_wan2_2_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py) | -| zimage | `Tongyi-MAI/Z-Image-Turbo` | T2I | [basic_zimage.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_zimage.py) | **Note (stable_audio)**: the Stable Audio Open pipelines generate audio (`StableAudioT2AConfig` / `StableAudioOpenSmallConfig`); they are registered @@ -168,8 +157,6 @@ optimizations: absence means **untested**, not incompatible. | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | | FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | -| DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P
720P | ❌ | ❌ | ✅ | ⭕ | ⭕ | | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P
720P | ❌ | ❌ | ✅ | ⭕ | ⭕ | @@ -179,10 +166,6 @@ optimizations: absence means **untested**, not incompatible. | Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ | | Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ | ⭕ | | Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ | -| TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P
720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| LongCat T2V 13.6B | `FastVideo/LongCat-Video-T2V-Diffusers` | 480P
720P | ❌ | ❌ | ❌ | ⭕ | ✅ | | Matrix Game 2.0 Base Distilled | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | Matrix Game 2.0 GTA Distilled | `FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | Matrix Game 2.0 TempleRun Distilled | `FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | @@ -251,11 +234,6 @@ listed under [Special requirements](#special-requirements). https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete - STA currently requires Hopper GPUs (H100s). -### TurboWan2.1 (TurboDiffusion) -- Uses TurboDiffusionPipeline with RCM scheduler for 1-4 step generation -- Requires SLA attention backend: `export FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN` -- Uses `guidance_scale=1.0` (no classifier-free guidance) - ### Matrix Game 2.0 - Image-to-video game world models with keyboard/mouse control input - Three variants available: Base (universal), GTA, and TempleRun diff --git a/examples/inference/basic/basic_dreamx_world.py b/examples/inference/basic/basic_dreamx_world.py deleted file mode 100644 index 9af4a02ed..000000000 --- a/examples/inference/basic/basic_dreamx_world.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_gamecraft.py b/examples/inference/basic/basic_gamecraft.py deleted file mode 100644 index c086d86e4..000000000 --- a/examples/inference/basic/basic_gamecraft.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_glm_image.py b/examples/inference/basic/basic_glm_image.py deleted file mode 100644 index 534218909..000000000 --- a/examples/inference/basic/basic_glm_image.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_longcat_i2v.py b/examples/inference/basic/basic_longcat_i2v.py deleted file mode 100644 index 001885247..000000000 --- a/examples/inference/basic/basic_longcat_i2v.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_longcat_t2v.py b/examples/inference/basic/basic_longcat_t2v.py deleted file mode 100644 index 58c87dc3a..000000000 --- a/examples/inference/basic/basic_longcat_t2v.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_longcat_vc.py b/examples/inference/basic/basic_longcat_vc.py deleted file mode 100644 index 36d5f0935..000000000 --- a/examples/inference/basic/basic_longcat_vc.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_turbodiffusion.py b/examples/inference/basic/basic_turbodiffusion.py deleted file mode 100644 index 6e85a734e..000000000 --- a/examples/inference/basic/basic_turbodiffusion.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_turbodiffusion_14b.py b/examples/inference/basic/basic_turbodiffusion_14b.py deleted file mode 100644 index ebaa39c3b..000000000 --- a/examples/inference/basic/basic_turbodiffusion_14b.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_turbodiffusion_i2v.py b/examples/inference/basic/basic_turbodiffusion_i2v.py deleted file mode 100644 index b28092f21..000000000 --- a/examples/inference/basic/basic_turbodiffusion_i2v.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/basic_zimage.py b/examples/inference/basic/basic_zimage.py deleted file mode 100644 index 732894163..000000000 --- a/examples/inference/basic/basic_zimage.py +++ /dev/null @@ -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() diff --git a/examples/inference/basic/edit_glm_image.py b/examples/inference/basic/edit_glm_image.py deleted file mode 100644 index 9fb95c75b..000000000 --- a/examples/inference/basic/edit_glm_image.py +++ /dev/null @@ -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() diff --git a/examples/train/configs/fine_tuning/longcat/t2v.yaml b/examples/train/configs/fine_tuning/longcat/t2v.yaml deleted file mode 100644 index 58aba7ae4..000000000 --- a/examples/train/configs/fine_tuning/longcat/t2v.yaml +++ /dev/null @@ -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 diff --git a/examples/train/configs/fine_tuning/wan/turbo_t2v_lora_sla.yaml b/examples/train/configs/fine_tuning/wan/turbo_t2v_lora_sla.yaml deleted file mode 100644 index 3f9daf02c..000000000 --- a/examples/train/configs/fine_tuning/wan/turbo_t2v_lora_sla.yaml +++ /dev/null @@ -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 diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 8dfb430e4..f771c6f3d 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -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, diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 521039785..1a8c96f98 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -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 diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py index ca3853992..6d03f1f9f 100644 --- a/fastvideo/attention/utils/flash_attn_cute.py +++ b/fastvideo/attention/utils/flash_attn_cute.py @@ -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)) diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 72b539f13..4ce594543 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -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" ] diff --git a/fastvideo/configs/models/dits/dreamx_world.py b/fastvideo/configs/models/dits/dreamx_world.py deleted file mode 100644 index bbd5177bb..000000000 --- a/fastvideo/configs/models/dits/dreamx_world.py +++ /dev/null @@ -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" diff --git a/fastvideo/configs/models/dits/glm_image.py b/fastvideo/configs/models/dits/glm_image.py deleted file mode 100644 index 8e83ab439..000000000 --- a/fastvideo/configs/models/dits/glm_image.py +++ /dev/null @@ -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" diff --git a/fastvideo/configs/models/dits/hunyuangamecraft.py b/fastvideo/configs/models/dits/hunyuangamecraft.py deleted file mode 100644 index 2cbeae22f..000000000 --- a/fastvideo/configs/models/dits/hunyuangamecraft.py +++ /dev/null @@ -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" diff --git a/fastvideo/configs/models/dits/longcat.py b/fastvideo/configs/models/dits/longcat.py deleted file mode 100644 index 735b84eea..000000000 --- a/fastvideo/configs/models/dits/longcat.py +++ /dev/null @@ -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" diff --git a/fastvideo/configs/models/dits/zimage.py b/fastvideo/configs/models/dits/zimage.py deleted file mode 100644 index ce9bf7848..000000000 --- a/fastvideo/configs/models/dits/zimage.py +++ /dev/null @@ -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" diff --git a/fastvideo/configs/models/encoders/base.py b/fastvideo/configs/models/encoders/base.py index 09ee57067..5dd4a531f 100644 --- a/fastvideo/configs/models/encoders/base.py +++ b/fastvideo/configs/models/encoders/base.py @@ -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 = { diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index 13210e3a2..acc5ca9c0 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -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", ] diff --git a/fastvideo/configs/models/vaes/gamecraftvae.py b/fastvideo/configs/models/vaes/gamecraftvae.py deleted file mode 100644 index 2780583ce..000000000 --- a/fastvideo/configs/models/vaes/gamecraftvae.py +++ /dev/null @@ -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) diff --git a/fastvideo/configs/models/vaes/glm_image.py b/fastvideo/configs/models/vaes/glm_image.py deleted file mode 100644 index a7bd1becd..000000000 --- a/fastvideo/configs/models/vaes/glm_image.py +++ /dev/null @@ -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 diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 34461ebf6..cf4a3e540 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -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" diff --git a/fastvideo/configs/pipelines/dreamx_world.py b/fastvideo/configs/pipelines/dreamx_world.py deleted file mode 100644 index 91158255f..000000000 --- a/fastvideo/configs/pipelines/dreamx_world.py +++ /dev/null @@ -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 diff --git a/fastvideo/configs/pipelines/glm_image.py b/fastvideo/configs/pipelines/glm_image.py deleted file mode 100644 index 7b2f46279..000000000 --- a/fastvideo/configs/pipelines/glm_image.py +++ /dev/null @@ -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 diff --git a/fastvideo/configs/pipelines/hunyuangamecraft.py b/fastvideo/configs/pipelines/hunyuangamecraft.py deleted file mode 100644 index d67846fe3..000000000 --- a/fastvideo/configs/pipelines/hunyuangamecraft.py +++ /dev/null @@ -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 diff --git a/fastvideo/configs/pipelines/longcat.py b/fastvideo/configs/pipelines/longcat.py deleted file mode 100644 index a8d778b98..000000000 --- a/fastvideo/configs/pipelines/longcat.py +++ /dev/null @@ -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}'") diff --git a/fastvideo/configs/pipelines/turbodiffusion.py b/fastvideo/configs/pipelines/turbodiffusion.py deleted file mode 100644 index ae073f41b..000000000 --- a/fastvideo/configs/pipelines/turbodiffusion.py +++ /dev/null @@ -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 diff --git a/fastvideo/configs/pipelines/zimage.py b/fastvideo/configs/pipelines/zimage.py deleted file mode 100644 index c4395fac8..000000000 --- a/fastvideo/configs/pipelines/zimage.py +++ /dev/null @@ -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 diff --git a/fastvideo/envs.py b/fastvideo/envs.py index e9d4c3be8..60ce84699 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -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", } diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index d57a179a7..4f54d6b07 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -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}) diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index 85effdce7..cd8cf59e5 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -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. diff --git a/fastvideo/models/camera/__init__.py b/fastvideo/models/camera/__init__.py deleted file mode 100644 index cb032e447..000000000 --- a/fastvideo/models/camera/__init__.py +++ /dev/null @@ -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"] diff --git a/fastvideo/models/camera/trajectory.py b/fastvideo/models/camera/trajectory.py deleted file mode 100644 index 796d363b8..000000000 --- a/fastvideo/models/camera/trajectory.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/dits/dreamx_world.py b/fastvideo/models/dits/dreamx_world.py deleted file mode 100644 index ba1d3ca8d..000000000 --- a/fastvideo/models/dits/dreamx_world.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/dits/dreamx_world_ar.py b/fastvideo/models/dits/dreamx_world_ar.py deleted file mode 100644 index cc001ff40..000000000 --- a/fastvideo/models/dits/dreamx_world_ar.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/dits/glm_image.py b/fastvideo/models/dits/glm_image.py deleted file mode 100644 index 6471713d2..000000000 --- a/fastvideo/models/dits/glm_image.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/dits/hunyuangamecraft.py b/fastvideo/models/dits/hunyuangamecraft.py deleted file mode 100644 index ec2def561..000000000 --- a/fastvideo/models/dits/hunyuangamecraft.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/dits/longcat.py b/fastvideo/models/dits/longcat.py deleted file mode 100644 index 603161b75..000000000 --- a/fastvideo/models/dits/longcat.py +++ /dev/null @@ -1,1141 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Native LongCat Video DiT implementation using FastVideo conventions. -""" - -from typing import Any -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange - -from fastvideo.configs.models.dits import LongCatVideoConfig -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.layernorm import RMSNorm, FP32LayerNorm -from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.rotary_embedding_3d import RotaryPositionalEmbedding3D -from fastvideo.attention.layer import DistributedAttention, LocalAttention -from fastvideo.models.dits.base import BaseDiT -from fastvideo.platforms import AttentionBackendEnum -from fastvideo.third_party.longcat_video.block_sparse_attention.bsa_interface import flash_attn_bsa_3d - -# ============================================================================ -# Embeddings -# ============================================================================ - - -class PatchEmbed3D(nn.Module): - """ - 3D patch embedding using Conv3d. - """ - - def __init__( - self, - patch_size: tuple[int, int, int] = (1, 2, 2), - in_channels: int = 16, - embed_dim: int = 4096, - ): - super().__init__() - self.patch_size = patch_size - self.in_channels = in_channels - self.embed_dim = embed_dim - - self.proj = nn.Conv3d( - in_channels, - embed_dim, - kernel_size=patch_size, - stride=patch_size, - bias=True, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Args: - x: [B, C, T, H, W] - Returns: - [B, N, C] where N = (T/pt) * (H/ph) * (W/pw) - """ - # Padding if needed - _, _, T, H, W = x.shape - if W % self.patch_size[2] != 0: - x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2])) - if H % self.patch_size[1] != 0: - x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1])) - if T % self.patch_size[0] != 0: - x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - T % self.patch_size[0])) - - x = self.proj(x) # [B, C, T', H', W'] - x = x.flatten(2).transpose(1, 2) # [B, N, C] - return x - - -class TimestepEmbedder(nn.Module): - """ - Sinusoidal timestep embedding + MLP projection. - """ - - def __init__( - self, - frequency_embedding_size: int = 256, - adaln_tembed_dim: int = 512, - dtype: torch.dtype | None = None, - ): - super().__init__() - self.frequency_embedding_size = frequency_embedding_size - - # Use FastVideo's ReplicatedLinear - self.linear_1 = ReplicatedLinear( - frequency_embedding_size, - adaln_tembed_dim, - bias=True, - params_dtype=dtype, - ) - self.act = nn.SiLU() - self.linear_2 = ReplicatedLinear( - adaln_tembed_dim, - adaln_tembed_dim, - bias=True, - params_dtype=dtype, - ) - - @staticmethod - def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: - """ - Create sinusoidal timestep embeddings. - """ - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * - torch.arange(start=0, end=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, latent_shape: tuple | None = None) -> torch.Tensor: - """ - Args: - t: [B] or [B, T] timesteps - latent_shape: (T, H, W) for temporal expansion - Returns: - [B, T, C] - """ - # Sinusoidal embedding in FP32 - t_freq = self.timestep_embedding(t.flatten(), self.frequency_embedding_size) - - # Cast to model dtype before MLP (matching original LongCat) - # Handle LoRA wrapper if present - linear_layer = self.linear_1.base_layer if hasattr(self.linear_1, 'base_layer') else self.linear_1 - target_dtype = linear_layer.weight.dtype - if t_freq.dtype != target_dtype: - t_freq = t_freq.to(target_dtype) - - # MLP projection - t_emb, _ = self.linear_1(t_freq) - t_emb = self.act(t_emb) - t_emb, _ = self.linear_2(t_emb) - - # Reshape if needed - if latent_shape is not None and len(t.shape) > 1: - B = t.shape[0] - T = latent_shape[0] - t_emb = t_emb.reshape(B, T, -1) - - return t_emb - - -class CaptionEmbedder(nn.Module): - """ - Caption embedding with MLP projection and optional text compaction. - """ - - def __init__( - self, - caption_channels: int = 4096, - hidden_size: int = 4096, - text_tokens_zero_pad: bool = True, - dtype: torch.dtype | None = None, - ): - super().__init__() - self.text_tokens_zero_pad = text_tokens_zero_pad - - # Two-layer MLP using ReplicatedLinear - # CRITICAL: Original LongCat uses GELU(approximate="tanh"), NOT SiLU! - self.linear_1 = ReplicatedLinear( - caption_channels, - hidden_size, - bias=True, - params_dtype=dtype, - ) - self.act = nn.GELU(approximate="tanh") # Match original LongCat - self.linear_2 = ReplicatedLinear( - hidden_size, - hidden_size, - bias=True, - params_dtype=dtype, - ) - - def forward( - self, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - """ - Args: - encoder_hidden_states: [B, N_text, C_text] or [B, 1, N_text, C_text] - encoder_attention_mask: [B, N_text] or [B, 1, 1, N_text] - Returns: - y: [B, N_text, C] - standard padded representation (like other models) - """ - # Handle extra dimension from wrapper - if len(encoder_hidden_states.shape) == 4: - encoder_hidden_states = encoder_hidden_states.squeeze(1) - - # Project - y, _ = self.linear_1(encoder_hidden_states) - y = self.act(y) - y, _ = self.linear_2(y) # [B, N_text, C] - - # Handle attention masking - just zero out padded tokens if requested - if encoder_attention_mask is not None: - # Remove extra dimensions - if len(encoder_attention_mask.shape) == 4: - encoder_attention_mask = encoder_attention_mask.squeeze(1).squeeze(1) - elif len(encoder_attention_mask.shape) == 3: - encoder_attention_mask = encoder_attention_mask.squeeze(1) - - seq_len = int(y.shape[1]) - mask_len = int(encoder_attention_mask.shape[1]) - if mask_len < seq_len: - encoder_attention_mask = F.pad( - encoder_attention_mask, - (0, seq_len - mask_len), - value=0, - ) - elif mask_len > seq_len: - encoder_attention_mask = encoder_attention_mask[:, :seq_len] - - # Zero out padded tokens if requested - if self.text_tokens_zero_pad: - y = y * encoder_attention_mask.unsqueeze(-1) - - # Return standard format [B, N_text, C] - no compaction! - return y - - -# ============================================================================ -# Attention Modules (Placeholders for now) -# ============================================================================ - - -class LongCatSelfAttention(nn.Module): - """ - Self-attention with 3D RoPE support and optional BSA. - """ - - def __init__( - self, - dim: int, - num_heads: int, - config: LongCatVideoConfig, - dtype: torch.dtype | None = None, - ): - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - - # Separate Q/K/V projections (not fused like original) - self.to_q = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - self.to_k = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - self.to_v = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - - # Per-head RMS normalization - self.q_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32) - self.k_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32) - - # Output projection - self.to_out = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - - # 3D RoPE - self.rope_3d = RotaryPositionalEmbedding3D(head_dim=self.head_dim) - - # BSA configuration - self.enable_bsa = getattr(config, 'enable_bsa', False) - self.bsa_params = getattr(config, 'bsa_params', None) - - # FastVideo attention backend (used when BSA is disabled) - self.attn = DistributedAttention( - num_heads=num_heads, - head_size=self.head_dim, - supported_attention_backends=config._supported_attention_backends, - ) - - def forward( - self, - x: torch.Tensor, # [B, N, C] - latent_shape: tuple, # (T, H, W) - num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) - return_kv: bool = False, # Return K/V for caching - **kwargs) -> torch.Tensor | tuple: - """ - Forward pass with 3D RoPE and optional BSA. - - For I2V mode (num_cond_latents > 0): - - Conditioned tokens only attend to themselves - - Noise tokens attend to ALL tokens (cond + noise) - - Args: - return_kv: If True, return (output, (k_cache, v_cache)) for KV caching - """ - B, N, C = x.shape - T, H, W = latent_shape - - # Project to Q/K/V - q, _ = self.to_q(x) - k, _ = self.to_k(x) - v, _ = self.to_v(x) - - # Reshape to heads: [B, N, num_heads, head_dim] - q = q.view(B, N, self.num_heads, self.head_dim) - k = k.view(B, N, self.num_heads, self.head_dim) - v = v.view(B, N, self.num_heads, self.head_dim) - - # Per-head RMS normalization - q = self.q_norm(q) - k = self.k_norm(k) - - # Save pre-RoPE K/V for cache if requested (before RoPE is applied) - if return_kv: - # [B, N, num_heads, head_dim] -> [B, num_heads, N, head_dim] - k_cache = k.transpose(1, 2).clone() - v_cache = v.transpose(1, 2).clone() - - # For RoPE: need [B, num_heads, N, head_dim] - q_rope = q.transpose(1, 2) - k_rope = k.transpose(1, 2) - - # Apply 3D RoPE - q_rope, k_rope = self.rope_3d(q_rope, k_rope, grid_size=latent_shape) - - # Transpose back: [B, N, num_heads, head_dim] or [B, H, N, D] for BSA - q = q_rope.transpose(1, 2) - k = k_rope.transpose(1, 2) - - # === I2V Split Attention === - # For I2V, conditioned tokens and noise tokens are processed separately - if num_cond_latents > 0: - # Calculate number of conditioned tokens (cond_latents * spatial_tokens_per_frame) - num_cond_tokens = num_cond_latents * (N // T) - - # Conditioned tokens: only attend to themselves (same seq length, use self.attn) - q_cond = q[:, :num_cond_tokens].contiguous() - k_cond = k[:, :num_cond_tokens].contiguous() - v_cond = v[:, :num_cond_tokens].contiguous() - out_cond, _ = self.attn(q_cond, k_cond, v_cond) - - # Noise tokens: attend to ALL tokens (different seq lengths!) - # Need to use flash attention directly since q has different length than k/v - q_noise = q[:, num_cond_tokens:].contiguous() # [B, N_noise, num_heads, head_dim] - # k, v are full: [B, N, num_heads, head_dim] - - # Transpose for flash attention: [B, num_heads, seq, head_dim] - q_noise_t = q_noise.transpose(1, 2) - k_t = k.transpose(1, 2) - v_t = v.transpose(1, 2) - - # Use scaled dot product attention (handles different q/kv lengths) - out_noise_t = torch.nn.functional.scaled_dot_product_attention( - q_noise_t, k_t, v_t, attn_mask=None, dropout_p=0.0, - is_causal=False) # [B, num_heads, N_noise, head_dim] - - # Transpose back: [B, N_noise, num_heads, head_dim] - out_noise = out_noise_t.transpose(1, 2) - - # Merge conditioned and noise outputs - out = torch.cat([out_cond, out_noise], dim=1) - - # Reshape and project out - out = out.reshape(B, N, C) - out, _ = self.to_out(out) - - if return_kv: - return out, (k_cache, v_cache) - return out - - # === Attention: BSA or standard === - if self.enable_bsa and T > 1: # Only use BSA for multi-frame videos - # BSA expects [B, H, S, D] format - q_bsa = q.transpose(1, 2).contiguous() # [B, num_heads, N, head_dim] - k_bsa = k.transpose(1, 2).contiguous() - v_bsa = v.transpose(1, 2).contiguous() - - # Handle SP split: BSA operates on per-rank spatial dimensions - # Replicate LongCat's cp_split_hw logic exactly - from fastvideo.distributed.parallel_state import get_sp_world_size - sp_size = get_sp_world_size() - if sp_size > 1: - # Calculate optimal 2D split (same as LongCat's get_optimal_split) - factors = [] - for i in range(1, int(sp_size**0.5) + 1): - if sp_size % i == 0: - factors.append([i, sp_size // i]) - cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1])) - - # Split H and W dimensions by their respective factors - T_bsa, H_bsa, W_bsa = latent_shape - assert H_bsa % cp_split_hw[0] == 0 and W_bsa % cp_split_hw[1] == 0, \ - f"H {H_bsa} must be divisible by {cp_split_hw[0]}, W {W_bsa} must be divisible by {cp_split_hw[1]}" - H_bsa = H_bsa // cp_split_hw[0] - W_bsa = W_bsa // cp_split_hw[1] - latent_shape_bsa = (T_bsa, H_bsa, W_bsa) - else: - latent_shape_bsa = latent_shape - - # Call BSA with per-rank latent shape - out = flash_attn_bsa_3d(q_bsa, - k_bsa, - v_bsa, - latent_shape_q=latent_shape_bsa, - latent_shape_k=latent_shape_bsa, - **self.bsa_params) # [B, num_heads, N, head_dim] - - # Transpose back: [B, N, num_heads, head_dim] - out = out.transpose(1, 2) - else: - # Standard attention: [B, N, num_heads, head_dim] - out, _ = self.attn(q, k, v) - - # Reshape and project out - out = out.reshape(B, N, C) - out, _ = self.to_out(out) - - if return_kv: - return out, (k_cache, v_cache) - return out - - def forward_with_kv_cache( - self, - x: torch.Tensor, # [B, N_noise, C] - only noise tokens - latent_shape: tuple, # (T_noise, H, W) - shape for noise only - num_cond_latents: int, # Number of conditioning latent frames - kv_cache: tuple, # (k_cond, v_cond) - [B, heads, N_cond, head_dim] - ) -> torch.Tensor: - """ - Forward using cached K/V from conditioning frames. - - x contains only NOISE tokens. - kv_cache contains pre-computed K/V for CONDITIONING tokens. - - CRITICAL: RoPE positions for noise tokens must start AFTER conditioning. - We achieve this by padding Q with dummy tokens for conditioning positions, - applying RoPE to the full sequence, then extracting only noise token Q. - """ - B, N, C = x.shape - T, H, W = latent_shape - - k_cache, v_cache = kv_cache - - # Handle batch size mismatch (cache might be smaller for CFG) - # When using CFG, latent_model_input is doubled [neg, pos], but cache is for original batch - if k_cache.shape[0] != B: - # Expand cache to match input batch size - # For CFG: repeat the cache for both negative and positive branches - repeat_factor = B // k_cache.shape[0] - k_cache = k_cache.repeat(repeat_factor, 1, 1, 1) - v_cache = v_cache.repeat(repeat_factor, 1, 1, 1) - - # Project to Q/K/V for noise tokens - q, _ = self.to_q(x) - k, _ = self.to_k(x) - v, _ = self.to_v(x) - - # Reshape to heads: [B, N, num_heads, head_dim] - q = q.view(B, N, self.num_heads, self.head_dim) - k = k.view(B, N, self.num_heads, self.head_dim) - v = v.view(B, N, self.num_heads, self.head_dim) - - # Per-head RMS normalization - q = self.q_norm(q) - k = self.k_norm(k) - - # Transpose for RoPE: [B, heads, N, head_dim] - q_rope = q.transpose(1, 2) - k_rope = k.transpose(1, 2) - v = v.transpose(1, 2) - - # CRITICAL: Apply RoPE with correct positional offset - # Noise frame queries need positions starting from num_cond_latents - # Following the original LongCat approach: - # 1. Pad Q with dummy tokens matching k_cache shape - # 2. Apply RoPE to full sequence (T_cond + T_noise) - # 3. Extract only the noise portion of Q - - # Create dummy Q padding to fill conditioning positions - # k_cache shape: [B, heads, N_cond, head_dim] - q_padding = torch.cat([torch.empty_like(k_cache), q_rope], dim=2).contiguous() - - # Concatenate cached K with noise K for RoPE - k_full = torch.cat([k_cache, k_rope], dim=2) - v_full = torch.cat([v_cache, v], dim=2) - - # Apply RoPE to full sequence (includes both cond and noise positions) - # Grid size: (T_cond + T_noise, H, W) - full_T = num_cond_latents + T - q_padding, k_full = self.rope_3d(q_padding, k_full, grid_size=(full_T, H, W)) - - # Extract only the noise portion of Q (last N tokens) - q_rope = q_padding[:, :, -N:].contiguous() - - # Run attention: Q_noise attends to full K/V (cond + noise) - out = torch.nn.functional.scaled_dot_product_attention(q_rope, - k_full, - v_full, - attn_mask=None, - dropout_p=0.0, - is_causal=False) # [B, heads, N_noise, head_dim] - - # Transpose back: [B, N_noise, heads, head_dim] - out = out.transpose(1, 2) - out = out.reshape(B, N, C) - out, _ = self.to_out(out) - - return out - - -class LongCatCrossAttention(nn.Module): - """ - Cross-attention for text conditioning (standard implementation like other models). - """ - - def __init__( - self, - dim: int, - num_heads: int, - config: LongCatVideoConfig, - dtype: torch.dtype | None = None, - ): - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - - # Separate Q/K/V projections - self.to_q = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - self.to_k = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - self.to_v = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - - # Per-head RMS normalization - self.q_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32) - self.k_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32) - - # Output projection - self.to_out = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype) - - # Cross-attention uses LocalAttention (FastVideo standard) - self.attn = LocalAttention( - num_heads=num_heads, - head_size=self.head_dim, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=config.arch_config._supported_attention_backends, - ) - - def forward( - self, - x: torch.Tensor, # [B, N_img, C] - context: torch.Tensor, # [B, N_text, C] - latent_shape: tuple = None, # (T, H, W) - needed for I2V - num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) - **kwargs) -> torch.Tensor: - """ - Forward pass for cross-attention (standard implementation). - - Args: - x: Image tokens [B, N_img, C] - context: Text tokens [B, N_text, C] (standard padded format) - latent_shape: (T, H, W) - needed for calculating num_cond_tokens - num_cond_latents: Number of conditioning latent frames (for I2V) - - For I2V mode (num_cond_latents > 0): - - Conditioned tokens get ZERO cross-attention output - - Only noise tokens get cross-attention with text - """ - B, N_img, C = x.shape - - # === I2V: Only noise tokens get cross-attention === - if num_cond_latents > 0 and latent_shape is not None: - T, H, W = latent_shape - num_cond_tokens = num_cond_latents * (N_img // T) - - # Only process noise tokens - x_noise = x[:, num_cond_tokens:] # [B, N_noise, C] - - # Project Q, K, V for noise tokens only - q, _ = self.to_q(x_noise) - k, _ = self.to_k(context) - v, _ = self.to_v(context) - - N_text = context.shape[1] - N_noise = x_noise.shape[1] - - # Reshape to heads - q = q.view(B, N_noise, self.num_heads, self.head_dim) - k = k.view(B, N_text, self.num_heads, self.head_dim) - v = v.view(B, N_text, self.num_heads, self.head_dim) - - # Per-head RMS normalization - q = self.q_norm(q) - k = self.k_norm(k) - - # Run cross-attention - out_noise = self.attn(q, k, v) # [B, N_noise, num_heads, head_dim] - out_noise = out_noise.reshape(B, N_noise, C) - out_noise, _ = self.to_out(out_noise) - - # Conditioned tokens get zero output - out_cond = torch.zeros((B, num_cond_tokens, C), dtype=out_noise.dtype, device=out_noise.device) - - # Merge - out = torch.cat([out_cond, out_noise], dim=1) - return out - - # === Standard cross-attention === - # Project Q, K, V (standard cross-attention like WanVideo/Cosmos) - q, _ = self.to_q(x) - k, _ = self.to_k(context) - v, _ = self.to_v(context) - - N_text = context.shape[1] - - # Reshape to heads - q = q.view(B, N_img, self.num_heads, self.head_dim) - k = k.view(B, N_text, self.num_heads, self.head_dim) - v = v.view(B, N_text, self.num_heads, self.head_dim) - - # Per-head RMS normalization - q = self.q_norm(q) - k = self.k_norm(k) - - # Run cross-attention using FastVideo's LocalAttention - # LocalAttention handles different q and k/v sequence lengths automatically - out = self.attn(q, k, v) # [B, N_img, num_heads, head_dim] - - # Reshape and project out - out = out.reshape(B, N_img, C) - out, _ = self.to_out(out) - - return out - - -# ============================================================================ -# Feed-Forward Network -# ============================================================================ - - -class LongCatSwiGLUFFN(nn.Module): - """ - SwiGLU feed-forward network using FastVideo's ReplicatedLinear. - - FFN(x) = down(gate(x) * SiLU(up(x))) - """ - - def __init__( - self, - dim: int, - hidden_dim: int, - dtype: torch.dtype | None = None, - ): - super().__init__() - - # Three projections for SwiGLU (no bias as per original) - self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False, params_dtype=dtype) # gate - self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False, params_dtype=dtype) # up - self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False, params_dtype=dtype) # down - - self.act = nn.SiLU() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass: SiLU(w1(x)) * w3(x) -> w2 (matching original LongCat) - """ - w1_out, _ = self.w1(x) - w3_out, _ = self.w3(x) - combined = self.act(w1_out) * w3_out - out, _ = self.w2(combined) - return out - - -# ============================================================================ -# Modulation Utilities -# ============================================================================ - - -def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - """ - Apply modulation in FP32 for numerical stability. - - Converts inputs to FP32 for the modulation operation, then casts back. - """ - orig_dtype = x.dtype - - # Convert to FP32 for numerical stability - shift_fp32 = shift.float() - scale_fp32 = scale.float() - - # Normalize and modulate in FP32 - x_norm = norm(x.to(torch.float32)) - x_mod = x_norm * (scale_fp32 + 1) + shift_fp32 - - return x_mod.to(orig_dtype) - - -# ============================================================================ -# Transformer Block -# ============================================================================ - - -class LongCatTransformerBlock(nn.Module): - """ - Single-stream transformer block with: - - AdaLN modulation (FP32) - - Self-attention - - Cross-attention - - SwiGLU FFN - """ - - def __init__( - self, - hidden_size: int, - num_heads: int, - mlp_ratio: int, - adaln_tembed_dim: int, - config: LongCatVideoConfig, - dtype: torch.dtype | None = None, - ): - super().__init__() - - self.hidden_size = hidden_size - self.num_heads = num_heads - - # AdaLN modulation (6 parameters: scale/shift for attn & ffn, gate for residual) - self.adaln_linear_1 = ReplicatedLinear( - adaln_tembed_dim, - 6 * hidden_size, - bias=True, - params_dtype=dtype, - ) - self.adaln_act = nn.SiLU() - - # Normalization layers (CRITICAL: Use LayerNorm not RMSNorm like original!) - # Original LongCat uses LayerNorm_FP32 with elementwise_affine=False - self.norm_attn = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) - self.norm_ffn = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) - # Cross-attention norm has elementwise_affine=True (has weight and bias) - self.norm_cross = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True) - - # Self-attention - self.self_attn = LongCatSelfAttention( - dim=hidden_size, - num_heads=num_heads, - config=config, - dtype=dtype, - ) - - # Cross-attention - self.cross_attn = LongCatCrossAttention( - dim=hidden_size, - num_heads=num_heads, - config=config, - dtype=dtype, - ) - - # SwiGLU FFN - ffn_hidden_dim = int(hidden_size * mlp_ratio * 2 / 3) - # Round up to nearest multiple of 256 - ffn_hidden_dim = 256 * ((ffn_hidden_dim + 255) // 256) - - self.ffn = LongCatSwiGLUFFN( - dim=hidden_size, - hidden_dim=ffn_hidden_dim, - dtype=dtype, - ) - - def forward( - self, - x: torch.Tensor, # [B, N, C] - context: torch.Tensor, # [B, N_text, C] - t: torch.Tensor, # [B, T, C_t] - latent_shape: tuple, # (T, H, W) - num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) - return_kv: bool = False, # Return K/V for caching - kv_cache: tuple | None = None, # Pre-computed K/V cache - skip_crs_attn: bool = False, # Skip cross-attention (for cache init) - **kwargs) -> torch.Tensor | tuple: - """ - Forward pass with AdaLN modulation. - - Args: - num_cond_latents: For I2V, number of conditioning latent frames. - These frames use split attention behavior. - return_kv: If True, return (x, (k_cache, v_cache)) - kv_cache: Pre-computed K/V from conditioning frames - skip_crs_attn: If True, skip cross-attention (used during cache init) - """ - B, N, C = x.shape - T, H, W = latent_shape - x_orig_dtype = x.dtype # Save for later casting - - # === AdaLN Modulation (CRITICAL: FP32 for stability like original) === - # Use autocast to compute modulation params in FP32 - with torch.amp.autocast(device_type='cuda', dtype=torch.float32): - t_mod = self.adaln_act(t) - mod_params, _ = self.adaln_linear_1(t_mod) - # Ensure FP32 output (needed when LoRA is applied) - if mod_params.dtype != torch.float32: - mod_params = mod_params.float() - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \ - mod_params.unsqueeze(2).chunk(6, dim=-1) # [B, T, 1, C] - - # === Self-Attention === - x_norm = modulate_fp32(self.norm_attn, x.view(B, T, -1, C), shift_msa, scale_msa) - x_norm = x_norm.view(B, N, C) - - # Handle KV cache - if kv_cache is not None: - # Move cache to device if offloaded - kv_cache = (kv_cache[0].to(x.device), kv_cache[1].to(x.device)) - attn_out = self.self_attn.forward_with_kv_cache( - x_norm, - latent_shape=latent_shape, - num_cond_latents=num_cond_latents, - kv_cache=kv_cache, - ) - kv_cache_new = None # Don't return cache when using cache - else: - attn_result = self.self_attn( - x_norm, - latent_shape=latent_shape, - num_cond_latents=num_cond_latents, - return_kv=return_kv, - ) - if return_kv: - attn_out, kv_cache_new = attn_result - else: - attn_out = attn_result - kv_cache_new = None - - # Residual with gating (CRITICAL: FP32 like original, then cast back) - with torch.amp.autocast(device_type='cuda', dtype=torch.float32): - x = x + (gate_msa * attn_out.view(B, T, -1, C)).view(B, N, C) - x = x.to(x_orig_dtype) - - # === Cross-Attention (skip if requested) === - if not skip_crs_attn: - x_norm_cross = self.norm_cross(x) - # When using KV cache, no need for num_cond_latents in cross-attn - cross_num_cond = 0 if kv_cache is not None else num_cond_latents - cross_out = self.cross_attn(x_norm_cross, - context, - latent_shape=latent_shape, - num_cond_latents=cross_num_cond) - x = x + cross_out - - # === FFN === - x_norm_ffn = modulate_fp32(self.norm_ffn, x.view(B, T, -1, C), shift_mlp, scale_mlp) - x_norm_ffn = x_norm_ffn.view(B, N, C) - - ffn_out = self.ffn(x_norm_ffn) - - # Residual with gating (CRITICAL: FP32 like original, then cast back) - with torch.amp.autocast(device_type='cuda', dtype=torch.float32): - x = x + (gate_mlp * ffn_out.view(B, T, -1, C)).view(B, N, C) - x = x.to(x_orig_dtype) - - if return_kv: - return x, kv_cache_new - return x - - -# ============================================================================ -# Final Layer -# ============================================================================ - - -class FinalLayer(nn.Module): - """ - Final output projection with AdaLN modulation. - """ - - def __init__( - self, - hidden_size: int, - out_channels: int, - adaln_tembed_dim: int, - patch_size: tuple[int, int, int], - dtype: torch.dtype | None = None, - ): - super().__init__() - - # AdaLN for final layer (2 parameters: scale and shift) - self.adaln_linear = ReplicatedLinear( - adaln_tembed_dim, - 2 * hidden_size, - bias=True, - params_dtype=dtype, - ) - self.adaln_act = nn.SiLU() - - # CRITICAL: Use LayerNorm not RMSNorm! (matches original) - self.norm = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) - - # Output projection - num_patch = patch_size[0] * patch_size[1] * patch_size[2] - self.proj = ReplicatedLinear( - hidden_size, - num_patch * out_channels, - bias=True, - params_dtype=dtype, - ) - - def forward( - self, - x: torch.Tensor, # [B, N, C] - t: torch.Tensor, # [B, T, C_t] - latent_shape: tuple, - ) -> torch.Tensor: - """ - Returns: [B, N, out_channels * patch_size^3] - """ - B, N, C = x.shape - T, _, _ = latent_shape - - # AdaLN modulation - t_mod = self.adaln_act(t) - mod_params, _ = self.adaln_linear(t_mod) - shift, scale = mod_params.unsqueeze(2).chunk(2, dim=-1) - - # Modulate (converts to FP32 internally for stability) - x = modulate_fp32(self.norm, x.view(B, T, -1, C), shift, scale) - x = x.reshape(B, N, C) - - # Project - x, _ = self.proj(x) - - return x - - -# ============================================================================ -# Main Model -# ============================================================================ - - -class LongCatTransformer3DModel(BaseDiT): - """ - Native LongCat Video Transformer using FastVideo layers. - """ - - # FSDP sharding: shard at each transformer block - _fsdp_shard_conditions = [ - lambda n, m: "blocks" in n and n.split(".")[-1].isdigit(), - ] - # torch.compile optimization: compile each transformer block for speedup - _compile_conditions = [ - lambda n, m: "blocks" in n and n.split(".")[-1].isdigit(), - ] - - # Parameter name mapping (for weight conversion) - param_names_mapping = {} # Will be defined in config - reverse_param_names_mapping = {} - lora_param_names_mapping = {} - - # Supported attention backends - _supported_attention_backends = ( - AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - ) - - def __init__(self, config: LongCatVideoConfig, hf_config: dict[str, Any]): - super().__init__(config=config, hf_config=hf_config) - - # Extract architecture parameters - self.hidden_size = config.hidden_size # 4096 - self.num_attention_heads = config.num_attention_heads # 32 - self.depth = config.depth # 48 - self.mlp_ratio = config.mlp_ratio # 4 - self.in_channels = config.in_channels # 16 - self.out_channels = config.out_channels # 16 - self.num_channels_latents = config.in_channels - self.patch_size = config.patch_size # [1, 2, 2] - - # Embeddings - self.patch_embed = PatchEmbed3D( - patch_size=self.patch_size, - in_channels=self.in_channels, - embed_dim=self.hidden_size, - ) - - self.time_embedder = TimestepEmbedder( - frequency_embedding_size=config.frequency_embedding_size, - adaln_tembed_dim=config.adaln_tembed_dim, - ) - - self.caption_embedder = CaptionEmbedder( - caption_channels=config.caption_channels, - hidden_size=self.hidden_size, - text_tokens_zero_pad=getattr(config, 'text_tokens_zero_pad', True), - ) - - # Transformer blocks (48 blocks) - self.blocks = nn.ModuleList([ - LongCatTransformerBlock( - hidden_size=self.hidden_size, - num_heads=self.num_attention_heads, - mlp_ratio=self.mlp_ratio, - adaln_tembed_dim=config.adaln_tembed_dim, - config=config, - ) for _ in range(self.depth) - ]) - - # Output projection - self.final_layer = FinalLayer( - hidden_size=self.hidden_size, - out_channels=self.out_channels, - adaln_tembed_dim=config.adaln_tembed_dim, - patch_size=self.patch_size, - ) - - def enable_bsa(self): - """Enable BSA for all self-attention layers.""" - for block in self.blocks: - block.self_attn.enable_bsa = True - - def disable_bsa(self): - """Disable BSA for all self-attention layers.""" - for block in self.blocks: - block.self_attn.enable_bsa = False - - def forward( - self, - hidden_states: torch.Tensor, # [B, C, T, H, W] - encoder_hidden_states: torch.Tensor | list[torch.Tensor], # [B, N_text, C_text] - timestep: torch.LongTensor, # [B] or [B, T] - encoder_attention_mask: torch.Tensor | None = None, # [B, N_text] - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, - guidance: float | None = None, # Unused, for API compatibility - num_cond_latents: int = 0, # For I2V: number of conditioning latent frames - # === KV Cache Parameters === - return_kv: bool = False, # If True, return (output, kv_cache_dict) - kv_cache_dict: dict | None = None, # Pre-computed {block_idx: (k, v)} - skip_crs_attn: bool = False, # Skip cross-attention (for cache init) - offload_kv_cache: bool = False, # Move cache to CPU after compute - **kwargs) -> torch.Tensor | tuple[torch.Tensor, dict]: - """ - Forward pass with FastVideo parameter ordering. - - NOTE: This follows FastVideo convention: - (hidden_states, encoder_hidden_states, timestep) - - Args: - num_cond_latents: For I2V, number of conditioning latent frames. - These frames are treated as "clean" (timestep=0) - and use split attention behavior. - return_kv: If True, return (output, kv_cache_dict) - kv_cache_dict: Pre-computed K/V cache {block_idx: (k, v)} - skip_crs_attn: If True, skip cross-attention (for cache init) - offload_kv_cache: If True, move cache to CPU after compute - """ - B, _, T, H, W = hidden_states.shape - - N_t = T // self.patch_size[0] - N_h = H // self.patch_size[1] - N_w = W // self.patch_size[2] - - # Handle list of encoder outputs (take first one) - if isinstance(encoder_hidden_states, list): - encoder_hidden_states = encoder_hidden_states[0] - - # 1. Patch embedding - x = self.patch_embed(hidden_states) # [B, N, C] - - # 2. Timestep embedding - # Expand timestep from [B] to [B, T] if needed - if timestep.ndim == 1: - timestep = timestep.unsqueeze(1).expand(-1, N_t) # [B, T] - - t = self.time_embedder(timestep.flatten(), latent_shape=(N_t, N_h, N_w)) - if t.ndim == 2: - t = t.reshape(B, N_t, -1) # [B, T, C_t] - - # 3. Caption embedding (standard format, no compaction) - context = self.caption_embedder(encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask) # [B, N_text, C] - - # 4. Transformer blocks with optional KV cache - kv_cache_dict_ret = {} if return_kv else None - - for i, block in enumerate(self.blocks): - # Get cache for this block if available - block_kv_cache = kv_cache_dict.get(i, None) if kv_cache_dict else None - - block_out = block( - x, - context, - t, - latent_shape=(N_t, N_h, N_w), - num_cond_latents=num_cond_latents, - return_kv=return_kv, - kv_cache=block_kv_cache, - skip_crs_attn=skip_crs_attn, - ) - - if return_kv: - x, kv_cache = block_out - # Store cache - if offload_kv_cache: - kv_cache_dict_ret[i] = (kv_cache[0].cpu(), kv_cache[1].cpu()) - else: - kv_cache_dict_ret[i] = (kv_cache[0].contiguous(), kv_cache[1].contiguous()) - else: - x = block_out - - # 5. Output projection - output = self.final_layer(x, t, latent_shape=(N_t, N_h, N_w)) - - # Reshape to [B, C_out, T, H, W] - output = self.unpatchify(output, N_t, N_h, N_w) - - # Cast to float32 for better accuracy (as per original) - output = output.to(torch.float32) - - if return_kv: - return output, kv_cache_dict_ret - return output - - def unpatchify(self, x: torch.Tensor, N_t: int, N_h: int, N_w: int) -> torch.Tensor: - """ - Args: - x: [B, N, C] where C = T_p * H_p * W_p * C_out - Returns: - [B, C_out, T, H, W] - """ - T_p, H_p, W_p = self.patch_size - x = rearrange( - x, - "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)", - N_t=N_t, - N_h=N_h, - N_w=N_w, - T_p=T_p, - H_p=H_p, - W_p=W_p, - C_out=self.out_channels, - ) - return x - - -# Entry point for model registry -EntryClass = LongCatTransformer3DModel diff --git a/fastvideo/models/dits/zimage.py b/fastvideo/models/dits/zimage.py deleted file mode 100644 index 219327790..000000000 --- a/fastvideo/models/dits/zimage.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/encoders/glm_image_ar_loader.py b/fastvideo/models/encoders/glm_image_ar_loader.py deleted file mode 100644 index 7d6be0f49..000000000 --- a/fastvideo/models/encoders/glm_image_ar_loader.py +++ /dev/null @@ -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() diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 1647c6a48..3e282d9a5 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -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.""" diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index e859f3e8c..f3baf6e28 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -25,18 +25,12 @@ logger = init_logger(__name__) _TEXT_TO_VIDEO_DIT_MODELS = { "MMAudioTransformer": ("dits", "mmaudio", "MMAudioTransformer"), "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), - "HunyuanGameCraftTransformer3DModel": ("dits", "hunyuangamecraft", "HunyuanGameCraftTransformer3DModel"), "HunyuanVideo15Transformer3DModel": ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"), "HYWorldTransformer3DModel": ("dits", "hyworld", "HYWorldTransformer3DModel"), "WanTransformer3DModel": ("wan", "transformer", "WanTransformer3DModel"), - "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), - "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), "CausalWanTransformer3DModel": ("wan", "causal_transformer", "CausalWanTransformer3DModel"), "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), - "LongCatVideoTransformer3DModel": - ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) - "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2) "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"), "SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"), "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"), @@ -54,8 +48,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = { _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("wan", "transformer", "WanTransformer3DModel"), - "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), - "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), "CausalWanTransformer3DModel": ("wan", "causal_transformer", "CausalWanTransformer3DModel"), "LingBotWorld2CausalFastTransformer3DModel": ( "dits", @@ -72,8 +64,6 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { # Text-to-image DiT models (2D image generation) _TEXT_TO_IMAGE_DIT_MODELS = { - "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"), - "ZImageTransformer2DModel": ("dits", "zimage", "ZImageTransformer2DModel"), } _TEXT_ENCODER_MODELS = { @@ -88,9 +78,6 @@ _TEXT_ENCODER_MODELS = { "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"), "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"), "Qwen2_5_VLForConditionalGeneration": ("encoders", "reason1", "Reason1TextEncoder"), - # Z-Image-Turbo's text_encoder/config.json declares architecture - # "Qwen3Model"; route it to the shared Qwen3 encoder (added for Flux2 Klein). - "Qwen3Model": ("encoders", "qwen3", "Qwen3ForCausalLM"), "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"), "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"), "Mistral3ForConditionalGeneration": ("encoders", "mistral3", "Mistral3ForConditionalGeneration"), @@ -111,7 +98,6 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = { _VAE_MODELS = { "AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), - "AutoencoderKLCausal3D": ("vaes", "gamecraftvae", "GameCraftVAE"), "AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"), "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"), "AutoencoderKLWan": ("wan", "vae", "AutoencoderKLWan"), @@ -142,7 +128,6 @@ _SCHEDULERS = { "FlowUniPCMultistepScheduler": ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"), "SelfForcingFlowMatchScheduler": ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"), - "RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"), } _UPSAMPLERS = { diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 5daffcdb1..7fad5b97f 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -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, diff --git a/fastvideo/models/schedulers/scheduling_rcm.py b/fastvideo/models/schedulers/scheduling_rcm.py deleted file mode 100644 index e48066af7..000000000 --- a/fastvideo/models/schedulers/scheduling_rcm.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/vaes/gamecraftvae.py b/fastvideo/models/vaes/gamecraftvae.py deleted file mode 100644 index e687c346c..000000000 --- a/fastvideo/models/vaes/gamecraftvae.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/vaes/gamecraftvae_blocks.py b/fastvideo/models/vaes/gamecraftvae_blocks.py deleted file mode 100644 index 7f6bedd7a..000000000 --- a/fastvideo/models/vaes/gamecraftvae_blocks.py +++ /dev/null @@ -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 diff --git a/fastvideo/models/wan/AGENTS.md b/fastvideo/models/wan/AGENTS.md index 46f478bba..e0f0aa32f 100644 --- a/fastvideo/models/wan/AGENTS.md +++ b/fastvideo/models/wan/AGENTS.md @@ -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 diff --git a/fastvideo/models/wan/definition.py b/fastvideo/models/wan/definition.py index b0b673b65..0758e5a8c 100644 --- a/fastvideo/models/wan/definition.py +++ b/fastvideo/models/wan/definition.py @@ -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) diff --git a/fastvideo/pipelines/AGENTS.md b/fastvideo/pipelines/AGENTS.md index 991b47867..e7b7b4639 100644 --- a/fastvideo/pipelines/AGENTS.md +++ b/fastvideo/pipelines/AGENTS.md @@ -20,15 +20,13 @@ pipelines/ │ ├── conditioning.py # CFG / negative prompt fan-out │ ├── denoising.py # Standard diffusion loop │ ├── sd35_conditioning.py # Per-model overrides (named by family) -│ ├── longcat_*.py # LongCat I2V/V2V/refine variants │ ├── gen3c_stages.py # Gen3C-specific stages -│ ├── gamecraft_denoising.py # GameCraft-specific │ └── matrixgame2_denoising.py # Matrix-Game 2.0-specific ├── basic/ # Per-model end-to-end pipelines -│ ├── cosmos/, dreamx_world/, flux/, flux_2/, gamecraft/, gen3c/, glm_image/ +│ ├── cosmos/, flux/, flux_2/, gen3c/ │ ├── hunyuan/, hunyuan15/, hyworld/, kandinsky5/, lingbot_video/, lingbotworld/ -│ ├── lingbotworld2/, longcat/, ltx2/, magi_human/, matrixgame2/, matrixgame3/ -│ ├── minimax_h3/, mmaudio/, sd35/, stable_audio/, turbodiffusion/, wan/, zimage/ +│ ├── lingbotworld2/, ltx2/, magi_human/, matrixgame2/, matrixgame3/ +│ ├── minimax_h3/, mmaudio/, sd35/, stable_audio/, wan/ │ └── /{_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`. diff --git a/fastvideo/pipelines/basic/dreamx_world/__init__.py b/fastvideo/pipelines/basic/dreamx_world/__init__.py deleted file mode 100644 index 38c7f6aa2..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/__init__.py +++ /dev/null @@ -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", -] diff --git a/fastvideo/pipelines/basic/dreamx_world/ar_denoising.py b/fastvideo/pipelines/basic/dreamx_world/ar_denoising.py deleted file mode 100644 index 874c44a5a..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/ar_denoising.py +++ /dev/null @@ -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, - ) diff --git a/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py b/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py deleted file mode 100644 index 13b87f154..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py +++ /dev/null @@ -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} diff --git a/fastvideo/pipelines/basic/dreamx_world/config.py b/fastvideo/pipelines/basic/dreamx_world/config.py deleted file mode 100644 index 760f5ed60..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/config.py +++ /dev/null @@ -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", -] diff --git a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py deleted file mode 100644 index 8cee19f2e..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py +++ /dev/null @@ -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 diff --git a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py deleted file mode 100644 index e4bf62e84..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py +++ /dev/null @@ -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 diff --git a/fastvideo/pipelines/basic/dreamx_world/presets.py b/fastvideo/pipelines/basic/dreamx_world/presets.py deleted file mode 100644 index 2d9d6e407..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/presets.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World model family pipeline presets.""" - -from fastvideo.api.presets import InferencePreset, PresetStageSpec - -_NEGATIVE_PROMPT_CN = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，" - "静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，" - "多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，" - "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，" - "背景人很多，倒着走") - -_DENOISE_STAGE = PresetStageSpec( - name="denoise", - kind="denoising", - description="DreamX-World camera-conditioned denoising pass", - allowed_overrides=frozenset({ - "num_inference_steps", - "guidance_scale", - }), -) - -DREAMX_WORLD_5B_CAM = InferencePreset( - name="dreamx_world_5b_cam", - version=1, - model_family="dreamx_world", - description="DreamX-World 5B camera-control video generation", - workload_type="i2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 480, - "width": 832, - "num_frames": 161, - "fps": 16, - "guidance_scale": 5.0, - "num_inference_steps": 30, - "negative_prompt": _NEGATIVE_PROMPT_CN, - }, -) - -DREAMX_WORLD_5B_AR = InferencePreset( - name="dreamx_world_5b_ar", - version=1, - model_family="dreamx_world", - description="DreamX-World 5B autoregressive camera-control generation", - workload_type="i2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 704, - "width": 1280, - "num_frames": 1005, - "fps": 16, - "guidance_scale": 1.0, - "num_inference_steps": 4, - "negative_prompt": _NEGATIVE_PROMPT_CN, - }, -) - -ALL_PRESETS = (DREAMX_WORLD_5B_CAM, DREAMX_WORLD_5B_AR) diff --git a/fastvideo/pipelines/basic/dreamx_world/stages.py b/fastvideo/pipelines/basic/dreamx_world/stages.py deleted file mode 100644 index 10e0b6741..000000000 --- a/fastvideo/pipelines/basic/dreamx_world/stages.py +++ /dev/null @@ -1,156 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World pipeline stages.""" - -from __future__ import annotations - -import torch - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.validators import VerificationResult - -from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import ( - build_dreamx_camera_condition, ) - -DREAMX_Y_CAMERA_KEY = "dreamx_y_camera" - -logger = init_logger(__name__) - - -class DreamXWorldCameraConditioningStage(PipelineStage): - """Build PRoPE camera conditioning for DreamX-World-5B-Cam.""" - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - del fastvideo_args - if DREAMX_Y_CAMERA_KEY in batch.extra: - return batch - - action_seq = batch.extra.get("dreamx_action_seq", batch.action_list) - action_speed_list = batch.extra.get("dreamx_action_speed_list", batch.action_speed_list) - if action_seq is None: - action_seq = ["w"] - if action_speed_list is None: - action_speed_list = [4] - - if isinstance(action_seq, str): - action_seq = [action_seq] - if isinstance(action_speed_list, int | float): - action_speed_list = [action_speed_list] - if len(action_speed_list) == 1 and len(action_seq) > 1: - action_speed_list = list(action_speed_list) * len(action_seq) - action_speed_list = [float(speed) for speed in action_speed_list] - - height = int(batch.height) if batch.height is not None else 704 - width = int(batch.width) if batch.width is not None else 1280 - num_frames = int(batch.num_frames) - dtype = batch.latents.dtype if torch.is_tensor(batch.latents) else torch.float32 - device = batch.latents.device if torch.is_tensor(batch.latents) else "cpu" - - y_camera = build_dreamx_camera_condition( - list(action_seq), - action_speed_list, - num_frames=num_frames, - height=height, - width=width, - dtype=dtype, - device=device, - ) - batch.extra[DREAMX_Y_CAMERA_KEY] = {key: value.unsqueeze(0) for key, value in y_camera.items()} - return batch - - def verify_output( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> VerificationResult: - del fastvideo_args - result = VerificationResult() - y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY) - result.add_check("dreamx_y_camera", y_camera, lambda value: isinstance(value, dict)) - if isinstance(y_camera, dict): - result.add_check("dreamx_y_camera.viewmats", y_camera.get("viewmats"), torch.is_tensor) - result.add_check("dreamx_y_camera.K", y_camera.get("K"), torch.is_tensor) - return result - - -class DreamXWorldImageVAEEncodingStage(PipelineStage): - """Encode the conditioning image into the first-frame latent. - - Official AR-forcing flow (AMAP-ML/DreamX-World inference_ar_forcing.py): - the input image is resized, normalized to [-1, 1], VAE-encoded - deterministically, and written into frame 0 of the noise — the causal - denoiser then treats frame 0 as clean context. This stage produces - ``batch.image_latent`` ([B, C, 1, H_lat, W_lat]); the injection into - the latents happens in DreamXWorldARCausalDenoisingStage. - """ - - def __init__(self, vae) -> None: - super().__init__() - self.vae = vae - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - if batch.pil_image is None: - # No conditioning image: the causal denoiser falls back to - # running from pure noise (frame 0 uninitialized). Warn loudly — - # this pipeline is registered I2V and the official flow always - # forces from a frame. - logger.warning("DreamXWorldARPipeline called without an input image; " - "first-frame context will be noise (T2V-style). Pass an " - "image for the official AR-forcing behavior.") - return batch - - from fastvideo.platforms import get_local_torch_device - from fastvideo.utils import PRECISION_TO_TYPE - - device = get_local_torch_device() - image = batch.pil_image - if not isinstance(image, torch.Tensor): - import numpy as np - import PIL.Image - assert isinstance(image, PIL.Image.Image) - width = batch.width if isinstance(batch.width, int) else batch.width[0] - height = batch.height if isinstance(batch.height, int) else batch.height[0] - image = image.convert("RGB").resize((width, height), PIL.Image.Resampling.LANCZOS) - arr = torch.from_numpy(np.asarray(image)).float().permute(2, 0, 1) / 255.0 - image = (arr - 0.5) / 0.5 # official Normalize([0.5], [0.5]) - image = image.unsqueeze(0) # [1, C, H, W] - if image.dim() == 4: - image = image.unsqueeze(2) # [B, C, 1, H, W] - elif image.dim() == 5: - image = image[:, :, :1] - image = image.to(device=device, dtype=torch.float32) - - vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - self.vae = self.vae.to(device) - with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): - if not vae_autocast_enabled: - image = image.to(vae_dtype) - encoder_output = self.vae.encode(image) - - # Official encode_to_latent is deterministic ((mean - mu) / sigma per - # channel); the posterior mean + shift/scale is the FastVideo - # equivalent of that normalization. - latent = encoder_output.mean - if getattr(self.vae, "shift_factor", None) is not None: - shift = self.vae.shift_factor - latent = latent - (shift.to(latent.device, latent.dtype) if isinstance(shift, torch.Tensor) else shift) - scale = self.vae.scaling_factor - latent = latent * (scale.to(latent.device, latent.dtype) if isinstance(scale, torch.Tensor) else scale) - - batch.image_latent = latent - return batch - - def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - result = VerificationResult() - return result diff --git a/fastvideo/pipelines/basic/flux_2/flux_2_latent_preparation.py b/fastvideo/pipelines/basic/flux_2/flux_2_latent_preparation.py index c30dd561c..7aa57de5a 100644 --- a/fastvideo/pipelines/basic/flux_2/flux_2_latent_preparation.py +++ b/fastvideo/pipelines/basic/flux_2/flux_2_latent_preparation.py @@ -120,8 +120,7 @@ class Flux2LatentPreparationStage(LatentPreparationStage): latents = latents * self.scheduler.init_noise_sigma else: latents = latents.to(device) - is_longcat_refine = (batch.refine_from is not None or batch.stage1_video is not None) - if (not is_longcat_refine) and hasattr(self.scheduler, "init_noise_sigma"): + if hasattr(self.scheduler, "init_noise_sigma"): latents = latents * self.scheduler.init_noise_sigma batch.latents = latents diff --git a/fastvideo/pipelines/basic/gamecraft/__init__.py b/fastvideo/pipelines/basic/gamecraft/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/fastvideo/pipelines/basic/gamecraft/gamecraft_pipeline.py b/fastvideo/pipelines/basic/gamecraft/gamecraft_pipeline.py deleted file mode 100644 index 00a5a8192..000000000 --- a/fastvideo/pipelines/basic/gamecraft/gamecraft_pipeline.py +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -HunyuanGameCraft video diffusion pipeline implementation. - -This module implements the HunyuanGameCraft pipeline for camera/action-conditioned -video generation with the modular pipeline architecture. -""" - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase -from fastvideo.pipelines.stages import ( - ConditioningStage, - DecodingStage, - InputValidationStage, - LatentPreparationStage, - TextEncodingStage, - TimestepPreparationStage, -) -from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage - -logger = init_logger(__name__) - - -class HunyuanGameCraftPipeline(ComposedPipelineBase): - """ - Pipeline for HunyuanGameCraft video generation. - - This pipeline supports: - - Text-to-video generation with camera/action conditioning - - Autoregressive generation with history frames - - 33-channel input (16 latent + 16 gt_latent + 1 mask) - - CameraNet for encoding Plücker coordinates - """ - - _required_config_modules = [ - "text_encoder", - "text_encoder_2", - "tokenizer", - "tokenizer_2", - "vae", - "transformer", - "scheduler", - ] - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): - """Set up pipeline stages with proper dependency injection.""" - - self.add_stage( - stage_name="input_validation_stage", - stage=InputValidationStage(), - ) - - self.add_stage( - stage_name="prompt_encoding_stage_primary", - stage=TextEncodingStage( - text_encoders=[ - self.get_module("text_encoder"), - self.get_module("text_encoder_2"), - ], - tokenizers=[ - self.get_module("tokenizer"), - self.get_module("tokenizer_2"), - ], - ), - ) - - self.add_stage( - stage_name="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"), - ), - ) - - self.add_stage( - stage_name="denoising_stage", - stage=GameCraftDenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"), - ), - ) - - self.add_stage( - stage_name="decoding_stage", - stage=DecodingStage(vae=self.get_module("vae")), - ) - - -EntryClass = HunyuanGameCraftPipeline diff --git a/fastvideo/pipelines/basic/gamecraft/presets.py b/fastvideo/pipelines/basic/gamecraft/presets.py deleted file mode 100644 index 1d2261ac0..000000000 --- a/fastvideo/pipelines/basic/gamecraft/presets.py +++ /dev/null @@ -1,33 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""HunyuanGameCraft model family pipeline presets.""" -from fastvideo.api.presets import InferencePreset, PresetStageSpec - -_DENOISE_STAGE = PresetStageSpec( - name="denoise", - kind="denoising", - description="Action-controlled denoising pass", - allowed_overrides=frozenset({ - "num_inference_steps", - "guidance_scale", - }), -) - -GAMECRAFT_I2V = InferencePreset( - name="gamecraft_i2v", - version=1, - model_family="gamecraft", - description="HunyuanGameCraft I2V at 704x1280", - workload_type="i2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 704, - "width": 1280, - "num_frames": 33, - "fps": 24, - "guidance_scale": 6.0, - "num_inference_steps": 50, - "negative_prompt": "", - }, -) - -ALL_PRESETS = (GAMECRAFT_I2V, ) diff --git a/fastvideo/pipelines/basic/glm_image/__init__.py b/fastvideo/pipelines/basic/glm_image/__init__.py deleted file mode 100644 index f85bca3cc..000000000 --- a/fastvideo/pipelines/basic/glm_image/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""GLM-Image pipeline package.""" - -from fastvideo.pipelines.basic.glm_image.glm_image_pipeline import ( - GlmImagePipeline, ) - -__all__ = ["GlmImagePipeline"] diff --git a/fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py b/fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py deleted file mode 100644 index d649414fe..000000000 --- a/fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -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.glm_image.stages import ( - GlmImageBeforeDenoisingStage, - GlmImageConditionEncodingStage, - GlmImageDecodingStage, - GlmImageDenoisingStage, -) -from fastvideo.pipelines.stages import InputValidationStage - -logger = init_logger(__name__) - - -class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "GlmImagePipeline" - - _required_config_modules = [ - "text_encoder", - "tokenizer", - "vae", - "transformer", - "scheduler", - "vision_language_encoder", - "processor", - ] - - _optional_config_modules: list[str] = [] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: - self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=1.0) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - self.add_stage( - stage_name="input_validation_stage", - stage=InputValidationStage(), - ) - - self.add_stage( - stage_name="glm_image_before_denoising_stage", - stage=GlmImageBeforeDenoisingStage( - vae=self.get_module("vae"), - text_encoder=self.get_module("text_encoder"), - tokenizer=self.get_module("tokenizer"), - processor=self.get_module("processor"), - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"), - vision_language_encoder=self.get_module("vision_language_encoder"), - ), - ) - - self.add_stage( - stage_name="glm_image_condition_encoding_stage", - stage=GlmImageConditionEncodingStage( - vae=self.get_module("vae"), - transformer=self.get_module("transformer"), - ), - ) - - self.add_stage( - stage_name="denoising_stage", - stage=GlmImageDenoisingStage( - 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=GlmImageDecodingStage( - vae=self.get_module("vae"), - pipeline=self, - ), - ) - - -EntryClass = GlmImagePipeline diff --git a/fastvideo/pipelines/basic/glm_image/stages/__init__.py b/fastvideo/pipelines/basic/glm_image/stages/__init__.py deleted file mode 100644 index c9f13124d..000000000 --- a/fastvideo/pipelines/basic/glm_image/stages/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""GLM-Image pipeline stages.""" - -from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (GlmImageBeforeDenoisingStage) -from fastvideo.pipelines.basic.glm_image.stages.condition_encoding import (GlmImageConditionEncodingStage) -from fastvideo.pipelines.basic.glm_image.stages.decoding import (GlmImageDecodingStage) -from fastvideo.pipelines.basic.glm_image.stages.denoising import (GlmImageDenoisingStage) - -__all__ = [ - "GlmImageBeforeDenoisingStage", - "GlmImageConditionEncodingStage", - "GlmImageDecodingStage", - "GlmImageDenoisingStage", -] diff --git a/fastvideo/pipelines/basic/glm_image/stages/before_denoising.py b/fastvideo/pipelines/basic/glm_image/stages/before_denoising.py deleted file mode 100644 index f3a0331d2..000000000 --- a/fastvideo/pipelines/basic/glm_image/stages/before_denoising.py +++ /dev/null @@ -1,269 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import re -from math import sqrt - -import numpy as np -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.validators import StageValidators as V -from fastvideo.pipelines.stages.validators import VerificationResult - -logger = init_logger(__name__) - - -def calculate_shift( - image_seq_len: int, - base_seq_len: int = 256, - base_shift: float = 0.25, - max_shift: float = 0.75, -) -> float: - return (image_seq_len / base_seq_len)**0.5 * max_shift + base_shift - - -def get_glyph_texts(prompt: str | list[str]) -> list[str] | list[list[str]]: - if isinstance(prompt, str): - prompts: list[str] = [prompt] - is_batch = False - else: - prompts = prompt - is_batch = True - out: list[list[str]] = [] - for p in prompts: - out.append( - re.findall(r"'([^']*)'", p) + re.findall(r"“([^“”]*)”", p) + re.findall(r'"([^"]*)"', p) + - re.findall(r"「([^「」]*)」", p)) - return out if is_batch else out[0] - - -def compute_glyph_embeds( - prompts: list[str], - tokenizer, - text_encoder, - device: torch.device, - dtype: torch.dtype, - max_sequence_length: int = 2048, -) -> torch.Tensor: - all_glyph_texts = get_glyph_texts(prompts) - all_glyph_embeds = [] - for glyph_texts in all_glyph_texts: - if len(glyph_texts) == 0: - glyph_texts = [""] - input_ids = tokenizer( - glyph_texts, - max_length=max_sequence_length, - truncation=True, - ).input_ids - input_ids = [[tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids for ids in input_ids] - max_length = max(len(ids) for ids in input_ids) - attention_mask = torch.tensor( - [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids], - device=device, - ) - input_ids_t = torch.tensor( - [ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids], - device=device, - ) - outputs = text_encoder(input_ids_t, attention_mask=attention_mask) - glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) - all_glyph_embeds.append(glyph_embeds) - - max_seq_len = max(emb.size(1) for emb in all_glyph_embeds) - padded = [] - for emb in all_glyph_embeds: - if emb.size(1) < max_seq_len: - pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype) - emb = torch.cat([pad, emb], dim=1) - padded.append(emb) - return torch.cat(padded, dim=0).to(device=device, dtype=dtype) - - -def _grid_dims(height: int, width: int) -> tuple[int, int, int, int]: - th, tw = height // 32, width // 32 - ratio = th / tw - pth = int(sqrt(ratio) * 16) - ptw = int(sqrt(1 / ratio) * 16) - return th, tw, pth, ptw - - -def _upsample_d32_to_d16(tokens: torch.Tensor, th: int, tw: int) -> torch.Tensor: - tokens = tokens.view(1, 1, th, tw).float() - tokens = torch.nn.functional.interpolate(tokens, scale_factor=2, mode="nearest").long() - return tokens.view(1, -1) - - -class GlmImageBeforeDenoisingStage(PipelineStage): - - def __init__(self, - vae, - text_encoder, - tokenizer, - processor, - transformer, - scheduler, - vision_language_encoder=None) -> None: - super().__init__() - self.vae = vae - self.text_encoder = text_encoder - self.tokenizer = tokenizer - self.processor = processor - self.transformer = transformer - self.scheduler = scheduler - if isinstance(vision_language_encoder, tuple): - self.vision_language_encoder, self.vl_processor = (vision_language_encoder[0], vision_language_encoder[1] - or processor) - else: - self.vision_language_encoder = vision_language_encoder - self.vl_processor = processor - - @torch.no_grad() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - device = get_local_torch_device() - dtype = torch.bfloat16 - th, tw, pth, ptw = _grid_dims(batch.height, batch.width) - - if batch.seed is not None: - torch.manual_seed(batch.seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(batch.seed) - - # 1-3. AR token generation. I2I prepends the condition image and uses a - # single-scale target grid; T2I is multi-scale. - is_t2i = batch.pil_image is None - if self.vision_language_encoder is not None: - content = [{"type": "text", "text": batch.prompt}] - if not is_t2i: - content.insert(0, {"type": "image", "image": batch.pil_image}) - messages = [{"role": "user", "content": content}] - inputs = self.vl_processor.apply_chat_template(messages, - tokenize=True, - target_h=batch.height, - target_w=batch.width, - return_dict=True, - return_tensors="pt").to(device) - - if is_t2i: - up_h, up_w = th, tw - large_start, large_count = pth * ptw, th * tw - max_new = large_count + (pth * ptw) + 1 - else: - # Condition grid(s) first, target grid last. - _, t_h, t_w = inputs["image_grid_thw"][-1].tolist() - up_h, up_w = int(t_h), int(t_w) - large_start, large_count = 0, up_h * up_w - max_new = large_count + 1 - - outputs = self.vision_language_encoder.generate(**inputs, max_new_tokens=max_new, do_sample=True) - gen_tokens = outputs[0][inputs.input_ids.shape[-1]:] - if gen_tokens.shape[0] >= large_start + large_count: - large_tokens = gen_tokens[large_start:large_start + large_count] - else: - available = gen_tokens[large_start:] - large_tokens = torch.zeros(large_count, dtype=gen_tokens.dtype, device=gen_tokens.device) - if available.shape[0] > 0: - large_tokens[:min(available.shape[0], large_count)] = available[:large_count] - logger.warning("AR generated %d tokens, expected %d. Padding with zeros.", gen_tokens.shape[0], - large_start + large_count) - batch.prior_token_id = _upsample_d32_to_d16(large_tokens, up_h, up_w) - batch.prior_token_drop = torch.zeros(batch.prior_token_id.shape, dtype=torch.bool, device=device) - - if not is_t2i: - self._compute_source_prior_tokens(batch, inputs) - else: - num_prior_tokens = 4 * th * tw - logger.warning("No vision_language_encoder provided; using random dropped priors.") - batch.prior_token_id = torch.randint(0, 16384, (1, num_prior_tokens), device=device) - batch.prior_token_drop = torch.ones(batch.prior_token_id.shape, dtype=torch.bool, device=device) - - # 4. Glyph T5 encoding. - prompts = [batch.prompt] if isinstance(batch.prompt, str) else list(batch.prompt) - prompt_embeds = compute_glyph_embeds(prompts, self.tokenizer, self.text_encoder, device, dtype) - - # 5. CFG-side negative encoding. - if batch.do_classifier_free_guidance: - neg_prompts = [batch.negative_prompt or ""] * len(prompts) - neg_embeds = compute_glyph_embeds(neg_prompts, self.tokenizer, self.text_encoder, device, dtype) - L_pos, L_neg = prompt_embeds.shape[1], neg_embeds.shape[1] - max_L = max(L_pos, L_neg) - if L_pos < max_L: - pad = torch.zeros(prompt_embeds.shape[0], - max_L - L_pos, - prompt_embeds.shape[2], - device=device, - dtype=dtype) - prompt_embeds = torch.cat([pad, prompt_embeds], dim=1) - if L_neg < max_L: - pad = torch.zeros(neg_embeds.shape[0], max_L - L_neg, neg_embeds.shape[2], device=device, dtype=dtype) - neg_embeds = torch.cat([pad, neg_embeds], dim=1) - # Row 0 conditional (positive), row 1 unconditional (negative). - prompt_embeds = torch.cat([prompt_embeds, neg_embeds], dim=0) - att_pos = torch.ones((1, max_L), device=device) - att_neg = torch.ones((1, max_L), device=device) - if L_pos < max_L: - att_pos[:, :max_L - L_pos] = 0 - if L_neg < max_L: - att_neg[:, :max_L - L_neg] = 0 - attention_mask = torch.cat([att_pos, att_neg], dim=0) - else: - attention_mask = torch.ones((1, prompt_embeds.shape[1]), device=device) - - batch.prompt_embeds = [prompt_embeds] - batch.attention_mask = attention_mask - - # 6. Latents + dynamic flow shift. - if batch.seed is not None: - torch.manual_seed(batch.seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(batch.seed) - batch.latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8), device=device, dtype=dtype) - - # Integer-cast linspace timesteps with resolution-dependent shift applied to - # sigmas only; the DiT is conditioned on the unshifted integer timesteps. - ntt = self.scheduler.config.num_train_timesteps - patch_size = self.transformer.patch_size - image_seq_len = ((batch.height // 8) * (batch.width // 8)) // (patch_size**2) - sched_timesteps = np.linspace(ntt, 1.0, batch.num_inference_steps + 1)[:-1].astype(np.int64).astype(np.float32) - sched_sigmas = sched_timesteps / ntt - self.scheduler.set_shift(calculate_shift(image_seq_len)) - self.scheduler.set_timesteps(batch.num_inference_steps, - device=device, - sigmas=sched_sigmas.tolist(), - timesteps=sched_timesteps.tolist()) - batch.timesteps = self.scheduler.timesteps - return batch - - @torch.no_grad() - def _compute_source_prior_tokens(self, batch: ForwardBatch, inputs) -> None: - image_grid_thw = inputs["image_grid_thw"] - num_condition_images = image_grid_thw.shape[0] - 1 - source_grids = image_grid_thw[:num_condition_images] - image_features = self.vision_language_encoder.get_image_features(inputs["pixel_values"], source_grids) - image_feature_parts = getattr(image_features, "pooler_output", image_features) - embed = torch.cat(image_feature_parts, dim=0) - src_ids_d32 = self.vision_language_encoder.get_image_tokens(embed, source_grids) - split_sizes = source_grids.prod(dim=-1).tolist() - upsampled = [ - _upsample_d32_to_d16(ids, int(grid[1]), int(grid[2])).squeeze(0) - for ids, grid in zip(torch.split(src_ids_d32, split_sizes), source_grids, strict=False) - ] - src_grids_up = source_grids.clone() - src_grids_up[:, 1] *= 2 - src_grids_up[:, 2] *= 2 - batch.extra["glm_prior_token_image_ids"] = torch.cat(upsampled, dim=0) - batch.extra["glm_source_image_grid_thw"] = src_grids_up - - def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - result = VerificationResult() - result.add_check("prompt", batch.prompt, V.string_not_empty) - return result - - def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - result = VerificationResult() - result.add_check("prior_token_id", batch.prior_token_id, V.is_tensor) - return result diff --git a/fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py b/fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py deleted file mode 100644 index 8fe8ef347..000000000 --- a/fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py +++ /dev/null @@ -1,67 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.image_processor import ImageProcessor -from fastvideo.models.dits.glm_image import GlmImageKVCache -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage - -_CONDITION_MULTIPLE_OF = 16 # vae_scale_factor (8) * DiT patch_size (2) - - -class GlmImageConditionEncodingStage(PipelineStage): - - def __init__(self, vae, transformer) -> None: - super().__init__() - self.vae = vae - self.transformer = transformer - self.image_processor = ImageProcessor(vae_scale_factor=_CONDITION_MULTIPLE_OF) - - @torch.no_grad() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - if batch.pil_image is None: - return batch - - device = get_local_torch_device() - dtype = torch.bfloat16 - self.vae.to(device) - - prior_ids = batch.extra["glm_prior_token_image_ids"].to(device) - if prior_ids.dim() == 1: - prior_ids = prior_ids.unsqueeze(0) - - # Latent patch count must match the source prior tokens; mismatch is fatal. - src_grid = batch.extra["glm_source_image_grid_thw"][0] - cond_h = int(src_grid[1]) * _CONDITION_MULTIPLE_OF - cond_w = int(src_grid[2]) * _CONDITION_MULTIPLE_OF - cond_img = self.image_processor.preprocess(batch.pil_image, cond_h, cond_w).to(device=device, - dtype=torch.float32) - latent = self.vae.encode(cond_img).latent_dist.mode() - # NOTE: at runtime self.vae.config is a diffusers FrozenDict with flat - # latents_mean/latents_std fields. Access the flat fields directly. - cfg = self.vae.config - mean = torch.tensor(cfg.latents_mean, device=device, dtype=torch.float32).view(1, -1, 1, 1) - std = torch.tensor(cfg.latents_std, device=device, dtype=torch.float32).view(1, -1, 1, 1) - latent = ((latent - mean) / std).to(dtype) - - kv_caches = GlmImageKVCache(num_layers=self.transformer.num_layers) - empty_text = batch.prompt_embeds[0][:1, :0, :].to(device=device, dtype=dtype) - with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=batch): - self.transformer( - hidden_states=latent, - encoder_hidden_states=empty_text, - prior_token_id=prior_ids, - prior_token_drop=torch.zeros((prior_ids.shape[0], ), dtype=torch.bool, device=device), - timestep=torch.zeros((1, ), device=device), - target_size=torch.tensor([tuple(cond_img.shape[-2:])], device=device, dtype=torch.long), - crop_coords=torch.zeros((1, 2), device=device, dtype=torch.long), - kv_caches=kv_caches, - kv_caches_mode="write", - ) - batch.extra["glm_kv_caches"] = kv_caches - return batch diff --git a/fastvideo/pipelines/basic/glm_image/stages/decoding.py b/fastvideo/pipelines/basic/glm_image/stages/decoding.py deleted file mode 100644 index 3b78255c8..000000000 --- a/fastvideo/pipelines/basic/glm_image/stages/decoding.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.pipelines.stages.decoding import DecodingStage -from fastvideo.utils import PRECISION_TO_TYPE - - -class GlmImageDecodingStage(DecodingStage): - - @torch.no_grad() - def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor: - self.vae.to(get_local_torch_device()) - latents = latents.to(get_local_torch_device()) - - vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] - vae_autocast = (vae_dtype != torch.float32 and not fastvideo_args.disable_autocast) - - latents = self._denormalize_latents(latents, fastvideo_args) - if latents.dim() == 5: - latents = latents.squeeze(2) - - with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast): - if fastvideo_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - if not vae_autocast: - latents = latents.to(vae_dtype) - decoded = self.vae.decode(latents) - - image = decoded.sample if hasattr(decoded, "sample") else decoded - image = (image / 2 + 0.5).clamp(0, 1) - return image.unsqueeze(2) diff --git a/fastvideo/pipelines/basic/glm_image/stages/denoising.py b/fastvideo/pipelines/basic/glm_image/stages/denoising.py deleted file mode 100644 index 57780138a..000000000 --- a/fastvideo/pipelines/basic/glm_image/stages/denoising.py +++ /dev/null @@ -1,192 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""CFG convention (both denoise paths): row 0 conditional (positive), row 1 unconditional.""" -from __future__ import annotations - -import torch - -from fastvideo.attention.backends.sdpa import SDPAMetadata -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.pipelines.stages.validators import StageValidators as V -from fastvideo.pipelines.stages.validators import VerificationResult -from fastvideo.platforms import AttentionBackendEnum - - -class GlmImageDenoisingStage(DenoisingStage): - - def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - result = VerificationResult() - result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) - latents = getattr(batch, "latent", getattr(batch, "latents", None)) - result.add_check("latents", latents, [V.is_tensor, V.with_dims(5)]) - result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) - return result - - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - device = get_local_torch_device() - dtype = torch.bfloat16 - guidance_scale = batch.guidance_scale - do_cfg = guidance_scale > 1.0 - - latents = getattr(batch, "latent", getattr(batch, "latents", None)) - if latents is None: - raise ValueError("No latents found in batch.") - if latents.dim() == 5: - latents = latents.squeeze(2) - - prompt_embeds = batch.prompt_embeds[0] - text_attention_mask = getattr(batch, "attention_mask", None) - timesteps = batch.timesteps - - patch_size = self.transformer.patch_size - _, _, h, w = latents.shape - image_seq_length = (h // patch_size) * (w // patch_size) - text_seq_length = prompt_embeds.shape[1] if prompt_embeds.dim() >= 2 else 0 - - first_block = self.transformer.transformer_blocks[0] - backend = getattr(first_block.attn1.attn, "backend", None) - sdpa = backend == AttentionBackendEnum.TORCH_SDPA and text_attention_mask is not None - - kv_caches = batch.extra.get("glm_kv_caches") - if kv_caches is None: - self._denoise_t2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale, - text_seq_length, image_seq_length, sdpa, device, dtype) - else: - self._denoise_i2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale, - text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype) - return batch - - def _denoise_t2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale, - text_seq_length, image_seq_length, sdpa, device, dtype) -> None: - num_inference_steps = batch.num_inference_steps - bs = 2 if do_cfg else 1 - target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long).repeat(bs, 1) - crop_coords = torch.zeros((bs, 2), device=device, dtype=torch.long) - - prior_token_id = batch.prior_token_id - if do_cfg and prior_token_id.shape[0] == 1: - prior_token_id = prior_token_id.repeat(2, 1) - if do_cfg: - prior_token_drop = torch.tensor([False, True], device=device) - else: - prior_token_drop = getattr(batch, "prior_token_drop", torch.tensor([False], device=device)) - - attention_mask_kv = None - if sdpa: - if (text_attention_mask.shape[0] == 1 and bs > 1): - text_attention_mask = text_attention_mask.repeat(bs, 1) - mix_attn_mask = torch.ones((bs, text_seq_length + image_seq_length), device=device, dtype=torch.float32) - mix_attn_mask[:, :text_seq_length] = (text_attention_mask.float().to(device)) - attention_mask_kv = (mix_attn_mask > 0).unsqueeze(1).unsqueeze(2) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - latent_model_input = torch.cat([latents] * 2) if do_cfg else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t).to(dtype) - t_expand = t.expand(latent_model_input.shape[0]) - 1 - - attn_metadata = (SDPAMetadata(current_timestep=i, attn_mask=attention_mask_kv) - if attention_mask_kv is not None else None) - - with torch.no_grad(), set_forward_context(current_timestep=i, - attn_metadata=attn_metadata, - forward_batch=batch): - noise_pred = self.transformer( - latent_model_input, - prompt_embeds, - prior_token_id, - prior_token_drop, - t_expand, - target_size, - crop_coords, - ) - - if do_cfg: - noise_pred_cond, noise_pred_uncond = noise_pred.float().chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) - - guidance_rescale = getattr(batch, "guidance_rescale", 0.0) - if guidance_rescale > 0.0: - dims = list(range(1, noise_pred_cond.ndim)) - std_text = noise_pred_cond.std(dim=dims, keepdim=True) - std_cfg = noise_pred.std(dim=dims, keepdim=True) - rescaled = noise_pred * (std_text / std_cfg) - noise_pred = (guidance_rescale * rescaled + (1 - guidance_rescale) * noise_pred) - - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - progress_bar.update() - - batch.latents = latents.unsqueeze(2) - - def _denoise_i2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale, - text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype) -> None: - """Two separate transformer calls (cond reads the cache, uncond skips it): - the cache mode is one global flag with batch-1 k/v, so a 2-row CFG call - cannot express both.""" - num_inference_steps = batch.num_inference_steps - target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long) - crop_coords = torch.zeros((1, 2), device=device, dtype=torch.long) - prior_token_id = batch.prior_token_id[:1] - drop_keep = torch.zeros((1, ), dtype=torch.bool, device=device) - drop_all = torch.ones((1, ), dtype=torch.bool, device=device) - cache_len = kv_caches[0].k_cache.shape[1] if kv_caches[0].k_cache is not None else 0 - - def _mask(row: int, with_cache: bool): - if not sdpa: - return None - prefix = cache_len if with_cache else 0 - m = torch.ones((1, prefix + text_seq_length + image_seq_length), device=device, dtype=torch.float32) - m[:, prefix:prefix + text_seq_length] = text_attention_mask[row:row + 1].float().to(device) - return (m > 0).unsqueeze(1).unsqueeze(2) - - cond_mask = _mask(0, with_cache=True) - uncond_mask = _mask(1, with_cache=False) if do_cfg else None - pos_embeds = prompt_embeds[:1] - neg_embeds = prompt_embeds[1:2] if do_cfg else None - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - latent_model_input = self.scheduler.scale_model_input(latents, t).to(dtype) - t_expand = t.expand(1) - 1 - - cond_meta = (SDPAMetadata(current_timestep=i, attn_mask=cond_mask) if cond_mask is not None else None) - with torch.no_grad(), set_forward_context(current_timestep=i, - attn_metadata=cond_meta, - forward_batch=batch): - noise_pred = self.transformer(latent_model_input, - pos_embeds, - prior_token_id, - drop_keep, - t_expand, - target_size, - crop_coords, - kv_caches=kv_caches, - kv_caches_mode="read") - - if do_cfg: - uncond_meta = (SDPAMetadata(current_timestep=i, attn_mask=uncond_mask) - if uncond_mask is not None else None) - with torch.no_grad(), set_forward_context(current_timestep=i, - attn_metadata=uncond_meta, - forward_batch=batch): - noise_pred_uncond = self.transformer(latent_model_input, - neg_embeds, - prior_token_id, - drop_all, - t_expand, - target_size, - crop_coords, - kv_caches=kv_caches, - kv_caches_mode="skip") - noise_pred = noise_pred_uncond.float() + guidance_scale * (noise_pred.float() - - noise_pred_uncond.float()) - - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - progress_bar.update() - - kv_caches.clear() - batch.latents = latents.unsqueeze(2) diff --git a/fastvideo/pipelines/basic/longcat/__init__.py b/fastvideo/pipelines/basic/longcat/__init__.py deleted file mode 100644 index db24fda1d..000000000 --- a/fastvideo/pipelines/basic/longcat/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""LongCat pipeline module.""" - -from fastvideo.pipelines.basic.longcat.longcat_pipeline import LongCatPipeline -from fastvideo.pipelines.basic.longcat.longcat_i2v_pipeline import LongCatImageToVideoPipeline -from fastvideo.pipelines.basic.longcat.longcat_vc_pipeline import LongCatVideoContinuationPipeline - -__all__ = ["LongCatPipeline", "LongCatImageToVideoPipeline", "LongCatVideoContinuationPipeline"] diff --git a/fastvideo/pipelines/basic/longcat/longcat_i2v_pipeline.py b/fastvideo/pipelines/basic/longcat/longcat_i2v_pipeline.py deleted file mode 100644 index ac37f6c6a..000000000 --- a/fastvideo/pipelines/basic/longcat/longcat_i2v_pipeline.py +++ /dev/null @@ -1,136 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat Image-to-Video pipeline implementation. - -This module implements I2V (Image-to-Video) generation for LongCat using Tier 3 -conditioning with timestep masking, num_cond_latents support, and RoPE skipping. - -Supports: -- Basic I2V (50 steps, guidance_scale=4.0) -- Distilled I2V with LoRA (16 steps, guidance_scale=1.0) -- Refinement I2V for 720p upscaling (with refinement LoRA + BSA) -""" - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import ( - DecodingStage, - InputValidationStage, - TextEncodingStage, - TimestepPreparationStage, -) -from fastvideo.pipelines.stages.longcat_image_vae_encoding import LongCatImageVAEEncodingStage -from fastvideo.pipelines.stages.longcat_i2v_latent_preparation import LongCatI2VLatentPreparationStage -from fastvideo.pipelines.stages.longcat_i2v_denoising import LongCatI2VDenoisingStage -from fastvideo.pipelines.stages.longcat_refine_init import LongCatRefineInitStage -from fastvideo.pipelines.stages.longcat_refine_timestep import LongCatRefineTimestepStage - -logger = init_logger(__name__) - - -class LongCatImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase): - """ - LongCat Image-to-Video pipeline. - - Generates video from a single input image using Tier 3 I2V conditioning: - - Per-frame timestep masking (timestep[:, 0] = 0) - - num_cond_latents parameter to transformer - - RoPE skipping for conditioning frames - - Selective denoising (skip first frame in scheduler) - """ - - _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - """Initialize LongCat-specific components.""" - # Same BSA initialization as base LongCat pipeline - pipeline_config = fastvideo_args.pipeline_config - transformer = self.get_module("transformer", None) - if transformer is None: - return - - # Enable BSA if configured - if pipeline_config.enable_bsa: - bsa_params_cfg = getattr(pipeline_config, 'bsa_params', None) or {} - sparsity = getattr(pipeline_config, 'bsa_sparsity', None) - cdf_threshold = getattr(pipeline_config, 'bsa_cdf_threshold', None) - chunk_q = getattr(pipeline_config, 'bsa_chunk_q', None) - chunk_k = getattr(pipeline_config, 'bsa_chunk_k', None) - - effective_bsa_params = dict(bsa_params_cfg) if isinstance(bsa_params_cfg, dict) else {} - if sparsity is not None: - effective_bsa_params['sparsity'] = sparsity - if cdf_threshold is not None: - effective_bsa_params['cdf_threshold'] = cdf_threshold - if chunk_q is not None: - effective_bsa_params['chunk_3d_shape_q'] = chunk_q - if chunk_k is not None: - effective_bsa_params['chunk_3d_shape_k'] = chunk_k - - # Provide defaults - effective_bsa_params.setdefault('sparsity', 0.9375) - effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4]) - effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4]) - - if hasattr(transformer, 'enable_bsa'): - logger.info("Enabling BSA for LongCat I2V transformer") - transformer.enable_bsa() - if hasattr(transformer, 'blocks'): - try: - for blk in transformer.blocks: - if hasattr(blk, 'self_attn'): - blk.self_attn.bsa_params = effective_bsa_params - except Exception as e: - logger.warning("Failed to set BSA params: %s", e) - logger.info("BSA parameters: %s", effective_bsa_params) - else: - if hasattr(transformer, 'disable_bsa'): - transformer.disable_bsa() - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): - """Set up I2V-specific pipeline stages.""" - - # 1. Input validation - self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) - - # 2. Text encoding (same as T2V) - self.add_stage(stage_name="prompt_encoding_stage", - stage=TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], - )) - - # 3. Image VAE encoding (for I2V - skipped in refinement mode) - self.add_stage(stage_name="image_vae_encoding_stage", - stage=LongCatImageVAEEncodingStage(vae=self.get_module("vae"))) - - # 4. Refinement initialization (skipped if not refining) - self.add_stage(stage_name="longcat_refine_init_stage", stage=LongCatRefineInitStage(vae=self.get_module("vae"))) - - # 5. Timestep preparation (generic) - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage(scheduler=self.get_module("scheduler"))) - - # 6. Refinement timestep override (skipped if not refining) - self.add_stage(stage_name="longcat_refine_timestep_stage", - stage=LongCatRefineTimestepStage(scheduler=self.get_module("scheduler"))) - - # 7. Latent preparation with I2V conditioning - self.add_stage(stage_name="latent_preparation_stage", - stage=LongCatI2VLatentPreparationStage(scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"))) - - # 8. Denoising with I2V support - self.add_stage(stage_name="denoising_stage", - stage=LongCatI2VDenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), - pipeline=self)) - - # 9. Decoding - self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) - - -EntryClass = LongCatImageToVideoPipeline diff --git a/fastvideo/pipelines/basic/longcat/longcat_pipeline.py b/fastvideo/pipelines/basic/longcat/longcat_pipeline.py deleted file mode 100644 index 933660860..000000000 --- a/fastvideo/pipelines/basic/longcat/longcat_pipeline.py +++ /dev/null @@ -1,123 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat video diffusion pipeline implementation. - -This module implements the LongCat video diffusion pipeline using FastVideo's -modular pipeline architecture. -""" - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import ( - DecodingStage, - InputValidationStage, - LatentPreparationStage, - TextEncodingStage, - TimestepPreparationStage, -) -from fastvideo.pipelines.stages.longcat_denoising import LongCatDenoisingStage -from fastvideo.pipelines.stages.longcat_refine_init import LongCatRefineInitStage -from fastvideo.pipelines.stages.longcat_refine_timestep import LongCatRefineTimestepStage - -logger = init_logger(__name__) - - -class LongCatPipeline(LoRAPipeline, ComposedPipelineBase): - """ - LongCat video diffusion pipeline with LoRA support. - """ - - _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - """Initialize LongCat-specific components.""" - - # Enable BSA (Block Sparse Attention) if configured - pipeline_config = fastvideo_args.pipeline_config - transformer = self.get_module("transformer", None) - if transformer is None: - raise RuntimeError("Transformer module not found during initializing LongCat pipeline.") - # If user toggles BSA via CLI/config - if pipeline_config.enable_bsa: - # Build effective BSA params: - # 1) from explicit CLI overrides if provided - # 2) else from pipeline_config.bsa_params - # 3) else fall back to reasonable defaults - bsa_params_cfg = pipeline_config.bsa_params - sparsity = pipeline_config.bsa_sparsity - cdf_threshold = pipeline_config.bsa_cdf_threshold - chunk_q = pipeline_config.bsa_chunk_q - chunk_k = pipeline_config.bsa_chunk_k - - effective_bsa_params = dict(bsa_params_cfg) if isinstance(bsa_params_cfg, dict) else {} - if sparsity is not None: - effective_bsa_params['sparsity'] = sparsity - if cdf_threshold is not None: - effective_bsa_params['cdf_threshold'] = cdf_threshold - if chunk_q is not None: - effective_bsa_params['chunk_3d_shape_q'] = chunk_q - if chunk_k is not None: - effective_bsa_params['chunk_3d_shape_k'] = chunk_k - # Provide defaults if still missing - effective_bsa_params.setdefault('sparsity', 0.9375) - effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4]) - effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4]) - - if hasattr(transformer, 'enable_bsa'): - logger.info("Enabling Block Sparse Attention (BSA) for LongCat transformer") - transformer.enable_bsa() - # Propagate params to all attention modules - if hasattr(transformer, 'blocks'): - try: - for blk in transformer.blocks: - if hasattr(blk, 'self_attn'): - blk.self_attn.bsa_params = effective_bsa_params - except Exception as e: - logger.warning("Failed to set BSA params on all blocks: %s", e) - logger.info("BSA parameters in effect: %s", effective_bsa_params) - else: - logger.warning("BSA is enabled in config but transformer does not support it") - else: - # Explicitly disable if present - if hasattr(transformer, 'disable_bsa'): - transformer.disable_bsa() - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - """Set up pipeline stages with proper dependency injection.""" - - 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")], - )) - - # Add refine initialization stage (will be skipped if not refining) - self.add_stage(stage_name="longcat_refine_init_stage", stage=LongCatRefineInitStage(vae=self.get_module("vae"))) - - # First prepare generic timesteps (for non-refine paths) - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage(scheduler=self.get_module("scheduler"))) - - # Then override timesteps for refinement (will be a no-op if not refining), - # matching LongCat's generate_refine schedule. - self.add_stage(stage_name="longcat_refine_timestep_stage", - stage=LongCatRefineTimestepStage(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="denoising_stage", - stage=LongCatDenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), - pipeline=self)) - - self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) - - -EntryClass = LongCatPipeline diff --git a/fastvideo/pipelines/basic/longcat/longcat_vc_pipeline.py b/fastvideo/pipelines/basic/longcat/longcat_vc_pipeline.py deleted file mode 100644 index 656d16fd1..000000000 --- a/fastvideo/pipelines/basic/longcat/longcat_vc_pipeline.py +++ /dev/null @@ -1,157 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat Video Continuation (VC) pipeline implementation. - -This module implements VC (Video Continuation) generation for LongCat with -KV cache optimization for 2-3x speedup. - -Supports: -- Basic VC (50 steps, guidance_scale=4.0) -- Distilled VC with LoRA (16 steps, guidance_scale=1.0) -- KV cache for conditioning frames -""" - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import ( - DecodingStage, - InputValidationStage, - TextEncodingStage, - TimestepPreparationStage, -) -from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVAEEncodingStage -from fastvideo.pipelines.stages.longcat_i2v_latent_preparation import LongCatI2VLatentPreparationStage -from fastvideo.pipelines.stages.longcat_kv_cache_init import LongCatKVCacheInitStage -from fastvideo.pipelines.stages.longcat_vc_denoising import LongCatVCDenoisingStage - -logger = init_logger(__name__) - - -class LongCatVideoContinuationPipeline(LoRAPipeline, ComposedPipelineBase): - """ - LongCat Video Continuation pipeline. - - Generates video continuation from multiple conditioning frames using - optional KV cache for 2-3x speedup. - - Key features: - - Takes video input (13+ frames typically) - - Encodes conditioning frames via VAE - - Optionally pre-computes KV cache for conditioning - - Uses cached K/V during denoising for speedup - - Concatenates conditioning back after denoising - """ - - _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - """Initialize LongCat-specific components.""" - pipeline_config = fastvideo_args.pipeline_config - transformer = self.get_module("transformer", None) - if transformer is None: - return - - # Enable BSA if configured (for VC, BSA may not be needed) - if getattr(pipeline_config, 'enable_bsa', False): - bsa_params_cfg = getattr(pipeline_config, 'bsa_params', None) or {} - sparsity = getattr(pipeline_config, 'bsa_sparsity', None) - cdf_threshold = getattr(pipeline_config, 'bsa_cdf_threshold', None) - chunk_q = getattr(pipeline_config, 'bsa_chunk_q', None) - chunk_k = getattr(pipeline_config, 'bsa_chunk_k', None) - - effective_bsa_params = dict(bsa_params_cfg) if isinstance(bsa_params_cfg, dict) else {} - if sparsity is not None: - effective_bsa_params['sparsity'] = sparsity - if cdf_threshold is not None: - effective_bsa_params['cdf_threshold'] = cdf_threshold - if chunk_q is not None: - effective_bsa_params['chunk_3d_shape_q'] = chunk_q - if chunk_k is not None: - effective_bsa_params['chunk_3d_shape_k'] = chunk_k - - # Provide defaults - effective_bsa_params.setdefault('sparsity', 0.9375) - effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4]) - effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4]) - - if hasattr(transformer, 'enable_bsa'): - logger.info("Enabling BSA for LongCat VC transformer") - transformer.enable_bsa() - if hasattr(transformer, 'blocks'): - try: - for blk in transformer.blocks: - if hasattr(blk, 'self_attn'): - blk.self_attn.bsa_params = effective_bsa_params - except Exception as e: - logger.warning("Failed to set BSA params: %s", e) - logger.info("BSA parameters: %s", effective_bsa_params) - else: - if hasattr(transformer, 'disable_bsa'): - transformer.disable_bsa() - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): - """Set up VC-specific pipeline stages.""" - - # 1. Input validation - self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) - - # 2. Text encoding - self.add_stage(stage_name="prompt_encoding_stage", - stage=TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], - )) - - # 3. Video VAE encoding (encodes conditioning frames) - self.add_stage(stage_name="video_vae_encoding_stage", - stage=LongCatVideoVAEEncodingStage(vae=self.get_module("vae"))) - - # 4. Timestep preparation - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage(scheduler=self.get_module("scheduler"))) - - # 5. Latent preparation (reuse I2V stage - it handles video_latent too) - self.add_stage(stage_name="latent_preparation_stage", - stage=LongCatVCLatentPreparationStage(scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"))) - - # 6. KV cache initialization (optional, based on config) - # This is always added but will skip if use_kv_cache=False - self.add_stage(stage_name="kv_cache_init_stage", - stage=LongCatKVCacheInitStage(transformer=self.get_module("transformer"))) - - # 7. Denoising with VC and KV cache support - self.add_stage(stage_name="denoising_stage", - stage=LongCatVCDenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), - pipeline=self)) - - # 8. Decoding - self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) - - -class LongCatVCLatentPreparationStage(LongCatI2VLatentPreparationStage): - """ - Prepare latents with video conditioning for first N frames. - - Extends I2V latent preparation to handle video_latent (multiple frames) - instead of image_latent (single frame). - """ - - def forward(self, batch, fastvideo_args): - """Prepare latents with VC conditioning.""" - - # Check if we have video_latent (from VC encoding stage) - video_latent = getattr(batch, 'video_latent', None) - if video_latent is not None: - # Set image_latent to video_latent for parent class compatibility - batch.image_latent = video_latent - - # Call parent class forward - return super().forward(batch, fastvideo_args) - - -EntryClass = LongCatVideoContinuationPipeline diff --git a/fastvideo/pipelines/basic/longcat/presets.py b/fastvideo/pipelines/basic/longcat/presets.py deleted file mode 100644 index 39a07652e..000000000 --- a/fastvideo/pipelines/basic/longcat/presets.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""LongCat model family pipeline presets.""" -from fastvideo.api.presets import InferencePreset, PresetStageSpec - -_DENOISE_STAGE = PresetStageSpec( - name="denoise", - kind="denoising", - description="Main denoising pass", - allowed_overrides=frozenset({ - "num_inference_steps", - "guidance_scale", - }), -) - -LONGCAT_T2V = InferencePreset( - name="longcat_t2v", - version=1, - model_family="longcat", - description="LongCat-Video T2V at 480p", - workload_type="t2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 480, - "width": 848, - "num_frames": 121, - "fps": 24, - "guidance_scale": 1.0, - "num_inference_steps": 50, - }, -) - -LONGCAT_I2V = InferencePreset( - name="longcat_i2v", - version=1, - model_family="longcat", - description="LongCat-Video I2V at 480p", - workload_type="i2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 480, - "width": 848, - "num_frames": 121, - "fps": 24, - "guidance_scale": 1.0, - "num_inference_steps": 50, - }, -) - -LONGCAT_VC = InferencePreset( - name="longcat_vc", - version=1, - model_family="longcat", - description="LongCat-Video continuation at 480p", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 480, - "width": 848, - "num_frames": 121, - "fps": 24, - "guidance_scale": 1.0, - "num_inference_steps": 50, - }, -) - -ALL_PRESETS = (LONGCAT_T2V, LONGCAT_I2V, LONGCAT_VC) diff --git a/fastvideo/pipelines/basic/magi_human/JOURNAL.md b/fastvideo/pipelines/basic/magi_human/JOURNAL.md index 2485fabe8..20d8e92a7 100644 --- a/fastvideo/pipelines/basic/magi_human/JOURNAL.md +++ b/fastvideo/pipelines/basic/magi_human/JOURNAL.md @@ -550,7 +550,7 @@ with hard rules around no-source-residue cleanup. | OQ-5 | **Basic-example output mp4 visual quality is impressionistic at 256x448.** Root cause identified: OQ-6 (pre-existing compounding bf16 drift over the 32-step denoise loop). Wave 2-3 investigation confirmed the 4-step pipeline parity shows 18.85x compounding ratio vs expected 4x linear. See OQ-6 for full details and mitigation candidates. | RESOLVED-ROOT-CAUSE-IDENTIFIED (see OQ-6) | | OQ-6 | **Video patch packing was spatial-major instead of channel-major.** Wave 14 (2026-05-02) ran upstream E2E and got coherent output; FV produced pure noise at same config. Oracle triage identified the bug in `_img2tokens` rearrange order: FV used `(pT pH pW C)` (spatial-major) but the DiT's `video_embedder` Linear weight was trained on the channel-major `(C pT pH pW)` layout that upstream's `UnfoldNd` (grouped-conv reshape, `unfoldNd/unfold.py:66`) produces. The pipeline parity test imported FV's `build_packed_inputs` for both sides at `test_magi_human_pipeline_parity.py:222`, so it consumed equally-permuted tokens on both sides and reported agreement on garbage. Fixed in `latent_preparation.py:_img2tokens` by changing the einops pattern from `(pT pH pW C)` to `(C pT pH pW)`. Validated end-to-end: `examples/inference/basic/basic_magi_human.py` now produces coherent video matching the prompt (woman on park bench reading a book, green trees). Earlier waves' production-side fixes (negative prompt, tokenizer padding, resolution defaults, silent-audio fallback) all still stand. | RESOLVED — Wave 14 | | OQ-11 | **Pipeline parity test imports FV's `build_packed_inputs` for the upstream side.** `tests/local_tests/magi_human/test_magi_human_pipeline_parity.py:222` calls FV's packer for both sides instead of upstream's real `MagiDataProxy.process_input`. This let the channel-major-vs-spatial-major bug (OQ-6, Wave 14) sit silent for weeks because both sides agreed on the wrong layout. Update the parity test to drive the upstream side through `MagiDataProxy.process_input` so future packing-layout regressions are caught at parity time, not at production E2E. | TRACKED FOLLOW-UP | -| OQ-7 | **Wan VAE shared fp32 op-order drift (MEDIUM PRIORITY).** FV uses `z * std + mean` at decode normalization; upstream uses `z / (1/std) + mean`. Bitwise non-equivalent in fp32. Affects all Wan-family pipelines (`fastvideo/configs/pipelines/wan.py`, `turbodiffusion.py`, `longcat.py`, magi-human). Magi VAE test loosened to `atol=1e-3, rtol=1e-3` (Wave 4) to defer. Tighten back to `atol=1e-4` once the Wan VAE op-order fix lands. Fix should be validated against Wan2.1, Wan2.2, and magi-human. Estimated 0.5-1 day to fix and validate. | TRACKED FOLLOW-UP | +| OQ-7 | **Wan VAE shared fp32 op-order drift (MEDIUM PRIORITY).** FV uses `z * std + mean` at decode normalization; upstream uses `z / (1/std) + mean`. Bitwise non-equivalent in fp32. Affects all Wan-family pipelines (`fastvideo/configs/pipelines/wan.py`, magi-human). Magi VAE test loosened to `atol=1e-3, rtol=1e-3` (Wave 4) to defer. Tighten back to `atol=1e-4` once the Wan VAE op-order fix lands. Fix should be validated against Wan2.1, Wan2.2, and magi-human. Estimated 0.5-1 day to fix and validate. | TRACKED FOLLOW-UP | | OQ-9 | **Upstream `dit_module.py` hardcoded bf16 casts block full fp32 parity validation.** `daVinci-MagiHuman/inference/model/dit/dit_module.py` has hardcoded `.to(torch.bfloat16)` casts at lines 619, 507, 650, 694, 696. FV's DiT forward is now dtype-agnostic (Wave 10), but parity tests against upstream still show bf16-noise residual drift in fp32 runs because upstream's intermediate tensors are bf16. Full fp32-clean parity would require patching the local upstream clone or adding a dtype-config flag upstream. Affects only fp32 parity testing, not production. | TRACKED FOLLOW-UP (LOW PRIORITY) | ## Troubleshooting diff --git a/fastvideo/pipelines/basic/turbodiffusion/__init__.py b/fastvideo/pipelines/basic/turbodiffusion/__init__.py deleted file mode 100644 index 0df303a96..000000000 --- a/fastvideo/pipelines/basic/turbodiffusion/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_pipeline import ( - TurboDiffusionPipeline, ) -from fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_i2v_pipeline import ( - TurboDiffusionI2VPipeline, ) - -__all__ = ["TurboDiffusionPipeline", "TurboDiffusionI2VPipeline"] diff --git a/fastvideo/pipelines/basic/turbodiffusion/presets.py b/fastvideo/pipelines/basic/turbodiffusion/presets.py deleted file mode 100644 index d71e18792..000000000 --- a/fastvideo/pipelines/basic/turbodiffusion/presets.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""TurboDiffusion model family pipeline presets.""" -from fastvideo.api.presets import InferencePreset, PresetStageSpec - -_DENOISE_STAGE = PresetStageSpec( - name="denoise", - kind="denoising", - description="Fast few-step denoising pass", - allowed_overrides=frozenset({ - "num_inference_steps", - "guidance_scale", - }), -) - -TURBO_T2V_1_3B = InferencePreset( - name="turbo_t2v_1_3b", - version=1, - model_family="turbodiffusion", - description="TurboWan 2.1 T2V 1.3B (4-step)", - workload_type="t2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 480, - "width": 832, - "num_frames": 81, - "fps": 16, - "guidance_scale": 1.0, - "num_inference_steps": 4, - "negative_prompt": "", - }, -) - -TURBO_T2V_14B = InferencePreset( - name="turbo_t2v_14b", - version=1, - model_family="turbodiffusion", - description="TurboWan 2.1 T2V 14B (4-step)", - workload_type="t2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 720, - "width": 1280, - "num_frames": 81, - "fps": 16, - "guidance_scale": 1.0, - "num_inference_steps": 4, - "negative_prompt": "", - }, -) - -TURBO_I2V_A14B = InferencePreset( - name="turbo_i2v_a14b", - version=1, - model_family="turbodiffusion", - description="TurboWan 2.2 I2V A14B (4-step)", - workload_type="i2v", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 720, - "width": 1280, - "num_frames": 81, - "fps": 16, - "guidance_scale": 1.0, - "num_inference_steps": 4, - "negative_prompt": "", - }, -) - -ALL_PRESETS = ( - TURBO_T2V_1_3B, - TURBO_T2V_14B, - TURBO_I2V_A14B, -) diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py deleted file mode 100644 index 0c591c715..000000000 --- a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -TurboDiffusion I2V (Image-to-Video) Pipeline Implementation. - -This module contains an implementation of the TurboDiffusion I2V pipeline -for 1-4 step image-to-video generation using rCM (recurrent Consistency Model) -sampling with SLA (Sparse-Linear Attention). - -Key differences from T2V: -- Uses dual models (high/low noise) with boundary switching -- sigma_max=200 (vs 80 for T2V) -- Mask conditioning with encoded first frame -""" - -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_rcm import RCMScheduler -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, ImageVAEEncodingStage, InputValidationStage, - LatentPreparationStage, TextEncodingStage, TimestepPreparationStage) - -logger = init_logger(__name__) - - -class TurboDiffusionI2VPipeline(LoRAPipeline, ComposedPipelineBase): - """ - TurboDiffusion I2V pipeline for 1-4 step image-to-video generation. - - Uses RCM scheduler, SLA attention, and dual model switching for - high-quality I2V generation. - """ - - _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "transformer_2", "scheduler"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - # Use RCM scheduler with higher sigma_max for I2V - logger.info("Initializing RCM scheduler for TurboDiffusion I2V (sigma_max=200)") - self.modules["scheduler"] = RCMScheduler(sigma_max=200.0) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - """Set up pipeline stages with proper dependency injection.""" - - 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))) - - # I2V: Encode initial image to latent space - self.add_stage(stage_name="image_latent_preparation_stage", - stage=ImageVAEEncodingStage(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)) - - -EntryClass = TurboDiffusionI2VPipeline diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py deleted file mode 100644 index f8a84a9ed..000000000 --- a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -TurboDiffusion Video Pipeline Implementation. - -This module contains an implementation of the TurboDiffusion video diffusion pipeline -for 1-4 step video generation using rCM (recurrent Consistency Model) sampling -with SLA (Sparse-Linear Attention). -""" - -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_rcm import RCMScheduler -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, LatentPreparationStage, - TextEncodingStage, TimestepPreparationStage) - -logger = init_logger(__name__) - - -class TurboDiffusionPipeline(LoRAPipeline, ComposedPipelineBase): - """ - TurboDiffusion video pipeline for 1-4 step generation. - - Uses RCM scheduler and SLA attention for fast, high-quality video generation. - """ - - _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - # Use RCM scheduler for TurboDiffusion - logger.info("Initializing RCM scheduler for TurboDiffusion") - self.modules["scheduler"] = RCMScheduler(sigma_max=80.0) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - """Set up pipeline stages with proper dependency injection.""" - - 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="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)) - - -EntryClass = TurboDiffusionPipeline diff --git a/fastvideo/pipelines/basic/zimage/__init__.py b/fastvideo/pipelines/basic/zimage/__init__.py deleted file mode 100644 index 9428b7d51..000000000 --- a/fastvideo/pipelines/basic/zimage/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Z-Image pipeline.""" - -from fastvideo.pipelines.basic.zimage.zimage_pipeline import ZImagePipeline - -__all__ = ["ZImagePipeline"] diff --git a/fastvideo/pipelines/basic/zimage/presets.py b/fastvideo/pipelines/basic/zimage/presets.py deleted file mode 100644 index 59e2e554c..000000000 --- a/fastvideo/pipelines/basic/zimage/presets.py +++ /dev/null @@ -1,40 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Z-Image inference presets.""" - -from fastvideo.api.presets import InferencePreset, PresetStageSpec - -_DENOISE_STAGE = PresetStageSpec( - name="denoise", - kind="denoising", - description="Z-Image denoising pass", - allowed_overrides=frozenset({ - "num_inference_steps", - "guidance_scale", - "cfg_normalization", - "cfg_truncation", - }), -) - -ZIMAGE_TURBO = InferencePreset( - name="zimage_turbo", - version=1, - model_family="zimage", - description="Z-Image-Turbo text-to-image generation", - workload_type="t2i", - stage_schemas=(_DENOISE_STAGE, ), - defaults={ - "height": 1024, - "width": 1024, - "num_frames": 1, - "fps": 1, - "seed": 42, - "guidance_scale": 0.0, - "num_inference_steps": 8, - "negative_prompt": "", - "max_sequence_length": 512, - "cfg_normalization": False, - "cfg_truncation": 1.0, - }, -) - -ALL_PRESETS = (ZIMAGE_TURBO, ) diff --git a/fastvideo/pipelines/basic/zimage/stages.py b/fastvideo/pipelines/basic/zimage/stages.py deleted file mode 100644 index 0c33fe437..000000000 --- a/fastvideo/pipelines/basic/zimage/stages.py +++ /dev/null @@ -1,337 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Pipeline stages for the native Z-Image text-to-image path.""" - -from __future__ import annotations - -import inspect - -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.hooks.activation_trace import trace_step -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.input_validation import InputValidationStage -from fastvideo.utils import PRECISION_TO_TYPE - - -class ZImageInputValidationStage(InputValidationStage): - """Validate the image geometry and reproduce the official device RNG.""" - - def _generate_seeds(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args - assert batch.seed is not None - batch.seeds = [batch.seed] - device = get_local_torch_device() - batch.generator = torch.Generator(device=device).manual_seed(batch.seed) - - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - if batch.do_classifier_free_guidance and batch.negative_prompt is None and not batch.negative_prompt_embeds: - batch.negative_prompt = "" - batch = super().forward(batch, fastvideo_args) - if batch.num_frames != 1: - raise ValueError(f"Z-Image is text-to-image and requires num_frames=1, got {batch.num_frames}") - if batch.height is None or batch.width is None: - raise ValueError("Z-Image requires height and width") - if batch.height % 16 or batch.width % 16: - raise ValueError("Z-Image height and width must be divisible by 16; " - f"got {batch.height}x{batch.width}") - return batch - - -class ZImageConditioningStage(PipelineStage): - """Trim padded Qwen states and materialize variable-length CFG streams.""" - - @staticmethod - def _trim_embeddings( - embeds: torch.Tensor, - attention_mask: torch.Tensor | None, - ) -> list[torch.Tensor]: - if attention_mask is None: - return list(embeds.unbind(0)) - return [ - sample[mask.to(device=sample.device, dtype=torch.bool)] - for sample, mask in zip(embeds, attention_mask, strict=True) - ] - - @staticmethod - def _repeat(items: list[torch.Tensor], count: int) -> list[torch.Tensor]: - return [item for item in items for _ in range(count)] - - @torch.no_grad() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - del fastvideo_args - if len(batch.prompt_embeds) != 1: - raise ValueError(f"Z-Image expects one text encoder, got {len(batch.prompt_embeds)}") - - prompt_mask = batch.prompt_attention_mask[0] if batch.prompt_attention_mask else None - prompt_embeds = self._trim_embeddings(batch.prompt_embeds[0], prompt_mask) - batch.extra["zimage_prompt_embeds"] = self._repeat(prompt_embeds, batch.num_videos_per_prompt) - - if batch.do_classifier_free_guidance: - if not batch.negative_prompt_embeds: - raise ValueError("Z-Image CFG requires negative prompt embeddings") - negative_mask = batch.negative_attention_mask[0] if batch.negative_attention_mask else None - negative_embeds = self._trim_embeddings(batch.negative_prompt_embeds[0], negative_mask) - batch.extra["zimage_negative_prompt_embeds"] = self._repeat( - negative_embeds, - batch.num_videos_per_prompt, - ) - else: - batch.extra["zimage_negative_prompt_embeds"] = [] - return batch - - -class ZImageLatentPreparationStage(PipelineStage): - """Create the official fp32 image latents on the transformer device.""" - - def __init__(self, transformer) -> None: - self.transformer = transformer - - @staticmethod - def _randn( - shape: tuple[int, ...], - generators: torch.Generator | list[torch.Generator] | None, - device: torch.device, - ) -> torch.Tensor: - if isinstance(generators, list): - if len(generators) != shape[0]: - raise ValueError(f"generator list length {len(generators)} does not match batch size {shape[0]}") - sample_shape = (1, *shape[1:]) - return torch.cat([ - torch.randn(sample_shape, generator=generator, device=device, dtype=torch.float32) - for generator in generators - ]) - return torch.randn(shape, generator=generators, device=device, dtype=torch.float32) - - @torch.no_grad() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - if batch.height is None or batch.width is None: - raise ValueError("Z-Image requires height and width before latent preparation") - - prompt_embeds = batch.extra.get("zimage_prompt_embeds") - if not isinstance(prompt_embeds, list) or not prompt_embeds: - raise ValueError("Z-Image conditioning must run before latent preparation") - - channels = int(getattr(self.transformer, "in_channels", 16)) - spatial_ratio = int(fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio) - shape = ( - len(prompt_embeds), - channels, - 1, - batch.height // spatial_ratio, - batch.width // spatial_ratio, - ) - device = get_local_torch_device() - if batch.latents is None: - latents = self._randn(shape, batch.generator, device) - else: - latents = batch.latents - if latents.ndim == 4: - latents = latents.unsqueeze(2) - if tuple(latents.shape) != shape: - raise ValueError(f"Expected Z-Image latents with shape {shape}, got {tuple(latents.shape)}") - latents = latents.to(device=device, dtype=torch.float32) - - batch.latents = latents - batch.raw_latent_shape = shape - return batch - - -class ZImageTimestepPreparationStage(PipelineStage): - """Apply the native scheduler's zero endpoint and discrete schedule.""" - - def __init__(self, scheduler) -> None: - self.scheduler = scheduler - - @staticmethod - def _calculate_shift( - image_seq_len: int, - base_seq_len: int, - max_seq_len: int, - base_shift: float, - max_shift: float, - ) -> float: - slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) - return image_seq_len * slope + base_shift - slope * base_seq_len - - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - if batch.latents is None: - raise ValueError("Z-Image latents must be prepared before timesteps") - if batch.timesteps is not None and batch.sigmas is not None: - raise ValueError("Only one of timesteps or sigmas may be supplied") - - scheduler = self.scheduler - sigma_min = float(fastvideo_args.pipeline_config.scheduler_sigma_min) - use_reference_timesteps = bool(fastvideo_args.pipeline_config.scheduler_use_reference_discrete_timesteps) - scheduler.sigma_min = sigma_min - scheduler.register_to_config( - sigma_min=sigma_min, - use_reference_discrete_timesteps=use_reference_timesteps, - ) - config = scheduler.config - image_seq_len = (batch.latents.shape[-2] // 2) * (batch.latents.shape[-1] // 2) - mu = self._calculate_shift( - image_seq_len, - int(config.get("base_image_seq_len", 256)), - int(config.get("max_image_seq_len", 4096)), - float(config.get("base_shift", 0.5)), - float(config.get("max_shift", 1.15)), - ) - device = get_local_torch_device() - - if batch.timesteps is not None: - if "timesteps" not in inspect.signature(scheduler.set_timesteps).parameters: - raise ValueError(f"{type(scheduler).__name__} does not accept custom timesteps") - scheduler.set_timesteps(timesteps=batch.timesteps, device=device, mu=mu) - elif batch.sigmas is not None: - if "sigmas" not in inspect.signature(scheduler.set_timesteps).parameters: - raise ValueError(f"{type(scheduler).__name__} does not accept custom sigmas") - scheduler.set_timesteps(sigmas=batch.sigmas, device=device, mu=mu) - else: - scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu) - - batch.timesteps = scheduler.timesteps - batch.num_inference_steps = len(batch.timesteps) - return batch - - -class ZImageDenoisingStage(PipelineStage): - """Run the native Z-Image flow-matching loop.""" - - performance_component_metric = "transformer_time_s" - - def __init__(self, transformer, scheduler) -> None: - self.transformer = transformer - self.scheduler = scheduler - - @torch.no_grad() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - if batch.latents is None or batch.timesteps is None: - raise ValueError("Z-Image denoising requires latents and timesteps") - - latents = batch.latents.float() - positive = batch.extra.get("zimage_prompt_embeds") - negative = batch.extra.get("zimage_negative_prompt_embeds", []) - if not isinstance(positive, list) or not positive: - raise ValueError("Z-Image denoising requires prompt embeddings") - - device = get_local_torch_device() - target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] - batch_size = latents.shape[0] - - for index, timestep_value in enumerate(batch.timesteps): - if timestep_value.item() == 0 and index == len(batch.timesteps) - 1: - continue - - timestep = timestep_value.expand(batch_size).to(device=device, dtype=torch.float32) - timestep = (1000.0 - timestep) / 1000.0 - current_guidance_scale = float(batch.guidance_scale) - if (batch.do_classifier_free_guidance and batch.cfg_truncation is not None and batch.cfg_truncation <= 1.0 - and timestep[0].item() > batch.cfg_truncation): - current_guidance_scale = 0.0 - apply_cfg = batch.do_classifier_free_guidance and current_guidance_scale > 0.0 - - if apply_cfg: - if not isinstance(negative, list) or len(negative) != batch_size: - raise ValueError("Z-Image CFG requires one negative embedding per image") - model_latents = latents.to(target_dtype).repeat(2, 1, 1, 1, 1) - model_embeddings = positive + negative - model_timestep = timestep.repeat(2) - else: - model_latents = latents.to(target_dtype) - model_embeddings = positive - model_timestep = timestep - - with ( - torch.autocast( - device_type=device.type, - enabled=False, - ), - trace_step(index), - set_forward_context( - current_timestep=index, - attn_metadata=None, - forward_batch=batch, - ), - ): - model_outputs = self.transformer( - hidden_states=model_latents, - encoder_hidden_states=model_embeddings, - timestep=model_timestep, - )[0] - - if apply_cfg: - positive_outputs = model_outputs[:batch_size] - negative_outputs = model_outputs[batch_size:] - guided_outputs = [] - for positive_output, negative_output in zip( - positive_outputs, - negative_outputs, - strict=True, - ): - positive_fp32 = positive_output.float() - prediction = positive_fp32 + current_guidance_scale * (positive_fp32 - negative_output.float()) - if batch.cfg_normalization: - positive_norm = torch.linalg.vector_norm(positive_fp32) - prediction_norm = torch.linalg.vector_norm(prediction) - if prediction_norm > positive_norm: - prediction = prediction * (positive_norm / prediction_norm) - guided_outputs.append(prediction) - noise_pred = torch.stack(guided_outputs) - else: - noise_pred = torch.stack([output.float() for output in model_outputs]) - - noise_pred = -noise_pred - latents = self.scheduler.step( - noise_pred, - timestep_value, - latents, - return_dict=False, - )[0].float() - - batch.latents = latents - return batch - - -class ZImageDecodingStage(PipelineStage): - """Apply the official latent transform and decode one image frame.""" - - performance_component_metric = "vae_decode_time_s" - - def __init__(self, vae) -> None: - self.vae = vae - - @torch.no_grad() - def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - if batch.latents is None: - raise ValueError("Z-Image decoding requires latents") - if fastvideo_args.output_type == "latent": - # FastVideo standardizes image and video latents as [B,C,T,H,W]. - # Tongyi's image-only API returns the equivalent tensor with T squeezed. - batch.output = batch.latents - return batch - - latents = batch.latents - if latents.ndim != 5 or latents.shape[2] != 1: - raise ValueError(f"Expected Z-Image latents [B,C,1,H,W], got {tuple(latents.shape)}") - - device = get_local_torch_device() - self.vae = self.vae.to(device) - vae_dtype = getattr(self.vae, "dtype", None) - if vae_dtype is None: - vae_dtype = next(self.vae.parameters()).dtype - config = self.vae.config - scaling_factor = float(config.scaling_factor) - shift_factor = float(config.shift_factor or 0.0) - latents_2d = latents.squeeze(2).to(device=device, dtype=vae_dtype) - latents_2d = latents_2d / scaling_factor + shift_factor - decoded = self.vae.decode(latents_2d, return_dict=False)[0] - decoded = (decoded / 2 + 0.5).clamp(0, 1) - batch.output = decoded.unsqueeze(2).float() - - if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") - return batch diff --git a/fastvideo/pipelines/basic/zimage/zimage_pipeline.py b/fastvideo/pipelines/basic/zimage/zimage_pipeline.py deleted file mode 100644 index ecf897c78..000000000 --- a/fastvideo/pipelines/basic/zimage/zimage_pipeline.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase -from fastvideo.pipelines.stages.text_encoding import TextEncodingStage - -from .stages import ( - ZImageConditioningStage, - ZImageDecodingStage, - ZImageDenoisingStage, - ZImageInputValidationStage, - ZImageLatentPreparationStage, - ZImageTimestepPreparationStage, -) - - -class ZImagePipeline(ComposedPipelineBase): - """Native Z-Image text-to-image pipeline.""" - - pipeline_config_cls: type[ZImagePipelineConfig] = ZImagePipelineConfig - _required_config_modules = [ - "scheduler", - "text_encoder", - "tokenizer", - "transformer", - "vae", - ] - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - scheduler = self.get_module("scheduler") - transformer = self.get_module("transformer") - - self.add_stage("input_validation_stage", ZImageInputValidationStage()) - self.add_stage( - "text_encoding_stage", - TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], - ), - ) - self.add_stage("zimage_conditioning_stage", ZImageConditioningStage()) - self.add_stage( - "latent_preparation_stage", - ZImageLatentPreparationStage(transformer=transformer), - ) - self.add_stage( - "timestep_preparation_stage", - ZImageTimestepPreparationStage(scheduler=scheduler), - ) - self.add_stage( - "denoising_stage", - ZImageDenoisingStage(transformer=transformer, scheduler=scheduler), - ) - self.add_stage( - "decoding_stage", - ZImageDecodingStage(vae=self.get_module("vae")), - ) - - -EntryClass = ZImagePipeline diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 7f85149df..c4a16d0e9 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -97,8 +97,8 @@ class ComposedPipelineBase(ABC): # component and loading it again is only safe when nothing outside the # loader has changed it. Two habits break that and neither raises: # - # * mutating a component after load. ``LongCatPipeline.initialize_pipeline`` - # turns on block-sparse attention and writes parameters into every + # * mutating a component after load, e.g. an ``initialize_pipeline`` that + # turns on a sparse-attention mode and writes parameters into every # transformer block. That runs once, so a re-materialized component # silently comes back with the feature off. # * reading a component's attributes while building stages. The shared diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 2a7d2533e..83289a586 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -92,12 +92,9 @@ class ForwardBatch: video_path: str | None = None video_latent: torch.Tensor | None = None - # Refine inputs (LongCat) - refine_from: str | None = None + # Refine / conditioning-frame inputs t_thresh: float = 0.5 - spatial_refine_only: bool = False num_cond_frames: int = 0 - stage1_video: list[PIL.Image.Image] | None = None # Loaded frames from refine_from # Primary encoder embeddings prompt_embeds: list[torch.Tensor] = field(default_factory=list) @@ -145,13 +142,6 @@ class ForwardBatch: # Camera control inputs (HYWorld) pose: str | None = None # Camera trajectory: pose string (e.g., 'w-31') or JSON file path - # Camera/action control inputs (GameCraft) - camera_states: torch.Tensor | None = None # Plücker coordinates [B, T, 6, H, W] - gt_latents: torch.Tensor | None = None # Ground truth latents for conditioning [B, 16, T, H, W] - conditioning_mask: torch.Tensor | None = None # Mask for conditioning [B, 1, T, H, W] - camera_trajectory: str | None = None # Camera trajectory file/identifier - action_list: list[str] | None = None # List of actions (e.g., ['forward', 'left']) - action_speed_list: list[float] | None = None # Speed for each action # Camera control inputs (LingBotWorld and LingBotWorld2) c2ws_plucker_emb: torch.Tensor | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat] action_path: str | None = None # Directory containing poses.npy and intrinsics.npy @@ -185,8 +175,6 @@ class ForwardBatch: 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 eta: float = 0.0 sigmas: list[float] | None = None diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 4667d2216..3b43658be 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -19,7 +19,6 @@ from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, Matri MatrixGame3ImageVAEEncodingStage, RefImageEncodingStage, ImageVAEEncodingStage, VideoVAEEncodingStage, Hy15ImageEncodingStage, HYWorldImageEncodingStage) -from fastvideo.pipelines.stages.gamecraft_image_encoding import (GameCraftImageVAEEncodingStage) from fastvideo.pipelines.stages.input_validation import InputValidationStage from fastvideo.pipelines.stages.latent_preparation import (Cosmos25LatentPreparationStage, CosmosLatentPreparationStage, Cosmos25AutoLatentPreparationStage, @@ -37,17 +36,11 @@ from fastvideo.pipelines.stages.kandinsky5 import (Kandinsky5DecodingStage, Kand from fastvideo.pipelines.stages.kandinsky6 import (Kandinsky6AudioDecodingStage, Kandinsky6DecodingStage, Kandinsky6DenoisingStage, Kandinsky6ImageEncodingStage, Kandinsky6LatentPreparationStage) -from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage from fastvideo.pipelines.stages.gen3c_stages import (Gen3CCFGPolicyStage, Gen3CConditioningStage, Gen3CDenoisingStage, Gen3CLatentPreparationStage) from fastvideo.pipelines.stages.text_encoding import (Cosmos25TextEncodingStage, TextEncodingStage) from fastvideo.pipelines.stages.timestep_preparation import (Cosmos25TimestepPreparationStage, TimestepPreparationStage) -# LongCat stages -from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVAEEncodingStage -from fastvideo.pipelines.stages.longcat_kv_cache_init import LongCatKVCacheInitStage -from fastvideo.pipelines.stages.longcat_vc_denoising import LongCatVCDenoisingStage - __all__ = [ "PipelineStage", "InputValidationStage", @@ -77,7 +70,6 @@ __all__ = [ "Kandinsky6DenoisingStage", "Kandinsky6ImageEncodingStage", "Kandinsky6LatentPreparationStage", - "GameCraftDenoisingStage", "Gen3CCFGPolicyStage", "Gen3CConditioningStage", "Gen3CLatentPreparationStage", @@ -101,13 +93,8 @@ __all__ = [ "RefImageEncodingStage", "ImageVAEEncodingStage", "VideoVAEEncodingStage", - "GameCraftImageVAEEncodingStage", "TextEncodingStage", "Cosmos25TextEncodingStage", - # LongCat stages - "LongCatVideoVAEEncodingStage", - "LongCatKVCacheInitStage", - "LongCatVCDenoisingStage", ] diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index b95abd287..27b5160a3 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -320,16 +320,6 @@ class DecodingStage(PipelineStage): # Convert to float32 for compatibility frames = frames.to(torch.float32) - # Crop padding if this is a LongCat refinement - if hasattr(batch, 'num_cond_frames_added') and hasattr(batch, 'new_frame_size_before_padding'): - num_cond_frames_added = batch.num_cond_frames_added - new_frame_size = batch.new_frame_size_before_padding - if num_cond_frames_added > 0 or frames.shape[2] != new_frame_size: - # frames is [B, C, T, H, W], crop temporal dimension - frames = frames[:, :, num_cond_frames_added:num_cond_frames_added + new_frame_size, :, :] - logger.info("Cropped LongCat refinement padding: %s:%s, final shape: %s", num_cond_frames_added, - num_cond_frames_added + new_frame_size, frames.shape) - # Update batch with decoded image batch.output = frames diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 5c867c969..f44d344a6 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -195,26 +195,6 @@ class DenoisingStage(PipelineStage): }, ) - camera_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "camera_states": batch.camera_states, - }, - ) - - dreamx_y_camera = batch.extra.get("dreamx_y_camera", batch.extra.get("y_camera")) - if isinstance(dreamx_y_camera, dict): - dreamx_y_camera = { - key: value.to(device=local_device, dtype=target_dtype) if torch.is_tensor(value) else value - for key, value in dreamx_y_camera.items() - } - dreamx_camera_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "y_camera": dreamx_y_camera, - }, - ) - for key in ("flux2_txt_ids", "flux2_img_ids"): value = batch.extra.get(key) if torch.is_tensor(value): @@ -417,8 +397,6 @@ class DenoisingStage(PipelineStage): **image_kwargs, **pos_cond_kwargs, **action_kwargs, - **camera_kwargs, - **dreamx_camera_kwargs, **timesteps_r_kwarg, **flux2_id_kwargs, ) @@ -461,8 +439,6 @@ class DenoisingStage(PipelineStage): **image_kwargs, **neg_cond_kwargs, **action_kwargs, - **camera_kwargs, - **dreamx_camera_kwargs, **timesteps_r_kwarg, **flux2_id_kwargs, ) diff --git a/fastvideo/pipelines/stages/gamecraft_denoising.py b/fastvideo/pipelines/stages/gamecraft_denoising.py deleted file mode 100644 index 3bac8d95a..000000000 --- a/fastvideo/pipelines/stages/gamecraft_denoising.py +++ /dev/null @@ -1,318 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -GameCraft denoising stage for camera/action-conditioned video generation. - -This stage implements the denoising loop for HunyuanGameCraft, which generates -game-like videos with camera and action conditioning via: -1. CameraNet - Encodes Plücker coordinates into features added to image embeddings -2. Concatenated input - 33 channels (16 latent + 16 gt_latent + 1 mask) -3. Mask-based conditioning for autoregressive generation -""" - -import torch - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger -from fastvideo.models.loader.component_loader import TransformerLoader -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.pipelines.stages.validators import StageValidators as V -from fastvideo.pipelines.stages.validators import VerificationResult - -logger = init_logger(__name__) - - -class GameCraftDenoisingStage(DenoisingStage): - """ - Denoising stage for HunyuanGameCraft with camera/action conditioning. - - This stage handles: - - Camera state encoding via CameraNet (Plücker coordinates) - - Concatenation of latents with gt_latents and mask (33 channels) - - Flow matching denoising with camera conditioning - - Support for autoregressive generation with history frames - """ - - def __init__( - self, - transformer, - scheduler, - pipeline=None, - transformer_2=None, - vae=None, - ) -> None: - super().__init__(transformer, scheduler, pipeline, transformer_2, vae) - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """ - Run the denoising loop with camera/action conditioning. - - Args: - batch: The current batch information. Must contain: - - latents: Noise latents [B, 16, T, H, W] - - camera_states: Plücker coordinates [B, T_video, 6, H_video, W_video] - - gt_latents (optional): Ground truth latents for conditioning [B, 16, T, H, W] - - conditioning_mask (optional): Mask for conditioning [B, 1, T, H, W] - fastvideo_args: The inference arguments. - - Returns: - The batch with denoised latents. - """ - pipeline = self.pipeline() if self.pipeline else None - if not fastvideo_args.model_loaded["transformer"]: - loader = TransformerLoader() - self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) - if pipeline: - pipeline.add_module("transformer", self.transformer) - fastvideo_args.model_loaded["transformer"] = True - - # Extract GameCraft-specific parameters - camera_states = getattr(batch, "camera_states", None) - if camera_states is None: - camera_states = batch.extra.get("camera_states", None) - - gt_latents = getattr(batch, "gt_latents", None) - if gt_latents is None: - gt_latents = batch.extra.get("gt_latents", None) - - conditioning_mask = getattr(batch, "conditioning_mask", None) - if conditioning_mask is None: - conditioning_mask = batch.extra.get("conditioning_mask", None) - - # Prepare extra step kwargs for scheduler - extra_step_kwargs = self.prepare_extra_func_kwargs( - self.scheduler.step, - { - "generator": batch.generator, - "eta": batch.eta, - }, - ) - - # Setup precision and autocast settings - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - - # Get timesteps - timesteps = batch.timesteps - if timesteps is None: - raise ValueError("Timesteps must be provided") - num_inference_steps = batch.num_inference_steps - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # Prepare image embeddings for I2V generation (if any) - image_embeds = batch.image_embeds - if len(image_embeds) > 0: - assert not torch.isnan(image_embeds[0]).any(), "image_embeds contains nan" - image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds] - - image_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - {"encoder_hidden_states_image": image_embeds}, - ) - - pos_cond_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "encoder_hidden_states_2": batch.clip_embedding_pos, - "encoder_attention_mask": batch.prompt_attention_mask, - }, - ) - - neg_cond_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "encoder_hidden_states_2": batch.clip_embedding_neg, - "encoder_attention_mask": batch.negative_attention_mask, - }, - ) - - # Get latents and embeddings - latents = batch.latents - prompt_embeds = batch.prompt_embeds - assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan" - - if batch.do_classifier_free_guidance: - neg_prompt_embeds = batch.negative_prompt_embeds - assert neg_prompt_embeds is not None - assert not torch.isnan(neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan" - - # Prepare gt_latents and mask for concatenation - # If not provided, use zeros (for unconditional generation) - gt_latents = torch.zeros_like(latents) if gt_latents is None else gt_latents.to(target_dtype) - - if conditioning_mask is None: - # Default mask: all zeros (generate everything) - conditioning_mask = torch.zeros( - latents.shape[0], - 1, - *latents.shape[2:], - device=latents.device, - dtype=target_dtype, - ) - else: - conditioning_mask = conditioning_mask.to(target_dtype) - - # Move camera states to device if provided - if camera_states is not None: - camera_states = camera_states.to(device=latents.device, dtype=target_dtype) - - # Debug logging - logger.debug("[GameCraft DEBUG] latents shape: %s, min/max: %.4f/%.4f", latents.shape, latents.min(), - latents.max()) - logger.debug("[GameCraft DEBUG] camera_states: %s", camera_states.shape if camera_states is not None else None) - logger.debug("[GameCraft DEBUG] prompt_embeds[0] shape: %s", prompt_embeds[0].shape) - - # Initialize lists for trajectory - trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] - - # For I2V: extract clean reference latents to inject at every step. - # Official GameCraft replaces conditioned frames of the denoising latents - # with the clean reference at EVERY timestep (not just the first). - # This keeps the conditioned frames noise-free so the model has a strong - # reference signal throughout the denoising process. - # - # ref_latent_for_injection: [B, 16, 1, H, W] or None - ref_latent_for_injection = getattr(batch, "_ref_latent_for_injection", None) - if (ref_latent_for_injection is None and gt_latents is not None and conditioning_mask.sum() > 0 - and gt_latents[:, :, 0].abs().sum() > 0): - # Extract the clean reference from gt_latents first frame - # gt_latents[:, :, 0] should have the VAE-encoded reference image - ref_latent_for_injection = gt_latents[:, :, 0:1].clone() # [B, 16, 1, H, W] - logger.info("[GameCraft I2V] Will inject ref latent at conditioned frames each step. " - "ref mean=%.4f", - ref_latent_for_injection.abs().mean()) - - # Run denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Skip if interrupted - if hasattr(self, "interrupt") and self.interrupt: - continue - - current_model = self.transformer - current_guidance_scale = batch.guidance_scale - - # I2V: replace conditioned frames with clean reference latent - # (matches official: latents[:,:,0,:,:] = last_latents[:,:,-1,:,:]) - if ref_latent_for_injection is not None: - # conditioning_mask shape: [B, 1, T, H, W] - # Where mask == 1, replace latents with the clean ref - cond_frames = conditioning_mask[0, 0, :, 0, 0] > 0.5 # [T] - for t_idx in range(cond_frames.shape[0]): - if cond_frames[t_idx]: - latents[:, :, t_idx, :, :] = ref_latent_for_injection[:, :, 0, :, :] - - # Prepare model input: concatenate latents, gt_latents, and mask - # [B, 33, T, H, W] = [B, 16, T, H, W] + [B, 16, T, H, W] + [B, 1, T, H, W] - latent_model_input = torch.cat( - [latents.to(target_dtype), gt_latents, conditioning_mask], - dim=1, - ) - - assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan" - - t_expand = t.repeat(latent_model_input.shape[0]) - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # Official GameCraft does NOT use embedded guidance (guidance=None) - # It uses standard CFG with guidance_scale instead - guidance_expand = None - - # Run transformer with camera conditioning - with torch.autocast( - device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled, - ): - batch.is_cfg_negative = False - with set_forward_context( - current_timestep=i, - attn_metadata=None, - forward_batch=batch, - ): - noise_pred = current_model( - latent_model_input, - prompt_embeds, - t_expand, - camera_states=camera_states, - guidance=guidance_expand, - **image_kwargs, - **pos_cond_kwargs, - ) - - # Debug: log first step output - if i == 0: - logger.info("[GameCraft DEBUG] Step 0 noise_pred: shape=%s, min/max=%.4f/%.4f", - noise_pred.shape, noise_pred.min(), noise_pred.max()) - - # Classifier-free guidance - if batch.do_classifier_free_guidance: - batch.is_cfg_negative = True - with set_forward_context( - current_timestep=i, - attn_metadata=None, - forward_batch=batch, - ): - noise_pred_uncond = current_model( - latent_model_input, - neg_prompt_embeds, - t_expand, - camera_states=camera_states, - guidance=guidance_expand, - **image_kwargs, - **neg_cond_kwargs, - ) - - noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred - noise_pred_uncond) - - # Compute the previous noisy sample - latents = self.scheduler.step( - noise_pred, - t, - latents, - **extra_step_kwargs, - return_dict=False, - )[0] - - # Store trajectory if requested - if batch.return_trajectory_latents: - trajectory_timesteps.append(t.clone()) - trajectory_latents.append(latents.clone()) - - # Update progress bar - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - - # Debug: log final latents - logger.info("[GameCraft DEBUG] Final latents: shape=%s, min/max=%.4f/%.4f", latents.shape, latents.min(), - latents.max()) - - # Store final latents and trajectory - batch.latents = latents - if batch.return_trajectory_latents: - batch.trajectory_timesteps = trajectory_timesteps - batch.trajectory_latents = torch.stack(trajectory_latents, dim=0) - - return batch - - def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - """Verify that required inputs are present.""" - result = VerificationResult() - result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) - result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) - result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) - return result - - def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - """Verify that outputs are properly set.""" - result = VerificationResult() - result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) - return result diff --git a/fastvideo/pipelines/stages/gamecraft_image_encoding.py b/fastvideo/pipelines/stages/gamecraft_image_encoding.py deleted file mode 100644 index 3f58e8bb4..000000000 --- a/fastvideo/pipelines/stages/gamecraft_image_encoding.py +++ /dev/null @@ -1,173 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -GameCraft image-to-video encoding stage. - -Encodes a reference image into gt_latents and conditioning_mask for -HunyuanGameCraft I2V generation. For T2V this stage is a no-op. -""" - -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.validators import VerificationResult -from fastvideo.utils import PRECISION_TO_TYPE - -logger = init_logger(__name__) - - -class GameCraftImageVAEEncodingStage(PipelineStage): - """ - Stage for encoding a reference image into gt_latents and conditioning_mask - for HunyuanGameCraft image-to-video generation. - - Official GameCraft I2V flow: - 1. VAE-encode the reference image -> [B, 16, 1, H_lat, W_lat] - 2. Scale by VAE scaling_factor (0.476986) - 3. Repeat to all temporal frames - 4. Zero out non-conditioned frames (first frame only for short videos, - first half for longer autoregressive generation) - 5. Build a binary mask (1 = conditioned, 0 = generate) - 6. Store gt_latents and conditioning_mask on the batch for the denoising stage - - If no image is provided (T2V mode), this stage is a no-op; the denoising - stage already falls back to zero gt_latents and zero mask. - """ - - def __init__(self, vae) -> None: - super().__init__() - self.vae = vae - - @torch.no_grad() - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Encode reference image for I2V, or skip for T2V.""" - - if batch.pil_image is None: - # T2V mode: nothing to do; denoising stage handles the fallback - return batch - - device = get_local_torch_device() - image = batch.pil_image - - # ------------------------------------------------------------------ - # 1. Preprocess image to tensor [B, 3, 1, H, W] (values in [-1, 1]) - # ------------------------------------------------------------------ - from torchvision import transforms - - target_height, target_width = batch.height, batch.width - - if isinstance(image, torch.Tensor): - # Already a tensor (e.g. from InputValidationStage causal path) - if image.dim() == 5: - # [B, C, F, H, W] -> take first frame - ref_pixel = image[:, :, :1] - elif image.dim() == 4: - ref_pixel = image.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] - else: - raise ValueError(f"Unexpected image tensor dims: {image.dim()}") - ref_pixel = ref_pixel.to(device=device, dtype=torch.float32) - else: - # PIL Image – resize, center-crop, normalize to [-1, 1] - from PIL import Image as PILImage - if not isinstance(image, PILImage.Image): - import numpy as np - if isinstance(image, np.ndarray): - image = PILImage.fromarray(image) - - original_w, original_h = image.size - scale = max(target_width / original_w, target_height / original_h) - resize_w = int(round(original_w * scale)) - resize_h = int(round(original_h * scale)) - - ref_transform = transforms.Compose([ - transforms.Resize( - (resize_h, resize_w), - interpolation=transforms.InterpolationMode.LANCZOS, - ), - transforms.CenterCrop((target_height, target_width)), - transforms.ToTensor(), - transforms.Normalize([0.5], [0.5]), - ]) - ref_pixel = ref_transform(image) # [3, H, W] - ref_pixel = ref_pixel.unsqueeze(0).unsqueeze(2) # [1, 3, 1, H, W] - ref_pixel = ref_pixel.to(device=device, dtype=torch.float32) - - # ------------------------------------------------------------------ - # 2. VAE-encode - # ------------------------------------------------------------------ - self.vae = self.vae.to(device) - - vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - - with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): - if fastvideo_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - if not vae_autocast_enabled: - ref_pixel = ref_pixel.to(vae_dtype) - encoder_output = self.vae.encode(ref_pixel) - - # Sample from the distribution (official uses .sample()) - ref_latents = encoder_output.latent_dist.sample().to(dtype=torch.float32) - # Scale by VAE scaling factor (0.476986 for GameCraft) - ref_latents.mul_(self.vae.config.scaling_factor) - # ref_latents: [B, 16, 1, H_lat, W_lat] - - # ------------------------------------------------------------------ - # 3. Build gt_latents + conditioning_mask - # ------------------------------------------------------------------ - # Get latent temporal dimension from batch latents (set by LatentPreparationStage) - latent_frames = batch.latents.shape[2] if batch.latents is not None else ((batch.num_frames - 1) // 4 + 1) - - # Repeat to all frames - gt_latents = ref_latents.repeat(1, 1, latent_frames, 1, 1) - # [B, 16, T_lat, H_lat, W_lat] - - # Mask construction following official GameCraft logic: - # - Short videos (latent_frames <= 10): first frame conditioned - # - Longer videos: first half conditioned (autoregressive) - mask = torch.ones( - gt_latents.shape[0], - 1, - gt_latents.shape[2], - gt_latents.shape[3], - gt_latents.shape[4], - device=gt_latents.device, - dtype=gt_latents.dtype, - ) - - if latent_frames <= 10: - # I2V: only first frame is conditioned - gt_latents[:, :, 1:, :, :] = 0.0 - mask[:, :, 1:, :, :] = 0.0 - else: - # Autoregressive: first half conditioned - half = latent_frames // 2 - gt_latents[:, :, half:, :, :] = 0.0 - mask[:, :, half:, :, :] = 0.0 - - batch.gt_latents = gt_latents.to(device=device) - batch.conditioning_mask = mask.to(device=device) - - # Offload - if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") - - return batch - - def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - result = VerificationResult() - # Stage is a no-op when pil_image is None, so nothing required - return result - - def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - result = VerificationResult() - # gt_latents and conditioning_mask are only set for I2V - return result diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index c01e6c7f8..04b0e7a19 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -137,12 +137,8 @@ class LatentPreparationStage(PipelineStage): if hasattr(self.scheduler, "init_noise_sigma"): latents = latents * self.scheduler.init_noise_sigma else: - # Pre-initialized latents: - # - For LongCat refine (refine_from or stage1_video present), we should not re-scale by init_noise_sigma. - # - For other models, keep the original behavior. latents = latents.to(device) - is_longcat_refine = (batch.refine_from is not None) or (batch.stage1_video is not None) - if (not is_longcat_refine) and hasattr(self.scheduler, "init_noise_sigma"): + if hasattr(self.scheduler, "init_noise_sigma"): latents = latents * self.scheduler.init_noise_sigma # Update batch with prepared latents diff --git a/fastvideo/pipelines/stages/longcat_denoising.py b/fastvideo/pipelines/stages/longcat_denoising.py deleted file mode 100644 index ef2363824..000000000 --- a/fastvideo/pipelines/stages/longcat_denoising.py +++ /dev/null @@ -1,161 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat-specific denoising stage implementing CFG-zero optimized guidance. -""" - -import torch -from tqdm import tqdm - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.forward_context import set_forward_context - -logger = init_logger(__name__) - - -class LongCatDenoisingStage(DenoisingStage): - """ - LongCat denoising stage with CFG-zero optimized guidance scale. - - Implements: - 1. Optimized CFG scale from CFG-zero paper - 2. Negation of noise prediction before scheduler step (flow matching convention) - 3. Batched CFG computation (unlike standard FastVideo separate passes) - """ - - def optimized_scale(self, positive_flat, negative_flat) -> torch.Tensor: - """ - Calculate optimized scale from CFG-zero paper. - - st_star = (v_cond^T * v_uncond) / ||v_uncond||^2 - - Args: - positive_flat: Conditional prediction, flattened [B, -1] - negative_flat: Unconditional prediction, flattened [B, -1] - - Returns: - st_star: Optimized scale [B, 1] - """ - # Calculate dot product - dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) - # Squared norm of uncondition - squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 - # st_star = v_cond^T * v_uncond / ||v_uncond||^2 - st_star = dot_product / squared_norm - return st_star - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """ - Run LongCat denoising loop with optimized CFG. - - Args: - batch: The current batch information. - fastvideo_args: The inference arguments. - - Returns: - The batch with denoised latents. - """ - if not fastvideo_args.model_loaded["transformer"]: - from fastvideo.models.loader.component_loader import TransformerLoader - loader = TransformerLoader() - self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) - pipeline = self.pipeline() if self.pipeline else None - if pipeline: - pipeline.add_module("transformer", self.transformer) - fastvideo_args.model_loaded["transformer"] = True - - # Inference dtype. We hardcode bf16 (matching the WanDenoisingStage - # pattern) rather than reading transformer.parameters().dtype: when - # a model is loaded with default_dtype=fp32 but its FSDP-wrapped - # submodules compute in bf16, the parameter-dtype heuristic - # mismatches the Conv3d weight/bias dtype and the patch_embed - # forward fails with "Input type (float) and bias type - # (c10::BFloat16) should be the same". - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - - # Extract batch parameters - latents = batch.latents - timesteps = batch.timesteps - prompt_embeds = batch.prompt_embeds[0] # LongCat uses single encoder - prompt_attention_mask = batch.prompt_attention_mask[0] if batch.prompt_attention_mask else None - guidance_scale = batch.guidance_scale - do_classifier_free_guidance = batch.do_classifier_free_guidance - - # Get negative prompts if doing CFG - if do_classifier_free_guidance: - negative_prompt_embeds = batch.negative_prompt_embeds[0] - negative_prompt_attention_mask = (batch.negative_attention_mask[0] - if batch.negative_attention_mask else None) - # Concatenate for batched processing - prompt_embeds_combined = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - if prompt_attention_mask is not None: - prompt_attention_mask_combined = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], - dim=0) - else: - prompt_attention_mask_combined = None - else: - prompt_embeds_combined = prompt_embeds - prompt_attention_mask_combined = prompt_attention_mask - - # Denoising loop - num_inference_steps = len(timesteps) - with tqdm(total=num_inference_steps, desc="LongCat Denoising") as progress_bar: - for i, t in enumerate(timesteps): - # Expand latents for CFG - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - latent_model_input = latent_model_input.to(target_dtype) - - # Expand timestep to match batch size - timestep = t.expand(latent_model_input.shape[0]).to(target_dtype) - - # Run transformer with context - batch.is_cfg_negative = False - with set_forward_context( - current_timestep=i, - attn_metadata=None, - forward_batch=batch, - ), torch.autocast(device_type='cuda', dtype=target_dtype, enabled=autocast_enabled): - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds_combined, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask_combined, - ) - - # Apply CFG with optimized scale - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) - - B = noise_pred_cond.shape[0] - positive = noise_pred_cond.reshape(B, -1) - negative = noise_pred_uncond.reshape(B, -1) - - # Calculate optimized scale (CFG-zero) - st_star = self.optimized_scale(positive, negative) - - # Reshape for broadcasting - st_star = st_star.view(B, 1, 1, 1, 1) - - # Apply optimized CFG formula - noise_pred = (noise_pred_uncond * st_star + guidance_scale * - (noise_pred_cond - noise_pred_uncond * st_star)) - - # CRITICAL: Negate noise prediction for flow matching scheduler - noise_pred = -noise_pred - - # Compute previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - progress_bar.update() - - # Update batch with denoised latents - batch.latents = latents - return batch diff --git a/fastvideo/pipelines/stages/longcat_i2v_denoising.py b/fastvideo/pipelines/stages/longcat_i2v_denoising.py deleted file mode 100644 index 27292be18..000000000 --- a/fastvideo/pipelines/stages/longcat_i2v_denoising.py +++ /dev/null @@ -1,151 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat I2V Denoising Stage with conditioning support. - -This stage implements Tier 3 I2V denoising: -1. Per-frame timestep masking (timestep[:, :num_cond_latents] = 0) -2. Passes num_cond_latents to transformer (for RoPE skipping) -3. Selective denoising (only updates non-conditioned frames) -4. CFG-zero optimized guidance -""" - -import torch -from tqdm import tqdm - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger -from fastvideo.models.loader.component_loader import TransformerLoader -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.longcat_denoising import LongCatDenoisingStage - -logger = init_logger(__name__) - - -class LongCatI2VDenoisingStage(LongCatDenoisingStage): - """ - LongCat denoising with I2V conditioning support. - - Key modifications from base LongCat denoising: - 1. Sets timestep=0 for conditioning frames - 2. Passes num_cond_latents to transformer - 3. Only applies scheduler step to non-conditioned frames - """ - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Run denoising loop with I2V conditioning.""" - - # Load transformer if needed - if not fastvideo_args.model_loaded["transformer"]: - loader = TransformerLoader() - self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) - fastvideo_args.model_loaded["transformer"] = True - - # Setup - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - - latents = batch.latents - timesteps = batch.timesteps - prompt_embeds = batch.prompt_embeds[0] - prompt_attention_mask = (batch.prompt_attention_mask[0] if batch.prompt_attention_mask else None) - guidance_scale = batch.guidance_scale - do_classifier_free_guidance = batch.do_classifier_free_guidance - - # Get num_cond_latents from batch - num_cond_latents = getattr(batch, 'num_cond_latents', 0) - - if num_cond_latents > 0: - logger.info("I2V Denoising: num_cond_latents=%s, latent_shape=%s", num_cond_latents, latents.shape) - - # Prepare negative prompts for CFG - if do_classifier_free_guidance: - negative_prompt_embeds = batch.negative_prompt_embeds[0] - negative_prompt_attention_mask = (batch.negative_attention_mask[0] - if batch.negative_attention_mask else None) - - prompt_embeds_combined = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - if prompt_attention_mask is not None: - prompt_attention_mask_combined = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], - dim=0) - else: - prompt_attention_mask_combined = None - else: - prompt_embeds_combined = prompt_embeds - prompt_attention_mask_combined = prompt_attention_mask - - # Denoising loop - num_inference_steps = len(timesteps) - - with tqdm(total=num_inference_steps, desc="I2V Denoising") as progress_bar: - for i, t in enumerate(timesteps): - - # 1. Expand latents for CFG - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - latent_model_input = latent_model_input.to(target_dtype) - - # 2. Expand timestep to match batch size - timestep = t.expand(latent_model_input.shape[0]).to(target_dtype) - - # 3. CRITICAL: Expand timestep to temporal dimension - # and set conditioning frames to timestep=0 - timestep = timestep.unsqueeze(-1).repeat(1, latent_model_input.shape[2]) - - # Mark conditioning frames as clean (timestep=0) - if num_cond_latents > 0: - timestep[:, :num_cond_latents] = 0 - - # 4. Run transformer with num_cond_latents - batch.is_cfg_negative = False - with set_forward_context( - current_timestep=i, - attn_metadata=None, - forward_batch=batch, - ), torch.autocast(device_type='cuda', dtype=target_dtype, enabled=autocast_enabled): - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds_combined, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask_combined, - num_cond_latents=num_cond_latents, - ) - - # 5. Apply CFG with optimized scale - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) - - B = noise_pred_cond.shape[0] - positive = noise_pred_cond.reshape(B, -1) - negative = noise_pred_uncond.reshape(B, -1) - - # CFG-zero optimized scale - st_star = self.optimized_scale(positive, negative) - st_star = st_star.view(B, 1, 1, 1, 1) - - noise_pred = (noise_pred_uncond * st_star + guidance_scale * - (noise_pred_cond - noise_pred_uncond * st_star)) - - # 6. CRITICAL: Negate for flow matching scheduler - noise_pred = -noise_pred - - # 7. CRITICAL: Only update non-conditioned frames - # The conditioning frames stay FIXED throughout denoising - if num_cond_latents > 0: - latents[:, :, num_cond_latents:] = self.scheduler.step(noise_pred[:, :, num_cond_latents:], - t, - latents[:, :, num_cond_latents:], - return_dict=False)[0] - else: - # No conditioning, update all frames - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - progress_bar.update() - - # Update batch with denoised latents - batch.latents = latents - return batch diff --git a/fastvideo/pipelines/stages/longcat_i2v_latent_preparation.py b/fastvideo/pipelines/stages/longcat_i2v_latent_preparation.py deleted file mode 100644 index a49f1cb93..000000000 --- a/fastvideo/pipelines/stages/longcat_i2v_latent_preparation.py +++ /dev/null @@ -1,95 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat I2V Latent Preparation Stage. - -This stage prepares latents with image conditioning for the first frame. -""" - -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage - -logger = init_logger(__name__) - - -class LongCatI2VLatentPreparationStage(LatentPreparationStage): - """ - Prepare latents with image conditioning for first frame. - - This stage: - 1. Generates random noise for all frames - 2. Replaces first latent frame with encoded image latent - 3. Marks conditioning information in batch - """ - - # Uses parent __init__ - no need for additional constructor - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Prepare latents with I2V conditioning.""" - - # IMPORTANT: Skip if latents already prepared (e.g., by refinement init stage) - # The refine_init stage encodes stage1 video and mixes with noise - don't overwrite! - if batch.latents is not None: - logger.info( - "I2V Latent Prep: Skipping - latents already prepared " - "(shape=%s), likely from refinement stage", batch.latents.shape) - return batch - - # 1. Calculate dimensions - num_frames = batch.num_frames - height = batch.height - width = batch.width - - # Get VAE compression factors - # IMPORTANT: Use VAE's temporal compression (4), NOT transformer's patch_size[0] (1) - vae_temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal - vae_spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial - - num_latent_frames = (num_frames - 1) // vae_temporal_scale + 1 - latent_height = height // vae_spatial_scale - latent_width = width // vae_spatial_scale - - num_channels = self.transformer.config.in_channels - - logger.info( - "I2V Latent Prep: num_frames=%s, num_latent_frames=%s " - "(vae_temporal_scale=%s), latent_shape=(%s, %s)", num_frames, num_latent_frames, vae_temporal_scale, - latent_height, latent_width) - - # 2. Generate random noise for all frames - # batch_size might not be set, default to 1 - batch_size = batch.batch_size if batch.batch_size is not None else 1 - shape = (batch_size, num_channels, num_latent_frames, latent_height, latent_width) - - # Handle generator - may be a list for batch handling - generator = batch.generator - if isinstance(generator, list): - generator = generator[0] if generator else None - - # torch.randn requires specific argument order: size, generator, dtype - latents = torch.randn(*shape, generator=generator).to(get_local_torch_device(), dtype=torch.float32) - - # 3. Replace first frame with conditioned image latent - if batch.image_latent is not None: - num_cond_latents = batch.num_cond_latents - latents[:, :, :num_cond_latents] = batch.image_latent[:, :, :num_cond_latents] - - logger.info("I2V: Replaced first %s latent frame(s) with image conditioning", num_cond_latents) - else: - logger.warning("No image_latent found in batch, proceeding without conditioning") - - # 4. Store in batch - batch.latents = latents - - # Required by base class output validator - batch.raw_latent_shape = (num_latent_frames, latent_height, latent_width) - - return batch diff --git a/fastvideo/pipelines/stages/longcat_image_vae_encoding.py b/fastvideo/pipelines/stages/longcat_image_vae_encoding.py deleted file mode 100644 index 4642f6fea..000000000 --- a/fastvideo/pipelines/stages/longcat_image_vae_encoding.py +++ /dev/null @@ -1,149 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat Image VAE Encoding Stage for I2V generation. - -This stage handles encoding a single input image to latent space with -LongCat-specific normalization for I2V conditioning. -""" - -import PIL -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.models.vision_utils import (normalize, numpy_to_pt, pil_to_numpy, resize) -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.utils import PRECISION_TO_TYPE - -logger = init_logger(__name__) - - -class LongCatImageVAEEncodingStage(PipelineStage): - """ - Encode input image to latent space for I2V conditioning. - - This stage: - 1. Preprocesses image to match target dimensions - 2. Encodes via VAE to latent space - 3. Applies LongCat-specific normalization - 4. Stores latent and calculates num_cond_latents - """ - - def __init__(self, vae): - super().__init__() - self.vae = vae - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Encode image to latent for I2V conditioning.""" - - # Skip image encoding for refinement tasks - we're refining an existing video - if getattr(batch, 'stage1_video', None) is not None or getattr(batch, 'refine_from', None) is not None: - logger.info("Skipping image encoding - refinement mode (using stage1_video)") - return batch - - # 1. Get image from batch - image = batch.pil_image # PIL.Image - if image is None: - raise ValueError("pil_image must be provided for I2V") - - if not isinstance(image, PIL.Image.Image): - raise TypeError(f"pil_image must be PIL.Image, got {type(image)}") - - # 2. Get target dimensions - height = batch.height - width = batch.width - - if height is None or width is None: - raise ValueError("height and width must be set for I2V") - - # 3. Preprocess image - image = resize(image, height, width, resize_mode="default") - image = pil_to_numpy(image) - image = numpy_to_pt(image) - image = normalize(image) # to [-1, 1] - - # 4. Add temporal dimension - # After numpy_to_pt: [1, C, H, W] (batch already added by pil_to_numpy) - # Add T dimension: [1, C, H, W] -> [1, C, 1, H, W] = [B, C, T, H, W] - image = image.unsqueeze(2) - image = image.to(get_local_torch_device(), dtype=torch.float32) - - # 5. Encode via VAE - self.vae = self.vae.to(get_local_torch_device()) - - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - - with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): - if fastvideo_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - - if not vae_autocast_enabled: - image = image.to(vae_dtype) - - with torch.no_grad(): - encoder_output = self.vae.encode(image) - latent = self.retrieve_latents(encoder_output, batch.generator) - - # 6. Apply LongCat-specific normalization - # Formula: (latents - mean) / std - latent = self.normalize_latents(latent) - - # 7. Calculate num_cond_latents - # Formula: 1 + (num_cond_frames - 1) // vae_temporal_scale - # For single image (num_cond_frames=1): always 1 latent frame - num_cond_frames = 1 # Single image - vae_temporal_scale = self.vae.config.scale_factor_temporal - batch.num_cond_latents = 1 + (num_cond_frames - 1) // vae_temporal_scale - - # 8. Store in batch - batch.image_latent = latent - batch.num_cond_frames = 1 - - logger.info("I2V: Encoded image to latent shape %s, num_cond_latents=%s", latent.shape, batch.num_cond_latents) - - # Offload VAE if needed - if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") - - return batch - - def retrieve_latents(self, encoder_output: object, generator: torch.Generator | None) -> torch.Tensor: - """Sample from VAE posterior.""" - # WAN VAE returns an object with .sample() method - if hasattr(encoder_output, 'sample'): - return encoder_output.sample(generator) - elif hasattr(encoder_output, 'latent_dist'): - return encoder_output.latent_dist.sample(generator) - elif hasattr(encoder_output, 'latents'): - return encoder_output.latents - else: - raise AttributeError("Could not access latents from encoder output") - - def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: - """ - Apply LongCat-specific latent normalization. - - Formula: (latents - mean) / std - - This matches the original LongCat implementation and is DIFFERENT - from standard VAE scaling (which uses scaling_factor). - """ - if not hasattr(self.vae.config, 'latents_mean') or not hasattr(self.vae.config, 'latents_std'): - raise ValueError("VAE config must have 'latents_mean' and 'latents_std' " - "for LongCat normalization") - - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, - 1).to(latents.device, latents.dtype) - - latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, - 1).to(latents.device, latents.dtype) - - return (latents - latents_mean) / latents_std diff --git a/fastvideo/pipelines/stages/longcat_kv_cache_init.py b/fastvideo/pipelines/stages/longcat_kv_cache_init.py deleted file mode 100644 index c85d0ce2b..000000000 --- a/fastvideo/pipelines/stages/longcat_kv_cache_init.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat KV Cache Initialization Stage for Video Continuation (VC). - -This stage pre-computes K/V cache for conditioning frames. - -""" - -import torch - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage - -logger = init_logger(__name__) - - -class LongCatKVCacheInitStage(PipelineStage): - """ - Pre-compute KV cache for conditioning frames. - - After this stage: - - batch.kv_cache_dict contains {block_idx: (k, v)} - - batch.cond_latents contains the conditioning latents - - batch.latents contains ONLY noise latents - """ - - def __init__(self, transformer): - super().__init__() - self.transformer = transformer - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Initialize KV cache from conditioning latents.""" - - # Check if KV cache is enabled - use_kv_cache = getattr(fastvideo_args.pipeline_config, 'use_kv_cache', True) - if not use_kv_cache: - batch.kv_cache_dict = {} - batch.use_kv_cache = False - logger.info("KV cache disabled, skipping initialization") - return batch - - batch.use_kv_cache = True - offload_kv_cache = getattr(fastvideo_args.pipeline_config, 'offload_kv_cache', False) - - # Get conditioning latents - num_cond_latents = batch.num_cond_latents - if num_cond_latents <= 0: - batch.kv_cache_dict = {} - logger.warning("num_cond_latents <= 0, skipping KV cache init") - return batch - - # Extract conditioning latents - cond_latents = batch.latents[:, :, :num_cond_latents].clone() - - logger.info("Initializing KV cache for %d conditioning latents, shape: %s", num_cond_latents, - cond_latents.shape) - - # Timestep = 0 for conditioning (they are "clean") - B = cond_latents.shape[0] - T_cond = cond_latents.shape[2] - timestep = torch.zeros(B, T_cond, device=cond_latents.device, dtype=cond_latents.dtype) - - # Empty prompt embeddings (cross-attn will be skipped) - max_seq_len = 512 - # Get caption dimension from transformer config - caption_dim = self.transformer.config.caption_channels - empty_embeds = torch.zeros(B, max_seq_len, caption_dim, device=cond_latents.device, dtype=cond_latents.dtype) - - # Get transformer dtype - if hasattr(self.transformer, 'module'): - transformer_dtype = next(self.transformer.module.parameters()).dtype - else: - transformer_dtype = next(self.transformer.parameters()).dtype - - # Run transformer with return_kv=True, skip_crs_attn=True - with ( - torch.no_grad(), - set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=batch, - ), - torch.autocast(device_type='cuda', dtype=transformer_dtype), - ): - _, kv_cache_dict = self.transformer( - hidden_states=cond_latents.to(transformer_dtype), - encoder_hidden_states=empty_embeds.to(transformer_dtype), - timestep=timestep.to(transformer_dtype), - return_kv=True, - skip_crs_attn=True, - offload_kv_cache=offload_kv_cache, - ) - - # Store cache and save cond_latents for later concatenation - batch.kv_cache_dict = kv_cache_dict - batch.cond_latents = cond_latents - - # Remove conditioning latents from main latents - # After this, batch.latents contains ONLY noise frames - batch.latents = batch.latents[:, :, num_cond_latents:] - - logger.info("KV cache initialized: %d blocks, offload=%s, remaining latents shape: %s", len(kv_cache_dict), - offload_kv_cache, batch.latents.shape) - - return batch diff --git a/fastvideo/pipelines/stages/longcat_refine_init.py b/fastvideo/pipelines/stages/longcat_refine_init.py deleted file mode 100644 index 25cac2b05..000000000 --- a/fastvideo/pipelines/stages/longcat_refine_init.py +++ /dev/null @@ -1,276 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat refinement initialization stage. - -This stage prepares the latent variables for LongCat's 480p->720p refinement by: -1. Loading the stage1 (480p) video -2. Upsampling it to 720p resolution -3. Encoding it with VAE -4. Mixing with noise according to t_thresh -""" - -import math -from pathlib import Path - -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.models.vision_utils import load_video -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.configs.pipelines.longcat import get_bucket_config - -logger = init_logger(__name__) - - -class LongCatRefineInitStage(PipelineStage): - """ - Stage for initializing LongCat refinement from a stage1 (480p) video. - - This replicates the logic from LongCatVideoPipeline.generate_refine(): - - Load stage1_video frames - - Upsample spatially and temporally - - VAE encode and normalize - - Mix with noise according to t_thresh - """ - - def __init__(self, vae) -> None: - super().__init__() - self.vae = vae - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """ - Initialize latents for refinement. - - Args: - batch: The current batch information. - fastvideo_args: The inference arguments. - - Returns: - The batch with initialized latents for refinement. - """ - refine_from = batch.refine_from - in_memory_stage1 = batch.stage1_video - - # Only run for refinement tasks: either a path (refine_from) or in-memory video is provided - if refine_from is None and in_memory_stage1 is None: - # Not a refinement task, skip - return batch - - # ------------------------------------------------------------------ - # 1. Obtain stage1 frames (either from disk or from in-memory input) - # ------------------------------------------------------------------ - if in_memory_stage1 is not None: - # User provided stage1 frames directly (e.g., from distilled stage output) - if len(in_memory_stage1) == 0: - raise ValueError("stage1_video is empty; expected a non-empty list of frames") - - if isinstance(in_memory_stage1[0], Image.Image): - pil_images = in_memory_stage1 - else: - # Assume numpy arrays or torch tensors with shape [H, W, C] - pil_images = [Image.fromarray(np.array(frame)) for frame in in_memory_stage1] - - logger.info("Initializing LongCat refinement from in-memory stage1_video (%s frames)", len(pil_images)) - else: - # Path-based refine: load video from disk (original design) - logger.info("Initializing LongCat refinement from file: %s", refine_from) - stage1_video_path = Path(refine_from) - if not stage1_video_path.exists(): - raise FileNotFoundError(f"Stage1 video not found: {refine_from}") - - # Load video frames as PIL Images - pil_images, original_fps = load_video(str(stage1_video_path), return_fps=True) - logger.info("Loaded stage1 video: %s frames @ %s fps", len(pil_images), original_fps) - - # Store in batch for reference (use PIL images, same as official demo) - batch.stage1_video = pil_images - - # Get parameters from batch - num_frames = len(pil_images) - spatial_refine_only = batch.spatial_refine_only - t_thresh = batch.t_thresh - num_cond_frames = batch.num_cond_frames if hasattr(batch, 'num_cond_frames') else 0 - - # Calculate new frame count (temporal upsampling if not spatial_refine_only) - new_num_frames = num_frames if spatial_refine_only else 2 * num_frames - logger.info("Refine mode: %s", 'spatial only' if spatial_refine_only else 'spatial + temporal') - - # Update batch.num_frames to reflect the upsampled count - batch.num_frames = new_num_frames - - # Use bucket system to select resolution (exactly like LongCat) - # Calculate scale_factor_spatial considering SP split - sp_size = fastvideo_args.sp_size if fastvideo_args.sp_size > 0 else 1 - vae_scale_factor_spatial = 8 # VAE spatial downsampling - patch_size_spatial = 2 # LongCat patch size - bsa_latent_granularity = 4 - scale_factor_spatial = vae_scale_factor_spatial * patch_size_spatial * bsa_latent_granularity # 64 - - # Calculate optimal split like LongCat (cp_split_hw logic) - # For sp_size=1: [1,1], max=1 - # For sp_size=2: [1,2], max=2 - # For sp_size=4: [2,2], max=2 - # For sp_size=8: [2,4], max=4 - if sp_size > 1: - # Get optimal 2D split factors (mimic context_parallel_util.get_optimal_split) - factors = [] - for i in range(1, int(sp_size**0.5) + 1): - if sp_size % i == 0: - factors.append([i, sp_size // i]) - cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1])) - scale_factor_spatial *= max(cp_split_hw) - logger.info("SP split: sp_size=%s, cp_split_hw=%s, max_split=%s", sp_size, cp_split_hw, max(cp_split_hw)) - else: - cp_split_hw = [1, 1] - - # Get bucket config and find closest bucket for the input aspect ratio - bucket_config = get_bucket_config('720p', scale_factor_spatial) - - # Get input aspect ratio from stage1 video - input_height, input_width = pil_images[0].height, pil_images[0].width - input_ratio = input_height / input_width - - # Find closest bucket - closest_ratio = min(bucket_config.keys(), key=lambda x: abs(float(x) - input_ratio)) - height, width = bucket_config[closest_ratio][0] - - logger.info("Input aspect ratio: %.2f (%sx%s)", input_ratio, input_width, input_height) - logger.info("Matched bucket ratio: %s -> resolution: %sx%s", closest_ratio, width, height) - logger.info("Target: %sx%s @ %s frames (sp_size=%s, scale_factor=%s)", width, height, new_num_frames, sp_size, - scale_factor_spatial) - - # Override batch height/width with bucket-selected resolution - batch.height = height - batch.width = width - - # Convert PIL images to tensor [T, C, H, W] - stage1_video_tensor = torch.stack([ - torch.from_numpy(np.array(img)).permute(2, 0, 1) # HWC -> CHW - for img in pil_images - ]).float() # [T, C, H, W] - - device = batch.prompt_embeds[0].device - dtype = batch.prompt_embeds[0].dtype - stage1_video_tensor = stage1_video_tensor.to(device=device, dtype=dtype) - - # Replicate LongCat's exact preprocessing (lines 1227-1235 in pipeline_longcat_video.py) - # First: spatial interpolation to target (height, width) on [T, C, H, W] - video_down = F.interpolate(stage1_video_tensor, size=(height, width), mode='bilinear', align_corners=True) - - # Rearrange to [C, T, H, W] and add batch dimension -> [1, C, T, H, W] - video_down = video_down.permute(1, 0, 2, 3).unsqueeze(0) # [1, C, T, H, W] - video_down = video_down / 255.0 # Normalize to [0, 1] - - # Then: temporal+spatial interpolation to (new_num_frames, height, width) - video_up = F.interpolate(video_down, size=(new_num_frames, height, width), mode='trilinear', align_corners=True) - - # Rescale to [-1, 1] for VAE - video_up = video_up * 2.0 - 1.0 - - logger.info("Upsampled video shape: %s", video_up.shape) - - # Padding logic (exactly like LongCat lines 1237-1255) - # Only pad temporal dimension to ensure BSA compatibility - vae_scale_factor_temporal = 4 - num_noise_frames = video_up.shape[2] - num_cond_frames - - num_cond_latents = 0 - num_cond_frames_added = 0 - if num_cond_frames > 0: - num_cond_latents = 1 + math.ceil((num_cond_frames - 1) / vae_scale_factor_temporal) - num_cond_latents = math.ceil(num_cond_latents / bsa_latent_granularity) * bsa_latent_granularity - num_cond_frames_added = 1 + (num_cond_latents - 1) * vae_scale_factor_temporal - num_cond_frames - num_cond_frames = num_cond_frames + num_cond_frames_added - - num_noise_latents = math.ceil(num_noise_frames / vae_scale_factor_temporal) - num_noise_latents = math.ceil(num_noise_latents / bsa_latent_granularity) * bsa_latent_granularity - num_noise_frames_added = num_noise_latents * vae_scale_factor_temporal - num_noise_frames - - if num_cond_frames_added > 0 or num_noise_frames_added > 0: - logger.info("Padding temporal dimension for BSA: cond_frames+=%s, noise_frames+=%s", num_cond_frames_added, - num_noise_frames_added) - pad_front = video_up[:, :, 0:1].repeat(1, 1, num_cond_frames_added, 1, 1) - pad_back = video_up[:, :, -1:].repeat(1, 1, num_noise_frames_added, 1, 1) - video_up = torch.cat([pad_front, video_up, pad_back], dim=2) - logger.info("Padded video shape: %s", video_up.shape) - - # Update batch with actual frame count after padding - batch.num_frames = video_up.shape[2] - - # Store padding info for later cropping (CRITICAL for correct output!) - batch.num_cond_frames_added = num_cond_frames_added - batch.num_noise_frames_added = num_noise_frames_added - batch.new_frame_size_before_padding = new_num_frames - - # Store num_cond_latents for denoising stage - if num_cond_latents > 0: - batch.num_cond_latents = num_cond_latents - logger.info("Will use num_cond_latents=%s during denoising", num_cond_latents) - - logger.info("Padding info: cond+=%s, noise+=%s, original=%s", num_cond_frames_added, num_noise_frames_added, - new_num_frames) - - # VAE encode with tiling for memory efficiency - logger.info("Encoding stage1 video with VAE (tiling enabled)...") - vae_dtype = next(self.vae.parameters()).dtype - vae_device = next(self.vae.parameters()).device - video_up = video_up.to(dtype=vae_dtype, device=vae_device) - - # Enable tiling for large video encoding - if hasattr(self.vae, 'enable_tiling'): - self.vae.enable_tiling() - - with torch.no_grad(): - latent_dist = self.vae.encode(video_up) - # Extract tensor from latent distribution - if hasattr(latent_dist, 'latent_dist'): - # Nested distribution wrapper - latent_up = latent_dist.latent_dist.sample() - elif hasattr(latent_dist, 'sample'): - # DiagonalGaussianDistribution or similar - latent_up = latent_dist.sample() - elif hasattr(latent_dist, 'latents'): - # Direct latents tensor - latent_up = latent_dist.latents - else: - # Assume it's already a tensor - latent_up = latent_dist - - # Normalize latents using VAE config (exactly like LongCat) - if hasattr(self.vae.config, 'latents_mean') and hasattr(self.vae.config, 'latents_std'): - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, - 1).to(latent_up.device, latent_up.dtype) - # LongCat uses: 1.0 / latents_std (equivalent to dividing by latents_std) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - latent_up.device, latent_up.dtype) - # LongCat: (latents - mean) * (1/std) - latent_up = (latent_up - latents_mean) * latents_std - - logger.info("Encoded latent shape: %s", latent_up.shape) - - # Mix with noise according to t_thresh - # latent_up = (1 - t_thresh) * latent_up + t_thresh * noise - noise = torch.randn_like(latent_up).contiguous() - latent_up = (1 - t_thresh) * latent_up + t_thresh * noise - - logger.info("Applied t_thresh=%s noise mixing", t_thresh) - - # Store in batch - ensure correct dtype and device - # The latents need to be on the same device as the transformer (CUDA) - target_device = batch.prompt_embeds[0].device - batch.latents = latent_up.to(device=target_device, dtype=dtype) - batch.raw_latent_shape = latent_up.shape - - logger.info("Latents device: %s, dtype: %s", batch.latents.device, batch.latents.dtype) - logger.info("LongCat refinement initialization complete") - - return batch diff --git a/fastvideo/pipelines/stages/longcat_refine_timestep.py b/fastvideo/pipelines/stages/longcat_refine_timestep.py deleted file mode 100644 index 0b292fd86..000000000 --- a/fastvideo/pipelines/stages/longcat_refine_timestep.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat refinement timestep preparation stage. - -This stage prepares special timesteps for LongCat refinement that start from t_thresh. -""" - -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage - -logger = init_logger(__name__) - - -class LongCatRefineTimestepStage(PipelineStage): - """ - Stage for preparing timesteps specific to LongCat refinement. - - For refinement, we need to start from t_thresh instead of t=1.0, so we: - 1. Generate normal timesteps for num_inference_steps - 2. Filter to only keep timesteps < t_thresh * 1000 - 3. Prepend t_thresh * 1000 as the first timestep - """ - - def __init__(self, scheduler) -> None: - super().__init__() - self.scheduler = scheduler - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """ - Prepare refinement-specific timesteps. - - Args: - batch: The current batch information. - fastvideo_args: The inference arguments. - - Returns: - The batch with refinement timesteps. - """ - # Only apply if this is a refinement task - # Trigger when either a refine_from path or in-memory stage1_video is provided - if batch.refine_from is None and batch.stage1_video is None: - return batch - - device = get_local_torch_device() - num_inference_steps = batch.num_inference_steps - t_thresh = batch.t_thresh - - logger.info("Preparing LongCat refinement timesteps (t_thresh=%s)", t_thresh) - - # ------------------------------------------------------------------ - # 1) Match LongCatVideoPipeline.get_timesteps_sigmas (non-distill): - # sigmas = linspace(1, 0.001, num_inference_steps) on CPU - # ------------------------------------------------------------------ - base_sigmas = torch.linspace( - 1.0, - 0.001, - num_inference_steps, - dtype=torch.float32, - device="cpu", # scheduler.set_timesteps expects CPU-convertible sigmas - ) - # Let the scheduler build its internal timestep schedule from sigmas - self.scheduler.set_timesteps(num_inference_steps, sigmas=base_sigmas, device=device) - base_timesteps = self.scheduler.timesteps - - # ------------------------------------------------------------------ - # 2) Apply t_thresh cropping exactly like generate_refine: - # timesteps = [t_thresh*1000] + [t for t in base_timesteps if t < t_thresh*1000] - # sigmas = timesteps / 1000 (with trailing zero) - # ------------------------------------------------------------------ - t_thresh_value = t_thresh * 1000.0 - t_thresh_tensor = torch.tensor(t_thresh_value, dtype=base_timesteps.dtype, device=device) - filtered_timesteps = base_timesteps[base_timesteps < t_thresh_tensor] - - timesteps = torch.cat([t_thresh_tensor.unsqueeze(0), filtered_timesteps]) - - # Update scheduler with these custom timesteps and corresponding sigmas - self.scheduler.timesteps = timesteps - sigmas = torch.cat([timesteps / 1000.0, torch.zeros(1, device=device)]) - self.scheduler.sigmas = sigmas - - logger.info("Refinement timesteps: %s steps starting from t=%s", len(timesteps), t_thresh) - logger.info("First few timesteps: %s", timesteps[:5].tolist()) - - # Store in batch so downstream stages (denoising) use the same schedule - batch.timesteps = timesteps - - return batch diff --git a/fastvideo/pipelines/stages/longcat_vc_denoising.py b/fastvideo/pipelines/stages/longcat_vc_denoising.py deleted file mode 100644 index ee6b4c4a2..000000000 --- a/fastvideo/pipelines/stages/longcat_vc_denoising.py +++ /dev/null @@ -1,191 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat VC Denoising Stage with KV cache support. - -This stage extends the I2V denoising stage to support: -1. KV cache for conditioning frames -2. Video continuation with multiple conditioning frames -""" - -import time -import torch -from tqdm import tqdm - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger -from fastvideo.models.loader.component_loader import TransformerLoader -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.longcat_denoising import LongCatDenoisingStage - -logger = init_logger(__name__) - - -class LongCatVCDenoisingStage(LongCatDenoisingStage): - """ - LongCat denoising with Video Continuation and KV cache support. - - Key differences from I2V denoising: - - Supports KV cache (reuses cached K/V from conditioning frames) - - Handles larger num_cond_latents - - Concatenates conditioning latents back after denoising - - When use_kv_cache=True: - - batch.latents contains ONLY noise frames (cond removed by KV cache init) - - batch.kv_cache_dict contains cached K/V - - batch.cond_latents contains conditioning latents for post-concat - - When use_kv_cache=False: - - batch.latents contains ALL frames (cond + noise) - - Timestep masking: timestep[:, :num_cond_latents] = 0 - - Selective denoising: only update noise frames - """ - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Run denoising loop with VC conditioning and optional KV cache.""" - - # Load transformer if needed - if not fastvideo_args.model_loaded["transformer"]: - loader = TransformerLoader() - self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) - fastvideo_args.model_loaded["transformer"] = True - - # Setup - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - - latents = batch.latents - timesteps = batch.timesteps - prompt_embeds = batch.prompt_embeds[0] - prompt_attention_mask = (batch.prompt_attention_mask[0] if batch.prompt_attention_mask else None) - guidance_scale = batch.guidance_scale - do_classifier_free_guidance = batch.do_classifier_free_guidance - - # Get VC-specific parameters - num_cond_latents = getattr(batch, 'num_cond_latents', 0) - use_kv_cache = getattr(batch, 'use_kv_cache', False) - kv_cache_dict = getattr(batch, 'kv_cache_dict', {}) - - logger.info("VC Denoising: num_cond_latents=%d, use_kv_cache=%s, latent_shape=%s", num_cond_latents, - use_kv_cache, latents.shape) - - # Prepare negative prompts for CFG - if do_classifier_free_guidance: - negative_prompt_embeds = batch.negative_prompt_embeds[0] - negative_prompt_attention_mask = (batch.negative_attention_mask[0] - if batch.negative_attention_mask else None) - - prompt_embeds_combined = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - if prompt_attention_mask is not None: - prompt_attention_mask_combined = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], - dim=0) - else: - prompt_attention_mask_combined = None - else: - prompt_embeds_combined = prompt_embeds - prompt_attention_mask_combined = prompt_attention_mask - - # Denoising loop - num_inference_steps = len(timesteps) - step_times = [] - - with tqdm(total=num_inference_steps, desc="VC Denoising") as progress_bar: - for i, t in enumerate(timesteps): - step_start = time.time() - - # 1. Expand latents for CFG - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - latent_model_input = latent_model_input.to(target_dtype) - - # 2. Expand timestep to match batch size - timestep = t.expand(latent_model_input.shape[0]).to(target_dtype) - - # 3. Expand timestep to temporal dimension - timestep = timestep.unsqueeze(-1).repeat(1, latent_model_input.shape[2]) - - # 4. Timestep masking (only when NOT using KV cache) - if not use_kv_cache and num_cond_latents > 0: - timestep[:, :num_cond_latents] = 0 - - # 5. Prepare transformer kwargs - # IMPORTANT: num_cond_latents is ALWAYS passed - needed for RoPE position offset - transformer_kwargs = { - 'num_cond_latents': num_cond_latents, - } - if use_kv_cache: - transformer_kwargs['kv_cache_dict'] = kv_cache_dict - - # 6. Run transformer - batch.is_cfg_negative = False - with set_forward_context( - current_timestep=i, - attn_metadata=None, - forward_batch=batch, - ), torch.autocast(device_type='cuda', dtype=target_dtype, enabled=autocast_enabled): - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds_combined, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask_combined, - **transformer_kwargs, - ) - - # 7. Apply CFG with optimized scale (CFG-zero) - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) - - B = noise_pred_cond.shape[0] - positive = noise_pred_cond.reshape(B, -1) - negative = noise_pred_uncond.reshape(B, -1) - - st_star = self.optimized_scale(positive, negative) - st_star = st_star.view(B, 1, 1, 1, 1) - - noise_pred = (noise_pred_uncond * st_star + guidance_scale * - (noise_pred_cond - noise_pred_uncond * st_star)) - - # 8. Negate for flow matching scheduler - noise_pred = -noise_pred - - # 9. Scheduler step - if use_kv_cache: - # All latents are noise frames (conditioning is in cache) - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - else: - # Only update noise frames (skip conditioning) - if num_cond_latents > 0: - latents[:, :, num_cond_latents:] = self.scheduler.step( - noise_pred[:, :, num_cond_latents:], - t, - latents[:, :, num_cond_latents:], - return_dict=False, - )[0] - else: - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - step_time = time.time() - step_start - step_times.append(step_time) - - # Log timing for first few steps - if i < 3: - logger.info("Step %d: %.2fs", i, step_time) - - progress_bar.update() - - # 10. If using KV cache, concatenate conditioning latents back - if use_kv_cache and hasattr(batch, 'cond_latents') and batch.cond_latents is not None: - latents = torch.cat([batch.cond_latents, latents], dim=2) - logger.info("Concatenated conditioning latents back, final shape: %s", latents.shape) - - # Log average timing - avg_time = sum(step_times) / len(step_times) - logger.info("Average step time: %.2fs (total: %.1fs)", avg_time, sum(step_times)) - - # Update batch with denoised latents - batch.latents = latents - return batch diff --git a/fastvideo/pipelines/stages/longcat_video_vae_encoding.py b/fastvideo/pipelines/stages/longcat_video_vae_encoding.py deleted file mode 100644 index 8905dc082..000000000 --- a/fastvideo/pipelines/stages/longcat_video_vae_encoding.py +++ /dev/null @@ -1,163 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -LongCat Video VAE Encoding Stage for Video Continuation (VC) generation. - -This stage handles encoding multiple video frames to latent space with -LongCat-specific normalization for VC conditioning. -""" - -from typing import Any - -import PIL.Image -import torch - -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.models.vision_utils import normalize, numpy_to_pt, pil_to_numpy, resize -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.utils import PRECISION_TO_TYPE - -logger = init_logger(__name__) - - -class LongCatVideoVAEEncodingStage(PipelineStage): - """ - Encode video frames to latent space for VC conditioning. - - This stage: - 1. Loads video frames from path or uses provided frames - 2. Takes the last num_cond_frames from the video - 3. Preprocesses and stacks frames - 4. Encodes via VAE to latent space - 5. Applies LongCat-specific normalization - 6. Calculates num_cond_latents - """ - - def __init__(self, vae): - super().__init__() - self.vae = vae - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """Encode video frames to latent for VC conditioning.""" - - # Get video from batch - can be path, list of PIL images, or already loaded - video = getattr(batch, 'video_frames', None) or getattr(batch, 'video_path', None) - num_cond_frames = getattr(batch, 'num_cond_frames', 13) # Default 13 for VC - - if video is None: - raise ValueError("video_frames or video_path must be provided for VC") - - # Load video if path - if isinstance(video, str): - from diffusers.utils import load_video - video = load_video(video) - logger.info("Loaded video from path: %d frames", len(video)) - - # Take last num_cond_frames - if len(video) > num_cond_frames: - video = video[-num_cond_frames:] - logger.info("Using last %d frames for conditioning", num_cond_frames) - elif len(video) < num_cond_frames: - logger.warning("Video has only %d frames, less than num_cond_frames=%d", len(video), num_cond_frames) - num_cond_frames = len(video) - - # Get target dimensions - height = batch.height - width = batch.width - - if height is None or width is None: - raise ValueError("height and width must be set for VC") - - # Preprocess and stack frames - processed_frames = [] - for frame in video: - if not isinstance(frame, PIL.Image.Image): - raise TypeError(f"Frame must be PIL.Image, got {type(frame)}") - - frame = resize(frame, height, width, resize_mode="default") - frame = pil_to_numpy(frame) # Returns [1, H, W, C] then converted - frame = numpy_to_pt(frame) # Returns [1, C, H, W] - frame = normalize(frame) # to [-1, 1] - processed_frames.append(frame) - - # Stack frames: [num_frames, C, H, W] -> [1, C, T, H, W] - video_tensor = torch.cat(processed_frames, dim=0) # [T, C, H, W] - video_tensor = video_tensor.permute(1, 0, 2, 3).unsqueeze(0) # [1, C, T, H, W] - video_tensor = video_tensor.to(get_local_torch_device(), dtype=torch.float32) - - logger.info("VC: Preprocessed video tensor shape: %s", video_tensor.shape) - - # Encode via VAE - self.vae = self.vae.to(get_local_torch_device()) - - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - - with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): - if fastvideo_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - - if not vae_autocast_enabled: - video_tensor = video_tensor.to(vae_dtype) - - with torch.no_grad(): - encoder_output = self.vae.encode(video_tensor) - latent = self.retrieve_latents(encoder_output, batch.generator) - - # Apply LongCat-specific normalization - latent = self.normalize_latents(latent) - - # Calculate num_cond_latents - # Formula: 1 + (num_cond_frames - 1) // vae_temporal_scale - vae_temporal_scale = self.vae.config.scale_factor_temporal - num_cond_latents = 1 + (num_cond_frames - 1) // vae_temporal_scale - - # Store in batch - batch.video_latent = latent - batch.num_cond_frames = num_cond_frames - batch.num_cond_latents = num_cond_latents - - logger.info("VC: Encoded %d frames to latent shape %s, num_cond_latents=%d", num_cond_frames, latent.shape, - num_cond_latents) - - # Offload VAE if needed - if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") - - return batch - - def retrieve_latents(self, encoder_output: Any, generator: torch.Generator | None) -> torch.Tensor: - """Sample from VAE posterior.""" - if hasattr(encoder_output, 'sample'): - return encoder_output.sample(generator) - elif hasattr(encoder_output, 'latent_dist'): - return encoder_output.latent_dist.sample(generator) - elif hasattr(encoder_output, 'latents'): - return encoder_output.latents - else: - raise AttributeError("Could not access latents from encoder output") - - def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: - """ - Apply LongCat-specific latent normalization. - - Formula: (latents - mean) / std - """ - if not hasattr(self.vae.config, 'latents_mean') or not hasattr(self.vae.config, 'latents_std'): - raise ValueError("VAE config must have 'latents_mean' and 'latents_std' " - "for LongCat normalization") - - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, - 1).to(latents.device, latents.dtype) - - latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, - 1).to(latents.device, latents.dtype) - - return (latents - latents_mean) / latents_std diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 8a8e2f068..b01150093 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -20,9 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import ( Cosmos25Config, Cosmos25_14BConfig, ) -from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig -from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.gen3c import Gen3CConfig from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyuan15I2V480PStepDistilledConfig, Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig, @@ -34,7 +32,6 @@ from fastvideo.configs.pipelines.kandinsky6_sr import Kandinsky6SRPipelineConfig from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig -from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.configs.pipelines.flux_2 import ( Flux2KleinPipelineConfig, @@ -44,18 +41,11 @@ from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig -from fastvideo.configs.pipelines.turbodiffusion import ( - TurboDiffusionI2V_A14B_Config, - TurboDiffusionT2V_14B_Config, - TurboDiffusionT2V_1_3B_Config, -) from fastvideo.models.wan import pipeline_config as wan_pipeline_config -from fastvideo.models.wan.definition import WAN_MODEL_DEFINITION_GROUPS, WanModelDefinition -from fastvideo.configs.pipelines.glm_image import GlmImageConfig +from fastvideo.models.wan.definition import WAN_MODEL_DEFINITIONS, WanModelDefinition from fastvideo.configs.pipelines.flux import FluxPipelineConfig from fastvideo.configs.pipelines.sd35 import SD35Config from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig) -from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig from fastvideo.api.sampling_param import SamplingParam from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam @@ -474,7 +464,7 @@ def _register_configs() -> None: default_preset="hunyuan15_sr_1080p", ) - # Hunyuan (excludes gamecraft, hyworld, and versioned models) + # Hunyuan (excludes hyworld and versioned models) register_configs( sampling_param_cls=None, pipeline_config_cls=HunyuanConfig, @@ -483,8 +473,8 @@ def _register_configs() -> None: "hunyuanvideo-community/HunyuanVideo", ], model_detectors=[ - lambda path: "hunyuan" in path.lower() and "gamecraft" not in path.lower() and "hyworld" not in path.lower( - ) and "1.5" not in path.lower() and "1-5" not in path.lower() + lambda path: "hunyuan" in path.lower() and "hyworld" not in path.lower() and "1.5" not in path.lower() and + "1-5" not in path.lower() ], model_family="hunyuan", default_preset="hunyuan_t2v", @@ -513,18 +503,6 @@ def _register_configs() -> None: default_preset="hyworld_t2v", ) - # HunyuanGameCraft - register_configs( - sampling_param_cls=None, - pipeline_config_cls=HunyuanGameCraftPipelineConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "FastVideo/HunyuanGameCraft-Diffusers", - ], - model_detectors=[lambda path: "gamecraft" in path.lower()], - model_family="gamecraft", - default_preset="gamecraft_i2v", - ) # LingBotWorld2 causal-fast register_configs( sampling_param_cls=None, @@ -838,43 +816,6 @@ def _register_configs() -> None: pipeline_cls_name="Kandinsky6TI2VAPipeline", ) - # LongCat (T2V, I2V, VC use same config; workload varies by path) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=LongCatT2V480PConfig, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=["FastVideo/LongCat-Video-T2V-Diffusers"], - model_detectors=[ - lambda path: "longcat" in path.lower() and "i2v" not in path.lower() and "imagetovideo" not in path.lower() - and "vc" not in path.lower() and "videocontinuation" not in path.lower(), - ], - model_family="longcat", - default_preset="longcat_t2v", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=LongCatT2V480PConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=["FastVideo/LongCat-Video-I2V-Diffusers"], - model_detectors=[ - lambda path: "longcatimagetovideo" in path.lower() or ("longcat" in path.lower() and "i2v" in path.lower()), - ], - model_family="longcat", - default_preset="longcat_i2v", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=LongCatT2V480PConfig, - workload_types=(), - hf_model_paths=["FastVideo/LongCat-Video-VC-Diffusers"], - model_detectors=[ - lambda path: "longcatvideocontinuation" in path.lower() or - ("longcat" in path.lower() and "vc" in path.lower()), - ], - model_family="longcat", - default_preset="longcat_vc", - ) - # MatrixGame 2.0 (I2V) register_configs( sampling_param_cls=MatrixGame2SamplingParam, @@ -994,76 +935,7 @@ def _register_configs() -> None: default_preset="cosmos_predict2_2b", ) - # TurboDiffusion - register_configs( - sampling_param_cls=None, - pipeline_config_cls=TurboDiffusionT2V_1_3B_Config, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers", - ], - model_detectors=[lambda path: "turbodiffusion" in path.lower() or "turbowan" in path.lower()], - model_family="turbodiffusion", - default_preset="turbo_t2v_1_3b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=TurboDiffusionT2V_14B_Config, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "loayrashid/TurboWan2.1-T2V-14B-Diffusers", - ], - model_family="turbodiffusion", - default_preset="turbo_t2v_14b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=TurboDiffusionI2V_A14B_Config, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "loayrashid/TurboWan2.2-I2V-A14B-Diffusers", - ], - model_family="turbodiffusion", - default_preset="turbo_i2v_a14b", - ) - - # Preserve first-match ordering around the DreamX registrations below. - _register_wan_configs(WAN_MODEL_DEFINITION_GROUPS[0]) - - register_configs( - sampling_param_cls=None, - pipeline_config_cls=DreamXWorld5BCamPipelineConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "FastVideo/DreamX-World-5B-Cam-Diffusers", - ], - model_detectors=[ - # Mutually exclusive with the AR detector below: Cam requires an - # explicit "cam" marker so hyphenated AR local paths (e.g. - # /ckpts/dreamx-world-5b-converted) don't first-match here — - # detector resolution is first-match in registration order. - lambda path: - ("dreamx-world" in path.lower() and "cam" in path.lower()) or "dreamxworldpipeline" in path.lower() - ], - model_family="dreamx_world", - default_preset="dreamx_world_5b_cam", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=DreamXWorld5BARPipelineConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "FastVideo/DreamX-World-5B-Diffusers", - ], - model_detectors=[ - lambda path: - ("dreamx-world-5b" in path.lower() and "cam" not in path.lower()) or "dreamxworldarpipeline" in path.lower( - ) - ], - model_family="dreamx_world", - default_preset="dreamx_world_5b_ar", - ) - _register_wan_configs(WAN_MODEL_DEFINITION_GROUPS[1]) + _register_wan_configs(WAN_MODEL_DEFINITIONS) # MiniMax H3 register_configs( @@ -1110,18 +982,6 @@ def _register_configs() -> None: default_preset="sd35_medium", ) - # GLM-Image - register_configs( - sampling_param_cls=None, - pipeline_config_cls=GlmImageConfig, - hf_model_paths=[ - "zai-org/GLM-Image", - ], - model_detectors=[lambda path: "glmimage" in path.lower() or "glm-image" in path.lower()], - workload_types=(WorkloadType.T2I, ), - model_family="glm_image", - ) - # FLUX.1-dev (Diffusers) register_configs( sampling_param_cls=FluxSamplingParam, @@ -1137,19 +997,6 @@ def _register_configs() -> None: ], ) - # Z-Image-Turbo - register_configs( - sampling_param_cls=None, - pipeline_config_cls=ZImagePipelineConfig, - workload_types=(WorkloadType.T2I, ), - hf_model_paths=["Tongyi-MAI/Z-Image-Turbo"], - model_detectors=[ - lambda path: "zimagepipeline" in path or "z-image" in path or "z_image" in path, - ], - model_family="zimage", - default_preset="zimage_turbo", - ) - # --- Part 3: Main Resolver --- @@ -1241,10 +1088,6 @@ def _register_presets() -> None: from fastvideo.api.presets import register_preset from fastvideo.pipelines.basic.cosmos.presets import ( ALL_PRESETS as COSMOS_PRESETS, ) - from fastvideo.pipelines.basic.dreamx_world.presets import ( - ALL_PRESETS as DREAMX_WORLD_PRESETS, ) - from fastvideo.pipelines.basic.gamecraft.presets import ( - ALL_PRESETS as GAMECRAFT_PRESETS, ) from fastvideo.pipelines.basic.gen3c.presets import ( ALL_PRESETS as GEN3C_PRESETS, ) from fastvideo.pipelines.basic.hunyuan.presets import ( @@ -1265,8 +1108,6 @@ def _register_presets() -> None: ALL_PRESETS as LINGBOTWORLD2_PRESETS, ) from fastvideo.pipelines.basic.lingbot_video.presets import ( ALL_PRESETS as LINGBOT_VIDEO_PRESETS, ) - from fastvideo.pipelines.basic.longcat.presets import ( - ALL_PRESETS as LONGCAT_PRESETS, ) from fastvideo.pipelines.basic.ltx2.presets import ( ALL_PRESETS as LTX2_PRESETS, ) from fastvideo.pipelines.basic.matrixgame2.presets import ( @@ -1281,20 +1122,14 @@ def _register_presets() -> None: ALL_PRESETS as SD35_PRESETS, ) from fastvideo.pipelines.basic.stable_audio.presets import ( ALL_PRESETS as STABLE_AUDIO_PRESETS, ) - from fastvideo.pipelines.basic.turbodiffusion.presets import ( - ALL_PRESETS as TURBODIFFUSION_PRESETS, ) from fastvideo.pipelines.basic.wan.presets import ( ALL_PRESETS as WAN_PRESETS, ) - from fastvideo.pipelines.basic.zimage.presets import ( - ALL_PRESETS as ZIMAGE_PRESETS, ) from fastvideo.pipelines.basic.flux_2.presets import ( ALL_PRESETS as FLUX2_PRESETS, ) all_preset_groups = ( COSMOS_PRESETS, - DREAMX_WORLD_PRESETS, FLUX2_PRESETS, - GAMECRAFT_PRESETS, GEN3C_PRESETS, HUNYUAN_PRESETS, HUNYUAN15_PRESETS, @@ -1305,7 +1140,6 @@ def _register_presets() -> None: LINGBOT_VIDEO_PRESETS, LINGBOTWORLD_PRESETS, LINGBOTWORLD2_PRESETS, - LONGCAT_PRESETS, LTX2_PRESETS, MATRIXGAME2_PRESETS, MATRIXGAME3_PRESETS, @@ -1313,9 +1147,7 @@ def _register_presets() -> None: MMAUDIO_PRESETS, SD35_PRESETS, STABLE_AUDIO_PRESETS, - TURBODIFFUSION_PRESETS, WAN_PRESETS, - ZIMAGE_PRESETS, ) for group in all_preset_groups: for preset in group: diff --git a/fastvideo/tests/api/test_ltx2_continuation.py b/fastvideo/tests/api/test_ltx2_continuation.py index 0ffca3473..d0714ea06 100644 --- a/fastvideo/tests/api/test_ltx2_continuation.py +++ b/fastvideo/tests/api/test_ltx2_continuation.py @@ -175,7 +175,7 @@ class TestValidation: """Invalid payloads error cleanly.""" def test_wrong_kind_rejected(self): - envelope = ContinuationState(kind="longcat.v1", payload={}) + envelope = ContinuationState(kind="other.v1", payload={}) with pytest.raises(ValueError, match="Expected ContinuationState.kind"): LTX2ContinuationState.from_continuation_state(envelope) diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index 9b0f2104e..53352cfa6 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -177,8 +177,6 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "grid_sizes": None, "c2ws_plucker_emb": None, "action_path": None, - "refine_from": None, - "stage1_video": None, }, "sampling": { "num_videos_per_prompt": 1, @@ -195,8 +193,6 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "guidance_scale": 1.0, "batch_cfg": False, "guidance_scale_2": None, - "cfg_normalization": False, - "cfg_truncation": 1.0, "guidance_rescale": 0.0, "true_cfg_scale": None, "use_embedded_guidance": None, diff --git a/fastvideo/tests/api/test_presets.py b/fastvideo/tests/api/test_presets.py index 66cf8be55..846f15ad7 100644 --- a/fastvideo/tests/api/test_presets.py +++ b/fastvideo/tests/api/test_presets.py @@ -480,30 +480,6 @@ class TestCosmosPresets: assert c25.defaults["fps"] == 24 -# ------------------------------------------------------------------- -# TurboDiffusion preset integration -# ------------------------------------------------------------------- - - -class TestTurboDiffusionPresets: - - def test_turbo_presets_registered(self) -> None: - import fastvideo.registry # noqa: F401 - presets = get_presets_for_family("turbodiffusion") - names = {p.name for p in presets} - assert names == { - "turbo_t2v_1_3b", - "turbo_t2v_14b", - "turbo_i2v_a14b", - } - - def test_turbo_4_step(self) -> None: - import fastvideo.registry # noqa: F401 - p = get_preset("turbo_t2v_14b", "turbodiffusion") - assert p.defaults["num_inference_steps"] == 4 - assert p.defaults["guidance_scale"] == 1.0 - - # ------------------------------------------------------------------- # SD35 preset integration # ------------------------------------------------------------------- @@ -663,12 +639,6 @@ class TestSinglePresetFamilies: p = get_preset("hyworld_t2v", "hyworld") assert p.workload_type == "t2v" - def test_gamecraft_registered(self) -> None: - import fastvideo.registry # noqa: F401 - p = get_preset("gamecraft_i2v", "gamecraft") - assert p.workload_type == "i2v" - assert p.defaults["num_frames"] == 33 - def test_gen3c_registered(self) -> None: import fastvideo.registry # noqa: F401 p = get_preset("gen3c_cosmos_7b", "gen3c") @@ -683,12 +653,6 @@ class TestSinglePresetFamilies: assert p3.defaults["num_inference_steps"] == 3 assert p3.defaults["height"] == 720 - def test_longcat_presets_registered(self) -> None: - import fastvideo.registry # noqa: F401 - presets = get_presets_for_family("longcat") - names = {p.name for p in presets} - assert names == {"longcat_t2v", "longcat_i2v", "longcat_vc"} - # ------------------------------------------------------------------- # Cross-family: total preset count @@ -718,8 +682,8 @@ class TestPresetDefaultTypes: ``ltx2_cfg_scale_video`` or ``ltx2_cfg_scale_audio`` is != 1.0, so any non-1.0 default silently forces CFG on for every model family that doesn't explicitly override these fields. Guard - against the regression that surfaced as the TurboDiffusion I2V - SSIM crash (``text_encoding.py:81`` assertion on + against the regression that surfaced as a Wan-based I2V SSIM + crash (``text_encoding.py:81`` assertion on ``negative_prompt``).""" from fastvideo.api.sampling_param import SamplingParam sp = SamplingParam() diff --git a/fastvideo/tests/api/test_wan_definitions.py b/fastvideo/tests/api/test_wan_definitions.py index eaa295fac..8019a7197 100644 --- a/fastvideo/tests/api/test_wan_definitions.py +++ b/fastvideo/tests/api/test_wan_definitions.py @@ -90,8 +90,6 @@ def test_legacy_aliases_and_sampling_defaults(model_path, config_name, preset_na # during a relocation, including when path and class detectors disagree. ("lucy-edit-custom", "WanPipeline", "wan", "wan_t2v_1_3b"), ("sfwan2.2-i2v-custom", "WanCausalDMDPipeline", "wan", "sf_wan_t2v_1_3b"), - ("dreamx-world-5b-lucy-edit-custom", "UnknownPipeline", "dreamx_world", "dreamx_world_5b_ar"), - ("dreamx-world-5b-custom", "WanPipeline", "wan", "wan_t2v_1_3b"), ]) def test_local_manifest_detectors_keep_first_match(tmp_path, directory, manifest_class, expected_family, expected_preset): checkpoint = tmp_path / directory diff --git a/fastvideo/tests/contract/test_env_policy.py b/fastvideo/tests/contract/test_env_policy.py index c84eb44f2..acdd2c4e2 100644 --- a/fastvideo/tests/contract/test_env_policy.py +++ b/fastvideo/tests/contract/test_env_policy.py @@ -188,7 +188,7 @@ KNOWN_VIOLATIONS: dict[str, int] = { 'fastvideo/tests/train/methods/test_minimax_h3_finetune.py: whole-environ': 2, 'fastvideo/tests/train/models/test_load_kandinsky5.py: write FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/tests/training/distill/test_anyflow_smoke.py: whole-environ': 2, - 'fastvideo/train/entrypoint/train.py: write FASTVIDEO_ATTENTION_BACKEND': 2, + 'fastvideo/train/entrypoint/train.py: write FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/utils.py: read ': 5, 'fastvideo/utils.py: write ': 1, 'fastvideo/worker/ray_distributed_executor.py: read ': 2, diff --git a/fastvideo/tests/golden_gate/test_dreamx.py b/fastvideo/tests/golden_gate/test_dreamx.py deleted file mode 100644 index a61e21c90..000000000 --- a/fastvideo/tests/golden_gate/test_dreamx.py +++ /dev/null @@ -1,127 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Golden-gate: DreamX-World-5B-Cam transformer block 0 (Wan2.2-TI2V-5B derivative -with PRoPE camera control, bidirectional diffusion variant). - -Spec notes: DreamXWorldArchConfig class defaults inherit Wan-14B dims and do NOT -match the 5B checkpoint — geometry comes from make_dreamx_world_5b_cam_dit_config() -(configs/pipelines/dreamx_world.py:19-34), verified equal to the mirror repo's -transformer/config.json. Checkpoint index is model.safetensors.index.json (saved -via save_torch_state_dict), hence weight_file="model.safetensors". temb is the 4D -[B, seq, 6, inner_dim] per-token shape (expand_timesteps=True pipeline path, -dreamx_world.py:273-282). y_camera dtype must equal hidden_states dtype (mixed -dtype einsum in _dreamx_prope_qkv errors); the camera builder is deterministic -(camera_conditioning.py:201-228). The ssim test runs TORCH_SDPA only. Forward -returns a plain [1, seq, 3072] tensor — no postprocess. The AR variant -(DreamX-World-5B, CausalWanAttentionBlock) is a different implementation and -needs its own gate recipe; not covered by this spec. -""" - -from __future__ import annotations - -import torch - -from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate - -__all__ = ["distributed_runtime"] - -INNER_DIM = 3072 # 24 heads x 128, DreamX-World-5B-Cam transformer/config.json -TEXT_LEN = 512 - - -def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.pipelines.dreamx_world import (make_dreamx_world_5b_cam_dit_config) - from fastvideo.models.dits.dreamx_world import DreamXWorldTransformerBlock - - arch = make_dreamx_world_5b_cam_dit_config().arch_config - # exactly as DreamXWorldTransformer3DModel.__init__ builds each block - # (dreamx_world.py:374-391) - return DreamXWorldTransformerBlock( - arch.num_attention_heads * arch.attention_head_dim, # 3072 - arch.ffn_dim, # 14336 - arch.num_attention_heads, # 24 - arch.qk_norm, # "rms_norm_across_heads" - arch.cross_attn_norm, # True - arch.eps, # 1e-6 - arch.added_kv_proj_dim, # None -> T2V cross-attn - arch._supported_attention_backends, - quant_config=None, - prefix=f"Wan.blocks.{layer}", - add_control_adapter=arch.add_control_adapter, # True - cam_method=arch.cam_method, # "prope" - attn_compress=arch.attn_compress, # 1 - cam_self_attn_layers=arch.cam_self_attn_layers, # None -> cam attn everywhere - layer_idx=layer, - ) - - -def _make_inputs(device: torch.device, seed: int) -> dict: - """Fixed batch mirroring the SSIM defaults (9 frames @ 64x64): - latent (3, 4, 4) -> patch (1,2,2) -> seq = 3*2*2 = 12, cameras = 3 latent frames.""" - from fastvideo.layers.rotary_embedding import get_rotary_pos_embed - from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (build_dreamx_camera_condition) - - generator = torch.Generator(device="cpu").manual_seed(seed) - num_frames, height, width = 9, 64, 64 - post_t = (num_frames - 1) // 4 + 1 # VAE 4x temporal - post_h, post_w = height // 16 // 2, width // 16 // 2 # VAE 16x spatial, patch (1,2,2) - seq = post_t * post_h * post_w # 12 - - # RoPE exactly as the model forward builds it (dreamx_world.py:429-439) - d = INNER_DIM // 24 # head_dim 128 - rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] # [44, 42, 42] - freqs_cos, freqs_sin = get_rotary_pos_embed((post_t, post_h, post_w), - INNER_DIM, - 24, - rope_dim_list, - rope_theta=10000, - dtype=torch.float64) - freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) - - hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) - encoder = torch.randn(1, TEXT_LEN, INNER_DIM, generator=generator, dtype=torch.float32) - # expand_timesteps=True pipeline -> per-token timesteps -> 4D temb (B, seq, 6, dim) - temb = torch.randn(1, seq, 6, INNER_DIM, generator=generator, dtype=torch.float32) - - # Deterministic camera condition, same builder + keys as the pipeline stage - # (stages.py:55-64). dtype MUST match hidden_states. - camera = build_dreamx_camera_condition(["w"], [4.0], - num_frames=num_frames, - height=height, - width=width, - dtype=torch.bfloat16, - device=device) - y_camera = {k: v.unsqueeze(0) for k, v in camera.items()} # viewmats (1,3,4,4), K (1,3,3,3) - - return { - "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), - "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), - "temb": temb.to(device=device, dtype=torch.bfloat16), - "freqs_cis": freqs_cis, - "original_seq_len": seq, - "y_camera": y_camera, - } - - -SPEC = GateSpec( - name="dreamx_world_5b_cam", - repo_id="FastVideo/DreamX-World-5B-Cam-Diffusers", - build_block=_build_block, - make_inputs=_make_inputs, - prefix_template="blocks.{N}.", - renames=( - # to_out rename must run before the generic attn1 strip - (r"\.attn1\.to_out\.0\.", ".to_out."), - (r"\.attn1\.", "."), # to_q/to_k/to_v/norm_q/norm_k -> block level - (r"\.attn2\.to_out\.0\.", ".attn2.to_out."), - (r"\.ffn\.net\.0\.proj\.", ".ffn.fc_in."), - (r"\.ffn\.net\.2\.", ".ffn.fc_out."), - (r"\.norm2\.", ".self_attn_residual_norm.norm."), - # identity: attn2.{to_q,to_k,to_v,norm_q,norm_k}, cam_self_attn.*, scale_shift_table - ), - attention_backend="TORCH_SDPA", - weight_file="model.safetensors", -) - - -def test_dreamx_world_cam_golden_gate(distributed_runtime) -> None: - run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_gamecraft.py b/fastvideo/tests/golden_gate/test_gamecraft.py deleted file mode 100644 index d9c5f36ac..000000000 --- a/fastvideo/tests/golden_gate/test_gamecraft.py +++ /dev/null @@ -1,87 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Golden-gate: HunyuanGameCraft double-stream block 0 (Hunyuan-video arch). - -Spec notes: the repeated block is MMDoubleStreamBlock from hunyuanvideo.py -(imported by hunyuangamecraft.py:24-28); ctor mirrors -HunyuanGameCraftTransformer3DModel.__init__ (hunyuangamecraft.py:252-261) with -hidden_size 3072 = 24 heads x 128. Checkpoint is a SINGLE FILE (no index.json) -under transformer/ and its keys are already in FastVideo form, so renames is -empty. RoPE is built over the video grid only, rope_dim_list (16,56,56), -rope_theta 256 (hunyuangamecraft.py:340-348) — applied inside -DistributedAttention to the img qkv, so freqs_cis length must equal the img -sequence length. Forward returns an (img, txt) TUPLE (hunyuanvideo.py:280); -postprocess concatenates along the sequence dim into one tensor. -""" - -from __future__ import annotations - -import torch - -import fastvideo.envs as envs -from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate - -__all__ = ["distributed_runtime"] - -HIDDEN = 3072 # 24 heads x 128, configs/models/dits/hunyuangamecraft.py:125 - - -def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftArchConfig - from fastvideo.models.dits.hunyuanvideo import MMDoubleStreamBlock - - arch = HunyuanGameCraftArchConfig() - return MMDoubleStreamBlock( - hidden_size=arch.hidden_size, # 3072 - num_attention_heads=arch.num_attention_heads, # 24 - mlp_ratio=arch.mlp_ratio, # 4.0 - dtype=arch.dtype, # None - supported_attention_backends=arch._supported_attention_backends, - prefix=f"hunyuan_gamecraft.double_blocks.{layer}", - ) - - -def _make_inputs(device: torch.device, seed: int) -> dict: - """3x4x5 = 60 video tokens + 32 text tokens, hidden 3072. - - Shapes per MMDoubleStreamBlock.forward (hunyuanvideo.py:193-200): - img [B, S_img, 3072], txt [B, S_txt, 3072], vec [B, 3072], - freqs_cis = (cos, sin) each [S_img, 128]. - """ - from fastvideo.layers.rotary_embedding import get_rotary_pos_embed - - generator = torch.Generator(device="cpu").manual_seed(seed) - tt, th, tw = 3, 4, 5 - img_seq, txt_seq = tt * th * tw, 32 - - freqs_cos, freqs_sin = get_rotary_pos_embed((tt, th, tw), HIDDEN, 24, [16, 56, 56], 256) # -> [60, 128] fp32 each - - img = torch.randn(1, img_seq, HIDDEN, generator=generator, dtype=torch.float32) - txt = torch.randn(1, txt_seq, HIDDEN, generator=generator, dtype=torch.float32) - vec = torch.randn(1, HIDDEN, generator=generator, dtype=torch.float32) - - return { - "img": img.to(device=device, dtype=torch.bfloat16), - "txt": txt.to(device=device, dtype=torch.bfloat16), - "vec": vec.to(device=device, dtype=torch.bfloat16), - "freqs_cis": (freqs_cos.to(device=device, - dtype=torch.bfloat16), freqs_sin.to(device=device, dtype=torch.bfloat16)), - } - - -SPEC = GateSpec( - name="gamecraft", - repo_id="FastVideo/HunyuanGameCraft-Diffusers", - build_block=_build_block, - make_inputs=_make_inputs, - prefix_template="double_blocks.{N}.", - renames=(), # checkpoint keys already match the block's state_dict 1:1 - attention_backend="FLASH_ATTN", - postprocess=lambda out: torch.cat(out, dim=1), # (img, txt) tuple -> one tensor - # The registry default is the Hugging Face repo id, so only an explicitly set value names a local checkout. - model_root=lambda: (envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.get() - if envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.is_set() else None), -) - - -def test_gamecraft_golden_gate(distributed_runtime) -> None: - run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_glm_image.py b/fastvideo/tests/golden_gate/test_glm_image.py deleted file mode 100644 index 1ecc65ab1..000000000 --- a/fastvideo/tests/golden_gate/test_glm_image.py +++ /dev/null @@ -1,86 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Golden-gate: GLM-Image (GlmImageTransformer2DModel) transformer block 0. - -Spec notes: inner_dim = 32 heads x 128 = 4096 (glm_image.py:631, arch config -defaults). RoPE is built outside the block by GlmImageRotaryPosEmbed -(glm_image.py:727-729) from only the shape/device of the pre-patchify latent -(B, 16, H, W); cos/sin are fp32 (image_seq, 64). temb is passed post-SiLU in -the hidden dtype (glm_image.py:747-750), so bf16 here. The block forward -returns a TUPLE (hidden_states, encoder_hidden_states) (glm_image.py:502) — -postprocess cats encoder first along seq, matching the attention layout -(glm_image.py:343). attn1.to_out is nn.ModuleList so "attn1.to_out.0.*" -matches natively; only the two ff.net renames apply. The SSIM test pins -TORCH_SDPA and loads local weights via FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR. -""" - -from __future__ import annotations - -import torch - -import fastvideo.envs as envs -from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate - -__all__ = ["distributed_runtime"] - -INNER_DIM = 4096 # 32 heads x 128, GlmImageDiTArchConfig defaults -TEXT_LEN = 32 -LATENT_H = LATENT_W = 64 # patch_size=2 -> 32x32 = 1024 image tokens - - -def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.models.dits.glm_image import GlmImageDiTArchConfig - from fastvideo.models.dits.glm_image import (GlmImageTransformer2DModel, GlmImageTransformerBlock) - - arch = GlmImageDiTArchConfig() - return GlmImageTransformerBlock( - arch.num_attention_heads * arch.attention_head_dim, # dim = 4096 - arch.num_attention_heads, # 32 - arch.attention_head_dim, # 128 - arch.time_embed_dim, # 512 - supported_attention_backends=GlmImageTransformer2DModel._supported_attention_backends, - prefix=f"transformer_blocks.{layer}", - ) - - -def _make_inputs(device: torch.device, seed: int) -> dict: - from fastvideo.models.dits.glm_image import GlmImageRotaryPosEmbed - - generator = torch.Generator(device="cpu").manual_seed(seed) - n_image = (LATENT_H // 2) * (LATENT_W // 2) # 1024 tokens - - # rope consumes only shape+device of the raw latent (glm_image.py:518-522) - rope = GlmImageRotaryPosEmbed(dim=128, patch_size=2, theta=10000.0) - cos, sin = rope(torch.empty(1, 16, LATENT_H, LATENT_W, device=device)) - - hidden = torch.randn(1, n_image, INNER_DIM, generator=generator, dtype=torch.float32) - encoder = torch.randn(1, TEXT_LEN, INNER_DIM, generator=generator, dtype=torch.float32) - temb = torch.randn(1, 512, generator=generator, dtype=torch.float32) - - return { - "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), - "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), - "temb": temb.to(device=device, dtype=torch.bfloat16), - "image_rotary_emb": (cos, sin), # fp32, already on device - } - - -SPEC = GateSpec( - name="glm_image", - repo_id="zai-org/GLM-Image", - build_block=_build_block, - make_inputs=_make_inputs, - prefix_template="transformer_blocks.{N}.", - renames=( - (r"\.ff\.net\.0\.proj\.", ".ff.fc_in."), - (r"\.ff\.net\.2\.", ".ff.fc_out."), - ), - attention_backend="TORCH_SDPA", - # forward returns (hidden_states, encoder_hidden_states); encoder first, - # like the attention concat layout (glm_image.py:343) - postprocess=lambda out: torch.cat([out[1], out[0]], dim=1), - model_root=lambda: envs.FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR.get(), -) - - -def test_glm_image_golden_gate(distributed_runtime) -> None: - run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_longcat.py b/fastvideo/tests/golden_gate/test_longcat.py deleted file mode 100644 index 004950aa6..000000000 --- a/fastvideo/tests/golden_gate/test_longcat.py +++ /dev/null @@ -1,80 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Golden-gate: LongCat-Video (native FastVideo impl) transformer block 0. - -Spec notes: LongCatVideoConfig() arch defaults match the mirror's -transformer/config.json (hidden 4096, 32 heads, mlp_ratio 4, adaln 512). -The block __init__ needs the FULL LongCatVideoConfig, not the arch config — -self_attn reads config._supported_attention_backends (longcat.py:273) while -cross_attn reads config.arch_config._supported_attention_backends (:548). -Checkpoint is a SINGLE fp32 file (transformer/model.safetensors, no -index.json) with zero renames — the mirror is pre-converted, keys match -module names. 3D RoPE is built inside the block from latent_shape; t is -per-latent-frame temb [B, T, 512]. Defaults (return_kv=False) return one -plain tensor, so no postprocess. Backend FLASH_ATTN per the ssim test -(tests/ssim/test_longcat_similarity.py:200). -""" - -from __future__ import annotations - - -import torch - -import fastvideo.envs as envs -from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate - -__all__ = ["distributed_runtime"] - -HIDDEN_SIZE = 4096 # configs/models/dits/longcat.py:77 -ADALN_DIM = 512 # configs/models/dits/longcat.py:93 - - -def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.models.dits.longcat import LongCatVideoConfig - from fastvideo.models.dits.longcat import LongCatTransformerBlock - - config = LongCatVideoConfig() - arch = config.arch_config - return LongCatTransformerBlock( - hidden_size=arch.hidden_size, - num_heads=arch.num_attention_heads, - mlp_ratio=arch.mlp_ratio, - adaln_tembed_dim=arch.adaln_tembed_dim, - config=config, - ) - - -def _make_inputs(device: torch.device, seed: int) -> dict: - # forward signature (longcat.py:773-783): x [B, N, 4096], context - # [B, N_text, 4096], t [B, T, 512], latent_shape (T, H, W), N == T*H*W. - generator = torch.Generator(device="cpu").manual_seed(seed) - t_frames, h, w = 3, 8, 8 - seq = t_frames * h * w - n_text = 32 # production pads to 512; any length works - - hidden = torch.randn(1, seq, HIDDEN_SIZE, generator=generator, dtype=torch.float32) - context = torch.randn(1, n_text, HIDDEN_SIZE, generator=generator, dtype=torch.float32) - temb = torch.randn(1, t_frames, ADALN_DIM, generator=generator, dtype=torch.float32) - - return { - "x": hidden.to(device=device, dtype=torch.bfloat16), - "context": context.to(device=device, dtype=torch.bfloat16), - "t": temb.to(device=device, dtype=torch.bfloat16), - "latent_shape": (t_frames, h, w), - } - - -SPEC = GateSpec( - name="longcat_t2v", - repo_id="FastVideo/LongCat-Video-T2V-Diffusers", - build_block=_build_block, - make_inputs=_make_inputs, - prefix_template="blocks.{N}.", - renames=(), - attention_backend="FLASH_ATTN", - model_root=lambda: envs.FASTVIDEO_TEST_LONGCAT_MODEL_ROOT.get(), - weight_file="model.safetensors", # single 54.3 GB fp32 file, no index.json -) - - -def test_longcat_golden_gate(distributed_runtime) -> None: - run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_wan_t2v.py b/fastvideo/tests/golden_gate/test_wan_t2v.py index 4bea1ee25..9fcc8dc7c 100644 --- a/fastvideo/tests/golden_gate/test_wan_t2v.py +++ b/fastvideo/tests/golden_gate/test_wan_t2v.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Golden-gate: Wan2.1-T2V-1.3B transformer block 0 (also the SFWan/TurboWan arch). +"""Golden-gate: Wan2.1-T2V-1.3B transformer block 0 (also the SFWan arch). Spec notes: arch-config defaults are the 14B model, so the 1.3B geometry is pinned explicitly from the checkpoint's transformer/config.json. RoPE is built diff --git a/fastvideo/tests/golden_gate/test_zimage.py b/fastvideo/tests/golden_gate/test_zimage.py deleted file mode 100644 index 9271a5b83..000000000 --- a/fastvideo/tests/golden_gate/test_zimage.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Golden-gate: Z-Image-Turbo (Tongyi-MAI) main-stack transformer block 0. - -Spec notes: ZImageDiTArchConfig defaults match the checkpoint's -transformer/config.json (dim=3840, n_heads=n_kv_heads=30, adaln_embed_dim=256). -Checkpoint keys match FastVideo module names exactly — no renames. Rotary is -built OUTSIDE the block: freqs_cis is complex64 from RopeEmbedder(theta=256, -axes_dims=(32,48,48)) and must stay complex64 (rotary math runs fp32 inside -attention, zimage.py:104-108). adaln_input is width min(dim, adaln_embed_dim) -= 256, not dim. Attention is plain F.scaled_dot_product_attention with a bool -[B, S] key-padding mask (True=attend); the block ignores fastvideo attention -backends, but TORCH_SDPA matches the ssim test and keys the golden filename. -Block returns a bare tensor [B, S, dim], no postprocess needed. -""" - -from __future__ import annotations - -import torch - -import fastvideo.envs as envs -from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate - -__all__ = ["distributed_runtime"] - -DIM = 3840 # 30 heads x 128, Z-Image-Turbo transformer/config.json - - -def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.models.dits.zimage import ZImageDiTArchConfig - from fastvideo.models.dits.zimage import ZImageTransformerBlock - - arch = ZImageDiTArchConfig() - return ZImageTransformerBlock( - layer, # main `layers` stack uses the plain index (zimage.py:337-338) - arch.dim, - arch.n_heads, - arch.n_kv_heads, - arch.norm_eps, - arch.qk_norm, - arch.adaln_embed_dim, - modulation=True, # main-stack blocks are modulated (zimage.py:337-338) - ) - - -def _make_inputs(device: torch.device, seed: int) -> dict: - """Fixed padded batch: 2 sequences over an 8x8 image grid, lengths [64, 48].""" - from fastvideo.models.dits.zimage import RopeEmbedder, ZImageTransformer2DModel - - generator = torch.Generator(device="cpu").manual_seed(seed) - batch, seq = 2, 64 - lengths = [64, 48] - - # freqs_cis exactly as the model builds it (zimage.py:342, 527-531); image - # position ids start at frame index cap_len+1 -> use start=(1, 0, 0). - rope = RopeEmbedder(theta=256.0, axes_dims=(32, 48, 48), axes_lens=(1536, 512, 512)) - pos_ids = ZImageTransformer2DModel.create_coordinate_grid((1, 8, 8), start=(1, 0, 0)).flatten(0, 2) # [64, 3] int32 - freqs_cis = rope(pos_ids).unsqueeze(0).expand(batch, seq, 64).contiguous() # complex64 [2, 64, 64] - - attn_mask = torch.zeros(batch, seq, dtype=torch.bool) - for i, length in enumerate(lengths): - attn_mask[i, :length] = True - - x = torch.randn(batch, seq, DIM, generator=generator, dtype=torch.float32) - adaln = torch.randn(batch, 256, generator=generator, dtype=torch.float32) - - return { - "x": x.to(device=device, dtype=torch.bfloat16), - "attn_mask": attn_mask.to(device), - "freqs_cis": freqs_cis.to(device), # stays complex64, never cast to bf16 - "adaln_input": adaln.to(device=device, dtype=torch.bfloat16), - } - - -SPEC = GateSpec( - name="zimage_turbo", - repo_id="Tongyi-MAI/Z-Image-Turbo", - build_block=_build_block, - make_inputs=_make_inputs, - prefix_template="layers.{N}.", - renames=(), - attention_backend="TORCH_SDPA", - # The registry default is the Hugging Face repo id, so only an explicitly set value names a local checkout. - model_root=lambda: (envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() - if envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.is_set() else None), -) - - -def test_zimage_golden_gate(distributed_runtime) -> None: - run_gate(SPEC) diff --git a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py index 4e2ec6800..add16e12e 100644 --- a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py +++ b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py @@ -60,9 +60,8 @@ LORA_CONFIGS = [{ } # TODO: Add a LoRA with lora_alpha values to test alpha scaling # - # Context: This change is mainly for an in-progress ticket porting over LongCat-Video, - # where they used an alpha value that is two times smaller than their rank. This fix - # ensures that LoRA weights are correctly scaled by the alpha/rank ratio when merged. + # Context: Some LoRAs use an alpha value that is two times smaller than their rank. + # This fix ensures that LoRA weights are correctly scaled by the alpha/rank ratio when merged. # # Issue: Currently, we cannot add a test for LoRA adapters with alpha values because: # - The existing public LoRAs for Wan-AI/Wan2.1-T2V-1.3B-Diffusers don't store lora_alpha @@ -70,7 +69,7 @@ LORA_CONFIGS = [{ # - This is why the alpha/rank scaling bug wasn't caught by existing tests # # The fix has been validated with: - # - LongCat-Video distilled LoRA (which includes alpha values) + # - A distilled LoRA that includes alpha values # - Manual testing shows correct alpha/rank scaling behavior # - Backward compatibility confirmed with LoRAs without alpha values # diff --git a/fastvideo/tests/loader/test_wan_family_imports.py b/fastvideo/tests/loader/test_wan_family_imports.py index 8ec1d747b..452b0b944 100644 --- a/fastvideo/tests/loader/test_wan_family_imports.py +++ b/fastvideo/tests/loader/test_wan_family_imports.py @@ -117,19 +117,14 @@ def test_legacy_vae_cache_context_is_shared_and_restored_after_errors(): @pytest.mark.parametrize("module_name", [ "fastvideo.models.dits.causal_wanvideo", - "fastvideo.models.dits.dreamx_world", "fastvideo.models.dits.matrixgame2.model", "fastvideo.models.dits.matrixgame2.causal_model", "fastvideo.models.dits.matrixgame3.model", "fastvideo.models.dits.lingbotworld.model", - "fastvideo.configs.models.dits.dreamx_world", "fastvideo.configs.models.dits.matrixgame2", "fastvideo.configs.models.dits.matrixgame3", "fastvideo.configs.pipelines.wan", - "fastvideo.configs.pipelines.dreamx_world", - "fastvideo.configs.pipelines.longcat", "fastvideo.configs.pipelines.lingbot_video", - "fastvideo.configs.pipelines.turbodiffusion", "fastvideo.configs.pipelines.lingbotworld2", ]) def test_downstream_wan_consumers_import(module_name): diff --git a/fastvideo/tests/ssim/test_dreamx_world_similarity.py b/fastvideo/tests/ssim/test_dreamx_world_similarity.py deleted file mode 100644 index e8648382e..000000000 --- a/fastvideo/tests/ssim/test_dreamx_world_similarity.py +++ /dev/null @@ -1,186 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import os -from pathlib import Path - -import pytest -from PIL import Image, ImageDraw - -import fastvideo.envs as envs -from fastvideo.logger import init_logger -from fastvideo.tests.ssim.inference_similarity_utils import ( - resolve_inference_device_reference_folder, - run_image_to_video_similarity_test, - run_text_to_video_similarity_test, -) -from fastvideo.tests.ssim.reference_utils import with_model_path - -logger = init_logger(__name__) - -REQUIRED_GPUS = 1 -pytestmark = pytest.mark.skip(reason="Disabled pending removal of DreamX World support.") - -device_reference_folder = resolve_inference_device_reference_folder(logger) - -# The tests add "model_path" from FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH and -# FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH. -DREAMX_WORLD_PARAMS = { - "num_gpus": 1, - "height": 64, - "width": 64, - "num_frames": 9, - "num_inference_steps": 1, - "guidance_scale": 1.0, - "seed": 1024, - "fps": 16, -} - -DREAMX_WORLD_FULL_QUALITY_PARAMS = { - **DREAMX_WORLD_PARAMS, - "height": 480, - "width": 832, - "num_frames": 161, - "num_inference_steps": 30, - "guidance_scale": 5.0, -} - -DREAMX_WORLD_AR_PARAMS = { - "num_gpus": 1, - "height": 192, - "width": 192, - "num_frames": 81, - "num_inference_steps": 4, - "guidance_scale": 1.0, - "seed": 2048, - "fps": 16, -} - -DREAMX_WORLD_AR_FULL_QUALITY_PARAMS = { - **DREAMX_WORLD_AR_PARAMS, - "height": 704, - "width": 1280, - "num_frames": 1005, -} - -DREAMX_WORLD_MODEL_TO_PARAMS = { - "DreamX-World-5B-Cam": DREAMX_WORLD_PARAMS, -} - -DREAMX_WORLD_AR_MODEL_TO_PARAMS = { - "DreamX-World-5B": DREAMX_WORLD_AR_PARAMS, -} - -FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS = { - "DreamX-World-5B-Cam": DREAMX_WORLD_FULL_QUALITY_PARAMS, -} - -FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS = { - "DreamX-World-5B": DREAMX_WORLD_AR_FULL_QUALITY_PARAMS, -} - -DREAMX_WORLD_TEST_CASES = [ - ( - "A cinematic first-person drive through a futuristic coastal city at sunrise, " - "reflective glass towers, clean streets, soft volumetric light.", - ("w", "d", "w"), - (4.0, 2.0, 4.0), - ), -] - -DREAMX_WORLD_AR_TEST_CASES = [ - ( - "A long autonomous drive through a futuristic coastal city at sunrise, " - "smooth forward camera motion, reflective glass towers, clean streets.", - ("w", "d", "w", "a"), - (2.0, 1.0, 2.0, 1.0), - ), -] - - -def _write_deterministic_reference_image(path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - image = Image.new("RGB", (96, 96), color=(42, 76, 112)) - draw = ImageDraw.Draw(image) - draw.rectangle((0, 56, 96, 96), fill=(36, 44, 52)) - draw.polygon([(0, 56), (48, 30), (96, 56)], fill=(120, 142, 158)) - draw.rectangle((34, 42, 62, 72), fill=(178, 198, 212)) - draw.line((0, 80, 96, 66), fill=(238, 209, 124), width=3) - image.save(path) - - -@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_TEST_CASES) -@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) -@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_MODEL_TO_PARAMS.keys())) -def test_dreamx_world_inference_similarity( - prompt: str, - action_list: tuple[str, ...], - action_speed_list: tuple[float, ...], - attention_backend_name: str, - model_id: str, - tmp_path: Path, -) -> None: - image_path = tmp_path / "dreamx_world_ssim_input.png" - _write_deterministic_reference_image(image_path) - model_path = envs.FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH.get() - - run_image_to_video_similarity_test( - logger=logger, - script_dir=os.path.dirname(os.path.abspath(__file__)), - device_reference_folder=device_reference_folder, - prompt=prompt, - image_path=str(image_path), - attention_backend_name=attention_backend_name, - model_id=model_id, - default_params_map=with_model_path(DREAMX_WORLD_MODEL_TO_PARAMS, model_path), - full_quality_params_map=with_model_path(FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS, model_path), - min_acceptable_ssim=0.98, - init_kwargs_override={ - "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", - }, - generation_kwargs_override={ - "action_list": list(action_list), - "action_speed_list": list(action_speed_list), - }, - ) - - -@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_AR_TEST_CASES) -@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) -@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_AR_MODEL_TO_PARAMS.keys())) -def test_dreamx_world_ar_inference_similarity( - prompt: str, - action_list: tuple[str, ...], - action_speed_list: tuple[float, ...], - attention_backend_name: str, - model_id: str, -) -> None: - model_path = envs.FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH.get() - run_text_to_video_similarity_test( - logger=logger, - script_dir=os.path.dirname(os.path.abspath(__file__)), - device_reference_folder=device_reference_folder, - prompt=prompt, - attention_backend_name=attention_backend_name, - model_id=model_id, - default_params_map=with_model_path(DREAMX_WORLD_AR_MODEL_TO_PARAMS, model_path), - full_quality_params_map=with_model_path(FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS, model_path), - min_acceptable_ssim=0.98, - init_kwargs_override={ - "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": "DreamXWorldARPipeline", - }, - generation_kwargs_override={ - "action_list": list(action_list), - "action_speed_list": list(action_speed_list), - }, - ) diff --git a/fastvideo/tests/ssim/test_gamecraft_similarity.py b/fastvideo/tests/ssim/test_gamecraft_similarity.py deleted file mode 100644 index e415eaaea..000000000 --- a/fastvideo/tests/ssim/test_gamecraft_similarity.py +++ /dev/null @@ -1,384 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -SSIM regression test for HunyuanGameCraft (T2V and I2V). - -Generates a video with deterministic seed and camera trajectory, -then compares against a device-specific reference video via MS-SSIM. - -Reference videos must be pre-generated and stored under: - reference_videos//_reference_videos/HunyuanGameCraft/ - / - -To create initial reference videos, run this test once and copy the -generated videos into the appropriate reference folder. -""" -import os - -import torch -import pytest - -import fastvideo.envs as envs -from fastvideo import VideoGenerator -from fastvideo.api.sampling_param import SamplingParam -from fastvideo.logger import init_logger -from fastvideo.tests.ssim.reference_utils import ( - build_generated_output_dir, - build_reference_folder_path, - get_cuda_device_name, - resolve_device_reference_folder, - select_ssim_params, - with_model_path, -) -from fastvideo.tests.utils import ( - compute_video_ssim_torchvision, - write_ssim_results, -) -from fastvideo.worker.multiproc_executor import MultiprocExecutor - -logger = init_logger(__name__) - -REQUIRED_GPUS = 1 -pytestmark = pytest.mark.skip(reason="Disabled pending removal of HunyuanGameCraft support.") - -# --------------------------------------------------------------------------- -# Device-dependent reference folder -# --------------------------------------------------------------------------- -device_reference_folder = resolve_device_reference_folder( - ( - ("A40", "A40"), - ("L40S", "L40S"), - ("A100", "A100"), - ("H100", "H100"), - ("GB200", "GB200"), - ("H200", "H200"), - ), - device_name=get_cuda_device_name(), - unknown_device_prefix="Unknown", - logger=logger, -) - -# --------------------------------------------------------------------------- -# Helpers – camera trajectory (self-contained, no official repo dependency) -# --------------------------------------------------------------------------- -from fastvideo.models.camera import create_camera_trajectory as _create_camera_trajectory - - -def _shutdown_executor(generator: VideoGenerator | None) -> None: - if generator is None: - return - if isinstance(generator.executor, MultiprocExecutor): - generator.executor.shutdown() - - -# --------------------------------------------------------------------------- -# Model parameters -# --------------------------------------------------------------------------- -# Same as basic_gamecraft.py: default HF path; set FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH for local weights. -# The tests add "model_path" from FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH. -GAMECRAFT_T2V_PARAMS = { - "num_gpus": 1, - "height": 480, - "width": 832, - "num_frames": 33, - "num_inference_steps": 20, - "guidance_scale": 6.0, - "seed": 1024, - "action": "forward", - "action_speed": 0.2, - "negative_prompt": "", -} - - -def _gamecraft_t2v_full_quality_params(model_path: str) -> dict[str, object]: - """Return the full-quality T2V params, taken from the default sampling params of ``model_path``.""" - full_quality_defaults = SamplingParam.from_pretrained(model_path) - return { - "num_gpus": GAMECRAFT_T2V_PARAMS["num_gpus"], - "model_path": model_path, - "height": full_quality_defaults.height, - "width": full_quality_defaults.width, - "num_frames": GAMECRAFT_T2V_PARAMS["num_frames"], # default num_frames: 33 - "num_inference_steps": full_quality_defaults.num_inference_steps, - "guidance_scale": full_quality_defaults.guidance_scale, - "seed": full_quality_defaults.seed, - "action": GAMECRAFT_T2V_PARAMS["action"], - "action_speed": GAMECRAFT_T2V_PARAMS["action_speed"], - "negative_prompt": full_quality_defaults.negative_prompt, - } - - -GAMECRAFT_I2V_PARAMS = { - **GAMECRAFT_T2V_PARAMS, - "image_path": ("https://huggingface.co/datasets/huggingface/documentation-images/" - "resolve/main/diffusers/astronaut.jpg"), -} - - -def _gamecraft_i2v_full_quality_params(model_path: str) -> dict[str, object]: - return { - **_gamecraft_t2v_full_quality_params(model_path), - "image_path": GAMECRAFT_I2V_PARAMS["image_path"], - } - - -MODEL_TO_PARAMS = { - "HunyuanGameCraft-T2V": GAMECRAFT_T2V_PARAMS, -} - -I2V_MODEL_TO_PARAMS = { - "HunyuanGameCraft-I2V": GAMECRAFT_I2V_PARAMS, -} - -TEST_PROMPTS = [ - "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.", -] - -I2V_TEST_PROMPTS = [ - "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background.", -] - -# --------------------------------------------------------------------------- -# T2V SSIM test -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("prompt", TEST_PROMPTS) -@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): - """Generate a T2V video with GameCraft and compare to reference via SSIM.""" - env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) - - script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = build_generated_output_dir( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - output_video_name = f"{prompt[:100].strip()}.mp4" - - os.makedirs(output_dir, exist_ok=True) - - model_path = envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.get() - params_map = select_ssim_params( - with_model_path(MODEL_TO_PARAMS, model_path), - {model_id: _gamecraft_t2v_full_quality_params(model_path)}, - ) - BASE_PARAMS = params_map[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] - - # Build camera trajectory - camera_states = _create_camera_trajectory( - action=BASE_PARAMS["action"], - height=BASE_PARAMS["height"], - width=BASE_PARAMS["width"], - num_frames=BASE_PARAMS["num_frames"], - action_speed=BASE_PARAMS["action_speed"], - dtype=torch.bfloat16, - ) - - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "use_fsdp_inference": True, - "dit_cpu_offload": True, - "vae_cpu_offload": True, - "text_encoder_cpu_offload": True, - "pin_cpu_memory": True, - } - - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "guidance_scale": BASE_PARAMS["guidance_scale"], - "seed": BASE_PARAMS["seed"], - "fps": 24, - "camera_states": camera_states, - "negative_prompt": BASE_PARAMS.get("negative_prompt", ""), - "save_video": True, - } - - generator: VideoGenerator | None = None - try: - generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - finally: - _shutdown_executor(generator) - - assert os.path.exists(output_dir), (f"Output video was not generated at {output_dir}") - - # Compare to reference - reference_folder = build_reference_folder_path( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - - if not os.path.exists(reference_folder): - logger.error("Reference folder missing") - raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") - - reference_video_name = None - for filename in os.listdir(reference_folder): - if filename.endswith(".mp4") and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - logger.error(f"Reference video not found for prompt: {prompt} " - f"with backend: {ATTENTION_BACKEND}") - raise FileNotFoundError("Reference video missing") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - generated_video_path = os.path.join(output_dir, output_video_name) - - logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - logger.info(f"Writing SSIM results to directory: {output_dir}") - - success = write_ssim_results( - output_dir, - ssim_values, - reference_video_path, - generated_video_path, - num_inference_steps, - prompt, - ) - - if not success: - logger.error("Failed to write SSIM results to file") - - min_acceptable_ssim = 0.93 - assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}") - - -# --------------------------------------------------------------------------- -# I2V SSIM test -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS) -@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -@pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys())) -def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): - """Generate an I2V video with GameCraft and compare to reference via SSIM.""" - env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) - - script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = build_generated_output_dir( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - output_video_name = f"{prompt[:100].strip()}.mp4" - - os.makedirs(output_dir, exist_ok=True) - - model_path = envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.get() - params_map = select_ssim_params( - with_model_path(I2V_MODEL_TO_PARAMS, model_path), - {model_id: _gamecraft_i2v_full_quality_params(model_path)}, - ) - BASE_PARAMS = params_map[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] - - # Build camera trajectory - camera_states = _create_camera_trajectory( - action=BASE_PARAMS["action"], - height=BASE_PARAMS["height"], - width=BASE_PARAMS["width"], - num_frames=BASE_PARAMS["num_frames"], - action_speed=BASE_PARAMS["action_speed"], - dtype=torch.bfloat16, - ) - - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "use_fsdp_inference": True, - "dit_cpu_offload": True, - "vae_cpu_offload": True, - "text_encoder_cpu_offload": True, - "pin_cpu_memory": True, - } - - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "image_path": BASE_PARAMS["image_path"], - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "guidance_scale": BASE_PARAMS["guidance_scale"], - "seed": BASE_PARAMS["seed"], - "fps": 24, - "camera_states": camera_states, - "negative_prompt": BASE_PARAMS.get("negative_prompt", ""), - "save_video": True, - } - - generator: VideoGenerator | None = None - try: - generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - finally: - _shutdown_executor(generator) - - assert os.path.exists(output_dir), (f"Output video was not generated at {output_dir}") - - # Compare to reference - reference_folder = build_reference_folder_path( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - - if not os.path.exists(reference_folder): - logger.error("Reference folder missing") - raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") - - reference_video_name = None - for filename in os.listdir(reference_folder): - if filename.endswith(".mp4") and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - logger.error(f"Reference video not found for prompt: {prompt} " - f"with backend: {ATTENTION_BACKEND}") - raise FileNotFoundError("Reference video missing") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - generated_video_path = os.path.join(output_dir, output_video_name) - - logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - logger.info(f"Writing SSIM results to directory: {output_dir}") - - success = write_ssim_results( - output_dir, - ssim_values, - reference_video_path, - generated_video_path, - num_inference_steps, - prompt, - ) - - if not success: - logger.error("Failed to write SSIM results to file") - - min_acceptable_ssim = 0.93 - assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}") diff --git a/fastvideo/tests/ssim/test_glm_image_similarity.py b/fastvideo/tests/ssim/test_glm_image_similarity.py deleted file mode 100644 index 31834733e..000000000 --- a/fastvideo/tests/ssim/test_glm_image_similarity.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""SSIM-based regression test for GLM-Image generation.""" -from __future__ import annotations - -import os -from pathlib import Path - -import pytest -import torch - -import fastvideo.envs as envs -from fastvideo.logger import init_logger -from fastvideo.tests.ssim.inference_similarity_utils import ( - run_text_to_video_similarity_test, ) -from fastvideo.tests.ssim.reference_utils import ( - get_cuda_device_name, - resolve_device_reference_folder, - with_model_path, -) - -logger = init_logger(__name__) - -REQUIRED_GPUS = 1 -pytestmark = pytest.mark.skip(reason="Disabled pending removal of GLM-Image support.") - -REPO_ROOT = Path(__file__).resolve().parents[3] - -device_reference_folder = resolve_device_reference_folder( - ( - ("A40", "A40"), - ("L40S", "L40S"), - ("H100", "H100"), - ("H200", "H200"), - # GB200 must precede B200: the lookup is a substring scan and - # "B200" matches "NVIDIA GB200", so the coarser pattern would - # otherwise claim every Grace-Blackwell host. - ("GB200", "GB200"), - ("B200", "B200"), - ), - device_name=get_cuda_device_name(), - fallback_device_prefix="L40S", - logger=logger, -) - -MODEL_ID = "zai-org__GLM-Image" - -TEST_PROMPTS = [ - "A beautiful landscape photography with rolling hills, " - "a winding river, and a vibrant sunset in the background. " - "Warm golden light, photorealistic style.", -] - -# The test adds "model_path" from FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR. -GLM_IMAGE_PARAMS = { - "num_gpus": 1, - "sp_size": 1, - "tp_size": 1, - "height": 256, - "width": 256, - "num_frames": 1, - "fps": 1, - "num_inference_steps": 4, - "guidance_scale": 1.5, - "seed": 0, - "neg_prompt": "", -} - -GLM_IMAGE_FULL_QUALITY_PARAMS = { - "num_gpus": 1, - "sp_size": 1, - "tp_size": 1, - "height": 1024, - "width": 1024, - "num_frames": 1, - "fps": 1, - "num_inference_steps": 50, - "guidance_scale": 1.5, - "seed": 0, - "neg_prompt": "", -} - -GLM_IMAGE_MODEL_TO_PARAMS = { - MODEL_ID: GLM_IMAGE_PARAMS, -} -GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS = { - MODEL_ID: GLM_IMAGE_FULL_QUALITY_PARAMS, -} - - -def _local_weights_dir() -> Path: - return Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or REPO_ROOT / "official_weights" / "glm_image") - - -def _has_weights() -> bool: - required = ["transformer", "vae", "text_encoder", "vision_language_encoder", "processor", "tokenizer", "scheduler"] - return all((_local_weights_dir() / r).exists() for r in required) - - -def _upstream_glm_image_available() -> bool: - try: - import transformers - except ImportError: - return False - return hasattr(transformers, "GlmImageForConditionalGeneration") - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="GLM-Image SSIM test requires CUDA", -) -@pytest.mark.skipif( - not _upstream_glm_image_available(), - reason="GLM-Image needs transformers>=5.0.0rc0 (ships the AR encoder).", -) -@pytest.mark.parametrize("prompt", TEST_PROMPTS) -@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) -@pytest.mark.parametrize("model_id", list(GLM_IMAGE_MODEL_TO_PARAMS.keys())) -def test_glm_image_similarity( - prompt: str, - attention_backend_name: str, - model_id: str, -) -> None: - if not _has_weights(): - pytest.skip(f"GLM-Image full weights not found at {_local_weights_dir()}.") - model_path = envs.FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR.get() or str(_local_weights_dir()) - - run_text_to_video_similarity_test( - logger=logger, - script_dir=os.path.dirname(os.path.abspath(__file__)), - device_reference_folder=device_reference_folder, - prompt=prompt, - attention_backend_name=attention_backend_name, - model_id=model_id, - default_params_map=with_model_path(GLM_IMAGE_MODEL_TO_PARAMS, model_path), - full_quality_params_map=with_model_path(GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS, model_path), - min_acceptable_ssim=0.98, - init_kwargs_override={ - "trust_remote_code": True, - "use_fsdp_inference": False, - }, - generation_kwargs_override={ - "save_video": True, - }, - ) diff --git a/fastvideo/tests/ssim/test_longcat_similarity.py b/fastvideo/tests/ssim/test_longcat_similarity.py deleted file mode 100644 index c68eddb6e..000000000 --- a/fastvideo/tests/ssim/test_longcat_similarity.py +++ /dev/null @@ -1,470 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -SSIM-based similarity tests for LongCat video generation. - -Tests three LongCat modes: -- T2V (Text-to-Video): 480p video from text prompt -- I2V (Image-to-Video): 480p video from image + text prompt -- VC (Video Continuation): 480p video continuation from input video + text prompt - -Sampling parameters are derived from: -- examples/inference/basic/basic_longcat_t2v.py -- examples/inference/basic/basic_longcat_i2v.py -- examples/inference/basic/basic_longcat_vc.py - -Note: num_inference_steps is reduced for CI speed (4 steps vs 50 in examples). -""" -import os - -import pytest -import torch - -import fastvideo.envs as envs -from fastvideo import VideoGenerator -from fastvideo.api.sampling_param import SamplingParam -from fastvideo.logger import init_logger -from fastvideo.tests.ssim.reference_utils import ( - build_generated_output_dir, - build_reference_folder_path, - get_cuda_device_name, - resolve_device_reference_folder, - select_ssim_params, -) -from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results - -logger = init_logger(__name__) - -REQUIRED_GPUS = 1 -pytestmark = pytest.mark.skip(reason="Disabled pending removal of LongCat support.") - -# Device-specific reference folder -device_reference_folder = resolve_device_reference_folder( - ( - ("A40", "A40"), - ("L40S", "L40S"), - ("H100", "H100"), - ("GB200", "GB200"), - ("H200", "H200"), - ), - device_name=get_cuda_device_name(), - logger=logger, -) - -# Common negative prompt from example scripts -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") - -# ============================================================================= -# LongCat T2V Parameters (from basic_longcat_t2v.py) -# ============================================================================= -LONGCAT_T2V_PARAMS = { - "num_gpus": 1, - "model_path": "FastVideo/LongCat-Video-T2V-Diffusers", - "height": 480, - "width": 480, - "num_frames": 43, - "num_inference_steps": 4, # Reduced from 50 for CI speed - "guidance_scale": 4.0, - "fps": 15, - "seed": 42, - "negative_prompt": NEGATIVE_PROMPT, -} -_LONGCAT_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LONGCAT_T2V_PARAMS["model_path"]) -LONGCAT_T2V_FULL_QUALITY_PARAMS = { - "num_gpus": LONGCAT_T2V_PARAMS["num_gpus"], - "model_path": LONGCAT_T2V_PARAMS["model_path"], - "height": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.height, - "width": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.width, - "num_frames": LONGCAT_T2V_PARAMS["num_frames"], # default num_frames: 125 - "num_inference_steps": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps, - "guidance_scale": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.guidance_scale, - "fps": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.fps, - "seed": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.seed, - "negative_prompt": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.negative_prompt, -} - -# ============================================================================= -# LongCat I2V Parameters (from basic_longcat_i2v.py) -# ============================================================================= -LONGCAT_I2V_PARAMS = { - "num_gpus": 1, - "model_path": "FastVideo/LongCat-Video-I2V-Diffusers", - "height": 480, - "width": 480, # Square for I2V - "num_frames": 43, - "num_inference_steps": 4, # Reduced from 50 for CI speed - "guidance_scale": 4.0, - "fps": 15, - "seed": 42, - "negative_prompt": NEGATIVE_PROMPT, -} -_LONGCAT_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LONGCAT_I2V_PARAMS["model_path"]) -LONGCAT_I2V_FULL_QUALITY_PARAMS = { - "num_gpus": LONGCAT_I2V_PARAMS["num_gpus"], - "model_path": LONGCAT_I2V_PARAMS["model_path"], - "height": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.height, - "width": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.width, - "num_frames": LONGCAT_I2V_PARAMS["num_frames"], # default num_frames: 125 - "num_inference_steps": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps, - "guidance_scale": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.guidance_scale, - "fps": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.fps, - "seed": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.seed, - "negative_prompt": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.negative_prompt, -} - -# ============================================================================= -# LongCat VC Parameters (from basic_longcat_vc.py) -# ============================================================================= -LONGCAT_VC_PARAMS = { - "num_gpus": 1, - "model_path": "FastVideo/LongCat-Video-VC-Diffusers", - "height": 480, - "width": 480, - "num_frames": 43, - "num_inference_steps": 4, # Reduced from 50 for CI speed - "guidance_scale": 4.0, - "fps": 15, - "seed": 42, - "num_cond_frames": 13, - "negative_prompt": NEGATIVE_PROMPT, -} -_LONGCAT_VC_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LONGCAT_VC_PARAMS["model_path"]) -LONGCAT_VC_FULL_QUALITY_PARAMS = { - "num_gpus": LONGCAT_VC_PARAMS["num_gpus"], - "model_path": LONGCAT_VC_PARAMS["model_path"], - "height": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.height, - "width": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.width, - "num_frames": LONGCAT_VC_PARAMS["num_frames"], # default num_frames: 125 - "num_inference_steps": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.num_inference_steps, - "guidance_scale": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.guidance_scale, - "fps": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.fps, - "seed": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.seed, - "num_cond_frames": LONGCAT_VC_PARAMS["num_cond_frames"], - "negative_prompt": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.negative_prompt, -} - -# Test prompts -T2V_TEST_PROMPTS = [ - "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.", -] - -I2V_TEST_PROMPTS = [ - "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.", -] - -I2V_IMAGE_PATHS = [ - "assets/girl.png", -] - -VC_TEST_PROMPTS = [ - "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.", -] - -VC_VIDEO_PATHS = [ - "assets/motorcycle.mp4", -] - - -def _resolve_asset_path(asset_path: str) -> str: - """Resolve asset path relative to FastVideo root.""" - # Check if absolute or already exists - if os.path.isabs(asset_path) or os.path.exists(asset_path): - return asset_path - # Try relative to workspace root - script_dir = os.path.dirname(os.path.abspath(__file__)) - repo_root = os.path.abspath(os.path.join(script_dir, "..", "..", "..")) - return os.path.join(repo_root, asset_path) - - -@pytest.mark.parametrize("prompt", T2V_TEST_PROMPTS) -@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): - """ - Test LongCat T2V inference and compare output to reference videos using SSIM. - - Parameters derived from examples/inference/basic/basic_longcat_t2v.py - """ - env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) - - params = select_ssim_params(LONGCAT_T2V_PARAMS, LONGCAT_T2V_FULL_QUALITY_PARAMS) - - script_dir = os.path.dirname(os.path.abspath(__file__)) - model_id = "LongCat-Video-T2V" - - output_dir = build_generated_output_dir( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - output_video_name = f"{prompt[:100].strip()}.mp4" - os.makedirs(output_dir, exist_ok=True) - - init_kwargs = { - "num_gpus": params["num_gpus"], - "use_fsdp_inference": True, - "dit_cpu_offload": True, - "vae_cpu_offload": True, - "text_encoder_cpu_offload": True, - "enable_bsa": False, - } - - generation_kwargs = { - "output_path": output_dir, - "height": params["height"], - "width": params["width"], - "num_frames": params["num_frames"], - "num_inference_steps": params["num_inference_steps"], - "guidance_scale": params["guidance_scale"], - "fps": params["fps"], - "seed": params["seed"], - "negative_prompt": params["negative_prompt"], - } - - generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - generator.shutdown() - - generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") - - # Find reference video - reference_folder = build_reference_folder_path( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - if not os.path.exists(reference_folder): - raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") - - reference_video_name = None - for filename in os.listdir(reference_folder): - if filename.endswith(".mp4") and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - raise FileNotFoundError( - f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - - logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - - write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, - params["num_inference_steps"], prompt) - - min_acceptable_ssim = 0.90 - assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}") - - -@pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS) -@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): - """ - Test LongCat I2V inference and compare output to reference videos using SSIM. - - Parameters derived from examples/inference/basic/basic_longcat_i2v.py - """ - env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) - - params = select_ssim_params(LONGCAT_I2V_PARAMS, LONGCAT_I2V_FULL_QUALITY_PARAMS) - - script_dir = os.path.dirname(os.path.abspath(__file__)) - model_id = "LongCat-Video-I2V" - - output_dir = build_generated_output_dir( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - output_video_name = f"{prompt[:100].strip()}.mp4" - os.makedirs(output_dir, exist_ok=True) - - # Get image path for this prompt - prompt_idx = I2V_TEST_PROMPTS.index(prompt) - image_path = _resolve_asset_path(I2V_IMAGE_PATHS[prompt_idx]) - - init_kwargs = { - "num_gpus": params["num_gpus"], - "use_fsdp_inference": True, - "dit_cpu_offload": True, - "vae_cpu_offload": True, - "text_encoder_cpu_offload": True, - "enable_bsa": False, - } - - generation_kwargs = { - "output_path": output_dir, - "image_path": image_path, - "height": params["height"], - "width": params["width"], - "num_frames": params["num_frames"], - "num_inference_steps": params["num_inference_steps"], - "guidance_scale": params["guidance_scale"], - "fps": params["fps"], - "seed": params["seed"], - "negative_prompt": params["negative_prompt"], - } - - generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - generator.shutdown() - - generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") - - # Find reference video - reference_folder = build_reference_folder_path( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - if not os.path.exists(reference_folder): - raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") - - reference_video_name = None - for filename in os.listdir(reference_folder): - if filename.endswith(".mp4") and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - raise FileNotFoundError( - f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - - logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - - write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, - params["num_inference_steps"], prompt) - - min_acceptable_ssim = 0.90 - assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}") - - -@pytest.mark.parametrize("prompt", VC_TEST_PROMPTS) -@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): - """ - Test LongCat VC (Video Continuation) inference and compare output to reference videos using SSIM. - - Parameters derived from examples/inference/basic/basic_longcat_vc.py - """ - env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) - - params = select_ssim_params(LONGCAT_VC_PARAMS, LONGCAT_VC_FULL_QUALITY_PARAMS) - - script_dir = os.path.dirname(os.path.abspath(__file__)) - model_id = "LongCat-Video-VC" - - output_dir = build_generated_output_dir( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - output_video_name = f"{prompt[:100].strip()}.mp4" - os.makedirs(output_dir, exist_ok=True) - - # Get video path for this prompt - prompt_idx = VC_TEST_PROMPTS.index(prompt) - video_path = _resolve_asset_path(VC_VIDEO_PATHS[prompt_idx]) - - if not os.path.exists(video_path): - pytest.skip(f"Input video not found at {video_path}") - - init_kwargs = { - "num_gpus": params["num_gpus"], - "use_fsdp_inference": False, - "dit_cpu_offload": False, - "vae_cpu_offload": True, - "text_encoder_cpu_offload": True, - "pin_cpu_memory": False, - "enable_bsa": False, - } - - generation_kwargs = { - "output_path": output_dir, - "video_path": video_path, - "num_cond_frames": params["num_cond_frames"], - "height": params["height"], - "width": params["width"], - "num_frames": params["num_frames"], - "num_inference_steps": params["num_inference_steps"], - "guidance_scale": params["guidance_scale"], - "fps": params["fps"], - "seed": params["seed"], - "negative_prompt": params["negative_prompt"], - } - - generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - generator.shutdown() - - generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") - - # Find reference video - reference_folder = build_reference_folder_path( - script_dir, - device_reference_folder, - model_id, - ATTENTION_BACKEND, - ) - if not os.path.exists(reference_folder): - raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") - - reference_video_name = None - for filename in os.listdir(reference_folder): - if filename.endswith(".mp4") and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - raise FileNotFoundError( - f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - - logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - - write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, - params["num_inference_steps"], prompt) - - min_acceptable_ssim = 0.90 - assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}") diff --git a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py deleted file mode 100644 index f73feea6b..000000000 --- a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py +++ /dev/null @@ -1,172 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import os - -import pytest - -from fastvideo.api.sampling_param import SamplingParam -from fastvideo.logger import init_logger -from fastvideo.tests.ssim.inference_similarity_utils import ( - DEVICE_MAPPINGS, - resolve_inference_device_reference_folder, - run_image_to_video_similarity_test, - run_text_to_video_similarity_test, -) -from fastvideo.tests.ssim.reference_utils import ( - get_cuda_device_name, - resolve_device_reference_folder, -) - -logger = init_logger(__name__) - -REQUIRED_GPUS = 4 -pytestmark = pytest.mark.skip(reason="Disabled pending removal of TurboDiffusion and TurboWan support.") - -device_name = get_cuda_device_name() -device_reference_folder = resolve_device_reference_folder( - DEVICE_MAPPINGS, - device_name=device_name, -) -if device_reference_folder is None: - raise ValueError(f"Unsupported device for ssim tests: {device_name}") - -TURBODIFFUSION_PARAMS = { - "num_gpus": 4, - "model_path": "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers", - "height": 480, - "width": 832, - "num_frames": 81, - "num_inference_steps": 4, - "guidance_scale": 1.0, - "seed": 42, - "sp_size": 4, - "tp_size": 1, - "fps": 24, -} -_TURBODIFFUSION_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(TURBODIFFUSION_PARAMS["model_path"]) -TURBODIFFUSION_FULL_QUALITY_PARAMS = { - "num_gpus": TURBODIFFUSION_PARAMS["num_gpus"], - "model_path": TURBODIFFUSION_PARAMS["model_path"], - "height": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.height, - "width": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.width, - "num_frames": TURBODIFFUSION_PARAMS["num_frames"], - "num_inference_steps": (_TURBODIFFUSION_FULL_QUALITY_DEFAULTS.num_inference_steps), - "guidance_scale": (_TURBODIFFUSION_FULL_QUALITY_DEFAULTS.guidance_scale), - "seed": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.seed, - "sp_size": TURBODIFFUSION_PARAMS["sp_size"], - "tp_size": TURBODIFFUSION_PARAMS["tp_size"], - "fps": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.fps, -} - -TURBODIFFUSION_MODEL_TO_PARAMS = { - "TurboWan2.1-T2V-1.3B-Diffusers": TURBODIFFUSION_PARAMS, -} -FULL_QUALITY_TURBODIFFUSION_MODEL_TO_PARAMS = { - "TurboWan2.1-T2V-1.3B-Diffusers": (TURBODIFFUSION_FULL_QUALITY_PARAMS), -} - -TURBODIFFUSION_TEST_PROMPTS = [ - "Will Smith casually eats noodles, his relaxed demeanor contrasting " - "with the energetic background of a bustling street food market. " - "The scene captures a mix of humor and authenticity. Mid-shot " - "framing, vibrant lighting.", -] - - -@pytest.mark.parametrize("prompt", TURBODIFFUSION_TEST_PROMPTS) -@pytest.mark.parametrize( - "model_id", - list(TURBODIFFUSION_MODEL_TO_PARAMS.keys()), -) -def test_turbodiffusion_inference_similarity( - prompt: str, - model_id: str, -) -> None: - run_text_to_video_similarity_test( - logger=logger, - script_dir=os.path.dirname(os.path.abspath(__file__)), - device_reference_folder=device_reference_folder, - prompt=prompt, - attention_backend_name="SLA_ATTN", - model_id=model_id, - default_params_map=TURBODIFFUSION_MODEL_TO_PARAMS, - full_quality_params_map=(FULL_QUALITY_TURBODIFFUSION_MODEL_TO_PARAMS), - min_acceptable_ssim=0.95, - init_kwargs_override={ - "override_pipeline_cls_name": "TurboDiffusionPipeline", - }, - ) - - -TURBODIFFUSION_I2V_PARAMS = { - "num_gpus": 4, - "model_path": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers", - "height": 480, - "width": 832, - "num_frames": 45, - "num_inference_steps": 4, - "guidance_scale": 1.0, - "seed": 42, - "sp_size": 4, - "tp_size": 1, - "fps": 24, -} -_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(TURBODIFFUSION_I2V_PARAMS["model_path"]) -TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS = { - "num_gpus": TURBODIFFUSION_I2V_PARAMS["num_gpus"], - "model_path": TURBODIFFUSION_I2V_PARAMS["model_path"], - "height": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.height, - "width": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.width, - "num_frames": TURBODIFFUSION_I2V_PARAMS["num_frames"], - "num_inference_steps": (_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps), - "guidance_scale": (_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.guidance_scale), - "seed": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.seed, - "sp_size": TURBODIFFUSION_I2V_PARAMS["sp_size"], - "tp_size": TURBODIFFUSION_I2V_PARAMS["tp_size"], - "fps": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.fps, -} - -TURBODIFFUSION_I2V_MODEL_TO_PARAMS = { - "TurboWan2.2-I2V-A14B-Diffusers": TURBODIFFUSION_I2V_PARAMS, -} -FULL_QUALITY_TURBODIFFUSION_I2V_MODEL_TO_PARAMS = { - "TurboWan2.2-I2V-A14B-Diffusers": (TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS), -} - -TURBODIFFUSION_I2V_TEST_PROMPTS = [ - "An astronaut hatching from an egg, on the surface of the moon, " - "the darkness and depth of space realised in the background. " - "High quality, ultrarealistic detail and breath-taking " - "movie-like camera shot.", -] - -TURBODIFFUSION_I2V_IMAGE_PATHS = [ - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", -] - - -@pytest.mark.skip(reason="Disabled: causes OOM too often in CI") -@pytest.mark.parametrize("prompt", TURBODIFFUSION_I2V_TEST_PROMPTS) -@pytest.mark.parametrize( - "model_id", - list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys()), -) -def test_turbodiffusion_i2v_inference_similarity( - prompt: str, - model_id: str, -) -> None: - image_path = TURBODIFFUSION_I2V_IMAGE_PATHS[TURBODIFFUSION_I2V_TEST_PROMPTS.index(prompt)] - run_image_to_video_similarity_test( - logger=logger, - script_dir=os.path.dirname(os.path.abspath(__file__)), - device_reference_folder=device_reference_folder, - prompt=prompt, - image_path=image_path, - attention_backend_name="SLA_ATTN", - model_id=model_id, - default_params_map=TURBODIFFUSION_I2V_MODEL_TO_PARAMS, - full_quality_params_map=(FULL_QUALITY_TURBODIFFUSION_I2V_MODEL_TO_PARAMS), - min_acceptable_ssim=0.95, - init_kwargs_override={ - "override_pipeline_cls_name": ("TurboDiffusionI2VPipeline"), - }, - ) diff --git a/fastvideo/tests/ssim/test_zimage_similarity.py b/fastvideo/tests/ssim/test_zimage_similarity.py deleted file mode 100644 index f739e9cbf..000000000 --- a/fastvideo/tests/ssim/test_zimage_similarity.py +++ /dev/null @@ -1,117 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""PNG regression for Tongyi-MAI/Z-Image@26f23eda's native Turbo recipe.""" - -from __future__ import annotations - -import os - -import pytest -import torch - -import fastvideo.envs as envs -from fastvideo.logger import init_logger -from fastvideo.tests.ssim.inference_similarity_utils import ( - run_text_to_video_similarity_test, ) -from fastvideo.tests.ssim.reference_utils import ( - get_cuda_device_name, - resolve_device_reference_folder, - with_model_path, -) - -logger = init_logger(__name__) - -REQUIRED_GPUS = 1 -pytestmark = pytest.mark.skip(reason="Disabled pending removal of Z-Image support.") -ZIMAGE_MIN_SSIM = 0.98 - -TEST_PROMPTS = [ - "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.", -] - -device_reference_folder = resolve_device_reference_folder( - ( - ("A40", "A40"), - ("L40S", "L40S"), - ("H100", "H100"), - ("H200", "H200"), - # GB200 must precede B200: the lookup is a substring scan and - # "B200" matches "NVIDIA GB200", so the coarser pattern would - # otherwise claim every Grace-Blackwell host. - ("GB200", "GB200"), - ("B200", "B200"), - ), - device_name=get_cuda_device_name(), - fallback_device_prefix="L40S", - logger=logger, -) - -# The test adds "model_path" from FASTVIDEO_TEST_ZIMAGE_MODEL_DIR. -ZIMAGE_MODEL_TO_PARAMS: dict[str, dict[str, object]] = { - "Tongyi-MAI__Z-Image-Turbo": { - "num_gpus": 1, - "sp_size": 1, - "tp_size": 1, - "height": 1024, - "width": 1024, - "num_frames": 1, - "fps": 1, - "num_inference_steps": 8, - "guidance_scale": 0.0, - "seed": 42, - }, -} - -ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS = {model_id: dict(params) for model_id, params in ZIMAGE_MODEL_TO_PARAMS.items()} - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="Z-Image SSIM test requires CUDA", -) -@pytest.mark.parametrize("prompt", TEST_PROMPTS) -@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) -@pytest.mark.parametrize("model_id", list(ZIMAGE_MODEL_TO_PARAMS)) -def test_zimage_similarity( - prompt: str, - attention_backend_name: str, - model_id: str, -) -> None: - zimage_model_path = envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() - is_hf_repo = "/" in zimage_model_path and not zimage_model_path.startswith("/") - if not is_hf_repo and not os.path.isdir(zimage_model_path): - pytest.skip(f"Z-Image weights not found at {zimage_model_path} " - "(set FASTVIDEO_TEST_ZIMAGE_MODEL_DIR to override)") - - run_text_to_video_similarity_test( - logger=logger, - script_dir=os.path.dirname(os.path.abspath(__file__)), - device_reference_folder=device_reference_folder, - prompt=prompt, - attention_backend_name=attention_backend_name, - model_id=model_id, - default_params_map=with_model_path(ZIMAGE_MODEL_TO_PARAMS, zimage_model_path), - full_quality_params_map=with_model_path(ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS, zimage_model_path), - min_acceptable_ssim=ZIMAGE_MIN_SSIM, - media_extension=".png", - init_kwargs_override={ - "revision": envs.FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION.get(), - "workload_type": "t2i", - "trust_remote_code": True, - "use_fsdp_inference": False, - "dit_cpu_offload": False, - "text_encoder_cpu_offload": False, - "vae_cpu_offload": False, - "image_encoder_cpu_offload": False, - "pin_cpu_memory": False, - }, - generation_kwargs_override={ - "save_video": True, - "negative_prompt": "", - "max_sequence_length": 512, - "cfg_normalization": False, - "cfg_truncation": 1.0, - }, - ) diff --git a/fastvideo/tests/stages/test_wan_pipeline_wiring.py b/fastvideo/tests/stages/test_wan_pipeline_wiring.py index 900f9aa59..74db356d9 100644 --- a/fastvideo/tests/stages/test_wan_pipeline_wiring.py +++ b/fastvideo/tests/stages/test_wan_pipeline_wiring.py @@ -19,9 +19,6 @@ from fastvideo.tests.stages._denoising_fixtures import RecordingDenoiser, TinyVA ("wan.wan_i2v_dmd_pipeline", "DmdDenoisingStage", False), ("wan.wan_causal_pipeline", "CausalDenoisingStage", False), ("wan.wan_causal_dmd_pipeline", "CausalDMDDenosingStage", True), - ("dreamx_world.dreamx_world_pipeline", "WanDenoisingStage", True), - ("turbodiffusion.turbodiffusion_pipeline", "WanDenoisingStage", False), - ("turbodiffusion.turbodiffusion_i2v_pipeline", "WanDenoisingStage", False), ]) def test_pipeline_wires_family_sampler_and_owns_scheduler(monkeypatch, env_overrides, module_name, sampler, first_frame): diff --git a/fastvideo/tests/train/fixtures/longcat_t2v_min.yaml b/fastvideo/tests/train/fixtures/longcat_t2v_min.yaml deleted file mode 100644 index 7969d8fad..000000000 --- a/fastvideo/tests/train/fixtures/longcat_t2v_min.yaml +++ /dev/null @@ -1,23 +0,0 @@ -# Minimum config to instantiate LongCatModel for the loading + forward -# smoke test. Real LongCat-Video-T2V checkpoint (~13.6B at bf16) loaded -# via load_module_from_path. flow_shift=12.0 matches the released -# LongCat scheduler config (0.0 would collapse FlowMatch timesteps). -# -# Forward-only: a single-GPU backward of the 13.6B transformer exceeds -# the L40S CI runner, so LongCat is covered by this loading smoke only -# (no per-method grad-norm test). - -models: - student: - _target_: fastvideo.train.models.longcat.LongCatModel - init_from: FastVideo/LongCat-Video-T2V-Diffusers - trainable: false - -method: - _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod - -training: - dit_precision: bf16 - -pipeline: - flow_shift: 12.0 diff --git a/fastvideo/tests/train/methods/test_matrixgame2_finetune.py b/fastvideo/tests/train/methods/test_matrixgame2_finetune.py index 5c5f33d8d..e353439e6 100644 --- a/fastvideo/tests/train/methods/test_matrixgame2_finetune.py +++ b/fastvideo/tests/train/methods/test_matrixgame2_finetune.py @@ -13,8 +13,7 @@ receive gradients. The goal is to validate the training grad path Matrix-Game-2.0-Base is ~14B at bf16. A single-GPU backward (params + grads, no FSDP) needs well over the L40S CI runner's 48 GB, so this test skips on memory-constrained GPUs and only exercises the grad path -on larger dev GPUs (H200 / GB200). LongCat (13.6B) is likewise covered -by a loading-only smoke for the same reason. +on larger dev GPUs (H200 / GB200). """ from __future__ import annotations diff --git a/fastvideo/tests/train/models/test_load_longcat.py b/fastvideo/tests/train/models/test_load_longcat.py deleted file mode 100644 index 7bf0f8a47..000000000 --- a/fastvideo/tests/train/models/test_load_longcat.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""GPU loading + forward smoke test for ``LongCatModel``. - -Loads the real LongCat-Video-T2V checkpoint (~13.6B at bf16) via -``LongCatModel.__init__`` and runs one transformer forward pass on -synthetic inputs. Catches loader or forward-signature regressions in -``fastvideo.train.models.longcat.LongCatModel`` and the underlying -``LongCatTransformer3DModel``. - -Forward-only (``trainable=False``, ``torch.no_grad``): a single-GPU -backward of the 13.6B transformer exceeds the L40S CI runner, so -LongCat training is not covered by a per-method grad-norm test. The -forward fits (the ~13B HunyuanModel loading test already runs on the -same runner). -""" - -from __future__ import annotations - -import fastvideo.envs as envs - -# Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. -envs.setdefault_external("MASTER_ADDR", "localhost") -envs.setdefault_external("MASTER_PORT", "29519") - -from pathlib import Path - -import pytest -import torch - -from fastvideo.forward_context import set_forward_context -from fastvideo.train.models.longcat import LongCatModel -from fastvideo.train.utils.config import load_run_config - -_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "longcat_t2v_min.yaml") - -# LongCat caption_channels (text embedding width). -_LONGCAT_TEXT_DIM = 4096 - - -@pytest.mark.usefixtures("distributed_setup") -def test_longcat_model_loads_and_forwards(): - if not torch.cuda.is_available(): - pytest.skip("requires CUDA") - - cfg = load_run_config(_FIXTURE) - model = LongCatModel( - init_from=cfg.models["student"]["init_from"], - training_config=cfg.training, - trainable=False, - ) - - transformer = model.transformer - assert isinstance(transformer, torch.nn.Module) - assert sum(p.numel() for p in transformer.parameters()) > 0 - - device = torch.device("cuda:0") - dtype = torch.bfloat16 - transformer = transformer.to(device=device, dtype=dtype).eval() - - # LongCat transformer takes [B, C, T, H, W] (in_channels=16, - # patch_size=(1,2,2)) and a [B, N_text, 4096] text embedding. - hidden_states = torch.randn(1, 16, 4, 32, 32, device=device, dtype=dtype) - encoder_hidden_states = torch.randn(1, 16, _LONGCAT_TEXT_DIM, device=device, dtype=dtype) - encoder_attention_mask = torch.ones(1, 16, device=device, dtype=dtype) - timestep = torch.tensor([500], device=device, dtype=dtype) - - with torch.no_grad(), set_forward_context( - current_timestep=0, - attn_metadata=None, - ): - out = transformer( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - encoder_attention_mask=encoder_attention_mask, - ) - - if isinstance(out, tuple): - out = out[0] - assert out.shape == hidden_states.shape, (f"output shape {tuple(out.shape)} != input shape " - f"{tuple(hidden_states.shape)}") - assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_wan_override_contract.py b/fastvideo/tests/train/models/test_wan_override_contract.py index fc391c6a5..30d07deef 100644 --- a/fastvideo/tests/train/models/test_wan_override_contract.py +++ b/fastvideo/tests/train/models/test_wan_override_contract.py @@ -6,11 +6,10 @@ import inspect import pytest from fastvideo.train.models.hunyuan.hunyuan import HunyuanModel -from fastvideo.train.models.longcat.longcat import LongCatModel from fastvideo.train.models.matrixgame2.matrixgame2 import MatrixGame2Model -@pytest.mark.parametrize("model_cls", [HunyuanModel, LongCatModel, MatrixGame2Model]) +@pytest.mark.parametrize("model_cls", [HunyuanModel, MatrixGame2Model]) def test_distill_kwargs_accept_wan_teacher_forcing_keywords(model_cls: type) -> None: parameters = inspect.signature(model_cls._build_distill_input_kwargs).parameters @@ -18,7 +17,7 @@ def test_distill_kwargs_accept_wan_teacher_forcing_keywords(model_cls: type) -> assert "aug_t" in parameters -@pytest.mark.parametrize("model_cls", [LongCatModel, MatrixGame2Model]) +@pytest.mark.parametrize("model_cls", [MatrixGame2Model]) def test_predict_noise_override_preserves_wan_teacher_forcing_keywords(model_cls: type) -> None: parameters = inspect.signature(model_cls.predict_noise).parameters diff --git a/fastvideo/tests/train/utils/test_finetune_config_construction.py b/fastvideo/tests/train/utils/test_finetune_config_construction.py index baaa5eed3..cb726daa4 100644 --- a/fastvideo/tests/train/utils/test_finetune_config_construction.py +++ b/fastvideo/tests/train/utils/test_finetune_config_construction.py @@ -40,4 +40,4 @@ def test_fine_tuning_recipe_constructs(recipe: Path) -> None: def test_recipe_glob_found_recipes() -> None: # Guard against the glob silently matching nothing after a move. - assert len(_RECIPES) >= 20 + assert len(_RECIPES) >= 15 diff --git a/fastvideo/tests/transformers/test_hunyuangamecraft.py b/fastvideo/tests/transformers/test_hunyuangamecraft.py deleted file mode 100644 index 66864a79f..000000000 --- a/fastvideo/tests/transformers/test_hunyuangamecraft.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Regression test for HunyuanGameCraft transformer. - -Tests that the FastVideo HunyuanGameCraft implementation produces consistent outputs. -""" -import os - -import pytest -import torch - -import fastvideo.envs as envs -from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig -from fastvideo.configs.pipelines import PipelineConfig -from fastvideo.distributed.parallel_state import ( - get_sp_parallel_rank, - get_sp_world_size, -) -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger -from fastvideo.models.loader.component_loader import TransformerLoader -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - -logger = init_logger(__name__) - -envs.setdefault_external("MASTER_ADDR", "localhost") -envs.setdefault_external("MASTER_PORT", "29507") - -# Path to converted weights (local path, not HuggingFace) -TRANSFORMER_PATH = "official_weights/hunyuan-gamecraft/transformer" - -# Reference latent computed from FastVideo HunyuanGameCraft model with: -# - seed=42, batch=1, frames=9, H=44, W=80, text_seq=32 -# - 33 camera frames at 352x640 resolution -# - timestep=500 -REFERENCE_LATENT = 42351.12903189659 - - -@pytest.mark.usefixtures("distributed_setup") -def test_hunyuangamecraft_transformer(env_overrides): - """Test HunyuanGameCraft transformer regression.""" - - if not os.path.exists(TRANSFORMER_PATH): - pytest.skip(f"Weights not found at {TRANSFORMER_PATH}") - - # Keep this model-specific override local to this test. Setting - # TORCHDYNAMO_DISABLE during module collection disables torch.compile for - # every later transformer test in the same pytest process; FakeTensor - # compile checks then execute Triton kernels against fake CUDA pointers. - env_overrides.enter_context(envs.override_external("TORCHDYNAMO_DISABLE", "1")) - - sp_rank = get_sp_parallel_rank() - sp_world_size = get_sp_world_size() - - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - precision = torch.bfloat16 - precision_str = "bf16" - - args = FastVideoArgs( - model_path=TRANSFORMER_PATH, - dit_cpu_offload=False, - use_fsdp_inference=False, - pipeline_config=PipelineConfig(dit_config=HunyuanGameCraftConfig(), dit_precision=precision_str), - ) - args.device = device - - loader = TransformerLoader() - model = loader.load(TRANSFORMER_PATH, args).to(device, dtype=precision) - model.eval() - - total_params = sum(p.numel() for p in model.parameters()) - weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters()) - weight_mean = weight_sum / total_params - logger.info("Total parameters: %s", total_params) - logger.info("Weight sum: %s", weight_sum) - logger.info("Weight mean: %s", weight_mean) - - torch.manual_seed(42) - - batch_size = 1 - latent_frames = 9 # GameCraft uses 9 latent frames - latent_height = 44 - latent_width = 80 - text_seq_len = 32 - - # Input latents [B, 33, T, H, W] - 16 latent + 16 gt_latent + 1 mask - hidden_states = torch.randn(batch_size, - 33, - latent_frames, - latent_height, - latent_width, - device=device, - dtype=precision) - - if sp_world_size > 1: - chunk_per_rank = hidden_states.shape[2] // sp_world_size - hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank] - - # Text embeddings (LLaMA) - text_states = torch.randn(batch_size, text_seq_len, 4096, device=device, dtype=precision) - - # CLIP pooled embeddings - text_states_2 = torch.randn(batch_size, 768, device=device, dtype=precision) - - # Text mask - text_mask = torch.ones(batch_size, text_seq_len, device=device, dtype=torch.long) - - # Timestep - timestep = torch.tensor([500], device=device, dtype=precision) - - # Camera states [B, num_frames, 6, H, W] - 33 video frames at full resolution - # For 9 latent frames, we use 33 camera frames (matching official model) - camera_states = torch.randn(batch_size, 33, 6, 352, 640, device=device, dtype=precision) - - encoder_hidden_states = [text_states, text_states_2] - - forward_batch = ForwardBatch(data_type="video") - - with torch.no_grad(): - with torch.amp.autocast('cuda', dtype=precision): - with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch): - output = model( - hidden_states, - encoder_hidden_states, - timestep, - camera_states=camera_states, - encoder_attention_mask=[text_mask], - ) - - latent = output.double().sum().item() - logger.info(f"Output shape: {output.shape}") - logger.info(f"Current latent: {latent}") - - if REFERENCE_LATENT is not None: - diff = abs(REFERENCE_LATENT - latent) - relative_diff = diff / abs(REFERENCE_LATENT) - logger.info(f"Reference latent: {REFERENCE_LATENT}") - logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%") - - assert relative_diff < 0.005, \ - f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)" - else: - logger.info(f"No reference latent set. To enable regression testing, set REFERENCE_LATENT = {latent}") diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/__init__.py b/fastvideo/third_party/longcat_video/block_sparse_attention/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py b/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py deleted file mode 100644 index 1b6bbd6c9..000000000 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py +++ /dev/null @@ -1,737 +0,0 @@ -import os -import torch -import torch.nn.functional as F -import triton -import triton.language as tl -import math - -from .common import _attn_fwd_gating, _attn_bwd_preprocess, configs_gating_preset -from .flash_attn_bsa_varlen_mask import (_attn_fwd_bsa_varlen, _attn_fwd_bsa_varlen_align, - _attn_bwd_dkdv_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_wrapper, - _attn_bwd_dq_bsa_varlen_align_wrapper, configs_fwd_bsa_varlen_preset, - configs_fwd_bsa_varlen_align_preset, configs_bwd_dkdv_bsa_varlen_preset, - configs_bwd_dq_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_align_preset) - -torch._dynamo.config.cache_size_limit = 32 - - -def is_cuda(): - return triton.runtime.driver.active.get_current_target().backend == "cuda" - - -def supports_tma(): - return is_cuda() and torch.cuda.get_device_capability()[0] >= 9 - - -HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl) - -if HAS_TMA_DESC: - print("TMA benchmarks will be running with experimental grid constant TMA descriptor.", ) -else: - print("TMA benchmarks will be running without grid constant TMA descriptor.", ) - - -# TmaAutoTuneHelper used in htyu's PR #5622 -class TmaAutoTuneHelper: - - # duck typing wrapper to implement the same interface as TmaDescKernelParam in Triton PR #4498 - class KernelParamWrapper: - - def __init__(self, desc): - self.desc = desc - - def tma_desc_cpu_ptr(self): - return self.desc.data_ptr() - - TMA_SIZE = 128 - - def __init__(self): - self.fill_1d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_1d_tma_descriptor) - self.fill_2d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_2d_tma_descriptor) - if HAS_TMA_DESC: - self.descriptors = {} - else: - self.cuda_descriptors = {} - - # Call this method outside of the lambda function for grid size - def init_tma_descriptor(self, name): - if HAS_TMA_DESC: - self.descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device="cpu", dtype=torch.int8) - else: - self.cuda_descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device="cuda", dtype=torch.int8) - - # Call this method inside the lambda function for grid size - def fill_1d_tma_descriptor(self, name, ptr, dim, block_dim, element_size): - if HAS_TMA_DESC: - desc_x = self.descriptors[name] - assert desc_x.data_ptr() % 64 == 0 - self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, desc_x.data_ptr()) - else: - desc_x = self.cuda_descriptors[name] - buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True) - self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, buf_x.data_ptr()) - desc_x.copy_(buf_x, non_blocking=True) - - # Call this method inside the lambda function for grid size - def fill_2d_tma_descriptor(self, name, ptr, dim1, dim0, block_dim1, block_dim0, element_size): - if HAS_TMA_DESC: - desc_x = self.descriptors[name] - assert desc_x.data_ptr() % 64 == 0 - self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, desc_x.data_ptr()) - else: - desc_x = self.cuda_descriptors[name] - buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True) - self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.data_ptr()) - desc_x.copy_(buf_x, non_blocking=True) - - def get_tma_descriptor_kernel_param(self, name): - if HAS_TMA_DESC: - assert self.descriptors[name] is not None - return self.KernelParamWrapper(self.descriptors[name]) - else: - assert self.cuda_descriptors[name] is not None - return self.cuda_descriptors[name] - - -@triton.jit -def create_mask_from_indices_kernel( - block_indices, - block_mask, - stride_bz, - stride_bh, - stride_bm, - stride_bs, - stride_mz, - stride_mh, - stride_mm, - stride_mn, - H, -): - i_zh, i_m, i_s = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_z, i_h = i_zh // H, i_zh % H - - off_b = i_z.to(tl.int64) * stride_bz + i_h.to(tl.int64) * stride_bh + i_m.to(tl.int64) * stride_bm + i_s.to( - tl.int64) * stride_bs - - b_i = tl.load(block_indices + off_b) - - off_m = i_z.to(tl.int64) * stride_mz + i_h.to(tl.int64) * stride_mh + i_m.to(tl.int64) * stride_mm + b_i.to( - tl.int64) * stride_mn - - b_m = 1 - tl.store(block_mask + off_m, b_m.to(block_mask.dtype.element_ty)) - - -def create_mask_from_indices_triton(block_indices, N_cols): - B, H, N_rows, S = block_indices.shape - block_mask = torch.zeros((B, H, N_rows, N_cols), dtype=torch.bool, device=block_indices.device) - create_mask_from_indices_kernel[(B * H, N_rows, S)]( - block_indices, - block_mask, - block_indices.stride(0), - block_indices.stride(1), - block_indices.stride(2), - block_indices.stride(3), - block_mask.stride(0), - block_mask.stride(1), - block_mask.stride(2), - block_mask.stride(3), - H, - ) - return block_mask - - -@torch.compile -def create_mask_from_indices_varlen(block_indices, N_cols_mask): - - B, H, M, _ = block_indices.shape - device = block_indices.device - - mask = torch.zeros((B, H, M, N_cols_mask), dtype=torch.bool, device=device) - - valid = block_indices < N_cols_mask - - b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(block_indices) - h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(block_indices) - m_idx = torch.arange(M, device=device)[None, None, :, None].expand_as(block_indices) - - valid_coords = (b_idx[valid], h_idx[valid], m_idx[valid], block_indices[valid]) - - mask[valid_coords] = True - - return mask - - -@torch.compile -def create_indices_k_from_indices_q_varlen( - block_indices, - N_cols_mask # indicate the number of the last dimension of the bool mask, since this information cannot be determined by block_indices, which may contain invalid elements -): - block_mask_qk = create_mask_from_indices_varlen(block_indices, N_cols_mask) - B, H, M, N = block_mask_qk.shape - block_mask_kq = block_mask_qk.permute(0, 1, 3, 2) - indices = torch.arange(M, device=block_indices.device).view(1, 1, 1, -1).expand_as(block_mask_kq) - block_indices_k = torch.where(block_mask_kq, indices, M) - block_indices_k, _ = torch.sort(block_indices_k, dim=-1) - - block_indices_k_lens = (block_indices_k < M).sum(dim=-1) - - return block_indices_k, block_indices_k_lens - - -@torch.compile -def mean_pooling_compression(x: torch.Tensor, block_size: int) -> torch.Tensor: - B, H, S = x.shape[:3] - num_block = math.ceil(S / block_size) - if S % block_size != 0: - x = F.pad(x, (0, 0, 0, num_block * block_size - S)) - x_cmp = x.view(B, H, num_block, block_size, -1).mean(dim=3) - return x_cmp - - -@torch.compile -def cal_score(q, k): - k_transposed = k.transpose(-1, -2) # [b, h, d, s_k] - score = torch.matmul(q, k_transposed) # [b, h, s_q, s_k] - return score - - -def cal_score_triton(q, k): - B, H, s_q, D = q.shape - s_k = k.shape[2] - - score = torch.empty(B, H, s_q, s_k, device=q.device, dtype=q.dtype) - - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_gating_preset['default'] - - grid = lambda args: (triton.cdiv(s_q, args["BLOCK_M"]), B * H, 1) - _attn_fwd_gating[grid](q, - k, - score, - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - k.stride(0), - k.stride(1), - k.stride(2), - k.stride(3), - score.stride(0), - score.stride(1), - score.stride(2), - score.stride(3), - H, - s_q, - s_k, - HEAD_DIM=D, - **kernel_config) - return score - - -@torch.compile -def get_select_indices_topk(q, k, sparsity): - score = cal_score(q, k) - block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity) - return block_indices, block_indices_lens - - -@torch.compile -def get_select_indices_topk_from_score(score, sparsity): - num_selected = int((1 - sparsity) * score.shape[-1]) - block_indices = torch.topk(score, num_selected)[1] - - block_indices_lens = torch.full((block_indices.shape[0], block_indices.shape[1], block_indices.shape[2]), - num_selected, - dtype=torch.int32, - device=block_indices.device) - - return block_indices, block_indices_lens - - -@torch.compile -def get_select_indices_cdf(q, k, cdf_threshold): - score = cal_score(q, k) - head_dim = q.shape[-1] - block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold, 1 / head_dim**0.5) - return block_indices, block_indices_lens - - -@torch.compile -def get_select_indices_cdf_from_score(score, cdf_threshold, sm_scale): - weights = torch.softmax(score * sm_scale, dim=-1) - - B, H, Sq, Sk = weights.shape - cdf_threshold = torch.full((H, ), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1) - weights_sorted = torch.sort(weights, dim=-1, descending=True) - cdf = torch.cumsum(weights_sorted.values, dim=-1) - num_selected = torch.searchsorted(cdf, cdf_threshold, right=True) - - return weights_sorted.indices, num_selected.squeeze(-1) - - -@torch.compile -def get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold): - score = cal_score(q, k) - head_dim = q.shape[-1] - block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, - 1 / head_dim**0.5) - return block_indices, block_indices_lens - - -@torch.compile -def get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_scale): - weights = torch.softmax(score * sm_scale, dim=-1) - - B, H, Sq, Sk = weights.shape - cdf_threshold = torch.full((H, ), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1) - weights_sorted = torch.sort(weights, dim=-1, descending=True) - cdf = torch.cumsum(weights_sorted.values, dim=-1) - num_selected = torch.searchsorted(cdf, cdf_threshold, right=True) - - # max(cdf, topk) - num_selected_topk = int((1 - sparsity) * score.shape[-1]) - num_selected[num_selected < num_selected_topk] = num_selected_topk - - return weights_sorted.indices, num_selected.squeeze(-1) - - -def get_select_indices(q, k, sparsity, cdf_threshold): - if sparsity is not None and cdf_threshold is None: - block_indices, block_indices_lens = get_select_indices_topk(q, k, sparsity) - elif sparsity is None and cdf_threshold is not None: - block_indices, block_indices_lens = get_select_indices_cdf(q, k, cdf_threshold) - elif sparsity is not None and cdf_threshold is not None: - block_indices, block_indices_lens = get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold) - else: - raise ValueError - return block_indices, block_indices_lens - - -def get_select_indices_from_score(score, sparsity, cdf_threshold): - if sparsity is not None and cdf_threshold is None: - block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity) - elif sparsity is None and cdf_threshold is not None: - block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold) - elif sparsity is not None and cdf_threshold is not None: - block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold) - else: - raise ValueError - return block_indices, block_indices_lens - - -def attn_fwd_bsa_varlen_triton(q, k, v, sm_scale, block_indices, block_indices_lens, chunk_size_q, chunk_size_k, - sparsity): - - B, H, Seq, D = q.shape - - o = torch.empty_like(q) - M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) - - grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1) - - config_key = 'BLOCK_N_LG=64' if chunk_size_k == 64 else 'default' - if chunk_size_k > 128: - fwd_func = _attn_fwd_bsa_varlen - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', - '0') == '1' else configs_fwd_bsa_varlen_preset[config_key] - else: - fwd_func = _attn_fwd_bsa_varlen_align - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', - '0') == '1' else configs_fwd_bsa_varlen_align_preset[config_key] - - block_indices = block_indices.contiguous() - block_indices_lens = block_indices_lens.contiguous() - - fwd_func[grid]( - q, - k, - v, - sm_scale, - M, - o, - block_indices, # [B, H, M_COMPRESS, S] - block_indices_lens, # [B, H, M_COMPRESS, S_MAX] - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - k.stride(0), - k.stride(1), - k.stride(2), - k.stride(3), - v.stride(0), - v.stride(1), - v.stride(2), - v.stride(3), - o.stride(0), - o.stride(1), - o.stride(2), - o.stride(3), - block_indices.stride(0), - block_indices.stride(1), - block_indices.stride(2), - block_indices.stride(3), - block_indices_lens.stride(0), - block_indices_lens.stride(1), - block_indices_lens.stride(2), - H, - Seq, - D, - BLOCK_M=chunk_size_q, - BLOCK_N_LG=chunk_size_k, - SPARSITY=sparsity, - **kernel_config) - - LN2 = 0.6931471824645996 - lse = M * LN2 # convert back to natural units (M is of base 2) - - return o, lse - - -def attn_bwd_bsa_varlen_triton(do, q, k, v, o, dq, dk, dv, sm_scale, M, block_indices, block_indices_lens, chunk_size_q, - chunk_size_k, sparsity): - RCP_LN2 = 1.4426950408889634 - M = M * RCP_LN2 # ln -> log2 - - do = do.contiguous() - # assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() - - BATCH, N_HEAD, N_CTX, HEAD_DIM = q.shape - N_CTX_KV = k.shape[-2] - - RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) # reciprocal - arg_k = k - arg_k = arg_k * (sm_scale * RCP_LN2) - - PRE_BLOCK = 128 if min(chunk_size_q, chunk_size_k) >= 128 else min(chunk_size_q, chunk_size_k) - - assert N_CTX % PRE_BLOCK == 0 - pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) - delta = torch.empty_like(M) - _attn_bwd_preprocess[pre_grid](o, do, delta, N_CTX, BLOCK_M=PRE_BLOCK, HEAD_DIM=HEAD_DIM) - - block_indices_k, block_indices_k_lens = create_indices_k_from_indices_q_varlen(block_indices=block_indices, - N_cols_mask=N_CTX_KV // chunk_size_k) - - block_indices = block_indices.contiguous() - block_indices_lens = block_indices_lens.contiguous() - block_indices_k = block_indices_k.contiguous() - block_indices_k_lens = block_indices_k_lens.contiguous() - - config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default' - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', - '0') == '1' else configs_bwd_dkdv_bsa_varlen_preset[config_key] - - grid_dkdv = lambda args: (triton.cdiv(arg_k.shape[2], args["BLOCK_N"]), 1, arg_k.shape[0] * arg_k.shape[1]) - _attn_bwd_dkdv_bsa_varlen_wrapper[grid_dkdv]( - q, - arg_k, - v, - sm_scale, # softmax scale - do, - dk, - dv, - M, # lse (log2) - delta, - block_indices_k, - block_indices_k_lens, - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - k.stride(0), - k.stride(1), - k.stride(2), - k.stride(3), - v.stride(0), - v.stride(1), - v.stride(2), - v.stride(3), - dk.stride(0), - dk.stride(1), - dk.stride(2), - dk.stride(3), - dv.stride(0), - dv.stride(1), - dv.stride(2), - dv.stride(3), - do.stride(0), - do.stride(1), - do.stride(2), - do.stride(3), - M.stride(0), - M.stride(1), - M.stride(2), - delta.stride(0), - delta.stride(1), - delta.stride(2), - block_indices_k.stride(0), - block_indices_k.stride(1), - block_indices_k.stride(2), - block_indices_k.stride(3), - block_indices_k_lens.stride(0), - block_indices_k_lens.stride(1), - block_indices_k_lens.stride(2), - N_HEAD, - N_CTX, - BLOCK_M=chunk_size_q, - BLOCK_N_DQ_LG=chunk_size_k, - HEAD_DIM=HEAD_DIM, - SPARSITY=sparsity, - **kernel_config) - - config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default' - if chunk_size_k > 128: - bwd_dq_func = _attn_bwd_dq_bsa_varlen_wrapper - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', - '0') == '1' else configs_bwd_dq_bsa_varlen_preset[config_key] - else: - bwd_dq_func = _attn_bwd_dq_bsa_varlen_align_wrapper - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', - '0') == '1' else configs_bwd_dq_bsa_varlen_align_preset[config_key] - - grid_dq = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), 1, q.shape[0] * q.shape[1]) - bwd_dq_func[grid_dq]( - q, - arg_k, - v, - do, - dq, - M, # lse (log2) - delta, - block_indices, - block_indices_lens, - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - k.stride(0), - k.stride(1), - k.stride(2), - k.stride(3), - v.stride(0), - v.stride(1), - v.stride(2), - v.stride(3), - dq.stride(0), - dq.stride(1), - dq.stride(2), - dq.stride(3), - do.stride(0), - do.stride(1), - do.stride(2), - do.stride(3), - M.stride(0), - M.stride(1), - M.stride(2), - delta.stride(0), - delta.stride(1), - delta.stride(2), - block_indices.stride(0), - block_indices.stride(1), - block_indices.stride(2), - block_indices.stride(3), - block_indices_lens.stride(0), - block_indices_lens.stride(1), - block_indices_lens.stride(2), - N_HEAD, - N_CTX, - BLOCK_M=chunk_size_q, - BLOCK_N_DQ_LG=chunk_size_k, - HEAD_DIM=HEAD_DIM, - SPARSITY=sparsity, - **kernel_config) - - -@torch.compile -def make_block_indices_varlen_cp_list(block_indices, cp_size, num_blocks_k_full): - """ - Args: - block_indices: [B, H, num_blocks_q_per_cp_rank, num_blocks_k_full] - - Return: - a list of [block_indices, block_indices_lens] for k from each cp_rank - - each block_indices starts from zero - - block_indices_lens indicates the valid number of elements in the last dimension of block_indices - """ - res = [] - num_blocks_per_rank = num_blocks_k_full // cp_size - for i in range(cp_size): - block_indices_tmp = block_indices.clone() - min_block_idx = i * num_blocks_per_rank - block_indices_tmp -= min_block_idx - block_indices_tmp[ - block_indices_tmp < - 0] = num_blocks_per_rank # block_indices_tmp < 0 indicate invalid indices, set them to num_blocks_per_rank in order to sort them to the tail, so that the first N elements of the block_indices indicated by block_indices_lens are valid - - block_indices_tmp, _ = torch.sort(block_indices_tmp, dim=-1) - - block_indices_tmp_lens = (block_indices_tmp < num_blocks_per_rank).sum(dim=-1) - - res.append([block_indices_tmp, block_indices_tmp_lens]) - - return res - - -@torch.compile -def flash_attn_fwd_softmax_lse_correction( - softmax_lse: torch.Tensor, - softmax_lse_per_step: torch.Tensor, -): - """Merge softmax stats of each step in Attention with context parallelism""" - max_scale = torch.max(softmax_lse, softmax_lse_per_step) - min_scale = torch.min(softmax_lse, softmax_lse_per_step) - lse_diff = min_scale - max_scale - lse_diff = lse_diff.nan_to_num( - nan=0. - ) # handle cases: tensor(-inf) - tensor(-inf) = tensor(nan); In the current cp implementation, it is possible that lses of 2 cp ranks are both -inf, if no block is selected from both cp ranks. In such cases, the finally corrected lse should remain -inf. - new_scale = max_scale + torch.log1p( - torch.exp(lse_diff)) # a + ln(1 + e^(b - a)) = ln(e^a) + ln(1 + e^(b - a)) = ln(e^a + e^b) - softmax_lse.copy_(new_scale) - - -@torch.compile -def flash_attn_fwd_out_correction_init( - out_init_step: torch.Tensor, # b h s d - softmax_lse: torch.Tensor, # b h s - softmax_lse_init_step: torch.Tensor, -): - """Merge partial outputs of the first step in Attention with context parallelism""" - softmax_lse_corrected_exp = torch.exp(softmax_lse_init_step - softmax_lse) - softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1) - out_corrected = out_init_step * softmax_lse_corrected_exp - return out_corrected.to(out_init_step.dtype) - - -@torch.compile -def flash_attn_fwd_out_correction( - out: torch.Tensor, - out_per_step: torch.Tensor, - softmax_lse: torch.Tensor, - softmax_lse_per_step: torch.Tensor, -): - """Merge partial outputs of each step in Attention with context parallelism""" - softmax_lse_corrected_exp = torch.exp(softmax_lse_per_step - softmax_lse) - softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1) - out_corrected = out_per_step * softmax_lse_corrected_exp - out.add_(out_corrected) - - -@torch.compile -def topk_sort(score, num_chunks_selected): - block_indices = torch.topk(score, num_chunks_selected)[1] - block_indices, _ = torch.sort(block_indices, dim=-1) - return block_indices - - -class _attention_bsa(torch.autograd.Function): - - @staticmethod - def forward(ctx, q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, sm_scale, use_tma=False): - # shape constraints - HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] - # when v is in float8_e5m2 it is transposed. - HEAD_DIM_V = v.shape[-1] - assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V - assert HEAD_DIM_K in {16, 32, 64, 128, 256} - - # ---------------------- gating ---------------------- - q_cmp = mean_pooling_compression(q, chunk_size_q) - k_cmp = mean_pooling_compression(k, chunk_size_k) - block_indices, block_indices_lens = get_select_indices(q_cmp, k_cmp, sparsity, cdf_threshold) - - # ---------------------- bsa ---------------------- - - o, lse = attn_fwd_bsa_varlen_triton(q, k, v, sm_scale, block_indices, block_indices_lens, chunk_size_q, - chunk_size_k, sparsity) - - ctx.save_for_backward(q, k, v, o, lse, block_indices, block_indices_lens) - ctx.sm_scale = sm_scale - ctx.HEAD_DIM = HEAD_DIM_K - ctx.chunk_size_q = chunk_size_q - ctx.chunk_size_k = chunk_size_k - ctx.use_tma = use_tma - ctx.sparsity = sparsity - - return o - - @staticmethod - def backward(ctx, do): - q, k, v, o, lse, block_indices, block_indices_lens = ctx.saved_tensors - - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - - attn_bwd_bsa_varlen_triton(do, q, k, v, o, dq, dk, dv, ctx.sm_scale, lse, block_indices, block_indices_lens, - ctx.chunk_size_q, ctx.chunk_size_k, ctx.sparsity) - - return dq, dk, dv, None, None, None, None, None, None - - -flash_attn_bsa = _attention_bsa.apply - - -def rearrange_THW_to_3d_block(x, Nt, Nh, Nw, t, h, w, D): - B, H, _, D = x.shape - x = x.view(B, H, Nt, t, Nh, h, Nw, w, D) - x = x.permute(0, 1, 2, 4, 6, 3, 5, 7, 8) # B H Nt Nh Nw t h w D - return x.contiguous().view(B, H, Nt * Nh * Nw * t * h * w, D) - - -def rearrange_3d_block_to_THW(x, Nt, Nh, Nw, t, h, w, D): - B, H, _, D = x.shape - x = x.view(B, H, Nt, Nh, Nw, t, h, w, D) - x = x.permute(0, 1, 2, 5, 3, 6, 4, 7, 8) # B H Nt t Nh h Nw w D - return x.contiguous().view(B, H, Nt * t * Nh * h * Nw * w, D) - - -def flash_attn_bsa_3d( - q: torch.Tensor, # [B, H, Sq, D] - k: torch.Tensor, # [B, H, Skv, D] - v: torch.Tensor, # [B, H, Skv, D] - latent_shape_q, - latent_shape_k, - # bsa_params - sparsity=0.875, - cdf_threshold=None, - chunk_3d_shape_q=[4, 4, 8], - chunk_3d_shape_k=[4, 4, 8], -) -> torch.Tensor: - _, _, Sq, head_dim_q = q.shape - _, _, Sk, head_dim_k = k.shape - - assert head_dim_q == head_dim_k - head_dim = head_dim_q - - Tq, Hq, Wq = latent_shape_q - Tk, Hk, Wk = latent_shape_k - - assert Tq * Hq * Wq == Sq - assert Tk * Hk * Wk == Sk - - tq, hq, wq = chunk_3d_shape_q - tk, hk, wk = chunk_3d_shape_k - - assert Tq % tq == 0 and Hq % hq == 0 and Wq % wq == 0 - assert Tk % tk == 0 and Hk % hk == 0 and Wk % wk == 0 - - Ntq = Tq // tq - Nhq = Hq // hq - Nwq = Wq // wq - - Ntk = Tk // tk - Nhk = Hk // hk - Nwk = Wk // wk - - q = rearrange_THW_to_3d_block(q, Ntq, Nhq, Nwq, tq, hq, wq, q.shape[-1]) - k = rearrange_THW_to_3d_block(k, Ntk, Nhk, Nwk, tk, hk, wk, k.shape[-1]) - v = rearrange_THW_to_3d_block(v, Ntk, Nhk, Nwk, tk, hk, wk, v.shape[-1]) - - chunk_size_q = tq * hq * wq - chunk_size_k = tk * hk * wk - - output = flash_attn_bsa(q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, 1 / head_dim**0.5) - - output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1]) - return output diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/common.py b/fastvideo/third_party/longcat_video/block_sparse_attention/common.py deleted file mode 100644 index 09502b67f..000000000 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/common.py +++ /dev/null @@ -1,130 +0,0 @@ -import triton -import triton.language as tl -import os - -if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1': - autotune = triton.autotune -else: - - def autotune(*args, **kwargs): - - def decorator(func): - return func - - return decorator - - -configs_gating_preset = { - 'default': { - 'BLOCK_M': 64, - 'BLOCK_N': 64, - 'num_stages': 3, - 'num_warps': 8, - } -} - -configs_gating = [ - triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \ - for BM in [64, 128] \ - for BN in [32, 64] \ - for s in [2, 3, 4, 5] \ - for w in [4, 8] \ -] - -gating_reevaluate_keys = ["M", "N"] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] - - -@autotune(configs_gating, key=gating_reevaluate_keys) -@triton.jit -def _attn_fwd_gating( - Q, - K, - Out, - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_oz, - stride_oh, - stride_om, - stride_on, - H, - M, - N, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - - tl.static_assert(BLOCK_N <= HEAD_DIM) - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh - o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh - - # block pointers - Q_block_ptr = tl.make_block_ptr( - base=Q + q_offset, - shape=(M, HEAD_DIM), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - - K_block_ptr = tl.make_block_ptr( - base=K + k_offset, - shape=(HEAD_DIM, N), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N), - order=(0, 1), - ) - O_block_ptr = tl.make_block_ptr( - base=Out + o_offset, - shape=(M, N), - strides=(stride_om, stride_on), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_N), - order=(1, 0), - ) - - # load q: it will stay in SRAM throughout - q = tl.load(Q_block_ptr, boundary_check=(0, )) - for start_n in range(0, N, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - k = tl.load(K_block_ptr, boundary_check=(1, )) - qk = tl.dot(q, k) - - tl.store(O_block_ptr, qk.to(Out.type.element_ty), boundary_check=(0, 1)) - - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) - O_block_ptr = tl.advance(O_block_ptr, (0, BLOCK_N)) - - -@triton.jit -def _attn_bwd_preprocess( - O, - DO, - Delta, # output - N_CTX, - BLOCK_M: tl.constexpr, - HEAD_DIM: tl.constexpr): - - off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) - off_hz = tl.program_id(1) - off_n = tl.arange(0, HEAD_DIM) - # load - o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]) - do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32) - delta = tl.sum(o * do, axis=1) - # write-back - tl.store(Delta + off_hz * N_CTX + off_m, delta) diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py b/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py deleted file mode 100644 index afbab5e1c..000000000 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - - -def p2p_communicate(rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm): - """Point-to-point communications of KV and dKV in Attention with context parallelism""" - send_recv_ops = [] - - if batch_p2p_comm: # int(os.getenv("NVTE_BATCH_MHA_P2P_COMM", "0")) or (cp_size == 2) 为啥呢 - if rank % 2 == 0: - send_op = torch.distributed.P2POp(torch.distributed.isend, send_tensor, send_dst, cp_group) - recv_op = torch.distributed.P2POp(torch.distributed.irecv, recv_tensor, recv_src, cp_group) - send_recv_ops.append(send_op) - send_recv_ops.append(recv_op) - else: - recv_op = torch.distributed.P2POp(torch.distributed.irecv, recv_tensor, recv_src, cp_group) - send_op = torch.distributed.P2POp(torch.distributed.isend, send_tensor, send_dst, cp_group) - send_recv_ops.append(recv_op) - send_recv_ops.append(send_op) - send_recv_reqs = torch.distributed.batch_isend_irecv(send_recv_ops) - else: - if rank % 2 == 0: - send_op = torch.distributed.isend(send_tensor, send_dst, cp_group) - recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group) - send_recv_ops.append(send_op) - send_recv_ops.append(recv_op) - else: - recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group) - send_op = torch.distributed.isend(send_tensor, send_dst, cp_group) - send_recv_ops.append(recv_op) - send_recv_ops.append(send_op) - send_recv_reqs = send_recv_ops - - return send_recv_reqs diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py b/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py deleted file mode 100644 index ea52e096f..000000000 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py +++ /dev/null @@ -1,1130 +0,0 @@ -import triton -import triton.language as tl -import os - -from .common import autotune -""" -TRITON_REEVALUATE_KEY=1 -- autotune whenever params in reevaluate keys change -- use in benchmark script to fine the best config - -TRITON_AUTOTUNE_ENBALE=1 -- if set to 0, autotune will not work, and the related params must be passed to the function call. -""" - -configs_fwd_bsa_varlen_preset = { - 'default': { - 'BLOCK_N': 64, - 'num_stages': 3, - 'num_warps': 8, - }, - 'BLOCK_N_LG=64': { - 'BLOCK_N': 64, - 'num_stages': 3, - 'num_warps': 4, - }, -} -configs_fwd_bsa_varlen = [ - triton.Config({'BLOCK_N': BN}, num_stages=s, num_warps=w) \ - for BN in [32, 64, 128] \ - for s in [2, 3, 4, 5] \ - for w in [4, 8] \ -] - -fwd_bsa_reevaluate_varlen_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get( - 'TRITON_REEVALUATE_KEY', '0') == '1' else [] - - -@autotune(list(configs_fwd_bsa_varlen), key=fwd_bsa_reevaluate_varlen_keys) -@triton.jit -def _attn_fwd_bsa_varlen( - Q, - K, - V, - sm_scale, - M, - Out, - block_indices, # [B, H, M_COMPRESS, S_MAX] - block_indices_lens, # [B, H, M_COMPRESS] - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - stride_oz, - stride_oh, - stride_om, - stride_ok, - stride_bz, - stride_bh, - stride_bm, - stride_bs, - stride_lz, - stride_lh, - stride_lm, - H, - N_CTX, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N_LG: tl.constexpr, - BLOCK_N: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking -): - - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - - q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh - v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh - o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh - b_offset = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh - l_offset = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh - - # block pointers - Q_block_ptr = tl.make_block_ptr( - base=Q + q_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - V_block_ptr = tl.make_block_ptr( - base=V + v_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_vn, stride_vk), - offsets=(0, 0), - block_shape=(BLOCK_N, HEAD_DIM), - order=(1, 0), - ) - KT_block_ptr = tl.make_block_ptr( - base=K + k_offset, - shape=(HEAD_DIM, N_CTX), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N), - order=(0, 1), - ) - O_block_ptr = tl.make_block_ptr( - base=Out + o_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_om, stride_ok), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - block_indices += b_offset + start_m * stride_bm - block_indices_lens += l_offset + start_m * stride_lm - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 - acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - # load scales - qk_scale = sm_scale - qk_scale *= 1.44269504 # 1/ln2; exp2(x/ln2) == exp2(ln(e^x) / ln2) == exp2(log2(e^x)) == exp(x) - # load q: it will stay in SRAM throughout - q = tl.load(Q_block_ptr) - S = tl.load(block_indices_lens) - for i in range(S): - block_id = tl.load(block_indices + i * stride_bs).to(tl.int32) - lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG - lo = tl.multiple_of(lo, BLOCK_N) - KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo)) - V_block_ptr_i = tl.advance(V_block_ptr, (lo, 0)) - - # loop over k, v and update accumulator - for start_n in range(lo, hi, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - kT = tl.load(KT_block_ptr_i) - qkT = tl.dot(q, kT) - - m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale) - qkT = qkT * qk_scale - m_ij[:, None] - p = tl.math.exp2(qkT) - - # -- update m_i and l_i - alpha = tl.math.exp2(m_i - m_ij) - l_ij = tl.sum(p, 1) - # -- update output accumulator -- - acc = acc * alpha[:, None] - # update acc - v = tl.load(V_block_ptr_i) - acc = tl.dot(p.to(v.dtype), v, acc) - # update m_i and l_i - # place this at the end of the loop to reduce register pressure: https://github.com/triton-lang/triton/commit/ee6abd9 - l_i = l_i * alpha + l_ij - m_i = m_ij - V_block_ptr_i = tl.advance(V_block_ptr_i, (BLOCK_N, 0)) - KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N)) - - # epilogue - m_i += tl.math.log2(l_i) - acc = acc / l_i[:, None] - m_ptrs = M + off_hz * N_CTX + offs_m - tl.store(m_ptrs, m_i) - tl.store(O_block_ptr, acc.to(Out.type.element_ty)) - - -configs_fwd_bsa_varlen_align_preset = { - 'default': { - 'num_stages': 3, - 'num_warps': 8, - }, - 'BLOCK_N_LG=64': { - 'num_stages': 3, - 'num_warps': 4, - }, -} -configs_fwd_bsa_varlen_align = [ - triton.Config({}, num_stages=s, num_warps=w) \ - for s in [2, 3, 4, 5] \ - for w in [4, 8] \ -] - -fwd_bsa_reevaluate_varlen_align_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get( - 'TRITON_REEVALUATE_KEY', '0') == '1' else [] - - -@autotune(list(configs_fwd_bsa_varlen_align), key=fwd_bsa_reevaluate_varlen_align_keys) -@triton.jit -def _attn_fwd_bsa_varlen_align( - Q, - K, - V, - sm_scale, - M, - Out, - block_indices, # [B, H, M_COMPRESS, S_MAX] - block_indices_lens, # [B, H, M_COMPRESS] - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - stride_oz, - stride_oh, - stride_om, - stride_on, - stride_bz, - stride_bh, - stride_bm, - stride_bs, - stride_lz, - stride_lh, - stride_lm, - H, - N_CTX, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N_LG: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking -): - - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - - q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh - v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh - o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh - b_offset = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh - l_offset = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh - - # block pointers - Q_block_ptr = tl.make_block_ptr( - base=Q + q_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - V_block_ptr = tl.make_block_ptr( - base=V + v_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_vn, stride_vk), - offsets=(0, 0), - block_shape=(BLOCK_N_LG, HEAD_DIM), - order=(1, 0), - ) - KT_block_ptr = tl.make_block_ptr( - base=K + k_offset, - shape=(HEAD_DIM, N_CTX), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N_LG), - order=(0, 1), - ) - O_block_ptr = tl.make_block_ptr( - base=Out + o_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_om, stride_on), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - block_indices += b_offset + start_m * stride_bm - block_indices_lens += l_offset + start_m * stride_lm - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 - acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - # load scales - qk_scale = sm_scale - qk_scale *= 1.44269504 # 1/ln2; exp2(x/ln2) == exp2(ln(e^x) / ln2) == exp2(log2(e^x)) == exp(x)；乘1/ln2后，exp2(x/ln2) == exp(x)，exp2速度更快 - # load q: it will stay in SRAM throughout - q = tl.load(Q_block_ptr) - S = tl.load(block_indices_lens) - for i in range(S): - block_id = tl.load(block_indices + i * stride_bs).to(tl.int32) - lo = block_id * BLOCK_N_LG - lo = tl.multiple_of(lo, BLOCK_N_LG) - KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo)) - V_block_ptr_i = tl.advance(V_block_ptr, (lo, 0)) - - # -- compute qk ---- - kT = tl.load(KT_block_ptr_i) - qkT = tl.dot(q, kT) - - m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale) - qkT = qkT * qk_scale - m_ij[:, None] - p = tl.math.exp2(qkT) - - # -- update m_i and l_i - alpha = tl.math.exp2(m_i - m_ij) - l_ij = tl.sum(p, 1) - # -- update output accumulator -- - acc = acc * alpha[:, None] - # update acc - v = tl.load(V_block_ptr_i) - acc = tl.dot(p.to(v.dtype), v, acc) # 没除se，fa2引入的优化 - # update m_i and l_i - # place this at the end of the loop to reduce register pressure: https://github.com/triton-lang/triton/commit/ee6abd9 - l_i = l_i * alpha + l_ij # 当前总se - m_i = m_ij - - # epilogue - m_i += tl.math.log2(l_i) - acc = acc / l_i[:, None] - m_ptrs = M + off_hz * N_CTX + offs_m - tl.store(m_ptrs, m_i) - tl.store(O_block_ptr, acc.to(Out.type.element_ty)) - - -# The main inner-loop logic for computing dK and dV. -@triton.jit -def _attn_bwd_dkdv_bsa_varlen( - dk, - dv, - k, - v, - Q, - DO, - M, - D, - block_indices, - block_indices_lens, - # shared by Q/K/V/DO. - # stride_tok, stride_d, - stride_qm, - stride_qk, - stride_dom, - stride_dok, - stride_mm, - stride_dm, - stride_bm, - N_CTX, - BLOCK_M1: tl.constexpr, - HEAD_DIM: tl.constexpr, -): - - QT_block_ptr = tl.make_block_ptr( - base=Q, - shape=(HEAD_DIM, N_CTX), - strides=(stride_qk, stride_qm), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_M1), - order=(0, 1), - ) - - DO_block_ptr = tl.make_block_ptr( - base=DO, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dom, stride_dok), - offsets=(0, 0), - block_shape=(BLOCK_M1, HEAD_DIM), - order=(1, 0), - ) - - S = tl.load(block_indices_lens) - for i in range(S): - block_id = tl.load(block_indices + i * stride_bm).to(tl.int32) - start_m = block_id * BLOCK_M1 - start_m = tl.multiple_of(start_m, BLOCK_M1) - - QT_block_ptr_i = tl.advance(QT_block_ptr, (0, start_m)) - DO_block_ptr_i = tl.advance(DO_block_ptr, (start_m, 0)) - - qT = tl.load(QT_block_ptr_i) - # Load m before computing qk to reduce pipeline stall. - offs_m = start_m + tl.arange(0, BLOCK_M1) * stride_mm - m = tl.load(M + offs_m) - kqT = tl.dot(k, qT) - pT = tl.math.exp2(kqT - m[None, :]) - - do = tl.load(DO_block_ptr_i) - # Compute dV. - ppT = pT - ppT = ppT.to(v.dtype) - dv += tl.dot(ppT, do) - # D (= delta) is pre-divided by ds_scale. - offs_d = start_m + tl.arange(0, BLOCK_M1) * stride_dm - Di = tl.load(D + offs_d) - # Compute dP and dS. - dpT = tl.dot(v, tl.trans(do)).to(tl.float32) - dsT = pT * (dpT - Di[None, :]) - dsT = dsT.to(v.dtype) - dk += tl.dot(dsT, tl.trans(qT)) - - return dk, dv - - -# the main inner-loop logic for computing dQ -@triton.jit -def _attn_bwd_dq_bsa_varlen( - dq, - q, - do, - m, - d, - K, - V, - N_CTX, - BLOCK_N2: tl.constexpr, - BLOCK_N_LG: tl.constexpr, - HEAD_DIM: tl.constexpr, - block_indices, - block_indices_lens, - stride_bn, - # stride_tok, stride_d, - stride_kn, - stride_kk, - stride_vn, - stride_vk, -): - - VT_block_ptr = tl.make_block_ptr( - base=V, - shape=(HEAD_DIM, N_CTX), - strides=(stride_vk, stride_vn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N2), - order=(0, 1), - ) - KT_block_ptr = tl.make_block_ptr( - base=K, - shape=(HEAD_DIM, N_CTX), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N2), - order=(0, 1), - ) - - S = tl.load(block_indices_lens) - for i in range(S): - block_id = tl.load(block_indices + i * stride_bn).to(tl.int32) - lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG - lo = tl.multiple_of(lo, BLOCK_N2) - - KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo)) - VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo)) - - for start_n in range(lo, hi, BLOCK_N2): - start_n = tl.multiple_of(start_n, BLOCK_N2) - - kT = tl.load(KT_block_ptr_i) - vT = tl.load(VT_block_ptr_i) - qkT = tl.dot(q, kT) - p = tl.math.exp2(qkT - m) - # Compute dP and dS. - dp = tl.dot(do, vT).to(tl.float32) - ds = p * (dp - d) - ds = ds.to( - kT.dtype - ) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128 - # Compute dQ. - # NOTE: We need to de-scale dq in the end, because kT was pre-scaled. - dq += tl.dot(ds, tl.trans(kT)) - # Increment pointers. - KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N2)) - VT_block_ptr_i = tl.advance(VT_block_ptr_i, (0, BLOCK_N2)) - - return dq - - -# the main inner-loop logic for computing dQ -@triton.jit -def _attn_bwd_dq_bsa_varlen_align( - dq, - q, - do, - m, - d, - K, - V, - N_CTX, - BLOCK_N_LG: tl.constexpr, - HEAD_DIM: tl.constexpr, - block_indices, - block_indices_lens, - stride_bn, - stride_kn, - stride_kk, - stride_vn, - stride_vk, -): - - VT_block_ptr = tl.make_block_ptr( - base=V, - shape=(HEAD_DIM, N_CTX), - strides=(stride_vk, stride_vn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N_LG), - order=(0, 1), - ) - KT_block_ptr = tl.make_block_ptr( - base=K, - shape=(HEAD_DIM, N_CTX), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N_LG), - order=(0, 1), - ) - - S = tl.load(block_indices_lens) - for i in range(S): - block_id = tl.load(block_indices + i * stride_bn).to(tl.int32) - lo = block_id * BLOCK_N_LG - lo = tl.multiple_of(lo, BLOCK_N_LG) - - KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo)) - VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo)) - - kT = tl.load(KT_block_ptr_i) - vT = tl.load(VT_block_ptr_i) - qkT = tl.dot(q, kT) - p = tl.math.exp2(qkT - m) - # Compute dP and dS. - dp = tl.dot(do, vT).to(tl.float32) - ds = p * (dp - d) - ds = ds.to( - kT.dtype - ) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128 - # Compute dQ. - # NOTE: We need to de-scale dq in the end, because kT was pre-scaled. - dq += tl.dot(ds, tl.trans(kT)) - - return dq - - -configs_bwd_dkdv_bsa_varlen_preset = { - 'default': { - 'BLOCK_N': 128, - 'num_stages': 2, - 'num_warps': 8, - }, - 'BLOCK_N_DQ_LG=64': { - 'BLOCK_N': 64, - 'num_stages': 2, - 'num_warps': 4, - } -} -configs_bwd_dkdv_bsa_varlen = [ - triton.Config({'BLOCK_N': BN}, num_stages=s, num_warps=w) \ - for BN in [32, 64, 128] \ - for s in [2, 3, 4, 5] \ - for w in [4, 8] \ -] -bwd_dkdv_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get( - 'TRITON_REEVALUATE_KEY', '0') == '1' else [] - - -@autotune(list(configs_bwd_dkdv_bsa_varlen), key=bwd_dkdv_bsa_varlen_reevaluate_keys) -@triton.jit -def _attn_bwd_dkdv_bsa_varlen_wrapper( - Q, - K, - V, - sm_scale, # softmax scale - DO, - DK, - DV, - M, # lse (log2) - D, - block_indices, - block_indices_lens, - # stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO. - # qkv - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - # dk dv do - stride_dkz, - stride_dkh, - stride_dkn, - stride_dkk, - stride_dvz, - stride_dvh, - stride_dvn, - stride_dvk, - stride_doz, - stride_doh, - stride_dom, - stride_dok, - # m, d - stride_mz, - stride_mh, - stride_mm, - stride_dz, - stride_dh, - stride_dm, - # - stride_bz, - stride_bh, - stride_bn, - stride_bm, # block_indices - stride_lz, - stride_lh, - stride_ln, # block_indices_lens - # - H, - N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_N_DQ_LG: tl.constexpr, # logical block size - HEAD_DIM: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking -): - - tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N == 0) - start_n = tl.program_id(0) - - off_hz = tl.program_id(2) - off_z = off_hz // H - off_h = off_hz % H - - off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh - off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh - - off_dk = off_z.to(tl.int64) * stride_dkz + off_h.to(tl.int64) * stride_dkh - off_dv = off_z.to(tl.int64) * stride_dvz + off_h.to(tl.int64) * stride_dvh - off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh - - off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh - off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh - - off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh - off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh - - # offset pointers for batch/head - Q += off_q - K += off_k - V += off_v - DO += off_do - DK += off_dk - DV += off_dv - - M += off_m - D += off_d - block_indices += off_block_incides - block_indices_lens += off_block_incides_lens - - # ---------------------------------------- [DKDV] ---------------------------------------- - - dv = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32) - dk = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32) - - # load K and V: they stay in SRAM throughout the inner loop. - - K_block_ptr = tl.make_block_ptr( - base=K, - shape=(N_CTX, HEAD_DIM), - strides=(stride_kn, stride_kk), - offsets=(start_n * BLOCK_N, 0), - block_shape=(BLOCK_N, HEAD_DIM), - order=(1, 0), - ) - V_block_ptr = tl.make_block_ptr( - base=V, - shape=(N_CTX, HEAD_DIM), - strides=(stride_vn, stride_vk), - offsets=(start_n * BLOCK_N, 0), - block_shape=(BLOCK_N, HEAD_DIM), - order=(1, 0), - ) - DK_block_ptr = tl.make_block_ptr( - base=DK, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dkn, stride_dkk), - offsets=(start_n * BLOCK_N, 0), - block_shape=(BLOCK_N, HEAD_DIM), - order=(1, 0), - ) - DV_block_ptr = tl.make_block_ptr( - base=DV, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dvn, stride_dvk), - offsets=(start_n * BLOCK_N, 0), - block_shape=(BLOCK_N, HEAD_DIM), - order=(1, 0), - ) - - k = tl.load(K_block_ptr) - v = tl.load(V_block_ptr) - - k_compress_idx = start_n * BLOCK_N // BLOCK_N_DQ_LG - block_indices_i = block_indices + k_compress_idx * stride_bn - block_indices_lens_i = block_indices_lens + k_compress_idx * stride_ln - - dk, dv = _attn_bwd_dkdv_bsa_varlen( - dk, - dv, - k, - v, - Q, - DO, - M, - D, - block_indices_i, - block_indices_lens_i, - # shared by Q/K/V/DO. - stride_qm, - stride_qk, - stride_dom, - stride_dok, - stride_mm, - stride_dm, - # - stride_bm, - N_CTX, - BLOCK_M, - HEAD_DIM, - ) - - # Write back dk - dk *= sm_scale # S = scale * QKT; dK = scale * QdST - tl.store(DK_block_ptr, dk.to(k.dtype)) - - # Write back dv - tl.store(DV_block_ptr, dv.to(v.dtype)) - - -configs_bwd_dq_bsa_varlen_preset = { - 'default': { - 'BLOCK_N_DQ': 64, - 'num_stages': 2, - 'num_warps': 8, - }, - 'BLOCK_N_DQ_LG=64': { - 'BLOCK_N_DQ': 64, - 'num_stages': 2, - 'num_warps': 4, - }, -} -configs_bwd_dq_bsa_varlen = [ - triton.Config({'BLOCK_N_DQ': BN}, num_stages=s, num_warps=w) \ - for BN in [32, 64, 128] \ - for s in [2, 3, 4, 5] \ - for w in [4, 8] \ -] -bwd_dq_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get( - 'TRITON_REEVALUATE_KEY', '0') == '1' else [] - - -@autotune(list(configs_bwd_dq_bsa_varlen), key=bwd_dq_bsa_varlen_reevaluate_keys) -@triton.jit -def _attn_bwd_dq_bsa_varlen_wrapper( - Q, - K, - V, # softmax scale - DO, - DQ, - M, # lse (log2) - D, - block_indices, - block_indices_lens, - # stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO. - # qkv - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - # dq do - stride_dqz, - stride_dqh, - stride_dqm, - stride_dqk, - stride_doz, - stride_doh, - stride_dom, - stride_dok, - # m, d - stride_mz, - stride_mh, - stride_mm, - stride_dz, - stride_dh, - stride_dm, - # - stride_bz, - stride_bh, - stride_bm, - stride_bn, # block_indices - stride_lz, - stride_lh, - stride_lm, # block_indices_lens - # - H, - N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_N_DQ_LG: tl.constexpr, # logical block size - BLOCK_N_DQ: tl.constexpr, - HEAD_DIM: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking -): - - tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ == 0) - tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0) - - LN2: tl.constexpr = 0.6931471824645996 # = ln(2) - - start_m = tl.program_id(0) - - off_hz = tl.program_id(2) - off_z = off_hz // H - off_h = off_hz % H - - off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh - off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh - - off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh - off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh - - off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh - off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh - - off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh - off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh - - # offset pointers for batch/head - Q += off_q - K += off_k - V += off_v - DO += off_do - DQ += off_dq - - M += off_m - D += off_d - block_indices += off_block_incides - block_indices_lens += off_block_incides_lens - - # ---------------------------------------- [DQ] ---------------------------------------- - - Q_block_ptr = tl.make_block_ptr( - base=Q, - shape=(N_CTX, HEAD_DIM), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - DO_block_ptr = tl.make_block_ptr( - base=DO, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dom, stride_dok), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - DQ_block_ptr = tl.make_block_ptr( - base=DQ, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dqm, stride_dqk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - - q = tl.load(Q_block_ptr) - do = tl.load(DO_block_ptr) - - dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - - start_m = start_m * BLOCK_M - - offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm - offs_d = start_m + tl.arange(0, BLOCK_M) * stride_dm - - m = tl.load(M + offs_m) - m = m[:, None] - - d = tl.load(D + offs_d) - d = d[:, None] - - block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm - block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm - - dq = _attn_bwd_dq_bsa_varlen( - dq, - q, - do, - m, - d, - K, - V, - N_CTX, - BLOCK_N_DQ, - BLOCK_N_DQ_LG, - HEAD_DIM, - block_indices_m, - block_indices_lens_m, - stride_bn, - stride_kn, - stride_kk, - stride_vn, - stride_vk, - ) - - # Write back dQ. - dq *= LN2 - tl.store(DQ_block_ptr, dq.to(q.dtype)) - - -configs_bwd_dq_bsa_varlen_align_preset = { - 'default': { - 'num_stages': 2, - 'num_warps': 8, - }, - 'BLOCK_N_DQ_LG=64': { - 'num_stages': 2, - 'num_warps': 4, - }, -} -configs_bwd_dq_bsa_varlen_align = [ - triton.Config({}, num_stages=s, num_warps=w) \ - for s in [2, 3, 4, 5] \ - for w in [4, 8] \ -] -bwd_dq_bsa_varlen_align_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get( - 'TRITON_REEVALUATE_KEY', '0') == '1' else [] - - -@autotune(list(configs_bwd_dq_bsa_varlen_align), key=bwd_dq_bsa_varlen_align_reevaluate_keys) -@triton.jit -def _attn_bwd_dq_bsa_varlen_align_wrapper( - Q, - K, - V, # softmax scale - DO, - DQ, - M, # lse (log2) - D, - block_indices, - block_indices_lens, - # qkv - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - # dq do - stride_dqz, - stride_dqh, - stride_dqm, - stride_dqk, - stride_doz, - stride_doh, - stride_dom, - stride_dok, - # m, d - stride_mz, - stride_mh, - stride_mm, - stride_dz, - stride_dh, - stride_dm, - # - stride_bz, - stride_bh, - stride_bm, - stride_bn, # block_indices - stride_lz, - stride_lh, - stride_lm, # block_indices_lens - # - H, - N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_N_DQ_LG: tl.constexpr, # logical block size - HEAD_DIM: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking -): - - tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ_LG == 0) - tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0) - - LN2: tl.constexpr = 0.6931471824645996 # = ln(2) - - start_m = tl.program_id(0) - - off_hz = tl.program_id(2) - off_z = off_hz // H - off_h = off_hz % H - - off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh - off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh - - off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh - off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh - - off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh - off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh - - off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh - off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh - - # offset pointers for batch/head - Q += off_q - K += off_k - V += off_v - DO += off_do - DQ += off_dq - - M += off_m - D += off_d - block_indices += off_block_incides - block_indices_lens += off_block_incides_lens - - # ---------------------------------------- [DQ] ---------------------------------------- - - Q_block_ptr = tl.make_block_ptr( - base=Q, - shape=(N_CTX, HEAD_DIM), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - DO_block_ptr = tl.make_block_ptr( - base=DO, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dom, stride_dok), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - DQ_block_ptr = tl.make_block_ptr( - base=DQ, - shape=(N_CTX, HEAD_DIM), - strides=(stride_dqm, stride_dqk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - - q = tl.load(Q_block_ptr) - do = tl.load(DO_block_ptr) - - dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - - start_m = start_m * BLOCK_M - - offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm - offs_d = start_m + tl.arange(0, BLOCK_M) * stride_dm - - m = tl.load(M + offs_m) - m = m[:, None] - - # D (= delta) is pre-divided by ds_scale. - d = tl.load(D + offs_d) - d = d[:, None] - - block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm - block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm - - dq = _attn_bwd_dq_bsa_varlen_align( - dq, - q, - do, - m, - d, - K, - V, - N_CTX, - BLOCK_N_DQ_LG, - HEAD_DIM, - block_indices_m, - block_indices_lens_m, - stride_bn, - stride_kn, - stride_kk, - stride_vn, - stride_vk, - ) - - # Write back dQ. - dq *= LN2 - tl.store(DQ_block_ptr, dq.to(q.dtype)) diff --git a/fastvideo/train/AGENTS.md b/fastvideo/train/AGENTS.md index 045e36d14..210fb337a 100644 --- a/fastvideo/train/AGENTS.md +++ b/fastvideo/train/AGENTS.md @@ -19,7 +19,7 @@ train/ │ └── rl/ # RL methods, sampling helpers, rewards ├── models/ │ ├── base.py # ModelBase / CausalModelBase wrappers -│ ├── wan/, hunyuan/, cosmos/, kandinsky5/, longcat/, ltx2/, matrixgame2/, minimax_h3/ # Per-family wrappers +│ ├── wan/, hunyuan/, cosmos/, kandinsky5/, ltx2/, matrixgame2/, minimax_h3/ # Per-family wrappers ├── callbacks/ # callback.py base + ema, grad_clip, validation └── utils/ ├── training_config.py # Hierarchical YAML config dataclasses diff --git a/fastvideo/train/entrypoint/train.py b/fastvideo/train/entrypoint/train.py index 3f71e2a72..8533e5c1f 100644 --- a/fastvideo/train/entrypoint/train.py +++ b/fastvideo/train/entrypoint/train.py @@ -54,20 +54,13 @@ def run_training_from_config( cfg = load_run_config(config_path, overrides=overrides) tc = cfg.training - model_path_lower = str(tc.model_path).lower() - - # Auto-set attention backend for model families that require a specific - # backend at load time, unless the user already overrode it explicitly. + # Auto-set the VSA attention backend when VSA sparsity is enabled, + # unless the user already overrode it explicitly. if tc.vsa_sparsity > 0.0: os.environ.setdefault( "FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN", ) - elif ("turbodiffusion" in model_path_lower or "turbowan" in model_path_lower): - os.environ.setdefault( - "FASTVIDEO_ATTENTION_BACKEND", - "SLA_ATTN", - ) maybe_init_distributed_environment_and_model_parallel( tc.distributed.tp_size, diff --git a/fastvideo/train/models/longcat/__init__.py b/fastvideo/train/models/longcat/__init__.py deleted file mode 100644 index c9aacfd95..000000000 --- a/fastvideo/train/models/longcat/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""LongCat model plugin package.""" - -from fastvideo.train.models.longcat.longcat import ( - LongCatModel as LongCatModel, ) diff --git a/fastvideo/train/models/longcat/longcat.py b/fastvideo/train/models/longcat/longcat.py deleted file mode 100644 index 8e78d10a8..000000000 --- a/fastvideo/train/models/longcat/longcat.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""LongCat model plugin (per-role instance).""" - -from __future__ import annotations - -from typing import Any, Literal, TYPE_CHECKING - -import torch - -from fastvideo.pipelines import TrainingBatch -from fastvideo.train.models.wan.wan import WanModel - -if TYPE_CHECKING: - from fastvideo.train.utils.training_config import TrainingConfig - - -class LongCatModel(WanModel): - """LongCat per-role model for training and distillation.""" - - _transformer_cls_name: str = "LongCatTransformer3DModel" - - @staticmethod - def _validate_flow_shift(flow_shift: float | None) -> float: - if flow_shift is None: - return 12.0 - - validated = float(flow_shift) - if validated == 0.0: - raise ValueError("LongCat training does not support flow_shift=0.0 because " - "it collapses FlowMatch training timesteps. Use 12.0 to " - "match the released LongCat scheduler config.") - return validated - - def __init__( - self, - *, - init_from: str, - training_config: TrainingConfig, - trainable: bool = True, - disable_custom_init_weights: bool = False, - flow_shift: float = 12.0, - enable_gradient_checkpointing_type: str | None = None, - transformer_override_safetensor: str | None = None, - ) -> None: - super().__init__( - init_from=init_from, - training_config=training_config, - trainable=trainable, - disable_custom_init_weights=disable_custom_init_weights, - flow_shift=self._validate_flow_shift(flow_shift), - enable_gradient_checkpointing_type=enable_gradient_checkpointing_type, - transformer_override_safetensor=transformer_override_safetensor, - ) - - def _init_timestep_mechanics(self) -> None: - assert self.training_config is not None - tc = self.training_config - flow_shift = getattr(tc.pipeline_config, "flow_shift", None) # type: ignore[union-attr] - self.timestep_shift = self._validate_flow_shift(flow_shift) - self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps) - self.min_timestep = 0 - self.max_timestep = self.num_train_timestep - - def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch: - training_batch.attn_metadata = None - return training_batch - - def _build_distill_input_kwargs( - self, - noise_input: torch.Tensor, - timestep: torch.Tensor, - text_dict: dict[str, torch.Tensor] | None, - clean_x: torch.Tensor | None = None, - aug_t: torch.Tensor | None = None, - ) -> dict[str, Any]: - if clean_x is not None or aug_t is not None: - raise NotImplementedError("LongCatModel does not support clean-history teacher forcing") - if text_dict is None: - raise ValueError("text_dict cannot be None for LongCat distillation") - - batch_size = int(noise_input.shape[0]) - if timestep.ndim == 0: - timestep = timestep.view(1).expand(batch_size) - elif timestep.ndim == 1 and int(timestep.shape[0]) == 1 and batch_size > 1: - timestep = timestep.expand(batch_size) - - return { - "hidden_states": noise_input.permute(0, 2, 1, 3, 4), - "encoder_hidden_states": text_dict["encoder_hidden_states"], - "encoder_attention_mask": text_dict["encoder_attention_mask"], - "timestep": timestep, - } - - def predict_noise( - self, - noisy_latents: torch.Tensor, - timestep: torch.Tensor, - batch: TrainingBatch, - *, - conditional: bool, - cfg_uncond: dict[str, Any] | None = None, - attn_kind: Literal["dense", "vsa"] = "dense", - clean_x: torch.Tensor | None = None, - aug_t: torch.Tensor | None = None, - ) -> torch.Tensor: - """Adapt LongCat's sign convention to FineTuneMethod's target. - - ``LongCatTransformer3DModel`` is pretrained to output the - ``clean - noise`` direction; ``LongCatDenoisingStage`` (the - bidirectional inference pipeline) explicitly negates the - transformer output before handing it to - ``FlowMatchEulerDiscreteScheduler.step``. Training methods on - the other hand (``FineTuneMethod``, - ``DiffusionForcingSFTMethod``) target ``noise - clean`` - directly (the standard rectified-flow velocity Wan uses). - - Without the negation here, the loss MSE pushes the transformer - toward ``noise - clean``, flipping its native output sign over - training. Inference then applies its own negation on top, so - the scheduler receives the wrong direction and produces noise - even while the training loss is dropping. Verified empirically - on a 100-step LongCat overfit run: step 0 generated meaningful - video, step 100 was pure noise despite low loss. - - Negating in ``predict_noise`` keeps the transformer's - pretrained sign convention intact while presenting the - training methods with a Wan-compatible - ``pred ≈ noise - clean`` for MSE. - """ - pred = super().predict_noise( - noisy_latents, - timestep, - batch, - conditional=conditional, - cfg_uncond=cfg_uncond, - attn_kind=attn_kind, - clean_x=clean_x, - aug_t=aug_t, - ) - return -pred diff --git a/mkdocs.yml b/mkdocs.yml index 53ed86a51..4f63adf40 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -166,12 +166,8 @@ nav: - Hunyuan: cookbook/hunyuan.md - Kandinsky 5: cookbook/kandinsky5.md - Kandinsky 6: cookbook/kandinsky6.md - - LongCat: cookbook/longcat.md - - TurboDiffusion: cookbook/turbodiffusion.md - Image: - FLUX: cookbook/flux.md - - GLM-Image: cookbook/glm-image.md - - Z-Image: cookbook/z-image.md - Stable Diffusion: cookbook/stable-diffusion.md - Audio: - Stable Audio: cookbook/stable-audio.md diff --git a/pyproject.toml b/pyproject.toml index d6ed836bb..ae8a69758 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,10 +21,9 @@ dependencies = [ "requests>=2.32.2", # Machine Learning & Transformers - # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in - # transformers 5.0.0. Floor raised from >=5.0.0 to >=5.15.0 for the public - # vision interpolation helper (get_vision_interpolation_indices_and_weights) - # that the MiniMax-H3 Qwen3-VL encoder cross-checks against. + # Floor raised from >=5.0.0 to >=5.15.0 for the public vision interpolation + # helper (get_vision_interpolation_indices_and_weights) that the MiniMax-H3 + # Qwen3-VL encoder cross-checks against. "transformers>=5.15.0", # <0.23: tokenizers 0.23 renamed RobertaProcessing's binding args, so # transformers' CLIP-style tokenizer loading dies with diff --git a/pyproject_other.toml b/pyproject_other.toml index e74e7a72a..cdf5db5db 100644 --- a/pyproject_other.toml +++ b/pyproject_other.toml @@ -21,8 +21,7 @@ dependencies = [ "requests>=2.32.2", # Machine Learning & Transformers - # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in - # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc). + # Floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc). "transformers>=5.0.0", "tokenizers>=0.20.1", "sentencepiece>=0.2.0", diff --git a/scripts/checkpoint_conversion/AGENTS.md b/scripts/checkpoint_conversion/AGENTS.md index 730ea9937..c5a62dd71 100644 --- a/scripts/checkpoint_conversion/AGENTS.md +++ b/scripts/checkpoint_conversion/AGENTS.md @@ -9,19 +9,11 @@ ``` checkpoint_conversion/ -├── convert_gamecraft_full.py # Combined DiT + VAE -├── convert_gamecraft_vae.py # VAE only -├── convert_gamecraft_weights.py # DiT only ├── convert_gen3c_to_fastvideo.py ├── convert_ltx2_weights.py ├── convert_minimax_h3_adaln_rank.py # Rank-reduces AdaLN in place; same family -├── convert_turbodiffusion_to_diffusers.py -├── convert_turbodiffusion_i2v_to_diffusers.py -├── extract_llava_text_encoder.py # Encoder extraction from a multimodal repo -├── longcat_to_fastvideo.py ├── stable_audio_to_diffusers.py ├── wan_to_diffusers.py -├── validate_longcat_weights.py # Post-conversion validation ├── pt_to_safetensors.py # Generic format flip └── create_hf_repo.py # Push to HF after conversion ``` diff --git a/scripts/checkpoint_conversion/convert_gamecraft_full.py b/scripts/checkpoint_conversion/convert_gamecraft_full.py deleted file mode 100644 index 20057c26e..000000000 --- a/scripts/checkpoint_conversion/convert_gamecraft_full.py +++ /dev/null @@ -1,239 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Convert the full Hunyuan-GameCraft-1.0 model to FastVideo diffusers format. - -This script converts: -1. Transformer (DiT) weights from DeepSpeed format -2. VAE weights -3. Text encoders (LLaVA-LLaMA-3-8B and CLIP) -4. Creates model_index.json and scheduler config - -Usage: - python scripts/checkpoint_conversion/convert_gamecraft_full.py \ - --input-dir Hunyuan-GameCraft-1.0/weights \ - --output-dir official_weights/hunyuan-gamecraft-diffusers - - # Or specify individual paths: - python scripts/checkpoint_conversion/convert_gamecraft_full.py \ - --transformer Hunyuan-GameCraft-1.0/weights/gamecraft_models/mp_rank_00_model_states.pt \ - --vae Hunyuan-GameCraft-1.0/weights/stdmodels/vae_3d/hyvae/checkpoint-step-270000.ckpt \ - --text-encoder Hunyuan-GameCraft-1.0/weights/stdmodels/llava-llama-3-8b-v1_1-transformers \ - --text-encoder-2 Hunyuan-GameCraft-1.0/weights/stdmodels/openai_clip-vit-large-patch14 \ - --output-dir official_weights/hunyuan-gamecraft-diffusers -""" -from __future__ import annotations - -import argparse -import json -import shutil -from pathlib import Path - -import torch -from safetensors.torch import save_file - -# Import conversion functions -from convert_gamecraft_weights import convert_weights as convert_transformer -from convert_gamecraft_vae import convert_gamecraft_vae - - -def create_model_index(output_dir: Path) -> None: - """Create model_index.json for the pipeline.""" - model_index = { - "_class_name": "HunyuanGameCraftPipeline", - "_diffusers_version": "0.30.0", - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["transformers", "LlamaModel"], - "text_encoder_2": ["transformers", "CLIPTextModel"], - "tokenizer": ["transformers", "AutoTokenizer"], - "tokenizer_2": ["transformers", "CLIPTokenizer"], - "transformer": ["fastvideo", "HunyuanGameCraftTransformer3DModel"], - "vae": ["fastvideo", "AutoencoderKLCausal3D"] - } - - with open(output_dir / "model_index.json", "w") as f: - json.dump(model_index, f, indent=2) - print(f"Created model_index.json") - - -def create_scheduler_config(output_dir: Path) -> None: - """Create scheduler config for FlowMatchEulerDiscreteScheduler.""" - scheduler_dir = output_dir / "scheduler" - scheduler_dir.mkdir(parents=True, exist_ok=True) - - scheduler_config = { - "_class_name": "FlowMatchEulerDiscreteScheduler", - "_diffusers_version": "0.30.0", - "base_image_seq_len": 256, - "base_shift": 0.5, - "invert_sigmas": False, - "max_shift": 1.15, - "num_train_timesteps": 1000, - "shift": 7.0, - "use_dynamic_shifting": True - } - - with open(scheduler_dir / "scheduler_config.json", "w") as f: - json.dump(scheduler_config, f, indent=2) - print(f"Created scheduler config") - - -def copy_text_encoder(src_dir: Path, dst_dir: Path, name: str) -> None: - """Copy text encoder files to output directory.""" - if not src_dir.exists(): - print(f"Warning: {name} not found at {src_dir}") - return - - dst_dir.mkdir(parents=True, exist_ok=True) - - # Copy all files - for src_file in src_dir.iterdir(): - if src_file.is_file(): - shutil.copy2(src_file, dst_dir / src_file.name) - - print(f"Copied {name} to {dst_dir}") - - -def copy_tokenizer(src_dir: Path, dst_dir: Path, name: str) -> None: - """Copy tokenizer files to output directory.""" - if not src_dir.exists(): - print(f"Warning: {name} not found at {src_dir}") - return - - dst_dir.mkdir(parents=True, exist_ok=True) - - # Copy tokenizer-related files - tokenizer_files = [ - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - "vocab.json", - "merges.txt", - "added_tokens.json", - ] - - for filename in tokenizer_files: - src_file = src_dir / filename - if src_file.exists(): - shutil.copy2(src_file, dst_dir / filename) - - print(f"Copied {name} to {dst_dir}") - - -def convert_full_model( - transformer_path: Path | None = None, - vae_path: Path | None = None, - text_encoder_path: Path | None = None, - text_encoder_2_path: Path | None = None, - input_dir: Path | None = None, - output_dir: Path = Path("official_weights/hunyuan-gamecraft-diffusers"), - verbose: bool = False, -) -> None: - """Convert all components of GameCraft to diffusers format.""" - - # Resolve paths from input_dir if individual paths not provided - if input_dir is not None: - if transformer_path is None: - transformer_path = input_dir / "gamecraft_models" / "mp_rank_00_model_states.pt" - if vae_path is None: - vae_path = input_dir / "stdmodels" / "vae_3d" / "hyvae" / "checkpoint-step-270000.ckpt" - if not vae_path.exists(): - vae_path = input_dir / "stdmodels" / "vae_3d" / "hyvae" / "pytorch_model.pt" - if text_encoder_path is None: - text_encoder_path = input_dir / "stdmodels" / "llava-llama-3-8b-v1_1-transformers" - if text_encoder_2_path is None: - text_encoder_2_path = input_dir / "stdmodels" / "openai_clip-vit-large-patch14" - - output_dir.mkdir(parents=True, exist_ok=True) - print(f"Converting GameCraft to diffusers format at {output_dir}") - print("=" * 60) - - # 1. Convert transformer - if transformer_path and transformer_path.exists(): - print(f"\n[1/5] Converting transformer from {transformer_path}") - convert_transformer( - input_path=transformer_path, - output_dir=output_dir, - save_config=True, - verbose=verbose, - ) - else: - print(f"\n[1/5] Skipping transformer (not found: {transformer_path})") - - # 2. Convert VAE - if vae_path and vae_path.exists(): - print(f"\n[2/5] Converting VAE from {vae_path}") - convert_gamecraft_vae( - input_path=vae_path, - output_dir=output_dir / "vae", - copy_config=True, - ) - else: - print(f"\n[2/5] Skipping VAE (not found: {vae_path})") - - # 3. Copy text encoder (LLaMA) - if text_encoder_path and text_encoder_path.exists(): - print(f"\n[3/5] Copying text encoder from {text_encoder_path}") - copy_text_encoder(text_encoder_path, output_dir / "text_encoder", "text_encoder") - copy_tokenizer(text_encoder_path, output_dir / "tokenizer", "tokenizer") - else: - print(f"\n[3/5] Skipping text_encoder (not found: {text_encoder_path})") - - # 4. Copy text encoder 2 (CLIP) - if text_encoder_2_path and text_encoder_2_path.exists(): - print(f"\n[4/5] Copying text encoder 2 from {text_encoder_2_path}") - copy_text_encoder(text_encoder_2_path, output_dir / "text_encoder_2", "text_encoder_2") - copy_tokenizer(text_encoder_2_path, output_dir / "tokenizer_2", "tokenizer_2") - else: - print(f"\n[4/5] Skipping text_encoder_2 (not found: {text_encoder_2_path})") - - # 5. Create model_index.json and scheduler - print(f"\n[5/5] Creating model_index.json and scheduler config") - create_model_index(output_dir) - create_scheduler_config(output_dir) - - print("\n" + "=" * 60) - print(f"Conversion complete! Output: {output_dir}") - print("\nDirectory structure:") - for item in sorted(output_dir.rglob("*")): - if item.is_file(): - rel_path = item.relative_to(output_dir) - size_mb = item.stat().st_size / (1024 * 1024) - print(f" {rel_path} ({size_mb:.1f} MB)") - - -def main(): - parser = argparse.ArgumentParser( - description="Convert full Hunyuan-GameCraft-1.0 model to FastVideo diffusers format.") - parser.add_argument("--input-dir", type=str, default=None, help="Path to Hunyuan-GameCraft-1.0/weights directory") - parser.add_argument("--transformer", - type=str, - default=None, - help="Path to transformer checkpoint (mp_rank_00_model_states.pt)") - parser.add_argument("--vae", type=str, default=None, help="Path to VAE checkpoint") - parser.add_argument("--text-encoder", type=str, default=None, help="Path to text encoder (LLaVA-LLaMA) directory") - parser.add_argument("--text-encoder-2", type=str, default=None, help="Path to text encoder 2 (CLIP) directory") - parser.add_argument("--output-dir", - type=str, - default="official_weights/hunyuan-gamecraft-diffusers", - help="Output directory for converted model") - parser.add_argument("--verbose", "-v", action="store_true", help="Print detailed conversion info") - - args = parser.parse_args() - - # Require either input-dir or individual paths - if args.input_dir is None and args.transformer is None: - parser.error("Either --input-dir or individual component paths must be provided") - - convert_full_model( - transformer_path=Path(args.transformer) if args.transformer else None, - vae_path=Path(args.vae) if args.vae else None, - text_encoder_path=Path(args.text_encoder) if args.text_encoder else None, - text_encoder_2_path=Path(args.text_encoder_2) if args.text_encoder_2 else None, - input_dir=Path(args.input_dir) if args.input_dir else None, - output_dir=Path(args.output_dir), - verbose=args.verbose, - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/convert_gamecraft_vae.py b/scripts/checkpoint_conversion/convert_gamecraft_vae.py deleted file mode 100644 index de44fb8a6..000000000 --- a/scripts/checkpoint_conversion/convert_gamecraft_vae.py +++ /dev/null @@ -1,93 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Convert GameCraft VAE weights to FastVideo safetensors format. - -The official GameCraft VAE (checkpoint-step-270000.ckpt or pytorch_model.pt) uses -a "vae." prefix. This script extracts and converts for FastVideo GameCraftVAE. - -Usage: - python scripts/checkpoint_conversion/convert_gamecraft_vae.py \ - --input Hunyuan-GameCraft-1.0/weights/stdmodels/vae_3d/hyvae/checkpoint-step-270000.ckpt \ - --output official_weights/hunyuan-gamecraft/vae -""" -from __future__ import annotations - -import argparse -import json -import shutil -from pathlib import Path - -import torch -from safetensors.torch import save_file - - -def convert_gamecraft_vae( - input_path: Path, - output_dir: Path, - copy_config: bool = True, -) -> dict: - """Convert official GameCraft VAE checkpoint to FastVideo format.""" - ckpt = torch.load(input_path, map_location="cpu", weights_only=True) - if "state_dict" in ckpt: - state_dict = ckpt["state_dict"] - else: - state_dict = ckpt - - # Extract VAE weights and strip "vae." prefix - vae_sd = {k.replace("vae.", ""): v for k, v in state_dict.items() if k.startswith("vae.")} - print(f"Extracted {len(vae_sd)} VAE parameters") - - output_dir.mkdir(parents=True, exist_ok=True) - save_file(vae_sd, output_dir / "diffusion_pytorch_model.safetensors") - print(f"Saved to {output_dir / 'diffusion_pytorch_model.safetensors'}") - - if copy_config: - config_src = input_path.parent / "config.json" - if config_src.exists(): - config_dst = output_dir / "config.json" - config = json.loads(config_src.read_text()) - config["_class_name"] = "AutoencoderKLCausal3D" - with open(config_dst, "w") as f: - json.dump(config, f, indent=2) - print(f"Saved config to {config_dst}") - else: - print(f"Warning: config.json not found at {config_src}") - - return {"total": len(vae_sd)} - - -def main(): - parser = argparse.ArgumentParser(description="Convert GameCraft VAE weights to FastVideo safetensors format.") - parser.add_argument( - "--input", - type=str, - required=True, - help="Path to official VAE checkpoint (.ckpt or .pt)", - ) - parser.add_argument( - "--output", - type=str, - default="official_weights/hunyuan-gamecraft/vae", - help="Output directory for converted weights", - ) - parser.add_argument( - "--no-config", - action="store_true", - help="Don't copy config.json", - ) - args = parser.parse_args() - - input_path = Path(args.input) - if not input_path.exists(): - raise FileNotFoundError(f"Input not found: {input_path}") - - stats = convert_gamecraft_vae( - input_path=input_path, - output_dir=Path(args.output), - copy_config=not args.no_config, - ) - print(f"\nConversion complete! Total parameters: {stats['total']}") - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/convert_gamecraft_weights.py b/scripts/checkpoint_conversion/convert_gamecraft_weights.py deleted file mode 100644 index 97a208454..000000000 --- a/scripts/checkpoint_conversion/convert_gamecraft_weights.py +++ /dev/null @@ -1,242 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Convert Hunyuan-GameCraft-1.0 weights from DeepSpeed format to FastVideo format. - -The official GameCraft checkpoint (mp_rank_00_model_states.pt) uses slightly different -naming conventions than FastVideo. This script converts the weights to be compatible -with FastVideo's HunyuanGameCraft model. - -Usage: - python scripts/checkpoint_conversion/convert_gamecraft_weights.py \ - --input Hunyuan-GameCraft-1.0/weights/gamecraft_models/mp_rank_00_model_states.pt \ - --output official_weights/hunyuan-gamecraft/transformer/diffusion_pytorch_model.safetensors -""" - -from __future__ import annotations - -import argparse -import json -import re -from collections import OrderedDict -from pathlib import Path - -import torch -from safetensors.torch import save_file - -# Mapping from official GameCraft naming to FastVideo naming -# GameCraft weights are already close to FastVideo format, but need minor adjustments -PARAM_NAME_MAP: dict[str, str] = { - # 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 - rename individual_token_refiner.blocks to refiner_blocks - # This is applied first, then subsequent mappings handle the sub-components - 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", - - # Text in c_embedder (context embedder) naming - r"^txt_in\.c_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1", - r"^txt_in\.c_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1", - - # Text in t_embedder (time embedder) naming - r"^txt_in\.t_embedder\.mlp\.0\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1", - r"^txt_in\.t_embedder\.mlp\.2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1", - - # Refiner block MLP naming - need to match the actual checkpoint names - 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", - - # Refiner block adaLN modulation naming - matches .adaLN_modulation.1.weight/bias - r"^txt_in\.refiner_blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$": - r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", - - # Camera net - keep as-is (already correct) - # r"^camera_net\.(.*)$": r"camera_net.\1", -} - - -def apply_mapping(key: str) -> str: - """Apply name mapping to convert official key to FastVideo key. - - Mappings are applied repeatedly until no more changes occur. - This allows for chained transformations. - """ - prev_key = None - while prev_key != key: - prev_key = key - for pattern, replacement in PARAM_NAME_MAP.items(): - if re.match(pattern, key): - key = re.sub(pattern, replacement, key) - break # Apply only one mapping per iteration - return key - - -def load_deepspeed_checkpoint(path: Path) -> dict[str, torch.Tensor]: - """Load weights from a DeepSpeed checkpoint file.""" - checkpoint = torch.load(path, map_location="cpu") - - if "module" in checkpoint: - state_dict = checkpoint["module"] - elif "state_dict" in checkpoint: - state_dict = checkpoint["state_dict"] - elif "model" in checkpoint: - state_dict = checkpoint["model"] - else: - state_dict = checkpoint - - if not isinstance(state_dict, dict): - raise TypeError(f"Expected dict state_dict, got {type(state_dict)}") - - return state_dict - - -def convert_weights( - input_path: Path, - output_dir: Path, - save_config: bool = True, - verbose: bool = False, -) -> dict[str, int]: - """ - Convert GameCraft weights to FastVideo format. - - Returns: - dict with conversion statistics - """ - print(f"Loading checkpoint from {input_path}") - state_dict = load_deepspeed_checkpoint(input_path) - - print(f"Found {len(state_dict)} parameters") - - # Convert weights - converted_state_dict: OrderedDict[str, torch.Tensor] = OrderedDict() - renamed_count = 0 - unchanged_count = 0 - - for key, value in state_dict.items(): - new_key = apply_mapping(key) - if new_key != key: - if verbose: - print(f" {key} -> {new_key}") - renamed_count += 1 - else: - unchanged_count += 1 - converted_state_dict[new_key] = value - - print(f"Renamed {renamed_count} parameters, {unchanged_count} unchanged") - - # Create output directory structure - transformer_dir = output_dir / "transformer" - transformer_dir.mkdir(parents=True, exist_ok=True) - - # Save as safetensors - output_path = transformer_dir / "diffusion_pytorch_model.safetensors" - print(f"Saving converted weights to {output_path}") - save_file(converted_state_dict, str(output_path)) - - # Save config if requested - if save_config: - config = { - "_class_name": "HunyuanGameCraftTransformer3DModel", - "_fastvideo_version": "0.1.0", - "in_channels": 33, # 16 + 16 + 1 for concat mode - "out_channels": 16, - "hidden_size": 3072, - "num_attention_heads": 24, - "attention_head_dim": 128, - "num_layers": 20, - "num_single_layers": 40, - "num_refiner_layers": 2, - "patch_size": [1, 2, 2], - "rope_axes_dim": [16, 56, 56], - "text_embed_dim": 4096, - "pooled_projection_dim": 768, - "mlp_ratio": 4.0, - "guidance_embeds": False, - "camera_net": True, - } - config_path = transformer_dir / "config.json" - with open(config_path, "w") as f: - json.dump(config, f, indent=2) - print(f"Saved config to {config_path}") - - # Verify camera_net weights - camera_keys = [k for k in converted_state_dict.keys() if "camera_net" in k] - print(f"\nCamera net weights ({len(camera_keys)}):") - for k in camera_keys: - print(f" {k}: {converted_state_dict[k].shape}") - - # Verify img_in weights (should be 33 channels) - if "img_in.proj.weight" in converted_state_dict: - img_in_shape = converted_state_dict["img_in.proj.weight"].shape - print(f"\nimg_in.proj.weight shape: {img_in_shape}") - expected_in_channels = 33 - actual_in_channels = img_in_shape[1] - if actual_in_channels != expected_in_channels: - print(f" WARNING: Expected {expected_in_channels} input channels, got {actual_in_channels}") - else: - print(f" OK: 33 input channels (16 latent + 16 gt_latent + 1 mask)") - - return { - "total": len(converted_state_dict), - "renamed": renamed_count, - "unchanged": unchanged_count, - } - - -def main(): - parser = argparse.ArgumentParser(description="Convert Hunyuan-GameCraft-1.0 weights to FastVideo format.") - parser.add_argument("--input", - type=str, - required=True, - help="Path to official GameCraft checkpoint (mp_rank_00_model_states.pt)") - parser.add_argument("--output", - type=str, - default="official_weights/hunyuan-gamecraft", - help="Output directory for converted weights") - parser.add_argument("--no-config", action="store_true", help="Don't save config.json") - parser.add_argument("--verbose", "-v", action="store_true", help="Print all renamed parameters") - - args = parser.parse_args() - - input_path = Path(args.input) - output_dir = Path(args.output) - - if not input_path.exists(): - raise FileNotFoundError(f"Input file not found: {input_path}") - - stats = convert_weights( - input_path=input_path, - output_dir=output_dir, - save_config=not args.no_config, - verbose=args.verbose, - ) - - print(f"\nConversion complete!") - print(f" Total parameters: {stats['total']}") - print(f" Renamed: {stats['renamed']}") - print(f" Unchanged: {stats['unchanged']}") - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py b/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py deleted file mode 100644 index ccc997457..000000000 --- a/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py +++ /dev/null @@ -1,197 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert TurboDiffusion I2V .pth checkpoint to Diffusers safetensors format. - -TurboDiffusion I2V uses two models: high-noise and low-noise. -This script converts both checkpoints to Diffusers format. - -Usage: - python convert_turbodiffusion_i2v_to_diffusers.py \ - --high_noise_path /path/to/TurboWan2.2-I2V-A14B-high-720P.pth \ - --low_noise_path /path/to/TurboWan2.2-I2V-A14B-low-720P.pth \ - --output_dir /path/to/output \ - --reference_repo Wan-AI/Wan2.1-I2V-14B-720P-Diffusers -""" - -import argparse -import os -import re -import json -import torch -import shutil -import glob -from safetensors import safe_open -from safetensors.torch import save_file -from huggingface_hub import snapshot_download - -# Weight mapping from TurboDiffusion -> Diffusers/FastVideo format -# Same as T2V but may need additional I2V-specific mappings -TURBODIFFUSION_WEIGHT_MAPPING = { - # Self attention - r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2", - r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2", - r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2", - r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2", - r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", - r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", - # Cross attention - r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", - r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", - r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", - r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", - # I2V-specific cross attention (add_k/v_proj for image context) - r"^blocks\.(\d+)\.cross_attn\.add_k\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2", - r"^blocks\.(\d+)\.cross_attn\.add_v\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_add_k\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_add_q\.(.*)$": r"blocks.\1.attn2.norm_added_q.\2", - # Norms and FFN - r"^blocks\.(\d+)\.norm1\.(.*)$": r"blocks.\1.norm1.\2", - r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", - r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.norm3.\2", - r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", - r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", - r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table", - # Embeddings - r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", - r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", - r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", - r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", - r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1", - # Head - r"^head\.head\.(.*)$": r"proj_out.\1", - r"^head\.norm\.(.*)$": r"norm_out.\1", - r"^head\.modulation$": r"scale_shift_table", - # SLA proj_l weights - r"^blocks\.(\d+)\.self_attn\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.attn_impl.proj_l.\2", -} - -SKIP_PATTERNS = [] - - -def should_skip_key(key: str) -> bool: - for pattern in SKIP_PATTERNS: - if re.match(pattern, key): - return True - return False - - -def convert_key(turbo_key: str) -> str: - for pattern, replacement in TURBODIFFUSION_WEIGHT_MAPPING.items(): - if re.match(pattern, turbo_key): - return re.sub(pattern, replacement, turbo_key) - return turbo_key - - -def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch.Tensor: - if len(tensor.shape) == 2 and len(target_shape) == 5: - return tensor.view(target_shape) - return tensor - - -def get_reference_shapes(reference_repo: str) -> dict: - print(f"Downloading reference model shapes from {reference_repo}...") - - local_dir = snapshot_download( - repo_id=reference_repo, - allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"], - local_dir_use_symlinks=False) - - weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors")) - shapes = {} - - for wf in weight_files: - with safe_open(wf, framework="pt") as f: - for key in f.keys(): - shapes[key] = f.get_tensor(key).shape - - return shapes, local_dir - - -def convert_checkpoint(input_path: str, output_dir: str, ref_shapes: dict, model_name: str) -> None: - """Convert a single TurboDiffusion checkpoint to Diffusers format.""" - print(f"\n{'='*60}") - print(f"Converting {model_name}: {input_path}") - print(f"{'='*60}") - - turbo_state_dict = torch.load(input_path, map_location="cpu", weights_only=True) - print(f"Loaded {len(turbo_state_dict)} keys") - - converted_state_dict = {} - skipped_keys = [] - - for turbo_key, tensor in turbo_state_dict.items(): - if should_skip_key(turbo_key): - skipped_keys.append(turbo_key) - continue - - new_key = convert_key(turbo_key) - - if "patch_embedding" in new_key and new_key in ref_shapes: - target_shape = ref_shapes[new_key] - if tensor.shape != target_shape: - print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}") - tensor = reshape_patch_embedding(tensor, target_shape) - - if new_key in ref_shapes: - if tensor.shape != ref_shapes[new_key]: - print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}") - - converted_state_dict[new_key] = tensor - - print(f"Converted: {len(converted_state_dict)} keys, Skipped: {len(skipped_keys)} keys") - - os.makedirs(output_dir, exist_ok=True) - output_path = os.path.join(output_dir, "diffusion_pytorch_model.safetensors") - print(f"Saving to {output_path}...") - save_file(converted_state_dict, output_path) - - return converted_state_dict - - -def main(): - parser = argparse.ArgumentParser(description="Convert TurboDiffusion I2V checkpoints to Diffusers format") - parser.add_argument("--high_noise_path", - type=str, - required=True, - help="Path to high-noise TurboDiffusion .pth checkpoint") - parser.add_argument("--low_noise_path", - type=str, - required=True, - help="Path to low-noise TurboDiffusion .pth checkpoint") - parser.add_argument("--output_dir", type=str, required=True, help="Output directory for converted safetensors") - parser.add_argument("--reference_repo", - type=str, - default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", - help="Reference HF repo to get expected tensor shapes") - - args = parser.parse_args() - - # Get reference shapes - ref_shapes, ref_local_dir = get_reference_shapes(args.reference_repo) - print(f"Got {len(ref_shapes)} reference shapes") - - # Convert high-noise model - high_noise_output = os.path.join(args.output_dir, "transformer_high") - convert_checkpoint(args.high_noise_path, high_noise_output, ref_shapes, "high-noise") - - # Convert low-noise model - low_noise_output = os.path.join(args.output_dir, "transformer_low") - convert_checkpoint(args.low_noise_path, low_noise_output, ref_shapes, "low-noise") - - # Copy config.json to both - src_config = os.path.join(ref_local_dir, "transformer", "config.json") - shutil.copy(src_config, os.path.join(high_noise_output, "config.json")) - shutil.copy(src_config, os.path.join(low_noise_output, "config.json")) - print("Copied config.json to both transformer directories") - - print(f"\n{'='*60}") - print(f"Conversion complete!") - print(f"High-noise model: {high_noise_output}") - print(f"Low-noise model: {low_noise_output}") - print(f"{'='*60}") - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py b/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py deleted file mode 100644 index 72f405a29..000000000 --- a/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py +++ /dev/null @@ -1,204 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert TurboDiffusion .pth checkpoint to Diffusers safetensors format. - -This script: -1. Loads the TurboDiffusion .pth checkpoint -2. Applies weight key renaming to match FastVideo/Diffusers format -3. Reshapes patch_embedding from [D, C*P] to [D, C, P_t, P_h, P_w] -4. Saves as sharded safetensors files compatible with diffusers - -Usage: - python convert_turbodiffusion_to_diffusers.py \ - --input_path /path/to/TurboWan2.1-T2V-1.3B-480P.pth \ - --output_dir /path/to/output/transformer \ - --reference_repo Wan-AI/Wan2.1-T2V-1.3B-Diffusers -""" - -import argparse -import os -import re -import json -import torch -import shutil -import glob -from safetensors import safe_open -from safetensors.torch import save_file -from huggingface_hub import snapshot_download - -# Weight mapping from TurboDiffusion -> Diffusers/FastVideo format -TURBODIFFUSION_WEIGHT_MAPPING = { - # Self attention - r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2", - r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2", - r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2", - r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2", - r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", - r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", - # Cross attention - r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", - r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", - r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", - r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", - # Norms and FFN - r"^blocks\.(\d+)\.norm1\.(.*)$": r"blocks.\1.norm1.\2", - r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", - r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.norm3.\2", - r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", - r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", - r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table", - # Embeddings - DON'T add .proj here! WanVideoArchConfig's param_names_mapping will add it - # patch_embedding.weight stays as patch_embedding.weight (HF format needs this) - r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", - r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", - r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", - r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", - r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1", - # Head - r"^head\.head\.(.*)$": r"proj_out.\1", - r"^head\.norm\.(.*)$": r"norm_out.\1", - r"^head\.modulation$": r"scale_shift_table", - # SLA proj_l weights - include them! They're the distilled attention weights - r"^blocks\.(\d+)\.self_attn\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.attn_impl.proj_l.\2", -} - -# No keys to skip - we want all weights including proj_l -SKIP_PATTERNS = [] - - -def should_skip_key(key: str) -> bool: - """Check if a key should be skipped (SLA-specific weights).""" - for pattern in SKIP_PATTERNS: - if re.match(pattern, key): - return True - return False - - -def convert_key(turbo_key: str) -> str: - """Convert TurboDiffusion key to Diffusers format.""" - for pattern, replacement in TURBODIFFUSION_WEIGHT_MAPPING.items(): - if re.match(pattern, turbo_key): - return re.sub(pattern, replacement, turbo_key) - return turbo_key # Return unchanged if no pattern matches - - -def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch.Tensor: - """Reshape patch_embedding from [D, C*P_t*P_h*P_w] to [D, C, P_t, P_h, P_w].""" - if len(tensor.shape) == 2 and len(target_shape) == 5: - return tensor.view(target_shape) - return tensor - - -def get_reference_shapes(reference_repo: str) -> dict: - """Download reference model and get expected shapes for patch_embedding.""" - print(f"Downloading reference model shapes from {reference_repo}...") - - # Download just the transformer config and a weight file to get shapes - local_dir = snapshot_download( - repo_id=reference_repo, - allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"], - local_dir_use_symlinks=False) - - # Load the first safetensors file to get shapes - weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors")) - shapes = {} - - for wf in weight_files: - with safe_open(wf, framework="pt") as f: - for key in f.keys(): - shapes[key] = f.get_tensor(key).shape - - return shapes - - -def main(): - parser = argparse.ArgumentParser(description="Convert TurboDiffusion checkpoint to Diffusers format") - parser.add_argument("--input_path", type=str, required=True, help="Path to TurboDiffusion .pth checkpoint") - parser.add_argument("--output_dir", type=str, required=True, help="Output directory for converted safetensors") - parser.add_argument("--reference_repo", - type=str, - default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - help="Reference HF repo to get expected tensor shapes") - parser.add_argument("--skip_sla_weights", - action="store_true", - default=False, - help="Skip SLA-specific weights (proj_l) that aren't in base model") - - args = parser.parse_args() - - # Load TurboDiffusion checkpoint - print(f"Loading TurboDiffusion checkpoint from {args.input_path}...") - turbo_state_dict = torch.load(args.input_path, map_location="cpu", weights_only=True) - print(f"Loaded {len(turbo_state_dict)} keys") - - # Get reference shapes for reshaping - ref_shapes = get_reference_shapes(args.reference_repo) - print(f"Got {len(ref_shapes)} reference shapes") - - # Convert keys and reshape tensors - converted_state_dict = {} - skipped_keys = [] - - for turbo_key, tensor in turbo_state_dict.items(): - # Skip SLA-specific weights if requested - if args.skip_sla_weights and should_skip_key(turbo_key): - skipped_keys.append(turbo_key) - continue - - # Convert key name - new_key = convert_key(turbo_key) - - # Reshape patch_embedding if needed - if "patch_embedding" in new_key and new_key in ref_shapes: - target_shape = ref_shapes[new_key] - if tensor.shape != target_shape: - print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}") - tensor = reshape_patch_embedding(tensor, target_shape) - - # Verify shape matches reference if available - if new_key in ref_shapes: - if tensor.shape != ref_shapes[new_key]: - print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}") - - converted_state_dict[new_key] = tensor - - print(f"\nConversion summary:") - print(f" Converted: {len(converted_state_dict)} keys") - print(f" Skipped (SLA): {len(skipped_keys)} keys") - - if skipped_keys: - print(f"\nSkipped SLA keys (first 5):") - for k in skipped_keys[:5]: - print(f" - {k}") - - # Create output directory - os.makedirs(args.output_dir, exist_ok=True) - - # Save as safetensors - output_path = os.path.join(args.output_dir, "diffusion_pytorch_model.safetensors") - print(f"\nSaving to {output_path}...") - save_file(converted_state_dict, output_path) - - # Copy config.json from reference - ref_local = snapshot_download(repo_id=args.reference_repo, - allow_patterns=["transformer/config.json"], - local_dir_use_symlinks=False) - src_config = os.path.join(ref_local, "transformer", "config.json") - dst_config = os.path.join(args.output_dir, "config.json") - shutil.copy(src_config, dst_config) - print(f"Copied config.json") - - print(f"\nDone! Converted weights saved to: {args.output_dir}") - print(f"\nNext steps:") - print(f" 1. Use create_hf_repo.py to create a complete diffusers repo:") - print(f" python scripts/checkpoint_conversion/create_hf_repo.py \\") - print(f" --repo_id Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\") - print(f" --local_dir /tmp/turbodiffusion-wan \\") - print(f" --checkpoint_dir {args.output_dir} \\") - print(f" --push_to_hub --upload_repo_id YOUR_USERNAME/TurboWan-Diffusers") - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py b/scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py deleted file mode 100755 index 5f3206c4e..000000000 --- a/scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -"""Convert DreamX-World-5B autoregressive weights to FastVideo layout. - -The HF repository stores one raw official ``model.safetensors`` whose keys match -FastVideo's native ``DreamXWorldARTransformer3DModel``. The converter writes a -Diffusers-like root with ``transformer/config.json`` and reusable Wan2.2 -components. Use ``--symlink-transformer`` locally to avoid duplicating the 21GB -AR tensor file. -""" - -from __future__ import annotations - -import argparse -import json -import shutil -from pathlib import Path - -TRANSFORMER_CONFIG: dict[str, object] = { - "_class_name": "DreamXWorldARTransformer3DModel", - "model_type": "ti2v", - "patch_size": [1, 2, 2], - "text_len": 512, - "num_attention_heads": 24, - "attention_head_dim": 128, - "in_channels": 48, - "out_channels": 48, - "text_dim": 4096, - "freq_dim": 256, - "ffn_dim": 14336, - "num_layers": 30, - "local_attn_size": 12, - "sink_size": 3, - "cross_attn_norm": True, - "qk_norm": True, - "eps": 1e-6, - "add_control_adapter": True, - "cam_method": "prope", - "attn_compress": 4, - "cam_self_attn_layers": list(range(30)), - "num_frames_per_block": 3, -} - -MODEL_INDEX: dict[str, object] = { - "_class_name": "DreamXWorldARPipeline", - "_diffusers_version": "0.31.0", - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["transformers", "UMT5EncoderModel"], - "tokenizer": ["transformers", "AutoTokenizer"], - "transformer": ["diffusers", "DreamXWorldARTransformer3DModel"], - "vae": ["diffusers", "AutoencoderKLWan"], -} - -REUSED_COMPONENTS = ("scheduler", "text_encoder", "tokenizer", "vae") - - -def _source_safetensors(source: Path) -> Path: - if source.is_file(): - return source - path = source / "model.safetensors" - if not path.exists(): - raise FileNotFoundError(f"Missing AR model.safetensors under {source}") - return path - - -def convert_transformer(source: Path, output: Path, symlink_transformer: bool) -> None: - src = _source_safetensors(source) - transformer_dir = output / "transformer" - transformer_dir.mkdir(parents=True, exist_ok=True) - (transformer_dir / "config.json").write_text(json.dumps(TRANSFORMER_CONFIG, indent=2) + "\n") - dst = transformer_dir / "model.safetensors" - if dst.is_symlink() and not dst.exists(): - dst.unlink() # dangling symlink: replace so re-conversion self-heals - if dst.exists(): - return - if symlink_transformer: - dst.symlink_to(src.resolve()) - else: - shutil.copy2(src, dst) - - -def _copy_or_link_component(component: str, component_source: Path, output: Path, symlink: bool) -> None: - src = component_source / component - dst = output / component - if not src.exists(): - raise FileNotFoundError(f"Missing reused component source: {src}") - if dst.is_symlink() and not dst.exists(): - dst.unlink() # dangling symlink: replace so re-conversion self-heals - if dst.exists(): - return - if symlink: - dst.symlink_to(src.resolve(), target_is_directory=src.is_dir()) - elif src.is_dir(): - shutil.copytree(src, dst) - else: - shutil.copy2(src, dst) - - -def write_model_index(output: Path, component_source: Path | None, symlink_components: bool) -> None: - if component_source is not None: - for component in REUSED_COMPONENTS: - _copy_or_link_component(component, component_source, output, symlink_components) - missing = [component for component in REUSED_COMPONENTS if not (output / component).exists()] - if missing: - print("skipping model_index.json; missing reused components: " + ", ".join(missing)) - return - (output / "model_index.json").write_text(json.dumps(MODEL_INDEX, indent=2) + "\n") - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--source", type=Path, required=True) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--component-source", type=Path) - parser.add_argument("--symlink-components", action="store_true") - parser.add_argument("--symlink-transformer", action="store_true") - args = parser.parse_args() - - convert_transformer(args.source, args.output, args.symlink_transformer) - write_model_index(args.output, args.component_source, args.symlink_components) - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py deleted file mode 100644 index 449f51afe..000000000 --- a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py +++ /dev/null @@ -1,226 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -"""Convert DreamX-World-5B-Cam raw transformer weights to FastVideo-loadable format. - -The GD-ML/DreamX-World-5B-Cam repository stores the transformer as raw -DreamX/Wan official shards. FastVideo's TransformerLoader expects a Diffusers-like -transformer folder with a config.json and safetensors whose keys can be mapped by -WanVideoConfig.param_names_mapping. This script performs the raw official -> -Diffusers-like key rename and writes the DreamX 5B-Cam transformer config. - -Example: - python scripts/checkpoint_conversion/dreamx_world_to_diffusers.py \ - --source official_weights/dreamx_world \ - --output converted_weights/dreamx_world -""" - -from __future__ import annotations - -import argparse -import json -import re -import shutil -from collections import OrderedDict -from pathlib import Path - -import torch -from huggingface_hub import save_torch_state_dict -from safetensors import safe_open -from safetensors.torch import load_file - -OFFICIAL_TO_DIFFUSERS_MAPPING: dict[str, str] = { - r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1", - r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1", - r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1", - r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1", - r"^time_projection\.1\.(.*)$": r"condition_embedder.time_proj.\1", - r"^img_emb\.proj\.0\.(.*)$": r"condition_embedder.image_embedder.norm1.\1", - r"^img_emb\.proj\.1\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1", - r"^img_emb\.proj\.3\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1", - r"^img_emb\.proj\.4\.(.*)$": r"condition_embedder.image_embedder.norm2.\1", - r"^head\.modulation": r"scale_shift_table", - r"^head\.head\.(.*)$": r"proj_out.\1", - r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", - r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", - r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", - r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2", - r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2", - r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2", - r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", - r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", - r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2", - r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", - r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2", - r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2", - r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2", - r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.net.2.\2", - r"^blocks\.(\d+)\.modulation": r"blocks.\1.scale_shift_table", - r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm2.\2", -} - -TRANSFORMER_CONFIG: dict[str, object] = { - "_class_name": "DreamXWorldTransformer3DModel", - "patch_size": [1, 2, 2], - "text_len": 512, - "num_attention_heads": 24, - "attention_head_dim": 128, - "in_channels": 48, - "out_channels": 48, - "text_dim": 4096, - "freq_dim": 256, - "ffn_dim": 14336, - "num_layers": 30, - "cross_attn_norm": True, - "qk_norm": "rms_norm_across_heads", - "eps": 1e-6, - "image_dim": None, - "added_kv_proj_dim": None, - "rope_max_seq_len": 1024, - "add_control_adapter": True, - "cam_method": "prope", - "attn_compress": 1, - "cam_self_attn_layers": None, -} - -MODEL_INDEX: dict[str, object] = { - "_class_name": "DreamXWorldPipeline", - "_diffusers_version": "0.31.0", - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["transformers", "UMT5EncoderModel"], - "tokenizer": ["transformers", "AutoTokenizer"], - "transformer": ["diffusers", "DreamXWorldTransformer3DModel"], - "vae": ["diffusers", "AutoencoderKLWan"], -} - -REUSED_COMPONENTS = ("scheduler", "text_encoder", "tokenizer", "vae") - - -def map_transformer_key(key: str) -> str: - for pattern, replacement in OFFICIAL_TO_DIFFUSERS_MAPPING.items(): - if re.match(pattern, key): - return re.sub(pattern, replacement, key) - return key - - -def _safetensor_files(source: Path) -> list[Path]: - if source.is_file(): - if source.suffix != ".safetensors": - raise ValueError(f"Only .safetensors files are supported, got {source}") - return [source] - - index_path = source / "diffusion_pytorch_model.safetensors.index.json" - if index_path.exists(): - index = json.loads(index_path.read_text()) - return sorted({source / shard for shard in index["weight_map"].values()}) - - files = sorted(source.glob("*.safetensors")) - if not files: - raise FileNotFoundError(f"No safetensors files found under {source}") - return files - - -def convert_transformer(source: Path, output: Path, max_shard_size: str) -> None: - transformer_dir = output / "transformer" - transformer_dir.mkdir(parents=True, exist_ok=True) - (transformer_dir / "config.json").write_text(json.dumps(TRANSFORMER_CONFIG, indent=2) + "\n") - - converted: OrderedDict[str, torch.Tensor] = OrderedDict() - for shard in _safetensor_files(source): - print(f"loading {shard}") - for key, tensor in load_file(shard, device="cpu").items(): - new_key = map_transformer_key(key) - if new_key in converted: - raise ValueError(f"Duplicate converted key: {new_key}") - converted[new_key] = tensor - - print(f"saving {len(converted)} tensors to {transformer_dir}") - save_torch_state_dict(converted, transformer_dir, max_shard_size=max_shard_size) - - -def _copy_or_link_component(component: str, component_source: Path, output: Path, symlink: bool) -> None: - src = component_source / component - dst = output / component - if not src.exists(): - raise FileNotFoundError(f"Missing reused component source: {src}") - if dst.is_symlink() and not dst.exists(): - dst.unlink() # dangling symlink: replace so re-conversion self-heals - if dst.exists(): - print(f"keeping existing {dst}") - return - if symlink: - dst.symlink_to(src.resolve(), target_is_directory=src.is_dir()) - print(f"linked {dst} -> {src}") - elif src.is_dir(): - shutil.copytree(src, dst) - print(f"copied {src} -> {dst}") - else: - shutil.copy2(src, dst) - print(f"copied {src} -> {dst}") - - -def write_model_index(output: Path, component_source: Path | None, symlink_components: bool) -> None: - if component_source is not None: - for component in REUSED_COMPONENTS: - _copy_or_link_component(component, component_source, output, symlink_components) - missing = [component for component in REUSED_COMPONENTS if not (output / component).exists()] - if missing: - print("skipping model_index.json; missing reused components: " + ", ".join(missing)) - print("pass --component-source to copy or link reused components") - return - (output / "model_index.json").write_text(json.dumps(MODEL_INDEX, indent=2) + "\n") - print(f"wrote {output / 'model_index.json'}") - - -def analyze(source: Path) -> None: - total = 0 - unchanged = 0 - examples: list[tuple[str, str]] = [] - for shard in _safetensor_files(source): - with safe_open(shard, framework="pt", device="cpu") as tensors: - for key in tensors: - total += 1 - new_key = map_transformer_key(key) - unchanged += int(new_key == key) - if len(examples) < 20 and new_key != key: - examples.append((key, new_key)) - print(f"total_keys={total} unchanged_keys={unchanged}") - for old, new in examples: - print(f"{old} -> {new}") - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--source", - type=Path, - required=True, - help="DreamX raw transformer directory or safetensors file") - parser.add_argument("--output", - type=Path, - required=True, - help="Output model root; transformer/ is created inside it") - parser.add_argument("--max-shard-size", default="10GB") - parser.add_argument( - "--component-source", - type=Path, - help="Optional Wan2.2 Diffusers root whose scheduler/text_encoder/tokenizer/vae components are reused.", - ) - parser.add_argument( - "--symlink-components", - action="store_true", - help="Symlink reused components from --component-source instead of copying them.", - ) - parser.add_argument("--analyze", action="store_true", help="Only print key mapping summary") - args = parser.parse_args() - - if args.analyze: - analyze(args.source) - else: - convert_transformer(args.source, args.output, args.max_shard_size) - write_model_index(args.output, args.component_source, args.symlink_components) - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/extract_llava_text_encoder.py b/scripts/checkpoint_conversion/extract_llava_text_encoder.py deleted file mode 100644 index b202aec4c..000000000 --- a/scripts/checkpoint_conversion/extract_llava_text_encoder.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Extract LlamaModel text encoder from LLaVA weights for GameCraft/HunyuanVideo. - -LLaVA stores the text encoder as `language_model.model.XXX`, but FastVideo/HunyuanVideo -expects just `XXX` (e.g., `layers.0.self_attn.q_proj.weight` instead of -`language_model.model.layers.0.self_attn.q_proj.weight`). - -This script extracts just the LlamaModel weights and renames them. - -Usage: - python scripts/checkpoint_conversion/extract_llava_text_encoder.py \ - --input Hunyuan-GameCraft-1.0/weights/stdmodels/llava-llama-3-8b-v1_1-transformers \ - --output official_weights/hunyuan-gamecraft-diffusers/text_encoder -""" -from __future__ import annotations - -import argparse -import json -from pathlib import Path -from collections import OrderedDict - -import torch -from safetensors import safe_open -from safetensors.torch import save_file - - -def extract_text_encoder( - input_dir: Path, - output_dir: Path, -) -> dict: - """Extract LlamaModel from LLaVA weights.""" - - # Find all safetensor shards - shard_files = sorted(input_dir.glob("model-*.safetensors")) - if not shard_files: - raise FileNotFoundError(f"No safetensor files found in {input_dir}") - - print(f"Found {len(shard_files)} weight shards") - - # Collect all weights, filtering for language_model.model prefix - extracted_weights: OrderedDict[str, torch.Tensor] = OrderedDict() - - for shard_file in shard_files: - print(f"Processing {shard_file.name}...") - with safe_open(str(shard_file), framework="pt") as f: - for key in f.keys(): - # Extract only language_model.model weights - if key.startswith("language_model.model."): - # Strip the prefix - new_key = key.replace("language_model.model.", "") - extracted_weights[new_key] = f.get_tensor(key) - - print(f"Extracted {len(extracted_weights)} weights") - - # Save to output directory - output_dir.mkdir(parents=True, exist_ok=True) - - # Save as single file (it should be <10GB which is fine for safetensors) - output_file = output_dir / "model.safetensors" - print(f"Saving to {output_file}...") - save_file(extracted_weights, str(output_file)) - - # Create config.json for LlamaModel - config = { - "_name_or_path": "llava-llama-3-8b-v1_1-text_encoder", - "architectures": ["LlamaModel"], - "attention_bias": False, - "attention_dropout": 0.0, - "bos_token_id": 128000, - "eos_token_id": 128001, - "head_dim": 128, - "hidden_act": "silu", - "hidden_size": 4096, - "initializer_range": 0.02, - "intermediate_size": 14336, - "max_position_embeddings": 8192, - "mlp_bias": False, - "model_type": "llama", - "num_attention_heads": 32, - "num_hidden_layers": 32, - "num_key_value_heads": 8, - "pretraining_tp": 1, - "rms_norm_eps": 1e-05, - "rope_scaling": None, - "rope_theta": 500000.0, - "tie_word_embeddings": False, - "torch_dtype": "float16", - "transformers_version": "4.48.0", - "use_cache": True, - "vocab_size": 128320 - } - - with open(output_dir / "config.json", "w") as f: - json.dump(config, f, indent=2) - print(f"Saved config.json") - - # Copy tokenizer files if they exist - tokenizer_files = [ - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - ] - for fname in tokenizer_files: - src = input_dir / fname - if src.exists(): - import shutil - shutil.copy2(src, output_dir / fname) - print(f"Copied {fname}") - - return {"total": len(extracted_weights)} - - -def extract_clip_text_encoder( - input_dir: Path, - output_dir: Path, -) -> dict: - """Extract CLIPTextModel from full CLIP model weights.""" - - # Find weight file - weight_files = list(input_dir.glob("model.safetensors")) + \ - list(input_dir.glob("pytorch_model.bin")) - - if not weight_files: - raise FileNotFoundError(f"No weight files found in {input_dir}") - - weight_file = weight_files[0] - print(f"Loading {weight_file}...") - - if weight_file.suffix == ".safetensors": - from safetensors.torch import load_file - state_dict = load_file(str(weight_file)) - else: - state_dict = torch.load(weight_file, map_location="cpu") - - # Extract only text_model weights - extracted_weights: OrderedDict[str, torch.Tensor] = OrderedDict() - - for key, value in state_dict.items(): - # CLIP full model has text_model.XXX, we want just XXX - if key.startswith("text_model."): - new_key = key.replace("text_model.", "") - extracted_weights[new_key] = value - - if not extracted_weights: - # Maybe it's already just the text model - extracted_weights = OrderedDict(state_dict) - - print(f"Extracted {len(extracted_weights)} weights") - - # Save to output directory - output_dir.mkdir(parents=True, exist_ok=True) - - output_file = output_dir / "model.safetensors" - print(f"Saving to {output_file}...") - save_file(extracted_weights, str(output_file)) - - # Create config.json for CLIPTextModel - config = { - "_name_or_path": "openai/clip-vit-large-patch14", - "architectures": ["CLIPTextModel"], - "attention_dropout": 0.0, - "bos_token_id": 0, - "dropout": 0.0, - "eos_token_id": 2, - "hidden_act": "quick_gelu", - "hidden_size": 768, - "initializer_factor": 1.0, - "initializer_range": 0.02, - "intermediate_size": 3072, - "layer_norm_eps": 1e-05, - "max_position_embeddings": 77, - "model_type": "clip_text_model", - "num_attention_heads": 12, - "num_hidden_layers": 12, - "pad_token_id": 1, - "projection_dim": 768, - "torch_dtype": "float16", - "transformers_version": "4.48.0", - "vocab_size": 49408 - } - - with open(output_dir / "config.json", "w") as f: - json.dump(config, f, indent=2) - print(f"Saved config.json") - - # Copy tokenizer files if they exist - tokenizer_files = [ - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - "vocab.json", - "merges.txt", - ] - for fname in tokenizer_files: - src = input_dir / fname - if src.exists(): - import shutil - shutil.copy2(src, output_dir / fname) - print(f"Copied {fname}") - - return {"total": len(extracted_weights)} - - -def main(): - parser = argparse.ArgumentParser(description="Extract text encoder from LLaVA or CLIP model weights.") - parser.add_argument("--input", type=str, required=True, help="Input directory containing model weights") - parser.add_argument("--output", type=str, required=True, help="Output directory for extracted weights") - parser.add_argument("--type", - type=str, - choices=["llama", "clip"], - default="llama", - help="Type of text encoder to extract") - - args = parser.parse_args() - - input_dir = Path(args.input) - output_dir = Path(args.output) - - if not input_dir.exists(): - raise FileNotFoundError(f"Input directory not found: {input_dir}") - - if args.type == "llama": - stats = extract_text_encoder(input_dir, output_dir) - else: - stats = extract_clip_text_encoder(input_dir, output_dir) - - print(f"\nExtraction complete! Total weights: {stats['total']}") - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/longcat_to_fastvideo.py b/scripts/checkpoint_conversion/longcat_to_fastvideo.py deleted file mode 100644 index 19a80601c..000000000 --- a/scripts/checkpoint_conversion/longcat_to_fastvideo.py +++ /dev/null @@ -1,706 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -""" -Convert LongCat weights to FastVideo native format. - -This script performs a complete conversion from original LongCat weights -to FastVideo native implementation in a single step: - -1. Converts transformer weights (with QKV/KV splitting) -2. Copies other components (VAE, text encoder, tokenizer, scheduler) -3. Converts LoRA weights (cfg_step_lora, refinement_lora) -4. Updates config files to point to native model - -Usage: - python scripts/checkpoint_conversion/longcat_to_fastvideo.py \ - --source /path/to/LongCat-Video/weights/LongCat-Video \ - --output weights/longcat-native \ - --validate -""" - -import argparse -import glob -import json -import re -import shutil -from pathlib import Path -from collections import OrderedDict - -import torch -from safetensors.torch import load_file, save_file -from tqdm import tqdm - - -def split_qkv(qkv_weight: torch.Tensor, qkv_bias: torch.Tensor | None = None): - """Split fused QKV projection into separate Q, K, V.""" - dim = qkv_weight.shape[0] // 3 - q, k, v = torch.chunk(qkv_weight, 3, dim=0) - - if qkv_bias is not None: - q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0) - else: - q_bias = k_bias = v_bias = None - - return (q, k, v), (q_bias, k_bias, v_bias) - - -def split_kv(kv_weight: torch.Tensor, kv_bias: torch.Tensor | None = None): - """Split fused KV projection into separate K, V.""" - dim = kv_weight.shape[0] // 2 - k, v = torch.chunk(kv_weight, 2, dim=0) - - if kv_bias is not None: - k_bias, v_bias = torch.chunk(kv_bias, 2, dim=0) - else: - k_bias = v_bias = None - - return (k, v), (k_bias, v_bias) - - -def convert_transformer_weights(source_weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Convert LongCat transformer weights to native FastVideo format. - - Main transformations: - 1. Split fused QKV projections (self-attention) - 2. Split fused KV projections (cross-attention) - 3. Rename parameters according to mapping - """ - converted = OrderedDict() - processed_keys = set() - - print(" Converting transformer weights...") - - for key, value in tqdm(source_weights.items(), desc=" Processing parameters"): - if key in processed_keys: - continue - - # === Embedders === - if key.startswith("x_embedder."): - new_key = key.replace("x_embedder.", "patch_embed.") - converted[new_key] = value - - elif key.startswith("t_embedder.mlp.0."): - new_key = key.replace("t_embedder.mlp.0.", "time_embedder.linear_1.") - converted[new_key] = value - - elif key.startswith("t_embedder.mlp.2."): - new_key = key.replace("t_embedder.mlp.2.", "time_embedder.linear_2.") - converted[new_key] = value - - elif key.startswith("y_embedder.y_proj.0."): - new_key = key.replace("y_embedder.y_proj.0.", "caption_embedder.linear_1.") - converted[new_key] = value - - elif key.startswith("y_embedder.y_proj.2."): - new_key = key.replace("y_embedder.y_proj.2.", "caption_embedder.linear_2.") - converted[new_key] = value - - # === Self-Attention QKV Splitting === - elif ".attn.qkv.weight" in key: - block_idx = key.split(".")[1] - qkv_weight = value - qkv_bias_key = key.replace(".weight", ".bias") - qkv_bias = source_weights.get(qkv_bias_key) - - (q, k, v), (q_bias, k_bias, v_bias) = split_qkv(qkv_weight, qkv_bias) - - converted[f"blocks.{block_idx}.self_attn.to_q.weight"] = q - converted[f"blocks.{block_idx}.self_attn.to_k.weight"] = k - converted[f"blocks.{block_idx}.self_attn.to_v.weight"] = v - - if q_bias is not None: - converted[f"blocks.{block_idx}.self_attn.to_q.bias"] = q_bias - converted[f"blocks.{block_idx}.self_attn.to_k.bias"] = k_bias - converted[f"blocks.{block_idx}.self_attn.to_v.bias"] = v_bias - - processed_keys.add(key) - if qkv_bias is not None: - processed_keys.add(qkv_bias_key) - - elif ".attn.qkv.bias" in key: - continue - - elif ".attn.proj." in key: - new_key = key.replace(".attn.proj.", ".self_attn.to_out.") - converted[new_key] = value - - elif ".attn.q_norm." in key or ".attn.k_norm." in key: - new_key = key.replace(".attn.", ".self_attn.") - converted[new_key] = value - - # === Cross-Attention === - elif ".cross_attn.q_linear." in key: - new_key = key.replace(".cross_attn.q_linear.", ".cross_attn.to_q.") - converted[new_key] = value - - elif ".cross_attn.kv_linear.weight" in key: - block_idx = key.split(".")[1] - kv_weight = value - kv_bias_key = key.replace(".weight", ".bias") - kv_bias = source_weights.get(kv_bias_key) - - (k, v), (k_bias, v_bias) = split_kv(kv_weight, kv_bias) - - converted[f"blocks.{block_idx}.cross_attn.to_k.weight"] = k - converted[f"blocks.{block_idx}.cross_attn.to_v.weight"] = v - - if k_bias is not None: - converted[f"blocks.{block_idx}.cross_attn.to_k.bias"] = k_bias - converted[f"blocks.{block_idx}.cross_attn.to_v.bias"] = v_bias - - processed_keys.add(key) - if kv_bias is not None: - processed_keys.add(kv_bias_key) - - elif ".cross_attn.kv_linear.bias" in key: - continue - - elif ".cross_attn.proj." in key: - new_key = key.replace(".cross_attn.proj.", ".cross_attn.to_out.") - converted[new_key] = value - - elif ".cross_attn.q_norm." in key or ".cross_attn.k_norm." in key: - converted[key] = value - - # === Final Layer (must come BEFORE general transformer block patterns) === - elif key.startswith("final_layer.adaLN_modulation.1."): - new_key = key.replace("final_layer.adaLN_modulation.1.", "final_layer.adaln_linear.") - converted[new_key] = value - - # === Transformer Block AdaLN === - elif ".adaLN_modulation.1." in key: - new_key = key.replace(".adaLN_modulation.1.", ".adaln_linear_1.") - converted[new_key] = value - - # === Transformer Block Normalization === - elif ".mod_norm_attn." in key or ".mod_norm_ffn." in key: - continue - - elif ".pre_crs_attn_norm.weight" in key: - new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.") - converted[new_key] = value - - elif ".pre_crs_attn_norm.bias" in key: - new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.") - converted[new_key] = value - - # === FFN (SwiGLU) === - elif ".ffn.w1." in key or ".ffn.w2." in key or ".ffn.w3." in key: - converted[key] = value - - elif key.startswith("final_layer.norm_final."): - continue - - elif key.startswith("final_layer.linear."): - new_key = key.replace("final_layer.linear.", "final_layer.proj.") - converted[new_key] = value - - else: - print(f" ⚠️ Unknown key: {key}") - converted[key] = value - - return converted - - -def validate_conversion(original: dict, converted: dict) -> bool: - """Validate that conversion preserved all parameters correctly.""" - print("\n Validating conversion...") - - orig_count = sum(p.numel() for p in original.values()) - conv_count = sum(p.numel() for p in converted.values()) - - dropped_count = 0 - for key, value in original.items(): - if ".mod_norm_attn." in key or ".mod_norm_ffn." in key: - dropped_count += value.numel() - elif "final_layer.norm_final." in key: - dropped_count += value.numel() - - expected_conv_count = orig_count - dropped_count - - print(f" Original parameters: {orig_count:,}") - print(f" Converted parameters: {conv_count:,}") - print(f" Dropped parameters (norms without params): {dropped_count:,}") - - if conv_count != expected_conv_count: - print(f" ⚠️ Parameter count mismatch!") - return False - - print(f" ✓ Parameter count matches") - - # Verify QKV/KV splits - print("\n Verifying QKV/KV splits...") - num_blocks = 48 - - for i in range(num_blocks): - orig_qkv_weight = original.get(f"blocks.{i}.attn.qkv.weight") - if orig_qkv_weight is not None: - conv_q = converted[f"blocks.{i}.self_attn.to_q.weight"] - conv_k = converted[f"blocks.{i}.self_attn.to_k.weight"] - conv_v = converted[f"blocks.{i}.self_attn.to_v.weight"] - reconstructed = torch.cat([conv_q, conv_k, conv_v], dim=0) - if not torch.allclose(orig_qkv_weight, reconstructed): - print(f" ❌ QKV weight mismatch in block {i}") - return False - - orig_kv_weight = original.get(f"blocks.{i}.cross_attn.kv_linear.weight") - if orig_kv_weight is not None: - conv_k = converted[f"blocks.{i}.cross_attn.to_k.weight"] - conv_v = converted[f"blocks.{i}.cross_attn.to_v.weight"] - reconstructed = torch.cat([conv_k, conv_v], dim=0) - if not torch.allclose(orig_kv_weight, reconstructed): - print(f" ❌ KV weight mismatch in block {i}") - return False - - print(f" ✓ All splits verified successfully") - return True - - -def copy_component(source_dir: Path, output_dir: Path, component: str, mapping: dict = None) -> bool: - """Copy a component directory, optionally with name mapping.""" - source_name = mapping.get(component, component) if mapping else component - source_path = source_dir / source_name - - if source_path.exists(): - output_path = output_dir / component - if output_path.exists(): - shutil.rmtree(output_path) - shutil.copytree(source_path, output_path) - print(f" ✓ {component} copied") - return True - else: - print(f" ⚠️ {component} not found, skipping") - return False - - -def create_model_index(): - """Create model_index.json for FastVideo native model.""" - return { - "_class_name": "LongCatPipeline", - "_diffusers_version": "0.32.0", - "workload_type": "video-generation", - "tokenizer": ["transformers", "AutoTokenizer"], - "text_encoder": ["transformers", "UMT5EncoderModel"], - "vae": ["diffusers", "AutoencoderKLWan"], - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "transformer": ["diffusers", "LongCatTransformer3DModel"] # Native model - } - - -def update_transformer_config(transformer_dir: Path): - """Update transformer config.json to point to native model.""" - config_path = transformer_dir / "config.json" - if not config_path.exists(): - print(" ⚠️ Transformer config not found, skipping") - return - - with open(config_path, 'r') as f: - config = json.load(f) - - if '_class_name' in config: - old_class = config['_class_name'] - config['_class_name'] = 'LongCatTransformer3DModel' - print(f" Updated _class_name: {old_class} → LongCatTransformer3DModel") - else: - config['_class_name'] = 'LongCatTransformer3DModel' - print(f" Added _class_name: LongCatTransformer3DModel") - - # Fix num_heads -> num_attention_heads for FastVideo compatibility - if 'num_heads' in config and 'num_attention_heads' not in config: - config['num_attention_heads'] = config.pop('num_heads') - print(f" Updated num_heads → num_attention_heads") - - with open(config_path, 'w') as f: - json.dump(config, f, indent=2) - - print(" ✓ Transformer config updated") - - -# ============================================================================ -# LoRA Conversion Functions -# ============================================================================ - - -def parse_lora_key(key: str) -> tuple[str, str]: - """Parse LongCat LoRA key into module path and weight type.""" - if key.startswith("lora___lorahyphen___"): - key = key[len("lora___lorahyphen___"):] - - key = key.replace("___lorahyphen___", ".") - - if ".lora_down.weight" in key: - return key.replace(".lora_down.weight", ""), "lora_down.weight" - elif ".lora_up.weight" in key: - return key.replace(".lora_up.weight", ""), "lora_up.weight" - elif ".lora_up.blocks." in key: - match = re.match(r"(.+)\.lora_up\.blocks\.(\d+)\.weight", key) - if match: - return match.group(1), f"lora_up.blocks.{match.group(2)}.weight" - elif ".alpha_scale" in key: - return key.replace(".alpha_scale", ""), "alpha_scale" - elif ".lora_alpha" in key: - return key.replace(".lora_alpha", ""), "lora_alpha" - - raise ValueError(f"Unknown LoRA key format: {key}") - - -def map_lora_module(module_path: str) -> list[tuple[str, str]]: - """Map LongCat module path to FastVideo paths. Returns [(path, component)].""" - # Self-attention QKV → Q, K, V - match = re.match(r"blocks\.(\d+)\.attn\.qkv", module_path) - if match: - b = match.group(1) - return [(f"blocks.{b}.self_attn.to_q", "q"), (f"blocks.{b}.self_attn.to_k", "k"), - (f"blocks.{b}.self_attn.to_v", "v")] - - # Self-attention output - match = re.match(r"blocks\.(\d+)\.attn\.proj", module_path) - if match: - return [(f"blocks.{match.group(1)}.self_attn.to_out", "single")] - - # Cross-attention Q - match = re.match(r"blocks\.(\d+)\.cross_attn\.q_linear", module_path) - if match: - return [(f"blocks.{match.group(1)}.cross_attn.to_q", "single")] - - # Cross-attention KV → K, V - match = re.match(r"blocks\.(\d+)\.cross_attn\.kv_linear", module_path) - if match: - b = match.group(1) - return [(f"blocks.{b}.cross_attn.to_k", "k"), (f"blocks.{b}.cross_attn.to_v", "v")] - - # FFN - match = re.match(r"blocks\.(\d+)\.ffn\.(w[123])", module_path) - if match: - return [(f"blocks.{match.group(1)}.ffn.{match.group(2)}", "single")] - - # AdaLN modulation - match = re.match(r"blocks\.(\d+)\.adaLN_modulation\.1", module_path) - if match: - return [(f"blocks.{match.group(1)}.adaln_linear_1", "single")] - - # Final layer - if module_path == "final_layer.adaLN_modulation.1": - return [("final_layer.adaln_linear", "single")] - if module_path == "final_layer.linear": - return [("final_layer.proj", "single")] - - raise ValueError(f"Unknown LoRA module: {module_path}") - - -def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str) -> dict[str, torch.Tensor]: - """Convert LongCat LoRA to FastVideo format.""" - print(f" Converting {lora_name}...") - print(f" Source keys: {len(source_weights)}") - - converted = OrderedDict() - - # Group by module - modules = {} - for key in source_weights.keys(): - try: - module_path, weight_type = parse_lora_key(key) - if module_path not in modules: - modules[module_path] = {} - modules[module_path][weight_type] = key - except ValueError: - continue - - # Process each module - for module_path, weight_keys in modules.items(): - try: - targets = map_lora_module(module_path) - except ValueError: - continue - - # Get alpha_scale if present (defaults to 1.0 if missing) - alpha_scale = 1.0 - if "alpha_scale" in weight_keys: - alpha_scale_tensor = source_weights[weight_keys["alpha_scale"]] - alpha_scale = alpha_scale_tensor.item() if alpha_scale_tensor.numel() == 1 else float( - alpha_scale_tensor.mean()) - - # Handle lora_down (lora_A) - if "lora_down.weight" in weight_keys: - lora_down = source_weights[weight_keys["lora_down.weight"]] - - if len(targets) == 1: - converted[f"{targets[0][0]}.lora_A"] = lora_down - # Compute alpha from alpha_scale and rank - rank = lora_down.shape[0] - alpha = alpha_scale * rank - converted[f"{targets[0][0]}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32) - else: - # Split for fused projections - n = len(targets) - rank = lora_down.shape[0] // n - for i, (path, _) in enumerate(targets): - converted[f"{path}.lora_A"] = lora_down[i * rank:(i + 1) * rank, :] - # Compute alpha from alpha_scale and rank for each split - alpha = alpha_scale * rank - converted[f"{path}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32) - - # Handle lora_up (lora_B) - may have multiple blocks - lora_up_blocks = [] - i = 0 - while f"lora_up.blocks.{i}.weight" in weight_keys: - lora_up_blocks.append(source_weights[weight_keys[f"lora_up.blocks.{i}.weight"]]) - i += 1 - - if lora_up_blocks: - # Multi-block LoRA: construct block-diagonal lora_B - # This is equivalent to the multi-block computation without modifying fastvideo - n_blocks = len(lora_up_blocks) - out_per_block, rank_per_block = lora_up_blocks[0].shape # e.g., [4096, 128] - - if len(targets) == 1: - # Single layer with multi-block: create block-diagonal matrix - total_out = out_per_block * n_blocks - total_rank = rank_per_block * n_blocks - lora_B_blockdiag = torch.zeros(total_out, total_rank, dtype=lora_up_blocks[0].dtype) - - for i in range(n_blocks): - lora_B_blockdiag[i * out_per_block:(i + 1) * out_per_block, - i * rank_per_block:(i + 1) * rank_per_block] = lora_up_blocks[i] - - converted[f"{targets[0][0]}.lora_B"] = lora_B_blockdiag - # Note: rank for alpha calculation should be total_rank (will be computed from lora_A.shape[0]) - else: - # Multi-block with split targets (e.g., QKV split) - # Each target gets one block - for i, (path, _) in enumerate(targets): - if i < n_blocks: - converted[f"{path}.lora_B"] = lora_up_blocks[i] - elif "lora_up.weight" in weight_keys: - lora_up = source_weights[weight_keys["lora_up.weight"]] - # Split if needed - if len(targets) == 1: - converted[f"{targets[0][0]}.lora_B"] = lora_up - else: - n = len(targets) - out_dim = lora_up.shape[0] // n - for i, (path, _) in enumerate(targets): - converted[f"{path}.lora_B"] = lora_up[i * out_dim:(i + 1) * out_dim, :] - else: - continue - - print(f" Output keys: {len(converted)} (including lora_alpha)") - # Count how many lora_alpha values were added - alpha_count = sum(1 for k in converted.keys() if "lora_alpha" in k) - print(f" Alpha values saved: {alpha_count}") - return converted - - -def convert_loras(source_dir: Path, output_dir: Path) -> bool: - """Convert all LoRA files in source directory.""" - lora_source = source_dir / "lora" - if not lora_source.exists(): - print(" No LoRA directory found, skipping") - return False - - lora_files = list(lora_source.glob("*.safetensors")) - if not lora_files: - print(" No LoRA files found, skipping") - return False - - print(f" Found {len(lora_files)} LoRA file(s)") - - # Map LoRA filenames to subdirectory names for FastVideo compatibility - # Each LoRA gets its own directory under lora/ - lora_subdir_mapping = { - "cfg_step_lora.safetensors": "distilled", - "refinement_lora.safetensors": "refinement", - } - - for lora_file in lora_files: - try: - # Determine output subdirectory - use mapping if available, otherwise generic name - if lora_file.name in lora_subdir_mapping: - lora_subdir_name = lora_subdir_mapping[lora_file.name] - else: - # For unknown LoRAs, create subdirectory based on filename - lora_subdir_name = lora_file.stem - - lora_output = output_dir / "lora" / lora_subdir_name - lora_output.mkdir(parents=True, exist_ok=True) - - # Load - source_weights = load_file(str(lora_file)) - - # Convert - converted = convert_lora_weights(source_weights, lora_file.stem) - - # Save - output_file = lora_output / lora_file.name - save_file(converted, str(output_file)) - - size_mb = output_file.stat().st_size / (1024**2) - print(f" ✓ {lora_file.name} → lora/{lora_subdir_name}/ ({size_mb:.1f} MB)") - - except Exception as e: - print(f" ❌ Failed to convert {lora_file.name}: {e}") - return False - - return True - - -def main(): - parser = argparse.ArgumentParser(description="Convert LongCat weights to FastVideo native format") - parser.add_argument( - "--source", - type=str, - required=True, - help="Path to original LongCat weights (LongCat-Video/weights/LongCat-Video/)", - ) - parser.add_argument( - "--output", - type=str, - required=True, - help="Path to output directory for native weights", - ) - parser.add_argument( - "--validate", - action="store_true", - help="Run validation after conversion", - ) - - args = parser.parse_args() - - source_dir = Path(args.source) - output_dir = Path(args.output) - - # Check source directory - if not source_dir.exists(): - print(f"❌ Error: Source directory not found: {source_dir}") - return 1 - - # Check for dit/transformer directory (original uses 'dit', we output to 'transformer') - transformer_source = source_dir / "dit" - if not transformer_source.exists(): - print(f"❌ Error: DiT directory not found in source") - return 1 - - print("=" * 60) - print("LongCat → FastVideo Native Conversion") - print("=" * 60) - print(f"Source: {source_dir}") - print(f"Output: {output_dir}") - print() - - # Step 1: Convert transformer weights - print("[Step 1/4] Converting transformer weights...") - - # Load source weights - shard_files = sorted(glob.glob(str(transformer_source / "*.safetensors"))) - if not shard_files: - print(f"❌ Error: No safetensors files found in {transformer_source}") - return 1 - - print(f" Found {len(shard_files)} shard(s)") - source_weights = {} - for shard_file in shard_files: - print(f" Loading {Path(shard_file).name}...") - source_weights.update(load_file(shard_file)) - - print(f" Loaded {len(source_weights)} parameters") - - # Convert - converted_weights = convert_transformer_weights(source_weights) - print(f" Converted to {len(converted_weights)} parameters") - - # Validate if requested - if args.validate: - if not validate_conversion(source_weights, converted_weights): - print("\n❌ Validation failed!") - return 1 - print("\n✓ Validation passed!") - - # Save - transformer_output = output_dir / "transformer" - transformer_output.mkdir(parents=True, exist_ok=True) - output_file = transformer_output / "model.safetensors" - print(f"\n Saving to {output_file}...") - save_file(converted_weights, str(output_file)) - size_gb = output_file.stat().st_size / (1024**3) - print(f" ✓ Saved ({size_gb:.2f} GB)") - print() - - # Step 2: Copy other components - print("[Step 2/5] Copying other components...") - output_dir.mkdir(parents=True, exist_ok=True) - - components = ["vae", "text_encoder", "tokenizer", "scheduler"] - for component in components: - copy_component(source_dir, output_dir, component) - - print() - - # Step 3: Convert LoRA weights - print("[Step 3/5] Converting LoRA weights...") - convert_loras(source_dir, output_dir) - print() - - # Step 4: Update transformer config - print("[Step 4/5] Updating transformer config...") - - # Copy config.json from source if exists - source_config = transformer_source / "config.json" - output_config = transformer_output / "config.json" - if source_config.exists(): - shutil.copy(source_config, output_config) - print(f" Copied config.json") - - update_transformer_config(transformer_output) - print() - - # Step 5: Create model_index.json - print("[Step 5/5] Creating model_index.json...") - model_index_path = output_dir / "model_index.json" - with open(model_index_path, 'w') as f: - json.dump(create_model_index(), f, indent=2) - print(f" ✓ Created {model_index_path}") - print() - - print("=" * 60) - print("✓ Conversion Complete!") - print("=" * 60) - print(f"Native weights ready at: {output_dir}") - print() - print("Converted components:") - print(" ✓ Transformer (native FastVideo implementation)") - print(" ✓ VAE, text encoder, tokenizer, scheduler") - if (output_dir / "lora").exists(): - lora_dirs = [d for d in (output_dir / "lora").iterdir() if d.is_dir()] - if lora_dirs: - print(f" ✓ LoRA weights ({len(lora_dirs)} adapters)") - for lora_dir in sorted(lora_dirs): - print(f" - lora/{lora_dir.name}/") - print() - print("Next steps:") - print() - print(" 1. Test basic generation:") - print(" from fastvideo import VideoGenerator") - print(f" generator = VideoGenerator.from_pretrained('{output_dir}')") - print(" video = generator.generate_video(") - print(" prompt='A cat playing piano',") - print(" num_inference_steps=50") - print(" )") - print() - if (output_dir / "lora" / "distilled").exists(): - print(" 2. Test distilled generation (16 steps with LoRA):") - print(f" generator = VideoGenerator.from_pretrained('{output_dir}',") - print(f" lora_path='{output_dir}/lora/distilled',") - print(" lora_nickname='distilled')") - print(" video = generator.generate_video(") - print(" prompt='A cat playing piano',") - print(" num_inference_steps=16,") - print(" guidance_scale=1.0)") - print() - print() - - return 0 - - -if __name__ == "__main__": - exit(main()) diff --git a/scripts/checkpoint_conversion/validate_longcat_weights.py b/scripts/checkpoint_conversion/validate_longcat_weights.py deleted file mode 100644 index d0416dc21..000000000 --- a/scripts/checkpoint_conversion/validate_longcat_weights.py +++ /dev/null @@ -1,237 +0,0 @@ -import os -import json -import argparse -from pathlib import Path -from safetensors import safe_open - - -def validate_components(model_path): - """Validate all model components exist.""" - print("=" * 60) - print("VALIDATING COMPONENTS") - print("=" * 60) - - components = { - "tokenizer": ["special_tokens_map.json", "tokenizer_config.json"], - "text_encoder": ["config.json", "model.safetensors.index.json"], - "vae": ["config.json", "diffusion_pytorch_model.safetensors"], - "scheduler": ["scheduler_config.json"], - "transformer": ["config.json", "diffusion_pytorch_model.safetensors.index.json"] - } - - all_valid = True - for component, required_files in components.items(): - component_path = os.path.join(model_path, component) - print(f"\n{component}:") - - if not os.path.exists(component_path): - print(f" ✗ Directory not found") - all_valid = False - continue - - for req_file in required_files: - file_path = os.path.join(component_path, req_file) - exists = os.path.exists(file_path) - symbol = "✓" if exists else "✗" - print(f" {symbol} {req_file}") - if not exists: - all_valid = False - - return all_valid - - -def validate_dit_weights(dit_path): - """Validate DiT weights structure.""" - print("\n" + "=" * 60) - print("VALIDATING DiT WEIGHTS") - print("=" * 60) - - # Load config - config_path = os.path.join(dit_path, "config.json") - with open(config_path) as f: - config = json.load(f) - - hidden_size = config["hidden_size"] - depth = config["depth"] - num_heads = config["num_heads"] - - print(f"\nArchitecture:") - print(f" - hidden_size: {hidden_size}") - print(f" - depth: {depth}") - print(f" - num_heads: {num_heads}") - - # Load weight index - index_path = os.path.join(dit_path, "diffusion_pytorch_model.safetensors.index.json") - with open(index_path) as f: - index = json.load(f) - - weight_map = index["weight_map"] - all_keys = list(weight_map.keys()) - - print(f"\nWeight statistics:") - print(f" - Total keys: {len(all_keys)}") - print(f" - Total size: {index['metadata']['total_size'] / 1e9:.2f} GB") - - # Check structure - embedder_keys = [k for k in all_keys if 'embedder' in k] - block_keys = [k for k in all_keys if k.startswith('blocks.')] - final_keys = [k for k in all_keys if k.startswith('final_layer.')] - - print(f"\nKey distribution:") - print(f" - Embedder layers: {len(embedder_keys)}") - print(f" - Transformer blocks: {len(block_keys)}") - print(f" - Final layer: {len(final_keys)}") - - # Verify all blocks present - block_nums = set() - for key in block_keys: - if key.startswith('blocks.'): - block_num = int(key.split('.')[1]) - block_nums.add(block_num) - - expected_blocks = set(range(depth)) - missing_blocks = expected_blocks - block_nums - - if missing_blocks: - print(f"\n✗ Missing blocks: {sorted(missing_blocks)}") - return False - else: - print(f"\n✓ All {depth} blocks present (0-{depth-1})") - - # Sample weights - first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors") - print(f"\nSampling weights from first shard:") - - with safe_open(first_shard, framework="pt", device="cpu") as f: - sample_keys = [k for k in f.keys() if k in all_keys][:5] - for key in sample_keys: - tensor = f.get_tensor(key) - print(f" - {key}") - print(f" Shape: {tuple(tensor.shape)}, Dtype: {tensor.dtype}") - - return True - - -def validate_shapes(dit_path): - """Validate weight shapes match expected architecture.""" - print("\n" + "=" * 60) - print("VALIDATING WEIGHT SHAPES") - print("=" * 60) - - # Load config - config_path = os.path.join(dit_path, "config.json") - with open(config_path) as f: - config = json.load(f) - - hidden_size = config["hidden_size"] - num_heads = config["num_heads"] - head_dim = hidden_size // num_heads - adaln_dim = config.get("adaln_tembed_dim", 512) - mlp_ratio = config.get("mlp_ratio", 4) - - # Calculate FFN hidden_dim using SwiGLU formula from blocks.py - # hidden_dim = int(2 * (hidden_size * mlp_ratio) / 3) - # rounded to multiple_of=256 - multiple_of = 256 - ffn_hidden = int(2 * hidden_size * mlp_ratio / 3) - ffn_hidden = multiple_of * ((ffn_hidden + multiple_of - 1) // multiple_of) - - # Expected shapes - expected = { - "x_embedder.proj.weight": (hidden_size, 16, 1, 2, 2), - "x_embedder.proj.bias": (hidden_size, ), - "t_embedder.mlp.0.weight": (adaln_dim, 256), - "t_embedder.mlp.2.weight": (adaln_dim, adaln_dim), - "y_embedder.y_proj.0.weight": (hidden_size, 4096), - "blocks.0.attn.qkv.weight": (3 * hidden_size, hidden_size), - "blocks.0.attn.q_norm.weight": (head_dim, ), - "blocks.0.attn.proj.weight": (hidden_size, hidden_size), - "blocks.0.cross_attn.q_linear.weight": (hidden_size, hidden_size), - "blocks.0.cross_attn.kv_linear.weight": (2 * hidden_size, hidden_size), - "blocks.0.ffn.w1.weight": (ffn_hidden, hidden_size), - "blocks.0.adaLN_modulation.1.weight": (6 * hidden_size, adaln_dim), - "final_layer.linear.weight": (64, hidden_size), - } - - # Load and check - first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors") - - all_valid = True - with safe_open(first_shard, framework="pt", device="cpu") as f: - for key, expected_shape in expected.items(): - if key in f.keys(): - tensor = f.get_tensor(key) - actual_shape = tuple(tensor.shape) - - if actual_shape == expected_shape: - print(f"✓ {key}: {actual_shape}") - else: - print(f"✗ {key}: expected {expected_shape}, got {actual_shape}") - all_valid = False - - return all_valid - - -def validate_model_index(model_path): - """Validate model_index.json exists and is correct.""" - print("\n" + "=" * 60) - print("VALIDATING MODEL INDEX") - print("=" * 60) - - model_index_path = os.path.join(model_path, "model_index.json") - - if not os.path.exists(model_index_path): - print("✗ model_index.json not found") - return False - - with open(model_index_path) as f: - index = json.load(f) - - required_keys = ["_class_name", "workload_type", "tokenizer", "text_encoder", "vae", "scheduler", "transformer"] - - all_valid = True - for key in required_keys: - if key in index: - print(f"✓ {key}: {index[key]}") - else: - print(f"✗ {key}: missing") - all_valid = False - - return all_valid - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Validate LongCat weights for FastVideo") - parser.add_argument("--model-path", type=str, required=True, help="Path to LongCat model directory") - parser.add_argument("--check-shapes", action="store_true", help="Also validate weight shapes (slower)") - - args = parser.parse_args() - - print(f"\nValidating: {args.model_path}\n") - - # Run validations - components_valid = validate_components(args.model_path) - model_index_valid = validate_model_index(args.model_path) - dit_valid = validate_dit_weights(os.path.join(args.model_path, "transformer")) - - if args.check_shapes: - shapes_valid = validate_shapes(os.path.join(args.model_path, "transformer")) - else: - shapes_valid = True - print("\nSkipping shape validation (use --check-shapes to enable)") - - # Summary - print("\n" + "=" * 60) - print("VALIDATION SUMMARY") - print("=" * 60) - - all_valid = components_valid and model_index_valid and dit_valid and shapes_valid - - if all_valid: - print("✓ All validations passed!") - print("✓ Model ready for FastVideo") - else: - print("✗ Some validations failed") - print("✗ Please check errors above") - - print("=" * 60) diff --git a/scripts/inference/inference_longcat.yaml b/scripts/inference/inference_longcat.yaml deleted file mode 100644 index 7887a28a7..000000000 --- a/scripts/inference/inference_longcat.yaml +++ /dev/null @@ -1,43 +0,0 @@ -# Run with: -# fastvideo generate --config scripts/inference/inference_longcat.yaml -generator: - model_path: FastVideo/LongCat-Video-T2V-Diffusers - engine: - num_gpus: 1 - parallelism: - tp_size: 1 - sp_size: 1 - offload: - dit: false - text_encoder: false - vae: false - pin_cpu_memory: false - pipeline: - experimental: - enable_bsa: false -request: - 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 - sampling: - seed: 42 - num_frames: 93 - height: 480 - width: 832 - fps: 15 - num_inference_steps: 50 - guidance_scale: 4.0 - output: - output_path: outputs_video/longcat_t2v diff --git a/scripts/inference/inference_longcat_distill.yaml b/scripts/inference/inference_longcat_distill.yaml deleted file mode 100644 index 8e6b77d3d..000000000 --- a/scripts/inference/inference_longcat_distill.yaml +++ /dev/null @@ -1,44 +0,0 @@ -# Run with: -# fastvideo generate --config scripts/inference/inference_longcat_distill.yaml -generator: - model_path: FastVideo/LongCat-Video-T2V-Diffusers - engine: - num_gpus: 1 - parallelism: - tp_size: 1 - sp_size: 1 - offload: - dit: false - pin_cpu_memory: false - pipeline: - components: - lora_path: FastVideo/LongCat-Video-T2V-Distilled-LoRA - experimental: - enable_bsa: false - lora_nickname: distilled -request: - 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 - sampling: - seed: 42 - num_frames: 93 - height: 480 - width: 832 - fps: 15 - num_inference_steps: 16 - guidance_scale: 1.0 - output: - output_path: outputs_video/longcat_distill diff --git a/scripts/inference/inference_longcat_i2v.yaml b/scripts/inference/inference_longcat_i2v.yaml deleted file mode 100644 index 7ab21af18..000000000 --- a/scripts/inference/inference_longcat_i2v.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Run with: -# fastvideo generate --config scripts/inference/inference_longcat_i2v.yaml -generator: - model_path: FastVideo/LongCat-Video-I2V-Diffusers - engine: - num_gpus: 1 - parallelism: - tp_size: 1 - sp_size: 1 - offload: - dit: false - pin_cpu_memory: false - pipeline: - workload_type: i2v - experimental: - enable_bsa: false -request: - 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 - inputs: - image_path: assets/girl.png - sampling: - seed: 42 - num_frames: 93 - height: 480 - width: 480 - fps: 15 - num_inference_steps: 50 - guidance_scale: 4.0 - output: - output_path: outputs_video/longcat_i2v diff --git a/scripts/inference/inference_longcat_refine_fromvideo.yaml b/scripts/inference/inference_longcat_refine_fromvideo.yaml deleted file mode 100644 index 29046b684..000000000 --- a/scripts/inference/inference_longcat_refine_fromvideo.yaml +++ /dev/null @@ -1,51 +0,0 @@ -# Run with: -# fastvideo generate --config scripts/inference/inference_longcat_refine_fromvideo.yaml -generator: - model_path: FastVideo/LongCat-Video-T2V-Diffusers - engine: - num_gpus: 1 - parallelism: - tp_size: 1 - sp_size: 1 - offload: - pin_cpu_memory: false - pipeline: - components: - lora_path: FastVideo/LongCat-Video-T2V-Refinement-LoRA - experimental: - enable_bsa: true - bsa_sparsity: 0.875 - bsa_chunk_q: [4, 4, 8] - bsa_chunk_k: [4, 4, 8] - lora_nickname: refinement -request: - 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 - inputs: - refine_from: outputs_video/longcat_distill/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 - sampling: - seed: 42 - height: 720 - width: 1280 - fps: 30 - num_inference_steps: 50 - guidance_scale: 1.0 - extensions: - t_thresh: 0.5 - spatial_refine_only: false - num_cond_frames: 0 - output: - output_path: outputs_video/longcat_refine_720p diff --git a/scripts/inference/inference_longcat_vc.yaml b/scripts/inference/inference_longcat_vc.yaml deleted file mode 100644 index 0fe5c78b3..000000000 --- a/scripts/inference/inference_longcat_vc.yaml +++ /dev/null @@ -1,43 +0,0 @@ -# Run with: -# fastvideo generate --config scripts/inference/inference_longcat_vc.yaml -generator: - model_path: FastVideo/LongCat-Video-VC-Diffusers - engine: - num_gpus: 1 - parallelism: - tp_size: 1 - sp_size: 1 - offload: - dit: false - pin_cpu_memory: false - pipeline: - experimental: - enable_bsa: false -request: - 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 - inputs: - video_path: assets/motorcycle.mp4 - sampling: - seed: 42 - num_frames: 93 - height: 480 - width: 832 - fps: 15 - num_inference_steps: 50 - guidance_scale: 4.0 - extensions: - num_cond_frames: 13 - output: - output_path: outputs_video/longcat_vc diff --git a/tests/local_tests/README.md b/tests/local_tests/README.md index 0a819a263..ffb8a5664 100644 --- a/tests/local_tests/README.md +++ b/tests/local_tests/README.md @@ -24,7 +24,6 @@ tests/local_tests// | Family | Workload | Dir | |---|---|---| | Flux2 | T2I | [`flux2/`](./flux2/) | -| Hunyuan GameCraft | T2V / I2V | [`gamecraft/`](./gamecraft/) | | GEN3C | T2V | [`gen3c/`](./gen3c/) | | Kandinsky-5 | T2V | [`kandinsky5/`](./kandinsky5/) | | Kandinsky-6 | TI2VA (T2V/I2V + audio) | [`kandinsky6/`](./kandinsky6/) | diff --git a/tests/local_tests/dreamx_world/PORT_STATUS.md b/tests/local_tests/dreamx_world/PORT_STATUS.md deleted file mode 100644 index 986a33e31..000000000 --- a/tests/local_tests/dreamx_world/PORT_STATUS.md +++ /dev/null @@ -1,189 +0,0 @@ -# DreamX World Port Status - -## Summary - -- model_family: `dreamx_world` -- workload_types: `I2V camera-control compatibility shim`; `I2V autoregressive camera-control forcing` -- official_ref: `https://github.com/AMAP-ML/DreamX-World` -- official_ref_dir: `DreamX-World/` -- hf_weights_path: `GD-ML/DreamX-World-5B-Cam` -- local_weights_dir: `official_weights/dreamx_world` -- source_layout: `raw_official` -- local_tests_readme: `tests/local_tests/dreamx_world/README.md` - -## Current Phase - -- phase: `phase_11_post_parity_handoff` -- status: `complete` -- owner: `orchestrator` -- last_updated: `2026-07-02` - -## Component Matrix - -| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | -|---|---|---|---|---|---|---|---|---|---| -| transformer | dit | ported_dedicated | `DreamX-World/models/wan_transformer3d.py`; PRoPE helpers in `DreamX-World/models/prope_utils.py` | `DreamX-World/inference_dreamx5b.py::setup_models`, `Wan2_2Transformer3DModel.from_pretrained(... cam_method=prope, add_control_adapter=True)` | `fastvideo/models/dits/dreamx_world.py`; `fastvideo/configs/models/dits/dreamx_world.py`; DreamX pipeline config helper | native_prope_pass | real_conversion_pass | strict_load_and_forward_parity_pass | none | -| vae | vae | reuse_pending | `DreamX-World/models/wan_vae3_8.py` | `DreamX-World/inference_dreamx5b.py::setup_models`, `AutoencoderKLWan3_8.from_pretrained(Wan2.2_VAE.pth)` | `fastvideo/models/vaes/wanvae.py`; DreamX VAE config helper | config_smoke_pass | raw_key_mapping_pass | encode_parity_pass | none | -| text_encoder/tokenizer | encoder | reuse_pending | `DreamX-World/models/wan_text_encoder.py`; tokenizer via Wan2.2 base model | `DreamX-World/inference_dreamx5b.py::setup_models`, `WanT5EncoderModel` + tokenizer subpaths | `fastvideo/models/encoders/t5.py::UMT5EncoderModel`; DreamX UMT5 config helper | config_smoke_pass | staged_weight_load_pass | hidden_state_parity_pass | none | -| scheduler | generic | reuse_proven | Diffusers `FlowMatchEulerDiscreteScheduler` | `DreamX-World/inference_dreamx5b.py::setup_models`, default `sampler_name=Flow` | `fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py` | pass | not_required | non_skip_pass | Q003 | -| camera_conditioning | generic | port_pending | `DreamX-World/utils/inference_utils.py`, `DreamX-World/models/prope_utils.py`, `DreamX-World/wan/modules/camera_prope.py` | `DreamX-World/inference_dreamx5b.py::get_camera_sequence`, `pipeline(... control_camera_video=...)` | `fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py` | pass | not_required | non_skip_pass | none | -| pipeline | pipeline | port_complete | `DreamX-World/pipeline/pipeline_dreamxworld.py` | `DreamX-World/inference_dreamx5b.py::process_inference_from_json` | `fastvideo/pipelines/basic/dreamx_world/` plus config/preset/registry | pipeline_load_generate_smoke_pass | model_index_and_config_consistency_smoke_pass | pipeline_api_vs_worker_forward_parity_pass | none | -| ar_transformer | dit | port_complete | `DreamX-World/wan/modules/causal_camera_model_2_2_prope_infinity.py::CausalWanModel` | `DreamX-World/inference_ar_forcing.py::load_pipeline` | `fastvideo/models/dits/dreamx_world_ar.py`; `fastvideo/configs/models/dits/dreamx_world.py::DreamXWorldARConfig` | tiny_official_forward_parity_pass | identity_conversion_pass | real_5b_strict_load_pass | none | -| ar_pipeline | pipeline | port_complete | `DreamX-World/pipeline/pipeline_causal_camera.py` | `DreamX-World/inference_ar_forcing.py::main` | `fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py`; `fastvideo/pipelines/basic/dreamx_world/ar_denoising.py`; registry/preset/config | config_registry_pass | symlink_model_index_pass | short_full_generation_pass | none | - -## Conversion State - -- conversion_script: `scripts/checkpoint_conversion/dreamx_world_to_diffusers.py` -- converted_weights_dir: `converted_weights/dreamx_world` -- source_layout: `raw_official` -- strict_load_status: `pass` -- conversion_script_status: `transformer_model_index_and_config_consistency_smoke_pass` -- model_index_status: `smoke_pass` -- passthrough_components: `Wan2.2 Diffusers scheduler, tokenizer, and text encoder are symlinked from official_weights/Wan2.2-TI2V-5B-Diffusers; VAE parity uses raw Wan2.2_VAE.pth with an explicit DreamX raw-to-FastVideo key mapper because official encode returns normalized latents.` -- retry_history: `none` - -## Parity Commands - -| Scope | Command | Last Result | Notes | -|---|---|---|---| -| transformer | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py -v -s` | strict_load_and_forward_parity_pass | 2026-07-01: converted real 5B-Cam transformer shards strict-load into dedicated `DreamXWorldTransformer3DModel` with 0 shape mismatches; official-vs-FastVideo small-input fp32 forward parity passes on CUDA (`diff_max=0.072533`, `diff_mean=0.008014`). | -| vae | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py -v -s` | encode_parity_pass | 2026-06-30: official DreamX raw `Wan2.2_VAE.pth` maps 196/196 keys into FastVideo Wan VAE and encode parity passes after applying the same official latent normalization (`(mu - mean) / std`). | -| text_encoder | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py -v -s` | hidden_state_parity_pass | 2026-06-30: official `WanT5EncoderModel` vs FastVideo `UMT5EncoderModel` hidden-state parity passes on CUDA using staged Wan2.2 text encoder/tokenizer weights and reference-only `xfuser` stubs. | -| scheduler | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py -v -s` | non_skip_pass | 2026-06-30: FastVideo FlowMatch scheduler matches official Diffusers timesteps and step output for DreamX default Flow sampler; `DreamXWorldPipeline` initializes FlowMatch with official `shift=3.0`. | -| camera_conditioning | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py -v -s` | non_skip_pass | 2026-06-29: 3 parameterized cases passed against official reference on CPU. | -| pipeline_config | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -v -s` | pipeline_entry_preset_scheduler_modelinfo_and_camera_stage_smoke_pass | DreamX 5B-Cam PipelineConfig wires DiT/VAE/UMT5/Flow/TI2V settings and official `shift=3.0`; default preset is registered for `GD-ML/DreamX-World-5B-Cam`; local converted-style `model_index.json` resolves to `DreamXWorldPipeline`; the pipeline initializes FlowMatch, camera conditioning writes `batch.extra["dreamx_y_camera"]`, and generic denoising can pass it as `y_camera`. | -| ar_transformer | `PYTHONPATH=/workspace/FastVideo python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py -q -rs` | 6_passed_0_skipped | 2026-07-02: AR converter writes symlinked transformer layout/model_index; tiny official `CausalWanModel` vs FastVideo `DreamXWorldARTransformer3DModel` forward parity passes; real 5B `model.safetensors` strict-loads with zero missing/unexpected keys from `/tmp/converted_dreamx_world_ar`. | -| ar_pipeline_config | `PYTHONPATH=/workspace/FastVideo python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -q -rs` | 10_passed_0_skipped | 2026-07-02: `DreamXWorld5BARPipelineConfig`, `DreamXWorldARPipeline`, registry config selection for `GD-ML/DreamX-World-5B`, and `dreamx_world_5b_ar` preset pass. | -| ar_full_generation | `PYTHONPATH=/workspace/FastVideo python /tmp/run_dreamx_ar_video_smoke.py` | generated_video_pass | 2026-07-02: A40 short full-generation smoke passed from `/tmp/converted_dreamx_world_ar` with 64x64, 9 frames, 4 denoise steps, `output_type=pil`, `save_video=True`; MP4 saved at `outputs_video/dreamx_world_ar_video_smoke/a quiet road through a futuristic coastal city at sunrise.mp4` and decoded as 9 frames of `(64, 64, 3)` uint8. | -| ar_long_horizon | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA python /tmp/run_dreamx_ar_long_horizon.py` | generated_video_pass | 2026-07-02: A40 long-horizon AR generation passed from `/tmp/converted_dreamx_world_ar` with 64x64, 1005 frames, 4 denoise steps, seed 4096; MP4 saved at `outputs_video/dreamx_world_ar_long_horizon/A long autonomous drive through a futuristic coastal city at sunrise, smooth forward camera motion,.mp4` and decoded as 1005 frames of `(64, 64, 3)`; end-to-end generation latency was 231.68s after load. | -| ar_ssim_default | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download` | 1_passed_0_skipped | 2026-07-02: A40 default AR SSIM reference seeded locally at `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/`; helper removes stale generated base MP4 before generation; default params are 192x192, 81 frames, 4 steps, seed 2048, min SSIM 0.98. | -| ar_ssim_modal_l40s | `modal run /tmp/modal_dreamx_ar_ssim_git.py` | 1_passed_0_skipped | 2026-07-02: Modal L40S run checked out `post-fix dreamx-world-5b-cam branch commit`, downloaded default `L40S_reference_videos` via `reference_videos_cli.py download`, used cached converted AR weights under `/root/data/dreamx_world_ar_converted`, seeded the missing AR L40S reference from generated output, reran a fresh generated-vs-reference compare successfully (`mean_ssim=1.0`), and exported the reference to Modal volume `hf-model-weights:dreamx_ar_ssim_l40s`. Downloaded local reference decodes as 81 frames of `(192, 192, 3)`. A40-vs-L40S reference spot check after deterministic AR noise fix: mean SSIM 0.7770, min 0.6139, max 0.9944. | -| pipeline_smoke | `python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py -q -rs` | 4_passed_0_skipped | 2026-06-30: combined smoke/parity passed. 2026-07-01: smoke alone passed with real `image_path` TI2V coverage (`3 passed`), validating image load, TI2V preprocessing, VAE first-frame encode under CPU offload, camera conditioning, and 1-step latent generation from `converted_weights/dreamx_world`. Tests force `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` to avoid the local FlashAttention-4 cute ABI mismatch. | -| basic_example | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA DREAMX_WORLD_MODEL_DIR=converted_weights/dreamx_world DREAMX_WORLD_IMAGE_PATH= DREAMX_WORLD_HEIGHT=64 DREAMX_WORLD_WIDTH=64 DREAMX_WORLD_NUM_FRAMES=9 DREAMX_WORLD_STEPS=1 DREAMX_WORLD_GUIDANCE=1.0 DREAMX_WORLD_OUTPUT_PATH=outputs_video/dreamx_world_example_smoke python examples/inference/basic/basic_dreamx_world.py` | generated_video_pass | 2026-06-30: example saved an MP4 under `outputs_video/dreamx_world_example_smoke`; imageio/ffmpeg decoded frame 0 as `(64, 64, 3)` uint8, fps 16, duration 0.56s. | - -## Open Questions - -| ID | Question | Owner | Needed By Phase | Status | Resolution | -|---|---|---|---|---|---| -| Q001 | Should first PR expose only the `DreamX-World-5B-Cam` 5s camera-control mode and exclude AR long-horizon forcing? | user | prep | resolved | User approved starting with `DreamX-World-5B-Cam`; AR long-horizon is out of first-PR scope. | -| Q002 | Does FastVideo's existing Wan2.2 TI2V transformer support DreamX PRoPE/control adapter with a small extension, or is a DreamX-specific DiT required? | component:transformer | Phase 3 | resolved | Project guidance prefers a separate DreamX DiT for maintainability. DreamX PRoPE/control adapter now lives in `fastvideo/models/dits/dreamx_world.py`; Wan DiT/config have no DreamX-specific fields or `y_camera` signature. | -| Q003 | Which sampler is in first-PR scope: official default `Flow` only, or also `Flow_Unipc` and `Flow_DPM++`? | orchestrator | Phase 3 | resolved | First PR should support official default `Flow` only. FastVideo FlowMatch scheduler parity is non-skip PASS; optional `Flow_Unipc` and `Flow_DPM++` are out of first-PR scope. | -| Q004 | Which HF token env var should be used if rate limits or gated Wan2.2 base weights require auth? | user | Phase 5 | resolved | No auth was required for the completed local downloads; keep using env var names only if future gated repos require auth. | -| Q005 | Should native FastVideo production code depend on DreamX reference-only packages such as `xfuser` or OpenCV? | user | Phase 3 | resolved | No. These packages may be used only for official reference/local parity setup; native FastVideo integration must remove that runtime requirement. | -| Q006 | Should AR handoff require a full generated long-horizon video in this no-HF-token/no-GPU-budget pass? | user/runtime | pipeline | resolved | A40 long-horizon generation passes with 1005 frames at 64x64/4 steps. AR default SSIM passes locally on A40 and on Modal L40S with a 192x192/81-frame reference. HF upload/publication remains a separate operation if the reference dataset should be updated upstream. | -| Q007 | Can A40 references stand in for L40S CI references? | quality | release | resolved | No. After deterministic AR noise fix, A40-vs-L40S reference mean SSIM is 0.7770, below the 0.98 same-device threshold. Publish the L40S-specific reference for CI. | - -## Issues And Blockers - -| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | -|---|---|---|---|---|---|---|---|---| -| I001 | prep | official_env | medium | Official import initially failed because `xfuser` was missing. | `ModuleNotFoundError: No module named 'xfuser'` from `python -c "import sys; sys.path.insert(0, 'DreamX-World'); import inference_dreamx5b"` | prep | resolved | Installed `xfuser==0.4.1`; import progressed. | -| I002 | prep | official_env | medium | Official import then failed because GUI OpenCV required missing system `libxcb.so.1`. | `ImportError: libxcb.so.1: cannot open shared object file` through `cv2` import in Diffusers ConsisID path. | prep | resolved | Installed `opencv-python-headless`; `import inference_dreamx5b` passed. | -| I003 | prep | weights | medium | HF repo has raw official transformer shards and no Diffusers `model_index.json`. | `inspect_hf_layout.py GD-ML/DreamX-World-5B-Cam --json` returned `source_layout=raw_official`, `needs_conversion=yes`, `model_index_class=null`. | conversion | resolved | Downloaded raw DreamX shards to `official_weights/dreamx_world`; converted transformer to `converted_weights/dreamx_world/transformer`; symlinked reusable Wan2.2 Diffusers components; real 5B transformer strict-load passes. | -| I004 | prep | dependencies | high | Official reference import required extra packages in the local environment, but FastVideo native runtime should not inherit those dependencies. | `xfuser==0.4.1` and `opencv-python-headless` were installed only to make `DreamX-World/inference_dreamx5b.py` import for reference/parity. | pipeline | resolved | Production DreamX FastVideo code uses native camera/image/video utilities and has no runtime `xfuser` or OpenCV import requirement; those packages remain reference-only local parity dependencies. | -| I005 | parity | transformer | medium | Transformer full forward parity initially failed in bf16 official harness. | Official CUDA bf16 LayerNorm path was unstable; fp32 small-input harness avoids that dtype issue and compares against FastVideo with single-process SP identity patches. | component:transformer | resolved | Official-vs-FastVideo forward parity now passes on CUDA with `diff_max=0.072533`, `diff_mean=0.008014`. | -| I006 | parity | vae/text_encoder | medium | VAE/text parity initially remained skipped after weights were staged. | Text official import needed a reference-only `xfuser` stub; VAE comparison initially used raw official normalized latents against FastVideo raw mu. | component:vae,component:text_encoder | resolved | Text hidden-state parity passes. VAE encode parity passes after raw key mapping and applying the official latent normalization to FastVideo output. | -| I007 | quality | pipeline_ti2v | medium | Real image-path TI2V smoke initially failed when `vae_cpu_offload=True` because DenoisingStage encoded the first frame while the VAE weights remained on CPU. | DreamX SSIM first run failed with `RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same` at `fastvideo/pipelines/stages/denoising.py` VAE encode. | pipeline | resolved | DenoisingStage now moves the VAE to `local_device` before TI2V first-frame encode; image-path pipeline smoke and DreamX SSIM both pass. | -| I008 | quality | ssim_helper | high | SSIM helper could compare against a stale generated base MP4 when a rerun saved the new video as `_1.mp4`. | Existing generated outputs made AR reference seeding appear to pass before a fresh generated-vs-reference compare. | quality | resolved | `run_text_to_video_similarity_test` and `run_image_to_video_similarity_test` now remove the stale generated base MP4 before generation. A40 and Modal L40S AR SSIM were rerun after the fix. | -| I009 | quality | ar_denoising | high | AR denoising added CUDA noise without using the request seed when the original generator was CPU-backed. | Fresh reruns against old AR references produced mean SSIM near 0.05. | pipeline | resolved | `DreamXWorldARCausalDenoisingStage` now derives a device-local generator from the request seed for AR noise. Fresh same-device A40 and L40S reruns pass with mean SSIM 1.0 after reseeding references. | - -## Escape Hatches - -| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | -|---|---|---|---|---|---|---| - -## Decisions - -| Date | Decision | Rationale | Impact | -|---|---|---|---| -| 2026-06-29 | First PR scope is `DreamX-World-5B-Cam` only. | Cam mode is closest to existing Wan2.2 TI2V support; AR forcing needs separate causal/KV pipeline work. | Component inventory and parity focus on `inference_dreamx5b.py` and `pipeline_dreamxworld.py`. | -| 2026-06-29 | Do not install full DreamX requirements during prep. | Full requirements pin core FastVideo stack packages. | Installed only `xfuser==0.4.1` and `opencv-python-headless` to make official imports work. | -| 2026-06-29 | Treat HF DreamX-World-5B-Cam weights as raw official transformer layout requiring conversion. | HF inspection found no `model_index.json`. | Phase 5 must create `scripts/checkpoint_conversion/dreamx_world_to_diffusers.py` after component prototype/key dumps. | -| 2026-06-29 | Do not add DreamX reference-only dependencies to FastVideo production requirements. | The current environment should remain the FastVideo environment; extra packages are only for official reference parity. | Native DreamX integration must avoid runtime `xfuser` and OpenCV requirements unless explicitly approved later. | -| 2026-06-29 | Implement DreamX camera conditioning as native FastVideo utility. | It is weightless and removes the need to import DreamX reference utilities at production runtime. | `fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py` now has non-skip parity against official action-to-PRoPE tensors. | -| 2026-06-29 | First PR supports DreamX default `Flow` sampler only. | FastVideo FlowMatch Euler scheduler matches the official Diffusers scheduler for DreamX defaults. | Pipeline work can use FastVideo native FlowMatch scheduler; UniPC and DPM++ are out of first-PR scope. | -| 2026-07-01 | Keep DreamX PRoPE/control adapter in a dedicated DreamX DiT class. | Project guidance is that putting too much DreamX behavior into Wan makes the model hard to manage. | `fastvideo/models/dits/dreamx_world.py` defines `DreamXWorldTransformer3DModel`, `DreamXWorldTransformerBlock`, and `DreamXPropeSelfAttention`; `fastvideo/configs/models/dits/dreamx_world.py` owns DreamX adapter config fields; Wan DiT/config are unchanged from DreamX. | -| 2026-06-30 | Camera parity test loads official camera functions by file instead of importing the official `utils` package. | Official package initialization pulls unrelated dependencies that can require GUI OpenCV system libraries. | Camera parity remains non-skip without adding DreamX reference-only dependencies to FastVideo production requirements. | -| 2026-06-30 | Add DreamX-World-5B-Cam model and pipeline config helpers plus a conversion script. | Official HF DreamX 5B-Cam transformer config is 30 layers, hidden size 3072, 24 heads, 48 latent channels, plus Wan2.2 48-channel VAE and UMT5-XXL text encoder. | DreamX helpers wire DiT/VAE/UMT5/Flow/TI2V settings; `dreamx_world_to_diffusers.py` writes a FastVideo-loadable transformer config plus renamed safetensors; strict-load smoke passes on a tiny official DreamX transformer and the real 5B converted shards. | -| 2026-06-30 | Pass DreamX camera PRoPE condition through the FastVideo batch/denoising path. | DreamX transformer expects `y_camera={"viewmats", "K"}` at denoising time. | `DreamXWorldPipeline` is registered as a basic pipeline entry and initializes the official default FlowMatch scheduler; `dreamx_world_5b_cam` preset mirrors official 5B-Cam defaults; `DreamXWorldCameraConditioningStage` writes `batch.extra["dreamx_y_camera"]`; generic denoising filters and forwards it as `y_camera` only for compatible transformers. | - -## Handoff Notes - -- Prep, component parity, pipeline smoke/parity, and the basic example validation are complete for `DreamX-World-5B-Cam`. -- Official reference clone is staged at `DreamX-World/` and ignored by git. -- Workspace-local weights are staged: DreamX raw transformer shards under `official_weights/dreamx_world`, Wan2.2 raw base artifacts under `official_weights/Wan2.2-TI2V-5B`, and Wan2.2 Diffusers reusable components under `official_weights/Wan2.2-TI2V-5B-Diffusers`. -- Camera conditioning parity is active and passing without weights. -- Default Flow scheduler parity is active and passing without weights. -- Transformer has corrected official 5B-Cam architecture in dedicated DreamX DiT/config files, native PRoPE/control-adapter, conversion mapping, real converted 5B strict-load, and official-vs-FastVideo forward parity passing on CUDA. VAE encode parity and text hidden-state parity pass on CUDA. Pipeline entry/registry, local model_info resolution, preset, config, FlowMatch scheduler init, camera stage, denoising `y_camera` kwarg smokes, independent CUDA pipeline smoke/parity, and a small saved-video basic example pass. -- Full local DreamX component suite is non-skip PASS: `python -m pytest tests/local_tests/dreamx_world/ -q -rs` returned `26 passed` on 2026-07-01. Pipeline smoke/parity and SSIM quality regression are also non-skip PASS locally. -- Keep `xfuser` and OpenCV as reference-only parity dependencies. Do not add them - to FastVideo requirements or production imports. - - -## Quality Regression - -- status: `added` -- test: `fastvideo/tests/ssim/test_dreamx_world_similarity.py` -- command: `PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B-Cam python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs` -- result: `1 passed, 0 skipped` on 2026-07-01 -- reference: Local A40/TORCH_SDPA reference seeded from the generated candidate under `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B-Cam/TORCH_SDPA/`. The test uses a deterministic generated input image, 64x64 request dimensions, 9 frames, 1 denoise step, seed 1024, and min SSIM 0.98. Full-quality params are present for 480x832/161 frames/30 steps. -- note: Modal L40S seeding passed using the configured Modal profile and unauthenticated HF public downloads. HF upload/publication still requires `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` with write access; no token values were used or recorded. - -## Final Handoff - -```text -final_handoff: - prep_handoff_complete: yes - conversion_status: pass - components: - - name: transformer - reuse_or_port: ported_dedicated_dit - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py - parity_status: non_skip_pass - concerns_or_unknowns: none - - name: vae - reuse_or_port: reused - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py - parity_status: non_skip_pass - concerns_or_unknowns: none - - name: text_encoder_tokenizer - reuse_or_port: reused - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py - parity_status: non_skip_pass - concerns_or_unknowns: none - - name: scheduler - reuse_or_port: reused - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py - parity_status: non_skip_pass - concerns_or_unknowns: none - - name: camera_conditioning - reuse_or_port: ported - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py - parity_status: non_skip_pass - concerns_or_unknowns: none - - name: ar_transformer - reuse_or_port: ported - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py - parity_status: non_skip_pass - concerns_or_unknowns: none - - name: ar_pipeline - reuse_or_port: ported - parity_test: tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py plus AR smoke/SSIM commands listed above - parity_status: non_skip_pass - concerns_or_unknowns: none - pipeline_smoke: pass - pipeline_parity: pass - example_status: pass - quality_regression: added - local_tests_readme: tests/local_tests/dreamx_world/README.md - port_state_file: tests/local_tests/dreamx_world/PORT_STATUS.md - token_values_committed: no - runtime_third_party_model_imports: none - blockers: none - escape_hatch: none -``` - -| 2026-07-02 | Add DreamX-World-5B autoregressive support. | Official AR repo is raw single-safetensors layout and needs a dedicated causal/KV stage. | Added native AR DiT/config, identity converter, AR pipeline config/preset/registry, and targeted non-skip tests. Raw AR weights are staged at `/tmp/dreamx_world_ar_weights`; converted symlink layout at `/tmp/converted_dreamx_world_ar`. | -| 2026-07-02 | Validate DreamX-World-5B AR short full generation on A40. | Targeted parity/config tests prove components, but end-to-end runtime can still fail at scheduler, RoPE cache, device, or decode boundaries. | `DreamXWorldARPipeline` generated and saved a 64x64/9-frame/4-step MP4 from `/tmp/converted_dreamx_world_ar`; the saved MP4 decodes to 9 frames. | -| 2026-07-02 | Validate DreamX-World-5B AR long-horizon and default SSIM on A40. | AR needs coverage beyond the 9-frame smoke to exercise longer KV/cache progression and a quality regression path. | 1005-frame 64x64/4-step generation passes and decodes; default AR SSIM test uses 192x192/81 frames because MS-SSIM requires short side >160. | -| 2026-07-02 | Validate DreamX-World-5B AR default SSIM on Modal L40S. | CI references are device-specific; A40 alone is not enough for L40S reference coverage. | Modal checked out `post-fix dreamx-world-5b-cam branch commit`, downloaded existing default L40S references, seeded the missing AR reference, reran SSIM successfully (`mean_ssim=1.0`), and exported the L40S reference back to the workspace. | diff --git a/tests/local_tests/dreamx_world/README.md b/tests/local_tests/dreamx_world/README.md deleted file mode 100644 index 2a836b17b..000000000 --- a/tests/local_tests/dreamx_world/README.md +++ /dev/null @@ -1,254 +0,0 @@ -# DreamX World Local Tests - -Local-only parity and smoke tests for the `dreamx_world` FastVideo port. These -tests compare FastVideo against the official DreamX-World reference -implementation and are not expected to run in CI unless explicitly promoted -later. - -Port progress, open questions, issues, and handoff notes live in -`tests/local_tests/dreamx_world/PORT_STATUS.md`. - -## Reference Assets - -| Field | Value | -|---|---| -| Model family | `dreamx_world` | -| First-PR scope | `DreamX-World-5B-Cam`; follow-up scope now includes `DreamX-World-5B` autoregressive forcing | -| Out-of-scope variants | none for the DreamX-World 5B/Cam paths currently ported | -| Workload types | I2V camera-control compatibility shim: image + prompt + action sequence to video | -| Official reference | `https://github.com/AMAP-ML/DreamX-World` | -| Local reference dir | `DreamX-World/` | -| Official commit/version | `221875811ba31f7eac6c3025b215c09ad2cefd1d` | -| HF weights | `GD-ML/DreamX-World-5B-Cam` | -| HF revision | default | -| Local weights dir | `official_weights/dreamx_world` | -| Source layout | `raw_official` | -| Needs conversion | `yes` | - -Do not write token values in this file. Current token env var detected during -prep: `none`. - -## Shared Environment Setup - -Run from the FastVideo repo root in the same conda/env used for FastVideo. Do -not create a separate upstream environment for parity tests. - -```bash -python ".agents/skills/add-model-01-prep/scripts/clone_reference_repo.py" \ - "https://github.com/AMAP-ML/DreamX-World.git" \ - "DreamX-World" \ - --commit "221875811ba31f7eac6c3025b215c09ad2cefd1d" \ - --update-gitignore -``` - -DreamX-World does not expose a packaging file for editable install. During prep -the official import check used `sys.path.insert(0, "DreamX-World")`. - -Additional official deps installed into the current environment for imports: - -```bash -uv pip install xfuser==0.4.1 -uv pip install opencv-python-headless -``` - -These packages are for running the official DreamX reference during local -parity only. They must not become FastVideo production/runtime dependencies for -the native `dreamx_world` pipeline. - -Do not install the full `DreamX-World/requirements.txt` without explicit -approval. It pins core FastVideo stack packages including `torch`, `torchvision`, -`triton`, `flash_attn`, and `diffusers`. - -## Official Environment Status - -```text -dependency_changes: installed official deps in current env -official_env_status: imports_ok -private_dep_stubs: none -blocked_on: none -``` - -Import check used during prep: - -```bash -python -c "import sys; sys.path.insert(0, 'DreamX-World'); import inference_dreamx5b; print('imports_ok')" -``` - -## Weight Setup - -HF layout inspection found no root `model_index.json`; the repo contains -`config.json`, a safetensors index, and three transformer safetensors shards. -This is a raw official transformer layout and requires conversion before -FastVideo can load it through `VideoGenerator.from_pretrained`. - -```bash -python ".agents/skills/add-model-01-prep/scripts/inspect_hf_layout.py" \ - "GD-ML/DreamX-World-5B-Cam" \ - --json -``` - -Weights have been staged workspace-locally. The raw DreamX transformer repo lives at `official_weights/dreamx_world`; Wan2.2 raw base artifacts live at `official_weights/Wan2.2-TI2V-5B`; Wan2.2 Diffusers reusable components live at `official_weights/Wan2.2-TI2V-5B-Diffusers`. To reproduce the DreamX download: - -```bash -python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \ - "GD-ML/DreamX-World-5B-Cam" \ - "official_weights/dreamx_world" -``` - - - -### DreamX-World-5B Autoregressive Setup - -The AR repository `GD-ML/DreamX-World-5B` is also raw official layout: no -`model_index.json`, root `config.json`, and a single `model.safetensors`. The -current environment has the raw AR checkpoint staged outside the workspace at -`/tmp/dreamx_world_ar_weights` to avoid workspace quota pressure. The converted -FastVideo layout is staged at `/tmp/converted_dreamx_world_ar` with the 21GB -transformer safetensors symlinked instead of copied. - -```bash -python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \ - "GD-ML/DreamX-World-5B" \ - "/tmp/dreamx_world_ar_weights" - -python scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py \ - --source /tmp/dreamx_world_ar_weights \ - --output /tmp/converted_dreamx_world_ar \ - --component-source official_weights/Wan2.2-TI2V-5B-Diffusers \ - --symlink-components \ - --symlink-transformer -``` - -AR production code uses `DreamXWorldARTransformer3DModel`, -`DreamXWorld5BARPipelineConfig`, `DreamXWorldARPipeline`, and -`DreamXWorldARCausalDenoisingStage`. The AR DiT is a native FastVideo port of -the official Apache-2.0 `CausalWanModel`; it has no production DreamX, Diffusers -model-class, Transformers model-class, `xfuser`, or OpenCV import. - -## Prototype And Conversion Artifacts - -State-dict key/shape dumps are generated after FastVideo native prototypes exist -and are used to build the conversion mapping. - -```text -official_key_dumps: - transformer: converted_weights/dreamx_world/_mapping/transformer_official_keys.json -fastvideo_key_dumps: - transformer: converted_weights/dreamx_world/_mapping/transformer_fastvideo_keys.json -conversion_script: scripts/checkpoint_conversion/dreamx_world_to_diffusers.py -conversion_script_status: transformer_model_index_and_config_consistency_smoke_pass -conversion_source_layout: raw_official -converted_weights_dir: converted_weights/dreamx_world -model_index_status: smoke_pass -strict_load_status: pass -``` - -The converter writes `transformer/` from raw DreamX shards. To create a full -FastVideo-loadable diffusers-style root, pass a Wan2.2 Diffusers directory as the -component source so reusable components are copied or symlinked before -`model_index.json` is emitted: - -```bash -python scripts/checkpoint_conversion/dreamx_world_to_diffusers.py \ - --source official_weights/dreamx_world \ - --output converted_weights/dreamx_world \ - --component-source /path/to/Wan2.2-TI2V-5B-Diffusers \ - --symlink-components -``` - -## Expected Parity Tests - -Planned local tests for this family: - -| Component | Official files / args | Test | Concerns | Status | -|---|---|---|---|---| -| transformer | `DreamX-World/models/wan_transformer3d.py`; instantiated in `DreamX-World/inference_dreamx5b.py` with `cam_method=prope`, `add_control_adapter=True` | `tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py` | FastVideo uses dedicated `DreamXWorldTransformer3DModel`/`DreamXWorldConfig` files for DreamX 5B-Cam config, PRoPE, conversion mapping, real converted 5B strict-load PASS, and official-vs-FastVideo small-input forward parity PASS on CUDA. Wan DiT/config have no DreamX-specific adapter fields. | strict_load_and_forward_parity_pass | -| vae | `DreamX-World/models/wan_vae3_8.py`; `vae_type=AutoencoderKLWan3_8`, `vae_subpath=Wan2.2_VAE.pth` | `tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py` | DreamX raw `Wan2.2_VAE.pth` maps 196/196 keys into FastVideo Wan VAE; encode parity passes after applying official latent normalization. | encode_parity_pass | -| text_encoder/tokenizer | `DreamX-World/models/wan_text_encoder.py`; T5 path from Wan2.2 base model | `tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py` | Official `WanT5EncoderModel` vs FastVideo `UMT5EncoderModel` hidden-state parity passes on CUDA with staged Wan2.2 weights/tokenizer. | hidden_state_parity_pass | -| scheduler | Diffusers `FlowMatchEulerDiscreteScheduler`; selected by default `sampler_name=Flow` | `tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py` | First PR can support official default `Flow`; optional `Flow_Unipc` and `Flow_DPM++` are out of first-PR scope. | non_skip_pass | -| camera_conditioning | `DreamX-World/utils/inference_utils.py`, `models/prope_utils.py`, `wan/modules/camera_prope.py` | `tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py` | Action sequence to PRoPE/control input must match official tensor shapes and values. | non_skip_pass | -| pipeline | `DreamX-World/pipeline/pipeline_dreamxworld.py`; call path in `DreamX-World/inference_dreamx5b.py` | `tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py`; `tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py`; `tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py` | DreamX PipelineConfig wires first-scope DiT/VAE/UMT5/Flow/TI2V settings, official `shift=3.0`, default preset values, FlowMatch scheduler initialization, and local `model_index.json` resolution; independent pipeline smoke covers real CUDA local load + latent generation, and parity compares public API output to worker-side explicit ForwardBatch execution. | pipeline_smoke_and_parity_pass | -| ar_transformer | `DreamX-World/wan/modules/causal_camera_model_2_2_prope_infinity.py`; instantiated by `DreamX-World/inference_ar_forcing.py` as `CausalWanModel` with `local_attn_size=12`, `sink_size=3`, `attn_compress=4` | `tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py` | Native `DreamXWorldARTransformer3DModel` keeps official identity key layout; tiny official-vs-FastVideo forward parity passes; real 5B AR safetensors strict-load passes from `/tmp/converted_dreamx_world_ar`. | tiny_forward_parity_and_real_strict_load_pass | -| ar_pipeline | `DreamX-World/pipeline/pipeline_causal_camera.py`; AR block/KV/context-noise loop | `tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py`; A40 short full-generation smoke | Dedicated `DreamXWorldARCausalDenoisingStage` implements blockwise KV forcing; raw HF repo must be converted before `VideoGenerator.from_pretrained` because it has no `model_index.json`; converted AR layout generated a 64x64/9-frame/4-step MP4 on A40. | config_registry_and_short_full_generation_pass | - -Include reused components in parity. Reuse is accepted only after the FastVideo -component definition and official instantiation arguments have both been checked -and the component parity test passes non-skip. - -Run the relevant tests with: - -```bash -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py -v -s -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py -v -s -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py -v -s -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py -v -s -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py -v -s -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -v -s -python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_conversion.py -v -s -python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py -q -rs -``` - -## Current Local Results - -```bash -python -m pytest tests/local_tests/dreamx_world/ -v -s -# 2026-07-01: 26 passed, 0 skipped -# 2026-07-02: AR targeted suite passed: 16 passed, 0 skipped - -PYTHONPATH=/workspace/FastVideo python /tmp/run_dreamx_ar_video_smoke.py -# 2026-07-02: DreamX-World-5B AR short full-generation smoke passed on A40 -# output: outputs_video/dreamx_world_ar_video_smoke/a quiet road through a futuristic coastal city at sunrise.mp4 -# decoded: 9 frames, (64, 64, 3), uint8 - -PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA python /tmp/run_dreamx_ar_long_horizon.py -# 2026-07-02: DreamX-World-5B AR long-horizon generation passed on A40 -# output: outputs_video/dreamx_world_ar_long_horizon/A long autonomous drive through a futuristic coastal city at sunrise, smooth forward camera motion,.mp4 -# decoded: 1005 frames, (64, 64, 3) - -PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download -# 2026-07-02: DreamX-World-5B AR default SSIM passed: 1 passed, 0 skipped -# local A40 reference: fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/ - -modal run /tmp/modal_dreamx_ar_ssim_git.py -# 2026-07-02: Modal L40S default SSIM passed: 1 passed, 0 skipped -# checked out post-fix dreamx-world-5b-cam branch commit -# first downloaded default L40S references with reference_videos_cli.py download -# seeded missing AR L40S reference, then reran a fresh generated-vs-reference compare -# Modal JSON: mean_ssim=1.0, min_ssim=1.0, max_ssim=1.0 -# local L40S reference: fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos/DreamX-World-5B/TORCH_SDPA/ -# A40-vs-L40S reference spot check after deterministic AR noise fix: mean SSIM 0.7770, min 0.6139, max 0.9944 - -python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py -q -rs -# 2026-06-30: 4 passed, 0 skipped -# 2026-07-01: smoke image-path TI2V coverage passed separately with 3 passed, 0 skipped - -PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B-Cam python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs -# 2026-07-01: 1 passed, 0 skipped -``` - -`camera_conditioning`, default `Flow` scheduler, DreamX component/pipeline configs, default preset, pipeline entry/registry, FlowMatch scheduler initialization, DreamX camera stage, denoising `y_camera` pass-through, conversion model-index/config-consistency checks, real converted 5B dedicated DreamX transformer strict-load and forward parity on CUDA, VAE encode parity on CUDA, text hidden-state parity on CUDA, native DreamX PRoPE branch structure smoke, AR transformer parity/strict-load, AR config/registry, and AR short full-generation smoke are non-skip PASS results. The full local DreamX component suite, independent pipeline smoke/parity suite, image-path TI2V smoke, and SSIM quality regression currently have zero skips. The basic 5B-Cam example was run against `converted_weights/dreamx_world` with a 64x64/9-frame/1-step saved-video smoke, and imageio decoded the generated MP4 successfully; the AR pipeline separately saved and decoded a 64x64/9-frame/4-step MP4 from `/tmp/converted_dreamx_world_ar`. - -## Review Notes - -- Required before handoff: non-skip PASS for each required component parity - test, including reused components that own weights or numerical behavior. -- First PR scope originally targeted `DreamX-World-5B-Cam`; scope was later - expanded to include `DreamX-World-5B` AR support with a separate causal/KV - pipeline. -- AR support has targeted parity/config coverage, short full-generation smoke, - 1005-frame A40 long-horizon generation, local A40 default SSIM coverage, - and Modal L40S default SSIM coverage. L40S validation was rerun after fixing - stale generated-output comparison in the SSIM helper and deterministic AR - noise seeding. HF reference publication remains a separate token-gated - release operation. -- FastVideo production code must not require `xfuser` or OpenCV just because the - official reference import needed them. Port camera/action preprocessing and - sequence-parallel behavior into existing FastVideo-native utilities or keep - reference-only imports inside local parity tests. -- Review agents should verify setup commands still match the PR, then run the - listed parity tests or report the exact blocker. - - -## Quality Regression - -Quality regression is added in `fastvideo/tests/ssim/test_dreamx_world_similarity.py`. The 5B-Cam default test uses the workspace converted model root, `TORCH_SDPA`, a deterministic generated conditioning image, 9 frames, 1 denoise step, seed 1024, and min SSIM 0.98. A local A40 reference was seeded under `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B-Cam/TORCH_SDPA/`, and the test passed non-skip locally on 2026-07-01. The AR default test uses `/tmp/converted_dreamx_world_ar` or `/root/data/dreamx_world_ar_converted`, `TORCH_SDPA`, 192x192, 81 frames, 4 steps, seed 2048, and min SSIM 0.98; local A40 and Modal L40S references were seeded under `fastvideo/tests/ssim/reference_videos/default/{A40,L40S}_reference_videos/DreamX-World-5B/TORCH_SDPA/`, and both tests passed non-skip on 2026-07-02 after fresh generated-output cleanup was added to the SSIM helper. Modal L40S validation checked out `post-fix dreamx-world-5b-cam branch commit`, seeded the missing AR L40S reference, reran the test, and wrote `mean_ssim=1.0`. A cross-device A40-vs-L40S reference spot check produced mean SSIM 0.7770, so CI should use the L40S-specific reference rather than the A40 artifact. Full-quality params are present for 5B-Cam 480x832/161 frames/30 steps and AR 704x1280/1005 frames/4 steps; publishing references to the HF dataset remains a release operation requiring a write-capable HF token env var, never a raw token value. diff --git a/tests/local_tests/dreamx_world/__init__.py b/tests/local_tests/dreamx_world/__init__.py deleted file mode 100644 index 988131360..000000000 --- a/tests/local_tests/dreamx_world/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py b/tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py deleted file mode 100644 index f80c41917..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World-5B autoregressive conversion smoke tests.""" - -import json - -from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_ar_dit_config -from fastvideo.models.registry import _LEGACY_FAST_VIDEO_MODELS -from scripts.checkpoint_conversion.dreamx_world_ar_to_diffusers import ( - MODEL_INDEX, - REUSED_COMPONENTS, - TRANSFORMER_CONFIG, - convert_transformer, - write_model_index, -) - - -def test_dreamx_world_ar_converter_writes_symlinked_transformer_and_model_index(tmp_path): - source = tmp_path / "raw" - source.mkdir() - raw_tensor = source / "model.safetensors" - raw_tensor.write_bytes(b"placeholder") - component_source = tmp_path / "wan22" - output = tmp_path / "dreamx_ar" - - for component in REUSED_COMPONENTS: - component_dir = component_source / component - component_dir.mkdir(parents=True) - (component_dir / "config.json").write_text("{}\n") - - convert_transformer(source, output, symlink_transformer=True) - write_model_index(output, component_source, symlink_components=True) - - assert (output / "transformer" / "model.safetensors").is_symlink() - model_index = json.loads((output / "model_index.json").read_text()) - assert model_index == MODEL_INDEX - assert model_index["_class_name"] == "DreamXWorldARPipeline" - assert model_index["transformer"] == ["diffusers", "DreamXWorldARTransformer3DModel"] - - -def test_dreamx_world_ar_transformer_config_matches_pipeline_dit_config(): - dit_config = make_dreamx_world_5b_ar_dit_config() - assert TRANSFORMER_CONFIG["_class_name"] == "DreamXWorldARTransformer3DModel" - assert TRANSFORMER_CONFIG["num_attention_heads"] == dit_config.num_attention_heads - assert TRANSFORMER_CONFIG["attention_head_dim"] == dit_config.attention_head_dim - assert TRANSFORMER_CONFIG["in_channels"] == dit_config.in_channels - assert TRANSFORMER_CONFIG["out_channels"] == dit_config.out_channels - assert TRANSFORMER_CONFIG["ffn_dim"] == dit_config.ffn_dim - assert TRANSFORMER_CONFIG["num_layers"] == dit_config.num_layers - assert TRANSFORMER_CONFIG["local_attn_size"] == dit_config.local_attn_size - assert TRANSFORMER_CONFIG["sink_size"] == dit_config.sink_size - assert TRANSFORMER_CONFIG["attn_compress"] == dit_config.attn_compress - assert tuple(TRANSFORMER_CONFIG["cam_self_attn_layers"]) == dit_config.cam_self_attn_layers - - -def test_dreamx_world_ar_model_index_component_classes_are_registered(): - for component in ("scheduler", "text_encoder", "transformer", "vae"): - class_name = MODEL_INDEX[component][1] - assert class_name in _LEGACY_FAST_VIDEO_MODELS diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py deleted file mode 100644 index 6d0f59518..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py +++ /dev/null @@ -1,189 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World-5B autoregressive transformer parity. - -Coverage scope: both. The tiny forward parity compares FastVideo's native AR DiT -against the official DreamX ``CausalWanModel`` implementation with identical -weights. The real 5B checkpoint gate strict-loads the downloaded safetensors -through the FastVideo model class. -""" -from __future__ import annotations - -import os -from pathlib import Path -import sys - -import pytest -import torch -from torch.testing import assert_close - -from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARArchConfig, DreamXWorldARConfig -from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_ar_dit_config -from fastvideo.models.dits.dreamx_world_ar import DreamXWorldARTransformer3DModel -from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict -from fastvideo.models.loader.utils import get_param_names_mapping -from fastvideo.models.loader.weight_utils import resolve_safetensors_files, safetensors_weights_iterator - -REPO_ROOT = Path(__file__).resolve().parents[3] -OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) -CONVERTED_AR_DIR = Path(os.getenv("DREAMX_WORLD_AR_CONVERTED_DIR", "/tmp/converted_dreamx_world_ar")) -PARITY_SCOPE = "both" - - -def _tiny_config() -> DreamXWorldARConfig: - return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig( - num_attention_heads=1, - attention_head_dim=8, - in_channels=4, - out_channels=4, - ffn_dim=16, - num_layers=1, - text_dim=8, - freq_dim=8, - text_len=4, - local_attn_size=2, - sink_size=1, - attn_compress=1, - cam_self_attn_layers=(0, ), - )) - - -def _load_official_tiny(): - if not OFFICIAL_REF_DIR.exists(): - pytest.skip(f"Official DreamX reference missing: {OFFICIAL_REF_DIR}") - sys.path.insert(0, str(OFFICIAL_REF_DIR)) - try: - from wan.modules import attention as official_attention - from wan.modules import causal_camera_model_2_2_prope_infinity as causal_module - from wan.modules import model_2_2 as official_model_2_2 - from wan.modules.causal_camera_model_2_2_prope_infinity import CausalWanModel - except Exception as exc: # noqa: BLE001 - pytest.skip(f"Cannot import official AR transformer: {exc}") - official_attention.FLASH_ATTN_2_AVAILABLE = False - official_attention.FLASH_ATTN_3_AVAILABLE = False - - def _sdpa_same_dtype(q, k, v, **kwargs): - del kwargs - out = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), - k.transpose(1, 2), - v.transpose(1, 2), - dropout_p=0.0) - return out.transpose(1, 2).contiguous() - - official_attention.attention = _sdpa_same_dtype - official_attention.flash_attention = _sdpa_same_dtype - official_model_2_2.flash_attention = _sdpa_same_dtype - causal_module.attention = _sdpa_same_dtype - return CausalWanModel( - model_type="ti2v", - patch_size=(1, 2, 2), - text_len=4, - in_dim=4, - dim=8, - ffn_dim=16, - freq_dim=8, - text_dim=8, - out_dim=4, - num_heads=1, - num_layers=1, - local_attn_size=2, - sink_size=1, - qk_norm=True, - cross_attn_norm=True, - add_control_adapter=True, - cam_method="prope", - attn_compress=1, - cam_self_attn_layers=(0, ), - ).eval() - - -def _make_inputs(): - torch.manual_seed(123) - x = [torch.randn(4, 1, 4, 4)] - t = torch.zeros(1, 4, dtype=torch.long) - context = [torch.randn(2, 8)] - camera = { - "viewmats": torch.eye(4).reshape(1, 1, 4, 4).repeat(1, 4, 1, 1), - "K": torch.eye(3).reshape(1, 1, 3, 3).repeat(1, 4, 1, 1), - } - kv_cache = [{ - "k": torch.zeros(1, 8, 1, 8), - "v": torch.zeros(1, 8, 1, 8), - "global_end_index": torch.tensor([0]), - "local_end_index": torch.tensor([0]), - "prope_k": torch.zeros(1, 8, 1, 8), - "prope_v": torch.zeros(1, 8, 1, 8), - "prope_global_end_index": torch.tensor([0]), - "prope_local_end_index": torch.tensor([0]), - }] - cross_cache = [{ - "k": torch.zeros(1, 4, 1, 8), - "v": torch.zeros(1, 4, 1, 8), - "is_init": False, - }] - return x, t, context, camera, kv_cache, cross_cache - - -def test_dreamx_world_ar_tiny_forward_matches_official(): - official = _load_official_tiny() - # The official init_weights zero-inits the output head (head.head.weight and - # biases), so both models would output exactly zero and the comparison would - # pass vacuously. Randomize the head (deterministically) before copying the - # state dict so the outputs reflect the internal computation. - generator = torch.Generator().manual_seed(7) - with torch.no_grad(): - official.head.head.weight.normal_(std=0.5, generator=generator) - official.head.head.bias.normal_(std=0.5, generator=generator) - fastvideo = DreamXWorldARTransformer3DModel(_tiny_config(), {}).eval() - fastvideo.load_state_dict(official.state_dict(), strict=True) - - x, t, context, camera, kv_cache, cross_cache = _make_inputs() - official_out = official(x=x, - t=t, - context=context, - seq_len=4, - y_camera=camera, - kv_cache=kv_cache, - crossattn_cache=cross_cache).detach() - x, t, context, camera, kv_cache, cross_cache = _make_inputs() - fastvideo_out = fastvideo(x=x, - t=t, - context=context, - seq_len=4, - y_camera=camera, - kv_cache=kv_cache, - crossattn_cache=cross_cache).detach() - assert official_out.abs().max() > 0, "official output is all-zero; parity comparison is vacuous" - assert_close(fastvideo_out, official_out, atol=1e-5, rtol=1e-5) - - -def test_dreamx_world_ar_5b_config_matches_official_shape(): - config = make_dreamx_world_5b_ar_dit_config() - assert config.num_layers == 30 - assert config.num_attention_heads == 24 - assert config.attention_head_dim == 128 - assert config.hidden_size == 3072 - assert config.ffn_dim == 14336 - assert config.local_attn_size == 12 - assert config.sink_size == 3 - assert config.attn_compress == 4 - assert config.cam_self_attn_layers == tuple(range(30)) - - -def test_dreamx_world_ar_converted_5b_transformer_strict_loads(): - transformer_dir = CONVERTED_AR_DIR / "transformer" - if not transformer_dir.exists(): - pytest.skip(f"Converted AR transformer missing: {transformer_dir}") - with torch.device("meta"): - model = DreamXWorldARTransformer3DModel(make_dreamx_world_5b_ar_dit_config(), {}) - incompatible = load_model_from_full_model_state_dict( - model, - safetensors_weights_iterator(resolve_safetensors_files(str(transformer_dir)), to_cpu=True), - device=torch.device("cpu"), - param_dtype=torch.bfloat16, - strict=True, - param_names_mapping=get_param_names_mapping(model.param_names_mapping), - training_mode=False, - ) - assert incompatible.missing_keys == [] - assert incompatible.unexpected_keys == [] - assert not any(param.is_meta for param in model.parameters()) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py deleted file mode 100644 index 09b9e5c2a..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py +++ /dev/null @@ -1,95 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World camera-conditioning parity against the official reference. - -Coverage scope: implementation_subcomponent. This verifies the weightless -action-sequence to PRoPE camera tensor path used by DreamX-World-5B-Cam. -""" -from __future__ import annotations - -import os -from pathlib import Path -import pytest -import torch -from torch.testing import assert_close - -from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import build_dreamx_camera_condition - -REPO_ROOT = Path(__file__).resolve().parents[3] -OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) -PARITY_SCOPE = "implementation_subcomponent" - - -def _load_official_functions(): - if not OFFICIAL_REF_DIR.exists(): - pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}") - try: - import importlib.util - - pose_path = OFFICIAL_REF_DIR / "utils" / "pose_utils.py" - pose_spec = importlib.util.spec_from_file_location("dreamx_world_pose_utils", pose_path) - if pose_spec is None or pose_spec.loader is None: - raise RuntimeError(f"Cannot load DreamX pose_utils: {pose_path}") - pose_module = importlib.util.module_from_spec(pose_spec) - pose_spec.loader.exec_module(pose_module) - - source = (OFFICIAL_REF_DIR / "utils" / "inference_utils.py").read_text() - source = source.replace("from .pose_utils import interpolate_camera_poses\n", "") - namespace = {"interpolate_camera_poses": pose_module.interpolate_camera_poses} - exec(compile(source, str(OFFICIAL_REF_DIR / "utils" / "inference_utils.py"), "exec"), namespace) - except Exception as exc: # noqa: BLE001 - local parity should skip missing reference deps. - pytest.skip(f"Cannot load DreamX camera reference: {exc}") - return namespace["ActionToPoseFromID"], namespace["GetPoseEmbedsFromPosesPrope"] - - -def _official_camera_condition( - action_seq: list[str], - action_speed_list: list[float], - *, - num_frames: int, - height: int, - width: int, - dtype: torch.dtype, -): - action_to_pose, get_pose_embeds = _load_official_functions() - duration = -(-num_frames // len(action_seq)) - poses = action_to_pose(action_seq, action_speed_list, duration=duration)[:num_frames] - condition, _ = get_pose_embeds(poses, height, width, len(poses), False, 0, dtype=dtype, device="cpu") - return condition - - -@pytest.mark.parametrize( - ("action_seq", "action_speed_list", "num_frames"), - [ - (["w"], [4], 81), - (["wj", "d"], [4, 6], 121), - (["i", "k", "l"], [3, 5, 2], 85), - ], -) -def test_dreamx_world_camera_conditioning_matches_official(action_seq, action_speed_list, num_frames): - dtype = torch.float32 - official = _official_camera_condition( - action_seq, - action_speed_list, - num_frames=num_frames, - height=704, - width=1280, - dtype=dtype, - ) - fastvideo = build_dreamx_camera_condition( - action_seq, - action_speed_list, - num_frames=num_frames, - height=704, - width=1280, - dtype=dtype, - device="cpu", - ) - - assert official.keys() == fastvideo.keys() == {"viewmats", "K"} - for key in ("viewmats", "K"): - assert official[key].shape == fastvideo[key].shape - diff = (official[key] - fastvideo[key]).abs() - print( - f"{key}: shape={tuple(fastvideo[key].shape)} diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}" - ) - assert_close(fastvideo[key], official[key], atol=1e-5, rtol=1e-5) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py b/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py deleted file mode 100644 index dbc3a26ab..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World conversion script smoke tests.""" - -import json - -from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_dit_config -from fastvideo.models.registry import _LEGACY_FAST_VIDEO_MODELS - -from scripts.checkpoint_conversion.dreamx_world_to_diffusers import ( - MODEL_INDEX, - REUSED_COMPONENTS, - TRANSFORMER_CONFIG, - _copy_or_link_component, - write_model_index, -) - - -def test_dreamx_world_converter_writes_full_model_index_with_reused_components(tmp_path): - component_source = tmp_path / "wan22" - output = tmp_path / "dreamx" - output.mkdir() - - for component in REUSED_COMPONENTS: - component_dir = component_source / component - component_dir.mkdir(parents=True) - (component_dir / "config.json").write_text("{}\n") - - write_model_index(output, component_source, symlink_components=True) - - model_index = json.loads((output / "model_index.json").read_text()) - assert model_index == MODEL_INDEX - assert model_index["_class_name"] == "DreamXWorldPipeline" - assert model_index["transformer"] == ["diffusers", "DreamXWorldTransformer3DModel"] - for component in REUSED_COMPONENTS: - assert (output / component).is_symlink() - - -def test_dreamx_world_transformer_config_has_camera_adapter_enabled(): - assert TRANSFORMER_CONFIG["_class_name"] == "DreamXWorldTransformer3DModel" - assert TRANSFORMER_CONFIG["add_control_adapter"] is True - assert TRANSFORMER_CONFIG["cam_method"] == "prope" - assert TRANSFORMER_CONFIG["num_layers"] == 30 - - -def test_dreamx_world_converter_transformer_config_matches_pipeline_dit_config(): - dit_config = make_dreamx_world_5b_cam_dit_config() - - assert TRANSFORMER_CONFIG["num_attention_heads"] == dit_config.num_attention_heads - assert TRANSFORMER_CONFIG["attention_head_dim"] == dit_config.attention_head_dim - assert TRANSFORMER_CONFIG["in_channels"] == dit_config.in_channels - assert TRANSFORMER_CONFIG["out_channels"] == dit_config.out_channels - assert TRANSFORMER_CONFIG["ffn_dim"] == dit_config.ffn_dim - assert TRANSFORMER_CONFIG["num_layers"] == dit_config.num_layers - assert TRANSFORMER_CONFIG["cross_attn_norm"] == dit_config.cross_attn_norm - assert TRANSFORMER_CONFIG["qk_norm"] == dit_config.qk_norm - assert TRANSFORMER_CONFIG["add_control_adapter"] == dit_config.add_control_adapter - assert TRANSFORMER_CONFIG["cam_method"] == dit_config.cam_method - assert TRANSFORMER_CONFIG["attn_compress"] == dit_config.attn_compress - assert TRANSFORMER_CONFIG["cam_self_attn_layers"] == dit_config.cam_self_attn_layers - - -def test_dreamx_world_model_index_component_classes_are_registered(): - for component in ("scheduler", "text_encoder", "transformer", "vae"): - class_name = MODEL_INDEX[component][1] - assert class_name in _LEGACY_FAST_VIDEO_MODELS - - -def test_dreamx_world_copy_or_link_component_keeps_broken_symlink(tmp_path): - component_source = tmp_path / "wan22" - src = component_source / "scheduler" - src.mkdir(parents=True) - output = tmp_path / "dreamx" - output.mkdir() - dst = output / "scheduler" - dst.symlink_to(tmp_path / "missing_scheduler", target_is_directory=True) - - _copy_or_link_component("scheduler", component_source, output, symlink=True) - - assert dst.is_symlink() diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py deleted file mode 100644 index da8f8bf53..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py +++ /dev/null @@ -1,324 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World pipeline config and conditioning smoke tests.""" - -from types import SimpleNamespace -import json - -import numpy as np -import pytest -import torch - -from fastvideo.api.presets import get_preset -from fastvideo.fastvideo_args import WorkloadType -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.pipeline_registry import PipelineType, import_pipeline_classes -from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import ( - DreamXCamera, - _interpolate_camera_poses, - build_dreamx_camera_condition, -) -from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig -from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage -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, -) -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.registry import get_default_preset, get_model_info, get_pipeline_config_cls_from_name - - -@pytest.fixture -def dreamx_model_manifests(monkeypatch): - """Registry/config checks do not need network access or model weights.""" - names = { - "GD-ML/DreamX-World-5B-Cam": "DreamXWorldPipeline", - "GD-ML/DreamX-World-5B": "DreamXWorldARPipeline", - } - - def manifest(model_path, revision=None): - name = names[model_path] - return {"_class_name": name, "pipeline_name": name, "_diffusers_version": "0.31.0"} - - monkeypatch.setattr("fastvideo.registry.maybe_download_model_index", manifest) - - -def test_dreamx_world_5b_cam_pipeline_config_wires_first_scope_components(): - config = DreamXWorld5BCamPipelineConfig() - - assert config.flow_shift == 3.0 - assert config.ti2v_task is True - assert config.expand_timesteps is True - assert config.dit_config.expand_timesteps is True - assert config.dit_config.num_layers == 30 - assert config.dit_config.add_control_adapter is True - assert config.dit_config.cam_method == "prope" - - assert config.vae_config.load_encoder is True - assert config.vae_config.load_decoder is True - assert config.vae_config.z_dim == 48 - assert config.vae_config.scale_factor_temporal == 4 - assert config.vae_config.scale_factor_spatial == 16 - - assert len(config.text_encoder_configs) == 1 - text_config = config.text_encoder_configs[0] - assert text_config.prefix == "umt5" - assert text_config.vocab_size == 256384 - assert text_config.d_model == 4096 - assert config.text_encoder_precisions == ("bf16", ) - - -def test_dreamx_world_pipeline_registry_discovers_entrypoint(): - pipelines = import_pipeline_classes(PipelineType.BASIC) - - assert pipelines["basic"]["DreamXWorldPipeline"] is DreamXWorldPipeline - - -def test_dreamx_world_local_model_index_resolves_model_info(tmp_path): - model_dir = tmp_path / "DreamX-World-5B-Cam-converted" - model_dir.mkdir() - for component in ("scheduler", "text_encoder", "tokenizer", "transformer", "vae"): - (model_dir / component).mkdir() - (model_dir / "model_index.json").write_text( - json.dumps({ - "_class_name": "DreamXWorldPipeline", - "_diffusers_version": "0.31.0", - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["transformers", "UMT5EncoderModel"], - "tokenizer": ["transformers", "AutoTokenizer"], - "transformer": ["diffusers", "DreamXWorldTransformer3DModel"], - "vae": ["diffusers", "AutoencoderKLWan"], - }) + "\n") - - info = get_model_info(str(model_dir), pipeline_type=PipelineType.BASIC, workload_type=WorkloadType.I2V) - - assert info.pipeline_cls is DreamXWorldPipeline - assert info.pipeline_config_cls is DreamXWorld5BCamPipelineConfig - - -def test_dreamx_world_model_path_resolves_pipeline_config(dreamx_model_manifests): - assert get_pipeline_config_cls_from_name("GD-ML/DreamX-World-5B-Cam") is DreamXWorld5BCamPipelineConfig - - -def test_dreamx_world_default_preset_is_registered(dreamx_model_manifests): - preset_name = get_default_preset("GD-ML/DreamX-World-5B-Cam") - preset = get_preset(preset_name, "dreamx_world") - - assert preset.name == "dreamx_world_5b_cam" - assert preset.workload_type == "i2v" - assert preset.defaults["height"] == 480 - assert preset.defaults["width"] == 832 - assert preset.defaults["num_frames"] == 161 - assert preset.defaults["num_inference_steps"] == 30 - assert preset.defaults["guidance_scale"] == 5.0 - - -def test_dreamx_world_pipeline_initializes_official_flow_scheduler(): - pipeline = DreamXWorldPipeline.__new__(DreamXWorldPipeline) - pipeline.modules = {} - fastvideo_args = SimpleNamespace(pipeline_config=DreamXWorld5BCamPipelineConfig()) - - pipeline.initialize_pipeline(fastvideo_args) - - scheduler = pipeline.modules["scheduler"] - assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler) - assert scheduler.config.shift == 3.0 - - -def test_dreamx_world_camera_conditioning_stage_sets_y_camera_extra(): - batch = ForwardBatch( - data_type="t2v", - action_list=["wj", "d"], - action_speed_list=[4, 6], - num_frames=17, - height=704, - width=1280, - latents=torch.zeros(1, 16, 5, 44, 80), - ) - stage = DreamXWorldCameraConditioningStage() - - out = stage.forward(batch, fastvideo_args=object()) - y_camera = out.extra[DREAMX_Y_CAMERA_KEY] - expected = build_dreamx_camera_condition( - ["wj", "d"], - [4, 6], - num_frames=17, - height=704, - width=1280, - dtype=torch.float32, - device="cpu", - ) - - assert set(y_camera) == {"viewmats", "K"} - for key, expected_value in expected.items(): - assert y_camera[key].shape == (1, *expected_value.shape) - torch.testing.assert_close(y_camera[key][0], expected_value) - assert stage.verify_output(out, fastvideo_args=object()).is_valid() - - -def test_dreamx_world_denoising_kwargs_filter_for_y_camera(): - y_camera = {"viewmats": torch.eye(4).reshape(1, 1, 4, 4), "K": torch.eye(3).reshape(1, 1, 3, 3)} - stage = DenoisingStage.__new__(DenoisingStage) - - def accepts_y_camera(hidden_states, encoder_hidden_states, timestep, y_camera=None): - return y_camera - - def no_y_camera(hidden_states, encoder_hidden_states, timestep): - return hidden_states - - assert stage.prepare_extra_func_kwargs(accepts_y_camera, {"y_camera": y_camera}) == {"y_camera": y_camera} - assert stage.prepare_extra_func_kwargs(no_y_camera, {"y_camera": y_camera}) == {} - - -def test_dreamx_world_ar_pipeline_config_wires_components(): - config = DreamXWorld5BARPipelineConfig() - assert config.is_causal is True - assert config.flow_shift == 5.0 - assert config.dmd_denoising_steps == (1000, 750, 500, 250) - assert config.warp_denoising_step is True - assert config.context_noise == 0.1 - assert config.dit_config.arch_config.local_attn_size == 12 - assert config.dit_config.arch_config.sink_size == 3 - assert config.dit_config.arch_config.attn_compress == 4 - - -def test_dreamx_world_ar_pipeline_registry_and_preset(dreamx_model_manifests): - from fastvideo.api.presets import get_preset - from fastvideo.registry import get_pipeline_config_cls_from_name, get_preset_selection - - assert get_pipeline_config_cls_from_name("GD-ML/DreamX-World-5B") is DreamXWorld5BARPipelineConfig - preset_name, family = get_preset_selection("GD-ML/DreamX-World-5B") - assert (preset_name, family) == ("dreamx_world_5b_ar", "dreamx_world") - preset = get_preset("dreamx_world_5b_ar", "dreamx_world") - assert preset.defaults["num_inference_steps"] == 4 - assert DreamXWorldARPipeline.pipeline_config_cls is DreamXWorld5BARPipelineConfig - - -def test_dreamx_world_camera_conditioning_stage_expands_scalar_speed(): - batch = ForwardBatch( - data_type="t2v", - action_list=["w", "d"], - action_speed_list=2.0, - num_frames=17, - height=704, - width=1280, - latents=torch.zeros(1, 16, 5, 44, 80), - ) - stage = DreamXWorldCameraConditioningStage() - - out = stage.forward(batch, fastvideo_args=object()) - - assert set(out.extra[DREAMX_Y_CAMERA_KEY]) == {"viewmats", "K"} - - -def test_dreamx_world_camera_interpolation_handles_single_camera(): - camera = DreamXCamera( - fx=0.8, - fy=0.8, - cx=0.5, - cy=0.5, - w2c_mat=np.eye(4, dtype=np.float64), - ) - - out = _interpolate_camera_poses( - [camera], - src_indices=np.array([0.0]), - tgt_indices=np.array([0.0, 1.0, 2.0]), - ) - - assert out == [camera, camera, camera] - - -def test_dreamx_world_ar_cache_initializes_camera_self_attention_entries(): - cam_self_attn = SimpleNamespace(num_heads=3, head_dim=5) - transformer = SimpleNamespace( - blocks=[SimpleNamespace(cam_self_attn=cam_self_attn), - SimpleNamespace(cam_self_attn=cam_self_attn)], - num_attention_heads=2, - attention_head_dim=4, - ) - stage = DreamXWorldARCausalDenoisingStage.__new__(DreamXWorldARCausalDenoisingStage) - stage.transformer = transformer - stage.num_transformer_blocks = 2 - stage.local_attn_size = 6 - - caches = stage._initialize_kv_cache( - batch_size=1, - dtype=torch.float32, - device=torch.device("cpu"), - frame_seq_length=7, - ) - - assert len(caches) == 2 - assert caches[0]["k"].shape == (1, 42, 2, 4) - assert caches[0]["prope_k"].shape == (1, 42, 3, 5) - assert caches[0]["prope_v"].shape == (1, 42, 3, 5) - assert int(caches[0]["prope_global_end_index"].item()) == 0 - assert int(caches[0]["prope_local_end_index"].item()) == 0 - - -def test_dreamx_world_ar_context_noise_fraction_maps_to_scheduler_timestep(): - assert DreamXWorldARCausalDenoisingStage._context_noise_timestep(0.1) == 100 - assert DreamXWorldARCausalDenoisingStage._context_noise_timestep(100) == 100 - - -def test_dreamx_world_ar_context_update_advances_camera_cache_indices(): - - class DummyTransformer: - - def __call__(self, *, hidden_states, encoder_hidden_states, timestep, y_camera, kv_cache, crossattn_cache, - current_start): - del encoder_hidden_states, y_camera, crossattn_cache - assert current_start == 0 - assert timestep.unique().tolist() == [100] - new_tokens = timestep.shape[1] - for cache in kv_cache: - cache["local_end_index"] += new_tokens - cache["global_end_index"] += new_tokens - cache["prope_local_end_index"] += new_tokens - cache["prope_global_end_index"] += new_tokens - cache["k"][:, :new_tokens] = 1 - cache["prope_k"][:, :new_tokens] = 1 - return hidden_states - - cam_self_attn = SimpleNamespace(num_heads=3, head_dim=5) - cache_transformer = SimpleNamespace( - blocks=[SimpleNamespace(cam_self_attn=cam_self_attn)], - num_attention_heads=2, - attention_head_dim=4, - ) - stage = DreamXWorldARCausalDenoisingStage.__new__(DreamXWorldARCausalDenoisingStage) - stage.transformer = cache_transformer - stage.num_transformer_blocks = 1 - stage.local_attn_size = 6 - caches = stage._initialize_kv_cache( - batch_size=1, - dtype=torch.float32, - device=torch.device("cpu"), - frame_seq_length=2, - ) - # Keep the cache allocation source separate from the callable transformer used by _update_context_cache. - stage.transformer = DummyTransformer() - - stage._update_context_cache( - block_latents=torch.zeros(1, 4, 3, 2, 2), - context=[torch.zeros(2, 4)], - camera_block={ - "viewmats": torch.eye(4).reshape(1, 1, 4, 4), - "K": torch.eye(3).reshape(1, 1, 3, 3) - }, - kv_cache=caches, - crossattn_cache=[{}], - start=0, - frame_seq_length=2, - target_dtype=torch.float32, - autocast_enabled=False, - context_noise=0.1, - ) - - assert int(caches[0]["local_end_index"].item()) == 6 - assert int(caches[0]["prope_local_end_index"].item()) == 6 - assert caches[0]["k"][:, :6].sum().item() == 48 - assert caches[0]["prope_k"][:, :6].sum().item() == 90 diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py deleted file mode 100644 index 367542995..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py +++ /dev/null @@ -1,52 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World default Flow scheduler parity. - -Coverage scope: implementation_subcomponent. DreamX-World-5B-Cam defaults to -Diffusers FlowMatchEulerDiscreteScheduler for sampler_name=Flow. This test -checks that FastVideo's native FlowMatchEulerDiscreteScheduler matches the -timestep schedule and Euler step used by the official default path. -""" -from __future__ import annotations - -from pathlib import Path -import inspect - -import torch -from diffusers import FlowMatchEulerDiscreteScheduler as OfficialFlowMatchEulerDiscreteScheduler -from omegaconf import OmegaConf -from torch.testing import assert_close - -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler, ) - -REPO_ROOT = Path(__file__).resolve().parents[3] -PARITY_SCOPE = "implementation_subcomponent" - - -def _scheduler_kwargs(cls): - config_path = REPO_ROOT / "DreamX-World" / "configs" / "wan2.2" / "wan_ti2v_5b.yaml" - config = OmegaConf.load(config_path) - raw_kwargs = OmegaConf.to_container(config["scheduler_kwargs"]) - signature = inspect.signature(cls) - return {key: value for key, value in raw_kwargs.items() if key in signature.parameters} - - -def test_dreamx_world_flow_scheduler_timesteps_and_step_match(): - official = OfficialFlowMatchEulerDiscreteScheduler(**_scheduler_kwargs(OfficialFlowMatchEulerDiscreteScheduler)) - fastvideo = FastVideoFlowMatchEulerDiscreteScheduler(**_scheduler_kwargs(FastVideoFlowMatchEulerDiscreteScheduler)) - - official.set_timesteps(50, device="cpu", mu=1) - fastvideo.set_timesteps(50, device="cpu", mu=1) - assert_close(fastvideo.timesteps, official.timesteps, atol=0, rtol=0) - assert_close(fastvideo.sigmas, official.sigmas, atol=0, rtol=0) - - torch.manual_seed(7) - sample = torch.randn(1, 4, 2, 8, 8) - model_output = torch.randn_like(sample) - timestep = official.timesteps[3] - - official_prev = official.step(model_output, timestep, sample, return_dict=False)[0] - fastvideo_prev = fastvideo.step(model_output, fastvideo.timesteps[3], sample, return_dict=False)[0] - diff = (official_prev - fastvideo_prev).abs() - print(f"scheduler diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") - assert_close(fastvideo_prev, official_prev, atol=0, rtol=0) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py deleted file mode 100644 index 9b3147581..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World Wan T5 encoder reuse parity scaffold. - -Coverage scope: implementation_subcomponent. It records the official -WanT5EncoderModel loading path and FastVideo T5 target for later activation -with staged Wan2.2 base text encoder/tokenizer weights. -""" -from __future__ import annotations - -import os -from pathlib import Path -import sys -import types - -import pytest -import torch -from omegaconf import OmegaConf -from torch.testing import assert_close -from transformers import AutoTokenizer - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.configs.pipelines.dreamx_world import ( - DreamXWorld5BCamPipelineConfig, - make_dreamx_world_5b_cam_text_encoder_config, -) -from fastvideo.models.loader.component_loader import TextEncoderLoader - -REPO_ROOT = Path(__file__).resolve().parents[3] -OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) -WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B")) -WAN_DIFFUSERS_DIR = Path( - os.getenv("DREAMX_WORLD_WAN_DIFFUSERS_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B-Diffusers")) -PARITY_SCOPE = "implementation_subcomponent" - - -def _add_official_to_path(): - if not OFFICIAL_REF_DIR.exists(): - pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}") - if str(OFFICIAL_REF_DIR) not in sys.path: - sys.path.insert(0, str(OFFICIAL_REF_DIR)) - - -def _install_xfuser_stub() -> None: - if "xfuser" in sys.modules: - return - xfuser = types.ModuleType("xfuser") - core = types.ModuleType("xfuser.core") - distributed = types.ModuleType("xfuser.core.distributed") - long_ctx = types.ModuleType("xfuser.core.long_ctx_attention") - distributed.get_sequence_parallel_rank = lambda: 0 - distributed.get_sequence_parallel_world_size = lambda: 1 - distributed.get_sp_group = lambda: None - distributed.get_world_group = lambda: types.SimpleNamespace(local_rank=0, rank=0) - distributed.init_distributed_environment = lambda *args, **kwargs: None - distributed.initialize_model_parallel = lambda *args, **kwargs: None - distributed.model_parallel_is_initialized = lambda: False - - class XFuserLongContextAttention: - - def __call__(self, *args, **kwargs): - raise RuntimeError("xfuser stub cannot execute attention") - - long_ctx.xFuserLongContextAttention = XFuserLongContextAttention - sys.modules.update({ - "xfuser": xfuser, - "xfuser.core": core, - "xfuser.core.distributed": distributed, - "xfuser.core.long_ctx_attention": long_ctx, - }) - - -def _text_kwargs(): - config = OmegaConf.load(OFFICIAL_REF_DIR / "configs" / "wan2.2" / "wan_ti2v_5b.yaml") - return OmegaConf.to_container(config["text_encoder_kwargs"]) - - -def _patch_single_process_text_parallel(monkeypatch): - import fastvideo.layers.linear as fastvideo_linear - import fastvideo.layers.vocab_parallel_embedding as fastvideo_embedding - import fastvideo.models.encoders.t5 as fastvideo_t5 - - for module in (fastvideo_t5, fastvideo_embedding, fastvideo_linear): - if hasattr(module, "get_tp_rank"): - monkeypatch.setattr(module, "get_tp_rank", lambda: 0) - if hasattr(module, "get_tp_world_size"): - monkeypatch.setattr(module, "get_tp_world_size", lambda: 1) - monkeypatch.setattr(fastvideo_embedding, "tensor_model_parallel_all_reduce", lambda x: x) - - -def _load_official_text_encoder(device, dtype): - _add_official_to_path() - _install_xfuser_stub() - text_path = WAN_BASE_DIR / "models_t5_umt5-xxl-enc-bf16.pth" - if not text_path.exists(): - pytest.skip(f"Wan2.2 text encoder weights missing: {text_path}") - try: - from models import WanT5EncoderModel - except Exception as exc: # noqa: BLE001 - pytest.skip(f"Cannot import official DreamX text encoder: {exc}") - model = WanT5EncoderModel.from_pretrained(str(text_path), - additional_kwargs=_text_kwargs(), - low_cpu_mem_usage=True, - torch_dtype=dtype) - return model.to(device=device, dtype=dtype).eval() - - -def _load_fastvideo_text_encoder(device, dtype, monkeypatch): - text_encoder_path = WAN_DIFFUSERS_DIR / "text_encoder" - if not text_encoder_path.exists(): - pytest.skip(f"Wan2.2 Diffusers text encoder missing: {text_encoder_path}") - _patch_single_process_text_parallel(monkeypatch) - pipeline_config = DreamXWorld5BCamPipelineConfig() - pipeline_config.text_encoder_configs[0]._fsdp_shard_conditions = [] - args = FastVideoArgs( - model_path=str(text_encoder_path), - pipeline_config=pipeline_config, - text_encoder_cpu_offload=(device.type == "cpu"), - ) - args.model_paths = {} - return TextEncoderLoader().load(str(text_encoder_path), args).to(device=device, dtype=dtype).eval() - - -def test_dreamx_world_text_encoder_config_matches_umt5_xxl_shape(): - config = make_dreamx_world_5b_cam_text_encoder_config() - assert config.vocab_size == 256384 - assert config.d_model == 4096 - assert config.d_kv == 64 - assert config.d_ff == 10240 - assert config.num_heads == 64 - assert config.num_layers == 24 - assert config.relative_attention_num_buckets == 32 - assert config.dropout_rate == 0.0 - assert config.text_len == 512 - assert config.prefix == "umt5" - - -def test_dreamx_world_fastvideo_text_encoder_loads_staged_weights(monkeypatch): - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - model = _load_fastvideo_text_encoder(device, torch.bfloat16, monkeypatch) - assert model.__class__.__name__ == "UMT5EncoderModel" - assert next(model.parameters()).device.type == device.type - assert next(model.parameters()).dtype == torch.bfloat16 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for text encoder parity.") -def test_dreamx_world_text_encoder_parity_scaffold(monkeypatch): - device = torch.device("cuda:0") - dtype = torch.bfloat16 - official = _load_official_text_encoder(device, dtype) - fastvideo = _load_fastvideo_text_encoder(device, dtype, monkeypatch) - tokenizer_path = WAN_BASE_DIR / "google" / "umt5-xxl" - if not tokenizer_path.exists(): - pytest.skip(f"Wan2.2 tokenizer missing: {tokenizer_path}") - tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_path)) - batch = tokenizer(["A quiet forest trail at sunrise."], padding="max_length", max_length=512, return_tensors="pt") - input_ids = batch.input_ids.to(device) - attention_mask = batch.attention_mask.to(device) - with torch.inference_mode(): - official_hidden = official(input_ids, attention_mask=attention_mask)[0].float().cpu() - fastvideo_hidden = fastvideo(input_ids, attention_mask=attention_mask).last_hidden_state.float().cpu() - assert official_hidden.shape == fastvideo_hidden.shape - assert_close(fastvideo_hidden, official_hidden, atol=1e-3, rtol=1e-3) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py deleted file mode 100644 index f7908fb1b..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py +++ /dev/null @@ -1,330 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World transformer parity scaffold. - -Coverage scope: both. The official side loads DreamX-World-5B-Cam through -Wan2_2Transformer3DModel.from_pretrained with PRoPE camera control enabled. -The FastVideo side strict-loads the converted DreamX transformer weights into -the native DreamX-World DiT implementation. -""" -from __future__ import annotations - -import math -import os -from pathlib import Path -import sys -import types - -import pytest -import torch -from omegaconf import OmegaConf -from torch.testing import assert_close - -from fastvideo.forward_context import set_forward_context -from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldArchConfig, DreamXWorldConfig) -from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import build_dreamx_camera_condition -from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_dit_config -from fastvideo.platforms import AttentionBackendEnum -from fastvideo.models.dits.dreamx_world import (DreamXPropeSelfAttention, DreamXWorldTransformer3DModel, - DreamXWorldTransformerBlock) -from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict -from fastvideo.models.loader.utils import get_param_names_mapping -from fastvideo.models.loader.weight_utils import resolve_safetensors_files, safetensors_weights_iterator -from scripts.checkpoint_conversion.dreamx_world_to_diffusers import map_transformer_key - -REPO_ROOT = Path(__file__).resolve().parents[3] -OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) -LOCAL_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "dreamx_world")) -WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B")) -CONVERTED_WEIGHTS_DIR = Path( - os.getenv("DREAMX_WORLD_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "dreamx_world")) -PARITY_SCOPE = "both" - - -def _install_xfuser_stub() -> None: - if "xfuser" in sys.modules: - return - xfuser = types.ModuleType("xfuser") - core = types.ModuleType("xfuser.core") - distributed = types.ModuleType("xfuser.core.distributed") - long_ctx = types.ModuleType("xfuser.core.long_ctx_attention") - distributed.get_sequence_parallel_rank = lambda: 0 - distributed.get_sequence_parallel_world_size = lambda: 1 - distributed.get_sp_group = lambda: None - distributed.get_world_group = lambda: types.SimpleNamespace(local_rank=0, rank=0) - distributed.init_distributed_environment = lambda *args, **kwargs: None - distributed.initialize_model_parallel = lambda *args, **kwargs: None - distributed.model_parallel_is_initialized = lambda: False - - class XFuserLongContextAttention: - - def __call__(self, *args, **kwargs): - raise RuntimeError("xfuser stub cannot execute attention") - - long_ctx.xFuserLongContextAttention = XFuserLongContextAttention - sys.modules.update({ - "xfuser": xfuser, - "xfuser.core": core, - "xfuser.core.distributed": distributed, - "xfuser.core.long_ctx_attention": long_ctx, - }) - - -def _make_tiny_dreamx_config() -> DreamXWorldConfig: - return DreamXWorldConfig(arch_config=DreamXWorldArchConfig( - num_attention_heads=1, - attention_head_dim=8, - in_channels=16, - out_channels=16, - ffn_dim=32, - num_layers=1, - 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 _add_official_to_path() -> None: - if not OFFICIAL_REF_DIR.exists(): - pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}") - if str(OFFICIAL_REF_DIR) not in sys.path: - sys.path.insert(0, str(OFFICIAL_REF_DIR)) - - -def _official_transformer_kwargs() -> dict: - config_path = OFFICIAL_REF_DIR / "configs" / "wan2.2" / "wan_ti2v_5b.yaml" - if not config_path.exists(): - pytest.skip(f"DreamX Wan config missing: {config_path}") - config = OmegaConf.load(config_path) - kwargs = OmegaConf.to_container(config["transformer_additional_kwargs"]) - kwargs["cam_method"] = "prope" - kwargs["add_control_adapter"] = True - return kwargs - - -def _load_official_transformer(device: torch.device, dtype: torch.dtype) -> torch.nn.Module: - _add_official_to_path() - _install_xfuser_stub() - if not LOCAL_WEIGHTS_DIR.exists(): - pytest.skip(f"DreamX transformer weights missing: {LOCAL_WEIGHTS_DIR}") - try: - from models import Wan2_2Transformer3DModel - except Exception as exc: # noqa: BLE001 - local parity should skip missing refs. - pytest.skip(f"Cannot import official DreamX transformer: {exc}") - model = Wan2_2Transformer3DModel.from_pretrained( - str(LOCAL_WEIGHTS_DIR), - transformer_additional_kwargs=_official_transformer_kwargs(), - torch_dtype=dtype, - ) - return model.to(device=device, dtype=dtype).eval() - - -def _load_fastvideo_transformer(device: torch.device, dtype: torch.dtype) -> torch.nn.Module: - model = _load_fastvideo_transformer_strict(torch.device("cpu"), dtype) - return model.to(device=device, dtype=dtype).eval() - - -def _load_fastvideo_transformer_strict(device: torch.device, dtype: torch.dtype) -> torch.nn.Module: - transformer_dir = CONVERTED_WEIGHTS_DIR / "transformer" - if not transformer_dir.exists(): - pytest.skip(f"Converted DreamX transformer missing: {transformer_dir}") - safetensors_files = resolve_safetensors_files(str(transformer_dir)) - config = make_dreamx_world_5b_cam_dit_config() - import fastvideo.models.dits.dreamx_world as fastvideo_dreamx - original_get_sp_world_size = fastvideo_dreamx.get_sp_world_size - fastvideo_dreamx.get_sp_world_size = lambda: 1 - try: - with torch.device("meta"): - model = DreamXWorldTransformer3DModel(config=config, hf_config={}) - finally: - fastvideo_dreamx.get_sp_world_size = original_get_sp_world_size - incompatible = load_model_from_full_model_state_dict( - model, - safetensors_weights_iterator(safetensors_files, to_cpu=True), - device=device, - param_dtype=dtype, - strict=True, - param_names_mapping=get_param_names_mapping(model.param_names_mapping), - training_mode=False, - ) - assert incompatible.missing_keys == [] - assert incompatible.unexpected_keys == [] - assert not any(param.is_meta for param in model.parameters()) - return model.to(device=device, dtype=dtype).eval() - - -def _make_inputs(device: torch.device, dtype: torch.dtype): - torch.manual_seed(1234) - num_frames = 5 - height = 64 - width = 64 - latent_frames = (num_frames - 1) // 4 + 1 - latent_h = height // 16 - latent_w = width // 16 - x = torch.randn(1, 48, latent_frames, latent_h, latent_w, device=device, dtype=dtype) - context = [torch.randn(16, 4096, device=device, dtype=dtype)] - seq_len = math.ceil((latent_h * latent_w) / 4 * latent_frames) - timestep = torch.full((1, seq_len), 250, device=device, dtype=torch.long) - camera = build_dreamx_camera_condition(["w"], [4], - num_frames=num_frames, - height=height, - width=width, - dtype=dtype, - device=device) - camera = {key: value.unsqueeze(0) for key, value in camera.items()} - return {"x": [x[0]], "context": context, "t": timestep, "seq_len": seq_len, "y_camera": camera} - - -def _run_official(model: torch.nn.Module, inputs: dict) -> torch.Tensor: - with torch.inference_mode(): - output = model(**inputs) - if isinstance(output, list): - output = torch.stack(output, dim=0) - assert torch.is_tensor(output), f"official output is not a tensor: {type(output)}" - return output.detach().float().cpu() - - -def _run_fastvideo(model: torch.nn.Module, inputs: dict) -> torch.Tensor: - hidden_states = torch.stack(inputs["x"], dim=0) - encoder_hidden_states = torch.stack( - [torch.cat([inputs["context"][0], inputs["context"][0].new_zeros(512 - inputs["context"][0].shape[0], 4096)])]) - with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): - output = model( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=inputs["t"], - y_camera=inputs["y_camera"], - ) - assert torch.is_tensor(output), f"FastVideo output is not a tensor: {type(output)}" - return output.detach().float().cpu() - - -def test_dreamx_world_conversion_mapping_strict_load_smoke(monkeypatch): - _add_official_to_path() - _install_xfuser_stub() - try: - from models import Wan2_2Transformer3DModel - except Exception as exc: # noqa: BLE001 - pytest.skip(f"Cannot import official DreamX transformer: {exc}") - - official = Wan2_2Transformer3DModel( - dim=8, - ffn_dim=32, - num_heads=1, - num_layers=1, - add_control_adapter=True, - cam_method="prope", - ) - official_state = official.state_dict() - diffusers_like_state = {map_transformer_key(key): value.detach().clone() for key, value in official_state.items()} - - import fastvideo.models.dits.dreamx_world as fastvideo_dreamx - monkeypatch.setattr(fastvideo_dreamx, "get_sp_world_size", lambda: 1) - with torch.device("meta"): - fastvideo = DreamXWorldTransformer3DModel(config=_make_tiny_dreamx_config(), hf_config={}) - - incompatible = load_model_from_full_model_state_dict( - fastvideo, - iter(diffusers_like_state.items()), - device=torch.device("cpu"), - param_dtype=torch.float32, - strict=True, - param_names_mapping=get_param_names_mapping(fastvideo.param_names_mapping), - training_mode=False, - ) - assert incompatible.missing_keys == [] - assert incompatible.unexpected_keys == [] - assert not any(param.is_meta for param in fastvideo.parameters()) - - -def test_dreamx_world_5b_cam_dit_config_matches_official_shape(): - config = make_dreamx_world_5b_cam_dit_config() - assert config.num_layers == 30 - assert config.num_attention_heads == 24 - assert config.attention_head_dim == 128 - assert config.hidden_size == 3072 - assert config.ffn_dim == 14336 - assert config.add_control_adapter is True - assert config.cam_method == "prope" - assert config.attn_compress == 1 - - -def test_dreamx_world_converted_5b_transformer_strict_loads(): - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - _load_fastvideo_transformer_strict(device, torch.bfloat16) - - -def test_dreamx_world_fastvideo_prope_branch_smoke(): - block = DreamXWorldTransformerBlock( - 8, - 32, - 1, - cross_attn_norm=True, - add_control_adapter=True, - cam_method="prope", - attn_compress=1, - layer_idx=0, - supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ), - ) - assert block.cam_self_attn is not None - assert [name for name, _ in block.named_parameters() if name.startswith("cam_self_attn.")][:8] == [ - "cam_self_attn.q_proj.weight", - "cam_self_attn.q_proj.bias", - "cam_self_attn.k_proj.weight", - "cam_self_attn.k_proj.bias", - "cam_self_attn.v_proj.weight", - "cam_self_attn.v_proj.bias", - "cam_self_attn.out_proj.weight", - "cam_self_attn.out_proj.bias", - ] - - module = DreamXPropeSelfAttention( - dim=8, - attn_dim=8, - num_heads=1, - qk_norm="rms_norm_across_heads", - ).eval() - assert module.num_heads == 1 - assert module.head_dim == 8 - assert tuple(module.out_proj.weight.shape) == (8, 8) - assert torch.count_nonzero(module.out_proj.weight) == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for transformer parity.") -def test_dreamx_world_transformer_parity_scaffold(monkeypatch): - device = torch.device("cuda:0") - dtype = torch.float32 - inputs = _make_inputs(device, dtype) - - official = _load_official_transformer(device, dtype) - official_out = _run_official(official, inputs) - del official - torch.cuda.empty_cache() - - import fastvideo.attention.layer as attention_layer - import fastvideo.distributed.communication_op as communication_op - import fastvideo.models.dits.dreamx_world as fastvideo_dreamx - monkeypatch.setattr(attention_layer, "get_sp_parallel_rank", lambda: 0) - monkeypatch.setattr(attention_layer, "get_sp_world_size", lambda: 1) - monkeypatch.setattr(attention_layer, - "sequence_model_parallel_all_to_all_4D", - lambda tensor, scatter_dim=2, gather_dim=1: tensor) - monkeypatch.setattr(attention_layer, "sequence_model_parallel_all_gather", lambda tensor, dim=-1: tensor) - monkeypatch.setattr(fastvideo_dreamx, "get_sp_world_size", lambda: 1) - monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 1) - monkeypatch.setattr(fastvideo_dreamx, - "sequence_model_parallel_shard", - lambda tensor, dim=1: (tensor, tensor.shape[dim])) - monkeypatch.setattr( - fastvideo_dreamx, - "sequence_model_parallel_all_gather_with_unpad", - lambda tensor, original_seq_len, dim=1: tensor.narrow(dim, 0, original_seq_len), - ) - fastvideo = _load_fastvideo_transformer(device, dtype) - fastvideo_out = _run_fastvideo(fastvideo, inputs) - assert official_out.shape == fastvideo_out.shape - diff = (official_out - fastvideo_out).abs() - print(f"diff_max={diff.max().item():.6f} diff_mean={diff.mean().item():.6f}") - assert_close(fastvideo_out, official_out, atol=1e-1, rtol=1e-1) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py deleted file mode 100644 index 114897b44..000000000 --- a/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py +++ /dev/null @@ -1,221 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World Wan2.2 VAE reuse parity scaffold. - -Coverage scope: implementation_subcomponent. The official side uses -AutoencoderKLWan3_8 from DreamX, while the FastVideo side targets the native -Wan VAE. This remains a scaffold until Wan2.2 base VAE weights are staged. -""" -from __future__ import annotations - -import os -from pathlib import Path -import sys -import re -import types - -import pytest -import torch -from omegaconf import OmegaConf -from torch.testing import assert_close - -from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_vae_config -from fastvideo.models.wan.vae import AutoencoderKLWan - -REPO_ROOT = Path(__file__).resolve().parents[3] -OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) -WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B")) -PARITY_SCOPE = "implementation_subcomponent" - - -def _add_official_to_path(): - if not OFFICIAL_REF_DIR.exists(): - pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}") - if str(OFFICIAL_REF_DIR) not in sys.path: - sys.path.insert(0, str(OFFICIAL_REF_DIR)) - - -def _install_xfuser_stub() -> None: - if "xfuser" in sys.modules: - return - xfuser = types.ModuleType("xfuser") - core = types.ModuleType("xfuser.core") - distributed = types.ModuleType("xfuser.core.distributed") - long_ctx = types.ModuleType("xfuser.core.long_ctx_attention") - distributed.get_sequence_parallel_rank = lambda: 0 - distributed.get_sequence_parallel_world_size = lambda: 1 - distributed.get_sp_group = lambda: None - distributed.get_world_group = lambda: types.SimpleNamespace(local_rank=0, rank=0) - distributed.init_distributed_environment = lambda *args, **kwargs: None - distributed.initialize_model_parallel = lambda *args, **kwargs: None - distributed.model_parallel_is_initialized = lambda: False - - class XFuserLongContextAttention: - - def __call__(self, *args, **kwargs): - raise RuntimeError("xfuser stub cannot execute attention") - - long_ctx.xFuserLongContextAttention = XFuserLongContextAttention - sys.modules.update({ - "xfuser": xfuser, - "xfuser.core": core, - "xfuser.core.distributed": distributed, - "xfuser.core.long_ctx_attention": long_ctx, - }) - - -def _map_residual_subkey(prefix: str, sub: str) -> str | None: - if sub == "residual.0.gamma": - return f"{prefix}.norm1.gamma" - match = re.match(r"^residual\.2\.(weight|bias)$", sub) - if match: - return f"{prefix}.conv1.{match.group(1)}" - if sub == "residual.3.gamma": - return f"{prefix}.norm2.gamma" - match = re.match(r"^residual\.6\.(weight|bias)$", sub) - if match: - return f"{prefix}.conv2.{match.group(1)}" - match = re.match(r"^shortcut\.(weight|bias)$", sub) - if match: - return f"{prefix}.conv_shortcut.{match.group(1)}" - return None - - -def _map_attention_subkey(prefix: str, sub: str) -> str | None: - if sub == "norm.gamma": - return f"{prefix}.norm.gamma" - match = re.match(r"^(to_qkv|proj)\.(weight|bias)$", sub) - if match: - return f"{prefix}.{match.group(1)}.{match.group(2)}" - return None - - -def _map_resample_subkey(prefix: str, sub: str) -> str | None: - match = re.match(r"^resample\.1\.(weight|bias)$", sub) - if match: - return f"{prefix}.resample.1.{match.group(1)}" - match = re.match(r"^time_conv\.(weight|bias)$", sub) - if match: - return f"{prefix}.time_conv.{match.group(1)}" - return None - - -def _map_dreamx_raw_vae_key(key: str) -> str | None: - match = re.match(r"^(conv1|conv2)\.(weight|bias)$", key) - if match: - prefix = "quant_conv" if match.group(1) == "conv1" else "post_quant_conv" - return f"{prefix}.{match.group(2)}" - match = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key) - if match: - return f"{match.group(1)}.conv_in.{match.group(2)}" - match = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key) - if match: - return f"{match.group(1)}.norm_out.gamma" - match = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key) - if match: - return f"{match.group(1)}.conv_out.{match.group(2)}" - match = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key) - if match: - return _map_residual_subkey(f"{match.group(1)}.mid_block.resnets.0", match.group(2)) - match = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key) - if match: - return _map_attention_subkey(f"{match.group(1)}.mid_block.attentions.0", match.group(2)) - match = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key) - if match: - return _map_residual_subkey(f"{match.group(1)}.mid_block.resnets.1", match.group(2)) - match = re.match(r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.(.*)$", key) - if match: - stage = int(match.group(1)) - block = int(match.group(2)) - sub = match.group(3) - if block in (0, 1): - return _map_residual_subkey(f"encoder.down_blocks.{stage}.resnets.{block}", sub) - if block == 2: - return _map_resample_subkey(f"encoder.down_blocks.{stage}.downsampler", sub) - return None - match = re.match(r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.(.*)$", key) - if match: - stage = int(match.group(1)) - block = int(match.group(2)) - sub = match.group(3) - if block in (0, 1, 2): - return _map_residual_subkey(f"decoder.up_blocks.{stage}.resnets.{block}", sub) - if block == 3: - return _map_resample_subkey(f"decoder.up_blocks.{stage}.upsampler", sub) - return None - return None - - -def _vae_kwargs(): - config = OmegaConf.load(OFFICIAL_REF_DIR / "configs" / "wan2.2" / "wan_ti2v_5b.yaml") - return OmegaConf.to_container(config["vae_kwargs"]) - - -def _load_official_vae(device, dtype): - _add_official_to_path() - _install_xfuser_stub() - vae_path = WAN_BASE_DIR / "Wan2.2_VAE.pth" - if not vae_path.exists(): - pytest.skip(f"Wan2.2 base VAE weights missing: {vae_path}") - try: - from models import AutoencoderKLWan3_8 - except Exception as exc: # noqa: BLE001 - pytest.skip(f"Cannot import official DreamX VAE: {exc}") - model = AutoencoderKLWan3_8.from_pretrained(str(vae_path), additional_kwargs=_vae_kwargs()) - return model.to(device=device, dtype=dtype).eval() - - -def _load_fastvideo_vae(device, dtype): - vae_path = WAN_BASE_DIR / "Wan2.2_VAE.pth" - if not vae_path.exists(): - pytest.skip(f"Wan2.2 raw VAE weights missing: {vae_path}") - config = make_dreamx_world_5b_cam_vae_config() - config.load_encoder = True - config.load_decoder = True - model = AutoencoderKLWan(config).to(device=device, dtype=dtype) - raw_state = torch.load(str(vae_path), map_location="cpu", weights_only=True) - mapped_state = {} - for key, value in raw_state.items(): - mapped_key = _map_dreamx_raw_vae_key(key) - if mapped_key is None: - raise AssertionError(f"Unmapped DreamX raw VAE key: {key}") - mapped_state[mapped_key] = value - model.load_state_dict(mapped_state, strict=True) - return model.eval() - - -def _normalize_fastvideo_vae_latent(latent: torch.Tensor) -> torch.Tensor: - config = make_dreamx_world_5b_cam_vae_config() - mean = torch.tensor(config.latents_mean, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) - std = torch.tensor(config.latents_std, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) - return (latent - mean) / std - - -def test_dreamx_world_vae_config_matches_wan22_shape(): - config = make_dreamx_world_5b_cam_vae_config() - assert config.z_dim == 48 - assert config.in_channels == 12 - assert config.out_channels == 12 - assert config.base_dim == 160 - assert config.decoder_base_dim == 256 - assert config.scale_factor_temporal == 4 - assert config.scale_factor_spatial == 16 - assert config.patch_size == 2 - assert config.is_residual is True - assert config.clip_output is False - assert len(config.latents_mean) == 48 - assert len(config.latents_std) == 48 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for VAE parity.") -def test_dreamx_world_vae_encode_parity_scaffold(): - device = torch.device("cuda:0") - dtype = torch.bfloat16 - official = _load_official_vae(device, dtype) - fastvideo = _load_fastvideo_vae(device, dtype) - torch.manual_seed(123) - video = torch.randn(1, 3, 5, 64, 64, device=device, dtype=dtype).clamp(-1, 1) - with torch.inference_mode(): - official_latent = official.encode(video).latent_dist.mean.float().cpu() - fastvideo_latent = _normalize_fastvideo_vae_latent(fastvideo.encode(video).mean).float().cpu() - assert official_latent.shape == fastvideo_latent.shape - assert_close(fastvideo_latent, official_latent, atol=5e-2, rtol=5e-2) diff --git a/tests/local_tests/gamecraft/README.md b/tests/local_tests/gamecraft/README.md deleted file mode 100644 index e11aa2f13..000000000 --- a/tests/local_tests/gamecraft/README.md +++ /dev/null @@ -1,67 +0,0 @@ -# Hunyuan GameCraft Local Tests - -Local-only parity tests for the `gamecraft` FastVideo port. They compare -FastVideo against the official `Hunyuan-GameCraft-1.0` implementation and are -skipped in CI; run locally on a single GPU (~40 GB+ memory). - -## Reference Assets - -| Field | Value | -|---|---| -| Model family | `gamecraft` | -| Workload types | `T2V`, `I2V` | -| Official reference | `Hunyuan-GameCraft-1.0` (local clone, see test files) | -| Local reference dir | `Hunyuan-GameCraft-1.0/` (env: `GAMECRAFT_OFFICIAL_PATH`) | -| Official commit/version | `` | -| HF weights | `` | -| HF revision | `` | -| Local weights dir | `Hunyuan-GameCraft-1.0/weights/` (env: `GAMECRAFT_WEIGHTS_PATH`) | -| Source layout | `` | -| Needs conversion | `` | - -> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value. - -## Shared Environment Setup - -Run from the FastVideo repo root in the same env used for FastVideo. See the -`add-model-prep` skill for canonical clone + install commands. - -```bash -# Clone the official reference under the FastVideo repo root. -git clone Hunyuan-GameCraft-1.0 -uv pip install --no-deps -e ./Hunyuan-GameCraft-1.0 -``` - -Do not change core dependency versions (`torch`, `diffusers`, `transformers`, -`flash-attn`, `triton`, CUDA packages) without explicit approval. - -## Official Environment Status - -```text -dependency_changes: -official_env_status: -private_dep_stubs: -blocked_on: -``` - -## Tests in this directory - -Run the whole family: - -```bash -DISABLE_SP=1 pytest tests/local_tests/gamecraft/ -v -s -``` - -| Component | Test | Concerns | Status | -|---|---|---|---| -| `text encoders` (LLaVA-LLaMA-3-8B + CLIP ViT-L/14) | [`test_gamecraft_encoders_parity.py`](./test_gamecraft_encoders_parity.py) | `` | `` | -| `transformer/DiT` | [`test_gamecraft_parity.py`](./test_gamecraft_parity.py) | `` | `` | -| `vae` (`AutoencoderKLCausal3D`) | [`test_gamecraft_vae_parity.py`](./test_gamecraft_vae_parity.py) | `` | `` | -| `pipeline` | [`test_gamecraft_pipeline_parity.py`](./test_gamecraft_pipeline_parity.py) | text encoding -> denoising -> VAE decode | `` | - -## Review Notes - -- Required before handoff: non-skip PASS for each component parity test, - including reused components that own weights or numerical behavior. -- Pipeline parity may start as a scaffold; final handoff requires non-skip - PASS or an explicit blocker accepted via the escape-hatch process. diff --git a/tests/local_tests/gamecraft/__init__.py b/tests/local_tests/gamecraft/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/local_tests/gamecraft/test_gamecraft_encoders_parity.py b/tests/local_tests/gamecraft/test_gamecraft_encoders_parity.py deleted file mode 100644 index 65640bd7a..000000000 --- a/tests/local_tests/gamecraft/test_gamecraft_encoders_parity.py +++ /dev/null @@ -1,244 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Numerical parity test for GameCraft text encoders. - -GameCraft uses: -1. LLaVA-LLaMA-3-8B for primary text encoding (4096-dim) -2. CLIP ViT-L/14 for secondary text encoding (768-dim pooled) - -Usage: - DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_encoders_parity.py -v -""" -import os -import sys -from pathlib import Path - -import pytest -import torch -from torch.testing import assert_close - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29516") -os.environ.setdefault("DISABLE_SP", "1") - -repo_root = Path(__file__).resolve().parents[3] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_gamecraft_llama_encoder(): - """Test LLaVA-LLaMA-3-8B text encoder.""" - torch.manual_seed(42) - - llama_path = Path( - os.getenv("GAMECRAFT_LLAMA_PATH", - repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "stdmodels" / "llava-llama-3-8b-v1_1-transformers")) - - if not llama_path.exists(): - pytest.skip(f"LLaMA encoder not found at {llama_path}") - - # Check for model weights - model_files = list(llama_path.glob("*.safetensors")) + list(llama_path.glob("*.bin")) - if not model_files: - pytest.skip(f"No model weights found in {llama_path}") - - device = torch.device("cuda:0") - dtype = torch.float16 # LLaMA uses fp16 - - print(f"\n[LLAMA TEST] Loading LLaVA-LLaMA-3-8B from {llama_path}") - - try: - from transformers import LlavaForConditionalGeneration, LlamaTokenizerFast - - # Load model - model = LlavaForConditionalGeneration.from_pretrained( - llama_path, - low_cpu_mem_usage=True, - torch_dtype=dtype, - ) - model = model.to(device) - model.eval() - - # Load tokenizer - tokenizer = LlamaTokenizerFast.from_pretrained(llama_path, padding_side="right") - - print(f"[LLAMA TEST] Model loaded successfully") - print(f"[LLAMA TEST] Model dtype: {model.dtype}") - - # Test encoding - test_prompt = "A medieval village with cobblestone streets" - inputs = tokenizer( - test_prompt, - return_tensors="pt", - padding="max_length", - max_length=256, - truncation=True, - ) - inputs = {k: v.to(device) for k, v in inputs.items()} - - with torch.no_grad(): - outputs = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - output_hidden_states=True, - ) - - # Get last hidden state - last_hidden_state = outputs.hidden_states[-1] - - print(f"[LLAMA TEST] Input tokens: {inputs['input_ids'].shape}") - print(f"[LLAMA TEST] Output hidden state shape: {last_hidden_state.shape}") - print(f"[LLAMA TEST] Hidden dim: {last_hidden_state.shape[-1]} (expected: 4096)") - print(f"[LLAMA TEST] Output stats: mean={last_hidden_state.mean():.6f}, std={last_hidden_state.std():.6f}") - - assert last_hidden_state.shape[-1] == 4096, f"Expected hidden dim 4096, got {last_hidden_state.shape[-1]}" - - print("[LLAMA TEST] LLaMA encoder test passed!") - - except ImportError as e: - pytest.skip(f"transformers not available: {e}") - except Exception as e: - pytest.fail(f"LLaMA encoder test failed: {e}") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_gamecraft_clip_encoder(): - """Test CLIP ViT-L/14 text encoder.""" - torch.manual_seed(42) - - clip_path = Path( - os.getenv("GAMECRAFT_CLIP_PATH", - repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "stdmodels" / "openai_clip-vit-large-patch14")) - - if not clip_path.exists(): - pytest.skip(f"CLIP encoder not found at {clip_path}") - - device = torch.device("cuda:0") - dtype = torch.float32 # CLIP typically uses fp32 - - print(f"\n[CLIP TEST] Loading CLIP ViT-L/14 from {clip_path}") - - try: - from transformers import CLIPTextModel, CLIPTokenizer - - # Load model - model = CLIPTextModel.from_pretrained(clip_path) - model = model.to(device=device, dtype=dtype) - model.eval() - - # Load tokenizer - tokenizer = CLIPTokenizer.from_pretrained(clip_path, max_length=77) - - print(f"[CLIP TEST] Model loaded successfully") - print(f"[CLIP TEST] Model dtype: {model.dtype}") - - # Test encoding - test_prompt = "A medieval village with cobblestone streets" - inputs = tokenizer( - test_prompt, - return_tensors="pt", - padding="max_length", - max_length=77, - truncation=True, - ) - inputs = {k: v.to(device) for k, v in inputs.items()} - - with torch.no_grad(): - outputs = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - ) - - # Get pooled output (used for text_states_2) - pooler_output = outputs.pooler_output - last_hidden_state = outputs.last_hidden_state - - print(f"[CLIP TEST] Input tokens: {inputs['input_ids'].shape}") - print(f"[CLIP TEST] Pooled output shape: {pooler_output.shape} (expected: [B, 768])") - print(f"[CLIP TEST] Last hidden state shape: {last_hidden_state.shape}") - print(f"[CLIP TEST] Pooled output stats: mean={pooler_output.mean():.6f}, std={pooler_output.std():.6f}") - - assert pooler_output.shape[-1] == 768, f"Expected pooled dim 768, got {pooler_output.shape[-1]}" - - print("[CLIP TEST] CLIP encoder test passed!") - - except ImportError as e: - pytest.skip(f"transformers not available: {e}") - except Exception as e: - pytest.fail(f"CLIP encoder test failed: {e}") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.skip( - reason= - "Official TextEncoder wrapper incompatible with installed transformers version (language_model attribute removed from LlavaForConditionalGeneration)" -) -def test_gamecraft_text_encoder_wrapper(): - """Test the official GameCraft TextEncoder wrapper.""" - torch.manual_seed(42) - - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - - sys.path.insert(0, str(official_path)) - - llama_path = official_path / "weights" / "stdmodels" / "llava-llama-3-8b-v1_1-transformers" - clip_path = official_path / "weights" / "stdmodels" / "openai_clip-vit-large-patch14" - - if not llama_path.exists() or not clip_path.exists(): - pytest.skip("Text encoder weights not found") - - device = torch.device("cuda:0") - - try: - from hymm_sp.text_encoder import TextEncoder - - # Test LLaMA encoder wrapper - print("\n[WRAPPER TEST] Testing LLaVA-LLaMA-3-8B wrapper...") - llama_encoder = TextEncoder( - text_encoder_type="llava-llama-3-8b", - max_length=256, - text_encoder_precision="fp16", - text_encoder_path=str(llama_path), - tokenizer_path=str(llama_path), - use_attention_mask=True, - device=device, - ) - - test_prompt = "A medieval village with cobblestone streets" - llama_output = llama_encoder(test_prompt) - - print(f"[WRAPPER TEST] LLaMA hidden state: {llama_output.hidden_state.shape}") - print( - f"[WRAPPER TEST] LLaMA attention mask: {llama_output.attention_mask.shape if llama_output.attention_mask is not None else 'None'}" - ) - - # Test CLIP encoder wrapper - print("\n[WRAPPER TEST] Testing CLIP wrapper...") - clip_encoder = TextEncoder( - text_encoder_type="clipL", - max_length=77, - text_encoder_precision="fp32", - text_encoder_path=str(clip_path), - tokenizer_path=str(clip_path), - use_attention_mask=True, - device=device, - ) - - clip_output = clip_encoder(test_prompt) - - print(f"[WRAPPER TEST] CLIP pooled output: {clip_output.hidden_state.shape}") - - print("[WRAPPER TEST] TextEncoder wrapper test passed!") - - except ImportError as e: - pytest.skip(f"Failed to import TextEncoder: {e}") - except Exception as e: - pytest.fail(f"TextEncoder wrapper test failed: {e}") - - -if __name__ == "__main__": - test_gamecraft_clip_encoder() - test_gamecraft_llama_encoder() - # test_gamecraft_text_encoder_wrapper() # Requires full model weights diff --git a/tests/local_tests/gamecraft/test_gamecraft_parity.py b/tests/local_tests/gamecraft/test_gamecraft_parity.py deleted file mode 100644 index b470f57d5..000000000 --- a/tests/local_tests/gamecraft/test_gamecraft_parity.py +++ /dev/null @@ -1,753 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Numerical parity test for HunyuanGameCraft transformer. - -This test compares the FastVideo implementation against the official -Hunyuan-GameCraft-1.0 implementation to ensure numerical alignment. - -Usage: - # Run with default paths - pytest tests/local_tests/gamecraft/test_gamecraft_parity.py -v - - # Run with custom paths - GAMECRAFT_OFFICIAL_PATH=path/to/Hunyuan-GameCraft-1.0 \ - GAMECRAFT_WEIGHTS_PATH=path/to/weights/gamecraft_models/mp_rank_00_model_states.pt \ - pytest tests/local_tests/gamecraft/test_gamecraft_parity.py -v - - # Enable debug logging - GAMECRAFT_DEBUG_LOGS=1 pytest tests/local_tests/gamecraft/test_gamecraft_parity.py -v -""" -import os -import sys -from pathlib import Path - -import pytest -import torch -from torch.testing import assert_close - -# Set up distributed environment defaults -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29514") -os.environ.setdefault("DISABLE_SP", "1") # Disable sequence parallelism for testing - -repo_root = Path(__file__).resolve().parents[3] - - -def _add_official_to_path(): - """Add official GameCraft implementation to Python path.""" - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - if official_path.exists() and str(official_path) not in sys.path: - sys.path.insert(0, str(official_path)) - return official_path - - -def _format_sum(tensor: torch.Tensor | None) -> str: - """Format tensor sum for logging.""" - if tensor is None: - return "None" - return f"{tensor.float().sum().item():.6f}" - - -def _attach_block_logging( - model: torch.nn.Module, - log_path: Path, - label: str, - enabled: bool, -) -> None: - """Attach forward hooks to log block outputs.""" - if not enabled: - return - - log_path.parent.mkdir(parents=True, exist_ok=True) - if log_path.exists(): - log_path.unlink() - - def _hook_factory(block_idx: int, block_type: str): - - def _hook(_module, _inputs, outputs): - if isinstance(outputs, tuple): - img_out, txt_out = outputs - img_sum = _format_sum(img_out) - txt_sum = _format_sum(txt_out) - out_str = f"img_sum={img_sum},txt_sum={txt_sum}" - else: - out_str = f"out_sum={_format_sum(outputs)}" - with log_path.open("a", encoding="utf-8") as f: - f.write(f"{label}:{block_type}:{block_idx}:{out_str}\n") - - return _hook - - # Attach to double blocks - if hasattr(model, "double_blocks"): - for idx, block in enumerate(model.double_blocks): - block.register_forward_hook(_hook_factory(idx, "double")) - - # Attach to single blocks - if hasattr(model, "single_blocks"): - for idx, block in enumerate(model.single_blocks): - block.register_forward_hook(_hook_factory(idx, "single")) - - -def _load_official_model(official_path: Path, weights_path: Path, device: torch.device, dtype: torch.dtype): - """Load the official GameCraft model.""" - try: - # Add official implementation to path - sys.path.insert(0, str(official_path)) - - from hymm_sp.modules.models import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG - from hymm_sp.config import parse_args - - # Create a minimal args object - class MinimalArgs: - text_projection = "single_refiner" - text_states_dim = 4096 - use_attention_mask = False - text_states_dim_2 = 768 - - args = MinimalArgs() - config = HUNYUAN_VIDEO_CONFIG['HYVideo-T/2'] - - # Initialize model - model = HYVideoDiffusionTransformer( - args=args, - patch_size=[1, 2, 2], - in_channels=16, # Will be adjusted for concat mode - out_channels=16, - hidden_size=config['hidden_size'], - mlp_width_ratio=config['mlp_width_ratio'], - num_heads=config['num_heads'], - depth_double_blocks=config['depth_double_blocks'], - depth_single_blocks=config['depth_single_blocks'], - rope_dim_list=config['rope_dim_list'], - guidance_embed=False, - dtype=dtype, - device='cpu', # Load on CPU first - multitask_mask_training_type="concat", # GameCraft uses concat mode - camera_in_channels=6, - camera_down_coef=8, - ) - - # Load weights - checkpoint = torch.load(weights_path, map_location='cpu') - state_dict = checkpoint['module'] - model.load_state_dict(state_dict, strict=True) - model = model.to(device=device, dtype=dtype) - model.eval() - - return model - - except ImportError as e: - pytest.skip(f"Failed to import official GameCraft: {e}") - except Exception as e: - pytest.skip(f"Failed to load official GameCraft model: {e}") - - -def _create_test_inputs(batch_size: int, device: torch.device, dtype: torch.dtype): - """Create test inputs for the transformer.""" - # Video latent shape: [B, C, T, H, W] - # For GameCraft: C=33 (16 latent + 16 gt_latent + 1 mask) - frames = 9 # 9 latent frames = 33 video frames with temporal compression - height = 88 # 704 / 8 - width = 160 # 1280 / 8 - - # Pure latents (what we're denoising) - latents = torch.randn(batch_size, 16, frames, height, width, device=device, dtype=dtype) - - # Ground truth latents (history/reference) - gt_latents = torch.randn(batch_size, 16, frames, height, width, device=device, dtype=dtype) - - # Mask (1 = use gt, 0 = generate) - mask = torch.zeros(batch_size, 1, frames, height, width, device=device, dtype=dtype) - mask[:, :, :frames // 2] = 1.0 # First half is conditioned - - # Concatenated input - x = torch.cat([latents, gt_latents, mask], dim=1) # [B, 33, T, H, W] - - # Timestep - t = torch.tensor([500.0] * batch_size, device=device, dtype=dtype) - - # Text embeddings (LLaMA) - text_states = torch.randn(batch_size, 77, 4096, device=device, dtype=dtype) - text_mask = torch.ones(batch_size, 77, device=device, dtype=dtype) - - # CLIP embeddings - text_states_2 = torch.randn(batch_size, 768, device=device, dtype=dtype) - - # Camera/action states: [B, T_video, 6, H_video, W_video] - # For 9 latent frames with temporal 4x compression -> ~36 video frames - # But camera net expects raw video resolution - video_frames = 33 # (9-1)*4 + 1 for 884 VAE - cam_latents = torch.randn(batch_size, video_frames, 6, 704, 1280, device=device, dtype=dtype) - - return { - 'x': x, - 't': t, - 'text_states': text_states, - 'text_mask': text_mask, - 'text_states_2': text_states_2, - 'cam_latents': cam_latents, - } - - -def _get_rotary_embeddings(shape, hidden_size, num_heads, rope_dim_list, device, dtype): - """Generate rotary position embeddings.""" - # Simplified RoPE generation for testing - # In production, this uses the full RoPE implementation - t, h, w = shape - seq_len = t * h * w - head_dim = hidden_size // num_heads - - freqs_cos = torch.ones(seq_len, head_dim, device=device, dtype=dtype) - freqs_sin = torch.zeros(seq_len, head_dim, device=device, dtype=dtype) - - return freqs_cos, freqs_sin - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_gamecraft_transformer_parity(): - """Test numerical parity between official and FastVideo GameCraft implementations.""" - torch.manual_seed(42) - - # Initialize distributed environment for FastVideo - from fastvideo.distributed.parallel_state import init_distributed_environment, initialize_model_parallel - init_distributed_environment(world_size=1, rank=0, local_rank=0) - initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1) - - # Get paths - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - weights_path = Path( - os.getenv("GAMECRAFT_WEIGHTS_PATH", - repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "gamecraft_models" / "mp_rank_00_model_states.pt")) - - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - if not weights_path.exists(): - pytest.skip(f"GameCraft weights not found at {weights_path}") - - device = torch.device("cuda:0") - dtype = torch.bfloat16 - batch_size = 1 - - # Load official model - print(f"\n[GameCraft TEST] Loading official model from {weights_path}") - official_model = _load_official_model(official_path, weights_path, device, dtype) - - # Enable debug logging if requested - debug_logs = os.getenv("GAMECRAFT_DEBUG_LOGS", "0") == "1" - _attach_block_logging( - official_model, - repo_root / "gamecraft_debug" / "official.log", - "official", - debug_logs, - ) - - # Create test inputs - # IMPORTANT: latent_len must be 9, 10, or 18 for official model's camera conditioning - frames = 9 # Valid latent temporal length - height = 44 # 352 / 8 (small spatial resolution for testing) - width = 80 # 640 / 8 - - latents = torch.randn(batch_size, 16, frames, height, width, device=device, dtype=dtype) - gt_latents = torch.randn(batch_size, 16, frames, height, width, device=device, dtype=dtype) - mask = torch.zeros(batch_size, 1, frames, height, width, device=device, dtype=dtype) - mask[:, :, :4] = 1.0 # First 4 frames are conditioned (history) - - x = torch.cat([latents, gt_latents, mask], dim=1) # [B, 33, T, H, W] - t = torch.tensor([500.0], device=device, dtype=dtype) - - text_states = torch.randn(batch_size, 32, 4096, device=device, dtype=dtype) - text_mask = torch.ones(batch_size, 32, device=device, dtype=dtype) - text_states_2 = torch.randn(batch_size, 768, device=device, dtype=dtype) - - # Camera states: for 9 latent frames -> 33 video frames with 884 VAE - video_frames = 33 - cam_latents = torch.randn(batch_size, video_frames, 6, 352, 640, device=device, dtype=dtype) - - # Get RoPE embeddings - after patchify with patch_size=[1,2,2]: T, H/2, W/2 - tt, th, tw = frames, height // 2, width // 2 - - # Use FastVideo's RoPE implementation for both models - from fastvideo.layers.rotary_embedding import get_rotary_pos_embed - - hidden_size = 3072 - num_heads = 24 - rope_dim_list = [16, 56, 56] # From GameCraft config - rope_theta = 256.0 - - freqs_cos, freqs_sin = get_rotary_pos_embed( - (tt, th, tw), - hidden_size, - num_heads, - rope_dim_list, - rope_theta, - ) - freqs_cos = freqs_cos.to(device=device, dtype=dtype) - freqs_sin = freqs_sin.to(device=device, dtype=dtype) - - print(f"[GameCraft TEST] Input shapes:") - print(f" x: {x.shape}") - print(f" text_states: {text_states.shape}") - print(f" cam_latents: {cam_latents.shape}") - - # Run official model - with torch.no_grad(): - official_output = official_model( - x=x, - t=t, - text_states=text_states, - text_mask=text_mask, - text_states_2=text_states_2, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - cam_latents=cam_latents, - return_dict=True, - ) - if isinstance(official_output, dict): - official_out = official_output['x'] - else: - official_out = official_output - - print(f"[GameCraft TEST] Official output shape: {official_out.shape}") - print(f"[GameCraft TEST] Official output stats: mean={official_out.mean():.6f}, std={official_out.std():.6f}") - - # Load FastVideo GameCraft model - print("[GameCraft TEST] Loading FastVideo GameCraft model...") - from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig - from fastvideo.models.dits.hunyuangamecraft import HunyuanGameCraftTransformer3DModel - - # Create FastVideo config - fastvideo_config = HunyuanGameCraftConfig() - fastvideo_config.dtype = dtype - - # Load converted weights - converted_weights_path = repo_root / "official_weights" / "hunyuan-gamecraft" / "transformer" - if not converted_weights_path.exists(): - # Try to convert weights on-the-fly - print(f"[GameCraft TEST] Converting weights from {weights_path}...") - from scripts.checkpoint_conversion.convert_gamecraft_weights import convert_gamecraft_weights - convert_gamecraft_weights( - input_path=str(weights_path), - output_dir=str(converted_weights_path), - ) - - # Load hf_config from converted weights - import json - config_path = converted_weights_path / "config.json" - if config_path.exists(): - with open(config_path) as f: - hf_config = json.load(f) - else: - # Create minimal hf_config - hf_config = { - "hidden_size": 3072, - "num_attention_heads": 24, - "num_layers": 20, - "num_single_layers": 40, - } - - # Initialize FastVideo model - fastvideo_model = HunyuanGameCraftTransformer3DModel( - config=fastvideo_config, - hf_config=hf_config, - ) - - # Load weights - from safetensors.torch import load_file - safetensor_files = list(converted_weights_path.glob("*.safetensors")) - if safetensor_files: - state_dict = {} - for sf_file in safetensor_files: - state_dict.update(load_file(sf_file)) - - # Apply param_names_mapping in reverse - reverse_mapping = fastvideo_config.reverse_param_names_mapping - mapped_state_dict = {} - for key, value in state_dict.items(): - mapped_key = key - for pattern, replacement in reverse_mapping.items(): - import re - if re.match(pattern, key): - mapped_key = re.sub(pattern, replacement, key) - break - mapped_state_dict[mapped_key] = value - - missing, unexpected = fastvideo_model.load_state_dict(mapped_state_dict, strict=False) - print(f"[GameCraft TEST] Missing keys: {len(missing)}") - print(f"[GameCraft TEST] Unexpected keys: {len(unexpected)}") - if missing: - print(f"[GameCraft TEST] First 5 missing: {missing[:5]}") - if unexpected: - print(f"[GameCraft TEST] First 5 unexpected: {unexpected[:5]}") - - fastvideo_model = fastvideo_model.to(device=device, dtype=dtype) - fastvideo_model.eval() - - _attach_block_logging( - fastvideo_model, - repo_root / "gamecraft_debug" / "fastvideo.log", - "fastvideo", - debug_logs, - ) - - # Prepare FastVideo inputs - # FastVideo expects encoder_hidden_states as [text_states, text_states_2] - fastvideo_encoder_hidden_states = [text_states, text_states_2] - - # Import forward context - from fastvideo.forward_context import set_forward_context - from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - - forward_batch = ForwardBatch(data_type="video") - - # Run FastVideo model - with torch.no_grad(): - with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch): - fastvideo_out = fastvideo_model( - x, # Already concatenated [B, 33, T, H, W] - fastvideo_encoder_hidden_states, - t, - camera_states=cam_latents, - encoder_attention_mask=[text_mask], - ) - - print(f"[GameCraft TEST] FastVideo output shape: {fastvideo_out.shape}") - print(f"[GameCraft TEST] FastVideo output stats: mean={fastvideo_out.mean():.6f}, std={fastvideo_out.std():.6f}") - - # Compare outputs using appropriate metrics for numerical parity - print("[GameCraft TEST] Comparing outputs...") - - # Compute various metrics - diff = (official_out - fastvideo_out).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - - # Correlation (most important for numerical alignment) - corr = torch.corrcoef(torch.stack([official_out.flatten().float(), fastvideo_out.flatten().float()]))[0, 1].item() - - # Cosine similarity - cos_sim = torch.nn.functional.cosine_similarity(official_out.flatten().unsqueeze(0).float(), - fastvideo_out.flatten().unsqueeze(0).float()).item() - - # Sum comparison (should match exactly) - off_sum = official_out.double().sum().item() - fv_sum = fastvideo_out.double().sum().item() - sum_diff_pct = abs(off_sum - fv_sum) / (abs(off_sum) + 1e-10) * 100 - - print(f"[GameCraft TEST] Max diff: {max_diff:.4f}") - print(f"[GameCraft TEST] Mean diff: {mean_diff:.6f}") - print(f"[GameCraft TEST] Correlation: {corr:.6f}") - print(f"[GameCraft TEST] Cosine similarity: {cos_sim:.6f}") - print(f"[GameCraft TEST] Sum diff: {sum_diff_pct:.4f}%") - - # Pass criteria: correlation > 0.99 and cosine_sim > 0.99 and sum_diff < 10% - # Sum diff can be larger due to floating point accumulation differences - assert corr > 0.99, f"Correlation too low: {corr}" - assert cos_sim > 0.99, f"Cosine similarity too low: {cos_sim}" - assert sum_diff_pct < 10.0, f"Sum difference too high: {sum_diff_pct}%" - - print("[GameCraft TEST] PASSED: Outputs are numerically aligned!") - print("[GameCraft TEST] Parity test complete!") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_gamecraft_layer_by_layer(): - """Test layer-by-layer parity to find divergence point.""" - import os - os.environ['DISABLE_SP'] = '1' - - torch.manual_seed(42) - - from fastvideo.distributed.parallel_state import init_distributed_environment, initialize_model_parallel - init_distributed_environment(world_size=1, rank=0, local_rank=0) - initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1) - - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - weights_path = Path( - os.getenv("GAMECRAFT_WEIGHTS_PATH", - repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "gamecraft_models" / "mp_rank_00_model_states.pt")) - - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - if not weights_path.exists(): - pytest.skip(f"GameCraft weights not found at {weights_path}") - - device = torch.device("cuda:0") - dtype = torch.bfloat16 - batch_size = 1 - frames = 9 - height, width = 44, 80 - text_seq_len = 32 - - # Load official model - print("\n[LAYER TEST] Loading official model...") - official_model = _load_official_model(official_path, weights_path, device, dtype) - - # Load FastVideo model - print("[LAYER TEST] Loading FastVideo model...") - from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig - from fastvideo.models.dits.hunyuangamecraft import HunyuanGameCraftTransformer3DModel - - fastvideo_config = HunyuanGameCraftConfig() - converted_weights_path = repo_root / "official_weights" / "hunyuan-gamecraft" / "transformer" - - import json - config_path = converted_weights_path / "config.json" - hf_config = {} - if config_path.exists(): - with open(config_path) as f: - hf_config = json.load(f) - - fastvideo_model = HunyuanGameCraftTransformer3DModel(config=fastvideo_config, hf_config=hf_config) - - from safetensors.torch import load_file - safetensor_files = list(converted_weights_path.glob("*.safetensors")) - if safetensor_files: - state_dict = {} - for sf_file in safetensor_files: - state_dict.update(load_file(sf_file)) - - reverse_mapping = fastvideo_config.reverse_param_names_mapping - mapped_state_dict = {} - for key, value in state_dict.items(): - mapped_key = key - for pattern, replacement in reverse_mapping.items(): - import re - if re.match(pattern, key): - mapped_key = re.sub(pattern, replacement, key) - break - mapped_state_dict[mapped_key] = value - - fastvideo_model.load_state_dict(mapped_state_dict, strict=False) - - fastvideo_model = fastvideo_model.to(device=device, dtype=dtype) - fastvideo_model.eval() - - # Create inputs - torch.manual_seed(42) - x = torch.randn(batch_size, 33, frames, height, width, device=device, dtype=dtype) - t = torch.tensor([500.0], device=device, dtype=dtype) - text_states = torch.randn(batch_size, text_seq_len, 4096, device=device, dtype=dtype) - text_states_2 = torch.randn(batch_size, 768, device=device, dtype=dtype) - cam_latents = torch.randn(batch_size, 33, 6, 352, 640, device=device, dtype=dtype) - - # RoPE - tt, th, tw = frames, height // 2, width // 2 - from fastvideo.layers.rotary_embedding import get_rotary_pos_embed - freqs_cos, freqs_sin = get_rotary_pos_embed((tt, th, tw), 3072, 24, [16, 56, 56], 256.0) - freqs_cos = freqs_cos.to(device=device, dtype=dtype) - freqs_sin = freqs_sin.to(device=device, dtype=dtype) - - print("[LAYER TEST] Testing embedding layers...") - - # Set up FastVideo context early (needed for txt_in) - from fastvideo.forward_context import set_forward_context - from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder - attn_metadata = FlashAttnMetadataBuilder().build(current_timestep=0, attn_mask=None) - - with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=attn_metadata): - # === img_in === - off_img = official_model.img_in(x) - fv_img = fastvideo_model.img_in(x) - diff = (off_img - fv_img).abs().max().item() - print(f" img_in: max_diff={diff:.6f} {'OK' if diff < 1e-3 else 'DIFF!'}") - - # === time_in === - off_time = official_model.time_in(t) - fv_time = fastvideo_model.time_in(t) - diff = (off_time - fv_time).abs().max().item() - print(f" time_in: max_diff={diff:.6f} {'OK' if diff < 1e-3 else 'DIFF!'}") - - # === vector_in === - off_vec = official_model.vector_in(text_states_2) - fv_vec = fastvideo_model.vector_in(text_states_2) - diff = (off_vec - fv_vec).abs().max().item() - print(f" vector_in: max_diff={diff:.6f} {'OK' if diff < 1e-3 else 'DIFF!'}") - - # === txt_in === - off_txt = official_model.txt_in(text_states, t, None) - fv_txt = fastvideo_model.txt_in(text_states, t) - diff = (off_txt - fv_txt).abs().max().item() - print(f" txt_in: max_diff={diff:.6f} {'OK' if diff < 1e-3 else 'DIFF!'}") - - # === camera_net === - off_cam = official_model.camera_net(cam_latents) - fv_cam = fastvideo_model.camera_net(cam_latents) - diff = (off_cam - fv_cam).abs().max().item() - print(f" camera_net: max_diff={diff:.6f} {'OK' if diff < 1e-3 else 'DIFF!'}") - - # Combined embeddings - off_img_combined = off_img + off_cam - fv_img_combined = fv_img + fv_cam - off_vec_combined = off_time + off_vec - fv_vec_combined = fv_time + fv_vec - - print("\n[LAYER TEST] Testing double blocks (first 5)...") - - # cu_seqlens for official model - seq_len = off_img_combined.shape[1] + off_txt.shape[1] - cu_seqlens = torch.tensor([0, seq_len, seq_len], device=device, dtype=torch.int32) - - off_img_cur, off_txt_cur = off_img_combined, off_txt - fv_img_cur, fv_txt_cur = fv_img_combined, fv_txt - - for i in range(min(5, len(official_model.double_blocks))): - # Official - off_img_out, off_txt_out = official_model.double_blocks[i](off_img_cur, - off_txt_cur, - off_vec_combined, - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=seq_len, - max_seqlen_kv=seq_len, - freqs_cis=(freqs_cos, freqs_sin), - use_sage=False) - - # FastVideo (context already set above) - fv_img_out, fv_txt_out = fastvideo_model.double_blocks[i](fv_img_cur, - fv_txt_cur, - fv_vec_combined, - freqs_cis=(freqs_cos, freqs_sin)) - - img_diff = (off_img_out - fv_img_out).abs().max().item() - txt_diff = (off_txt_out - fv_txt_out).abs().max().item() - img_mean_diff = (off_img_out - fv_img_out).abs().mean().item() - - status = 'OK' if img_diff < 0.1 else 'DIFF!' - print( - f" double_block[{i}]: img_max_diff={img_diff:.4f}, img_mean_diff={img_mean_diff:.6f}, txt_max_diff={txt_diff:.4f} {status}" - ) - - # Use outputs for next layer - off_img_cur, off_txt_cur = off_img_out, off_txt_out - fv_img_cur, fv_txt_cur = fv_img_out, fv_txt_out - - if img_diff > 1.0: - print(f" [STOPPING] Large divergence at double_block[{i}]") - break - - print("\n[LAYER TEST] Done.") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_gamecraft_cameranet_parity(): - """Test CameraNet component in isolation.""" - torch.manual_seed(42) - - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - weights_path = Path( - os.getenv("GAMECRAFT_WEIGHTS_PATH", - repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "gamecraft_models" / "mp_rank_00_model_states.pt")) - - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - if not weights_path.exists(): - pytest.skip(f"GameCraft weights not found at {weights_path}") - - device = torch.device("cuda:0") - dtype = torch.bfloat16 - - # Add official to path - _add_official_to_path() - - try: - from hymm_sp.modules.cameranet import CameraNet - except ImportError as e: - pytest.skip(f"Failed to import CameraNet: {e}") - - # Load CameraNet weights - checkpoint = torch.load(weights_path, map_location='cpu') - state_dict = checkpoint['module'] - - # Extract CameraNet weights - camera_state_dict = {} - for k, v in state_dict.items(): - if k.startswith('camera_net.'): - new_key = k.replace('camera_net.', '') - camera_state_dict[new_key] = v - - # Initialize CameraNet - camera_net = CameraNet( - in_channels=6, - downscale_coef=8, - out_channels=16, - patch_size=[1, 2, 2], - hidden_size=3072, - ) - camera_net.load_state_dict(camera_state_dict, strict=True) - camera_net = camera_net.to(device=device, dtype=dtype) - camera_net.eval() - - # Test input - batch_size = 1 - num_frames = 17 - cam_input = torch.randn(batch_size, num_frames, 6, 352, 640, device=device, dtype=dtype) - - with torch.no_grad(): - cam_output = camera_net(cam_input) - - print(f"[CameraNet TEST] Official input shape: {cam_input.shape}") - print(f"[CameraNet TEST] Official output shape: {cam_output.shape}") - print(f"[CameraNet TEST] Official output stats: mean={cam_output.mean():.6f}, std={cam_output.std():.6f}") - - # Load FastVideo CameraNet - print("[CameraNet TEST] Loading FastVideo CameraNet...") - from fastvideo.models.dits.hunyuangamecraft import CameraNet as FastVideoCameraNet - - # Initialize FastVideo CameraNet - fastvideo_camera_net = FastVideoCameraNet( - in_channels=6, - downscale_coef=8, - out_channels=16, - patch_size=[1, 2, 2], - hidden_size=3072, - dtype=dtype, - prefix="camera_net", - ) - - # Map official weights to FastVideo format - fastvideo_camera_state_dict = {} - for k, v in camera_state_dict.items(): - new_key = k - # Map camera_in -> PatchEmbed format - if "camera_in.proj" in k: - new_key = k # Keep as-is for now - fastvideo_camera_state_dict[new_key] = v - - missing, unexpected = fastvideo_camera_net.load_state_dict(fastvideo_camera_state_dict, strict=False) - print(f"[CameraNet TEST] Missing keys: {len(missing)}, Unexpected keys: {len(unexpected)}") - if missing: - print(f"[CameraNet TEST] Missing: {missing}") - if unexpected: - print(f"[CameraNet TEST] Unexpected: {unexpected}") - - fastvideo_camera_net = fastvideo_camera_net.to(device=device, dtype=dtype) - fastvideo_camera_net.eval() - - # Run FastVideo CameraNet - with torch.no_grad(): - fastvideo_cam_output = fastvideo_camera_net(cam_input) - - print(f"[CameraNet TEST] FastVideo output shape: {fastvideo_cam_output.shape}") - print( - f"[CameraNet TEST] FastVideo output stats: mean={fastvideo_cam_output.mean():.6f}, std={fastvideo_cam_output.std():.6f}" - ) - - # Compare outputs - print("[CameraNet TEST] Comparing outputs...") - try: - assert_close(cam_output, fastvideo_cam_output, atol=0.0, rtol=0.0) - print("[CameraNet TEST] PASSED: CameraNet outputs match!") - except AssertionError as e: - diff = (cam_output - fastvideo_cam_output).abs() - print(f"[CameraNet TEST] Max diff: {diff.max():.6f}") - print(f"[CameraNet TEST] Mean diff: {diff.mean():.6f}") - print(f"[CameraNet TEST] FAILED: {e}") - raise - - print("[CameraNet TEST] CameraNet parity test complete!") - - -if __name__ == "__main__": - # Run tests directly - test_gamecraft_cameranet_parity() - test_gamecraft_transformer_parity() diff --git a/tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py b/tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py deleted file mode 100644 index 9cf21adec..000000000 --- a/tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py +++ /dev/null @@ -1,508 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -End-to-end pipeline parity test for HunyuanGameCraft. - -This test compares the FastVideo pipeline against the official -Hunyuan-GameCraft-1.0 implementation to ensure numerical alignment -across the complete pipeline (text encoding -> denoising -> VAE decode). - -Usage: - # Run with default paths - DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py -v -s - - # Run with custom paths - GAMECRAFT_OFFICIAL_PATH=path/to/Hunyuan-GameCraft-1.0 \ - GAMECRAFT_WEIGHTS_PATH=path/to/weights \ - DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py -v -s - -Notes: - - This test requires the official Hunyuan-GameCraft-1.0 repo - - This test requires significant GPU memory (~40GB+) - - The test compares denoised latents, not final decoded video (for speed) -""" -import math -import os -import sys -from pathlib import Path - -import pytest -import torch -import numpy as np -from torch.testing import assert_close - -# Set up distributed environment defaults -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29520") -os.environ.setdefault("DISABLE_SP", "1") - -repo_root = Path(__file__).resolve().parents[3] - - -def _log_tensor_stats(label: str, tensor: torch.Tensor) -> None: - """Log tensor statistics for debugging.""" - tensor_f32 = tensor.float() - print(f"[GAMECRAFT PIPELINE] {label}: shape={tuple(tensor.shape)} " - f"dtype={tensor.dtype} device={tensor.device} " - f"min={tensor_f32.min().item():.6f} max={tensor_f32.max().item():.6f} " - f"mean={tensor_f32.mean().item():.6f} std={tensor_f32.std().item():.6f}") - - -def _add_official_to_path(): - """Add official GameCraft implementation to Python path.""" - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - if official_path.exists() and str(official_path) not in sys.path: - sys.path.insert(0, str(official_path)) - return official_path - - -def _create_camera_trajectory_from_action( - action_id: str, - height: int, - width: int, - num_frames: int, - action_speed: float = 0.2, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.bfloat16, -): - """ - Create camera trajectory (Plücker coordinates) from an action ID. - - This replicates the official GameCraft camera trajectory generation. - - Args: - action_id: One of 'w' (forward), 'a' (left), 'd' (right), 's' (backward), - 'left_rot', 'right_rot', 'up_rot', 'down_rot' - height: Video height in pixels - width: Video width in pixels - num_frames: Number of video frames - action_speed: Speed of motion - device: Torch device - dtype: Torch dtype - - Returns: - plucker_embedding: [1, num_frames, 6, height, width] - uncond_plucker_embedding: [1, num_frames, 6, height, width] - """ - official_path = _add_official_to_path() - - try: - from hymm_sp.sample_inference import ( - ActionToPoseFromID, - GetPoseEmbedsFromPoses, - ) - except ImportError as e: - pytest.skip(f"Cannot import official camera trajectory functions: {e}") - - # Generate poses from action - poses = ActionToPoseFromID(action_id, value=action_speed, duration=num_frames) - - # Convert to Plücker embeddings - plucker_embedding, uncond_plucker_embedding, _ = GetPoseEmbedsFromPoses(poses, - height, - width, - num_frames, - flip=False, - start_index=0) - - # Add batch dimension and convert to target dtype/device - plucker_embedding = plucker_embedding.unsqueeze(0).to(device=device, dtype=dtype) - uncond_plucker_embedding = uncond_plucker_embedding.unsqueeze(0).to(device=device, dtype=dtype) - - return plucker_embedding, uncond_plucker_embedding - - -def _run_official_pipeline( - prompt: str, - action_id: str, - height: int, - width: int, - num_frames: int, - num_inference_steps: int, - seed: int, - device: torch.device, - dtype: torch.dtype, - official_path: Path, - weights_path: Path, -): - """ - Run the official GameCraft pipeline and return denoised latents. - - Returns: - denoised_latents: The final denoised latents - timesteps: The timesteps used - """ - sys.path.insert(0, str(official_path)) - - try: - from hymm_sp.config import parse_args - from hymm_sp.sample_inference import HunyuanVideoSampler - from hymm_sp.constants import PRECISION_TO_TYPE - from hymm_sp.vae import load_vae - from hymm_sp.modules import load_model - from hymm_sp.text_encoder import TextEncoder - from hymm_sp.constants import PROMPT_TEMPLATE - except ImportError as e: - pytest.skip(f"Cannot import official GameCraft modules: {e}") - - # Build args for official model - we need to parse default args - # and override specific values - sys.argv = [ - "test", - "--model", - "HYVideo-T/2", - "--prompt", - prompt, - "--vae", - "884-16c-hy", - "--vae-precision", - "fp16", - "--precision", - "bf16", - "--dit-weight", - str(weights_path / "gamecraft_models" / "mp_rank_00_model_states.pt"), - "--video-length", - str(num_frames + 1), # Official uses +1 - "--video-size", - str(height), - str(width), - "--seed", - str(seed), - "--infer-steps", - str(num_inference_steps), - "--flow-shift", - "7.0", - "--embedded-cfg-scale", - "6.0", - "--cfg-scale", - "1.0", - "--text-encoder", - str(weights_path / "text_encoder"), - "--text-encoder-2", - str(weights_path / "text_encoder_2"), - "--tokenizer", - str(weights_path / "text_encoder"), - "--tokenizer-2", - str(weights_path / "text_encoder_2"), - "--prompt-template-video", - "dit-llm-encode-video", - "--load-key", - "module", - ] - - args = parse_args() - args.cpu_offload = 0 - args.use_fp8 = False - - # Create the sampler - sampler = HunyuanVideoSampler.from_pretrained( - pretrained_model_path=str(weights_path / "gamecraft_models"), - args=args, - device=device, - ) - - # Run inference - result = sampler.predict( - prompt=prompt, - is_image=True, # 33 frames mode - size=(height, width), - video_length=num_frames + 1, - seed=seed, - infer_steps=num_inference_steps, - guidance_scale=1.0, # CFG scale - flow_shift=7.0, - action_id=action_id, - action_speed=0.2, - start_index=0, - return_latents=True, - output_type="latent", - ) - - denoised_latents = result.get("denoised_lantents") # Note: typo in official code - timesteps = result.get("timesteps") - - # Clean up - del sampler - torch.cuda.empty_cache() - - return denoised_latents, timesteps - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="GameCraft pipeline test requires CUDA.", -) -def test_gamecraft_pipeline_latent_parity(): - """ - Test that FastVideo GameCraft pipeline produces matching latents. - - This test runs both pipelines with identical: - - Prompt - - Camera trajectory (action) - - Random seed - - Number of inference steps - - And compares the denoised latents (before VAE decode) for numerical parity. - """ - # Check for required paths - official_path = _add_official_to_path() - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - - weights_path = Path( - os.getenv("GAMECRAFT_WEIGHTS_PATH", repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "stdmodels")) - if not weights_path.exists(): - pytest.skip(f"GameCraft weights not found at {weights_path}") - - # Check for required weight files - dit_weights = weights_path / "gamecraft_models" / "mp_rank_00_model_states.pt" - if not dit_weights.exists(): - pytest.skip(f"DiT weights not found at {dit_weights}") - - # Test parameters - small size for faster testing - device = torch.device("cuda:0") - dtype = torch.bfloat16 - prompt = "A first-person view of walking through a lush forest." - action_id = "w" # Forward motion - height = 256 # Smaller for testing - width = 256 - num_frames = 33 # Standard GameCraft frame count - num_inference_steps = 5 # Reduced for testing - seed = 42 - - print(f"\n[GAMECRAFT PIPELINE TEST] Configuration:") - print(f" - Prompt: {prompt}") - print(f" - Action: {action_id} (forward)") - print(f" - Size: {height}x{width}") - print(f" - Frames: {num_frames}") - print(f" - Steps: {num_inference_steps}") - print(f" - Seed: {seed}") - - # ========================================================================= - # Run Official Pipeline - # ========================================================================= - print("\n[GAMECRAFT PIPELINE TEST] Running official pipeline...") - - try: - official_latents, official_timesteps = _run_official_pipeline( - prompt=prompt, - action_id=action_id, - height=height, - width=width, - num_frames=num_frames, - num_inference_steps=num_inference_steps, - seed=seed, - device=device, - dtype=dtype, - official_path=official_path, - weights_path=weights_path, - ) - _log_tensor_stats("Official latents", official_latents) - except Exception as e: - pytest.skip(f"Failed to run official pipeline: {e}") - - # Clear memory - torch.cuda.empty_cache() - - # ========================================================================= - # Run FastVideo Pipeline - # ========================================================================= - print("\n[GAMECRAFT PIPELINE TEST] Running FastVideo pipeline...") - - # Import FastVideo components - from fastvideo import VideoGenerator - - # Get diffusers-format model path - diffusers_path = os.getenv("GAMECRAFT_DIFFUSERS_PATH", str(repo_root / "official_weights" / "hunyuan-gamecraft")) - if not os.path.exists(diffusers_path): - pytest.skip(f"FastVideo GameCraft weights not found at {diffusers_path}") - - # Create camera trajectory for FastVideo - plucker_embedding, uncond_plucker_embedding = _create_camera_trajectory_from_action( - action_id=action_id, - height=height, - width=width, - num_frames=num_frames, - action_speed=0.2, - device=device, - dtype=dtype, - ) - - # Create generator - generator = VideoGenerator.from_pretrained( - diffusers_path, - num_gpus=1, - use_fsdp_inference=False, - dit_cpu_offload=False, - vae_cpu_offload=True, # Save memory - text_encoder_cpu_offload=True, - pin_cpu_memory=False, - ) - - # Run generation - result = generator.generate_video( - prompt=prompt, - output_path="outputs_video/gamecraft_test", - save_video=False, - height=height, - width=width, - num_frames=num_frames, - num_inference_steps=num_inference_steps, - seed=seed, - guidance_scale=1.0, - camera_states=plucker_embedding, - ) - - fastvideo_latents = result.get("samples") - - _log_tensor_stats("FastVideo latents", fastvideo_latents) - - generator.shutdown() - - # ========================================================================= - # Compare Results - # ========================================================================= - print("\n[GAMECRAFT PIPELINE TEST] Comparing latents...") - - # Ensure same shape - assert official_latents.shape == fastvideo_latents.shape, (f"Shape mismatch: official {official_latents.shape} vs " - f"fastvideo {fastvideo_latents.shape}") - - # Compute differences - official_f32 = official_latents.float() - fastvideo_f32 = fastvideo_latents.float() - - abs_diff = (official_f32 - fastvideo_f32).abs() - max_diff = abs_diff.max().item() - mean_diff = abs_diff.mean().item() - - # Compute correlation - official_flat = official_f32.flatten() - fastvideo_flat = fastvideo_f32.flatten() - correlation = torch.corrcoef(torch.stack([official_flat, fastvideo_flat]))[0, 1].item() - - print(f"\n[GAMECRAFT PIPELINE TEST] Results:") - print(f" - Max absolute difference: {max_diff:.6f}") - print(f" - Mean absolute difference: {mean_diff:.6f}") - print(f" - Correlation: {correlation:.6f}") - - # Assertions - # For full pipeline, we expect higher tolerance due to accumulated errors - assert correlation > 0.99, f"Correlation {correlation} is too low (expected > 0.99)" - assert max_diff < 0.5, f"Max diff {max_diff} is too high (expected < 0.5)" - - print("\n[GAMECRAFT PIPELINE TEST] PASSED!") - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="GameCraft pipeline test requires CUDA.", -) -def test_gamecraft_pipeline_camera_trajectory(): - """ - Test that camera trajectory (Plücker embedding) generation matches. - - This is a simpler test that just validates camera conditioning matches - between official and FastVideo implementations. - """ - official_path = _add_official_to_path() - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - - device = torch.device("cuda:0") - dtype = torch.bfloat16 - - # Test parameters - height = 256 - width = 256 - num_frames = 33 - action_id = "w" - action_speed = 0.2 - - print("\n[CAMERA TEST] Generating camera trajectories...") - - # Generate using official code path - plucker, uncond_plucker = _create_camera_trajectory_from_action( - action_id=action_id, - height=height, - width=width, - num_frames=num_frames, - action_speed=action_speed, - device=device, - dtype=dtype, - ) - - _log_tensor_stats("Plücker embedding", plucker) - _log_tensor_stats("Uncond Plücker embedding", uncond_plucker) - - # Basic shape validation - expected_shape = (1, num_frames, 6, height, width) - assert plucker.shape == expected_shape, (f"Plücker shape {plucker.shape} != expected {expected_shape}") - assert uncond_plucker.shape == expected_shape, ( - f"Uncond Plücker shape {uncond_plucker.shape} != expected {expected_shape}") - - # Validate uncond is identity (zeros for translation, identity for rotation) - # In practice, uncond should have specific structure - assert not torch.allclose(plucker, uncond_plucker), ("Conditional and unconditional embeddings should differ") - - print("\n[CAMERA TEST] PASSED!") - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="GameCraft pipeline test requires CUDA.", -) -def test_gamecraft_pipeline_smoke(): - """ - Smoke test for GameCraft pipeline - just verify it runs without errors. - - This is a lighter test that doesn't require the official implementation, - just validates that the FastVideo pipeline can run. - """ - diffusers_path = os.getenv("GAMECRAFT_DIFFUSERS_PATH", str(repo_root / "official_weights" / "hunyuan-gamecraft")) - - if not os.path.exists(diffusers_path): - pytest.skip(f"FastVideo GameCraft weights not found at {diffusers_path}") - - if not os.path.isfile(os.path.join(diffusers_path, "model_index.json")): - pytest.skip(f"model_index.json not found in {diffusers_path}") - - device = torch.device("cuda:0") - - print("\n[SMOKE TEST] Running FastVideo GameCraft pipeline...") - - from fastvideo import VideoGenerator - - # Very small test - height = 128 - width = 128 - num_frames = 9 # Minimum - num_inference_steps = 2 - - generator = VideoGenerator.from_pretrained( - diffusers_path, - num_gpus=1, - use_fsdp_inference=False, - dit_cpu_offload=True, - vae_cpu_offload=True, - text_encoder_cpu_offload=True, - pin_cpu_memory=False, - ) - - result = generator.generate_video( - prompt="A test video.", - output_path="outputs_video/gamecraft_smoke", - save_video=False, - height=height, - width=width, - num_frames=num_frames, - num_inference_steps=num_inference_steps, - seed=123, - guidance_scale=1.0, - ) - - assert result is not None, "Pipeline returned None" - assert "samples" in result or "latents" in result, "No output in result" - - generator.shutdown() - - print("\n[SMOKE TEST] PASSED!") diff --git a/tests/local_tests/gamecraft/test_gamecraft_vae_parity.py b/tests/local_tests/gamecraft/test_gamecraft_vae_parity.py deleted file mode 100644 index 878a9bffd..000000000 --- a/tests/local_tests/gamecraft/test_gamecraft_vae_parity.py +++ /dev/null @@ -1,274 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Numerical parity test for GameCraft VAE (AutoencoderKLCausal3D). - -Ported from official Hunyuan-GameCraft-1.0/hymm_sp/vae. Compares FastVideo's -GameCraftVAE against the official implementation. - -Follows the same pattern as test_ltx2_vae.py: -- Loads same weights into both models (strict key checking) -- Encodes same input, compares latents -- Decodes same latents through both decoders, compares output -- Checks for non-finite values - -Usage: - DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_vae_parity.py -v -""" -import json -import os -import sys -from pathlib import Path - -import pytest -import torch -from torch.testing import assert_close - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29515") -os.environ.setdefault("DISABLE_SP", "1") - -repo_root = Path(__file__).resolve().parents[3] -if str(repo_root) not in sys.path: - sys.path.insert(0, str(repo_root)) - - -def _load_weights(weights_path: Path) -> dict[str, torch.Tensor]: - """Load VAE weights from checkpoint and strip 'vae.' prefix.""" - print(f"[VAE TEST] Loading weights from {weights_path}") - state_dict = torch.load(weights_path, map_location="cpu", weights_only=True) - if "state_dict" in state_dict: - state_dict = state_dict["state_dict"] - # Strip 'vae.' prefix - vae_sd = {k.replace("vae.", ""): v for k, v in state_dict.items() if k.startswith("vae.")} - print(f"[VAE TEST] Loaded {len(vae_sd)} VAE tensors") - return vae_sd - - -def _load_into_model(model: torch.nn.Module, weights: dict[str, torch.Tensor], name: str) -> tuple[int, list[str]]: - """Load weights into a model, checking for missing/unexpected keys.""" - model_state = model.state_dict() - filtered = {k: v for k, v in weights.items() if k in model_state and model_state[k].shape == v.shape} - missing = [k for k in model_state.keys() if k not in filtered] - unexpected = [k for k in weights.keys() if k not in model_state] - print(f"[VAE TEST] {name}: Loading {len(filtered)}/{len(model_state)} tensors " - f"({len(missing)} missing, {len(unexpected)} unexpected)") - if missing: - print(f"[VAE TEST] {name} missing keys: {missing[:5]}{'...' if len(missing) > 5 else ''}") - if not filtered: - return 0, missing - model.load_state_dict(filtered, strict=False) - return len(filtered), missing - - -def _load_official_vae(official_path: Path, vae_sd: dict[str, torch.Tensor], device: torch.device, dtype: torch.dtype): - """Load the official GameCraft VAE from hymm_sp/vae.""" - sys.path.insert(0, str(official_path)) - - try: - from hymm_sp.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D - except ImportError as e: - pytest.skip(f"Failed to import official VAE: {e}") - - config_path = (official_path / "weights" / "stdmodels" / "vae_3d" / "hyvae" / "config.json") - if not config_path.exists(): - pytest.skip(f"VAE config not found at {config_path}") - - with open(config_path) as f: - config = json.load(f) - - vae = AutoencoderKLCausal3D( - in_channels=config.get("in_channels", 3), - out_channels=config.get("out_channels", 3), - down_block_types=tuple(config.get("down_block_types", [])), - up_block_types=tuple(config.get("up_block_types", [])), - block_out_channels=tuple(config.get("block_out_channels", [])), - layers_per_block=config.get("layers_per_block", 2), - act_fn=config.get("act_fn", "silu"), - latent_channels=config.get("latent_channels", 16), - norm_num_groups=config.get("norm_num_groups", 32), - sample_size=config.get("sample_size", 256), - sample_tsize=config.get("sample_tsize", 64), - scaling_factor=config.get("scaling_factor", 0.476986), - time_compression_ratio=config.get("time_compression_ratio", 4), - mid_block_add_attention=config.get("mid_block_add_attention", True), - mid_block_causal_attn=config.get("mid_block_causal_attn", True), - ) - - loaded, missing = _load_into_model(vae, vae_sd, "Official VAE") - if loaded == 0: - pytest.skip("Failed to load weights into official VAE") - if missing: - pytest.skip(f"Official VAE has {len(missing)} missing keys; cannot ensure parity") - - if not hasattr(vae, "use_trt_decoder"): - vae.use_trt_decoder = False - return vae.to(device=device, dtype=dtype).eval() - - -def _load_fastvideo_vae(vae_sd: dict[str, torch.Tensor], device: torch.device, dtype: torch.dtype): - """Load FastVideo GameCraftVAE with official weights.""" - from fastvideo.configs.models.vaes import GameCraftVAEConfig - from fastvideo.models.vaes.gamecraftvae import GameCraftVAE - - config = GameCraftVAEConfig() - vae = GameCraftVAE(config) - - loaded, missing = _load_into_model(vae, vae_sd, "FastVideo VAE") - if loaded == 0: - pytest.skip("Failed to load weights into FastVideo VAE") - if missing: - pytest.skip(f"FastVideo VAE has {len(missing)} missing keys; cannot ensure parity") - - return vae.to(device=device, dtype=dtype).eval() - - -def test_gamecraft_vae_parity(): - """Test VAE encode/decode parity between official and FastVideo. - - Key: decodes the SAME latents through both decoders to isolate - decode-only differences (following LTX2 VAE test pattern). - """ - torch.manual_seed(42) - - official_path = Path(os.getenv("GAMECRAFT_OFFICIAL_PATH", repo_root / "Hunyuan-GameCraft-1.0")) - vae_weights_path = Path( - os.getenv( - "GAMECRAFT_VAE_PATH", - repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "stdmodels" / "vae_3d" / "hyvae" / "pytorch_model.pt", - )) - - if not official_path.exists(): - pytest.skip(f"Official GameCraft repo not found at {official_path}") - if not vae_weights_path.exists(): - alt_paths = [ - vae_weights_path.parent / "checkpoint-step-270000.ckpt", - vae_weights_path.parent / "diffusion_pytorch_model.safetensors", - ] - for alt in alt_paths: - if alt.exists(): - vae_weights_path = alt - break - else: - pytest.skip(f"VAE weights not found. Tried: {vae_weights_path}, {alt_paths}") - - # Use fp32 for precise comparison (bf16 amplifies small differences) - use_cuda = torch.cuda.is_available() - try: - if use_cuda: - _ = torch.zeros(1, device="cuda:0") - del _ - torch.cuda.empty_cache() - except Exception: - use_cuda = False - device = torch.device("cuda:0" if use_cuda else "cpu") - # Use fp32 for parity test precision (bf16 gives ~0.2 max diff due to precision) - dtype = torch.float32 - - # Load weights once, share between both models - vae_sd = _load_weights(vae_weights_path) - - batch_size = 1 - frames = 5 - height = 64 - width = 64 - video_input = torch.randn(batch_size, 3, frames, height, width, device=device, dtype=dtype) - print(f"\n[VAE TEST] Input shape: {video_input.shape}, dtype: {dtype}") - - # --- Encode comparison --- - print("\n[VAE TEST] Loading official VAE...") - official_vae = _load_official_vae(official_path, vae_sd, device, dtype) - with torch.no_grad(): - official_latent_dist = official_vae.encode(video_input) - if hasattr(official_latent_dist, "latent_dist"): - official_latent_dist = official_latent_dist.latent_dist - official_latents = official_latent_dist.mode() - - print("[VAE TEST] Loading FastVideo GameCraftVAE...") - fastvideo_vae = _load_fastvideo_vae(vae_sd, device, dtype) - with torch.no_grad(): - fv_latent_dist = fastvideo_vae.encode(video_input) - fv_latents = fv_latent_dist.latent_dist.mode() - - print(f"[VAE TEST] Official latents: {official_latents.shape}") - print(f"[VAE TEST] FastVideo latents: {fv_latents.shape}") - assert official_latents.shape == fv_latents.shape, "Latent shape mismatch" - assert torch.isfinite(official_latents).all(), "Official encoder produced non-finite latents" - assert torch.isfinite(fv_latents).all(), "FastVideo encoder produced non-finite latents" - - enc_diff = (official_latents - fv_latents).abs().max().item() - enc_mean = (official_latents - fv_latents).abs().mean().item() - print(f"[VAE TEST] Encode max diff: {enc_diff:.6f}, mean diff: {enc_mean:.6f}") - assert_close(official_latents, fv_latents, atol=1e-2, rtol=1e-2) - - # --- Decode comparison (SAME latents through both decoders) --- - # Use official latents as the shared input to isolate decode-only differences - shared_latents = official_latents - - with torch.no_grad(): - official_decoded = official_vae.decode(shared_latents) - if hasattr(official_decoded, "sample"): - official_decoded = official_decoded.sample - - # Free official VAE memory before loading FastVideo decode - del official_vae - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - with torch.no_grad(): - fv_decoded = fastvideo_vae.decode(shared_latents) - if hasattr(fv_decoded, "sample"): - fv_decoded = fv_decoded.sample - - print(f"[VAE TEST] Official decoded: {official_decoded.shape}") - print(f"[VAE TEST] FastVideo decoded: {fv_decoded.shape}") - assert official_decoded.shape == fv_decoded.shape, "Decoded shape mismatch" - assert torch.isfinite(official_decoded).all(), "Official decoder produced non-finite output" - assert torch.isfinite(fv_decoded).all(), "FastVideo decoder produced non-finite output" - - dec_diff = (official_decoded - fv_decoded).abs().max().item() - dec_mean = (official_decoded - fv_decoded).abs().mean().item() - print(f"[VAE TEST] Decode max diff: {dec_diff:.6f}, mean diff: {dec_mean:.6f}") - assert_close(official_decoded, fv_decoded, atol=1e-2, rtol=1e-2) - - print("[VAE TEST] PASSED: Numerical parity between official and FastVideo!") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_gamecraft_vae_config_compatibility(): - """Test that GameCraft VAE config is compatible with FastVideo's GameCraftVAE.""" - vae_config_path = (repo_root / "Hunyuan-GameCraft-1.0" / "weights" / "stdmodels" / "vae_3d" / "hyvae" / - "config.json") - - if not vae_config_path.exists(): - pytest.skip(f"VAE config not found at {vae_config_path}") - - with open(vae_config_path) as f: - gc_config = json.load(f) - - print("\n[VAE CONFIG TEST] GameCraft VAE config:") - for k, v in gc_config.items(): - print(f" {k}: {v}") - - # Expected HunyuanVideo VAE config - expected_config = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 16, - "block_out_channels": [128, 256, 512, 512], - "time_compression_ratio": 4, - "scaling_factor": 0.476986, - } - - # Verify compatibility - for key, expected in expected_config.items(): - if key in gc_config: - actual = gc_config[key] - assert actual == expected, f"Config mismatch for {key}: expected {expected}, got {actual}" - print(f"[VAE CONFIG TEST] {key}: OK (matches expected)") - - print("[VAE CONFIG TEST] Config compatibility check passed!") - - -if __name__ == "__main__": - test_gamecraft_vae_config_compatibility() - test_gamecraft_vae_parity() diff --git a/tests/local_tests/glm_image/README.md b/tests/local_tests/glm_image/README.md deleted file mode 100644 index 32037b77e..000000000 --- a/tests/local_tests/glm_image/README.md +++ /dev/null @@ -1,74 +0,0 @@ -# GLM-Image Local Tests - -Local-only parity and smoke tests for the `glm_image` port: FastVideo vs the -official HF `transformers` + `diffusers` GLM-Image implementation. These are -**not run in CI** — they need the full weights, CUDA, and (for the reference -side) the diffusers GLM-Image classes; each test skips with an actionable -message when a prerequisite is missing. - -## What you need - -| | | -|---|---| -| Weights | `zai-org/GLM-Image` → `official_weights/glm_image` (~30 GB: AR encoder 4 shards, transformer 3 shards, vae + text_encoder) | -| Runtime deps | `transformers>=5.15.0` (repo floor; the AR encoder `GlmImageForConditionalGeneration` first ships in 5.0.0) and `diffusers>=0.38.0` — both committed in `pyproject.toml` | -| Test-only deps | the diffusers `GlmImageTransformer2DModel` / `GlmImagePipeline` reference classes need `diffusers>=0.37.0.dev0` (used **only** by the parity tests) | - -Download weights (from the repo root; cache via `HF_HOME` if `/` is small): - -```bash -python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \ - "zai-org/GLM-Image" "official_weights/glm_image" -``` - -Optionally install the diffusers reference classes to run the parity side: - -```bash -uv pip install -U "git+https://github.com/huggingface/diffusers.git@ff3b86b4" -``` - -Parity was checked against transformers `a65bf6c9` + diffusers `ff3b86b4`. - -## Parity results (GB200) - -| Component | Test | Result | -|---|---|---| -| `transformer` (DiT) | `test_transformer_parity.py` | fp32 cosine `1.000000` (100% within 5e-3); bf16 cosine `0.999969` (element-wise rounding over 30 layers vs a different SDPA kernel — not a bug; RoPE proven exactly equivalent). fp32 + bf16 + strict-load | -| `vae` (`AutoencoderKL`) | `test_vae_parity.py` | decode + encode + config, 3/3 | -| `text_encoder` (T5) + **ByT5** tokenizer | `test_t5_parity.py` | tokenizer + glyph extract + embeds, 3/3 | -| `vision_language_encoder` (GLM-4-9B AR) | `test_ar_parity.py` | surface + to/eval + greedy determinism, 3/3 | -| `pipeline` (T2I) | `test_pipeline_parity.py` | injected-input determinism: latent cosine `0.999997`, decoded-image MAE `0.159/255`. validity + deterministic parity, 2/2 | -| `pipeline` (I2I / edit) | `test_edit_pipeline_parity.py` | validity / wiring gate (condition image enters via a KV-cache write pass, not noise-the-latent). diffusers distribution parity deferred (stochastic AR) | - -`processor` (`GlmImageProcessor`) and `scheduler` (`FlowMatchEulerDiscreteScheduler`, -shared flow-matching) are exercised via the AR and pipeline tests. - -Run: - -```bash -pytest tests/local_tests/glm_image/test_transformer_parity.py -v -s -pytest tests/local_tests/glm_image/test_vae_parity.py -v -s -pytest tests/local_tests/glm_image/test_t5_parity.py -v -s -pytest tests/local_tests/glm_image/test_ar_parity.py -v -s -pytest tests/local_tests/glm_image/test_pipeline_parity.py -v -s -pytest tests/local_tests/glm_image/test_edit_pipeline_parity.py -v -s -``` - -## Design notes (load-bearing — don't regress) - -- **Pipeline parity is by injection.** Full-pipeline pixel parity is ill-posed - (stochastic AR + independent latent RNG), so the pipeline tests inject matched - `(prompt_embeds, prior_token_ids, latents)` into both the diffusers pipeline - and the real `GlmImageDenoisingStage`, then compare denoised latents + decoded - image — not raw end-to-end pixels. -- **HF imports are confined** to the lazy-wrapper loader - `fastvideo/models/encoders/glm_image_ar_loader.py`, never the production - loaders (production-boundary rule). -- **Tokenizer is ByT5**, matching the diffusers reference (not plain T5); - `compute_glyph_embeds` mirrors diffusers `_get_glyph_embeds`. -- **Model-specific stages** live under - `fastvideo/pipelines/basic/glm_image/stages/` per the `add-model` Files Map. -- **Param mapping** is handled in-place at load via `param_names_mapping` - (`fastvideo/configs/models/dits/glm_image.py`) + the native VAE/encoder - configs — no conversion script; the loader raises on any unmatched param and - `test_transformer_strict_load.py` asserts completeness with `strict=True`. diff --git a/tests/local_tests/glm_image/test_ar_parity.py b/tests/local_tests/glm_image/test_ar_parity.py deleted file mode 100644 index 3c4f4f0cf..000000000 --- a/tests/local_tests/glm_image/test_ar_parity.py +++ /dev/null @@ -1,107 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import os -from pathlib import Path - -import pytest -import torch - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") -os.environ.setdefault("DISABLE_SP", "1") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -AR_DIR = LOCAL_WEIGHTS_DIR / "vision_language_encoder" -PROCESSOR_DIR = LOCAL_WEIGHTS_DIR / "processor" - - -def _has_weights() -> bool: - return AR_DIR.exists() and any(AR_DIR.glob("*.safetensors")) - - -def _has_glm_image_transformers() -> bool: - import importlib.util - spec = importlib.util.find_spec("transformers") - if spec is None: - return False - import transformers - return hasattr(transformers, "GlmImageForConditionalGeneration") - - -pytestmark = [ - pytest.mark.skipif( - not _has_weights(), - reason=f"GLM-Image AR encoder weights not found at {AR_DIR}.", - ), - pytest.mark.skipif( - not _has_glm_image_transformers(), - reason=("transformers in this env lacks " - "`GlmImageForConditionalGeneration` (main pins 4.57.3; needs " - "5.0.0rc0+). Bump transformers locally to exercise the AR " - "lazy-wrapper."), - ), -] - -SAMPLE_PROMPT = ("A beautiful landscape photography with rolling hills, a winding river, " - "and a vibrant sunset in the background. Photorealistic style.") - - -@pytest.fixture(scope="module") -def device() -> torch.device: - if not torch.cuda.is_available(): - pytest.skip("CUDA required for GLM-Image AR encoder parity.") - return torch.device("cuda") - - -@pytest.fixture(scope="module") -def fastvideo_ar(device): - pytest.importorskip("fastvideo") - try: - from fastvideo.models.encoders.glm_image_ar_loader import (GlmImageARLoader) - except ImportError as e: - pytest.skip(f"FastVideo lazy-wrapper not yet at target path: {e}") - loader = GlmImageARLoader(str(AR_DIR), str(PROCESSOR_DIR), torch_dtype=torch.bfloat16) - loader.to(device).eval() - return loader - - -def test_ar_wrapper_exposes_expected_surface(fastvideo_ar): - assert callable(getattr(fastvideo_ar, "generate", None)) - assert getattr(fastvideo_ar, "processor", None) is not None - assert getattr(fastvideo_ar, "config", None) is not None - assert getattr(fastvideo_ar, "generation_config", None) is not None - from transformers import GlmImageForConditionalGeneration - assert isinstance(fastvideo_ar._model, GlmImageForConditionalGeneration) - - -def test_ar_to_and_eval_propagate(fastvideo_ar, device): - assert next(fastvideo_ar._model.parameters()).device.type == device.type - assert not fastvideo_ar._model.training - - -def test_ar_generate_is_deterministic_under_do_sample_false(fastvideo_ar, device): - processor = fastvideo_ar.processor - messages = [{ - "role": "user", - "content": [{ - "type": "text", - "text": SAMPLE_PROMPT - }], - }] - inputs = processor.apply_chat_template(messages, - tokenize=True, - target_h=1024, - target_w=1024, - return_dict=True, - return_tensors="pt").to(device) - with torch.no_grad(): - out_a = fastvideo_ar.generate(**inputs, max_new_tokens=32, do_sample=False) - out_b = fastvideo_ar.generate(**inputs, max_new_tokens=32, do_sample=False) - assert out_a.shape == out_b.shape - assert torch.equal(out_a, out_b) diff --git a/tests/local_tests/glm_image/test_edit_pipeline_parity.py b/tests/local_tests/glm_image/test_edit_pipeline_parity.py deleted file mode 100644 index a727d7516..000000000 --- a/tests/local_tests/glm_image/test_edit_pipeline_parity.py +++ /dev/null @@ -1,124 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""End-to-end scaffold for the GLM-Image image-to-image (edit) pipeline. - -Drives FastVideo `VideoGenerator` in edit mode (a condition image is passed) and -checks that the unified pipeline routes through the condition-encoding stage + -KV-cache denoising path and emits a well-formed image. Like the T2I parity test, -pixel-exact parity vs diffusers is NOT asserted: the AR prior is sampled -(`do_sample=True`) and the diffusion latents draw from independent RNG streams, -so the codebases produce different valid samples of the same edit. We gate on -validity (and, when diffusers is present, a comparable brightness regime). - -Skips cleanly until weights are available; GPU-heavy. -""" -from __future__ import annotations - -import os -from pathlib import Path - -import numpy as np -import pytest -import torch -from PIL import Image - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29520") -os.environ.setdefault("DISABLE_SP", "1") -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -CONDITION_IMAGE = REPO_ROOT / "assets" / "images" / "couple.jpg" - - -def _has_weights() -> bool: - required = ["transformer", "vae", "text_encoder", "vision_language_encoder", "processor", "tokenizer", "scheduler"] - return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required) - - -def _upstream_glm_image_available() -> bool: - try: - import transformers - except ImportError: - return False - return hasattr(transformers, "GlmImageForConditionalGeneration") - - -pytestmark = [ - pytest.mark.skipif( - not _has_weights(), - reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.", - ), - pytest.mark.skipif( - not _upstream_glm_image_available(), - reason=("Edit pipeline needs transformers>=5.0.0rc0 " - "(GlmImageForConditionalGeneration); main pin predates it. " - "Bump locally to run."), - ), -] - -EDIT_PROMPT = "Make the scene a snowy winter landscape." -SEED = 0 -HEIGHT = 512 -WIDTH = 512 -STEPS = 8 # low step count keeps the test under 5 min on a single GPU - - -@pytest.fixture(scope="module") -def device(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required for GLM-Image edit pipeline.") - return torch.device("cuda") - - -def _to_uint8_hwc(a) -> np.ndarray: - """Coerce a frame buffer (np or torch, [0,1] float or uint8, with optional - leading batch/time dims, HWC or CHW) to a single (H, W, 3) uint8 image.""" - if torch.is_tensor(a): - a = a.float().cpu().numpy() - a = np.asarray(a) - while a.ndim > 3: - a = a[0] - if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3): - a = np.transpose(a, (1, 2, 0)) - if a.dtype != np.uint8: - scale = 255.0 if float(a.max()) <= 1.5 else 1.0 - a = np.clip(a * scale, 0, 255).astype(np.uint8) - return a - - -def _fastvideo_edit_image(device) -> np.ndarray: - pytest.importorskip("fastvideo") - try: - from fastvideo import VideoGenerator - except ImportError as e: - pytest.skip(f"FastVideo VideoGenerator unavailable: {e}") - condition = Image.open(CONDITION_IMAGE).convert("RGB") - gen = VideoGenerator.from_pretrained(str(LOCAL_WEIGHTS_DIR), num_gpus=1, trust_remote_code=True) - result = gen.generate_video(prompt=EDIT_PROMPT, - pil_image=condition, - save_video=False, - return_frames=True, - height=HEIGHT, - width=WIDTH, - num_inference_steps=STEPS, - guidance_scale=1.5, - seed=SEED) - gen.shutdown() - return _to_uint8_hwc(result["frames"][0]) - - -def test_edit_pipeline_produces_valid_image(device): - """Wiring gate for the edit path: passing a condition image routes through - the condition-encoding (KV-cache write) stage and the read/skip denoising - path, and the pipeline emits a well-formed, non-degenerate image of the - requested size. Component numerical correctness is covered by the (green) - DiT / VAE / T5 / AR component-parity tests.""" - fv = _fastvideo_edit_image(device) - assert fv.shape == (HEIGHT, WIDTH, 3), f"unexpected shape {fv.shape}" - assert fv.dtype == np.uint8 and np.isfinite(fv).all() - assert fv.std() > 10.0, f"image is near-constant (std={fv.std():.2f})" diff --git a/tests/local_tests/glm_image/test_pipeline_parity.py b/tests/local_tests/glm_image/test_pipeline_parity.py deleted file mode 100644 index 2ce23a6e6..000000000 --- a/tests/local_tests/glm_image/test_pipeline_parity.py +++ /dev/null @@ -1,308 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import os -from pathlib import Path - -import numpy as np -import pytest -import torch - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") -os.environ.setdefault("DISABLE_SP", "1") -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer" - - -def _has_weights() -> bool: - required = ["transformer", "vae", "text_encoder", "vision_language_encoder", "processor", "tokenizer", "scheduler"] - return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required) - - -def _upstream_glm_image_available() -> bool: - try: - import transformers - import diffusers - except ImportError: - return False - return (hasattr(transformers, "GlmImageForConditionalGeneration") and hasattr(diffusers, "GlmImagePipeline")) - - -pytestmark = [ - pytest.mark.skipif( - not _has_weights(), - reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.", - ), - pytest.mark.skipif( - not _upstream_glm_image_available(), - reason=("Pipeline parity needs transformers>=5.0.0rc0 and " - "diffusers>=0.37.0.dev0; main pins predate both. Bump locally " - "to run."), - ), -] - -SAMPLE_PROMPT = ("A landscape photo with rolling green hills under a clear blue sky.") -SEED = 0 -HEIGHT = 512 -WIDTH = 512 -STEPS = 8 - -# bf16 + a different SDPA kernel across 30 DiT layers x STEPS steps leaves a -# small residual; a real wiring/schedule bug is far larger than these bounds. -LATENT_COSINE_MIN = 0.995 -IMAGE_MAE_MAX = 5.0 # /255 - - -@pytest.fixture(scope="module") -def device(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required for GLM-Image pipeline parity.") - return torch.device("cuda") - - -def _to_uint8_hwc(a) -> np.ndarray: - if torch.is_tensor(a): - a = a.detach().float().cpu().numpy() - a = np.asarray(a) - while a.ndim > 3: - a = a[0] - if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3): - a = np.transpose(a, (1, 2, 0)) - if a.dtype != np.uint8: - scale = 255.0 if float(a.max()) <= 1.5 else 1.0 - a = np.clip(a * scale, 0, 255).astype(np.uint8) - return a - - -# --------------------------------------------------------------------------- # -# FastVideo native DiT loader (mirrors test_transformer_parity.py). -# --------------------------------------------------------------------------- # -def _load_state_dict(dir_path: Path) -> dict[str, torch.Tensor]: - safetensors = pytest.importorskip("safetensors.torch") - sd: dict[str, torch.Tensor] = {} - for shard in sorted(dir_path.glob("*.safetensors")): - sd.update(safetensors.load_file(str(shard))) - return sd - - -def _apply_param_mapping(sd, mapping): - import re - out = {} - for k, v in sd.items(): - new_k = k - for pat, repl in mapping.items(): - if re.match(pat, k): - new_k = re.sub(pat, repl, k) - break - out[new_k] = v - return out - - -def _ensure_distributed(): - """The denoising stage calls get_local_torch_device(), which needs - FastVideo's world/TP groups. Initialize a single-rank (world_size=1) group - in-process (mirrors tests/local_tests/sd35/test_sd35_component_parity.py).""" - import torch.distributed as dist - from fastvideo.distributed.parallel_state import (get_tp_group, init_distributed_environment, - initialize_model_parallel) - try: - get_tp_group() - return - except Exception: - pass - if not dist.is_initialized(): - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - os.environ.setdefault("LOCAL_RANK", "0") - torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) - store_path = f"/tmp/fastvideo_glm_pg_{os.getpid()}.store" - dist.init_process_group(backend="nccl", init_method=f"file://{store_path}", rank=0, world_size=1) - init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") - try: - get_tp_group() - except Exception: - initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1, data_parallel_size=1) - - -def _load_fastvideo_transformer(device, dtype): - from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig - from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel - cfg = GlmImageDiTConfig() - model = GlmImageTransformer2DModel(cfg, {"_class_name": "GlmImageTransformer2DModel"}) - sd = _load_state_dict(TRANSFORMER_DIR) - sd = _apply_param_mapping(sd, cfg.arch_config.param_names_mapping) - missing, unexpected = model.load_state_dict(sd, strict=False) - assert not missing, f"FastVideo DiT missing keys: {missing[:10]}" - assert not unexpected, f"FastVideo DiT unexpected keys: {unexpected[:10]}" - return model.to(device, dtype=dtype).eval() - - -def _fastvideo_denoise_latents(device, dtype, *, prompt_embeds, prior_token_ids, init_latents): - """Drive the real GlmImageDenoisingStage with injected, matched inputs and - return the denoised latents (1, 16, H/8, W/8).""" - from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (FlowMatchEulerDiscreteScheduler) - from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (calculate_shift) - from fastvideo.pipelines.basic.glm_image.stages.denoising import (GlmImageDenoisingStage) - from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - - _ensure_distributed() - transformer = _load_fastvideo_transformer(device, dtype) - scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0) - stage = GlmImageDenoisingStage(transformer=transformer, scheduler=scheduler) - - # Reproduce the production timestep/sigma schedule (before_denoising.py): - # integer-cast linspace timesteps, resolution-dependent shift on sigmas. - ntt = scheduler.config.num_train_timesteps - patch = transformer.patch_size - image_seq_len = ((HEIGHT // 8) * (WIDTH // 8)) // (patch**2) - sched_t = np.linspace(ntt, 1.0, STEPS + 1)[:-1].astype(np.int64).astype(np.float32) - scheduler.set_shift(calculate_shift(image_seq_len)) - scheduler.set_timesteps(STEPS, device=device, sigmas=(sched_t / ntt).tolist(), timesteps=sched_t.tolist()) - - batch = ForwardBatch(data_type="image") - batch.prompt_embeds = [prompt_embeds.to(device, dtype)] - batch.attention_mask = None # match diffusers (no text padding mask) - batch.prior_token_id = prior_token_ids.to(device) - batch.prior_token_drop = torch.zeros_like(prior_token_ids, dtype=torch.bool, device=device) - batch.latents = init_latents.clone().unsqueeze(2).to(device, dtype) - batch.timesteps = scheduler.timesteps - batch.height, batch.width = HEIGHT, WIDTH - batch.num_inference_steps = STEPS - batch.guidance_scale = 1.5 - batch.do_classifier_free_guidance = True - batch.seed = SEED - batch.extra = {} # no glm_kv_caches -> T2I denoise path - - out = stage.forward(batch, fastvideo_args=None) - latents = out.latents - if latents.dim() == 5: - latents = latents.squeeze(2) - latents = latents.detach().float().cpu() - del transformer, stage - torch.cuda.empty_cache() - import torch.distributed as dist - if dist.is_initialized(): - dist.destroy_process_group() - return latents - - -def _fastvideo_image(device) -> np.ndarray: - pytest.importorskip("fastvideo") - try: - from fastvideo import VideoGenerator - except ImportError as e: - pytest.skip(f"FastVideo VideoGenerator unavailable: {e}") - gen = VideoGenerator.from_pretrained(str(LOCAL_WEIGHTS_DIR), num_gpus=1, trust_remote_code=True) - result = gen.generate_video(prompt=SAMPLE_PROMPT, - save_video=False, - return_frames=True, - height=HEIGHT, - width=WIDTH, - num_inference_steps=STEPS, - guidance_scale=1.5, - seed=SEED) - gen.shutdown() - return _to_uint8_hwc(result["frames"][0]) - - -def test_fastvideo_pipeline_produces_valid_image(device): - fv = _fastvideo_image(device) - assert fv.shape == (HEIGHT, WIDTH, 3), f"unexpected shape {fv.shape}" - assert fv.dtype == np.uint8 and np.isfinite(fv).all() - assert fv.std() > 10.0, f"image is near-constant (std={fv.std():.2f})" - - -def test_pipeline_denoise_parity_deterministic(device): - """Inject identical (prompt_embeds, prior_token_ids, init_latents) into the - diffusers pipeline and FastVideo's GlmImageDenoisingStage, then compare the - denoised latents and the decoded images. With the stochastic AR and the - independent latent RNG removed, the two implementations must agree to within - bf16/SDPA-kernel noise.""" - # Inference only: the T5 glyph encoder leaves a grad graph on the embeds, - # which later breaks image_processor.postprocess()'s .numpy() on the decoded - # image. Disabling grad for the whole test also trims memory. - torch.set_grad_enabled(False) - diffusers = pytest.importorskip("diffusers") - from diffusers.utils.torch_utils import randn_tensor - - dtype = torch.bfloat16 - pipe = diffusers.GlmImagePipeline.from_pretrained(str(LOCAL_WEIGHTS_DIR), torch_dtype=dtype).to(device) - - gen = torch.Generator(device=device).manual_seed(SEED) - - # --- shared, deterministic inputs (computed once, fed to both) ---------- # - pos, neg = pipe.encode_prompt(SAMPLE_PROMPT, do_classifier_free_guidance=True, device=device, dtype=dtype) - assert pos.shape[1] == neg.shape[1], (f"quote-free prompt should give equal pos/neg glyph lengths, got " - f"{pos.shape[1]} vs {neg.shape[1]}") - prior_token_ids, _, _ = pipe.generate_prior_tokens(SAMPLE_PROMPT, - HEIGHT, - WIDTH, - image=None, - device=device, - generator=gen) - latent_ch = pipe.transformer.config.in_channels - init_latents = randn_tensor((1, latent_ch, HEIGHT // 8, WIDTH // 8), generator=gen, device=device, dtype=dtype) - - # --- official denoised latents ----------------------------------------- # - # prompt=None: diffusers check_inputs rejects passing prompt + prompt_embeds - # together; the embeds (and prior tokens) fully specify the run. - official = pipe(prompt=None, - prompt_embeds=pos, - negative_prompt_embeds=neg, - prior_token_ids=prior_token_ids, - latents=init_latents.clone(), - height=HEIGHT, - width=WIDTH, - num_inference_steps=STEPS, - guidance_scale=1.5, - output_type="latent").images.float().cpu() - - # FastVideo packs CFG as a single [pos; neg] 2-row tensor (no padding here - # because L_pos == L_neg, asserted above). - packed = torch.cat([pos, neg], dim=0) - - # --- FastVideo denoised latents (real GlmImageDenoisingStage) ---------- # - fv = _fastvideo_denoise_latents(device, - dtype, - prompt_embeds=packed, - prior_token_ids=prior_token_ids, - init_latents=init_latents) - - assert official.shape == fv.shape, (f"latent shape mismatch: {official.shape} vs {fv.shape}") - - # --- latent-level agreement -------------------------------------------- # - cos = torch.nn.functional.cosine_similarity(official.flatten(), fv.flatten(), dim=0).item() - lat_mae = (official - fv).abs().mean().item() - lat_scale = official.abs().mean().item() - print(f"[glm-image denoise parity] latent cosine={cos:.6f} " - f"MAE={lat_mae:.5f} (|latent| mean={lat_scale:.5f}, " - f"rel={lat_mae / max(lat_scale, 1e-6):.4f})") - - # --- decoded-image agreement (same VAE both sides isolates denoise) ----- # - def _decode(lat): - lat = lat.to(device, dtype) - mean = torch.tensor(pipe.vae.config.latents_mean).view(1, -1, 1, 1).to(device, dtype) - std = torch.tensor(pipe.vae.config.latents_std).view(1, -1, 1, 1).to(device, dtype) - img = pipe.vae.decode((lat * std + mean), return_dict=False)[0] - return _to_uint8_hwc(pipe.image_processor.postprocess(img, output_type="np")[0]) - - off_img, fv_img = _decode(official), _decode(fv) - img_mae = np.abs(off_img.astype(np.float32) - fv_img.astype(np.float32)).mean() - print(f"[glm-image denoise parity] decoded image MAE={img_mae:.3f}/255 " - f"(diffusers mean={off_img.mean():.1f}, fv mean={fv_img.mean():.1f})") - - del pipe - torch.cuda.empty_cache() - - assert cos > LATENT_COSINE_MIN, (f"denoised-latent cosine {cos:.6f} < {LATENT_COSINE_MIN}: the FastVideo " - "and official denoise diverge well beyond kernel noise (wiring bug).") - assert img_mae < IMAGE_MAE_MAX, (f"decoded-image MAE {img_mae:.3f}/255 >= {IMAGE_MAE_MAX}: outputs are " - "not equivalent under matched inputs (wiring/schedule bug).") diff --git a/tests/local_tests/glm_image/test_t5_parity.py b/tests/local_tests/glm_image/test_t5_parity.py deleted file mode 100644 index ef189eca3..000000000 --- a/tests/local_tests/glm_image/test_t5_parity.py +++ /dev/null @@ -1,161 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Parity scaffold for GLM-Image glyph T5 encoder. - -The text_encoder is a small byte-level T5 (d_model=1472, gated-gelu, ByT5 -tokenizer, vocab_size=384). This test exercises the glyph-only encoding path -that the diffusers pipeline implements as `_get_glyph_embeds` and compares -against the FastVideo glyph-encoding stage logic. Until the FastVideo stage is -refactored to mirror diffusers (per-prompt loop, even-length pad-prefix, -attention-mask flattening, left-padded batch), the parity test asserts the -high-impact pieces individually so failures point to the exact path that -diverged. -""" -from __future__ import annotations - -import os -from pathlib import Path - -import pytest -import torch - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") -os.environ.setdefault("DISABLE_SP", "1") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -TEXT_ENCODER_DIR = LOCAL_WEIGHTS_DIR / "text_encoder" -TOKENIZER_DIR = LOCAL_WEIGHTS_DIR / "tokenizer" - - -def _has_weights() -> bool: - return (TEXT_ENCODER_DIR / "model.safetensors").exists() and TOKENIZER_DIR.exists() - - -pytestmark = pytest.mark.skipif( - not _has_weights(), - reason=f"GLM-Image text_encoder/tokenizer not found under {LOCAL_WEIGHTS_DIR}.", -) - -SAMPLE_PROMPT = ('A photo of a coffee shop with a wooden sign reading "Daily Grind" ' - "and a chalk menu next to it") - - -@pytest.fixture(scope="module") -def device() -> torch.device: - if not torch.cuda.is_available(): - pytest.skip("CUDA required for GLM-Image glyph encoder parity.") - return torch.device("cuda") - - -@pytest.fixture(scope="module") -def official_tokenizer(): - transformers = pytest.importorskip("transformers") - return transformers.ByT5Tokenizer.from_pretrained(str(TOKENIZER_DIR)) - - -@pytest.fixture(scope="module") -def official_text_encoder(device): - transformers = pytest.importorskip("transformers") - return transformers.T5EncoderModel.from_pretrained(str(TEXT_ENCODER_DIR), - torch_dtype=torch.float32).to(device).eval() - - -def test_tokenizer_is_byt5(official_tokenizer): - """Glyph tokenizer must be ByT5 (byte-level), not generic T5.""" - assert type(official_tokenizer).__name__ == "ByT5Tokenizer" - assert official_tokenizer.vocab_size <= 400, ("ByT5 vocab is byte-level (~256+special)") - - -def test_glyph_text_extraction(): - """Glyph extraction must pull single+double+CJK quoted spans.""" - pytest.importorskip("fastvideo") - try: - from fastvideo.pipelines.basic.glm_image.stages.before_denoising import ( # noqa - get_glyph_texts) - except ImportError as e: - pytest.skip(f"FastVideo glyph extractor not yet at target path: {e}") - out = get_glyph_texts(SAMPLE_PROMPT) - assert "Daily Grind" in out - - -def _diffusers_reference_get_glyph_embeds(prompts, tokenizer, text_encoder, device, dtype, max_sequence_length=2048): - """Inlined copy of diffusers `_get_glyph_embeds` + `get_glyph_texts`. - - Source: `reference/diffusers/src/diffusers/pipelines/glm_image/` - `pipeline_glm_image.py::GlmImagePipeline._get_glyph_embeds` at commit - ff3b86b4755b46a7b5656dfcf84d25bd25ad4740 (diffusers main, 0.37.0.dev0). - Inlined here because the installed diffusers package (`__version__` - typically <0.37.0) does not yet ship `GlmImagePipeline`. This function is - the reference oracle for FastVideo's `compute_glyph_embeds`. - """ - import re - - def get_glyph_texts(prompt): - if isinstance(prompt, str): - prompt = [prompt] - out = [] - for p in prompt: - out.append( - re.findall(r"'([^']*)'", p) + re.findall(r"“([^“”]*)”", p) + re.findall(r'"([^"]*)"', p) + - re.findall(r"「([^「」]*)」", p)) - return out - - all_glyph_texts = get_glyph_texts(prompts) - all_glyph_embeds = [] - for glyph_texts in all_glyph_texts: - if len(glyph_texts) == 0: - glyph_texts = [""] - input_ids = tokenizer( - glyph_texts, - max_length=max_sequence_length, - truncation=True, - ).input_ids - input_ids = [[tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids for ids in input_ids] - max_length = max(len(ids) for ids in input_ids) - attention_mask = torch.tensor( - [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids], - device=device, - ) - input_ids_t = torch.tensor( - [ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids], - device=device, - ) - outputs = text_encoder(input_ids_t, attention_mask=attention_mask) - glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) - all_glyph_embeds.append(glyph_embeds) - - max_seq_len = max(emb.size(1) for emb in all_glyph_embeds) - padded = [] - for emb in all_glyph_embeds: - if emb.size(1) < max_seq_len: - pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype) - emb = torch.cat([pad, emb], dim=1) - padded.append(emb) - return torch.cat(padded, dim=0).to(device=device, dtype=dtype) - - -def test_get_glyph_embeds_matches_diffusers(official_tokenizer, official_text_encoder, device): - """FastVideo `compute_glyph_embeds` must equal the diffusers reference.""" - pytest.importorskip("fastvideo") - try: - from fastvideo.pipelines.basic.glm_image.stages.before_denoising import ( # noqa - compute_glyph_embeds) - except ImportError as e: - pytest.skip(f"FastVideo glyph encoder not yet at target path: {e}") - - with torch.no_grad(): - official_embeds = _diffusers_reference_get_glyph_embeds([SAMPLE_PROMPT], official_tokenizer, - official_text_encoder, device, torch.float32) - fv_embeds = compute_glyph_embeds([SAMPLE_PROMPT], - tokenizer=official_tokenizer, - text_encoder=official_text_encoder, - device=device, - dtype=torch.float32, - max_sequence_length=2048) - assert fv_embeds.shape == official_embeds.shape, (f"shape mismatch: {fv_embeds.shape} vs {official_embeds.shape}") - torch.testing.assert_close(fv_embeds, official_embeds, atol=1e-4, rtol=1e-4) diff --git a/tests/local_tests/glm_image/test_transformer_parity.py b/tests/local_tests/glm_image/test_transformer_parity.py deleted file mode 100644 index 95aac33e0..000000000 --- a/tests/local_tests/glm_image/test_transformer_parity.py +++ /dev/null @@ -1,193 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Component parity scaffold for GLM-Image DiT (GlmImageTransformer2DModel). - -Compares the FastVideo-native port at `fastvideo.models.dits.glm_image` against -the diffusers reference `GlmImageTransformer2DModel` from -`diffusers.models.transformers.transformer_glm_image`. - -Skips cleanly until both the FastVideo class and the local weights exist; the -test is real (loads weights, forwards a fixed input, compares output tensors). -""" -from __future__ import annotations - -import os -from pathlib import Path - -import pytest -import torch - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") -os.environ.setdefault("DISABLE_SP", "1") -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer" - - -def _has_weights() -> bool: - return TRANSFORMER_DIR.exists() and any(TRANSFORMER_DIR.glob("*.safetensors")) - - -def _has_glm_image_diffusers() -> bool: - """`GlmImageTransformer2DModel` lands in diffusers>=0.37.0.dev0. - - FastVideo main pins `diffusers>=0.33.1`; the installed env is 0.36.0. The - numerical parity test needs the diffusers class as an oracle — until the - installed diffusers is upgraded, this test skips. The structural surface - is still verified by `test_glm_image_transformer_strict_load.py`. - """ - import importlib.util - spec = importlib.util.find_spec("diffusers") - if spec is None: - return False - import diffusers - return hasattr(diffusers, "GlmImageTransformer2DModel") - - -pytestmark = [ - pytest.mark.skipif( - not _has_weights(), - reason=(f"GLM-Image transformer weights not found at {TRANSFORMER_DIR}. " - "Download via " - "`python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py " - "zai-org/GLM-Image official_weights/glm_image`."), - ), - pytest.mark.skipif( - not _has_glm_image_diffusers(), - reason=("diffusers in this env lacks `GlmImageTransformer2DModel` " - "(main allows >=0.33.1; installed 0.36.0; class lands in " - "0.37.0.dev0). Bump diffusers locally to run numerical parity. " - "Structural surface is covered by " - "`test_glm_image_transformer_strict_load.py`."), - ), -] - - -@pytest.fixture(scope="module") -def device() -> torch.device: - if not torch.cuda.is_available(): - pytest.skip("CUDA required for GLM-Image transformer parity.") - return torch.device("cuda") - - -def _load_official(device, dtype): - diffusers = pytest.importorskip("diffusers") - cls = diffusers.GlmImageTransformer2DModel - return cls.from_pretrained(str(TRANSFORMER_DIR), torch_dtype=dtype).to(device).eval() - - -def _load_fastvideo(device, dtype): - pytest.importorskip("fastvideo") - try: - from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig - from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel - except ImportError as e: - pytest.skip(f"FastVideo GLM-Image DiT not yet ported: {e}") - cfg = GlmImageDiTConfig() - model = GlmImageTransformer2DModel(cfg, {"_class_name": "GlmImageTransformer2DModel"}) - sd = _load_state_dict(TRANSFORMER_DIR) - sd = _apply_param_mapping(sd, cfg.arch_config.param_names_mapping) - missing, unexpected = model.load_state_dict(sd, strict=False) - assert not missing, f"FastVideo DiT missing keys: {missing[:10]}" - assert not unexpected, f"FastVideo DiT unexpected keys: {unexpected[:10]}" - return model.to(device, dtype=dtype).eval() - - -def _forward_pair(device, dtype): - """Run both DiTs on identical synthetic inputs; return (fv_out, ref_out).""" - from fastvideo.forward_context import set_forward_context - ref = _load_official(device, dtype) - fv = _load_fastvideo(device, dtype) - inputs = _make_inputs(device, dtype=dtype) - with torch.no_grad(): - ref_out = ref(**inputs, return_dict=False)[0] - with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=None): - fv_out = fv(**inputs) - if isinstance(fv_out, tuple): - fv_out = fv_out[0] - assert ref_out.shape == fv_out.shape, (f"shape mismatch: {ref_out.shape} vs {fv_out.shape}") - # Free the two 7B models before the next dtype runs (single-GPU budget). - del ref, fv - torch.cuda.empty_cache() - return fv_out.float(), ref_out.float() - - -def _load_state_dict(dir_path: Path) -> dict[str, torch.Tensor]: - safetensors = pytest.importorskip("safetensors.torch") - sd: dict[str, torch.Tensor] = {} - for shard in sorted(dir_path.glob("*.safetensors")): - sd.update(safetensors.load_file(str(shard))) - return sd - - -def _apply_param_mapping(sd, mapping): - import re - out = {} - for k, v in sd.items(): - new_k = k - for pat, repl in mapping.items(): - if re.match(pat, k): - new_k = re.sub(pat, repl, k) - break - out[new_k] = v - return out - - -def _make_inputs(device, - batch_size=1, - height=64, - width=64, - text_seq_len=32, - dtype=torch.bfloat16) -> dict[str, torch.Tensor]: - """Synthetic inputs matching the diffusers `GlmImageTransformer2DModel.forward` - contract. Latent shape is `(B, 16, H, W)`; post-patch grid is `(H/2, W/2)` - so `prior_token_id` has `(H/2)*(W/2)` entries per batch element.""" - torch.manual_seed(0) - patch_size = 2 - num_patches = (height // patch_size) * (width // patch_size) - return { - "hidden_states": torch.randn(batch_size, 16, height, width, device=device, dtype=dtype), - "encoder_hidden_states": torch.randn(batch_size, text_seq_len, 1472, device=device, dtype=dtype), - "prior_token_id": torch.randint(0, 16384, (batch_size, num_patches), device=device, dtype=torch.long), - "prior_token_drop": torch.zeros(batch_size, device=device, dtype=torch.bool), - "timestep": torch.tensor([500] * batch_size, device=device, dtype=torch.long), - "target_size": torch.tensor([[height * 8, width * 8]] * batch_size, device=device, dtype=torch.long), - "crop_coords": torch.zeros(batch_size, 2, device=device, dtype=torch.long), - } - - -# fp32 is the structural-correctness gate: with identical math the two -# implementations agree to ~1e-3 (residual is SDPA reduction-order noise). -ATOL_FP32 = 1e-2 -RTOL_FP32 = 1e-2 -# bf16 is the inference dtype. Rounding accumulates across 30 layers and the -# attention kernels differ, so element-wise atol is meaningless; cosine -# similarity is the right statistical bar for "same function, bf16 noise". -COS_MIN_BF16 = 0.999 - - -def test_transformer_forward_matches_diffusers_fp32(device): - """Structural gate: in fp32 the FastVideo DiT is numerically faithful to - the diffusers reference (cosine ~1.0, element-wise close to ~1e-3).""" - fv_out, ref_out = _forward_pair(device, torch.float32) - cos = torch.nn.functional.cosine_similarity(fv_out.flatten(), ref_out.flatten(), dim=0) - assert cos > 0.9999, f"fp32 cosine similarity too low: {cos:.6f}" - torch.testing.assert_close(fv_out, ref_out, atol=ATOL_FP32, rtol=RTOL_FP32) - - -def test_transformer_forward_matches_diffusers_bf16(device): - """Inference-dtype sanity: bf16 output tracks the diffusers reference up to - precision accumulation (high cosine similarity). Element-wise atol is not a - meaningful bar for a 30-layer bf16 transformer vs a different attention - kernel (~28% of elements exceed atol=5e-3 purely from rounding) — see - PORT_STATUS I020.""" - fv_out, ref_out = _forward_pair(device, torch.bfloat16) - cos = torch.nn.functional.cosine_similarity(fv_out.flatten(), ref_out.flatten(), dim=0) - assert cos > COS_MIN_BF16, (f"bf16 cosine similarity {cos:.6f} below {COS_MIN_BF16}; " - "indicates a structural divergence, not just precision.") diff --git a/tests/local_tests/glm_image/test_transformer_strict_load.py b/tests/local_tests/glm_image/test_transformer_strict_load.py deleted file mode 100644 index 0aab4d32b..000000000 --- a/tests/local_tests/glm_image/test_transformer_strict_load.py +++ /dev/null @@ -1,93 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Strict-load smoke for GLM-Image DiT. - -Loads all three transformer shards from `official_weights/glm_image/transformer/`, -applies the FastVideo `param_names_mapping`, and calls -`GlmImageTransformer2DModel.load_state_dict(..., strict=True)`. Asserts the -config-driven model instantiates and every checkpoint key has a matching -FastVideo parameter with the same shape and dtype. - -This is weight-load verification only — it does not run a forward pass. Numerical -parity against the diffusers reference lives in -`test_glm_image_transformer_parity.py` and requires `diffusers >= 0.37.0.dev0`. -""" -from __future__ import annotations - -import os -import re -from pathlib import Path - -import pytest -import torch - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") -os.environ.setdefault("DISABLE_SP", "1") -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer" - - -def _has_shards() -> bool: - return TRANSFORMER_DIR.exists() and any(TRANSFORMER_DIR.glob("*.safetensors")) - - -pytestmark = pytest.mark.skipif( - not _has_shards(), - reason=f"GLM-Image transformer shards not at {TRANSFORMER_DIR}.", -) - - -def test_dit_strict_load_against_hf_checkpoint(): - safetensors = pytest.importorskip("safetensors.torch") - from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig - from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel - - cfg = GlmImageDiTConfig() - mapping = cfg.arch_config.param_names_mapping - - # 1. Build the FastVideo model on CPU/meta-free path. - torch.set_default_device("cpu") - hf_config = {"_class_name": "GlmImageTransformer2DModel"} - model = GlmImageTransformer2DModel(cfg, hf_config) - - # 2. Load every shard and apply param_names_mapping. - raw_sd: dict[str, torch.Tensor] = {} - for shard in sorted(TRANSFORMER_DIR.glob("*.safetensors")): - raw_sd.update(safetensors.load_file(str(shard))) - - def rename(k: str) -> str: - for pat, repl in mapping.items(): - if re.match(pat, k): - return re.sub(pat, repl, k) - return k - - renamed_sd = {rename(k): v for k, v in raw_sd.items()} - - # 3. Compare key sets (sanity, redundant with strict=True but gives a - # clearer error message if it fails). - fv_keys = set(model.state_dict().keys()) - ckpt_keys = set(renamed_sd.keys()) - missing = fv_keys - ckpt_keys - unexpected = ckpt_keys - fv_keys - assert not missing, f"FastVideo DiT missing {len(missing)} keys, e.g. {sorted(missing)[:5]}" - assert not unexpected, f"checkpoint has {len(unexpected)} unexpected keys, e.g. {sorted(unexpected)[:5]}" - - # 4. Shape and dtype sanity — strict_load with shape mismatch raises early. - fv_sd = model.state_dict() - shape_mismatches = [] - for k, v in renamed_sd.items(): - if v.shape != fv_sd[k].shape: - shape_mismatches.append((k, tuple(fv_sd[k].shape), tuple(v.shape))) - assert not shape_mismatches, (f"shape mismatches: {shape_mismatches[:5]}") - - # 5. Strict load. - incompatible = model.load_state_dict(renamed_sd, strict=True) - assert not incompatible.missing_keys - assert not incompatible.unexpected_keys diff --git a/tests/local_tests/glm_image/test_vae_parity.py b/tests/local_tests/glm_image/test_vae_parity.py deleted file mode 100644 index 673e724c8..000000000 --- a/tests/local_tests/glm_image/test_vae_parity.py +++ /dev/null @@ -1,104 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Component parity scaffold for GLM-Image VAE (AutoencoderKL). - -Compares the FastVideo-native AutoencoderKL port against diffusers' -`AutoencoderKL` loaded from `zai-org/GLM-Image/vae`. The published checkpoint -uses `block_out_channels=[128, 512, 1024, 1024]`, `latent_channels=16`, and -per-channel `latents_mean[16]` + `latents_std[16]` normalization (not a scalar -`scaling_factor`). - -Skips cleanly until both the FastVideo class and the VAE shard exist. -""" -from __future__ import annotations - -import json -import os -from pathlib import Path - -import pytest -import torch - -import fastvideo.envs as envs - -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") -os.environ.setdefault("DISABLE_SP", "1") - -REPO_ROOT = Path(__file__).resolve().parents[3] -FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or - REPO_ROOT / "official_weights" / FAMILY) -VAE_DIR = LOCAL_WEIGHTS_DIR / "vae" - - -def _has_weights() -> bool: - return (VAE_DIR / "diffusion_pytorch_model.safetensors").exists() - - -pytestmark = pytest.mark.skipif( - not _has_weights(), - reason=f"GLM-Image VAE weights not found at {VAE_DIR}.", -) - - -@pytest.fixture(scope="module") -def device() -> torch.device: - if not torch.cuda.is_available(): - pytest.skip("CUDA required for GLM-Image VAE parity.") - return torch.device("cuda") - - -@pytest.fixture(scope="module") -def vae_config(): - with open(VAE_DIR / "config.json") as f: - return json.load(f) - - -@pytest.fixture(scope="module") -def official_vae(device): - diffusers = pytest.importorskip("diffusers") - vae = diffusers.AutoencoderKL.from_pretrained(str(VAE_DIR), torch_dtype=torch.float32).to(device).eval() - return vae - - -@pytest.fixture(scope="module") -def fastvideo_vae(device): - pytest.importorskip("fastvideo") - try: - from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig - from fastvideo.models.vaes.autoencoder_kl import AutoencoderKL - except ImportError as e: - pytest.skip(f"FastVideo native AutoencoderKL not yet ported: {e}") - cfg = GlmImageVAEConfig() - vae = AutoencoderKL(cfg) - safetensors = pytest.importorskip("safetensors.torch") - sd = safetensors.load_file(str(VAE_DIR / "diffusion_pytorch_model.safetensors")) - missing, unexpected = vae.load_state_dict(sd, strict=False) - assert not missing, f"FastVideo VAE missing keys: {missing[:10]}" - assert not unexpected, f"FastVideo VAE unexpected keys: {unexpected[:10]}" - return vae.to(device, dtype=torch.float32).eval() - - -def test_vae_config_matches_real_checkpoint(vae_config): - assert vae_config["block_out_channels"] == [128, 512, 1024, 1024] - assert vae_config["latent_channels"] == 16 - assert "latents_mean" in vae_config and len(vae_config["latents_mean"]) == 16 - assert "latents_std" in vae_config and len(vae_config["latents_std"]) == 16 - - -def test_vae_decode_matches_diffusers(official_vae, fastvideo_vae, device): - torch.manual_seed(0) - latents = torch.randn(1, 16, 32, 32, device=device, dtype=torch.float32) - with torch.no_grad(): - official_out = official_vae.decode(latents, return_dict=False)[0] - fv_out = fastvideo_vae.decode(latents, return_dict=False)[0] - torch.testing.assert_close(fv_out, official_out, atol=1e-3, rtol=1e-3) - - -def test_vae_encode_matches_diffusers(official_vae, fastvideo_vae, device): - torch.manual_seed(0) - pixels = torch.randn(1, 3, 256, 256, device=device, dtype=torch.float32) - with torch.no_grad(): - official_z = official_vae.encode(pixels).latent_dist.mode() - fv_z = fastvideo_vae.encode(pixels).latent_dist.mode() - torch.testing.assert_close(fv_z, official_z, atol=1e-3, rtol=1e-3) diff --git a/tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py b/tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py deleted file mode 100644 index f10288994..000000000 --- a/tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py +++ /dev/null @@ -1,105 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""DreamX-World pipeline parity checks. - -This test compares the FastVideo pipeline's DreamX-specific conditioning and -single-step scheduler path against an explicit hand-rolled pass using the same -loaded modules. It is intentionally local and deterministic: component parity -against the official DreamX repository lives in ``tests/local_tests/dreamx_world``. -""" -from __future__ import annotations - -import gc -import os -from pathlib import Path -from typing import Any, cast - -import pytest -import torch -from torch.testing import assert_close - -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -MODEL_DIR = Path(os.getenv("DREAMX_WORLD_MODEL_DIR", "converted_weights/dreamx_world")) - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), - reason="DreamX-World pipeline parity requires CUDA", -) - - -def _run_worker_forward_batch(worker_wrapper: Any, request_kwargs: dict[str, Any]) -> torch.Tensor: - from fastvideo.api.sampling_param import SamplingParam - from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - from fastvideo.utils import shallow_asdict - - fastvideo_args = worker_wrapper.worker.fastvideo_args - sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path) - sampling_param.update({key: value for key, value in request_kwargs.items() if key not in {"prompt", "output_path"}}) - sampling_param.prompt = request_kwargs["prompt"] - - latents_size = [ - (sampling_param.num_frames - 1) // 4 + 1, - sampling_param.height // 8, - sampling_param.width // 8, - ] - n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - batch = ForwardBatch( - **shallow_asdict(sampling_param), - eta=0.0, - n_tokens=n_tokens, - VSA_sparsity=fastvideo_args.VSA_sparsity, - ) - output_batch = worker_wrapper.worker.pipeline.forward(batch, fastvideo_args) - assert output_batch.output is not None - return output_batch.output.detach().cpu() - - -def _close_generator(generator: Any) -> None: - generator.shutdown() - gc.collect() - torch.cuda.empty_cache() - - -def test_dreamx_world_one_step_pipeline_latent_matches_manual_pass() -> None: - if not MODEL_DIR.exists(): - pytest.fail(f"DreamX-World converted model directory is missing: {MODEL_DIR}") - - from fastvideo import VideoGenerator - common_kwargs = dict( - prompt="a quiet road through a futuristic city at sunrise", - output_path="outputs_video/dreamx_world_parity", - save_video=False, - return_frames=True, - height=64, - width=64, - num_frames=9, - num_inference_steps=1, - guidance_scale=1.0, - action_list=["w"], - action_speed_list=[2.0], - seed=123, - ) - - generator = VideoGenerator.from_pretrained( - str(MODEL_DIR), - num_gpus=1, - use_fsdp_inference=False, - dit_cpu_offload=False, - vae_cpu_offload=True, - text_encoder_cpu_offload=True, - pin_cpu_memory=False, - output_type="latent", - override_pipeline_cls_name="DreamXWorldPipeline", - ) - try: - result = cast(dict[str, Any], generator.generate_video(**common_kwargs)) - pipeline_latents = cast(torch.Tensor, result["samples"]).detach().cpu() - - manual_latents = generator.executor.collective_rpc( - _run_worker_forward_batch, - kwargs={"request_kwargs": common_kwargs}, - )[0] - finally: - _close_generator(generator) - - assert_close(pipeline_latents, manual_latents, atol=0.0, rtol=0.0) diff --git a/tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py b/tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py deleted file mode 100644 index 29b80036b..000000000 --- a/tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py +++ /dev/null @@ -1,155 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Smoke tests for the DreamX-World-5B-Cam pipeline.""" -from __future__ import annotations - -import os -from pathlib import Path -from typing import Any, cast - -import pytest -import torch -from PIL import Image, ImageDraw - -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -MODEL_DIR = Path(os.getenv("DREAMX_WORLD_MODEL_DIR", "converted_weights/dreamx_world")) - - -def _write_smoke_image(path: Path) -> None: - image = Image.new("RGB", (96, 96), color=(42, 76, 112)) - draw = ImageDraw.Draw(image) - draw.rectangle((0, 56, 96, 96), fill=(36, 44, 52)) - draw.polygon([(0, 56), (48, 30), (96, 56)], fill=(120, 142, 158)) - draw.rectangle((34, 42, 62, 72), fill=(178, 198, 212)) - image.save(path) - - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), - reason="DreamX-World pipeline smoke requires CUDA", -) - - -def test_dreamx_world_typed_surface_preflight() -> None: - import fastvideo.registry as registry - from fastvideo.api.presets import get_preset, get_presets_for_family - from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig - from fastvideo.fastvideo_args import WorkloadType - from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import ( - DreamXWorldPipeline, - EntryClass, - ) - - assert DreamXWorldPipeline.__name__ == "DreamXWorldPipeline" - assert EntryClass is DreamXWorldPipeline - assert DreamXWorldPipeline._required_config_modules == [ - "text_encoder", - "tokenizer", - "vae", - "transformer", - "scheduler", - ] - - default_preset, model_family = registry.get_preset_selection("GD-ML/DreamX-World-5B-Cam") - assert model_family == "dreamx_world" - assert default_preset == "dreamx_world_5b_cam" - - info = registry.get_model_info( - "GD-ML/DreamX-World-5B-Cam", - workload_type=WorkloadType.I2V, - override_pipeline_cls_name="DreamXWorldPipeline", - ) - assert info.pipeline_cls is DreamXWorldPipeline - assert info.pipeline_config_cls is DreamXWorld5BCamPipelineConfig - - names = {p.name for p in get_presets_for_family("dreamx_world")} - assert "dreamx_world_5b_cam" in names - preset = get_preset("dreamx_world_5b_cam", "dreamx_world") - assert preset.defaults["num_inference_steps"] == 30 - assert preset.defaults["height"] == 480 - assert preset.defaults["width"] == 832 - assert preset.defaults["num_frames"] == 161 - assert preset.defaults["guidance_scale"] == 5.0 - - cfg = DreamXWorld5BCamPipelineConfig() - assert cfg.flow_shift == 3.0 - assert cfg.ti2v_task is True - assert cfg.expand_timesteps is True - assert cfg.dit_config.arch_config.add_control_adapter is True - assert cfg.dit_config.arch_config.cam_method == "prope" - - -def test_dreamx_world_camera_stage_writes_y_camera() -> None: - from types import SimpleNamespace - - from fastvideo.pipelines.basic.dreamx_world.stages import ( - DREAMX_Y_CAMERA_KEY, - DreamXWorldCameraConditioningStage, - ) - from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - - batch = ForwardBatch( - data_type="video", - prompt="camera smoke", - latents=torch.zeros(1, 48, 3, 8, 8, dtype=torch.bfloat16, device="cuda"), - num_frames=9, - height=64, - width=64, - action_list=["w", "d"], - action_speed_list=[2.0, 1.0], - ) - out = DreamXWorldCameraConditioningStage().forward(batch, cast(Any, SimpleNamespace())) - - y_camera = out.extra[DREAMX_Y_CAMERA_KEY] - assert set(y_camera) == {"viewmats", "K"} - assert y_camera["viewmats"].shape == (1, 3, 4, 4) - assert y_camera["K"].shape == (1, 3, 3, 3) - assert y_camera["viewmats"].device.type == "cuda" - assert y_camera["viewmats"].dtype == torch.bfloat16 - - -def test_dreamx_world_pipeline_load_generate_latent_smoke(tmp_path: Path) -> None: - if not MODEL_DIR.exists(): - pytest.fail(f"DreamX-World converted model directory is missing: {MODEL_DIR}") - - from fastvideo import VideoGenerator - - image_path = tmp_path / "dreamx_world_smoke_input.png" - _write_smoke_image(image_path) - - generator = VideoGenerator.from_pretrained( - str(MODEL_DIR), - num_gpus=1, - use_fsdp_inference=False, - dit_cpu_offload=False, - vae_cpu_offload=True, - text_encoder_cpu_offload=True, - pin_cpu_memory=False, - output_type="latent", - override_pipeline_cls_name="DreamXWorldPipeline", - ) - try: - result = generator.generate_video( - prompt="a quiet road through a futuristic city at sunrise", - output_path="outputs_video/dreamx_world_smoke", - save_video=False, - return_frames=True, - height=64, - width=64, - num_frames=9, - num_inference_steps=1, - guidance_scale=1.0, - image_path=str(image_path), - action_list=["w"], - action_speed_list=[2.0], - seed=0, - ) - finally: - generator.shutdown() - - assert isinstance(result, dict) - samples = cast(dict[str, Any], result)["samples"] - assert torch.is_tensor(samples) - assert samples.ndim == 5 - assert samples.shape[1] == 48 - assert torch.isfinite(samples).all() diff --git a/tests/local_tests/pipelines/test_zimage_pipeline_parity.py b/tests/local_tests/pipelines/test_zimage_pipeline_parity.py deleted file mode 100644 index b3aadddb2..000000000 --- a/tests/local_tests/pipelines/test_zimage_pipeline_parity.py +++ /dev/null @@ -1,326 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Native Z-Image stage contract and pinned-repository pipeline parity.""" - -from __future__ import annotations - -import gc -import importlib -import os -import subprocess -import sys -from contextlib import contextmanager -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch -from torch.testing import assert_close - -import fastvideo.envs as envs -from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, ) -from fastvideo.pipelines.basic.zimage.stages import ( - ZImageConditioningStage, - ZImageDecodingStage, - ZImageDenoisingStage, - ZImageInputValidationStage, - ZImageLatentPreparationStage, - ZImageTimestepPreparationStage, -) -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - -REPO_ROOT = Path(__file__).resolve().parents[3] -REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" -REFERENCE_REPO = Path(os.getenv("ZIMAGE_REFERENCE_REPO", REPO_ROOT / "Z-Image")) -MODEL_DIR = Path(envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() if envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.is_set() else - REPO_ROOT / "official_weights" / "Z-Image") - - -class _ConstantTransformer: - in_channels = 16 - - def __init__(self) -> None: - self.calls: list[tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]] = [] - - def __call__( - self, - *, - hidden_states: torch.Tensor, - encoder_hidden_states: list[torch.Tensor], - timestep: torch.Tensor, - ) -> tuple[list[torch.Tensor], dict]: - self.calls.append((hidden_states, encoder_hidden_states, timestep)) - return [torch.ones_like(sample) for sample in hidden_states.unbind(0)], {} - - -class _RecordingVAE(torch.nn.Module): - - def __init__(self) -> None: - super().__init__() - self.anchor = torch.nn.Parameter(torch.zeros(())) - self.config = SimpleNamespace(scaling_factor=0.3611, shift_factor=0.1159) - self.decode_input: torch.Tensor | None = None - - def decode(self, latents: torch.Tensor, return_dict: bool = False) -> tuple[torch.Tensor]: - assert return_dict is False - self.decode_input = latents.detach().clone() - batch, _, height, width = latents.shape - return (torch.zeros(batch, 3, height * 8, width * 8, device=latents.device), ) - - -def _stage_args(*, output_type: str = "pil") -> SimpleNamespace: - return SimpleNamespace( - pipeline_config=ZImagePipelineConfig(), - output_type=output_type, - vae_cpu_offload=False, - disable_autocast=False, - ) - - -def test_zimage_native_default_stage_math(monkeypatch: pytest.MonkeyPatch) -> None: - """Exercise the pinned native 1-step formulas without model assets.""" - autocast_calls: list[dict[str, object]] = [] - - @contextmanager - def record_autocast(**kwargs): - autocast_calls.append(kwargs) - yield - - monkeypatch.setattr( - "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device", - lambda: torch.device("cpu"), - ) - monkeypatch.setattr("fastvideo.pipelines.basic.zimage.stages.torch.autocast", record_autocast) - args = _stage_args() - transformer = _ConstantTransformer() - scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) - vae = _RecordingVAE() - batch = ForwardBatch( - data_type="image", - prompt="test prompt", - negative_prompt="", - prompt_embeds=[torch.arange(32, dtype=torch.float32).reshape(1, 4, 8)], - prompt_attention_mask=[torch.tensor([[True, True, False, False]])], - height=64, - width=64, - num_frames=1, - num_inference_steps=1, - guidance_scale=0.0, - seed=42, - latents=torch.zeros(1, 16, 1, 8, 8), - ) - - ZImageInputValidationStage().forward(batch, args) - ZImageConditioningStage().forward(batch, args) - ZImageLatentPreparationStage(transformer).forward(batch, args) - ZImageTimestepPreparationStage(scheduler).forward(batch, args) - - assert len(batch.extra["zimage_prompt_embeds"]) == 1 - assert batch.extra["zimage_prompt_embeds"][0].shape == (2, 8) - assert_close(batch.timesteps, torch.tensor([1000.0])) - assert scheduler.sigma_min == 0.0 - assert scheduler.config.use_reference_discrete_timesteps is True - - ZImageDenoisingStage(transformer, scheduler).forward(batch, args) - assert autocast_calls == [{"device_type": "cpu", "enabled": False}] - assert_close(batch.latents, torch.ones_like(batch.latents)) - _, encoded, normalized_timestep = transformer.calls[0] - assert encoded[0].shape == (2, 8) - assert_close(normalized_timestep, torch.tensor([0.0])) - - ZImageDecodingStage(vae).forward(batch, args) - assert vae.decode_input is not None - expected_vae_input = torch.full((1, 16, 8, 8), 1 / 0.3611 + 0.1159) - assert_close(vae.decode_input, expected_vae_input) - assert batch.output is not None - assert batch.output.shape == (1, 3, 1, 64, 64) - assert_close(batch.output, torch.full_like(batch.output, 0.5)) - - -def test_zimage_batch_rng_matches_one_native_generator(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device", - lambda: torch.device("cpu"), - ) - args = _stage_args(output_type="latent") - transformer = _ConstantTransformer() - batch = ForwardBatch( - data_type="image", - prompt=["one", "two"], - negative_prompt="", - height=64, - width=64, - num_frames=1, - num_inference_steps=1, - guidance_scale=0.0, - seed=42, - ) - ZImageInputValidationStage().forward(batch, args) - batch.extra["zimage_prompt_embeds"] = [torch.zeros(1, 8), torch.zeros(1, 8)] - ZImageLatentPreparationStage(transformer).forward(batch, args) - - expected = torch.randn( - (2, 16, 1, 8, 8), - generator=torch.Generator("cpu").manual_seed(42), - ) - assert_close(batch.latents, expected) - - -def test_zimage_text_stage_honors_request_max_length(monkeypatch: pytest.MonkeyPatch) -> None: - from fastvideo.pipelines.stages.text_encoding import TextEncodingStage - - args = _stage_args() - stage = TextEncodingStage(text_encoders=[object()], tokenizers=[object()]) - observed: list[int | None] = [] - - def fake_encode_text( - text, - fastvideo_args, - encoder_index, - return_attention_mask, - max_length=None, - ): - del text, fastvideo_args, encoder_index, return_attention_mask - observed.append(max_length) - return [torch.zeros(1, 4, 8)], [torch.ones(1, 4, dtype=torch.long)] - - monkeypatch.setattr(stage, "encode_text", fake_encode_text) - batch = ForwardBatch( - data_type="image", - prompt="test", - negative_prompt="", - max_sequence_length=64, - guidance_scale=0.0, - ) - stage.forward(batch, args) - assert observed == [64] - - -def test_zimage_cfg_none_uses_native_empty_negative_prompt(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device", - lambda: torch.device("cpu"), - ) - batch = ForwardBatch( - data_type="image", - prompt="test", - negative_prompt=None, - height=64, - width=64, - num_frames=1, - num_inference_steps=1, - guidance_scale=2.0, - seed=42, - ) - ZImageInputValidationStage().forward(batch, _stage_args()) - assert batch.negative_prompt == "" - - -def _require_pinned_reference() -> Path: - if not REFERENCE_REPO.is_dir(): - pytest.skip(f"Pinned Tongyi Z-Image clone not found: {REFERENCE_REPO}") - result = subprocess.run( - ["git", "-C", str(REFERENCE_REPO), "rev-parse", "HEAD"], - check=True, - capture_output=True, - text=True, - ) - actual = result.stdout.strip() - assert actual == REFERENCE_REVISION, (f"Tongyi Z-Image oracle must be pinned at {REFERENCE_REVISION}, got {actual}") - return REFERENCE_REPO / "src" - - -def _import_reference_modules() -> tuple[object, object]: - source_root = _require_pinned_reference().resolve() - if str(source_root) not in sys.path: - sys.path.insert(0, str(source_root)) - pipeline_module = importlib.import_module("zimage.pipeline") - utils_module = importlib.import_module("utils") - for module in (pipeline_module, utils_module): - module_path = Path(module.__file__ or "").resolve() - assert module_path.is_relative_to(source_root), ( - f"Z-Image oracle import escaped the pinned repository: {module_path}") - return pipeline_module, utils_module - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Full Z-Image pipeline parity requires CUDA") -def test_zimage_pipeline_latents_match_pinned_native_repo(tmp_path: Path) -> None: - """Compare one complete native step; the oracle never imports Diffusers' pipeline.""" - if not (MODEL_DIR / "model_index.json").is_file(): - pytest.skip(f"Pinned Z-Image weights not found: {MODEL_DIR}") - pipeline_module, utils_module = _import_reference_modules() - prompt = "Young Chinese woman in red Hanfu, intricate embroidery." - - previous_fastvideo_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") - try: - # Exercise the shared PyTorch SDPA path with each process's native - # kernel selection; the official and FastVideo implementations both - # pass the same padding mask into scaled_dot_product_attention. - utils_module.set_attention_backend(None) - components = utils_module.load_from_local_dir( - MODEL_DIR, - device="cuda", - dtype=torch.bfloat16, - compile=False, - ) - reference = pipeline_module.generate( - prompt=prompt, - **components, - height=64, - width=64, - num_inference_steps=1, - guidance_scale=0.0, - negative_prompt="", - max_sequence_length=512, - generator=torch.Generator("cuda").manual_seed(42), - output_type="latent", - ).detach().float().cpu() - del components - gc.collect() - torch.cuda.empty_cache() - - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" - from fastvideo import VideoGenerator - - generator = VideoGenerator.from_pretrained( - str(MODEL_DIR), - num_gpus=1, - tp_size=1, - sp_size=1, - use_fsdp_inference=False, - dit_cpu_offload=False, - dit_layerwise_offload=False, - text_encoder_cpu_offload=False, - vae_cpu_offload=False, - pin_cpu_memory=False, - output_type="latent", - ) - try: - result = generator.generate_video( - prompt, - negative_prompt="", - output_path=str(tmp_path), - save_video=False, - return_frames=True, - height=64, - width=64, - num_frames=1, - fps=1, - num_inference_steps=1, - guidance_scale=0.0, - max_sequence_length=512, - seed=42, - ) - finally: - generator.shutdown() - actual = result["samples"].squeeze(2).float() - assert actual.shape == reference.shape - diff = (actual - reference).abs() - print(f"Z-Image pipeline latent parity: max={diff.max():.6f}, mean={diff.mean():.6f}") - assert_close(actual, reference, atol=1e-2, rtol=1e-2) - finally: - if previous_fastvideo_backend is None: - os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) - else: - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous_fastvideo_backend diff --git a/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py b/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py deleted file mode 100644 index 61f4ccc49..000000000 --- a/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py +++ /dev/null @@ -1,155 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Import/registry preflight and optional real Z-Image-Turbo smoke.""" - -from __future__ import annotations - -from contextlib import nullcontext -from pathlib import Path -from typing import Any, cast - -import pytest -import torch - -import fastvideo.envs as envs - -MODEL_DIR_ENV = envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() if envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.is_set() else None - - -def test_zimage_typed_surface_preflight() -> None: - from fastvideo.api import GeneratorConfig, PipelineSelection - from fastvideo.api.compat import generator_config_to_fastvideo_args - import fastvideo.registry as registry - from fastvideo.api.presets import get_preset, get_presets_for_family - from fastvideo.configs.models.dits.zimage import ZImageDiTConfig - from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig - from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig - from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig - from fastvideo.fastvideo_args import WorkloadType - from fastvideo.pipelines.basic.zimage.zimage_pipeline import EntryClass, ZImagePipeline - - assert ZImagePipeline.__name__ == "ZImagePipeline" - assert EntryClass is ZImagePipeline - assert ZImagePipeline.pipeline_config_cls is ZImagePipelineConfig - assert ZImagePipeline._required_config_modules == [ - "scheduler", - "text_encoder", - "tokenizer", - "transformer", - "vae", - ] - - default_preset, model_family = registry.get_preset_selection("Tongyi-MAI/Z-Image-Turbo") - assert (default_preset, model_family) == ("zimage_turbo", "zimage") - info = registry.get_model_info( - "Tongyi-MAI/Z-Image-Turbo", - workload_type=WorkloadType.T2I, - override_pipeline_cls_name="ZImagePipeline", - ) - assert info.pipeline_cls is ZImagePipeline - assert info.pipeline_config_cls is ZImagePipelineConfig - - assert {preset.name for preset in get_presets_for_family("zimage")} == {"zimage_turbo"} - preset = get_preset("zimage_turbo", "zimage") - assert preset.defaults == { - "height": 1024, - "width": 1024, - "num_frames": 1, - "fps": 1, - "seed": 42, - "guidance_scale": 0.0, - "num_inference_steps": 8, - "negative_prompt": "", - "max_sequence_length": 512, - "cfg_normalization": False, - "cfg_truncation": 1.0, - } - - config = ZImagePipelineConfig() - assert isinstance(config.dit_config, ZImageDiTConfig) - assert isinstance(config.vae_config, AutoencoderKLVAEConfig) - assert isinstance(config.text_encoder_configs[0], Qwen3TextConfig) - assert config.text_encoder_configs[0].chat_template_enable_thinking is True - assert config.text_encoder_configs[0].output_hidden_states is True - assert config.text_encoder_configs[0].tokenizer_kwargs["max_length"] == 512 - assert config.scheduler_sigma_min == 0.0 - assert config.scheduler_use_reference_discrete_timesteps is True - - args = generator_config_to_fastvideo_args( - GeneratorConfig( - model_path="Tongyi-MAI/Z-Image-Turbo", - pipeline=PipelineSelection(workload_type="t2i"), - )) - assert isinstance(args.pipeline_config, ZImagePipelineConfig) - assert args.workload_type is WorkloadType.T2I - - -def test_model_download_honors_pinned_revision(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - import fastvideo.utils as utils - - observed: dict[str, object] = {} - - def fake_snapshot_download(**kwargs): - observed.update(kwargs) - return str(tmp_path) - - monkeypatch.setattr(utils, "snapshot_download", fake_snapshot_download) - monkeypatch.setattr(utils, "get_lock", lambda _key: nullcontext()) - resolved = utils.maybe_download_model( - "Tongyi-MAI/Z-Image-Turbo", - revision="f332072aa78be7aecdf3ee76d5c247082da564a6", - ) - - assert resolved == str(tmp_path) - assert observed["repo_id"] == "Tongyi-MAI/Z-Image-Turbo" - assert observed["revision"] == "f332072aa78be7aecdf3ee76d5c247082da564a6" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Z-Image load/generate smoke requires CUDA") -def test_zimage_pipeline_load_generate_smoke() -> None: - if MODEL_DIR_ENV is None: - pytest.skip("Set FASTVIDEO_TEST_ZIMAGE_MODEL_DIR to activate the real load/generate smoke") - model_dir = Path(MODEL_DIR_ENV) - if not (model_dir / "model_index.json").is_file(): - pytest.skip(f"Z-Image model_index.json not found under {model_dir}") - - from fastvideo import VideoGenerator - - generator = VideoGenerator.from_pretrained( - str(model_dir), - num_gpus=1, - tp_size=1, - sp_size=1, - use_fsdp_inference=False, - dit_cpu_offload=False, - dit_layerwise_offload=False, - text_encoder_cpu_offload=False, - vae_cpu_offload=False, - pin_cpu_memory=False, - output_type="latent", - ) - try: - result = generator.generate_video( - prompt="a red panda reading a book", - negative_prompt="", - output_path="outputs/zimage/smoke", - save_video=False, - return_frames=True, - height=64, - width=64, - num_frames=1, - fps=1, - num_inference_steps=1, - guidance_scale=0.0, - max_sequence_length=64, - cfg_normalization=False, - cfg_truncation=1.0, - seed=42, - ) - finally: - generator.shutdown() - - result_dict = cast(dict[str, Any], result) - samples = result_dict["samples"] - assert torch.is_tensor(samples) - assert samples.ndim in (4, 5) - assert torch.isfinite(samples).all() diff --git a/tests/local_tests/zimage/PORT_STATUS.md b/tests/local_tests/zimage/PORT_STATUS.md deleted file mode 100644 index 1a3b7d085..000000000 --- a/tests/local_tests/zimage/PORT_STATUS.md +++ /dev/null @@ -1,162 +0,0 @@ -# Z-Image Port Status - -## Summary - -- model_family: `zimage` -- variant: `Z-Image-Turbo` -- workload_types: `T2I` -- PR scope: required components plus same-PR FastVideo-native pipeline, example/parity, SSIM test, and L40S seed workflow -- official_ref: `https://github.com/Tongyi-MAI/Z-Image` -- official_ref_dir: `/Z-Image/src` -- official_ref_commit: `26f23eda626ffadda020b04ff79488e1d72004cd` -- hf_weights_path: `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` -- local_weights_dir: `/official_weights/Z-Image/` -- source_layout: official repo-native `src/zimage/`; HF component subfolders provide assets only -- weight_layout: HF per-component subfolders; these are assets, not the reference implementation -- local_tests_readme: `tests/local_tests/zimage/README.md` - -The pinned text encoder has 36 decoder blocks, hidden size 2560, -intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension -128. FastVideo obtains those values from the pinned `config.json`. - -## Current Phase - -- phase: Phase 11 (post-parity review and handoff); quality reference published -- status: external_ci_gate -- owner: infra/reviewer -- last_updated: 2026-07-15 - -## Pinned Official Pipeline Context - -- official_pipeline: `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`, - `src/zimage/pipeline.py::generate` and `src/config/inference.py` -- oracle_policy: official Z-Image source only; Diffusers is not the implementation oracle -- workload_and_output: `T2I`; prompt plus optional negative prompt to PIL images, or latents when requested -- model_index_contract: `_class_name=ZImagePipeline`; required modules are `transformer`, `vae`, - `text_encoder`, `tokenizer`, and `scheduler` -- official_defaults: `height=1024`, `width=1024`, `num_inference_steps=8`, - `guidance_scale=0.0`, `cfg_truncation=1.0`, `max_sequence_length=512` -- reproducible_rng: `torch.Generator("cuda").manual_seed(42)`; the API accepts a generator and has no fixed seed default -- scheduler_contract: `use_reference_discrete_timesteps=True`, runtime `sigma_min=0.0`, endpoint-preserving - `num_steps + 1` schedule, and no DiT call for the final zero timestep -- decode_contract: apply `(latents / 0.3611) + 0.1159` before official VAE decode -- implementation_status: native pipeline/config/preset/registry/example and the 8-step 1024x1024 PNG SSIM test are - implemented; one-step native GPU parity is bit-exact; the L40S candidate generated at inference head `9a6dbcadc` - was author-approved and published at dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c` - -## Component Matrix - -| Component | Type | Reuse/Port | Official Definition / Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | -|---|---|---|---|---|---|---|---|---| -| Production text encoder (Qwen3 body) | text_encoder | reused production passthrough | Official `src/utils/loader.py` uses Transformers with typed load then device-only move; official pipeline consumes `hidden_states[-2]` | `TextEncoderLoader` -> `Qwen3ForCausalLM.from_pretrained_local` -> body-only Transformers `AutoModel` | DONE | not_needed for current HF subfolder | PINNED GB200 PASS; reference/production outputs exact | none | -| Native Qwen quality | implementation_subcomponent | shared native implementation | Same pinned text-encoder weights and official `hidden_states[-2]` boundary; not loaded by the pipeline | direct `fastvideo.models.encoders.qwen3.Qwen3ForCausalLM` construction | DONE | not_needed | PINNED GB200 PASS; fp32 `atol=2e-4`, `rtol=1e-4`; bf16 thresholds unchanged | none | -| Tokenizer | tokenizer | reused | Official loader's `AutoTokenizer` from `/tokenizer/`; official thinking template and length 512 | `TokenizerLoader` | DONE | not_needed | PINNED GB200 PASS (2 passed) | none | -| VAE | vae | reused with accepted wrapper exception | Pinned official `src/zimage/autoencoder.py`; raw decode is the component boundary | Existing Diffusers-backed `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` through production `VAELoader`, plus direct raw-decode coverage | DONE | not_needed | PINNED GB200 PASS (2 passed) | none | -| Scheduler | scheduler | reused with Z-Image extension | Pinned official `src/zimage/scheduler.py`; official pipeline sets `sigma_min=0.0` | `FlowMatchEulerDiscreteScheduler(use_reference_discrete_timesteps=True, sigma_min=0.0)` | DONE | not_needed | PINNED GB200 PASS (5 passed) | none | -| Transformer (`ZImageTransformer2DModel`) | dit | native port | Pinned official `src/zimage/transformer.py`; HF `transformer/` supplies weights only | `fastvideo.models.dits.zimage.ZImageTransformer2DModel` through `TransformerLoader` | DONE | not_needed; exact 521-key surface | PINNED GB200 PASS (5 passed; real-weight full forward exact) | none | -| Pipeline | pipeline | native same-PR implementation | Pinned official `src/zimage/pipeline.py::generate` | FastVideo `ZImagePipeline` plus config/preset/registry/example/parity | DONE | not_needed | RTX 5090 ONE-STEP LATENTS BIT-EXACT; L40S REFERENCE PUBLISHED | I005 (CI verification only) | - -## Conversion State - -- conversion_script: none for the reused components or native transformer -- converted_weights_dir: none -- source_layout: official repo-native Python (`Tongyi-MAI/Z-Image`) -- hf_weight_layout: per-component subfolders; not a Diffusers implementation oracle -- transformer_conversion: not_needed; official, HF index, and native names align exactly -- transformer_key_surface: 521 tensors; sorted-key SHA256 `3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354` -- strict_load_status: transformer real-weight strict load passed; native Qwen strictness regressions passed -- passthrough_components: body-only Qwen3 `AutoModel`, tokenizer assets -- direct_load_components: native Qwen3 parity target, VAE, scheduler -- native_qwen_exclusion: `lm_head.weight` only; the native encoder owns the body, not an LM head -- native_qwen_fused_contract: every Q/K/V and gate/up shard is required unless the exact fused destination is present -- native_qwen_auxiliary_contract: mapped quantization auxiliaries, including scale parameters, must load and may not be silently discarded -- retry_history: none - -## Parity Commands - -| Scope | Command | Last Result | Notes | -|---|---|---|---| -| Scheduler | `pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v -s` | 5 PASSED, no skips, GB200, 2026-07-15 | Enforces official-source SHA/import origin, fp32 scalar contract, exact zero endpoint, and positional/default regressions. | -| Tokenizer | `pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -s` | 2 PASSED as part of the pinned full GB200 run, 2026-07-15 | Uses `enable_thinking=True`, `max_length=512`, and fails if chat templating is unavailable. | -| VAE | `pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v -s` | 2 PASSED, no skips, GB200, 2026-07-15 | Direct raw decode plus production `VAELoader`; config values are checked, while latent scale/shift application is pipeline behavior. | -| Production text encoder | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | REFERENCE/PRODUCTION CHECKS EXACT, GB200, 2026-07-15 | The required pipeline path matches an independent official-loader materialization exactly. | -| Native Qwen fp32 quality | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | PASS, GB200, 2026-07-15 | Z-Image-only `atol=2e-4`, `rtol=1e-4`; production equality and shared Qwen math are unchanged. | -| Native Qwen bf16 quality | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16] -v -s` | PASS, GB200, 2026-07-15 | Existing mean/median distribution thresholds pass unchanged against an independent materialization of the official loader contract. | -| Native Qwen per-layer bf16 diagnostic | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | PASS, GB200, 2026-07-15 | Pinned architecture has 36 blocks and 37 hidden-state entries under the tested contract. | -| Transformer parity | `ZIMAGE_OFFICIAL_REF_DIR= pytest tests/local_tests/zimage/test_zimage_transformer_parity.py -v -s` | 5 PASSED, no skips, GB200, 2026-07-15 | Production 521-key/shape surface, tiny CPU parity, mask behavior, SP guard, real-weight strict load, and 6.15B full forward pass; real-weight output was exact (`max=0`, `mean=0`). | -| Complete component closure | `ZIMAGE_OFFICIAL_REF_DIR= ZIMAGE_TRANSFORMER_DIR= pytest tests/local_tests/zimage/ -v -s` | 34 PASSED, no skips, GB200, 2026-07-15, `f310aacb1` | Official source `26f23eda...`, HF snapshot `f332072a...`; production/native encoder, scheduler, tokenizer, transformer, and VAE all passed. | -| Pipeline typed surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS, hardware-free, 2026-07-14 | Covers exact `EntryClass`, config, preset, registry, model detection, and official defaults. | -| Native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS, hardware-free, 2026-07-14 | Covers conditioning trim, endpoint-preserving timestep setup, denoising sign/timestep math, and VAE scale/shift. | -| Full pipeline parity | `ZIMAGE_REFERENCE_REPO= FASTVIDEO_TEST_ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS, RTX 5090, `fec2baa6c` | Real FastVideo one-step latents match pinned native Tongyi exactly (`max=0`, `mean=0`). Later commits through `f310aacb1` change only API inventory/tests, not inference math. | -| PNG SSIM quality | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | L40S REFERENCE UPLOADED; END-TO-END SSIM CI PENDING | Dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`; remote SHA-256 `6c2be648...f7a6`; upload scoped to `Tongyi-MAI__Z-Image-Turbo/TORCH_SDPA`. | - -## Open Questions - -| ID | Question | Owner | Needed By Phase | Status | Resolution | -|---|---|---|---|---|---| -| Q001 | Which exact Z-Image source revision is the parity oracle? | prep | Phase 1 | resolved | Pinned `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`; scheduler/VAE tests fail on a wrong HEAD or import origin. | -| Q002 | Which immutable HF snapshot supplies component weights? | prep | Phase 1 | resolved | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6`. | -| Q003 | What prompt/tokenization/hidden-state contract does the official pipeline use? | pipeline | Phase 6 | resolved | `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)`, tokenization length 512, and valid-token `hidden_states[-2]`. | -| Q004 | Does the native transformer require conversion? | conversion | Phase 5 | resolved | No. Pinned HF index, official production meta model, and FastVideo production meta model expose the same 521 keys; official/native key shapes match exactly. | - -## Issues And Blockers - -| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | -|---|---|---|---|---|---|---|---|---| -| I001 | Phase 6 | native_qwen_quality | medium | The optional native fp32 `hidden_states[-2]` tail narrowly exceeded `atol=1e-4` after smooth cross-kernel accumulation. | Exact layer trace: embeddings exact, layer-0 max `3.815e-6`, first misses at layer 33, largest failing boundary difference `1.686811e-4`; final RMSNorm passes. The required production `AutoModel` path is exact. | parity | resolved | Author approved Z-Image-only `_FP32_ATOL=2e-4` with `_FP32_RTOL=1e-4`; complete pinned GB200 suite passed 34/34 with bf16 and all other gates unchanged. | -| I002 | Phase 7 | scheduler/pipeline | high | Published `scheduler_config.json` omits the reference-schedule flag and official zero endpoint. | The native stage applies both values; hardware-free stage coverage and pinned native GPU parity pass. | pipeline | resolved | The Z-Image pipeline applies `use_reference_discrete_timesteps=True` and `sigma_min=0.0`; one-step latents match official Tongyi exactly. | -| I003 | Phase 6 | tokenizer/text_encoder | medium | Prompt formatting and target length previously differed from the official call. | Pinned source uses thinking chat template, 512 tokens, and `hidden_states[-2]`. | parity | resolved | Corrected behavior passed the pinned GB200 run. | -| I004 | Phase 4 | transformer | high | `ZImageTransformer2DModel` was not ported. | Native config/model/registry and deterministic parity test now exist. | port | resolved | Production meta surface matches all 521 pinned HF keys; tiny full-forward parity passes exactly. | -| I005 | Phase 10 | pipeline/quality | high | Verify the published L40S PNG through end-to-end SSIM CI. | Author approved the reviewed candidate; the guarded upload published exactly one file at dataset revision `7b27d000...f01c`, and remote bytes match SHA-256 `6c2be648...f7a6`. | infra | external_ci_gate | Reference publication is resolved. The full-suite launcher remains blocked by the external-fork documentation workflow's `action_required` state; merge is separate. | -| I006 | Phase 6 | text_encoder | high | HF passthrough loading inherited an ambient target-device context, and the earlier test harness incorrectly cast float32 RoPE buffers during bf16 placement. | HF loading now occurs outside the ambient context; typed load is followed by device-only movement, matching official `src/utils/loader.py`; production/independent outputs match exactly. | encoder | resolved | Hardware-free placement/default-dtype tests and the pinned GB200 production path pass. | -| I007 | Phase 6 | text_encoder | high | Lenient native loading could miss one fused source shard or an auxiliary quantization scale. | Loader completeness covers every destination, required split shard, exact fused destination, and mapped scale parameter. | encoder | resolved | Strictness regressions and pinned checkpoint loading pass. | -| I008 | Phase 6 | vae | medium | Direct decode alone did not cover production registry/config/device/strict-load behavior; the shared target is a runtime Diffusers wrapper. | VAE parity now covers direct raw decode and production `VAELoader` against official `src/zimage/autoencoder.py`. | parity | resolved | Pinned GB200 subset passed 2/2. The wrapper is an implementation exception, not the oracle; official latent normalization remains Phase 7 pipeline scope. | -| I009 | Phase 2 | reference assets | high | Adding a nonexistent clone path to `sys.path` could import another installed `zimage` package and yield false parity. | Tests now verify exact clone HEAD and module paths under `Z-Image/src`; `/Z-Image/` is ignored. | prep | resolved | Wrong SHA/path/import is a failure; absent local assets may skip. | -| I010 | Phase 6 | tokenizer | medium | Missing `apply_chat_template` previously skipped and could hide an incompatible tokenizer. | Tokenizer parity now asserts the API and exact official kwargs. | parity | resolved | Missing chat-template behavior fails; pinned asset-backed parity passed. | -| I011 | Phase 6 | transformer | high | Production real-weight strict load and full-forward parity were not verified. | Pinned 24.6 GB weights and official source were exercised on GB200. | parity | resolved | Strict load and full forward passed exactly (`max=0`, `mean=0`). | -| I012 | Phase 6 | required_components | high | Record a complete non-skip run after the loader/oracle corrections and approved native-Qwen tolerance. | Exact PR head `f310aacb1`, pinned official source `26f23eda...`, immutable HF snapshot `f332072a...`, GB200. | parity | resolved | Complete suite passed 34/34 in 94.21s with no skips; real-weight transformer output was exact (`max=0`, `mean=0`). | - -## Escape Hatches - -| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | -|---|---|---|---|---|---|---| -| E001 | Phase 6 | cost | Approve one standard 8-hour 4×GB200 hold to run the complete pinned component and real-weight transformer parity suite at PR #1339? | Approve `model/r9`; reuse a suitable existing allocation if available. | resolved | Author approved. Existing job `1545932` was reused, so no new hold was created. Excludes Modal, SSIM seeding/upload, and merge. | -| E002 | Phase 6 | numerical | Accept the traced Z-Image-only native encoder fp32 tail without weakening shared Qwen or bf16 gates? | Approve `model/r10`: change only `_FP32_ATOL` from `1e-4` to `2e-4`, retain `_FP32_RTOL=1e-4`, then rerun native parity at the exact resulting head. | resolved | Author approved. The one-line test threshold changed at `f310aacb1`; complete pinned GB200 suite passed 34/34 with production equality, bf16, strict-load, and official-oracle checks unchanged. | - -## Decisions - -| Date | Decision | Rationale | Impact | -|---|---|---|---| -| 2026-07-15 | Publish the visually approved L40S PNG reference to `FastVideo/ssim-reference-videos`. | The author confirmed the staged image looks good and explicitly requested the HF upload. | Guarded no-force upload added one Z-Image PNG at dataset revision `7b27d000...f01c`; remote bytes match the reviewed artifact. | -| 2026-07-15 | Accept the traced Z-Image-only native encoder fp32 absolute tolerance of `2e-4` while retaining `rtol=1e-4`. | Exact tracing showed gradual fused/split-kernel accumulation rather than a semantic break; the largest failing absolute tail was `1.686811e-4`. | One test constant changed. The full pinned suite passed 34/34; production equality, bf16 thresholds, strict loading, and official Tongyi oracle are unchanged. | -| 2026-07-14 | Carry the native FastVideo pipeline, example/parity, SSIM test, and L40S seed workflow on PR #1339. | The author explicitly requested the pipeline and quality regression on the same PR rather than a follow-up. | Implementation, component closure, native GPU parity, and reference publication are complete. Later commits do not change inference math. | -| 2026-07-14 | Treat the exact production `AutoModel` loader as the required pipeline text encoder and direct native-Qwen parity as a separate quality row. | `TextEncoderLoader` calls `Qwen3ForCausalLM.from_pretrained_local`, which returns body-only Transformers `AutoModel`; the pipeline does not direct-construct FastVideo-native Qwen. | The production component is exact. The separately approved native-Qwen tolerance changed no production math. | -| 2026-07-13 | Use pinned `Tongyi-MAI/Z-Image@26f23eda...` as the implementation oracle; treat HF subfolders only as weight/config assets. | The official repo defines transformer, VAE, scheduler, pipeline, and loader semantics. A runtime component may inherit Diffusers without making Diffusers the reference. | Tests verify official clone SHA/import origin; docs distinguish the oracle from storage/runtime implementation details. | -| 2026-07-13 | Keep official latent scale/shift application out of the VAE component assertion. | `(latents / 0.3611) + 0.1159` is pipeline-stage behavior, while VAE component parity compares raw decode and production loading. | The Z-Image stage implements the exact formula; pinned native GPU parity passed. | -| 2026-07-13 | Preserve the published transformer key surface and use raw masked SDPA at SP=1. | The checkpoint already matches the official class, while padded variable-length attention needs a key mask unavailable in current distributed wrappers. | No converter is needed; runtime rejects SP>1 until a mask-aware distributed path has parity evidence. | -| 2026-07-04 | Pin both external sources and treat all earlier GPU results as historical until rerun. | Earlier evidence did not record an immutable HF revision, and source imports were not origin-checked. | README has executable clone/download commands; tests enforce the source SHA; all asset-backed parity rows were revalidated. | -| 2026-07-04 | Use a body-only Transformers `AutoModel` for production and a separate `AutoModel` object to materialize the official loader contract independently. | The pinned official Z-Image source is the oracle. Every FastVideo consumer needs hidden states, not full-vocabulary logits or an LM head, and comparing an object with itself is not independent evidence. | Production avoids the LM head while parity covers production loading and native implementation separately. | -| 2026-07-04 | Make fused-source and quantization-scale completeness part of native Qwen strict loading. | A loaded destination name alone did not prove that every required Q/K/V or gate/up shard and auxiliary scale arrived. | Missing split shards or mapped scale parameters fail; an exact fused destination remains valid. | -| 2026-07-04 | Reuse the existing Diffusers-backed AutoencoderKL wrapper for the Z-Image VAE component. | The target is established shared code with production-loader parity against pinned official `src/zimage/autoencoder.py`. | The wrapper remains an implementation detail for this VAE only; official Z-Image source defines pipeline behavior and numerical parity. | -| 2026-07-04 | Treat absent chat-template behavior as incompatibility, not an optional test environment. | Thinking-format prompt construction is part of the official numerical contract. | Tokenizer parity fails when `apply_chat_template` is unavailable. | -| 2026-07-04 | Record the pinned Qwen architecture as 36 blocks, hidden size 2560, intermediate size 9728, 32 attention heads, and 8 KV heads. | Earlier notes described a different architecture and could mis-size the native target. | Config/review evidence now matches `f332072aa78be7aecdf3ee76d5c247082da564a6`; expected hidden-state length is 37. | -| 2026-06-21 | Expand generated Qwen `position_ids` to `[batch_size, seq_len]`. | The rotary layer flattens positions; `[1, seq_len]` misaligned RoPE for batch sizes greater than one. | Batch-one behavior is unchanged; pinned batch-two parity passed. | -| 2026-06-21 | Reuse the shared native `Qwen3ForCausalLM` implementation rather than maintain a second Z-Image encoder. | The shared implementation supports the same config-driven GQA, fused QKV/gate-up, and hidden-state output contract. | Registry/native parity uses the shared class; architecture values come from the pinned config. | -| 2026-05-12 | Use mean and median drift checks for deep bf16 encoder parity. | Cross-kernel fused/unfused bf16 arithmetic produces a long max-error tail; fp32 and per-layer diagnostics distinguish this from structural divergence. | Existing thresholds passed unchanged on the pinned 36-block snapshot. | -| 2026-05-12 | Forward the complete scheduler config except loader metadata. | Hand-picking three fields could silently omit future scheduler behavior. | Scheduler parity consumes all on-disk constructor fields and explicitly supplies Z-Image runtime overrides. | - -## Handoff Notes - -- The implementation oracle is the pinned official `Tongyi-MAI/Z-Image` clone, - never Diffusers. Transformers is used for text components only because the - official loader does so; HF component subfolders provide immutable assets. -- The required production text encoder is the exact body-only `AutoModel` path. - `model/r10` is resolved: only the Z-Image test's fp32 absolute tolerance is - `2e-4`; relative tolerance, shared Qwen math, bf16, and production loading are unchanged. -- I012 is closed by the exact-head 34/34 non-skip GB200 run at `f310aacb1`. -- Native pipeline/config/preset/registry/example, parity tests, and the 8-step - 1024x1024 PNG SSIM test are implemented. Native RTX 5090 one-step parity is - bit-exact. The reviewed L40S PNG is live at dataset revision - `7b27d00004b723ef0cf4923e77ae57128573f01c` with matching remote bytes. -- End-to-end SSIM/full-suite execution remains an infra/reviewer gate; reference - publication does not authorize PR approval or merge. diff --git a/tests/local_tests/zimage/README.md b/tests/local_tests/zimage/README.md deleted file mode 100644 index e7c24ff3c..000000000 --- a/tests/local_tests/zimage/README.md +++ /dev/null @@ -1,200 +0,0 @@ -# Z-Image Local Tests - -Local component coverage and quality handoff for the `zimage` FastVideo port -(`T2I`, Z-Image-Turbo only). The same PR is authorized to carry the components, -FastVideo-native `ZImagePipeline`, example, pipeline parity, SSIM test, and L40S -reference seed. The native pipeline surface and 8-step 1024x1024 PNG SSIM test -are implemented. Native one-step GPU parity is bit-exact, and the L40S candidate -has been published to `FastVideo/ssim-reference-videos`. - -> **Status:** Implementation and quality-reference seeding complete; Phase 11 -> final handoff. The -> Phase 6 exact-head closure is complete: the full pinned component suite passed -> 34/34 with no skips on GB200 at `f310aacb1`. The required production text -> encoder is the official-loader-compatible Transformers `AutoModel` path, which -> matches an independently materialized reference exactly. Scheduler, tokenizer, -> VAE, native Qwen fp32/bf16, and the 6.15B real-weight transformer all pass. -> Native one-step pipeline parity is bit-exact on RTX 5090. The reviewed -> 1024x1024 L40S PNG generated at inference head `9a6dbcadc` is published at -> dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`. See -> [`PORT_STATUS.md`](./PORT_STATUS.md) for the live state. - -## Reference assets and scope - -| Field | Value | -|---|---| -| Model family / variant | `zimage` / `Z-Image-Turbo` | -| Workload | `T2I` | -| PR scope | required components, native pipeline, example/parity, SSIM test, and L40S seed | -| Required component scope | production `AutoModel` text encoder, tokenizer, VAE, scheduler, native transformer | -| Additional quality scope | direct FastVideo-native Qwen parity (`model/r10`; not used by the pipeline loader) | -| Official reference | `https://github.com/Tongyi-MAI/Z-Image` | -| Local reference dir | `/Z-Image/src` | -| Official commit | `26f23eda626ffadda020b04ff79488e1d72004cd` | -| HF weights | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` | -| Local weights dir | `/official_weights/Z-Image/` (component subfolders; transformer is 24.6 GB) | -| Reference implementation layout | Official repo-native `src/zimage/` Python; Diffusers is not the parity oracle | -| Weight layout | HF per-component subfolders (`text_encoder/`, `vae/`, `scheduler/`, `transformer/`) | -| Conversion | not needed; the native and official production transformer expose the same 521 checkpoint keys/shapes | -| HF token env | `HF_TOKEN` (the pinned repository is public; never record a token value) | - -The pinned text-encoder config declares `Qwen3ForCausalLM` with 36 decoder -blocks, hidden size 2560, intermediate size 9728, 32 attention heads, 8 KV -heads, and head dimension 128. FastVideo loads those values from -`text_encoder/config.json`; they are not hard-coded in the Z-Image integration. - -## Shared environment setup - -Run from the FastVideo repository root in the same environment used for -FastVideo. Keep the reference clone and weights inside the ignored paths shown -below. - -Clone and pin the official reference: - -```bash -python .agents/skills/add-model-01-prep/scripts/clone_reference_repo.py \ - https://github.com/Tongyi-MAI/Z-Image.git \ - Z-Image \ - --commit 26f23eda626ffadda020b04ff79488e1d72004cd \ - --update-gitignore -git -C Z-Image rev-parse HEAD -``` - -The final command must print -`26f23eda626ffadda020b04ff79488e1d72004cd`. The scheduler and VAE tests also -enforce that exact HEAD and verify that imported `zimage` modules resolve under -`/Z-Image/src`. A missing clone may skip local parity; a wrong SHA or -wrong import origin fails rather than silently testing another installation. - -Download the immutable HF snapshot: - -```bash -python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py \ - Tongyi-MAI/Z-Image-Turbo \ - official_weights/Z-Image \ - --revision f332072aa78be7aecdf3ee76d5c247082da564a6 \ - --allow-pattern 'text_encoder/*' \ - --allow-pattern 'tokenizer/*' \ - --allow-pattern 'vae/*' \ - --allow-pattern 'scheduler/*' -``` - -The real-weight transformer test additionally needs `transformer/*` from the -same pinned revision (24.6 GB). Do not download it without approving that cost; -point `ZIMAGE_TRANSFORMER_DIR` at an existing pinned transformer subfolder. - -Do not change core dependency versions (`torch`, `diffusers`, `transformers`, -`flash-attn`, `triton`, or CUDA packages) without explicit approval. - -```text -dependency_changes: none -official_env_status: official source pinned and imported; GB200 component validation run completed -private_dep_stubs: none -blocked_on: repo-wide full-suite/SSIM launch gate and PR review; model implementation work is complete -``` - -## Run the tests - -```bash -pytest tests/local_tests/zimage/ -v -s -``` - -| Component | Coverage scope and contract | Status | -|---|---|---| -| Scheduler ([`test_zimage_scheduler_parity.py`](./test_zimage_scheduler_parity.py)) | `implementation_subcomponent`; exact official `sigma_min=0.0` plus `use_reference_discrete_timesteps=True`; positional/default-path regressions remain asset-free | PINNED GB200 PASS (5 passed) | -| Tokenizer ([`test_zimage_tokenizer_parity.py`](./test_zimage_tokenizer_parity.py)) | `production_loader`; exact `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)` and `max_length=512`; absence of `apply_chat_template` is a failure, not a skip | PINNED GB200 PASS (2 passed) | -| VAE ([`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py)) | `both`; direct raw decode plus production `VAELoader`, config values, and raw decode | PINNED GB200 PASS (2 passed) | -| Production text encoder ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `production_loader`; `TextEncoderLoader` materializes a body-only Transformers `AutoModel` matching official `src/utils/loader.py`; required by the pipeline | PINNED GB200 PASS (reference/production outputs exact) | -| Native Qwen quality ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `implementation_subcomponent`; direct FastVideo-native construction, strict loading, and fp32/bf16 diagnostics; not loaded by the pipeline | PINNED GB200 PASS; fp32 `atol=2e-4`, `rtol=1e-4`; bf16 thresholds unchanged | -| Transformer ([`test_zimage_transformer_parity.py`](./test_zimage_transformer_parity.py)) | `both`; pinned production meta key/shape surface, additive-mask attention parity, deterministic tiny CPU full-forward parity, explicit `sp_size=1` guard, and real-weight production-loader/full-forward parity | PINNED GB200 PASS (5 passed; real-weight full forward exact) | - -## Pinned official pipeline context - -| Field | Contract | -|---|---| -| Oracle | `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`, `src/zimage/pipeline.py::generate` and `src/config/inference.py`; never Diffusers | -| Pipeline/modules | `_class_name=ZImagePipeline`; `transformer`, `vae`, `text_encoder`, `tokenizer`, `scheduler` | -| Inputs/outputs | text-to-image prompt with optional negative prompt; PIL images by default or latents when requested | -| Official defaults | `1024x1024`, `num_inference_steps=8`, `guidance_scale=0.0`, `cfg_truncation=1.0`, `max_sequence_length=512` | -| Reproducible RNG | `torch.Generator("cuda").manual_seed(42)`; `generate()` accepts the generator rather than fixing a seed internally | -| Scheduler | `use_reference_discrete_timesteps=True`, runtime `sigma_min=0.0`, and the endpoint-preserving `num_steps + 1` schedule with the final zero-timestep DiT call skipped | -| Decode | apply `(latents / 0.3611) + 0.1159` before official VAE decode | -| FastVideo status | Native pipeline/config/preset/registry/example and PNG SSIM test are implemented. One-step RTX 5090 parity is bit-exact; the reviewed L40S PNG generated at `9a6dbcadc` is published at dataset revision `7b27d000...f01c`. Later PR changes affect only tests/docs. | - -## Pipeline and quality verification - -| Scope | Command | Status | -|---|---|---| -| Hardware-free pipeline surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS | -| Hardware-free native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS | -| Pinned native full GPU pipeline parity | `ZIMAGE_REFERENCE_REPO= FASTVIDEO_TEST_ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS on RTX 5090; one-step latents bit-exact (`max=0`, `mean=0`) | -| 8-step 1024x1024 PNG SSIM | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | Reviewed L40S PNG uploaded; remote SHA-256 `6c2be648...f7a6`; end-to-end SSIM CI pending | - -## Component contracts - -### Text encoder and tokenizer - -- The production `TextEncoderLoader` path deliberately loads Transformers - `AutoModel`, even when the checkpoint advertises `Qwen3ForCausalLM`. FastVideo - consumes hidden states only, so the production path is body-only and does not - materialize or execute an LM head. -- The parity oracle is the pinned official Z-Image source and pipeline contract. - A separate Transformers `AutoModel` instance materializes the official - loader behavior independently; it is not the object returned by FastVideo's - production loader. The native `Qwen3ForCausalLM` implementation is compared - separately. -- CPU/MPS text-encoder offload remains on the requested device. The loader records - `_fastvideo_input_device`, and `TextEncodingStage` sends token tensors there. -- Native loading must account for every destination parameter and every required - fused source shard. Q/K/V and gate/up shards are all required unless the exact - destination is supplied already fused; quantized auxiliary parameters such as - `scale_weight` are not silently ignored. -- The pinned official prompt path enables thinking, tokenizes to 512, and consumes - `hidden_states[-2]` at valid-token positions. A 36-block Qwen model exposes 37 - hidden-state entries in this contract (embedding/intermediate entries plus the - final normalized state). -- `TextEncoderConfig.chat_template_enable_thinking` is keyword-only and defaults - to `False` for existing model families. The Z-Image pipeline config opts in - with `True`. - -### VAE - -The test covers both direct implementation parity and production `VAELoader` -resolution/strict loading, confirms the official config's scaling/shift values, -and compares raw decode outputs. The official pipeline's -`(latents / 0.3611) + 0.1159` transformation is pipeline behavior, so it remains -a Phase 7 Z-Image pipeline-stage and pipeline-parity requirement. This port -reuses the -existing shared `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` wrapper, -which subclasses Diffusers `AutoencoderKL`, as a narrowly scoped exception to -the native-component boundary. That runtime inheritance does not make Diffusers -the reference implementation: parity is against pinned -`Tongyi-MAI/Z-Image/src/zimage/autoencoder.py`. This is not precedent for the -transformer or pipeline implementation. - -### Scheduler - -The pinned official pipeline mutates `scheduler.sigma_min = 0.0` before building -its `num_steps + 1` schedule. The Z-Image pipeline applies both -`use_reference_discrete_timesteps=True` and `sigma_min=0.0`; the published -`scheduler_config.json` alone does not encode the complete runtime contract. - -### Transformer - -- The FastVideo-native transformer preserves all 521 published checkpoint keys. - The sorted key digest from the pinned HF safetensors index is - `3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354`. -- Tiny CPU parity uses the pinned official class, identical deterministic state, - variable-length image/text batches, and concrete full-forward output checks. -- Attention uses raw torch SDPA because the padded variable-length stream needs - a key mask not exposed by FastVideo's distributed wrappers. Runtime rejects - initialized sequence-parallel world sizes above one. -- The production-loader/full-forward test runs when the pinned transformer - subfolder and CUDA are available. The pinned GB200 run strict-loaded the - 6.15B model and matched the official full-forward output exactly. - -## Remaining work - -- Run the end-to-end SSIM/full suite after the external-fork documentation - approval gate permits the repo-wide launcher. -- Keep PR approval and merge separate from the completed reference upload. diff --git a/tests/local_tests/zimage/__init__.py b/tests/local_tests/zimage/__init__.py deleted file mode 100644 index 988131360..000000000 --- a/tests/local_tests/zimage/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/zimage/test_zimage_encoder_parity.py b/tests/local_tests/zimage/test_zimage_encoder_parity.py deleted file mode 100644 index caf04902f..000000000 --- a/tests/local_tests/zimage/test_zimage_encoder_parity.py +++ /dev/null @@ -1,716 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Parity test for Z-Image Qwen3 text encoder support in FastVideo. - -This compares: -1) direct transformers AutoModel output, and -2) FastVideo's production ``TextEncoderLoader`` passthrough, and -3) FastVideo's shared native Qwen3 encoder (``Qwen3ForCausalLM``, added for Flux2 - Klein and reused here — Z-Image-Turbo's ``Qwen3Model`` checkpoint routes - to it via the model registry), - -using identical local checkpoint, tokenization, and inputs. - -Usage: - pytest tests/local_tests/zimage/test_zimage_encoder_parity.py -v -s -""" - -from __future__ import annotations - -from pathlib import Path -from types import SimpleNamespace -import gc -import json -import os - -import pytest -import torch - - -def _reclaim_vram() -> None: - """Actually reclaim VRAM after the caller has ``del``'d its model refs. - - HF transformer models hold reference cycles, so ``del`` alone does not free - them — a ``gc.collect()`` is required before the CUDA caching allocator will - release the blocks. Without this the parity test keeps two full encoder - copies resident and OOMs on a 44 GB GPU. - """ - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - -def _gpu_mem_mb() -> float: - if torch.cuda.is_available(): - return torch.cuda.memory_allocated() / (1024**2) - return 0.0 - - -from torch.testing import assert_close -from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer -from safetensors.torch import safe_open - -from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig -from fastvideo.distributed.parallel_state import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel -from fastvideo.layers.quantization.absmax_fp8 import AbsMaxFP8MergedParameter -from fastvideo.models.encoders.qwen3 import Qwen3ForCausalLM -from fastvideo.models.loader.component_loader import TextEncoderLoader -from fastvideo.models.loader.utils import set_default_torch_dtype -from fastvideo.models.registry import ModelRegistry - -PARITY_SCOPE = "both" - -# Strict-load contract: the shared Qwen3 encoder is body-only (embed_tokens + -# layers + norm, no lm_head). Z-Image-Turbo ships a full Qwen3 checkpoint, so -# `lm_head.weight` is the only key that may go unmatched; anything else means a -# real silent drop. Enforced here in the test since the shared encoder's loader -# is intentionally lenient (it's used by multiple models). -_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"} - -REPO_ROOT = Path(__file__).resolve().parents[3] -ZIMAGE_TEXT_ENCODER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "text_encoder" -ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer" - - -@pytest.fixture(scope="module", autouse=True) -def _init_dist_and_tp_groups(): - os.environ.setdefault("MASTER_ADDR", "localhost") - os.environ.setdefault("MASTER_PORT", "29531") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - os.environ.setdefault("LOCAL_RANK", "0") - - maybe_init_distributed_environment_and_model_parallel(1, 1) - yield - cleanup_dist_env_and_memory() - - -def _load_json(path: Path) -> dict: - with path.open("r", encoding="utf-8") as f: - return json.load(f) - - -def _iter_safetensors(path: Path): - with safe_open(str(path), framework="pt", device="cpu") as f: - for k in f.keys(): - yield k, f.get_tensor(k) - - -def _iter_pretrained_safetensors(model_dir: Path): - single = model_dir / "model.safetensors" - if single.exists(): - yield from _iter_safetensors(single) - return - - index = model_dir / "model.safetensors.index.json" - if index.exists(): - idx = _load_json(index) - shard_names = sorted(set(idx["weight_map"].values())) - for shard in shard_names: - yield from _iter_safetensors(model_dir / shard) - return - - raise FileNotFoundError( - f"Missing safetensors checkpoint in {model_dir} (expected model.safetensors or model.safetensors.index.json)") - - -def _pretrained_safetensor_keys(model_dir: Path) -> set[str]: - return {name for name, _ in _iter_pretrained_safetensors(model_dir)} - - -def _load_qwen3_config() -> Qwen3TextConfig: - cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json") - for key in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"): - cfg_raw.pop(key, None) - config = Qwen3TextConfig() - config.update_model_arch(cfg_raw) - assert config.num_hidden_layers == 36 - assert config.hidden_size == 2560 - assert config.intermediate_size == 9728 - assert config.num_attention_heads == 32 - assert config.num_key_value_heads == 8 - return config - - -def _loader_args(cpu_offload: bool) -> SimpleNamespace: - return SimpleNamespace( - text_encoder_cpu_offload=cpu_offload, - override_text_encoder_quant=None, - override_text_encoder_safetensors=None, - pin_cpu_memory=False, - disable_offload_on_unified_memory=lambda device_id, offload_flag=None: False, - ) - - -def _precision_name(dtype: torch.dtype) -> str: - if dtype == torch.float32: - return "fp32" - if dtype == torch.bfloat16: - return "bf16" - raise ValueError(f"Unsupported parity dtype: {dtype}") - - -def _tokenize_official_prompts(tokenizer, prompts: list[str]): - formatted = [ - tokenizer.apply_chat_template( - [{ - "role": "user", - "content": prompt - }], - tokenize=False, - add_generation_prompt=True, - enable_thinking=True, - ) for prompt in prompts - ] - return tokenizer( - formatted, - padding="max_length", - truncation=True, - max_length=512, - return_tensors="pt", - ) - - -def _assert_native_load_surface( - model, - checkpoint_keys: set[str], - loaded_params: set[str], -) -> None: - param_names = {name for name, _ in model.named_parameters()} - missing = param_names - loaded_params - assert not missing, f"FastVideo parameters not loaded: {sorted(missing)}" - - normalized_keys = {name[len("model."):] if name.startswith("model.") else name for name in checkpoint_keys} - mapped_keys: set[str] = set() - for checkpoint_name in normalized_keys: - for param_name, weight_name, _ in model.config.arch_config.stacked_params_mapping: - if weight_name in checkpoint_name: - mapped_keys.add(checkpoint_name.replace(weight_name, param_name)) - break - else: - mapped_keys.add(checkpoint_name) - - unexpected = mapped_keys - param_names - { - "rotary_emb.inv_freq", - "rotary_emb.cos_cached", - "rotary_emb.sin_cached", - } - assert unexpected <= _ALLOWED_UNEXPECTED_KEYS, ("Unexpected checkpoint keys not in allowlist: " - f"{sorted(unexpected - _ALLOWED_UNEXPECTED_KEYS)}") - - -def _fake_qwen3_weight_target(include_quant_scales: bool = False): - stacked = Qwen3TextConfig().arch_config.stacked_params_mapping - params = { - "layers.0.self_attn.qkv_proj.weight": torch.nn.Parameter(torch.empty(1)), - "layers.0.mlp.gate_up_proj.weight": torch.nn.Parameter(torch.empty(1)), - } - if include_quant_scales: - qkv_scale = AbsMaxFP8MergedParameter(torch.zeros(3), requires_grad=False) - qkv_scale.output_partition_sizes = [1, 1, 1] - gate_up_scale = AbsMaxFP8MergedParameter(torch.zeros(2), requires_grad=False) - gate_up_scale.output_partition_sizes = [1, 1] - params.update({ - "layers.0.self_attn.qkv_proj.scale_weight": qkv_scale, - "layers.0.mlp.gate_up_proj.scale_weight": gate_up_scale, - }) - for param in params.values(): - if not hasattr(param, "weight_loader"): - param.weight_loader = (lambda target, loaded, *args: target.data.copy_(loaded)) - - model = SimpleNamespace( - config=SimpleNamespace(arch_config=SimpleNamespace(stacked_params_mapping=stacked)), - named_parameters=lambda: params.items(), - ) - return model, params, stacked - - -# The provisional bf16 thresholds below come from the historical unpinned -# 35-layer run. Per add-model-02-parity's calibration block, bf16 max-based -# asserts are meaningless for deep encoders — fused-vs-unfused linear order -# alone produces a long tail, so the test uses mean + median instead. The -# pinned snapshot has 36 blocks and must be rerun before these thresholds are -# treated as current evidence. -# -# Historical empirical numbers (unrecorded snapshot, 35 layers, CUDA): -# last_hidden_state (post-norm): mean=0.0152, median=0.0117 -# hidden_states[-2] (pre-norm): mean=0.0754, median=0.0625 -# Thresholds below add ~1.6x headroom over observed. -_BF16_MEAN_DRIFT_POST_NORM = 0.025 -_BF16_MEAN_DRIFT_PRE_NORM = 0.120 -_BF16_MEDIAN_DRIFT_POST_NORM = 0.020 -_BF16_MEDIAN_DRIFT_PRE_NORM = 0.100 -_FP32_ATOL = 2e-4 -_FP32_RTOL = 1e-4 - - -def _print_diag(label: str, ref: torch.Tensor, fv: torch.Tensor) -> torch.Tensor: - diff = (ref.float() - fv.float()).abs() - flat = diff.flatten() - p99 = flat.kthvalue(max(1, int(0.99 * flat.numel()))).values - print( - f"[{label}] max_diff={diff.max():.4f} mean_diff={diff.mean():.4f} " - f"median_diff={diff.median():.4f} p99_diff={p99:.4f}", - flush=True, - ) - return diff - - -@pytest.mark.parametrize( - ("cpu_offload", "target_device", "expected_device"), - [ - (True, torch.device("cuda"), torch.device("cpu")), - (False, torch.device("meta"), torch.device("meta")), - ], -) -def test_qwen3_production_loader_avoids_device_context_and_honors_placement( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - cpu_offload: bool, - target_device: torch.device, - expected_device: torch.device, -): - """HF loading stays ambient while final placement honors target/offload.""" - import fastvideo.platforms as platforms - - placements: list[torch.device] = [] - auto_model_calls: list[dict] = [] - - class FakeHFModel: - - def eval(self): - return self - - def to(self, device: torch.device): - placements.append(torch.device(device)) - return self - - fake_model = FakeHFModel() - - ambient_device = torch.empty(0).device - default_devices: list[torch.device] = [] - - def fake_auto_model_from_pretrained(*args, **kwargs): - auto_model_calls.append(kwargs) - default_devices.append(torch.empty(0).device) - return fake_model - - def fail_causal_lm_from_pretrained(*args, **kwargs): - raise AssertionError("Qwen text encoding must not load AutoModelForCausalLM") - - monkeypatch.setattr( - AutoModel, - "from_pretrained", - staticmethod(fake_auto_model_from_pretrained), - ) - monkeypatch.setattr( - AutoModelForCausalLM, - "from_pretrained", - staticmethod(fail_causal_lm_from_pretrained), - ) - monkeypatch.setattr( - platforms, - "_current_platform", - SimpleNamespace( - is_mps=lambda: False, - verify_model_arch=lambda _arch: None, - ), - ) - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr("fastvideo.models.loader.component_loader.get_local_torch_device", - lambda: torch.device("cuda")) - - config = Qwen3TextConfig() - # The pinned checkpoint names the causal wrapper even though the official - # Z-Image loader intentionally asks Transformers for the body-only model. - config.arch_config.architectures = ["Qwen3ForCausalLM"] - loaded = TextEncoderLoader().load_model( - str(tmp_path), - config, - target_device, - _loader_args(cpu_offload=cpu_offload), - dtype="fp32", - use_text_encoder_override=True, - ) - - assert loaded is fake_model - assert auto_model_calls == [{ - "local_files_only": True, - "dtype": torch.float32, - "low_cpu_mem_usage": True, - }] - assert default_devices == [ambient_device] - assert placements == [expected_device] - assert loaded._fastvideo_input_device == expected_device - - -def test_set_default_torch_dtype_restores_after_exception(): - original = torch.get_default_dtype() - - with pytest.raises(RuntimeError, match="expected failure"): - with set_default_torch_dtype(torch.bfloat16): - raise RuntimeError("expected failure") - - assert torch.get_default_dtype() == original - - -@pytest.mark.parametrize("checkpoint_prefix", ["", "model."]) -def test_qwen3_native_load_accepts_already_fused_weights_and_scales(checkpoint_prefix: str, ): - """Direct FastVideo/state-dict keys satisfy each fused target atomically.""" - fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True) - expected = {name: torch.arange(1, param.numel() + 1, dtype=param.dtype) for name, param in params.items()} - weights = [(f"{checkpoint_prefix}{name}", expected[name]) for name in params] - - loaded = Qwen3ForCausalLM.load_weights(fake_model, weights) - - assert loaded == set(params) - for name, param in params.items(): - assert torch.equal(param.data, expected[name]) - _assert_native_load_surface( - fake_model, - {name - for name, _ in weights}, - loaded, - ) - - -def test_qwen3_native_load_rejects_incompatible_fused_quant_scale_shape(): - fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True) - weights = [(name, torch.ones_like(param)) for name, param in params.items() - if name != "layers.0.self_attn.qkv_proj.scale_weight"] - weights.append(("layers.0.self_attn.qkv_proj.scale_weight", torch.ones(1))) - - with pytest.raises(AssertionError, match="Attempted to load weight"): - Qwen3ForCausalLM.load_weights(fake_model, weights) - - -def test_qwen3_native_load_uses_real_shard_loader_for_split_quant_scales(): - fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True) - weights = [(name, torch.ones_like(param)) for name, param in params.items() if name.endswith(".weight")] - weights.extend((f"model.layers.0.self_attn{weight_name}.scale_weight", torch.tensor(float(index + 1))) - for index, (_, weight_name, _) in enumerate(stacked[:3])) - weights.extend((f"model.layers.0.mlp{weight_name}.scale_weight", torch.tensor(float(index + 4))) - for index, (_, weight_name, _) in enumerate(stacked[3:])) - - loaded = Qwen3ForCausalLM.load_weights(fake_model, weights) - - assert loaded == set(params) - assert torch.equal( - params["layers.0.self_attn.qkv_proj.scale_weight"].data, - torch.tensor([1.0, 2.0, 3.0]), - ) - assert torch.equal( - params["layers.0.mlp.gate_up_proj.scale_weight"].data, - torch.tensor([4.0, 5.0]), - ) - - -@pytest.mark.parametrize( - ("missing_weight_name", "missing_shard_id"), - [ - (".q_proj", "q"), - (".k_proj", "k"), - (".v_proj", "v"), - (".gate_proj", 0), - (".up_proj", 1), - ], -) -def test_qwen3_native_load_rejects_missing_fused_weight_shard( - missing_weight_name: str, - missing_shard_id: str | int, -): - fake_model, _, stacked = _fake_qwen3_weight_target() - weights = [(f"model.layers.0.self_attn{weight_name}.weight", torch.empty(1)) for _, weight_name, _ in stacked[:3] - if weight_name != missing_weight_name] - weights.extend((f"model.layers.0.mlp{weight_name}.weight", torch.empty(1)) for _, weight_name, _ in stacked[3:] - if weight_name != missing_weight_name) - - with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info: - Qwen3ForCausalLM.load_weights(fake_model, weights) - - destination = ("layers.0.self_attn.qkv_proj.weight" - if isinstance(missing_shard_id, str) else "layers.0.mlp.gate_up_proj.weight") - assert f"{destination}[{missing_shard_id}]" in str(exc_info.value) - - -@pytest.mark.parametrize( - ("missing_weight_name", "missing_shard_id"), - [ - (".q_proj", "q"), - (".k_proj", "k"), - (".v_proj", "v"), - (".gate_proj", 0), - (".up_proj", 1), - ], -) -def test_qwen3_native_load_rejects_missing_fused_quant_scale_shard( - missing_weight_name: str, - missing_shard_id: str | int, -): - fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True) - weights = [(name, torch.empty(1)) for name in params if name.endswith(".weight")] - weights.extend((f"model.layers.0.self_attn{weight_name}.scale_weight", torch.empty(1)) - for _, weight_name, _ in stacked[:3] if weight_name != missing_weight_name) - weights.extend((f"model.layers.0.mlp{weight_name}.scale_weight", torch.empty(1)) - for _, weight_name, _ in stacked[3:] if weight_name != missing_weight_name) - - with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info: - Qwen3ForCausalLM.load_weights(fake_model, weights) - - destination = ("layers.0.self_attn.qkv_proj.scale_weight" - if isinstance(missing_shard_id, str) else "layers.0.mlp.gate_up_proj.scale_weight") - assert f"{destination}[{missing_shard_id}]" in str(exc_info.value) - - -@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required") -@pytest.mark.parametrize( - "dtype", - [ - pytest.param(torch.float32, id="fp32"), - pytest.param( - torch.bfloat16, - id="bf16", - marks=pytest.mark.skipif( - not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()), - reason="bf16 parity requires a bf16-capable CUDA device", - ), - ), - ], -) -def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype): - if not ZIMAGE_TOKENIZER_DIR.exists(): - pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - torch.manual_seed(11) - - tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True) - - prompts = [ - "A cinematic shot of a rainy neon street at night.", - "A watercolor illustration of a fox in autumn leaves.", - ] - toks = _tokenize_official_prompts(tokenizer, prompts) - input_ids = toks["input_ids"].to(device) - attention_mask = toks["attention_mask"].to(device) - config = _load_qwen3_config() - - ref = AutoModel.from_pretrained( - str(ZIMAGE_TEXT_ENCODER_DIR), - local_files_only=True, - trust_remote_code=True, - dtype=dtype, - low_cpu_mem_usage=True, - ).eval().to(device=device) - - with torch.no_grad(): - ref_out = ref( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - return_dict=True, - ) - ref_last = ref_out.last_hidden_state.detach().float().cpu() - assert ref_out.hidden_states is not None - ref_hs_m2 = ref_out.hidden_states[-2].detach().float().cpu() - - print(f"[mem] after HF ref forward: {_gpu_mem_mb():.0f} MiB", flush=True) - # Keep only CPU outputs so each full encoder is resident one at a time. - del ref, ref_out - _reclaim_vram() - print(f"[mem] after freeing HF ref: {_gpu_mem_mb():.0f} MiB", flush=True) - - production = TextEncoderLoader().load_model( - str(ZIMAGE_TEXT_ENCODER_DIR), - config, - device, - _loader_args(cpu_offload=False), - dtype=_precision_name(dtype), - use_text_encoder_override=True, - ) - assert production.__class__.__name__ == "Qwen3Model" - assert not hasattr(production, "lm_head") - assert production._fastvideo_input_device == device - - with torch.no_grad(): - production_out = production( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - return_dict=True, - ) - production_last = production_out.last_hidden_state.detach().float().cpu() - assert production_out.hidden_states is not None - production_hs_m2 = production_out.hidden_states[-2].detach().float().cpu() - - # Production loading must preserve the independent Transformers oracle - # exactly; native-kernel drift is assessed separately below. - assert_close(ref_last, production_last, atol=0.0, rtol=0.0) - assert_close(ref_hs_m2, production_hs_m2, atol=0.0, rtol=0.0) - del production, production_out, production_last, production_hs_m2 - _reclaim_vram() - - fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model") - fv = fv_cls(config).eval() - loaded = fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR)) - assert loaded, "No Qwen3 weights were loaded into FastVideo model" - fv = fv.to(device=device, dtype=dtype) - print(f"[mem] after FastVideo load: {_gpu_mem_mb():.0f} MiB", flush=True) - - _assert_native_load_surface( - fv, - _pretrained_safetensor_keys(ZIMAGE_TEXT_ENCODER_DIR), - loaded, - ) - - with torch.no_grad(): - fv_out = fv( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - fv_last = fv_out.last_hidden_state.detach().float().cpu() - assert fv_out.hidden_states is not None - fv_hs_m2 = fv_out.hidden_states[-2].detach().float().cpu() - - # Z-Image uses only valid token embeddings selected by attention mask. - # Compare parity on that exact path to avoid padded-token artifacts. - mask_cpu = attention_mask.detach().bool().cpu() - is_bf16 = dtype == torch.bfloat16 - - for i in range(mask_cpu.shape[0]): - valid = mask_cpu[i] - ref_last_v = ref_last[i][valid] - fv_last_v = fv_last[i][valid] - ref_hs_m2_v = ref_hs_m2[i][valid] - fv_hs_m2_v = fv_hs_m2[i][valid] - - last_diff = _print_diag(f"batch{i} last_hidden_state {dtype}", ref_last_v, fv_last_v) - hs_m2_diff = _print_diag(f"batch{i} hidden_states[-2] {dtype}", ref_hs_m2_v, fv_hs_m2_v) - - if is_bf16: - # Distribution checks instead of element-wise assert_close: cross- - # kernel bf16 produces a long max tail (fused vs unfused linears), - # but a real op bug pushes mean *and* median high. Mean catches - # systematic drift; median catches "wrong weights / swapped - # layers" where most elements diverge. Pre-norm tensors get a - # looser bound because the final RMSNorm compresses drift ~5x. - assert last_diff.mean().item() < _BF16_MEAN_DRIFT_POST_NORM, ( - f"batch{i} last_hidden_state bf16 mean drift " - f"{last_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_POST_NORM}") - assert last_diff.median().item() < _BF16_MEDIAN_DRIFT_POST_NORM, ( - f"batch{i} last_hidden_state bf16 median drift " - f"{last_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_POST_NORM}") - assert hs_m2_diff.mean().item() < _BF16_MEAN_DRIFT_PRE_NORM, ( - f"batch{i} hidden_states[-2] bf16 mean drift " - f"{hs_m2_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_PRE_NORM}") - assert hs_m2_diff.median().item() < _BF16_MEDIAN_DRIFT_PRE_NORM, ( - f"batch{i} hidden_states[-2] bf16 median drift " - f"{hs_m2_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_PRE_NORM}") - else: - assert_close(ref_last_v, fv_last_v, atol=_FP32_ATOL, rtol=_FP32_RTOL) - assert_close(ref_hs_m2_v, fv_hs_m2_v, atol=_FP32_ATOL, rtol=_FP32_RTOL) - - -@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required") -@pytest.mark.skipif( - not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()), - reason="bf16 per-layer diagnostic requires a bf16-capable CUDA device", -) -def test_zimage_qwen3_encoder_per_layer_bf16_diagnostic(): - """Diagnostic-only: prints per-layer FastVideo-vs-HF drift for the bf16 - forward to distinguish monotonic accumulation (bf16 tail) from a - single-layer spike (real op mismatch). - - This test always passes; review the captured stdout to interpret. - Expected pattern for healthy bf16 accumulation: max/mean/median grow - smoothly with layer depth. Suspect pattern: one layer's diff is - >>2x the previous, suggesting a real bug at that block. - """ - if not ZIMAGE_TOKENIZER_DIR.exists(): - pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(11) - - tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True) - toks = _tokenize_official_prompts( - tokenizer, - ["A cinematic shot of a rainy neon street at night."], - ) - input_ids = toks["input_ids"].to(device) - attention_mask = toks["attention_mask"].to(device) - - ref = AutoModel.from_pretrained( - str(ZIMAGE_TEXT_ENCODER_DIR), - local_files_only=True, - trust_remote_code=True, - dtype=dtype, - low_cpu_mem_usage=True, - ).eval().to(device=device) - - with torch.no_grad(): - ref_out = ref( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - return_dict=True, - ) - # Cache the reference hidden states on CPU, then free the HF model so - # only one encoder is resident when FastVideo runs (avoids OOM). - ref_hidden_cpu = [h.detach().float().cpu() for h in ref_out.hidden_states] - del ref, ref_out - _reclaim_vram() - - fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model") - cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json") - for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"): - cfg_raw.pop(k, None) - cfg = Qwen3TextConfig() - cfg.update_model_arch(cfg_raw) - fv = fv_cls(cfg).eval() - fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR)) - fv = fv.to(device=device, dtype=dtype) - - with torch.no_grad(): - fv_out = fv( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - fv_hidden_cpu = [h.detach().float().cpu() for h in fv_out.hidden_states] - - assert fv_out.hidden_states is not None - assert len(ref_hidden_cpu) == 37, (f"expected embeddings plus 36 block outputs, got {len(ref_hidden_cpu)} entries") - assert len(ref_hidden_cpu) == len(fv_hidden_cpu), ( - f"len mismatch: HF={len(ref_hidden_cpu)} FV={len(fv_hidden_cpu)}") - - mask_cpu = attention_mask.detach().bool().cpu()[0] - print( - f"\n[per-layer bf16 diagnostic] num_hidden_states={len(ref_hidden_cpu)} " - f"valid_tokens={mask_cpu.sum().item()}/{mask_cpu.numel()}", - flush=True, - ) - print( - f"{'idx':>3} {'kind':<10} {'max':>8} {'mean':>8} {'median':>8} {'p99':>8}", - flush=True, - ) - for idx, (ref_h, fv_h) in enumerate(zip(ref_hidden_cpu, fv_hidden_cpu)): - ref_v = ref_h[0][mask_cpu] - fv_v = fv_h[0][mask_cpu] - diff = (ref_v - fv_v).abs() - flat = diff.flatten() - p99 = torch.quantile(flat, 0.99).item() - if idx == 0: - kind = "embedding" - elif idx == len(ref_hidden_cpu) - 1: - kind = "post-norm" - else: - kind = f"layer{idx - 1}-out" - print( - f"{idx:>3} {kind:<10} " - f"{diff.max().item():>8.4f} {diff.mean().item():>8.4f} " - f"{diff.median().item():>8.4f} {p99:>8.4f}", - flush=True, - ) diff --git a/tests/local_tests/zimage/test_zimage_scheduler_parity.py b/tests/local_tests/zimage/test_zimage_scheduler_parity.py deleted file mode 100644 index ef4807929..000000000 --- a/tests/local_tests/zimage/test_zimage_scheduler_parity.py +++ /dev/null @@ -1,224 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Parity tests for Z-Image FlowMatchEulerDiscreteScheduler. - -These tests compare FastVideo's scheduler implementation against the -reference scheduler shipped in the local Z-Image repository. - -Usage: - pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v -""" - -from __future__ import annotations - -import importlib -import json -import subprocess -import sys -from pathlib import Path - -import numpy as np -import pytest -import torch -from torch.testing import assert_close - -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler as FastVideoScheduler, ) - -REPO_ROOT = Path(__file__).resolve().parents[3] -ZIMAGE_REPO = REPO_ROOT / "Z-Image" -ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src" -ZIMAGE_SCHEDULER_CFG = (REPO_ROOT / "official_weights" / "Z-Image" / "scheduler" / "scheduler_config.json") -ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" -PARITY_SCOPE = "implementation_subcomponent" - - -def _require_pinned_reference_module(module_name: str, source_file: Path): - if not ZIMAGE_REPO.exists(): - pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}") - if not source_file.is_file(): - pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}") - - try: - result = subprocess.run( - ["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"], - check=True, - capture_output=True, - text=True, - ) - except (FileNotFoundError, subprocess.CalledProcessError) as exc: - pytest.fail(f"Cannot verify Z-Image reference revision: {exc}") - - actual_revision = result.stdout.strip() - assert actual_revision == ZIMAGE_REFERENCE_REVISION, ( - "Z-Image reference clone is not at the pinned revision: " - f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}") - - if str(ZIMAGE_SRC) not in sys.path: - sys.path.insert(0, str(ZIMAGE_SRC)) - try: - module = importlib.import_module(module_name) - except Exception as exc: - pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}") - - module_file = Path(module.__file__ or "").resolve() - assert module_file.is_relative_to( - ZIMAGE_SRC.resolve()), (f"{module_name} resolved outside the pinned clone: {module_file}") - return module - - -@pytest.fixture(scope="module") -def reference_scheduler_cls(): - module = _require_pinned_reference_module( - "zimage.scheduler", - ZIMAGE_SRC / "zimage" / "scheduler.py", - ) - return module.FlowMatchEulerDiscreteScheduler - - -@pytest.fixture(scope="module") -def scheduler_config() -> dict: - if not ZIMAGE_SCHEDULER_CFG.exists(): - pytest.skip(f"Z-Image scheduler config not found: {ZIMAGE_SCHEDULER_CFG}") - with ZIMAGE_SCHEDULER_CFG.open("r", encoding="utf-8") as f: - return json.load(f) - - -# Diffusers-style scheduler-config keys we do NOT want to forward — they are -# loader/state metadata, not constructor kwargs. Anything else in -# scheduler_config.json (including future keys like `time_shift_type`, -# `invert_sigmas`, etc.) must be forwarded so parity actually exercises the -# config-on-disk and not a hand-curated subset. -_SCHEDULER_CONFIG_LOADER_KEYS = frozenset({"_class_name", "_diffusers_version", "_name_or_path"}) - - -def _scheduler_kwargs_from_config(scheduler_config: dict) -> dict: - return {k: v for k, v in scheduler_config.items() if k not in _SCHEDULER_CONFIG_LOADER_KEYS} - - -def _make_scheduler_pair(reference_scheduler_cls, scheduler_config: dict, **overrides): - kwargs = _scheduler_kwargs_from_config(scheduler_config) - kwargs.update(overrides) - - ref = reference_scheduler_cls(**kwargs) - # The pinned Z-Image pipeline mutates this immediately before scheduling. - ref.sigma_min = 0.0 - - # These values must be serialized by the future production pipeline config. - # setdefault also allows this test to consume that config once it lands. - fv_kwargs = dict(kwargs) - fv_kwargs.setdefault("use_reference_discrete_timesteps", True) - fv_kwargs.setdefault("sigma_min", 0.0) - fv = FastVideoScheduler(**fv_kwargs) - return ref, fv - - -@pytest.fixture(scope="module") -def scheduler_pair(reference_scheduler_cls, scheduler_config: dict): - return _make_scheduler_pair(reference_scheduler_cls, scheduler_config) - - -def _run_step_loop(ref_scheduler, fv_scheduler, sample_shape=(2, 4, 32, 32)): - torch.manual_seed(123) - - ref_sample = torch.randn(sample_shape, dtype=torch.float32) - fv_sample = ref_sample.clone() - - for t in ref_scheduler.timesteps: - model_output = torch.randn_like(ref_sample) - - ref_next = ref_scheduler.step( - model_output, - t, - ref_sample, - return_dict=False, - )[0] - fv_next = fv_scheduler.step( - model_output, - t, - fv_sample, - return_dict=False, - )[0] - - assert_close(ref_next, fv_next, atol=1e-6, rtol=1e-6) - - ref_sample = ref_next - fv_sample = fv_next - - -def test_zimage_scheduler_parity_default_schedule_and_step(scheduler_pair): - ref, fv = scheduler_pair - - num_inference_steps = 8 - ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu") - fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu") - - assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6) - assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6) - - _run_step_loop(ref, fv) - - -def test_scheduler_positional_args_keep_existing_bindings_and_default_schedule(): - scheduler = FastVideoScheduler(1000, 1.0, False, 0.7) - - assert scheduler.config.base_shift == 0.7 - assert scheduler.config.use_reference_discrete_timesteps is False - - scheduler.set_timesteps(num_inference_steps=8, device="cpu") - assert_close( - scheduler.sigmas[-2], - torch.tensor(scheduler.sigma_min, dtype=torch.float32), - ) - - -def test_scheduler_honors_explicit_zero_sigma_min(): - scheduler = FastVideoScheduler( - sigma_min=0.0, - use_reference_discrete_timesteps=True, - ) - - assert scheduler.sigma_min == 0.0 - scheduler.set_timesteps(num_inference_steps=8, device="cpu") - assert_close( - scheduler.timesteps[-1], - torch.tensor(125.0, dtype=torch.float32), - ) - - -def test_reference_schedule_preserves_float64_linspace_rounding(): - scheduler = FastVideoScheduler( - sigma_min=0.0, - use_reference_discrete_timesteps=True, - ) - scheduler.set_timesteps(num_inference_steps=9, device="cpu") - - expected = torch.from_numpy(np.linspace(1000.0, 0.0, 10)[:-1] / 1000.0, ).to(torch.float32) * 1000.0 - float32_regression = torch.from_numpy( - np.linspace(1000.0, 0.0, 10, dtype=np.float32)[:-1] / np.float32(1000.0), ) * 1000.0 - - # The fifth timestep differs by one float32 ULP depending on where the - # linspace is rounded, so this assertion fails if dtype=np.float32 returns. - assert expected[4].item() != float32_regression[4].item() - assert scheduler.timesteps[4].item() == expected[4].item() - - -def test_zimage_scheduler_parity_dynamic_shifting_with_mu( - reference_scheduler_cls, - scheduler_config: dict, -): - ref, fv = _make_scheduler_pair( - reference_scheduler_cls, - scheduler_config, - use_dynamic_shifting=True, - ) - - mu = 0.75 - num_inference_steps = 8 - ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu) - fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu) - - assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6) - assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6) - - _run_step_loop(ref, fv) diff --git a/tests/local_tests/zimage/test_zimage_tokenizer_parity.py b/tests/local_tests/zimage/test_zimage_tokenizer_parity.py deleted file mode 100644 index 32e9160bb..000000000 --- a/tests/local_tests/zimage/test_zimage_tokenizer_parity.py +++ /dev/null @@ -1,127 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Parity test for Z-Image tokenizer loading and tokenization behavior. - -This compares: -1) direct transformers AutoTokenizer loading from local tokenizer dir, and -2) FastVideo TokenizerLoader loading path, - -using the exact chat-template and tokenization settings from the pinned -Z-Image pipeline. - -Usage: - pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path - -import pytest -import torch -from torch.testing import assert_close -from transformers import AutoTokenizer - -from fastvideo.models.loader.component_loader import TokenizerLoader - -REPO_ROOT = Path(__file__).resolve().parents[3] -ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer" -PARITY_SCOPE = "production_loader" -OFFICIAL_MAX_SEQUENCE_LENGTH = 512 -OFFICIAL_CHAT_TEMPLATE_KWARGS = { - "tokenize": False, - "add_generation_prompt": True, - "enable_thinking": True, -} - - -@dataclass -class _DummyPipelineConfig: - text_encoder_configs: tuple = field(default_factory=tuple) - - -@dataclass -class _DummyFastVideoArgs: - pipeline_config: _DummyPipelineConfig = field(default_factory=_DummyPipelineConfig) - - -def _load_reference_tokenizer(): - if not ZIMAGE_TOKENIZER_DIR.exists(): - pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") - return AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True) - - -def _load_fastvideo_tokenizer(): - if not ZIMAGE_TOKENIZER_DIR.exists(): - pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") - loader = TokenizerLoader() - return loader.load(str(ZIMAGE_TOKENIZER_DIR), _DummyFastVideoArgs()) - - -def test_zimage_tokenizer_loader_and_tokenization_parity(): - ref_tok = _load_reference_tokenizer() - fv_tok = _load_fastvideo_tokenizer() - - prompts = [ - "A close-up portrait of a cat in warm studio lighting.", - "An oil painting of a lighthouse at sunset.", - ] - - tok_kwargs = { - "padding": "max_length", - "max_length": OFFICIAL_MAX_SEQUENCE_LENGTH, - "truncation": True, - "return_tensors": "pt", - } - - ref_out = ref_tok(prompts, **tok_kwargs) - fv_out = fv_tok(prompts, **tok_kwargs) - - assert_close(ref_out["input_ids"], fv_out["input_ids"], atol=0, rtol=0) - assert_close(ref_out["attention_mask"], fv_out["attention_mask"], atol=0, rtol=0) - - # Compare key tokenizer attributes that affect generation-time behavior. - assert ref_tok.padding_side == fv_tok.padding_side - assert ref_tok.pad_token_id == fv_tok.pad_token_id - assert ref_tok.eos_token_id == fv_tok.eos_token_id - - -@pytest.mark.skipif(not ZIMAGE_TOKENIZER_DIR.exists(), reason="Z-Image tokenizer assets are required") -def test_zimage_tokenizer_chat_template_parity(): - ref_tok = _load_reference_tokenizer() - fv_tok = _load_fastvideo_tokenizer() - - assert callable(getattr(ref_tok, "apply_chat_template", - None)), ("Pinned Z-Image tokenizer must expose apply_chat_template") - assert callable(getattr(fv_tok, "apply_chat_template", - None)), ("FastVideo TokenizerLoader dropped the required apply_chat_template API") - - messages = [{"role": "user", "content": "Describe a futuristic city skyline."}] - - ref_text = ref_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS) - fv_text = fv_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS) - - assert ref_text == fv_text - - # Z-Image deliberately enables Qwen's thinking template. Guard against a - # pipeline silently using the generic non-thinking prompt path. - non_thinking_text = fv_tok.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True, - enable_thinking=False, - ) - assert fv_text != non_thinking_text - - tokenization_kwargs = { - "padding": "max_length", - "max_length": OFFICIAL_MAX_SEQUENCE_LENGTH, - "truncation": True, - "return_tensors": "pt", - } - ref_tokens = ref_tok(ref_text, **tokenization_kwargs) - fv_tokens = fv_tok(fv_text, **tokenization_kwargs) - - assert_close(ref_tokens["input_ids"], fv_tokens["input_ids"], atol=0, rtol=0) - assert_close(ref_tokens["attention_mask"], fv_tokens["attention_mask"], atol=0, rtol=0) diff --git a/tests/local_tests/zimage/test_zimage_transformer_parity.py b/tests/local_tests/zimage/test_zimage_transformer_parity.py deleted file mode 100644 index d67537394..000000000 --- a/tests/local_tests/zimage/test_zimage_transformer_parity.py +++ /dev/null @@ -1,343 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Z-Image transformer parity. - -Coverage scope: both. The tiny CPU tests compare the pinned implementation and -FastVideo model directly. The production-loader test activates when the real -transformer subfolder is available. -""" - -from __future__ import annotations - -import hashlib -import importlib.util -import json -import os -from pathlib import Path -import subprocess -import sys -from types import ModuleType - -import pytest -import torch -from torch.testing import assert_close - -from fastvideo.configs.models.dits import ZImageDiTConfig -from fastvideo.models.dits.zimage import ZImageAttention, ZImageTransformer2DModel, _prepare_attention_mask - -REPO_ROOT = Path(__file__).resolve().parents[3] -OFFICIAL_REF_DIR = Path(os.getenv("ZIMAGE_OFFICIAL_REF_DIR", REPO_ROOT / "Z-Image")) -OFFICIAL_SRC = OFFICIAL_REF_DIR / "src" -TRANSFORMER_DIR = Path(os.getenv("ZIMAGE_TRANSFORMER_DIR", REPO_ROOT / "official_weights" / "Z-Image" / "transformer")) -REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" -HF_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6" -HF_KEY_COUNT = 521 -HF_KEY_SHA256 = "3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354" -PARITY_SCOPE = "both" - -PRODUCTION_CONFIG = { - "all_patch_size": (2, ), - "all_f_patch_size": (1, ), - "in_channels": 16, - "dim": 3840, - "n_layers": 30, - "n_refiner_layers": 2, - "n_heads": 30, - "n_kv_heads": 30, - "norm_eps": 1e-5, - "qk_norm": True, - "cap_feat_dim": 2560, - "rope_theta": 256.0, - "t_scale": 1000.0, - "axes_dims": (32, 48, 48), - "axes_lens": (1536, 512, 512), -} - -TINY_CONFIG = { - "all_patch_size": (2, ), - "all_f_patch_size": (1, ), - "in_channels": 2, - "dim": 12, - "n_layers": 1, - "n_refiner_layers": 1, - "n_heads": 2, - "n_kv_heads": 2, - "norm_eps": 1e-5, - "qk_norm": True, - "cap_feat_dim": 8, - "rope_theta": 256.0, - "t_scale": 1000.0, - "axes_dims": (2, 2, 2), - "axes_lens": (192, 16, 16), -} - - -def _load_source_module(name: str, path: Path): - spec = importlib.util.spec_from_file_location(name, path) - if spec is None or spec.loader is None: - pytest.fail(f"Cannot build import spec for {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture(scope="module") -def reference_transformer_cls(): - source_file = OFFICIAL_SRC / "zimage" / "transformer.py" - if not OFFICIAL_REF_DIR.exists(): - pytest.skip(f"Pinned Z-Image reference clone not found: {OFFICIAL_REF_DIR}") - if not source_file.is_file(): - pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}") - - try: - revision = subprocess.run( - ["git", "-C", str(OFFICIAL_REF_DIR), "rev-parse", "HEAD"], - check=True, - capture_output=True, - text=True, - ).stdout.strip() - except (FileNotFoundError, subprocess.CalledProcessError) as exc: - pytest.fail(f"Cannot verify Z-Image reference revision: {exc}") - assert revision == REFERENCE_REVISION, (f"expected Z-Image reference {REFERENCE_REVISION}, got {revision}") - - module_names = ( - "config", - "config.inference", - "config.model", - "utils", - "utils.import_utils", - "utils.attention", - "_zimage_reference_transformer", - ) - saved_modules = {name: sys.modules.get(name) for name in module_names} - added_path = str(OFFICIAL_SRC) not in sys.path - if added_path: - sys.path.insert(0, str(OFFICIAL_SRC)) - - try: - for name in module_names: - sys.modules.pop(name, None) - utils_package = ModuleType("utils") - utils_package.__path__ = [str(OFFICIAL_SRC / "utils")] - sys.modules["utils"] = utils_package - _load_source_module("utils.import_utils", OFFICIAL_SRC / "utils" / "import_utils.py") - _load_source_module("utils.attention", OFFICIAL_SRC / "utils" / "attention.py") - module = _load_source_module("_zimage_reference_transformer", source_file) - assert Path(module.__file__).resolve() == source_file.resolve() - yield module.ZImageTransformer2DModel - finally: - for name in module_names: - sys.modules.pop(name, None) - if saved_modules[name] is not None: - sys.modules[name] = saved_modules[name] - if added_path: - sys.path.remove(str(OFFICIAL_SRC)) - - -def _build_fastvideo(config_values: dict) -> ZImageTransformer2DModel: - config = ZImageDiTConfig() - config.update_model_arch(config_values) - return ZImageTransformer2DModel( - config=config, - hf_config={ - "_class_name": "ZImageTransformer2DModel", - **config_values - }, - ) - - -def _state_shapes(model: torch.nn.Module) -> dict[str, tuple[int, ...]]: - return {name: tuple(tensor.shape) for name, tensor in model.state_dict().items()} - - -def test_zimage_transformer_production_meta_key_surface(reference_transformer_cls): - from fastvideo.models.registry import ModelRegistry - - model_cls, _ = ModelRegistry.resolve_model_cls("ZImageTransformer2DModel") - assert model_cls is ZImageTransformer2DModel - - with torch.device("meta"): - reference = reference_transformer_cls(**PRODUCTION_CONFIG) - fastvideo = _build_fastvideo(PRODUCTION_CONFIG) - - reference_shapes = _state_shapes(reference) - fastvideo_shapes = _state_shapes(fastvideo) - assert fastvideo_shapes == reference_shapes - keys = sorted(reference_shapes) - key_digest = hashlib.sha256(("\n".join(keys) + "\n").encode()).hexdigest() - assert len(keys) == HF_KEY_COUNT - assert key_digest == HF_KEY_SHA256 - print(f"Z-Image production transformer key/shape surface: {len(reference_shapes)} tensors") - - -def test_zimage_transformer_tiny_cpu_parity(reference_transformer_cls): - reference = reference_transformer_cls(**TINY_CONFIG).eval() - generator = torch.Generator(device="cpu").manual_seed(123) - with torch.no_grad(): - for parameter in reference.parameters(): - parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) - - fastvideo = _build_fastvideo(TINY_CONFIG).eval() - assert _state_shapes(fastvideo) == _state_shapes(reference) - fastvideo.load_state_dict(reference.state_dict(), strict=True) - - generator.manual_seed(456) - images = [ - torch.randn((2, 1, 4, 4), generator=generator), - torch.randn((2, 1, 8, 10), generator=generator), - ] - caption_features = [ - torch.randn((3, 8), generator=generator), - torch.randn((35, 8), generator=generator), - ] - timestep = torch.tensor([0.25, 0.75], dtype=torch.float32) - - with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): - reference_outputs = reference(images, timestep, caption_features)[0] - positional_outputs = fastvideo(images, caption_features, timestep)[0] - keyword_outputs = fastvideo( - hidden_states=images, - encoder_hidden_states=caption_features, - timestep=timestep, - )[0] - - assert len(positional_outputs) == len(keyword_outputs) == len(reference_outputs) == 2 - for positional, actual, expected in zip(positional_outputs, keyword_outputs, reference_outputs): - assert actual.shape == expected.shape - assert expected.abs().max() > 0 - assert_close(positional, expected, atol=1e-6, rtol=1e-6) - assert_close(actual, expected, atol=1e-6, rtol=1e-6) - - -def test_zimage_attention_additive_mask_parity(reference_transformer_cls): - reference_module = sys.modules[reference_transformer_cls.__module__] - reference = reference_module.ZImageAttention(12, 2, 2).eval() - generator = torch.Generator(device="cpu").manual_seed(321) - with torch.no_grad(): - for parameter in reference.parameters(): - parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) - - fastvideo = ZImageAttention(12, 2, 2).eval() - fastvideo.load_state_dict(reference.state_dict(), strict=True) - hidden_states = torch.randn((2, 5, 12), generator=generator) - attention_mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, True]]) - additive_mask = _prepare_attention_mask(attention_mask, hidden_states.dtype) - assert additive_mask is not None - assert additive_mask.dtype == hidden_states.dtype - assert additive_mask.shape == (2, 1, 1, 5) - assert additive_mask[0, 0, 0, 0] == 0 - assert torch.isneginf(additive_mask[0, 0, 0, -1]) - - with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): - expected = reference(hidden_states, attention_mask=attention_mask) - actual = fastvideo(hidden_states, attention_mask=attention_mask) - assert_close(actual, expected, atol=1e-6, rtol=1e-6) - - -def test_zimage_transformer_rejects_sequence_parallelism(monkeypatch: pytest.MonkeyPatch): - model = _build_fastvideo(TINY_CONFIG) - monkeypatch.setattr("fastvideo.models.dits.zimage.model_parallel_is_initialized", lambda: True) - monkeypatch.setattr("fastvideo.models.dits.zimage.get_sp_world_size", lambda: 2) - - with pytest.raises(NotImplementedError, match="sp_size=1"): - model( - hidden_states=[torch.zeros(2, 1, 4, 4)], - encoder_hidden_states=[torch.zeros(3, 8)], - timestep=torch.tensor([0.5]), - ) - - -def _has_real_weights() -> bool: - return any(TRANSFORMER_DIR.glob("*.safetensors")) - - -def _load_real_config() -> dict: - with (TRANSFORMER_DIR / "config.json").open(encoding="utf-8") as file: - config = json.load(file) - assert config.pop("_class_name") == "ZImageTransformer2DModel" - config.pop("_diffusers_version", None) - config.pop("_name_or_path", None) - for key, expected in PRODUCTION_CONFIG.items(): - actual = config[key] - if isinstance(expected, tuple): - actual = tuple(actual) - assert actual == expected, f"unexpected production config {key}={actual!r}" - assert set(config) == set(PRODUCTION_CONFIG) - return config - - -def _run_fastvideo_production(inputs: tuple, dtype: torch.dtype) -> list[torch.Tensor]: - from fastvideo.configs.pipelines.base import PipelineConfig - from fastvideo.fastvideo_args import FastVideoArgs - from fastvideo.models.loader.component_loader import TransformerLoader - - precision = "bf16" if dtype == torch.bfloat16 else "fp16" - args = FastVideoArgs( - model_path=str(TRANSFORMER_DIR), - dit_cpu_offload=False, - dit_layerwise_offload=False, - use_fsdp_inference=False, - pipeline_config=PipelineConfig(dit_config=ZImageDiTConfig(), dit_precision=precision), - ) - model = TransformerLoader().load(str(TRANSFORMER_DIR), args).eval() - assert isinstance(model, ZImageTransformer2DModel) - assert args.model_paths["transformer"] == str(TRANSFORMER_DIR) - assert next(model.parameters()).device.type == "cuda" - with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): - outputs = model( - hidden_states=inputs[0], - encoder_hidden_states=inputs[1], - timestep=inputs[2], - )[0] - result = [output.detach().float().cpu() for output in outputs] - del model - torch.cuda.empty_cache() - return result - - -def _run_reference_production(reference_transformer_cls, config: dict, inputs: tuple, - dtype: torch.dtype) -> list[torch.Tensor]: - from safetensors.torch import load_file - - state_dict = {} - for shard in sorted(TRANSFORMER_DIR.glob("*.safetensors")): - state_dict.update(load_file(str(shard), device="cpu")) - with torch.device("meta"): - model = reference_transformer_cls(**config) - model.load_state_dict(state_dict, strict=True, assign=True) - del state_dict - model = model.to(device="cuda", dtype=dtype).eval() - with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): - outputs = model(inputs[0], inputs[2], inputs[1])[0] - result = [output.detach().float().cpu() for output in outputs] - del model - torch.cuda.empty_cache() - return result - - -@pytest.mark.skipif( - not _has_real_weights(), - reason=("Pinned Z-Image transformer weights are absent; set ZIMAGE_TRANSFORMER_DIR " - f"to Tongyi-MAI/Z-Image-Turbo@{HF_REVISION}/transformer"), -) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Real-weight Z-Image transformer parity requires CUDA") -def test_zimage_transformer_production_loader_forward_parity(reference_transformer_cls): - config = _load_real_config() - dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 - generator = torch.Generator(device="cuda").manual_seed(789) - inputs = ( - [torch.randn((16, 1, 4, 4), device="cuda", dtype=dtype, generator=generator)], - [torch.randn((3, 2560), device="cuda", dtype=dtype, generator=generator)], - torch.tensor([0.25], device="cuda", dtype=dtype), - ) - - fastvideo_outputs = _run_fastvideo_production(inputs, dtype) - reference_outputs = _run_reference_production(reference_transformer_cls, config, inputs, dtype) - assert len(fastvideo_outputs) == len(reference_outputs) == 1 - - actual, expected = fastvideo_outputs[0], reference_outputs[0] - diff = (actual - expected).abs() - mean_scale = expected.abs().mean().clamp_min(1e-6) - print(f"Z-Image real-weight parity: max={diff.max():.6f}, mean={diff.mean():.6f}") - assert diff.mean() / mean_scale < 0.01 - assert_close(actual, expected, atol=1e-2, rtol=1e-2) diff --git a/tests/local_tests/zimage/test_zimage_vae_parity.py b/tests/local_tests/zimage/test_zimage_vae_parity.py deleted file mode 100644 index 5c71bd674..000000000 --- a/tests/local_tests/zimage/test_zimage_vae_parity.py +++ /dev/null @@ -1,195 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Parity test for Z-Image VAE (AutoencoderKL decode path). - -This compares FastVideo's AutoencoderKL wrapper against the Z-Image -reference AutoencoderKL implementation using the same local weights. - -Usage: - pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v -""" - -from __future__ import annotations - -import importlib -import json -import subprocess -import sys -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch -from safetensors.torch import load_file as safetensors_load_file -from torch.testing import assert_close - -from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig -from fastvideo.models.loader.component_loader import VAELoader -from fastvideo.models.vaes.autoencoder_kl import AutoencoderKL as FastVideoAutoencoderKL - -REPO_ROOT = Path(__file__).resolve().parents[3] -ZIMAGE_REPO = REPO_ROOT / "Z-Image" -ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src" -ZIMAGE_VAE_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "vae" -ZIMAGE_VAE_CFG = ZIMAGE_VAE_DIR / "config.json" -ZIMAGE_VAE_WEIGHTS = ZIMAGE_VAE_DIR / "diffusion_pytorch_model.safetensors" -ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" -ZIMAGE_VAE_SCALING_FACTOR = 0.3611 -ZIMAGE_VAE_SHIFT_FACTOR = 0.1159 -PARITY_SCOPE = "both" - - -def _require_pinned_reference_module(module_name: str, source_file: Path): - if not ZIMAGE_REPO.exists(): - pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}") - if not source_file.is_file(): - pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}") - - try: - result = subprocess.run( - ["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"], - check=True, - capture_output=True, - text=True, - ) - except (FileNotFoundError, subprocess.CalledProcessError) as exc: - pytest.fail(f"Cannot verify Z-Image reference revision: {exc}") - - actual_revision = result.stdout.strip() - assert actual_revision == ZIMAGE_REFERENCE_REVISION, ( - "Z-Image reference clone is not at the pinned revision: " - f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}") - - if str(ZIMAGE_SRC) not in sys.path: - sys.path.insert(0, str(ZIMAGE_SRC)) - try: - module = importlib.import_module(module_name) - except Exception as exc: - pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}") - - module_file = Path(module.__file__ or "").resolve() - assert module_file.is_relative_to( - ZIMAGE_SRC.resolve()), (f"{module_name} resolved outside the pinned clone: {module_file}") - return module - - -@pytest.fixture(scope="module") -def reference_autoencoder_cls(): - module = _require_pinned_reference_module( - "zimage.autoencoder", - ZIMAGE_SRC / "zimage" / "autoencoder.py", - ) - return module.AutoencoderKL - - -def _load_cfg() -> dict: - if not ZIMAGE_VAE_CFG.exists(): - pytest.skip(f"Z-Image VAE config not found: {ZIMAGE_VAE_CFG}") - with ZIMAGE_VAE_CFG.open("r", encoding="utf-8") as f: - cfg = json.load(f) - cfg.pop("_class_name", None) - cfg.pop("_diffusers_version", None) - cfg.pop("_name_or_path", None) - return cfg - - -def _load_weights() -> dict[str, torch.Tensor]: - if not ZIMAGE_VAE_WEIGHTS.exists(): - pytest.skip(f"Z-Image VAE weights not found: {ZIMAGE_VAE_WEIGHTS}") - return safetensors_load_file(str(ZIMAGE_VAE_WEIGHTS), device="cpu") - - -def _build_reference(reference_autoencoder_cls, cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module: - ref = reference_autoencoder_cls( - in_channels=cfg["in_channels"], - out_channels=cfg["out_channels"], - down_block_types=tuple(cfg["down_block_types"]), - up_block_types=tuple(cfg["up_block_types"]), - block_out_channels=tuple(cfg["block_out_channels"]), - layers_per_block=cfg["layers_per_block"], - latent_channels=cfg["latent_channels"], - norm_num_groups=cfg["norm_num_groups"], - scaling_factor=cfg["scaling_factor"], - shift_factor=cfg["shift_factor"], - use_quant_conv=cfg["use_quant_conv"], - use_post_quant_conv=cfg["use_post_quant_conv"], - mid_block_add_attention=cfg["mid_block_add_attention"], - ).eval() - ref.load_state_dict(weights, strict=True) - return ref - - -def _build_fastvideo(cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module: - fv_cfg = AutoencoderKLVAEConfig() - fv_cfg.update_model_arch(cfg) - fv = FastVideoAutoencoderKL(fv_cfg).eval() - fv.load_state_dict(weights, strict=True) - return fv - - -def _load_fastvideo_production_vae(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "fastvideo.models.loader.component_loader.get_local_torch_device", - lambda: torch.device("cpu"), - ) - pipeline_config = SimpleNamespace( - vae_config=AutoencoderKLVAEConfig(), - vae_precision="fp32", - ) - fastvideo_args = SimpleNamespace( - model_paths={}, - pipeline_config=pipeline_config, - vae_cpu_offload=False, - ) - vae = VAELoader().load(str(ZIMAGE_VAE_DIR), fastvideo_args) - assert isinstance(vae, FastVideoAutoencoderKL) - assert fastvideo_args.model_paths["vae"] == str(ZIMAGE_VAE_DIR) - return vae - - -def test_zimage_vae_decode_parity(reference_autoencoder_cls): - torch.manual_seed(7) - - cfg = _load_cfg() - weights = _load_weights() - - ref = _build_reference(reference_autoencoder_cls, cfg, weights) - fv = _build_fastvideo(cfg, weights) - - # Decode-only parity is the critical path for Z-Image pipeline integration. - latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32) - - with torch.no_grad(): - ref_out = ref.decode(latents, return_dict=False)[0].detach().float() - fv_out = fv.decode(latents, return_dict=False)[0].detach().float() - - assert ref_out.shape == fv_out.shape - assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4) - - -def test_zimage_vae_production_loader_and_raw_decode_parity( - reference_autoencoder_cls, - monkeypatch: pytest.MonkeyPatch, -): - torch.manual_seed(17) - - cfg = _load_cfg() - weights = _load_weights() - assert cfg["scaling_factor"] == ZIMAGE_VAE_SCALING_FACTOR - assert cfg["shift_factor"] == ZIMAGE_VAE_SHIFT_FACTOR - ref = _build_reference(reference_autoencoder_cls, cfg, weights) - fv = _load_fastvideo_production_vae(monkeypatch) - - assert ref.config.scaling_factor == cfg["scaling_factor"] - assert fv.config.scaling_factor == cfg["scaling_factor"] - assert ref.config.shift_factor == cfg["shift_factor"] - assert fv.config.shift_factor == cfg["shift_factor"] - - decode_latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32) - - with torch.no_grad(): - ref_out = ref.decode(decode_latents, return_dict=False)[0].detach().float() - fv_out = fv.decode(decode_latents, return_dict=False)[0].detach().float() - - assert ref_out.shape == fv_out.shape - assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4)