1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Commit Graph

62 Commits

Author SHA1 Message Date
Sayak Paul
7ecb1e1b4f Flux3 image (#15046)
* Bump the actions group with 17 updates (#12)

Bumps the actions group with 17 updates:

| Package | From | To |
| --- | --- | --- |
| [actions/checkout](https://github.com/actions/checkout) | `6` | `7` |
| [actions/upload-artifact](https://github.com/actions/upload-artifact) | `4.6.2` | `7.0.1` |
| [docker/setup-buildx-action](https://github.com/docker/setup-buildx-action) | `3.12.0` | `4.2.0` |
| [docker/login-action](https://github.com/docker/login-action) | `3.7.0` | `4.6.0` |
| [docker/build-push-action](https://github.com/docker/build-push-action) | `6.19.2` | `7.3.0` |
| [huggingface/doc-builder/.github/workflows/build_main_documentation.yml](https://github.com/huggingface/doc-builder) | `e60a538eea9817ab312196d0d233604b01697265` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` |
| [actions/setup-python](https://github.com/actions/setup-python) | `5.6.0` | `7.0.0` |
| [huggingface/doc-builder/.github/workflows/build_pr_documentation.yml](https://github.com/huggingface/doc-builder) | `e60a538eea9817ab312196d0d233604b01697265` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` |
| [anthropics/claude-code-action](https://github.com/anthropics/claude-code-action) | `1.0.155` | `1.0.193` |
| [huggingface/security-workflows/.github/workflows/codeql-reusable.yml](https://github.com/huggingface/security-workflows) | `1.1.0` | `1.2.0` |
| [actions/download-artifact](https://github.com/actions/download-artifact) | `4.3.0` | `8.0.1` |
| [actions/labeler](https://github.com/actions/labeler) | `5.0.0` | `7.0.0` |
| [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) | `1.13.0` | `1.14.2` |
| [huggingface/tailscale-action](https://github.com/huggingface/tailscale-action) | `7d53c9737e53934c30290b5524d1c9b4a7c98c8a` | `07d36e4c3b68707bd3fe070dbb96f03f993a6efd` |
| [trufflesecurity/trufflehog](https://github.com/trufflesecurity/trufflehog) | `3.94.2` | `3.97.0` |
| [crate-ci/typos](https://github.com/crate-ci/typos) | `1.42.1` | `1.49.0` |
| [huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml](https://github.com/huggingface/doc-builder) | `9ad2de8582b56c017cb530c1165116d40433f1c6` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` |


Updates `actions/checkout` from 6 to 7
- [Release notes](https://github.com/actions/checkout/releases)
- [Commits](https://github.com/actions/checkout/compare/v6...v7)

Updates `actions/upload-artifact` from 4.6.2 to 7.0.1
- [Release notes](https://github.com/actions/upload-artifact/releases)
- [Commits](https://github.com/actions/upload-artifact/compare/v4.6.2...v7.0.1)

Updates `docker/setup-buildx-action` from 3.12.0 to 4.2.0
- [Release notes](https://github.com/docker/setup-buildx-action/releases)
- [Commits](8d2750c68a...bb05f3f551)

Updates `docker/login-action` from 3.7.0 to 4.6.0
- [Release notes](https://github.com/docker/login-action/releases)
- [Commits](c94ce9fb46...dbcb813823)

Updates `docker/build-push-action` from 6.19.2 to 7.3.0
- [Release notes](https://github.com/docker/build-push-action/releases)
- [Commits](10e90e3645...53b7df96c9)

Updates `huggingface/doc-builder/.github/workflows/build_main_documentation.yml` from e60a538eea9817ab312196d0d233604b01697265 to af49b8dc3859c3d80c7220c8c0ba146431132e49
- [Release notes](https://github.com/huggingface/doc-builder/releases)
- [Commits](e60a538eea...af49b8dc38)

Updates `actions/setup-python` from 5.6.0 to 7.0.0
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/v5.6.0...v7)

Updates `huggingface/doc-builder/.github/workflows/build_pr_documentation.yml` from e60a538eea9817ab312196d0d233604b01697265 to af49b8dc3859c3d80c7220c8c0ba146431132e49
- [Release notes](https://github.com/huggingface/doc-builder/releases)
- [Commits](e60a538eea...af49b8dc38)

Updates `anthropics/claude-code-action` from 1.0.155 to 1.0.193
- [Release notes](https://github.com/anthropics/claude-code-action/releases)
- [Commits](80b3182633...9d7150bc8a)

Updates `huggingface/security-workflows/.github/workflows/codeql-reusable.yml` from 1.1.0 to 1.2.0
- [Release notes](https://github.com/huggingface/security-workflows/releases)
- [Commits](dc6ca34688...e0d40aaf3c)

Updates `actions/download-artifact` from 4.3.0 to 8.0.1
- [Release notes](https://github.com/actions/download-artifact/releases)
- [Commits](https://github.com/actions/download-artifact/compare/v4.3.0...3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c)

Updates `actions/labeler` from 5.0.0 to 7.0.0
- [Release notes](https://github.com/actions/labeler/releases)
- [Commits](8558fd7429...bf12e9b00b)

Updates `pypa/gh-action-pypi-publish` from 1.13.0 to 1.14.2
- [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases)
- [Commits](ed0c53931b...dc37677b2e)

Updates `huggingface/tailscale-action` from 7d53c9737e53934c30290b5524d1c9b4a7c98c8a to 07d36e4c3b68707bd3fe070dbb96f03f993a6efd
- [Release notes](https://github.com/huggingface/tailscale-action/releases)
- [Commits](7d53c9737e...07d36e4c3b)

Updates `trufflesecurity/trufflehog` from 3.94.2 to 3.97.0
- [Release notes](https://github.com/trufflesecurity/trufflehog/releases)
- [Commits](6bd2d14f7a...bcfcf73aaf)

Updates `crate-ci/typos` from 1.42.1 to 1.49.0
- [Release notes](https://github.com/crate-ci/typos/releases)
- [Changelog](https://github.com/crate-ci/typos/blob/master/CHANGELOG.md)
- [Commits](65120634e7...8a48f81b6c)

Updates `huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml` from 9ad2de8582b56c017cb530c1165116d40433f1c6 to af49b8dc3859c3d80c7220c8c0ba146431132e49
- [Release notes](https://github.com/huggingface/doc-builder/releases)
- [Commits](9ad2de8582...af49b8dc38)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: '7'
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: actions/upload-artifact
  dependency-version: 7.0.1
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: docker/setup-buildx-action
  dependency-version: 4.2.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: docker/login-action
  dependency-version: 4.6.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: docker/build-push-action
  dependency-version: 7.3.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: huggingface/doc-builder/.github/workflows/build_main_documentation.yml
  dependency-version: af49b8dc3859c3d80c7220c8c0ba146431132e49
  dependency-type: direct:production
  dependency-group: actions
- dependency-name: actions/setup-python
  dependency-version: 7.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml
  dependency-version: af49b8dc3859c3d80c7220c8c0ba146431132e49
  dependency-type: direct:production
  dependency-group: actions
- dependency-name: anthropics/claude-code-action
  dependency-version: 1.0.193
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: actions
- dependency-name: huggingface/security-workflows/.github/workflows/codeql-reusable.yml
  dependency-version: 1.2.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: actions/download-artifact
  dependency-version: 8.0.1
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: actions/labeler
  dependency-version: 7.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: pypa/gh-action-pypi-publish
  dependency-version: 1.14.2
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: huggingface/tailscale-action
  dependency-version: 07d36e4c3b68707bd3fe070dbb96f03f993a6efd
  dependency-type: direct:production
  dependency-group: actions
- dependency-name: trufflesecurity/trufflehog
  dependency-version: 3.97.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: crate-ci/typos
  dependency-version: 1.49.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml
  dependency-version: af49b8dc3859c3d80c7220c8c0ba146431132e49
  dependency-type: direct:production
  dependency-group: actions
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Scope GITHUB_TOKEN permissions per job (#16)

A job with no `permissions:` block inherits whatever the repository
hands out. Each block added here grants what that job's own steps
need and nothing more.

Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>

* Pin GitHub Actions to commit SHAs (#17)

A tag is mutable: whoever controls the action's repository can move
`v4` to different code, and the next run picks it up with no diff to
review. Pinning to the commit SHA freezes the code that runs today.

Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>

* Bump doc-builder workflow pin (#22)

* Bump doc-builder pin to 9fc8a41 in build_documentation.yml

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Bump doc-builder pin to 9fc8a41 in build_pr_documentation.yml

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Bump doc-builder pin to 9fc8a41 in upload_pr_documentation.yml

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

* Add FLUX.3 image: Flux3Transformer2DModel and modular text-to-image / image-to-image pipeline

Adds the FLUX.3 image transformer, the modular pipeline (text-to-image and image-to-image with one
prompt per call and one or more reference images), Flux3ImageProcessor, the converter for BFL's
model package, tests and docs. Verified against BFL's reference on the real checkpoint
(integration-notes/parity_bf16.py, temporary folder).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* integration-notes: ship the knight example image with the doc examples

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* integration-notes: drop the tokenizer config workaround (not needed with transformers 5.x)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Address review: rename the mode channel dict and the final norm, name the reference workflow
image_conditioned, document the guidance scale

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Add the step-distilled FLUX.3 image model (flow map)

- transformer: `flow_map` config and `target_timestep_embedder`; forward takes
  `target_timestep` (the timestep each step jumps to), added to the modulation
  of every mode and the final layer
- convert script: map `ref_time_in.*` and read `flow_map` from the package config
- modular: dedicated `Flux3Distilled*` blocks (text encoder, text input, text
  RoPE inputs, denoiser) with no guider and no negative prompt, 16 steps by
  default; `Flux3DistilledAutoBlocks` and `Flux3DistilledModularPipeline`
- tests for the flow-map transformer path and the distilled modular pipeline
- docs: step-distilled section and autodoc entries

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Add the step-distilled examples script

Same four cases as flux3_examples.py, run against the distilled Hub repo.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* feat: add a single file support in Flux3 image. (#32)

* Warn about the denoiser input fields that are not passed to the transformer

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Document reproducible generation with user-provided latents

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Run the QK norm and RoPE in fp32 with a single cast back, matching the BFL implementation

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* parity script: run against the Hub repos and print the imported diffusers source

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Wrap a forward docstring line for the doc-builder style check

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* support lora loading (including ostris) in flux3 image pipeline. (#25)

* feat: support loading and performing inference with single-file formatted FP8 and NVFP4 checkpoints (#24)

* feat: support loading siungle file format checkpoints including quants.

* remove updates.patch

* remove fluff.

* feat: add a single file support in Flux3 image.

* bfl prefix

* version guard.

* prpoagate changes.

* add tests

* doc mods and remove integration notes

* remove script for quant

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

* add LoRA training for Flux 3 image small

* tests: clarify FLUX3_TINY_REPO comment

The tiny repo is uploaded to hf-internal-testing and the env var is a
local-dev escape hatch, not a pre-PR TODO.

* examples/flux3: drop unused args and dead class

Remove the text-encoder weight-decay arg from both trainers (neither
trains the text encoder) and, from the img2img trainer, the
prior-preservation leftovers (--class_data_dir, --sample_batch_size,
PromptDataset) that it never uses.

* bump yarn example rank to 32 in flux3 small readme

* Document <ref_image_n> references and layout prompts for FLUX.3

Matches the prompt grammar in BFL's reference inference code (bfl-hf-flux3-confidential#3): references are addressed as <ref_image_n> in the order passed to image, and regions can be placed with a trailing JSON list of id / bbox / desc (from / src_bbox / tgt_bbox when editing).

* examples/flux3: set klein model id, document ref image prompts, clean leftover refs

* point flux3 image docs at the klein release model ids

* fix: TestFlux3DistilledModularPipelineWorkflow::test_workflow_defaults test (#35)

* fix checkpoint ids in the flux3 docs. (#38)

* [docs] update pipe loading with vae fp32 dtype (#39)

* update pipe loading with vae fp32 dtype

* updating `load_componenets`

* use actual ckpts. (#42)

* flux3 klein example: rename from "small" and fix caption dropout for long prompts (#41)

* examples/flux3: rename small trainer, tests and readme files to klein

* examples/flux3: fix caption dropout shape mismatch for prompts longer than the pad window

* examples/flux3: add caption dropout to the i2i trainer

* examples/flux3: complete the klein rename (test class names, example output dirs)

* examples/flux3: guard against train_batch_size>1 with variable-length captions

* update flux3.md (#43)

* examples/flux3: support train_batch_size>1 and prior preservation with variable-length captions (#44)

* doc

* autodocstrings.

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>
Co-authored-by: Pauline Bailly-Masson <155966238+paulinebm@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: linoytsaban <linoy.tsaban@gmail.com>
Co-authored-by: apolinário <joaopaulo.passos@gmail.com>
Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
Co-authored-by: Aritra Roy Gosthipaty <aritra.born2fly@gmail.com>
Co-authored-by: Yumeng Li <109459218+YumengLi007@users.noreply.github.com>
2026-10-15 22:00:50 +05:30
Jingya HUANG
ee849197ce [core] Tensor parallelism for Qwen-Image-2.1 (#14865)
* feat: qwen image 21 tp plan

* test: vallidated on neuron

* fix: lru_cache_unless_export leak

* fix: use complex RoPE on TPU for Qwen Image 2.1

TorchTPU reports TPU tensors as "tpu" and supports complex dtypes, so
only Neuron needs the angle-based RoPE path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor: inline the Neuron device check for Qwen Image 2.1 RoPE

* review: cache Qwen Image 2.1 RoPE device freqs with lru_cache

* review: drop redundant torch.device

* fix: use torch.accelerator.device_count() when resolving the parallel device

* refactor: drop the Neuron branch in Qwen Image 2.1 block ids

* restore: neuron  repeat_interleave lowering error fix

* review: add TensorParallelNeuronTesterMixin for Neuron TP tests

* review: revert Qwen Image 2.1 RoPE freqs cache to a per-instance dict

* review: drop the Neuron RoPE path in Qwen Image 2.1

Neuron now runs complex tensors natively, so Qwen Image 2.1 uses the
default complex RoPE on every device. This removes
`apply_rotary_emb_qwen_neuron`, `ROPE_PER_DEVICE` and the per-device
freqs cache, and keeps `self.freqs = [freq.to(device) ...]` as on main.

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-10-15 17:07:50 +05:30
Jingya HUANG
1d5d056ecf [TPU] TorchTPU backend integration - eager / torch.compile / tp (#14039)
* feat: add torchtpu

* feat:draft TorchTPU support

* fix: wan overflow issue + compile mode error on sdxl

* doc: enhance with TorchTPU doc

* doc: enhance with TorchTPU doc

* style: remove unused imports flagged by ruff

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* docs: remove Debug Eager and Fused Eager sections from tpu.md

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* feat: add TP for TPU

* style: fix import sorting in TPU test scripts

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* style: ruff format TPU test scripts

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: style

* fix: test for native 4 devices

* fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check

- Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload
  support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their
  `# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow,
  chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of
  StableDiffusionXLPipeline.encode_prompt are intentionally left untouched;
  they'll be handled in a follow-up PR that fixes device placement for every
  pipeline component (not just text encoders).
- Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking
  the docs build: "not present in the table of contents").
- Remove the redundant "prefer non-CPU, non-meta component" loop from
  DiffusionPipeline._execution_device: PR #14383 already fixed this in
  DiffusionPipeline.device, which _execution_device falls back to. Verified
  on TPU hardware that _execution_device still resolves correctly for a
  split-placement pipeline after the removal.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* tests: cleanup

* fix: fix compile mode

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* test: remove flux2 e2e test

* doc: apply suggestions

* doc: apply suggestions

* review: remove monkey patch

* review: revert neuron-specific changes in the tests

* review: apply suggestions

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* doc: add tpu to tp doc

* review: remove unnecessary for tpu

* test: assert every _tp_plan parameter is sharded after a TP load

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* removal: delete redundant tp shard helpers

* removal: drop TPU workarounds no longer needed after recent main changes

* doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example

* doc: encode the prompt on CPU in the TPU TP example so it runs end to end

* test: tighten TPU TP tolerance and simplify the TPU TP worker

* doc: use enable_model_cpu_offload() in the TPU eager example

* doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example

* doc: run the text encoders on TPU in the compiled example

* test: drop the redundant TPU sync in the TP worker

* test: trim comments in the TPU TP tests

* test: run the TPU TP test with mp.spawn like the CUDA one

* test: pin the TPU TP test to 4 chips

* Update docs/source/en/optimization/tpu.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples

* test: skip the TPU TP test for models without a _tp_plan

* doc: correct the warning of recompilation

* fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet

* doc: update advice for compile mode when one chip doesn't fit

* test: pass the TPU TP tolerances as test arguments

* review: apply suggestions to the doc

* review: view compilation patch for tpu-only

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: JingyaHuang <huang_jingya@outlok.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-10-09 16:16:23 +05:30
Christopher
78bf2c6266 [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163)
* [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein)

Adds loading of LoKr (LyCORIS Kronecker product) adapters:

- `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`,
  inferred from the tensor shapes via `_create_lokr_config` (decompose
  factor, per-module rank/alpha patterns).
- State dict conversions for the formats in the wild: ai-toolkit Z-Image
  (dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2
  (fused qkv), LyCORIS underscore format, and bare dotted diffusers paths.
- BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker
  factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's
  QKV projections and maps the adapter 1:1 (exact).
- Alpha follows the LyCORIS convention: scaling applies only to
  rank-decomposed factors and is baked into the weights at conversion.

Fixes #13221

* [LoRA] LoKr review follow-ups: adapter-agnostic error message, refuse fused-QKV load over an unfused adapter

- "Invalid adapter checkpoint. We currently support LoRA and LoKr." replaces the
  message that still said "LoRA checkpoint" and described the substring check.
- Loading a fused-QKV LoKr checkpoint now refuses when an adapter is already
  injected on to_q/to_k/to_v or add_{q,k,v}_proj: fuse_qkv_projections() would
  replace those modules and orphan it. Covered by a test.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* [LoRA] rename the LoKr tests to test_lokr_conversion.py

The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [LoRA] LoKr review follow-ups: scope to Z-Image/Flux2, model-level fused QKV, LoKrTesterMixin

- Only the Z-Image and Flux2 loader mixins accept LoKr checkpoints, the two that
  convert them; the other mixins are back to their LoRA-only check.
- The fused-QKV handling moves to `_maybe_fuse_qkv_projections_for_lokr` in
  peft_utils, called from `load_lora_adapter`, so the model fuses its projections
  whichever entry point loads the adapter. The quantized-model error explains why.
- `_bake_lokr_alpha` becomes `_bake_lokr_alpha_` as it edits the dict in place.
- The BFL Flux2 converter no longer accepts expanded diffusers block names under
  `diffusion_model.`: no LoKr checkpoint uses that layout; diffusers-named
  checkpoints have no prefix and take the generic converter.
- The LyCORIS converter names the keys it does not recognize.
- Tests move from tests/lora to a `LoKrTesterMixin` in tests/models/testing_utils,
  with the checkpoint formats tested on the Flux2 and Z-Image model test classes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [LoRA] mark ZImageLoraLoaderMixin.load_lora_weights as copied from Flux2

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-10-09 12:56:56 +05:30
Sayak Paul
c2798cc785 [tests] refactor the pipeline-level attention processor tests (#14701)
* refactor the pipeline-level attention processor tests

* simplify stuff.

* remove utils.py
2026-10-05 15:38:23 +05:30
Jingya HUANG
863092ffec [core] Shard tensor-parallel checkpoints on load and save (#14544)
* [core] Shard tensor-parallel checkpoints on load and save

Stream each rank's slice of a tensor-parallel checkpoint straight off disk
instead of materializing the full checkpoint on every rank and resharding it
afterwards, and gather the shards back on save.

- `from_pretrained(..., parallel_config=TensorParallelConfig(...))` resolves
  the shard specs on the still-meta model, then slices each safetensors tensor
  before the dtype cast, so host memory peaks at ~1/tp_degree of the checkpoint.
- `save_pretrained` all-gathers the DTensors into an ordinary checkpoint, or
  writes a distributed checkpoint with `dcp=True` so no full tensor is ever
  formed. The writing `tp_degree` is recorded, since a packed weight's stored
  layout is interleaved by it.
- Factor the plan interpretation out of the Neuron pre-shard path into shared
  `TPShardSpec` / `resolve_tp_shard_specs` / `_local_shard` / `_hooks_only_styles`
  helpers, so both backends and both the load and save paths shard identically.

* Raise when tensor parallelism is combined with quantization, offloading or LoRA

Addresses the remaining two items of the review on #13718: tensor parallelism was rejected
alongside quantization and `device_map` only on the `from_pretrained` streaming path, while
`enable_parallelism` — which the quantization error message itself recommended — accepted a
quantized, offloaded or adapter-injected model and sharded it anyway.

- Add `_check_tp_model_state`, called from `apply_tensor_parallel`, the one chokepoint every TP
  entry point funnels through. It rejects a model that is quantized, group-offloaded, placed by
  accelerate (`device_map` or CPU offload), or has PEFT layers injected. Placed before the
  device-type check so the reported reason is the useful one.
- Guard the reverse order too: `enable_group_offload`, the two pipeline CPU-offload methods, and
  `load_lora_adapter` now refuse a tensor-parallel model.
- `save_pretrained` refuses a quantized tensor-parallel model. Previously the `dcp=True` branch
  returned before the quantizer's serialization step, writing shards with no quantization
  metadata and no error.
- The DCP load guard checked the `quantization_config` kwarg only, so a pre-quantized checkpoint
  directory loaded silently; check the config's own entry too, and add the missing `_tp_plan`
  check that otherwise surfaced as a raw `AttributeError`.
- Correct the `from_pretrained` message and the doc sentence that pointed at `enable_parallelism`
  as a way to shard a quantized model.

The new tests are the first tensor-parallel tests that need neither an accelerator nor more than
one rank: every case asserts a raise before any collective, so they run single-process on gloo.

* review: remove tp save/ dcp related

* review: add a new helper .

* review: moved the TP checks into enable_parallelism

* review: move the warning under caution block

* review: add loading time numbers

* review: move tp size dividende check

* review: delete test

* review: consolidate tp checks

* review: consolidate the tp checks, all in tensor_parallel.py

* review: improve conditionals

* review: debug info

* review: check weights format before instantiating the model

* review: restore duplicated _find_mismatched_keys since irrelevant

* revert: delete duplicated helper

* review: update doc for loading benchmark of flux2-dev

* review: improve the raising check

* review: add an args section

* Test tensor-parallel loading from sharded checkpoints

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-10-01 10:16:29 +05:30
Sayak Paul
fef717ffb0 [core] propagate sage attention updates. (#14584)
* propagate sage attention updates.

* add sage blackwell.

* up

* style.

* address feedback

* style.
2026-09-29 18:47:23 +05:30
yzhautouskay
344d6e7300 Add SeaCache support for Cosmos3 pipelines (#14663)
* Add optional FP32 sampling state for Cosmos3

* Add cache execution contexts to Cosmos3 pipelines

* Add SeaCache support for Cosmos3

* Cap consecutive SeaCache reuses

* Align SeaCache defaults with Cosmos3 inference

* Enable SeaCache by default for Cosmos3

* Document SeaCache for Cosmos3

* Use FP32 sampling state by default for Cosmos3

* Fix SeaCache reconfiguration, transfer state isolation, and FP32 sampling

* Remove SeaCache stats and ablation instrumentation

* Move SeaCache to transformer-level and disable by default

* Align SeaCache and tests with model-level cache patterns

* Make Cosmos3 sampling state always FP32

* Add pipe definition to the docs snippet

* Remove redundant hook registry cache resets

* Use raw vision latents for SeaCache indicators

* Make Cosmos3 SeaCache safe for regional compilation

* Fix SeaCache with Cosmos3 model parallelism

* Support single-stream models in SeaCache

* Add Wan T2V support to SeaCache

* pass denoising step info to cache hooks via cache_context

Replaces SeaCache's pipeline-state callbacks: the pipeline attaches the loop
coordinates (step_index, sigma, num_inference_steps) to the denoiser call with
`cache_context(name, ...)`, carried by a frozen CacheContext dataclass and stored
on each stateful hook's StateManager (the registry no longer holds context).

Co-Authored-By: Dhruv Nair <dhruv.nair@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Fix Wan cache metadata and Cosmos modular cleanup

Pass denoising step metadata to Wan T2V conditional and unconditional cache contexts.

Restore stateful cache reset after Cosmos modular pipeline calls, restore the required Literal import, and update regression tests and documentation.

* Refactor Cosmos3 SeaCache decoder hook

* Apply CI formatting updates

* Fix distilled system prompt fallback

* Preserve Cosmos3 fp32 sampling state

---------

Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-09-18 11:25:32 +05:30
Sayak Paul
d77d530445 [torchao]: fix how FqnConfig is resolved and handled (#14686)
* fix: how fqnconfig passed to torchaoconfig is resolved.

* testing with quantize_()

* change torchao version in setup.py as well.
2026-09-15 22:01:21 +03:00
Sayak Paul
daad99a04a [tests] remove quanto tests. (#14254)
* remove quanto tests.

* deprecate quanto and remove tests
2026-08-27 13:59:40 +05:30
Sayak Paul
0aa743c08d [tests] migrate lora tests for pipelines to use new mixins (#14268)
* start migrating lora tests.

* massive refactor.

* up

* remove planning doc

* address review: tighten lora bias tests

* move run_pipe()

* collapse.

* up

* done

* update

* only use diffusers api contract in tests.

* fix: pass tmp_path through AuraFlow test_lora_B_bias override

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-20 17:19:50 +05:30
Dhruv Nair
ac56fa2f25 [CI] Fix authentication on forked PRs (#14534)
fix auth for forked PRs
2026-08-19 17:41:04 +05:30
Jingya HUANG
7d2e86a85a [core] Support tensor parallelism for model inference (CUDA, Neuron) (#13718)
* draft:add neuron as a legit backend

* feat: neuron-specific changes in the pipeline

* tests: eager tests

* draft: start with tp for flux2

* fix: style

* fix:apr_02 beta

* feat:add wan

* fix:pixart

* fix: rewrite flux swiglu activation to avoid gather op in neuron IR

* test: pixart compile mode on neuron

* cleanup & fix style

* merge: another change

* review: cleanup+suggestions

* fix: CIs style

* tests: add test units for tp

* fix: in case of text-encoder(s) on CPU

* review:cleanup+add test

* fix: style

* doc: remove it for now

* Add from_single_file support for SkyReelsV2 and ChronoEdit transformers (#13946)

SkyReels-V2 and ChronoEdit are both built on Wan, and their transformers have
the same keys as WanTransformer3DModel, so they reuse
convert_wan_transformer_to_diffusers (like WanVACE / WanAnimate). This lets the
community GGUF builds load directly.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

* multi-GPU VAE Fix for Cosmos 3 (#13924)

fix(cosmos3): pin VAE latent norm buffers to encode output device

Under sharded placement (device_map="balanced"), vae.encode() runs on the
VAE's own device while the mean/inv_std buffers were pinned to x.device,
causing a cross-device RuntimeError. Compute raw_mu first, then pin the
normalization buffers to its device so all tensors share one device.

Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* docs: fix repeated word typo in set_timesteps docstring (#13876)

* docs: fix repeated word typo in set_timesteps docstring

Removed the duplicate word "schedule" from the docstring for the sigmas argument in EulerDiscreteScheduler.set_timesteps.

* Update scheduling_euler_discrete.py

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

* clean some stuff to simplify code.

* clean more to remove permutation related shenanigans.

* revert: put torch.chunk back

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Address Sayak's TP review nits (docstrings + stale comments)

- _modeling_parallel: fix stale comment claiming attention processors read
  tp_degree at runtime (no longer true after the PackedColwise/Rowwise
  refactor; processors are TP-agnostic via head_dim).
- tensor_parallel: single backticks instead of double in docstrings, add
  input/output examples to _blocks_to_block_sizes and _resolve_tp_plan, and
  drop the stale "applies the Flux2 fused-weight permutations" line from
  apply_tensor_parallel (permuters were removed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix stale permuter reference in Flux2ParallelSelfAttnProcessor comment

The proportional QKV/MLP slice is now guaranteed by PackedColwiseParallel
block-by-block sharding, not the removed fused-weight permuter.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Derive TP backend inside apply_tensor_parallel from the mesh device type

Drops the `backend` parameter: apply_tensor_parallel already holds the TP
mesh, so it derives neuron-vs-default from `config._mesh.device_type` itself.
The caller no longer computes it. Documents why torch_device can't be used
(Neuron reports as "cpu" via torch's accelerator API). Removes the now-unused
is_torch_neuronx_available import from modeling_utils.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Simplify TP backend derivation to the mesh device type alone

Drop the redundant is_torch_neuronx_available() guard: a DeviceMesh("neuron", ...)
only exists when running on Neuron, so tp_mesh.device_type is the single source
of truth.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove incorrect "Neuron reports as cpu" rationale from TP backend docstring

Verified on the current native torch_neuronx drop that Neuron does surface as
the torch accelerator (torch._C._get_accelerator().type == "neuron",
get_device() == "neuron"); the old "reports as cpu" claim was stale. The
mesh-device-type derivation is unchanged — it's the device being sharded onto.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop stale "XLA" labels from Neuron device/runtime comments

Neuron is a native torch backend now (not XLA-routed): remove the "(XLA)"
parenthetical from the randn-on-CPU note in torch_utils and the "XLA runtime"
mention from the Neuron TP test docstring. Behavior unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* make style: reflow tensor_parallel docstrings

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Revert "revert: put torch.chunk back"

This reverts commit 7ea75f771a.

* fix: sharding for neuron + validate flux 1

* fix: change for check_repository_consistency

* feat: support qwen image as well

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* review: keep bria and nucleusmoe out

* review: keep bria and nucleusmoe out

* review: keep bria and nucleusmoe out

* review: address agent review

* review: apply suggestions

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* review: change example tp size to 4

* review: add flux tp test

* review: address comments

* review: improve the doc

* review:apply qwen rotary_emb per device

* review:remove redundant device def after PR #14383

* review: doc title suggestion

* review: restore copy comment

* review: remove cp divisibility explanation from doc on parallelism guide

* fix style

---------

Co-authored-by: Hz_Zhang <47402297+HaozheZhang6@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Ramkumar R <ramkumarashvanth09@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
2026-08-19 16:20:59 +05:30
Dhruv Nair
aeda394db7 [CI] Allow single file tests to run as fast tests (#14324)
* update

* update

* update

* update

* [FA2] Use `kernels-community/aiter-flash-attn-ck` Hub kernel for ROCm FlashAttention (drop `aiter` dependency) (#14436)

Use  Hub kernel

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* add HF_TOKEN to pr workflow

* fix token

* clean up

* Apply suggestions from code review

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: BADAOUI Abdennacer <106801897+Abdennacer-Badaoui@users.noreply.github.com>
2026-08-19 10:32:03 +05:30
Sayak Paul
425b113c21 [tests] refactor pipeline-level quantization tests (#14435)
* fix(torchao): route root-level tensors around safetensors reconstruction

Models with parameters at the root of the module tree (e.g. Wan's
`scale_shift_table`) crashed torchao's `unflatten_tensor_state_dict`
when loading serialized checkpoints, since flattened tensor names are
assumed to carry a module prefix. Filter such tensors and their
metadata entries out of the reconstruction and merge them back
unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(bnb): defer 8-bit weights split from their SCB stats across shards

Sharded serialization can place an 8-bit weight and its `SCB`
statistics in different shard files, in which case the shard-by-shard
loader failed with "Missing quantization component `SCB`". Hold the
incomplete half of the pair back until its counterpart arrives with a
later shard, mirroring the torchao pending mechanism, and disable
parallel shard loading for prequantized 8-bit checkpoints.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* [tests] extend model-level quantization tester mixins

Migrate remaining model-level coverage from tests/quantization into the
tester mixins so it runs for every wired model:

- base: buffer-placement assertions in the device-map test, and an
  opt-in sharded-serialization test enabled by setting
  `sharded_serialization_config`
- bnb: serialization across all configs (sharded included), dtype
  assignment and adapter training for 8-bit, device moves preserving
  the memory footprint, corrupted-state-dict loading error, and a fixed
  modules-to-not-convert test (BitsAndBytesConfig only exposes
  llm_int8_skip_modules; the old test passed an unsupported kwarg and
  only survived by being skipped)
- torchao: custom device maps with cpu/disk offload, generalized from
  the Flux-specific test
- gguf: the diffusers-format single-file loading path, wired into the
  Flux model tests

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* [tests] fix model-level quantization test wiring

- SD3.5: the quantized testers reused the random-init dummy inputs
  (4 latent channels, fp32) while the tiny Hub checkpoint has
  in_channels=8 and the quantizers load the model in half precision;
  give them matching inputs and relax the 4-bit memory expectation for
  the tiny checkpoint.
- QwenImage / Flux2: the quantized testers had no Hub checkpoint wired
  at all, so every test errored; point them at
  hf-internal-testing/tiny-qwenimage-pipe and tiny-flux2 with matching
  inputs.
- NucleusMoE: no tiny checkpoint exists on the Hub yet; comment the
  testers out like the LTX ones.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* [tests] split tests/quantization into pipeline-level and backend-level tiers

Pipeline-level quantization tests (pipeline quality slices, cpu offload,
LoRA loading, compile, PipelineQuantizationConfig) move to
tests/pipelines/testing_utils/quantization.py, marked per backend so the
nightly CI can select them with `pytest -m`. tests/quantization keeps
only backend-level tests that fit neither tier: config validation,
utility warnings, and GGUF CUDA kernel correctness. Tests already
covered by the model-level mixins are dropped.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* [ci] select nightly quantization tests by marker across test tiers

Each backend job now runs `pytest -m <marker>` over tests/models,
tests/quantization, and tests/pipelines/testing_utils/quantization.py,
giving the model-level mixin tests a nightly home with the backend
dependencies installed. The torchao job additionally installs mslk.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: sayak@huggingface.co <sayak@ip-10-53-90-45.ec2.internal>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-18 20:53:32 +02:00
Sayak Paul
d5baa4fb54 tests: fix cuda model tests. (#14464)
* tests: fix cuda model tests.

* up

* up

* up

* add to docstrings.

* restore previous beehaviour.

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-08-15 16:56:06 +02:00
Sayak Paul
9558d918ba [tests] tighten keep_in_32 modules tests (#14399)
* tighten keep_in_32 modules tests

* remove unneeded tests

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-08-13 08:49:35 +02:00
Sayak Paul
7564fb016d Ltx 2.5 (#14447)
* Add LTX-2.4 support to the LTX2 transformer, pipelines, and conversion script

LTX-2.4 reuses the existing LTX2 model and pipeline classes. The delta over
2.3 is a small set of additive config flags plus a Gemma 3 -> Gemma 4 text
encoder swap, so no new model or pipeline classes are introduced.

- transformer_ltx2.py: add `ff_bias`/`audio_ff_bias` (2.4's video FFN drops
  its bias) and `use_prompt_adaln_single` (toggles timestep-dependent prompt
  cross-attention modulation; when off, cross-attention K/V becomes
  timestep-independent and cacheable across denoising steps for a given
  prompt). Both default to their 2.3 behavior. The flag is read back from
  `self.config` rather than mirrored onto an instance attribute, and all
  three new constructor args are documented.

- pipeline_ltx2.py / pipeline_ltx2_image2video.py: add an optional
  `prompt_enhancer` component to both pipelines (previously T2V-only).
  LTX-2.4's fine-tuned text encoder is conditioning-only, so enhancement
  uses a separate off-the-shelf google/gemma-4-E2B-it checkpoint with its
  own message format and decoding recipe -- unlike LTX-2.0/2.3, where one
  checkpoint serves both roles. `enhance_prompt()` resolves the format and
  `.generate` kwargs from whichever model is active, and frees a dedicated
  enhancer from GPU memory right after use (guarded so it does not interfere
  with accelerate's offload hooks), mirroring the existing text_encoder
  handling.

  Both `__call__`s gain `enable_prompt_enhancement: bool | None = None`,
  which resolves to `True` when a dedicated `prompt_enhancer` is configured
  (LTX-2.4) or when `system_prompt` was passed explicitly (matching prior
  LTX-2.0/2.3 behavior exactly), and `False` otherwise. Explicit `False`
  disables enhancement even on 2.4. When enabled with no `system_prompt` on
  a 2.4 pipeline, the matching default system prompt is injected.

  `max_new_tokens`/`seed` keep their literal defaults (512/10)
  unconditionally: greedy decoding consumes no randomness, so `seed` is
  inert for the dedicated-enhancer case, and 512 tokens comfortably covers
  the target caption length. No public API or default changes were needed
  for LTX-2.0/2.3.

  Validation of the enhancement arguments happens in `check_inputs`, so an
  unsatisfiable request fails before the prompt is encoded rather than
  partway through generation. Text encoder, tokenizer, processor and
  enhancer type hints name the concrete Gemma classes they accept.

- utils.py: add a `PromptEnhancementConfig` dataclass plus
  `GEMMA3_PROMPT_ENHANCEMENT_CONFIG`/`GEMMA4_PROMPT_ENHANCEMENT_CONFIG` as
  the single source of truth for each model's message prefix and `.generate`
  kwargs, shared by both pipelines. Add the validated "capstyle_plus"
  LTX2_4_T2V/I2V_DEFAULT_SYSTEM_PROMPT strings, marked `docstyle-ignore` so
  `doc-builder style` cannot re-wrap them -- the prompts must stay
  byte-for-byte identical to the reference, newlines included.

- convert_ltx2_to_diffusers.py: add a "2.4" branch to all five
  get_ltx2_*_config functions. Transformer/VAE/vocoder configs are
  structurally identical to 2.3 (verified against the checkpoint's own
  safetensors metadata); only `ff_bias=False` differs. Connector
  `caption_channels`/`text_proj_in_factor` are now derived from the live
  Gemma text config rather than hardcoded, since 2.4's text encoder is not
  pinned to a single checkpoint the way Gemma-3-12B is for 2.0/2.3. Swap
  `Gemma3ForConditionalGeneration`/`Gemma3Processor` for
  `AutoModelForImageTextToText`/`AutoProcessor` so one path covers Gemma 3
  and Gemma 4. Raise a clear error when --version 2.4 is requested without
  pointing --text_encoder_model_id at a Gemma 4 (gemma4_unified) checkpoint,
  and add --prompt_enhancer_model_id, required whenever --version 2.4 is
  combined with --add_processor, since falling back to
  --text_encoder_model_id (correct for 2.0/2.3) would pair 2.4 with the
  wrong enhancement model. Also fixes a pre-existing vocoder
  class-selection bug that only checked for "2.3", and `processor` never
  being passed into the --full_pipeline LTX2Pipeline(...) construction.

- docs/source/en/api/pipelines/ltx2.md: add an "LTX-2.4" section covering
  what carries over from 2.3 unchanged (guidance recommendations, aside
  from a different STG block index) and what does not (a single-stage
  checkpoint only, with no two-stage or distilled workflow yet), plus the
  corrected prompt-enhancement recipe and its enabled-by-default behavior
  for both LTX2Pipeline and LTX2ImageToVideoPipeline.

* Fix resolution-dependent timestep shift being a no-op in all LTX2 pipelines

Every LTX2Pipeline variant's `mu = calculate_shift(...)` call passed the
scheduler's `max_image_seq_len` config value as the `image_seq_len`
argument instead of the current generation's actual packed sequence
length. Since calculate_shift's formula is a line through
(base_seq_len, base_shift) and (max_seq_len, max_shift), passing
image_seq_len == max_seq_len always evaluates to exactly max_shift --
so `mu` was pinned to a constant (2.05 for the LTX-2.4 scheduler config)
regardless of height/width/num_frames, even though `use_dynamic_shifting:
true` is set specifically to make this resolution-dependent.

Found while benchmarking LTX-2.4 diffusers output against the reference
pipeline with bit-identical starting noise: the reference computes this
shift from the real video token count (ltx_core's LTX2Scheduler.execute),
which diverges substantially from the constant diffusers was using at
any resolution other than exactly the checkpoint's max-anchor token
count. For a 768x512, 121-frame video (6144 tokens), the reference lands
on mu ~= 2.78 versus diffusers' constant 2.05.

Fixed by passing the current call's actual packed video latent length
(`latents.shape[1]`) as `image_seq_len`, matching the pattern already
used correctly in pipeline_flux.py (the source this was copied from).
Applies to all 5 LTX2 pipeline variants, each with their own independent
`__call__` (not linked via `# Copied from` for this method).

* Add LTX-2.4 duration head for automatic num_frames prediction

LTX-2.4 checkpoints ship a small regression head (~1.9M params) that predicts
the natural duration of the shot implied by a caption, from the same text
connector output the transformer is conditioned on. With it converted, a caller
can let the model choose the video length instead of picking `num_frames`.

- duration_head.py: `LTX2DurationHead` (a `ModelMixin` optional pipeline
  component) plus the `LTX2AutoDuration` request object. Modality-specific
  projections map the video and audio connector streams into a shared pooler
  dim, learnable modality embeddings tag them, one learnable query cross-attends
  the concatenation, and a small MLP regresses a log-duration. `forward` returns
  seconds as a tensor; `predict_num_frames` clamps to bounds and snaps to the
  VAE's causal temporal grid.

  The attention pooler uses explicit to_q/to_k/to_v/to_out with
  `dispatch_attention_fn` rather than `torch.nn.MultiheadAttention`, which both
  reference implementations use because that is the layout the checkpoint ships.
  `nn.MultiheadAttention` is documented in diffusers as breaking
  `enable_sequential_cpu_offload`; the split form also gets backend dispatch.

  Two details are load-bearing for numerical parity: the GELU must be
  tanh-approximated (the exact GELU gives different numbers against the
  JAX-trained head), and the clamp must precede the grid snap (a clamped frame
  count is not necessarily grid-aligned). Where narrow bounds convert to a frame
  window containing no grid point -- at 24 fps [1.0s, 1.02s] rounds to [24, 24],
  and 24 is not 8k + 1 -- the nearest grid point is used and a warning logged,
  rather than refusing to generate over a rounding artifact.

  The output MLP's config argument is `mlp_hidden_dim`, not the reference's
  `mlp_hidden`: the submodule keeps the checkpoint's `mlp_hidden` name, and
  `ModelMixin.__getattr__` resolves config keys ahead of submodules, so the two
  colliding would shadow the `nn.Linear` with an `int`.

- pipeline_ltx2.py / pipeline_ltx2_image2video.py: `num_frames` becomes
  `int | LTX2AutoDuration | None`. Omitting it auto-predicts when the checkpoint
  ships a head and keeps the legacy 121 otherwise, mirroring the reference --
  whose CLI also defaults to auto-prediction -- and matching the resolution
  these pipelines already do for `enable_prompt_enhancement`, which likewise
  switches on the presence of an optional LTX-2.4 component. Pre-2.4 pipelines
  have no head, so nothing changes for them.

  The prediction runs immediately after the connectors and before `num_frames`
  is first read. Only the positive half of the CFG-concatenated batch is used,
  and rows past the first are `num_videos_per_prompt` duplicates.

  Auto-duration is rejected from `check_inputs` -- before prompt enhancement and
  encoding, so a bad request costs nothing -- when there is no head, and when
  more than one prompt is supplied: the batch carries a single temporal
  dimension, so prompts with different natural lengths cannot share one frame
  count. Batched prompts with an explicit integer `num_frames` are unaffected.

- convert_ltx2_to_diffusers.py: `--duration_head` (with
  `--duration_head_prefix`) alongside the other per-component flags. The head's
  keys sit at the checkpoint top level rather than under the DiT prefix, so the
  existing prefix helper handles them. The checkpoint's fused `in_proj_weight`
  is split into separate q/k/v, and dimensions are read back from weight shapes
  since checkpoint metadata carries no duration_head config; only
  `num_pooler_heads` is not recoverable and is fixed at 4. Pre-2.4 checkpoints
  yield no keys and are skipped rather than failing.

Verified against ltx_core on the real 2.4 checkpoint: fed the reference's
recorded connector tokens, the converted head predicts 12.9375s (video only),
2.671875s (audio only) and 3.515625s (both) -- bit-identical to the reference --
and the same frame counts through the clamp/snap path at 24/25/30/8 fps. The
snapping arithmetic matches the reference `seconds_to_clamped_num_frames` across
260 combinations of duration, frame rate and bounds.

* Add LTX-2 modular video and audio decoder blocks

Add the first blocks of the LTX-2 modular pipeline under
modular_pipelines/ltx2/decoders.py:

- LTX2VaeDecoderStep: unpacks and decodes video latents (or returns
  latents for output_type="latent"), applying the optional decode-time
  noise on normalized latents before denormalizing, matching the
  standard LTX2Pipeline decode stage.
- LTX2AudioDecoderStep: unpacks and decodes audio latents into a
  waveform via the audio VAE and vocoder in a single block.

Pack/unpack/denormalize helpers are redefined at module level rather
than imported, since modular blocks must not import from
diffusers.pipelines.* (the vocoder class is imported from the pipelines
path for now, flagged for relocation to models/).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add LTX-2 modular pipeline (t2v + i2v) and parity harnesses

Add the LTX-2 modular pipeline package under modular_pipelines/ltx2/,
covering the joint video+audio text-to-video and image-to-video
workflows for LTX-2.4:

- encoders.py: dedicated Gemma-4 prompt enhancer (t2v/i2v), Gemma text
  encoder, text connectors, and the i2v image VAE encoder.
- before_denoise.py: text-input expansion, flow-match timesteps (with a
  deep-copied audio scheduler), video/audio latent prep, and RoPE coords.
- denoise.py: the joint video+audio denoise loop with manual guidance
  (CFG + spatio-temporal + modality-isolation), shared across t2v/i2v.
- decoders.py: video VAE decode and audio VAE + vocoder decode.
- modular_blocks_ltx2.py: LTX2Blocks (t2v), LTX2ImageToVideoBlocks (i2v),
  and LTX2AutoBlocks (both, default), plus the auto/conditional wrappers.
- modular_pipeline.py: LTX2ModularPipeline with the compression-ratio and
  patch-size properties the blocks read.

Wire up lazy imports and register the pipeline (top-level diffusers
exports, modular_pipelines __init__, MODULAR_PIPELINE_MAPPING, and dummy
objects).

Also add integrations/ (temporary, for-visibility) parity harnesses that
compare the modular t2v/i2v blocksets against the standard LTX-2
pipelines by sharing the same loaded components. This directory is meant
to be removed before the final integration.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add script to create tiny LTX-2.4 test pipeline ckpt (under assumption that it inherits all configs from LTX-2.3)

* make fix-copies

* Run modular connector once on CFG-concatenated batch for bitwise parity

Run the LTX-2 text connector once on the CFG-concatenated `[uncond, cond]`
batch (as `LTX2Pipeline` does) instead of once per branch, then split the
outputs back into uncond/cond. The connector is applied per batch element, so
both forms are mathematically equivalent, but its GEMM/attention kernels round
identically for a given row only at batch >= 2; running the branches separately
diverged from the standard pipeline by ~1e-6 at `num_videos_per_prompt=1`. The
modular path is now bitwise-identical to the standard pipeline at any batch size.

Also extend the T2V/I2V parity harnesses with a `--num_videos_per_prompt`
argument and a `--check_tensor_stats` flag for per-output min/mean/std/max.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Use tiny LTX2VocoderWithBWE in the LTX-2.4 test checkpoint

Swap the placeholder LTX2Vocoder for a scaled-down LTX2VocoderWithBWE that
mirrors LTX-2.3's vocoder (snakebeta + antialiasing, no final activation/bias,
16kHz -> 48kHz bandwidth extension) while keeping the same in/out channel
shapes. Dimensions are reduced but the shape invariants the two-stage forward
requires are preserved: in_channels = audio_vae.output_channels * mel_bins,
bwe_in_channels = out_channels * num_mel_channels, filter_length == window_length,
and prod(bwe_upsample_factors) == (output_sr // input_sr) * hop_length so the BWE
residual and the resampled stage-1 skip line up.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Match LTX-2.4 transformer flags in the test checkpoint

Only `ff_bias=False` is a transformer-level delta from LTX-2.3 per the
authoritative LTX-2.4 config in `scripts/convert_ltx2_to_diffusers.py`;
`audio_ff_bias` and `use_prompt_adaln_single` keep their `True` defaults.
The test checkpoint was incorrectly overriding both to `False`, so drop
those two overrides. T2V/I2V modular-vs-standard parity remains bitwise.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Add CPU offload flag to parity scripts to be able to load full checkpoints

* Support dynamic timestep shift and duration head in LTX-2 modular pipeline

Track two recent additions to the standard LTX-2.4 pipelines in the
modular blocks (t2v + i2v):

- Resolution-aware timestep shift: LTX2SetTimestepsStep now computes `mu`
  from the actual packed video sequence length (derived from
  height/width/num_frames and the transformer patch sizes) instead of a
  constant, matching the standard pipeline's `latents.shape[1]`-based
  shift. Uses the compute-from-dims approach (like LTX-1 / Flux2), so no
  block reordering is needed.

- Optional duration head: add LTX2DurationStep, which predicts a concrete
  `num_frames` from the connector text conditioning via the `duration_head`
  component when `num_frames` is an `LTX2AutoDuration` request, and
  re-emits it. Wrapped in the LTX2AutoDurationStep conditional (skipped for
  an integer `num_frames`) and wired into all three blocksets after the
  connector step, so `num_frames` is resolved before the shift and latent
  prep run.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add duration predictor support for tiny checkpoint and parity scripts

* Refactor LTX-2 modular guidance into video/audio guiders

Replace the manual multi-term guidance in `LTX2LoopDenoiser` with a
`LTX2Guidance` guider (in `guider.py`), instantiated once as the video
`guider` and once as the `audio_guider`. Each combines CFG + spatio-temporal
guidance (STG) + modality-isolation via the delta formulation in x0 space; the
denoiser owns a `plan_guidance_passes` union plan across the two guiders, runs
each transformer pass, converts velocity->x0, and delegates the per-modality
combine to the guiders. Guidance scales are now guider config, not `__call__`
kwargs. Parity harnesses updated to configure the guiders accordingly.

Parity investigation (not yet resolved):
- The refactor runs every guidance pass as its own single-batch transformer
  forward, whereas the standard `LTX2Pipeline` batches the cond+uncond CFG pair
  into one forward and runs STG/modality-isolation as separate single-batch
  passes. STG and modality already match (single-batch in both).
- The change is mathematically equivalent, not a logic bug: in fp32 the
  denoised latents match to ~8e-6 mean abs diff (sparse outliers up to
  ~3.5e-4), and disabling STG does not move the diff.
- But GPU matmul is not batch-invariant, so cond computed alone differs from
  cond computed inside a batch-of-2. Negligible in fp32 (~1e-6/op); ~1e-2/op in
  bf16, where amplification by the CFG delta and accumulation over sampler steps
  drives the modular vs. standard bf16 latents to ~10% mean-relative divergence.
- Net: numerical, but the modular pipeline does NOT reproduce the standard
  pipeline bitwise in bf16 (the real inference dtype). Restoring parity would
  require re-batching the cond+uncond pair into a single forward to match the
  reference execution, keeping STG/modality single-batch.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Batch cond+uncond CFG forward to mirror standard LTX-2 pipeline

Restore fp32 bitwise parity with `LTX2Pipeline` in the modular denoiser by
running the cond+uncond CFG pair as a single batched transformer forward
(`torch.cat([latents] * 2)` + `.chunk(2)`), keeping STG and modality-isolation
as separate single-batch conditional forwards -- matching the reference
op-for-op (batch sizes, repeated coords, cache-context names).

`plan_guidance_passes` now emits forward-groups (`identifiers` / `conditioning`
aligned lists + `flags` + `cache_context`) instead of one entry per pass; the
denoiser runs each group once and chunks the CFG forward back into its
`[uncond, cond]` identifiers. The `LTX2Guidance` combine is unchanged -- only
how the four x0 tensors are obtained changed.

Parity results:
- fp32: bitwise (0.0 max abs diff), verified at full-checkpoint scale including
  under CPU offload. This is the authoritative parity gate.
- bf16: the previous single-batch-per-pass design diverged ~10% mean-relative
  from the standard pipeline (GPU matmul is not batch-invariant: cond alone vs.
  cond in a batch-of-2). Batching the CFG pair removes that gap; a smaller
  ~1% (tiny) / ~5% (full) bf16 gap remains. It is not a logic difference (fp32
  is bitwise across scale and offload); it is a bf16-kernel effect -- coarser
  mantissa amplifying non-associative accumulation order, plus bf16 using
  different kernels than fp32 (tensor-core GEMM algorithm selection, fused
  attention). bf16 is therefore a close-but-not-bitwise check, not a gate.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drive LTX-2 modular guidance through the guider API (per-pass inputs)

Replace the batched-CFG denoiser with one that runs each guidance pass as its
own single-batch transformer forward, driven end-to-end through the standard
guider API. `LTX2Guidance.prepare_inputs(guider_inputs)` builds one identifier-
tagged batch per active pass from a dict whose values are 4-tuples indexed by
pass [cond, uncond, stg, modality]; the per-pass model flags
(`spatio_temporal_guidance_blocks`, `isolate_modalities`) ride in those tuples
alongside the encoder inputs, so a pass fully describes its own forward. The
denoiser unions both guiders' passes by identifier, runs each once (storing
video+audio x0 on the batch), and combines each modality via its guider's
`forward`/`__call__`, filtered to that guider's active passes so the batch count
matches `num_conditions`.

Removes the bespoke `plan_guidance_passes` union helper and the empty-dict
`prepare_inputs_from_block_state` call: the plan is now expressed as the
`guider_inputs` tuples + `active_predictions()`, so guidance logic lives behind
the guider API rather than in the denoiser.

Parity trade vs. the previous batched-CFG design:
- Batched CFG matched the standard pipeline op-for-op and was fp32-bitwise.
  Running every pass single-batch is mathematically equivalent but, since GPU
  matmul is not batch-invariant, cond computed alone differs from cond inside a
  batch-of-2: ~1e-4 mean-relative in fp32 on a full checkpoint (sparse outliers
  up to ~3.5e-4 max), ~10% mean-relative in bf16.
- This is numerical, not a logic difference, and fp32-within-tolerance (not
  bitwise) is the modular-ecosystem norm. The trade buys end-to-end guider-API
  usage (swappable within the LTX-2 guidance family, per-pass flags carried the
  same way as encoder inputs) at the cost of the bitwise guarantee.

Parity harnesses: gate on magnitude-aware stats (mean abs diff relative to mean
magnitude, plus a loose max-abs ceiling) instead of assert_close's near-bitwise
fp32 defaults, which no single-batch design can clear. fp32 is the authoritative
gate (1e-3/1e-3); bf16 is a loose sanity check (0.15/0.5). `--atol`/`--rtol`
become `--mean_rel_tol`/`--max_abs_tol`; motivation documented in-file.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drive LTX-2 modular guidance via prepare_inputs_from_block_state

Move `LTX2LoopDenoiser` onto the standard Wan/Z-Image guider idiom: the
denoiser now owns a `guider_input_fields` map (transformer arg -> per-pass
block-state attribute names, indexed [cond, uncond, stg, modality]) and calls
`guider.prepare_inputs_from_block_state(block_state, guider_input_fields)`
instead of hand-building a literal `guider_inputs` dict in `__call__`. This
lifts the cond/uncond/stg/modality field mapping to a construction-time arg
(swappable per workflow) and resolves the connector_*->encoder_hidden_states
name mismatch via the map keys.

The two per-pass model flags (`spatio_temporal_guidance_blocks`,
`isolate_modalities`) are pass-identity constants, not block-state
conditioning, so they can't ride the name-referenced field map; the denoiser
sets them on each batch by identifier after preparation, via a
`pass_flags.get(identifier, (None, False))` lookup. The plain-conditional
default keeps this correct for any guider that emits a subset of passes (e.g.
a swapped-in `ClassifierFreeGuidance` -> just pred_cond/pred_uncond gets no STG
and no modality isolation).

`LTX2Guidance` gains `prepare_inputs_from_block_state` (names, via the base
`_prepare_batch_from_block_state` helper); `prepare_inputs` (literals) is
retained so both halves of the guider data-prep API are implemented. Pass
structure and numerics are unchanged (still four single-batch forwards), so the
fp32/bf16 parity story is untouched; user-confirmed parity OK.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Route LTX-2 batch-invariant denoiser kwargs via denoiser_input_fields

Tag the three upstream-produced, batch-invariant transformer kwargs
(`audio_num_frames` from the audio-latents step; `video_coords` / `audio_coords`
from the coords step) with `kwargs_type="denoiser_input_fields"`, and have
`LTX2LoopDenoiser` collect them from `block_state.denoiser_input_fields` filtered
against the transformer's forward signature (à la qwenimage/cosmos3) instead of
listing them as explicit inputs. This drops three explicit `InputParam`s in favor
of one `denoiser_input_fields` template input.

Per-pass conditioning (cond/uncond/stg/modality) stays on the guider field map --
the tag only delivers a flat dict, so it can't do the cond/uncond split or the
connector_*->encoder_hidden_states rename. The locally-computed latent dims
(num_frames/height/width/fps) are still supplied in-denoiser: they aren't upstream
outputs and their names would clash with the pixel-space values in state.

Parity unchanged (validated T2V + multi-frame I2V): the tagged values reach the
transformer identically; the first-forward capture shows them bitwise-equal.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Fix LTX-2 parity harnesses: guidance overrides reach both pipelines

The CLI guidance flags (`--guidance_scale`, etc.) only fed the modular guiders;
the standard pipeline was always called with the hardcoded GUIDANCE dict. So
`--guidance_scale 1.0` disabled CFG on the modular side only, comparing
standard-with-full-CFG (batch-2) against modular-no-CFG (batch-1) -- an
apples-to-oranges run that manifested as a huge (but spurious) I2V mismatch.

Add `_resolve_guidance(args)` (CLI overrides on top of GUIDANCE) and drive BOTH
the standard call and `_make_guiders` from the one resolved dict, in both the t2v
and i2v harnesses. With CFG correctly disabled on both sides, multi-frame I2V is
bitwise-ish (~5e-6 mean-rel), confirming the default-guidance ~8e-3 divergence is
the documented cond/uncond batch-invariance (amplified by I2V's per-token masked
timestep + clean anchor frame), not a logic bug.

Also:
- i2v: loosen the fp32 gate to (2e-2, 1.5e-1) to fit that amplified-but-numerical
  multi-frame divergence, with a comment pointing at the CFG-off run as the tight
  bug-catching gate.
- i2v: add `--debug_forward`, which diffs the first transformer forward
  (inputs + outputs) between the two runs via a forward hook -- the diagnostic
  that pinpointed the batch-shape mismatch above.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Document denoiser_input_fields standalone-run caveat for LTX-2 denoiser

`audio_num_frames` / `video_coords` / `audio_coords` reach `LTX2LoopDenoiser`
via the `denoiser_input_fields` tag, not as named inputs. Note in the docstring
that a standalone run (without the upstream tagging blocks) must pass them
through `denoiser_input_fields={...}`; plain named kwargs are silently ignored
(modular.md's kwargs_type standalone gotcha).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Address LTX-2 modular self-review: fix dead enhancer trigger + docstrings

Blocking fix: `enable_prompt_enhancement` was declared only as a
`block_trigger_inputs` entry / `select_block` param, never as an `InputParam`,
so it was not an accepted pipeline input -- `pipe(enable_prompt_enhancement=True)`
was dropped as "unexpected" and the prompt enhancer could never run. Declare it
on both enhancer sub-blocks so it reaches `select_block` (mirrors how the `image`
trigger is declared).

Also from the self-review:
- Regenerate the modular auto-docstrings: the guidance knobs (guidance_scale,
  stg_scale, ...) that moved onto the guider no longer show as block inputs, and
  the enhancer trigger now appears.
- Add descriptions to the 16 `InputParam`s that rendered as "TODO: Add
  description." in the generated docstrings (conversion checklist requires none).
- Drop the defensive `getattr(tokenizer, "padding_side", "left")` (a declared
  tokenizer always has it; gotcha #7).
- Rewrite two ephemeral comments (the connector-output "reconcile when denoise.py
  is written" NOTE, now resolved; the guider "batched-CFG variant in git history"
  pointer) into standing rationale.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Fix stale I2V parity-harness --help text for the loosened fp32 gate

The `--mean_rel_tol` / `--max_abs_tol` help strings still advertised the old
fp32 defaults (1e-3, 1e-3); the I2V gate is now (2e-2, 1.5e-1). Update the help
to match the actual DTYPE_TOLERANCES so `--help` doesn't misreport the gate.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* make style and make quality

* Add the LTX-2.4 diffusion VAE decoder with a native-parity harness

* Add a converter for the LTX-2.4 diffusion VAE decoder

* Initialise use_slicing and document the diffusion VAE decisions

* Record real-weight parity against the reference decoder

* Add model-level tests for the LTX-2.4 diffusion VAE decoder

The generated test file needed three decisions rather than fill-ins: the class had no forward() for the
mixins to call, the decoder denoises so forward has to take a generator for any output comparison to mean
anything, and MemoryTesterMixin.test_group_offloading reused one inputs_dict across four forwards without
re-seeding it. The last is fixed in the mixin, reusing the helper test_group_offloading_with_disk already
had for the same reason, now module-level and reading the signature off the class (offloading replaces
model.forward with a *args wrapper).

* Fold the diffusion VAE converter into convert_ltx2_to_diffusers.py

--diffusion_vae now sits beside --vae instead of shipping a second script. It also drops the standalone
script's --encoder-vae: the encoder is in the same checkpoint and goes through the conv VAE's own rename
rules, with its config pinned in get_ltx2_diffusion_video_vae_config. Output verified bitwise identical to
the standalone script's on the rc2 checkpoint, 491/491 tensors and the same config.json.

* Make the diffusion decoder work through LTX2Pipeline

The documented usage -- LTX2Pipeline.from_pretrained(repo, vae=diffusion_vae) -- had never been run: every
real-weight run so far called vae.decode() directly. Through the pipeline it raised on
vae.config.timestep_conditioning, which this decoder has no reason to carry, and would have bound the
positional decode timestep to the decoder's generator. The decode step now branches on the decoder type,
skipping the decode_timestep pre-noising and passing the generator instead, which also makes decoding
seed-reproducible; two full runs from one seed agree bitwise. The conv path's own statements are unchanged.

* Make the diffusion decoder work through the modular decode step

Same fix as the standard pipeline, in LTX2VaeDecoderStep: skip the decode_timestep pre-noising and pass the
generator instead of a timestep. Verified on rc2 weights by injecting the decoder into
LTX2Blocks().init_pipeline() with update_components, the route the modular parity harness already uses -- two
runs at 320x448x17 from one seed agree bitwise, so no modular_model_index.json repo was needed to run this.

* Note that the pipeline decode block predates the LTX-2.4 integration PR

* Address diffusers self-review on the diffusion VAE decoder

Four things the project's own rules catch:

- models.md forbids unconditional torch.float64 in a model (MPS/NPU/Neuron cannot run it). The RoPE
  frequency base now goes through maybe_adjust_dtype_for_device, the helper flux/flux2/wan use for exactly
  this. Measured cost of the downcast where it applies: 1.5e-08 on the frequencies, 1e-06 on the angles.
- rope_dim_split was threaded through five classes but was never a config entry, so it was always None and
  always fell back to the default split. Removed, and the default-split helper inlined into its one caller.
- FromOriginalModelMixin was declared with no SINGLE_FILE_LOADABLE_CLASSES entry, so from_single_file raised
  rather than working. Dropped; single-file support can come with its own mapping later.
- set_attention_backend() overwrites the processor's _attention_backend, which handed the neighborhood
  BlockMask to backends that cannot read it. Guarded like AnyFlowCausalAttnProcessor does.

Also runs make fix-copies, which was missing the dummy object for the new class.

* Add the API doc page for AutoencoderKLLTX2VideoDiffusionDecoder

* Mark the parity harness transient and drop the notes doc from the diff

Matches the header the other integrations/ harnesses on this branch already carry, and states that the file
imports the native reference package. The notes doc goes to the PR description instead: the self-review skill
is explicit that it never ships in the diff.

* Accept the diffusion decoder in the image-to-video pipeline too

Same branch as the text-to-video pipeline: skip the decode_timestep pre-noising and pass the generator instead
of a timestep. The remaining LTX2 pipelines that share this decode block (condition, hdr_lora, ic_lora,
latent_upsample) are deliberately left for later, so the surface stays the two pipelines this PR actually
exercises.

* Reorganize LTX-2 modular blocks into coarser component-aligned phases

Regroup the top-level LTX-2 block assemblies so each direct child of
LTX2AutoBlocks is a coherent, component-aligned phase, mirroring Flux 2's
[text_encoder, vae_encoder, denoise, decode] shape:

- New LTX2TextConditioningStep wraps [text_encoder, text_input, connectors] as
  one "text_encoder" phase (LTX-2's text conditioning is genuinely a three-step
  model chain; text_input must precede connectors for batch-invariant parity, so
  the whole chain lives here rather than pushing text_input into denoise).
- New LTX2DecoderStep wraps [video_decode, audio_decode] as one "decode" phase.
- duration stays a top-level child: it is component-aligned (duration_head) and a
  ConditionalPipelineBlocks that self-skips for concrete num_frames, so it is a
  cleanly optional phase alongside prompt_enhancer / vae_encoder -- keeping it
  out of the (mandatory) denoise block.

Top-level children become:
  t2v:  [text_encoder, duration, denoise, decode]
  i2v:  [text_encoder, duration, vae_encoder, denoise, decode]
  auto: [prompt_enhancer, text_encoder, duration, vae_encoder, denoise, decode]

Also reorder the class definitions to follow this flow, keeping the core-denoise
trio (LTX2CoreDenoiseStep, LTX2Image2VideoCoreDenoiseStep, LTX2AutoCoreDenoiseStep)
contiguous instead of split by LTX2AutoVaeEncoderStep.

Pure re-grouping: no block bodies change and intermediate outputs still propagate,
so t2v/i2v parity is bit-identical to before (user-confirmed at 17 frames).
Docstrings regenerated.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Enable LTX-2 prompt enhancement by default when the enhancer is loaded

Mirror the standard pipeline's tri-state prompt-enhancement default
(pipeline_ltx2.py:559-560), where `enable_prompt_enhancement` defaults to
auto-on when a `prompt_enhancer` component is present (LTX-2.4) and off
otherwise, and wire the enhancer into the t2v/i2v assemblies.

- Both enhancer sub-blocks: default `enable_prompt_enhancement=None` (tri-state).
  Resolve `None` against component presence inside `__call__` -- select_block
  cannot see components, so the gate dispatches for None/True and the sub-block
  decides. `None` + no enhancer -> pass the prompt through unchanged (the
  Cosmos3 optional-component pattern); explicit `True` + no enhancer -> a clear
  ValueError instead of the previous cryptic `NoneType` crash from within the
  enhance helper.
- LTX2AutoPromptEnhancerStep.select_block: skip only on an explicit `False`.
- Add LTX2AutoPromptEnhancerStep to LTX2Blocks and LTX2ImageToVideoBlocks so all
  three assemblies share the prompt_enhancer-first shape (previously only
  LTX2AutoBlocks had it; passing the flag to the others was silently ignored).
- Parity harnesses: force `enable_prompt_enhancement=False` on the modular run
  too (the modular sets now carry an enhancer), and fix the now-stale
  "contains no enhancer block" docstrings.

Only the dedicated-enhancer (LTX-2.4) path is wired; the standard pipeline's
text_encoder-as-enhancer fallback for LTX-2.0/2.3 remains a follow-up. Parity
unchanged (both pipelines run with enhancement disabled in the harnesses).
Docstrings regenerated.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Rename LTX-2.4 product identifiers to LTX-2.5.

Update conversion --version, model ids, system-prompt constant names, and docs/comments to the public 2.5 name. Behavior is unchanged in this commit.

* Add image-conditioning H.264 CRF matching Lightricks training.

Resolve CRF 18 for LTX-2.5 (Gemma 4) and 33 for earlier generations, and apply the PyAV re-compress path for I2V and single-frame Condition/IC keyframes.

* Align LTX-2 pipelines with the Lightricks reference behavior.

Extract shared prompt-enhancement and check_inputs mixins; enable Condition duration+enhancement and IC enhancement (HDR remains encode-only); match enhancement recipe and LTX-2.3/2.5 sampling defaults.

* Update LTX-2 docs for 2.5 naming, defaults, and opt-in enhancement.

Document LTX_2_3/2.5 sampling defaults, STG block 28, and enable_prompt_enhancement=True as the opt-in path.

* Replace LTX2AutoDuration with min_seconds/max_seconds call args.

Drop the dataclass wrapper and expose duration bounds directly on the T2V, I2V, and Condition pipelines.

* Align LTX-2 modular blocks with the merged standard-pipeline changes

The standard LTX-2 pipelines moved to shared `LTX2PromptEnhancementMixin` /
`LTX2CheckInputsMixin`, dropped `LTX2AutoDuration`, adopted the Lightricks
reference defaults, and gained H.264 CRF re-compression of image
conditionings. Port all of that to the modular blocks.

- Prompt enhancement: re-derive `_enhance_prompt` from the new mixin -- the
  `user prompt:` / `User Raw Input Prompt:` templates, 896px long-side image
  prep, left-padding to a multiple of 8 for Flash Attention, `clean_response`
  on the decoded output, and `max_new_tokens=None` resolving to the Gemma-4
  budget of 600.
- Revert `enable_prompt_enhancement` to a plain `False` opt-in, matching the
  reference pipelines (the tri-state auto-on default is gone upstream). The
  friendly error for `True` with no enhancer loaded is kept.
- Auto-duration: trigger on an omitted `num_frames` instead of an
  `LTX2AutoDuration` sentinel, with `min_seconds`/`max_seconds` as call args.
  `LTX2DurationStep` raises when no `duration_head` is loaded rather than
  silently falling back to 121 frames.
- Every `num_frames` InputParam default becomes `None`: input defaults are
  seeded into the pipeline state before any block runs, so a default of 121
  would leave the auto-duration branch permanently unreachable.
- Defaults sweep: `num_inference_steps` 40 -> 30, `use_cross_timestep`
  False -> True, and the video/audio guider ComponentSpec configs updated to
  the reference guidance stack.
- `LTX2VaeEncoderStep` gains `image_crf` and re-compresses the conditioning
  image before preprocess, declaring `text_encoder` only to resolve the model
  default CRF.
- Rename LTX-2.4 -> LTX-2.5 throughout, following the upstream product rename.

Parity harnesses updated for the new duration arguments, plus an `--image_crf`
flag on the i2v harness defaulting to 0 so the comparison isolates block logic
from the PyAV codec round-trip. Verified against a real checkpoint: t2v, i2v at
`--image_crf 0`, and i2v at `--image_crf -1` all pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Evaluate the decoder's SwiGLU in token tiles

Decode-only peak VRAM at 121 frames and 512x768 was 24.6 GiB, and 17.0 GiB of that was one transient
inside the first stage-5 block: `w_gate(x)` and `w_up(x)` are both hidden-width and their product makes
a third, so the whole video's worth of hidden activations was live three times over. At that size each
one is 2,973,696 tokens x 1024 x 2 bytes = 5.67 GiB.

The MLP is pointwise across tokens, so tiling it is exact -- it changes only how many hidden-width
elements exist at once. Measured on a full-size decode: peak 24.8 -> 19.4 GiB (-22%), wall clock 2.79 ->
2.84 s (+1.6%), output bit-identical (max abs diff 0.000e+00). That also puts the diffusion decoder
below the conv decoder, which needs 23.8 GiB at the same size.

Tile size 16384 tokens is the reference decoder's own default (`DEFAULT_SWIGLU_TILE_SIZE`). Bounding a
token count rather than a number of tiles keeps peak memory independent of resolution.

This is the first of the two transients the reference's default decode mode avoids; the other is the
full-volume QKV projection plus RoPE's fp32 upcast, which it addresses by chunking attention along
width. That one is not bit-exact (it reassociates NATTEN's accumulation), so it is left for a follow-up
rather than changed here.

* Declare _no_split_modules on the diffusion decoder

Both block types close over a residual add that combines outputs from different children, which is the
case models.md says has to stay co-located, so neither can be split across devices.

Unskips three memory tests that were gated on the attribute: the file goes from 35 passed / 7 skipped to
38 passed / 4 skipped, and the three (test_cpu_offload, test_disk_offload_with_safetensors,
test_disk_offload_without_safetensors) pass.

* Match the encoder defaults to the published LTX-2.4 config

block_out_channels was (256, 512, 1024, 2048) and layers_per_block was (4, 6, 6, 2, 2); both were
inherited from the 2.0 conv VAE. The published vae_diffusion/config.json has (256, 512, 1024, 1024) and
(4, 6, 4, 2, 2). Real loads were unaffected because the config sets these explicitly, but the defaults
described a model that does not exist.

All 27 __init__ defaults that appear in the published config now match it.

* Cover the token-tiled SwiGLU path in the tests

The dummy video is 9x48x48, so its stage-5 grid is 5184 tokens against a 16384-token tile size: every
other test in the file takes the untiled branch and the tiling loop ran nowhere in CI. This shrinks the
tile size to force ~41 tiles and requires torch.equal against the untiled evaluation, since the MLP is
pointwise across tokens and tiling must be exact rather than close.

Checked the test is not vacuous: sabotaging the tiled branch alone (silu on w_up instead of w_gate)
fails it at 8.83e-01.

* Let the parity harness run against the LTX-2.5 reference too

The harness only spoke LTX-2.4's decoder API and died on `forward_pre_diffusion` against 2.5, so it could
not check the release the port is now targeted at. Three shims absorb the differences:

* 2.5 selects a decode pathway by swapping `__class__` on every block, so an as-constructed decoder has no
  `forward_combined`. Install the COMBINED pathway explicitly, uncompiled, since that is the pathway this
  port implements.
* stages 1-4 split into `forward_stages_1_to_3` + `forward_stage_4`, so the chunked pathway can defer
  stage 4.
* the trailing NATTEN ghost pad moved out to the caller while its crop stayed behind
  `forward_stage_4(pad_trailing=True)`; the two have to be paired or the context loses real frames (3 of
  17 survive at this config, which is how the mismatch first showed up).

Same numbers against both references -- context 5.367e-04, step 7.337e-04 -- which is the check that the
shims are semantically right and not just shape-compatible.

* Add LTX-2.X modular condition blocks

Ports LTX2ConditionPipeline to the modular pipeline system, including the
LTX-2.5 additions (per-condition CRF, conditions-sourced prompt enhancement,
auto-duration, the Lightricks reference defaults).

New blocks:
  - encoders.py: LTX2ConditionEncoderStep (resize/center-crop, single-frame
    H.264 re-compression at the model CRF, temporal trim, VAE encode) and
    LTX2ConditionPromptEnhancerStep (grounds the rewrite in the first PIL
    frame found in `conditions`, text-only otherwise).
  - before_denoise.py: LTX2ConditionPrepareLatentsStep (first-frame overwrite
    + keyframe token append with their own RoPE coords, emitting
    conditioning_mask / clean_latents / appended_coords / base_token_count),
    LTX2ConditionSetTimestepsStep, LTX2ConditionPrepareAudioLatentsStep,
    LTX2ConditionPrepareCoordsStep, and the _prepare_keyframe_coords helper.
  - denoise.py: LTX2ConditionLoopBeforeDenoiser (mask-scaled per-token video
    timestep), LTX2ConditionLoopAfterDenoiser (x0 blend against the clean
    condition latents, then a full-sequence scheduler step),
    LTX2ConditionDenoiseStep.
  - decoders.py: LTX2TrimConditionTokensStep.
  - modular_blocks_ltx2.py: LTX2ConditionAutoPromptEnhancerStep,
    LTX2ConditionCoreDenoiseStep, LTX2ConditionDecoderStep, LTX2ConditionBlocks.

LTX2LoopDenoiser is reused unchanged: the two LTX2Guidance guiders already
express the whole CFG + STG + modality-isolation stack the condition pipeline
uses, and the transformer takes the longer packed sequence with the base
num_frames/height/width exactly as the standard pipeline does. The standard
pipeline's CFG batch duplication (conditioning mask, coords) has no modular
counterpart, since each guidance pass is its own single-batch forward.

Block ordering: prepare-latents runs *before* set-timesteps, unlike the
text-to-video and image-to-video core denoise steps. The resolution-aware
shift `mu` is computed from the packed latent sequence length, which here
includes the appended keyframe tokens, so the latents must exist first. This
mirrors LTX2ConditionPipeline (its section 4 precedes section 5). The
text-to-video and image-to-video blocksets are unaffected: with no appended
tokens the grid-derived sequence length already equals latents.shape[1].

Two subtleties the parity harness surfaced:

  - The condition pipeline samples audio noise directly in packed shape
    [B, L, C * M], while the text-to-video and image-to-video pipelines sample
    unpacked [B, C, L, M] and pack afterwards. Both draw the same number of
    values from the generator but lay them out differently, so reusing the
    shared LTX2PrepareAudioLatentsStep silently desynchronizes the audio noise
    and, through the joint attention, the video too (1.4 mean-rel on audio,
    0.13 on video before the fix). Hence a standalone condition audio block.

  - `noise_scale` means different things across workflows: 0.0 for
    text-to-video/image-to-video, `None -> sigmas[0] or 1.0` for conditions.
    A blockset keeps the first non-None default across its blocks, so a 0.0
    would have shadowed the condition resolution and left the latents
    unnoised. Both condition-path blocks declare None; prepare-latents
    resolves once and writes the value back for the audio step.

Keyframe tokens, masks and coords are expanded to the generation batch (the
pattern the in-context pipeline already uses for reference tokens). The
standard pipeline builds them at batch 1 and cats them onto a batch-B
sequence, so it raises for num_videos_per_prompt > 1 with any condition at
latent index > 0; the modular blocks handle that case.

Verification: integrations/ltx2_condition_parity.py (transient, removed before
merge) compares LTX2ConditionBlocks against LTX2ConditionPipeline on shared
component objects. Eight cases on a tiny checkpoint -- first-frame only,
keyframe only, both plus a multi-frame video condition, two keyframes, no
conditions, --crf -1, --predict_duration, num_videos_per_prompt=2 -- all pass
at 3e-7 to 1.3e-6 mean-rel. On a full checkpoint at 768x512x17 with two
conditions through the CRF path: video 8.7e-04 mean-rel / 1.1e-02 max-abs,
audio 1.5e-04 / 2.1e-03, in line with the image-to-video figures at the same
frame count.

make style, make quality and make fix-copies are clean; auto-docstrings
regenerated with no TODO placeholders.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Add LTX-2.X modular in-context (IC-LoRA) blocks

Ports LTX2InContextPipeline to the modular pipeline system on top of the
condition blocks, and fixes two bugs in the standard in-context pipeline that
the parity harness surfaced.

New blocks:
  - encoders.py: LTX2ReferenceEncoderStep (preprocess, VAE-encode and pack each
    reference video, compute the coords that map its tokens into the target
    space, and build the per-token cross-attention strengths), plus the
    _downsample_mask_to_latent helper.
  - before_denoise.py: LTX2InContextPrepareLatentsStep (frame conditions as in
    the condition step, then the reference tokens appended last) and
    LTX2BuildVideoSelfAttentionMaskStep.
  - modular_blocks_ltx2.py: LTX2AutoReferenceEncoderStep,
    LTX2AutoBuildVideoSelfAttentionMaskStep, LTX2InContextCoreDenoiseStep,
    LTX2InContextBlocks.

LTX2ConditionDenoiseStep is reused unchanged: with the blend gate fixed below,
reference tokens are pinned by exactly the same x0 blend as frame conditions,
matching the reference implementation where VideoConditionByReferenceLatent and
VideoConditionByKeyframeIndex produce the same (denoise_mask, clean_latent) pair
and post_process_latent runs unconditionally. The condition encoder, timesteps,
audio-latents, coords and decoder blocks are shared as well.
video_self_attention_mask reaches the transformer through the
denoiser_input_fields tag, so LTX2LoopDenoiser needed no change.

`reference_conditions` is optional, matching LTX2InContextPipeline: IC-LoRAs
that carry their behavior in the adapter weights (camera control, style) take
no reference video, which is the shape of the standard pipeline's own docstring
example. The reference encoder and the attention-mask step are both skipped when
it is absent. There is no duration step (the in-context pipeline ships without a
duration_head), so LTX2ConditionEncoderStep now raises a concise error if
`num_frames` is still None by the time the conditions are encoded.

Per-reference token counts come from the encoder rather than an equal split of
the total, so references of differing lengths (a reference video shorter than
num_frames) get the right strengths.

Standard-pipeline fixes in pipeline_ltx2_ic_lora.py:

  - The x0 blend was gated on `has_conditions` alone, so with reference
    conditions and no frame conditions the reference tokens were seeded clean
    and given a zeroed timestep but never re-pinned. Their x0 prediction error,
    divided by sigma each step, then accumulated and drifted them away from the
    encoded reference. Now gated on `has_conditions or num_ref_tokens > 0`, in
    line with the surrounding gates and the reference implementation.

  - video_self_attention_mask was rebound in the loop rather than expanded into
    a per-pass local. The main pass expanded it to the CFG-doubled batch, after
    which the STG pass's expand() to `latents.shape[0]` raised
    "RuntimeError: The expanded size of the tensor (1) must match the existing
    size (2)". Any run combining an attention mask with STG or modality
    guidance -- i.e. the default guidance settings -- died on the first step.

Also removed the unused has_appended_tokens local and replaced the "# - TODO"
placeholder above the attention-mask construction with the block structure,
which the reference confirms.

Verification: integrations/ltx2_in_context_parity.py (transient, removed before
merge). Twelve cases on a tiny checkpoint -- reference only, reference at
partial strength, scalar attention strength, pixel-space attention mask,
downscale factor 2, references plus first-frame and keyframe conditions, two
references, the CRF path, and four reference-free variants -- all pass at 3.4e-7
to 1.3e-6 mean-rel. The seven-case condition suite still passes. On a full
checkpoint at 768x512x17 with a reference plus two conditions: video 4.4e-05
mean-rel / 5.6e-04 max-abs, audio 3.2e-05 / 2.4e-04.

Known limitation, not addressed here: transformer_ltx2.py converts the
multiplicative [0, 1] self-attention mask to an additive bias linearly, as
(1 - mask) * -10000, whereas the reference uses log space (bias = log(mask),
finfo.min at zero). Intermediate strengths therefore saturate to "fully masked",
making conditioning_attention_strength effectively binary and spatially varying
masks inert beyond their zero/non-zero pattern. Parity is unaffected -- both
pipelines feed the same mask into the same transformer -- but the feature does
not currently work as documented.

make style, make quality and make fix-copies are clean; auto-docstrings
regenerated with no TODO placeholders.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Make the diffusion decoder a decode pipeline instead of a VAE

Addresses the review: a full diffusion model should not be an AutoencoderKL that hides its denoising
behind `.decode()`, and it should not have to be threaded into every LTX-2 pipeline as a `vae`.

AutoencoderKLLTX2VideoDiffusionDecoder becomes LTX2VideoDiffusionDecoderModel, a plain ModelMixin with
no encoder. Encoding stays with AutoencoderKLLTX2Video, whose latent space this consumes unchanged, so
latents remain interchangeable between the two decoders. Dropping the encoder also removes the
converter's awkwardness of sourcing encoder weights and the real block_out_channels from the conv
`vae/` folder: ten encoder-only config keys are gone.

LTX2VideoDiffusionDecodePipeline drives it, shaped like LTX2LatentUpsamplePipeline: run any LTX-2
pipeline with output_type="latent", then decode. `vae` is an optional component, since it is only ever
consulted for the latent statistics, and the decoder carries its own `latents_mean` / `latents_std`
buffers to fall back on. That means a decode-only workflow never loads a second autoencoder.

The `vae=` path is removed accordingly: the isinstance branches in LTX2Pipeline,
LTX2ImageToVideoPipeline and the modular decode step are gone, along with the fixtures that built a
diffusion decoder to pass as `vae`. The model's slicing/tiling test class went with them, since
enable_tiling was an AutoencoderMixin concern and this is no longer an autoencoder.

Tests: 37 passed / 3 skipped on the model, 83 passed / 2 skipped across the t2v and i2v pipelines.

Still to come: mapping the 2.5 checkpoint to a blockset that selects the diffusion decode block, and
the KandinskyCombinedPipeline pattern for running the two-step flow by default in the standard API.

* Add a modular decode block for the LTX-2 diffusion decoder

The previous commit removed the `isinstance(vae, ...)` branch that let the diffusion decoder stand in
as the modular `vae`, which was the right call twice over: it is what the review objected to, and
modular.md says a new case gets a new block rather than a branch inside an existing one. But removing
it without a replacement left modular with no route to the diffusion decoder at all.

LTX2DiffusionVaeDecoderStep is that replacement. It declares the decoder as its own
`diffusion_decoder` component rather than borrowing the `vae` slot, and passes a generator instead of
a decode timestep, since this decoder draws the noise it denoises. `LTX2VaeDecoderStep` is untouched,
so the convolutional path is unaffected.

What is still missing is the wiring: a blockset using this block, and an `_ltx2_map_fn` so a
checkpoint that ships the diffusion decoder selects it automatically instead of the user choosing.
That lands in modular_blocks_ltx2.py, which #12 is reorganizing, so it needs sequencing against that
PR rather than racing it.

* Test the diffusion decode pipeline, and give it a video-only output

The decode pipeline had no coverage: it was never actually run. Writing the tests immediately turned
up a bug, which is that it returned `LTX2PipelineOutput`, whose `audio` field is required. This
pipeline decodes video only, so it now returns a video-only `LTX2VideoDecodeOutput` instead of
passing a meaningless `audio=None`.

Four tests: decode with no `vae` (the fallback to the decoder's own latent statistics), decode with a
`vae` supplied (its statistics take precedence, checked by making the two disagree rather than
asserting a shape), reproducibility under a seeded generator plus divergence under a different seed,
and `denormalize=False` actually skipping the statistics. The decoder's stat buffers are set to
non-trivial values in the fixture so a run that silently skipped denormalization could not pass by
coincidence.

* Fix the decode example: output_type="latent" is already denormalized

The example in the docs and the pipeline docstring both showed `output_type="latent"` feeding straight
into the decode pipeline on its defaults. That path denormalizes twice: `LTX2Pipeline` applies the
latent statistics before returning latents (pipeline_ltx2.py), and the decode pipeline applies them
again because `denormalize` defaults to True. Every channel ends up scaled by its own std a second
time, so e.g. channel 0 (std 0.238) comes out roughly 4x too small.

Passing `denormalize=False` in the example and saying so in both docstrings. Caught by running the
flow end to end on real rc4 weights, which is also where the two decoders agree at 0.9978 once the
statistics are applied only once.

Whether the *default* should flip, or `output_type="latent"` should stop denormalizing instead, is a
contract question across both pipelines and is still open.

* Fix IC-LoRA video self-attention mask grouping

The video self-attention mask deviated from the reference implementation
(`ltx_core.conditioning.mask_utils.build_attention_mask`) in two ways. Both
predate the modular blocks and applied to `LTX2InContextPipeline` as well; the
modular block had ported the same construction.

1. Keyframe conditions cross-attended with reference tokens. `num_noisy_tokens`
   was derived as `latents.shape[1] - num_ref_tokens`, which folds the appended
   keyframe tokens into the noisy group, so they received the reference cross
   mask. The reference takes `num_noisy_tokens` from the target latent shape --
   generated-video tokens only -- and tracks the extras offset separately, so
   keyframe <-> reference is 0.0. Split the two quantities apart:
   `_build_video_self_attention_mask` now takes `num_prefix_tokens` alongside
   `num_noisy_tokens`.

2. Multiple reference conditions fully attended to each other. All references
   were passed as a single attention group with a 1.0 self-block. The reference
   wraps each `LTX2ReferenceCondition` in its own
   `ConditioningItemAttentionStrengthWrapper`, building the mask once per item,
   so reference_i <-> reference_j is 0.0. The call site now splits
   `ref_cross_mask` into one group per reference before building the mask.

Threading per-reference token counts out of `_encode_reference_conditions` also
removes the equal-split assumption in `prepare_latents` and
`prepare_reference_latents`, which assigned the wrong per-token strengths when
references encoded to different token counts (reachable when a reference video
is shorter than `num_frames`).

API impact
----------

`prepare_latents` returns a 7-tuple instead of a 6-tuple, adding
`ref_token_counts` (`list[int]`) as the new final element. Elements 1-6 keep the
types they have on `main`; in particular `ref_cross_mask` remains the
concatenated `[1, num_ref_tokens]` tensor. Callers that unpack the tuple fail
immediately with a `ValueError` rather than silently receiving a changed type.

`_encode_reference_conditions` (private) likewise gains `reference_token_counts`
as a fourth return element. `prepare_reference_latents` keeps its documented
4-tuple unchanged.

Verification
------------

Correctness against the reference. Both the standard helper and the modular
block were checked for exact tensor equality (`torch.equal`) against
`build_attention_mask`, driven the way `ic_lora.py` drives it -- one
`ConditioningItemAttentionStrengthWrapper` per reference, so the reference
builds iteratively with `existing_mask` carried forward. Five layouts, all
matching, exercised through the same concatenate-then-split path the call site
uses:

  - references only, 1 reference
  - keyframes + 1 reference          (exercises fix 1)
  - keyframes + 2 references         (exercises fixes 1 and 2)
  - 2 references of unequal length   (exercises fix 2 and the token counts)
  - no keyframes, 3 references       (exercises fix 2)

Standard/modular parity on a real checkpoint, 512x512x17, fp32, via
`integrations/ltx2_in_context_parity.py`. Video latents (1, 128, 3, 16, 16),
audio latents (1, 8, 18, 16); gates are mean-rel 2e-2 and max-abs 1.5e-1:

  --reference 1.0 9 --condition 2 0.5 9 --conditioning_attention_strength 0.5
    video: max abs 4.822e-05, mean rel 7.851e-06
    audio: max abs 1.875e-04, mean rel 1.773e-05

  --reference 1.0 9 --reference 0.8 9 --conditioning_attention_strength 0.5
    video: max abs 1.953e-04, mean rel 9.101e-05
    audio: max abs 9.251e-05, mean rel 1.026e-05

The first case is the layout where `num_prefix_tokens != num_noisy_tokens`; a
wrong prefix/base split would have placed the group blocks at the wrong offsets
with valid shapes, surfacing as a parity mismatch. The returned video shape also
unpacks from exactly 768 base tokens (3 x 16 x 16), confirming
`base_token_count` is the quantity now fed as `num_noisy_tokens`.

Scope of the parity runs: both sides share this construction, so parity shows
they did not desynchronize, not that either matches the reference -- that rests
on the equality test above. Both references in the second case are 9 frames and
so encode to equal token counts, meaning the equal-split and per-count paths
coincide; the token-count change is not discriminated by these runs. A reference
pair of differing lengths (e.g. `--reference 1.0 9 --reference 0.8 17`, giving
512 vs 768 tokens) would separate them, and would have failed parity before this
commit since the modular block already used per-reference counts while the
standard pipeline split evenly.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Add LTX25AutoBlocks and LTX25ModularPipeline

Following the review: LTX-2.5 is a new checkpoint, so it gets its own blockset and pipeline rather
than a config flag inside the LTX-2 ones.

LTX25AutoBlocks subclasses LTX2AutoBlocks and overrides exactly one entry, `video_decode`, to
LTX2DiffusionVaeDecoderStep. Every other leaf block is reused, so the two blocksets have identical
block_names and cannot drift apart. That makes the diffusion decoder the default for LTX-2.5, which
matches what the checkpoint ships natively, without the user having to pass a decoder at all.

Going the other way stays a one-line swap rather than a second official blockset:

    blocks.sub_blocks["video_decode"] = LTX2VaeDecoderStep()

LTX25ModularPipeline carries `default_blocks_name = "LTX25AutoBlocks"` and a new "ltx2.5" key in
MODULAR_PIPELINE_MAPPING, so a repo shipping `modular_model_index.json` routes straight to the new
blockset.

The condition and in-context workflows belong in this blockset too, but those blocks live in the
modular condition/IC PR, so folding them in belongs there rather than as a competing change here.

* Drop the encoder half in the diffusion decoder converter

The converter still assembled the whole VAE checkpoint, encoder included, and handed it to a strict
`load_state_dict`. That was correct while the class was an autoencoder; after it became decoder-only
the ~84 `encoder.*` keys have nowhere to land and the load raises. Nothing caught it because no test
converts anything and every other check loads pre-converted weights, where unexpected keys are only
warned about.

Now the encoder is discarded rather than remapped, keeping the `decoder.` half and the per-channel
statistics that become `latents_mean` / `latents_std`. `strict=True` stays, since it is what guards
against a botched rename rule.

Verified by actually running it: pulled the 395 `vae.*` tensors out of the native rc4 checkpoint by
byte range, converted, and diffed against the already-published folder. 407/407 parameters, max abs
difference 0.000e+00, so the fix changes nothing about the weights it produces.

Also added a regression test, since the absence of any converter coverage is why this slipped.

* Share the neighborhood attention mask and parameterise the AdaLN chunk count

Two of the resolved review threads. Both are internal to the decoder modules and are unaffected by
the pipeline-vs-model question still open on the thread above them.

Build the FlexAttention mask once per stage instead of once per block. The mask depends only on the
grid and the kernel, both fixed within a stage, so 24 builds collapse to 5; measured at 768x512x121
that is 582ms to 132ms. build_block_mask returns None for NATTEN, which encodes the window in its
kernel. That is not merely an optimisation: create_block_mask materialises O(N^2) uncompiled, so a
69x64x96 stage would ask for 167 GiB and stage 5 for 8.2 TiB, and the mask must never be built for a
path that would not read it.

Make the AdaLN chunk count a constructor argument, num_chunks on the shared AdaLN and num_mod_params
on the diffusion block, with the decoder passing shared_adaln.num_chunks through so the two cannot
drift. num_mod_params is already the name LTX2AdaLayerNormSingle uses. The block's forward still
destructures exactly seven chunks and reads four of them, which is the reference's shape.

Decoding the test fixture is bitwise identical to before (max abs diff 0.000e+00). 39 passed, 4
skipped.

* Export the LTX-2.5 modular classes from the top-level namespace

`ModularPipeline.from_pretrained` resolves a repo's `_class_name` with
`getattr(importlib.import_module("diffusers"), class_name)`, so a class that is only reachable from
`diffusers.modular_pipelines.ltx2` cannot be selected by a checkpoint. `LTX25ModularPipeline` and
`LTX25AutoBlocks` were exported there but not from `diffusers` itself, which left the whole
checkpoint-to-blockset mapping dead: loading a repo whose index names them raised
`AttributeError: module diffusers has no attribute LTX25ModularPipeline`.

Exported alongside `LTX2ModularPipeline` / `LTX2AutoBlocks`, which is how the loader already finds
those. Verified against a real repo: `ModularPipeline.from_pretrained(...)` now resolves
LTX25ModularPipeline -> LTX25AutoBlocks -> LTX2DiffusionVaeDecoderStep, with the decoder supplied by
the repo index and nothing passed by hand.

Found by loading an actual checkpoint rather than constructing the blocks directly; calling
`LTX25AutoBlocks().init_pipeline(repo)` goes blocks-to-pipeline and never exercises the lookup.

* Name the diffusion decoder LTX-2.5 and save it to diffusion_decoder/

Two naming fixes, one of them load-bearing.

The converter wrote the decoder to a `vae_diffusion/` subfolder, a leftover from when the class was an
autoencoder. The folder name matters now: `DiffusionPipeline.from_pretrained` resolves each component
from the subfolder named after it, so `LTX2VideoDiffusionDecodePipeline` (component name
`diffusion_decoder`, same as the modular `ComponentSpec`) can never load itself out of a repo whose
folder is called `vae_diffusion`. The converter now saves to `diffusion_decoder/`, matching the layout
the LTX-2.5 weight repos already use, and the docs example loads from that name.

The docstrings also called it the LTX-2.4 decoder. The diffusion decoder ships with LTX-2.5; the
converter handles those checkpoints under the `--version 2.4` architecture family, which the version
guard's message now spells out instead of misattributing the feature.

Verified by re-running the converter on the rc4 distilled checkpoint after the change: the output lands
in `diffusion_decoder/` and is bit-identical (sha256) to the pre-change conversion, so this renames and
redocuments without touching what is produced. The converter regression test passes.

* Move the LTX-2 text-input expansion into the core denoise step

`LTX2TextInputStep` expanded the prompt embeddings to the effective batch
size inside `LTX2TextConditioningStep`, so running the text-conditioning
stage on its own baked `num_videos_per_prompt` into its outputs. Feeding
them to `LTX2CoreDenoiseStep` with a different `num_videos_per_prompt`
meant re-running the text stage or reshaping its outputs by hand.

`LTX2TextInputStep` now expands the connector conditioning (the tensors
the denoiser actually reads through `guider_input_fields`) instead of the
raw Gemma embeds, and runs as the first block of both core denoise steps.
The raw `prompt_embeds` expansion is dropped -- it only existed to feed
the connector at the expanded batch. `batch_size` / `dtype` move up to
`LTX2TextEncoderStep`, since `LTX2DurationStep` needs `batch_size` before
the denoise stage.

Two behavior notes:
- The connectors now run at one row per prompt rather than
  `num_videos_per_prompt` rows. They are applied per batch element, so
  this is mathematically identical (verified to ~3.6e-7, float32
  rounding) and N times cheaper, but not bitwise identical to
  `LTX2Pipeline` when `num_videos_per_prompt > 1`.
- Expansion uses `repeat_interleave` for the masks as well as the embeds.
  The previous code (inherited from the standard pipeline) interleaved
  the embeds but tiled the masks, which misaligned them whenever both
  `batch_size > 1` and `num_videos_per_prompt > 1`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Describe the connector inputs on LTX2TextInputStep

`LTX2TextInputStep` declares the connector conditioning tensors and now
runs first in the denoise sequence, so its declarations -- not
`LTX2LoopDenoiser`'s -- are the ones the docstring generator picks up.
They carried no `description`, which left `TODO: Add description.`
placeholders on the three core denoise blocks.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Build the condition and in-context workflows from LTX2AutoBlocks

`LTX2AutoBlocks` covered text-to-video and image-to-video only; the
condition and in-context (IC-LoRA) workflows were reachable solely
through their own blocksets. Fold them into the auto blockset so all
four are selected from the inputs, the way t2v and i2v already were.

- Merge `LTX2ConditionAutoPromptEnhancerStep` into
  `LTX2AutoPromptEnhancerStep`, which now picks the condition, image or
  text enhancer off `conditions` / `image`.
- Add `LTX2AutoConditionEncoderStep`, triggered by `conditions` or
  `reference_conditions` so a reference-only in-context run still gets
  the empty per-condition lists `LTX2InContextPrepareLatentsStep`
  requires.
- Extend `LTX2AutoCoreDenoiseStep` to four branches and add
  `LTX2AutoDecoderStep`, which trims the appended keyframe/reference
  tokens whenever `base_token_count` is present.

Two fixes this exposed:

- `noise_scale` means an interpolation factor (t2v/i2v, 0.0) in one
  workflow and an initial noise level (condition/in-context, `None` ->
  `sigmas[0]` or 1.0) in another. A blockset keeps the first non-`None`
  default across its blocks, so sharing one blockset published 0.0 and
  silently zeroed the noise for condition workflows that omitted the
  argument. `LTX2PrepareLatentsStep` and `LTX2PrepareAudioLatentsStep`
  now declare `None` and resolve to 0.0 in the block, mirroring what the
  condition steps already do.
- `LTX2AutoBuildVideoSelfAttentionMaskStep.select_block` compared
  `conditioning_attention_strength` against 1.0, but
  `get_execution_blocks` passes unset trigger inputs as an explicit
  `None`, so resolving any in-context workflow raised `TypeError`.

The parity harnesses passed `noise_scale=None` explicitly, which took
the argument-present path and hid the shadowing bug; `--omit_noise_scale`
drops the key entirely so the default is exercised.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Run the text-input step in the condition and in-context denoise steps

`LTX2TextInputStep` is the only block that expands the connector text
conditioning by `num_videos_per_prompt`, and it was missing from
`LTX2ConditionCoreDenoiseStep` and `LTX2InContextCoreDenoiseStep`. The
prepare-latents/audio/coords steps all build their tensors at
`batch_size * num_videos_per_prompt`, so the video latents arrived at the
transformer at batch B*N while `encoder_hidden_states` was still at B —
a shape error in cross-attention for any `num_videos_per_prompt > 1`.

The step goes at the head of both sequences, before prepare-latents, so
the prepare-latents-before-set-timesteps ordering these workflows need
(`mu` is read off the full packed sequence length) is untouched.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Broadcast keyframe and reference conditioning across the generation batch

Frame conditions are VAE-encoded at batch 1, but `latents`,
`conditioning_mask` and `clean_latents` are built at
`batch_size * num_videos_per_prompt`. Appending a keyframe condition
(latent index > 0) therefore failed in `prepare_latents` with
"Sizes of tensors must match except in dimension 1" whenever the batch
was greater than 1. First-frame conditions were unaffected because they
are applied by indexed assignment, which broadcasts.

Expand the keyframe tokens, mask and coords to `batch_size` before the
concatenation, and do the same for the IC-LoRA reference coords, which
feed the same `appended_coords` block that is later concatenated onto
batch-`batch_size` `video_coords`. This is what the modular LTX-2 blocks
already do, and matches the `[B, 3, ...]` shape both docstrings promise.

Also picks up a pending `make style` docstring rewrap in the same file.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* feat: use natten for processing. (#15)

* feat: use natten for processing.

* address dhruv's comments.

* remove unneeded test

* remove integration file

* make vae properties accessible without loading it. (#17)

* Address review: replace LTX2 mixins with Copied-from pipeline methods.

Inline check_inputs/enhance_prompt on LTX2Pipeline and sync I2V/Condition/InContext via # Copied from, and update docs for shared 2.5 defaults.

(cherry picked from commit c183923cdd9c951f9694f9f6cd903714330cff03)

* Point the modular enhance-prompt comment at LTX2Pipeline

`LTX2PromptEnhancementMixin` was inlined into the pipelines in c183923cd,
so the mirrored-from reference no longer names an existing class.

* Remove the integrations/ dev scripts

The LTX-2 parity harnesses and the tiny-checkpoint builder were local
development scaffolding and are not part of the shipped library.

* moular tests (#19)

* Keep LTX-2.5.1 keyframes_abs_pos_embedding through convert. (#20)

Distilled 2.5.1 monoliths ship this (1, inner_dim) buffer with
use_keyframes_abs_pos_embedding=True; store it on the transformer for
load/save without wiring it into the regular distilled forward yet.

* remove test_ltx2_keyframes_abs_pos_embedding_load_and_forward

* Update the LTX-2.5 docs: two stage generation, distilled default, repo id (#21)

* feat: implement tiling for the diffusion decoder. (#22)

* add a note on tiling (#23)

* fix documentation.

---------

Co-authored-by: aartyomov <aartyomov@lightricks.com>
Co-authored-by: Daniel Gu <daniel.gu@huggingface.co>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: linoytsaban <linoy.tsaban@gmail.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: DN6 <dhruv.nair@gmail.com>
Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
2026-08-11 07:10:00 -10:00
Sayak Paul
50e7158093 [tests] fix torchao tests (#14258)
fix torchao tests

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-08-07 20:10:08 +08:00
Sayak Paul
8a29f3139a [tests] standardize model-level quant tests (#14332)
standardize model-level quant tests
2026-08-01 07:53:56 +07:00
Álvaro Somoza
8b33b47332 [quantization] SDNQ core loading (#14277)
* initial

* fix code quality

* dependency test fix

* apply review suggestions

* add extras and version check

* fix gating and prefer local cache for modular

* simplify transformers registration and make the version warning silenceable

* remove pipeline level test

* only model level testing

* layers to quantizze, change em dashes

* env flag for transformers

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-07-30 13:09:11 +05:30
Sayak Paul
3c468926ff [tests] unify keep_in_32 tests in quantization test suites (#14257)
unify keep_in_32 tests in quantization test suites
2026-07-30 07:50:09 +05:30
jiqing-feng
52110fbbf6 Fix LoRA hot-swapping recompilation with different_shapes_for_compilation (#14297)
[tests] fix hot-swapping recompilation with different_shapes_for_compilation

`test_hotswapping_compiled_model_linear` and
`test_hotswapping_compiled_model_both_linear_and_other` failed for every model
setting `different_shapes_for_compilation`, e.g.:

    RecompileError: tensor 'hidden_states' size mismatch at index 1.
    expected 16, actual 32

The dummy inputs are `(batch, height * width, channels)` with `channels = 16`
and the first traced shape is `(4, 4)`, so the image sequence length equals the
channel count. Duck shaping assigns both dims the same symbol, which `img_in`
(an `nn.Linear` with constant `in_features`) then specializes to 16, forcing a
recompile on the next shape.

`use_duck_shape = False` was already set in `test_compile_on_different_shapes`
and `test_hotswapping_compile_on_different_shapes` (#11327) but the other two
multi-shape hot-swapping tests were missed. Move it into
`_check_model_hotswap` so all of them are covered, and drop the now-redundant
duplicate.

Also drop the two `xfail(strict=True)` markers on
`TestQwenImageTransformerLoRAHotSwap` that were masking this bug.

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-07-29 15:16:41 +05:30
Steven Liu
23b0004330 torch.dtype alias (#14162)
* alias

* add tests

* feedback

* Fix failing test_dtype_alias pipeline tests

Cosmos pipelines keep the safety_checker in float32 and require it to be
passed explicitly on reload, mirroring their test_torch_dtype_dict
overrides. PRX skips dtype-dict/save-load tests because the custom
T5GemmaEncoder is not compatible with transformers v5; do the same for
test_dtype_alias.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-14 10:08:01 -07:00
Vinh H. Pham
60ec6f7242 Add Nunchaku Lite single-file quantization (#14100)
* Add Nunchaku Lite single-file quantization

* Support config-backed Nunchaku Lite loading

* Remove Nunchaku runtime manifest metadata loading

* Simplify Nunchaku compact config loading

* Add Nunchaku Lite quantization tests

* Document Nunchaku Lite checkpoint loading

* Refine Nunchaku Lite quantization docs

* Remove unused Nunchaku smooth factor original weights

* Update docs/source/en/quantization/nunchaku.md

Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>

* Update src/diffusers/quantizers/nunchaku/utils.py

Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>

* Update src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py

Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>

* Address Nunchaku Lite review feedback

* Update Nunchaku Lite quantizer

* Rename Nunchaku Lite fp4 precision to nvfp4

* Allow torch_dtype to override Nunchaku compute dtype

* Keep Nunchaku replacement modules on meta during loading

* Require kernels when importing Nunchaku runtime utils

* Use kernels requirement decorators for Nunchaku tests

* Validate Nunchaku CUDA capability requirements

* Apply suggestions from code review

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* docs: update nunchaku quantization guide

* test: remove nunchaku quantization tests

* fix: reject hopper for nunchaku lite

* Apply suggestion from @rootonchair

* docs: update nunchaku lite compile example

* Apply style fixes

* Apply suggestions from code review

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Vinh H. Pham  <phamvinh257@gmail.com>

* update doc to be more reasonable

* fix import that requires kernels

---------

Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
2026-07-14 10:09:08 +05:30
Cheung Ka Wai
59d4e5ac86 add SP support for _flash_3_varlen_hub backend (#13809)
* add `_flash_3_varlen_hub` support for SP

* make signiture clear

* update

* add accuracy test

* fix dtype

* loosen threshold because of bfloat16

* fix the break change of https://github.com/huggingface/kernels-community/commit/1471f94

* fix stale error message

* fix merge

* revert change

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-07-08 09:12:29 +05:30
Sayak Paul
5f33428194 update licensing for others to 2026. (#14128)
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-07-08 07:40:20 +05:30
Sayak Paul
3017d887ad [tests] fix autoencoderdc ml training fix. (#14129)
fix autoencoderdc ml training fix.
2026-07-07 07:48:44 +05:30
Cheung Ka Wai
5be68ced49 fix _flash_3_varlen_hub mask handling (#14115)
* fix fa3 mask

* add test case

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-07-06 10:33:11 +05:30
GiGi Koneti
cbdb63798b Fix model offloading and training tests + prevent examples timeout (#14091)
Fix AutoencoderVidTok output format and skip training tests on unsupported hardware
2026-07-04 08:32:40 +05:30
Sayak Paul
72eb60c2da [tests] fix test_from_save_pretrained_dtype_inference (#13872)
* fix test_from_save_pretrained_dtype_inference

* fix

* add warning to to()

* address feedback

* fix

* fix

* fix

* logger.debug -> logger.warning

* Update src/diffusers/models/modeling_utils.py

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

* fix

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-07-03 09:16:56 +05:30
GiGi Koneti
f09adee565 [Tests] Skip layerwise casting tests on devices without float8_e4m3fn support (#14073)
* Skip layerwise casting tests on devices without float8_e4m3fn support

* Xfail float8 layerwise casting tests on MPS

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-06-29 06:38:00 -07:00
Sayak Paul
9d52996c4d [tests] implement base model output caching in model-level tests (#14059)
* implement base model output caching in model-level tests

* single quotes

* memory

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-06-24 22:01:19 -07:00
Sayak Paul
6a71b6e332 Fix model cuda tests (#13975)
* port final set of model tests and others

* fix extracter.

* fix cuda tests for models.

* Revert "fix extracter."

This reverts commit 178c4cb99b.

* Revert "port final set of model tests and others"

This reverts commit a92c70c081.

* fix more

* address reviewer feedback

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-06-23 17:42:36 +05:30
Sayak Paul
d1f8e55c3b [tests] fix consistency decoder tests (#13905)
* fix consistency decoder tests

* address feedback

* feedback

* up
2026-06-15 14:38:52 +05:30
hlky
0b838125b9 Incorporate safetensors support to TorchAO (#13719)
* Incorporate safetensors support to TorchAO

* Address TorchAO safetensors review comments

* Address TorchAO safetensors review feedback

* Address follow-up TorchAO review feedback

* Simplify TorchAO metadata loading guard

* Address TorchAO test utility review feedback

* Scope TorchAO safetensors to quantizer hooks

* Use tensor assertion helper in TorchAO tests

* ruff

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-06-12 10:01:38 +05:30
jiqing-feng
784fa62652 Use device_map="auto" in single file tests to support large models on limited GPU memory (#13816)
* fix flux tests OOM on 24G GPU

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* revert wrong change

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-06-11 09:09:13 +05:30
Sayak Paul
1c7f759d84 [tests] refactor UNet model tests to align with the new pattern (#13153)
* refactor unet2d condition model tests.

* fix tests

* up

* fix

* Revert "fix"

This reverts commit 46d44b73d8.

* up

* recompile limit

* [tests] refactor test_models_unet_1d.py to use modular testing mixins

Refactor UNet1D model tests to follow the modern testing pattern using
BaseModelTesterConfig and focused mixin classes (ModelTesterMixin,
MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin).

Both UNet1D standard and RL variants now have separate config classes
and dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* [tests] refactor test_models_unet_2d.py to use modular testing mixins

Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the
modern testing pattern. Each variant gets its own config class and
dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* [tests] refactor test_models_unet_3d_condition.py to use modular testing mixins

Refactor UNet3DConditionModel tests to follow the modern testing pattern
with separate classes for core, attention, memory, training, and LoRA.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* [tests] refactor test_models_unet_controlnetxs.py to use modular testing mixins

Refactor UNetControlNetXSModel tests to follow the modern testing
pattern with separate classes for core, memory, training, and LoRA.
Specialized tests (from_unet, freeze_unet, forward_no_control,
time_embedding_mixing) remain in the core test class.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* [tests] refactor test_models_unet_spatiotemporal.py to use modular testing mixins

Refactored the spatiotemporal UNet test file to follow the modern modular testing
pattern with BaseModelTesterConfig and focused test classes:

- UNetSpatioTemporalTesterConfig: Base configuration with model setup
- TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin)
- TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin)
- TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin)
- TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin)
- TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b989.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-06-10 12:33:34 +05:30
Xin He
c31bb1c148 Integrate AutoRound into Diffusers (#13552)
* support auto_round

Signed-off-by: Xin He <xin3.he@intel.com>

* add document and unit tests

Signed-off-by: Xin He <xin3.he@intel.com>

* fix CI

Signed-off-by: Xin He <xin3.he@intel.com>

* Apply suggestions from code review

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* update document and overwrite the default quantization_config with specified backend.

Signed-off-by: Xin He <xin3.he@intel.com>

* add UT and fix bug

Signed-off-by: Xin He <xin3.he@intel.com>

* update per comments

Signed-off-by: Xin He <xin3.he@intel.com>

* update per comments

Signed-off-by: Xin He <xin3.he@intel.com>

* fix compile error in doc

Signed-off-by: Xin He <xin3.he@intel.com>

* Apply style fixes

* small nits

* Add auto_round dependency to the versions table

Signed-off-by: Xin He <xin3.he@intel.com>

* fix make deps_table_check_updated

Signed-off-by: Xin He <xin3.he@intel.com>

* fix CI

Signed-off-by: Xin He <xin3.he@intel.com>

---------

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-06-10 11:41:18 +05:30
Cheung Ka Wai
86dab152ea Fix the QwenImage Attention mask under Ulysses SP (#13756)
* fix mask

* handle mask locally

* update according to comment

* fix ulysses_anything as well

* Update src/diffusers/models/transformers/transformer_qwenimage.py

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* add accuracy test

* fix controlnet

* refactor worker

* refactor correctness test

* fix style

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-06-08 15:23:06 +05:30
Sayak Paul
ed0711878c [tests] refactor caching tests. (#13235)
* refactor magcache tests.

* include taylorseer in the caching mixin.

* up

* add back magcache and migrate to pytest

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-06-02 07:44:23 +05:30
jiqing-feng
907c0c2c76 Implement _dequantize for TorchAO quantizer (#13538)
* Implement _dequantize for TorchAO quantizer

- Add _dequantize() method in TorchAoHfQuantizer that dequantizes
  TorchAOBaseTensor weights back to standard nn.Parameter
- Fix _verify_if_layer_quantized to check isinstance(weight, TorchAOBaseTensor)
  so dequantized layers are correctly detected as non-quantized

* enable dequantize for TorchAO tester mixin

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* check dequantize

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* fix dequantize: clear is_quantized flag and cast dtype after dequantize

* fix

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* fix error report

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-05-18 18:39:50 +05:30
jiqing-feng
2f4a7177f0 Fix training gradient underflow in quantization tests (#13539)
* Fix training gradient underflow in quantization tests

Change autocast dtype from float16 to bfloat16 in _test_quantization_training.
Float16's limited dynamic range causes gradients to underflow to zero when
passing through quantized tensor subclass operations.

* fix autocast dtype check

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-05-18 15:20:07 +05:30
Sayak Paul
5fd27277b8 [tests] fix bitsandbytes compile tests for flux. (#13750)
fix bitsandbytes compile tests for flux.
2026-05-15 14:24:15 +09:00
Sayak Paul
cbdedbaf03 [tests] add attention backend tests. (#13174)
* add attention backend tests.

* remove existing tests/others/test_attention_backends.py file

* modify generate_model_tests.py

* remove native.

* account for _keep_in_fp32_modules

* don't skip when exception is raised.

* use is_kernels_available()

* mark with compile.

* move rtol and atol to methods as defaults.

* Apply suggestions from code review

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* up

* up
2026-05-12 17:00:00 +09:00
Cheung Ka Wai
72ea12132e add SP support for flash_varlen_hub backend (#13479)
* add mask support for flash backend

* fix test case

* refactor test

* add protection

* fix comment

* update according to suggestion

* revert change

* fix according to claude review

* add test converage for QwenImage

* add SP support and fix non-contiguous mask for flash_varlen kernel

* revert change

* Update tests/models/testing_utils/parallelism.py

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Update tests/models/testing_utils/parallelism.py

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* drop `_padded_to_unpad`

* follow `if _parallel_config is None` pattern

* rename `attn_mask_2d`

* move check to the top

* make comment clear

* move non-contiguous-attention-mask as default dummy data

* revert and update

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-05-10 08:41:08 +09:00
jiqing-feng
8ee10d8536 Enable TorchAO int4wo quantization tests on XPU (#13537)
* Enable TorchAO int4wo quantization tests on XPU

- Remove _int4wo_skip marker that restricted int4wo tests to CUDA only
- Add XPU-specific int4_packing_format='plain_int32' for Int4WeightOnlyConfig

* add xpu to not skip

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* Apply style fixes

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2026-05-06 12:31:49 +09:00
Sayak Paul
f7fd76adcd [attention backends] fix ring CP for flash and flash 3 (#13182)
* tests: add cp backend and attention backend tests.

* up

* up

* up

* fix ring for flash and flash_3

* generate.

* Apply suggestions from code review

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

* up

* up

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-04-24 20:35:08 -03:00
Sayak Paul
7c88e5fe21 [tests] fix group offloading with disk tests (#13491)
fix group offloading with disk tests
2026-04-22 11:41:37 -03:00
Dhruv Nair
50987b1290 [CI] Fix BnB tests (#13481)
* update

* update

* update
2026-04-22 13:19:32 +05:30