* [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>
* Fix TaylorSeer cache crash when the feature shape changes between steps
* Add QwenImage 2.1 x TaylorSeer tests covering the KV-cache prefill shape change
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* [LoRA] add LoRA training for Qwen Image 2.1
Add DreamBooth LoRA training for Qwen-Image 2.1, text-to-image and
image-to-image, with fast tests and a README section.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* default rank and alpha to 16 and document the ratio
`--rank` and `--lora_alpha` are independent arguments, so raising the rank
alone leaves the update scaled by `lora_alpha / rank`. Default both to 16,
which keeps the scale at 1, and add a README section explaining the ratio.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* fix validation when embeddings are precomputed or only the final pass runs
Two fixes from review:
`QwenImage21ValidationPipeline` substituted the cached image-pad mask after
calling the base `encode_prompt`, but that call raises for a missing mask when
it is handed `prompt_embeds` together with a condition image, so the
substitution never ran and image-to-image validation failed. Pass the mask in
before the call instead.
`--final_validation_prompt` was accepted by the final-inference guard, but the
prompt embeddings were only built under `--validation_prompt`, so the final pass
reached a pipeline whose text encoder is `None` with no prompt at all. Build the
embeddings for whichever prompt is set.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* fix encoding under --offload without --cache_latents
`vae.encode` sat outside the `offload_models` block, so the context manager had
already moved the VAE back to the CPU by the time it ran, while the prepared
dataloader hands the batch over on the accelerator:
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type
(torch.FloatTensor) should be the same
Move the call inside the block and cover that flag combination, which is the
only path that encodes pixels inside the training loop. The image-to-image
trainer already had it right.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* add LoRA tests for the Qwen-Image 2.1 pipeline
Follow the Flux layout and reuse `LoraTesterMixin` and `LoraMemoryTesterMixin`,
so the adapters these trainers produce are covered by the pipeline's own
loading, fusing and memory-offload tests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* fix caching a prompt mask that `encode_prompt` returns as None
`encode_prompt` drops the prompt mask when nothing in the batch is padded,
which is the common case for `--caption_column` datasets since captions in a
bucket often tokenize to the same length. The cache is filled per sample, so
slicing that `None` raised during latent caching:
prompt_embeds_mask_cache[idx] = prompt_embeds_mask[i : i + 1]
TypeError: 'NoneType' object is not subscriptable
Materialize the mask before filling the cache, and cover the path with a test:
none of the existing ones passed `--caption_column`, which is why this only
showed up on a real dataset.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* pass rank and alpha explicitly in the README dog example
The dog example was written for the old 4/4 defaults; now that both default
to 16, spell out 4/4 so the documented command keeps training the size it
was tuned for.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* drop copy-paste leftovers from other trainers
`--output_dir` defaulted to `hidream-dreambooth-lora`, the weight-decay help
mentioned UNet params, and the VAE cast comment talked about the Flux VAE.
All three were inherited verbatim from the trainer this was derived from.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* lower the README dog example to lr 1e-4
At 2e-4 the example memorizes the five training images: prompted scenes
collapse to the training backdrop rather than following the prompt. 1e-4
is what the image-to-image example already uses.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ
* skip the LoRA scale test for the Qwen-Image 2.1 dummy
Halving the LoRA scale moves this dummy transformer's output by less than the
tolerance the assertion allows, so the two outputs compare equal and the test
fails on CPU. Skip it with the measurement recorded rather than loosening a
tolerance shared with every other pipeline's LoRA tests.
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* 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>
* feat: add Qwen-Image 2.1 pipeline with block-causal attention and KV cache
New classes:
- QwenImage21Transformer2DModel: single-stream transformer with block-causal
attention and t=0 modulation for the text and condition-image prefix
- AutoencoderKLQwenImage21: 64-channel VAE (z_dim=64, decoder_base_dim=144)
- QwenImage21Pipeline: text-to-image and image-conditioned generation
Attention:
- Block-causal: the joint text/image sequence is causal while each image block
(condition and target) stays internally bidirectional. Built as a compiled
flex_attention BlockMask, which keeps the score matrix block-sparse and makes
2048x2048 feasible.
- flex_attention is optional. Without it the mask is approximated by a two-pass
prefill (prefix causally, then the target image over the cached prefix). Exact
for text and for the target image, approximate for condition images.
KV cache:
- The text and condition-image prefix is modulated from t=0, so its activations
do not change across denoising steps and its keys and values are cached after
the first step. Later steps only recompute the target image's tokens.
Also: separate text-to-image and image-conditioned prompt templates with
image-pad token downsampling, and plain classifier-free guidance.
Includes model tests, docs, and full registration.
* fix style
* refactor: address PR review feedback for Qwen-Image 2.1
- Remove `causal_block` config flag (always on for released checkpoint)
- Split attention into QwenImage21FlexAttnProcessor and QwenImage21SDPAAttnProcessor
- Replace two-pass approximate SDPA prefill with exact multi-pass prefill
(each image block gets bidirectional attention, text gets causal mask)
- Refactor KV cache to QwenImage21KVCache/QwenImage21KVLayerCache classes
with explicit kv_cache_mode="extract"/"cached"/"extend"
- Lazy-compile flex_attention on first use (fixes OOM on uncompiled path)
- Fix edit pipeline: remove broken _downsample_image_pad_tokens, add
mm_token_type_ids for transformers 5.x, auto-convert RGB to RGBA
- Use @apply_lora_scale decorator, extract _IMG_TOKENS_PER_SLOT constant
- Add # Copied from markers for retrieve_latents and _encode_vae_image
- Fix mutable default feat_idx=[0] in all 8 VAE forward methods
- Update docs: add usage snippet, remove stale causal_block references
- Delete examples/qwenimage21/ (snippet moved to docs)
* fix: KV cache pinned the whole prefill sequence at batch size 1
The "extract" branch stored the prefix as `key[:, cache_write_slice].contiguous()`.
At batch size 1 that slice already counts as contiguous, because PyTorch ignores
size-1 dimensions in the check, so `contiguous()` returned the same view and the
cache pinned the full prefill K/V for every step of the denoising loop: 8.0 GiB of
resident memory at 2048x2048 across 32 layers. At batch size 2 and above the slice
is not contiguous, the copy happens, and the leak disappears, so no test caught it.
Store a `clone()` instead, and assert in the tests that the cached prefix owns its
storage.
Measured at 2048x2048 on one H100, bf16, batch 1, 20 steps: peak memory for a full
pipeline call drops from 64.5 GiB to 56.5 GiB.
* move the block-causal segmentation into the SDPA processor
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WpU4nuugzCw8T1c4tL6N2A
* refactor: let the processor pick the block-causal path
Builds on the previous commit by @yiyixuxu, which moved the segmentation into the
processor.
- Drop the internal `torch.compile(flex_attention)`. The convention is that the
caller compiles, so `QwenImage21FlexAttnProcessor` goes through
`dispatch_attention_fn(..., backend="flex")` and warns once when it finds an
uncompiled `flex_attention` — that falls back to a dense fp32 score matrix and
runs out of memory at high resolution.
- `_attention_backend` is `None` on the flex processor. It is what
`set_attention_backend()` sets and only applies to the cached decode steps; the
prefill needs the flex kernel for its `BlockMask` and is not configurable. The
processor raises from `__init__` when flex_attention is unavailable.
- Pad the sequence axis directly instead of transposing around `F.pad`, so the
padded tensors stay contiguous, which the compiled flex kernel requires.
- Rename `QwenImage21SDPAAttnProcessor` to `QwenImage21AttnProcessor` and make it
the default. Once compiled, flex is 1.5% faster end to end at 2048x2048 (31.8s vs
32.3s over 20 steps) and 3.8% faster with two condition images, in exchange for
37s of compilation; uncompiled it cannot render 2048x2048 at all. A default that
only works when the caller compiles is the wrong trade, so flex is documented as
the opt-in path instead.
- Derive the prefix segment boundaries once per forward rather than once per layer.
They only depend on `image_ids` and `prefix_len`, so the per-layer version repeated
the same `tolist()` device sync 32 times. `forward` passes down whichever
representation the installed processors read, and builds neither for a processor
that does not need it.
- Replace `test_non_flex_backend_rejected_when_causal`, which no longer describes the
intended behaviour, with a check that `set_attention_backend` only affects decode.
Measured on one H100, bf16, batch 1, 20 steps at 2048x2048: the default path runs out
of the box in 32.7s at 56.5 GiB peak, and `set_attn_processor(QwenImage21AttnProcessor())`
now works at all — it used to hand the flex `BlockMask` to a non-flex kernel and raise.
* fix: VAE class defaults did not describe the released model
`scale_factor_spatial` was 8 while the encoder applies four spatial downsamples:
`encode` takes a 1024x1024 image to a (1, 64, 1, 64, 64) latent, so the ratio is 16.
Every tile-to-latent conversion divides by it, so tiling silently produced a
wrong-shaped latent — a 2048x2048 encode came out as 168x168 instead of 128x128. The
in/out channel defaults were 3 while this VAE takes four channels, so the class could
not be instantiated from its own defaults.
Add the model test file that was missing, covering the ratio against the architecture
and the shape of a tiled encode. Tile values are not compared: each tile starts the
causal convolution feature cache fresh, which is a property of the tiling
implementation rather than of these defaults.
* address review feedback: copies, pipeline, docs
VAE:
- Mark the classes that are byte-identical to their Wan originals with `# Copied from`
(`DupUp3D`, `WanUpsample`, `WanRMS_norm`, `WanAttentionBlock`, `patchify`,
`unpatchify`). Adopting the upstream `RMS_norm.forward` in the process also picks up
a fix we had missed: it normalizes in fp32 for fp16/bf16/fp8 inputs.
- Drop the `non_linearity` argument, which was always "silu", from the blocks that
take it, and validate `AvgDown3D`'s channel divisibility with a `ValueError` before
the fields are assigned rather than with an `assert` after.
Pipeline:
- Take `calculate_dimensions` verbatim from the edit pipeline so it can carry a
`# Copied from`, which is what fixes `check_repository_consistency`: the marker on
`_encode_vae_image` was one blank line out of sync.
- Move the `QwenImage21KVCache` import to the top of the module.
Docs:
- Apply @stevhliu's suggestions. `models.autoencoders.vae.AutoencoderKLOutput` does
not exist, so that autodoc reference was broken; the module is `autoencoder_kl`.
- Describe the two attention processors instead of the old "with and without the flex
backend" split, and add the snippet for opting into flex, which has to be compiled.
`make style` also reflowed a few docstrings from the previous commit.
* address review feedback: validate before doing work
The `kv_cache` checks in the transformer's `forward` ran after the input projections,
the joint sequence build, the rotary embeddings and the modulation, so a bad
`kv_cache_mode` was only reported once that work had been done. They now run first.
Following the same point through the pipeline turned up a check that could never fire:
`check_inputs` warns when `height` and `width` are not divisible by
`vae_scale_factor * 2`, but the rounding happened before the call, so the values it saw
were always divisible. It now runs before the rounding, and `output_resolution=1000`
warns and yields 992x992.
* fix: num_images_per_prompt > 1 raised in the pipeline
`encode_prompt` expands `prompt_embeds` and its mask to
`batch_size * num_images_per_prompt`, but returns `image_pad_mask` unexpanded, and the
target slots appended to that mask were sized from `latents`, which is expanded. Any
call with `num_images_per_prompt > 1` therefore died in the concatenation:
RuntimeError: Sizes of tensors must match except in dimension 1.
Expected size 1 but got size 2 for tensor number 1 in the list.
Size the slots from each mask's own batch instead. The transformer reads the layout
from row 0 because samples share it, so the mask does not need expanding.
Verified end to end: `num_images_per_prompt=2`, a list of two prompts, both together,
and each of those with classifier-free guidance and with a condition image.
* Use the recommended sampling defaults: 40 steps, no guidance
Qwen-Image 2.1 is meant to be sampled in 40 steps without classifier-free guidance, so
`num_inference_steps` defaults to 40 and `true_cfg_scale` to 1.0. The other QwenImage
pipelines default to 50 and 4.0, which is why this differs from its siblings.
It also removes a warning from every default call: `true_cfg_scale=4.0` with no negative
prompt took the "guidance is not enabled" branch. Passing a `negative_prompt` without
raising `true_cfg_scale` still warns, which is the case worth warning about.
The docs and the example docstring rely on the defaults now instead of passing a step
count, and the docs state the recommendation.
* docs: complete the forward and __call__ docstrings
`utils/check_forward_call_docstrings.py` on main checks that every argument in a
forward/__call__ signature has a docstring entry and that a non-None return type has a
Returns section. It landed after this branch's base, so it only started running here
once the copy check stopped failing ahead of it.
Add the missing entries: `sample_posterior` and `generator` plus Returns on the VAE's
forward, `attention_kwargs` and `return_dict` plus Returns on the transformer's, and
the four embedding arguments and `callback_on_step_end_tensor_inputs` on the pipeline's
__call__.
* add pipeline tests for qwenimage 2.1 (#6)
* add more copied froms (#5)
* docs: describe multiple condition images, and inline the flex warning
The pipeline page now has a section on passing several condition images, which is where
@sayakpaul asked for it, and the flex section emphasises that the processor wants a
compiled model.
`_warn_if_flex_attention_is_uncompiled()` is inlined at its only call site, as
requested. A class-level flag keeps it to one warning per process: the default
processor is constructed per attention module, so 32 instances would otherwise each
warn, and the logger has no `warning_once`.
* fix-copies: drop the AvgDown3D marker
`AvgDown3D` validates its channel divisibility with a `ValueError` before assigning its
fields, where Wan still asserts after, so the copy is not consistent and
`check_copies` fails on it. The other ten markers from #5 are fine and stay.
* fix: feed the transformer the pre-norm text hidden state
The transformer was trained on the last decoder layer's output of the text encoder,
before the encoder's final RMSNorm. Up to transformers 4.x that is what
`hidden_states[-1]` holds. From transformers 5.0 the output capturing ties that entry
to `last_hidden_state`, so it comes back normalized instead, and the transformer reads
something a third of the way off — visible first as garbled text in the rendered image.
Neutralize the final norm for the encoder call with a forward hook that returns the
module's input, so `hidden_states[-1]` is the layer output on either version. Nothing
else is touched: the weights stay untouched, which matters because they are on `meta`
under offloading, and no version check is needed.
On the released checkpoint the prompt embeddings now match the pre-norm value exactly
and the rendered image is pixel-identical to it, where before the whole image shifted
by 5.35/255 on average.
* fix: match the checkpoint's text conditioning, and repair the unexercised paths
Aligned with the text encoder the checkpoint was trained with:
- The image marker is `<image1>`, `<image2>`, … as in training, not `Picture 1: `. The two
tokenize to different lengths (4 tokens against 5), so the conditioning the transformer
read was a sequence that never occurred in training, and the rendered image moves
4.02/255 on average. With the marker corrected the output is pixel-identical to the
training template. This also retires the `random.choice` over four spellings of that
word, which ran on the global RNG and left image-conditioned generation irreproducible
from `generator`.
- Condition images reach the vision encoder with their alpha composited over white, as in
training. The VAE still reads all four channels.
- The processor pads on the left, as in training. The joint sequence is re-padded on the
right either way, so this only changes the positions the encoder itself sees for a batch
of prompts of different lengths.
- An empty prompt becomes a space. Qwen has no bos token, so the encoder would otherwise
have nothing to read.
Prompt embeddings:
- The 2D prompt mask was repeated with `repeat(1, n, 1)`, which prepends an axis and tiles
the rows where the 3D embeddings interleave theirs. With more than one prompt and more
than one image per prompt, each sample was denoised against another prompt's padding.
- Supplying `prompt_embeds` raised: `image_pad_mask` only comes from `encode_prompt`'s own
encoding, and the target slots appended to it were sized from `latents`. It is synthesized
for text embeddings now, and required when the embeddings cover condition images.
Supplying embeddings without a mask raised too.
- `has_neg_prompt` no longer requires `negative_prompt_embeds_mask`, so a caller who passes
`encode_prompt`'s own output back in keeps guidance. An unpadded prompt returns `None` for
the mask, the pattern `pipeline_qwenimage.py` uses, since a mask that carries no
information costs the backends that reject one.
Condition images:
- A tensor or ndarray `image` raised on `image.size`. They are normalized to PIL up front; a
latents tensor is rejected with a message, because the text encoder has to see the image,
and a per-prompt nested list with another, because one flat set applies to the whole
batch. That also removes the half-wired path where a latents tensor was silently dropped
from both the prompt and the latents.
- A list of prompts with a condition image raised a bare `StopIteration`: every prompt's
template repeats the placeholders, but the processor was handed one set of images.
Denoising and validation:
- Interrupting on the first step used to `continue`, skipping the step that prefills the KV
cache and leaving the next one to decode from an empty one. It breaks out now.
- `kv_cache_mode` without a `kv_cache` is rejected instead of failing later on a shape.
- `prepare_latents` checks the generator list before spending a VAE encode per image.
- The VAE's image convolution names its limitation instead of asserting: it folds the single
frame away and has no temporal context, so it cannot take a feature cache.
Removed: `_downsample_image_pad_tokens`, `_max_length`, the `_drop_idx_ti2i` alias, and the
KV cache's `is_populated` and `clear()`, none of which anything reached. Training collapses
each run of `<|image_pad|>` to one token; skipping that is equivalent here, since those
positions are overwritten by the VAE latents either way.
The text encoder hook now points at huggingface/transformers#48087, which lets the config
untie `hidden_states[-1]` from 5.18 and will make the hook unnecessary.
* fix callback test (#2)
* Add a TODO at the text encoder hook
It can be replaced with `tie_last_hidden_states=False` in the text encoder's config once
huggingface/transformers#48087 ships in a stable release.
---------
Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* resolve modular pipeline classes for families outside the auto-pipeline mappings and fail clearly without blocks
* test modular class resolution outside the auto-pipeline mappings
* register ltx in the auto video mappings instead of resolving by pipeline folder
* drop the unreachable workflow check and rename the resolution test
* move the resolution asserts into the fallback test
---------
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* Invalidate HookRegistry child-registries cache on enable/disable cache
* Test the enable-cache-after-inference flow through a pipeline
* Make the FirstBlockCache regression test drive two explicit cache_context forwards
* Drop the pipeline-level enable-cache-after-inference test
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* add modular blockset for wan 2.2 vace
* drop transformer from vace encoder components, adapt tests to new mixin structure
* regenerate stale vace blocks docstring
* generalize wan22 loop denoiser for vace, drop the vace copy
* Delete image.png
* fold trim_latents into the vace core denoise step
---------
Co-authored-by: YiYi Xu <yixu310@gmail.com>
[Stable unCLIP] Fix Stable unCLIP group offloading tests
`PriorTransformer` holds `clip_mean` / `clip_std` as parameters of the model
itself rather than of a submodule, so group offloading gathers them into the
"unmatched group" whose hook wraps the top-level `forward`. They are therefore
onloaded only for the duration of `forward`, while `post_process_latents` runs
after the denoising loop and saw them back on the offload device:
RuntimeError: Expected all tensors to be on the same device, but found at
least two devices, cuda:0 and cpu!
Move them to the latents' device explicitly.
This also fixes the same failure in the three Kandinsky prior pipelines, whose
strict xfail markers now XPASS, so remove them along with the overrides that
existed only to carry them. Their `PIPELINE_GROUP_OFFLOAD_XFAIL_REASON` was a
misdiagnosis: the model-level parameters are onloaded for `forward`, and the
pipeline-level test was failing on `post_process_latents` like the others.
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: split block/leaf group offloading and derive the offloaded components
`GroupOffloadTesterMixin.test_group_offloading_inference` ran both offload
levels in one test body, so a pipeline that failed at one level had to skip
both. Split it into `test_group_offloading_inference_block_level` and
`test_group_offloading_inference_leaf_level`, sharing the helpers the test
body used to define inline, and compare against the class-scoped
`base_pipe_output` rather than rebuilding a baseline per level.
The set of components to offload was a hardcoded list of eight names, so a
pipeline with a component under any other name had it silently left on CPU
for the forward pass to trip over. Derive the set instead: every `nn.Module`
component is offloaded unless the config lists it in
`group_offloading_leaf_level_exclude_modules`, the new
`group_offloading_exclude_modules`, or `group_offloading_onload_component_names`.
A name in either exclusion list that matches no component on the pipeline
fails as the typo it is.
Ideogram4's `unconditional_transformer` was one of the silently dropped
components, which is why its group offload test was skipped; it now needs no
declaration at all, and `text_encoder` picks up block-level coverage it never
had. LTX2's skip goes the same way, with `audio_vae` declared alongside the
other VAEs the tests keep on the accelerator.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* style: fix `create_pipe` indentation in the group offload tester
The method body was indented one level too deep, which `ruff format` rewrites.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: give group offloading one exclusion list per level
`group_offloading_onload_component_names` named an implementation detail —
components the tests keep resident — and justified itself with a mechanism
that cannot occur: lazy prefetch mis-tracing a tiled VAE decode needs
`use_stream=True`, which no test on this mixin sets, and most of the
pipelines it covers never enable tiling either.
What it actually encodes is a level-specific incapability, the mirror of
`group_offloading_leaf_level_exclude_modules`. Leaf-level offloading onloads
each leaf on its own `forward`, so it breaks on compute that reads a leaf's
`.weight` directly. Block-level onloads a group when the group's leader runs
its `forward`, so it breaks on compute that re-enters submodules without
going through that leader — which is what a VAE decode path does. Rename it
to `group_offloading_block_level_exclude_modules` and scope it to the level
it describes. `group_offloading_exclude_modules` is then redundant, since a
component that fails at both levels goes in both lists, and it never had a
user. `vqvae` leaves the default: no pipeline in the new harness has one.
Measured across the suite, offloading the `vae` at block level breaks 28 of
86 test classes, while all 86 offload it at leaf level with no failures --
coverage the component-scoped tests were skipping and the pipeline-level
test has always had. Sweep before and after: the same 9 pre-existing
failures, 782 -> 810 passing.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Move group offloading changes to shared LTX-2 test mixins
* tests: cover the group offloading stream prefetch path
`use_stream=True` is the only way into `LazyPrefetchGroupOffloadingHook`,
which traces the group execution order on the first forward pass, wires each
group to its successor and flips `non_blocking` for the passes after. Nothing
in the pipeline suite reached that code, which is how a stale comment about
streams mis-tracing a tiled VAE decode survived unexamined for so long.
Add `test_group_offloading_inference_block_level_streaming` and
`..._leaf_level_streaming` alongside the synchronous pair, as separate methods
rather than a parametrized test so a pipeline can skip one level without
losing the other. Both route through `_run_group_offload_inference`, so they
compare against the same `base_pipe_output` and honour the same per-level
exclusion lists.
`apply_group_offloading` raises rather than degrading when there is nowhere to
put a stream, and `@require_torch_accelerator` still passes on MPS, so the two
tests also check for CUDA or XPU before running.
Guard against the flag quietly ceasing to have an effect: after enabling
offloading with `use_stream=True`, assert at least one group across the
pipeline actually carries a stream. The count is pipeline-wide because a
component with no `ModuleList`/`Sequential` children has nothing to prefetch
and legitimately streams nothing, and it asserts rather than skips because a
regression that drops `use_stream` would zero every count and a skip would
turn the whole streaming suite green by omission.
Sweep over `tests/pipelines -k group_offloading`: 31 failed, 1145 passed,
identical to the run without these tests. All 13 streaming failures have a
non-streaming twin that fails the same way, so the prefetch path adds no
failure of its own.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: port the group offload overrides orphaned by the level split
Splitting `test_group_offloading_inference` into a block-level and a
leaf-level test removed the name six other test files override, so their
markers guarded a method the mixin no longer defines while the real tests ran
unprotected. The `xfail(strict=True)` ones hid it: `super()` raised
`AttributeError`, which counts as the expected failure, so they reported green
while testing nothing.
Rechecking each reason rather than renaming mechanically, three of the six
turned out to describe behaviour that no longer exists:
`LTX2DFRPipelineTesterConfig` subclasses `BasePipelineTesterConfig` rather
than `LTX2BaseTesterConfig`, so the family's `audio_vae` block-level exclusion
never reached it and its audio VAE was being offloaded. It needs the
declaration, not a skip, and its group offload tests now pass.
AudioLDM2 and Motif Video only fail at block level. `get_text_features()` and
the vision tower bypass the `forward` that block-level offloading gates the
group on, while leaf level onloads each leaf on its own `forward` and passes,
so the skip and the xfail are scoped to block level and the leaf-level tests
now run.
The Kandinsky prior xfails stay at both levels but their reason was describing
the hardcoded component list this branch removed. The actual cause is that the
pipeline calls `PriorTransformer.post_process_latents()` after the denoising
loop, reading `clip_mean` / `clip_std` that group offloading onloads only for
the duration of `forward`.
Also drop the class-count measurement from the block-level exclusion comment:
it was already stale one merge later. The mechanism it was evidence for —
`vae.decode()` never running `vae.forward()` — is the durable part.
Sweep over `tests/pipelines -k group_offloading`: 31 failed -> 13, all of them
failing identically on `main`, except `kolors`'s pipeline-level test, which
fails in isolation on `main` too and only passes there because of the test
order the split changed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Clarify the Ideogram 4 block-level group offloading module exclusion comment
* tests: parametrize `use_stream` instead of splitting the streaming tests
The streaming tests were separate methods so a pipeline could skip or xfail
them independently of their synchronous twins. Nothing has ever needed that:
across three full sweeps every streaming failure had a synchronous twin at the
same level and vice versa, and none of the sixteen downstream overrides marked
a streaming variant differently from its twin. The two *levels* do get scoped
apart — AudioLDM2 and Motif Video are block-level only — so those keep their
own methods.
Fold `use_stream` into a parameter on each level's test. Four mixin tests
become two and sixteen override methods become eight, with the same coverage:
a full sweep of the parametrized and the separate-method versions on the same
tree gives 15 failed, 1263 passed, 152 skipped, 22 xfailed, 1 xpassed on both,
with identical failure sets once `[stream]` and the `_streaming` suffix are
normalised.
Marks do not carry through a subclass override, so an override that forgets to
re-declare the parametrize collapses to one un-parametrized test that errors on
the missing argument — which an `xfail` then reports as green, the same failure
mode the orphaned overrides had. `_USE_STREAM` carries a note saying so, and the
six overrides re-declare it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: point the PNDM marker comment at the split group offload tests
`test_group_offloading_inference` no longer exists — it became one test per
offload level — so the comment listing which tests share the class-level xfail
named a test that is not there, and undercounted them by one.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* docs: record the group offload streaming parameter in the testing guide
The shared memory mixin now runs each offload level with and without
`use_stream`, which is the only path that reaches group offloading's lazy
prefetch. Say so next to the exclusion-list guidance, along with the reason it
is a parameter rather than a separate test, and the requirement that a subclass
overriding either test re-declare the `parametrize` — marks do not carry
through an override, and forgetting leaves a test that errors on the missing
argument and is reported as green by an `xfail`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: inline the single-caller group offload helpers
`_split_group_offload_components` and `_count_streamed_groups` each had one
caller, so per `code_style.md` they read better at the call site. The reasons
they existed survive as comments: that a component is offloaded unless the
level's list names it, and that several submodules share one `ModuleGroup` so
the stream count has to dedupe on group identity.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* docs: condense the group offload guidance to one bullet
The three group offload bullets in the testing guide had grown to explain
the hook mechanism at a level of detail that answers a reviewer's questions
rather than telling a test author what to write. Collapse them into one
bullet keeping the parts that change what a PR does: which remedy to reach
for, which exclusion list matches which hazard, that an unmatched name fails
as a typo, and that an override must re-declare the `use_stream` parametrize.
The worked examples below the bullet already carry the concrete diagnosis,
and the parametrize trap is explained in full at the code site.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Explicitly declares group offloading leaf level excluded components for LTX-2
* Simplify comment regarding _USE_STREAM in GroupOffloadTesterMixin
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Add LTX-2.5 DFR pipeline with generated keyframe slots
Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus
single-pixel-frame keyframe slots at a fraction of the requested resolution on
a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2
re-denoises at twice that resolution with the slots re-attached and an optional
spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4
refine rounds tile the canvas at shared keyframes and densify with ancestral
Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows
the rounds.
The transformer already stored keyframes_abs_pos_embedding for load/save; this
wires it into the forward through a new video_keyframes_mask argument, which
only a keyframes-aware pipeline passes, so other pipelines are unaffected.
The epilogue denoises the whole canvas in one loop and tiles the transformer
call inside it, so every Euler step steps a canvas whose tiles have already
agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since
neither side of a height or width border holds a known answer. Temporal tiles
are cut on the keyframe seams the last refine round stitched on: both windows
reproduce a shared keyframe there, so the later one drops its run-up under a
rectangular mask rather than averaging it. Conditionings are attached once on
the whole canvas and filtered per tile at the token level, and a keyframe two
windows share is one token they both read.
The epilogue is handed its keyframes rather than asked to generate them. Each
carry plane is decoded on its own -- the VAE is causal, so a stacked decode
would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at
the output resolution, and pinned fully clean. Only the video latent is
spatially upsampled.
Conditioning fps is snapped to 60 above 30 rather than merely capped there, at
every stage. RoPE time is pixel_frame / fps, and the transformer is trained
around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between
those, and it shows as stutter at the latent borders. Playback fps is
unchanged, so 24 fps with one round still ships 48 fps.
A condition's index is read on the canvas num_frames asks for, and the moment
it names is carried onto each refine round's longer canvas by scaling its pixel
position by 2**round. The scaled position does not generally land on a latent
boundary, so it travels as a pixel index rather than through the public latent
index; a keyframe conditioning is appended as extra tokens instead of being
spliced into the base grid, so it does not need to.
height and width must be divisible by 2**spatial_upscalings times the VAE's
spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That
rule is checked ahead of the looser one every LTX-2 pipeline applies, so the
error names the divisor a DFR caller actually has to satisfy.
Four details are easy to get wrong, and each is covered by a test after showing
up as a visible seam at a tile handover:
- The ancestral step injects noise into every token, so the conditioning blend
has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor
keyframes erode over the schedule, and those anchors are the only thing
pinning adjacent tiles onto the same content.
- Velocity is converted to x0 with each token's own noise level, not the scalar
schedule sigma: a token held at strength s sits at (1 - s) * sigma.
- Each tile draws its ancestral noise from a generator seeded
seed + 1000 * round + tile, kept separate from the main generator so the
draws do not consume state the next tile's initial noising reads.
- Two tiles invent the slot that falls in the later one's dropped lead-in. The
stitch keeps the earlier tile's frames there, so the earlier tile's copy is
the one the canvas holds, and the one the next round must anchor on.
* Split the DFR pipeline into composable stages
Addresses the review on #14567: use the same compose pattern as the other LTX
two-stage pipelines, get ancestral Euler from an existing scheduler, and put the
temporal rounds in their own pipeline so the schedule is not switched mid-call.
`LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers
compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round,
and a documented recipe is the copy-paste 1080p path. The recipe knobs
(`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and
the required upsampler components are gone; `height`/`width` are this pass, not
the final output.
`ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a
re-application of the conditioning blend, which ancestral noise would otherwise
erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline`
owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on
`FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than
silently taking a deterministic step and returning a softer canvas.
Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to
`LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are
raw on both sides of every boundary, `output_type="latent"` returns the untrimmed
canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim
before decode. `LTX2DFRPipelineOutput` carries `keyframes` and
`keyframe_positions`, which cannot be re-derived after a round.
Verified against the pre-split implementation: bit-exact on dummy components in
fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds;
and within one bf16 ulp on the real checkpoint, where the only difference is that
the split normalizes upsampled latents in fp32 rather than bf16.
Also in this pass:
- `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a
resolved token plan. Resolving one needs the RoPE coordinates that only exist
once `prepare_latents` has run, so a caller could not build the plan at all.
- `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The
composed epilogue needs it, so it was public API in practice while named
private, and its normalized return forced callers into `_denormalize_latents`.
- Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline.
Enhancement belongs to stage 1 -- re-running it would denoise the canvas under
a different prompt than the one that generated it -- and the batch is set by
the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised.
- Fix the docs recipe: `requested_frames` counted latent frames where
`trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the
three pipelines share components, so place them together instead of offloading
one and leaving `temporal_latent_upsampler` off the device; and the detailing
IC-LoRA is applied at 0.5, the strength the reference hardcodes.
- Pass `crf=0` on the conditions in the refine-round test. The default CRF sends
the image through H.264 re-compression, which needs PyAV, so the test failed on
any environment without it while testing nothing about re-compression.
* Drop LTX2DFRCoreMixin in favor of # Copied from
Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the
classes stay self-contained. Shared helpers copy from LTX2Pipeline /
LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline.
dfr_core.py keeps only the constants and canvas helpers.
* inline _finalize_output
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* remove dfr_core.py
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* small refactors to the ltx2 dfr testing suite
* save
* remove dfr_dummies.py.
* remove tests/pipelines/ltx2/test_ltx2_dfr_layout.py
---------
Co-authored-by: aartyomov <aartyomov@lightricks.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* refactor c* pipeline tests.
Migrate the remaining `tests/pipelines/` directories starting with C to the config class +
composable pytest mixins introduced in #14113: chroma, chronoedit, cogview3, consisid,
consistency_models, controlnet, controlnet_flux, controlnet_hunyuandit, controlnet_sd3 and
cosmos. cogvideo and cogview4 were already migrated.
Notable points beyond the mechanical conversion:
- Expected slices switch from the `"np"` layout (`[0, -3:, -3:, -1]`) to the `"pt"` one
(`[0, -1, -3:, -3:]`); the values themselves are unchanged.
- The Cosmos `safety_checker` overrides (save/load, variants, dtype dict, encode_prompt)
were duplicated across five files and now live in a shared
`CosmosSafetyCheckerTesterMixin` in `tests/pipelines/cosmos/testing_utils.py`.
- chroma and controlnet_flux import `FluxIPAdapterTesterMixin` from `..flux.testing_utils`
rather than the old `..test_pipelines_common` one.
- `PipelineTesterMixin.test_save_load_float16` now passes the non-deprecated `dtype=` to
`from_pretrained` instead of `torch_dtype=`, which the Cosmos2.5 `from_pretrained`
wrappers read when casting the safety checker they inject.
- Dropped along the way, matching #14619: `PipelineLatentTesterMixin`,
`PipelineKarrasSchedulerTesterMixin`, the attention-slicing and xformers tests, and the
`test_ip_adapter` overrides that only supplied an `expected_pipe_slice`.
`pytest tests/pipelines/{chroma,chronoedit,cogview3,consisid,consistency_models,controlnet,controlnet_flux,controlnet_hunyuandit,controlnet_sd3,cosmos}`
gives 14 failed, 1014 passed, 82 skipped. All 14 failures are in cosmos and reproduce
unchanged on main: `DummyCosmosSafetyChecker` is a parameter-less `ModelMixin` with a single
non-persistent buffer, so accelerate does not hook it for sequential offload and a reloaded
pipeline gets a fresh fp32 one.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* fix
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Add LTX-2.5 DFR pipeline with generated keyframe slots
Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus
single-pixel-frame keyframe slots at a fraction of the requested resolution on
a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2
re-denoises at twice that resolution with the slots re-attached and an optional
spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4
refine rounds tile the canvas at shared keyframes and densify with ancestral
Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows
the rounds.
The transformer already stored keyframes_abs_pos_embedding for load/save; this
wires it into the forward through a new video_keyframes_mask argument, which
only a keyframes-aware pipeline passes, so other pipelines are unaffected.
The epilogue denoises the whole canvas in one loop and tiles the transformer
call inside it, so every Euler step steps a canvas whose tiles have already
agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since
neither side of a height or width border holds a known answer. Temporal tiles
are cut on the keyframe seams the last refine round stitched on: both windows
reproduce a shared keyframe there, so the later one drops its run-up under a
rectangular mask rather than averaging it. Conditionings are attached once on
the whole canvas and filtered per tile at the token level, and a keyframe two
windows share is one token they both read.
The epilogue is handed its keyframes rather than asked to generate them. Each
carry plane is decoded on its own -- the VAE is causal, so a stacked decode
would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at
the output resolution, and pinned fully clean. Only the video latent is
spatially upsampled.
Conditioning fps is snapped to 60 above 30 rather than merely capped there, at
every stage. RoPE time is pixel_frame / fps, and the transformer is trained
around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between
those, and it shows as stutter at the latent borders. Playback fps is
unchanged, so 24 fps with one round still ships 48 fps.
A condition's index is read on the canvas num_frames asks for, and the moment
it names is carried onto each refine round's longer canvas by scaling its pixel
position by 2**round. The scaled position does not generally land on a latent
boundary, so it travels as a pixel index rather than through the public latent
index; a keyframe conditioning is appended as extra tokens instead of being
spliced into the base grid, so it does not need to.
height and width must be divisible by 2**spatial_upscalings times the VAE's
spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That
rule is checked ahead of the looser one every LTX-2 pipeline applies, so the
error names the divisor a DFR caller actually has to satisfy.
Four details are easy to get wrong, and each is covered by a test after showing
up as a visible seam at a tile handover:
- The ancestral step injects noise into every token, so the conditioning blend
has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor
keyframes erode over the schedule, and those anchors are the only thing
pinning adjacent tiles onto the same content.
- Velocity is converted to x0 with each token's own noise level, not the scalar
schedule sigma: a token held at strength s sits at (1 - s) * sigma.
- Each tile draws its ancestral noise from a generator seeded
seed + 1000 * round + tile, kept separate from the main generator so the
draws do not consume state the next tile's initial noising reads.
- Two tiles invent the slot that falls in the later one's dropped lead-in. The
stitch keeps the earlier tile's frames there, so the earlier tile's copy is
the one the canvas holds, and the one the next round must anchor on.
* Split the DFR pipeline into composable stages
Addresses the review on #14567: use the same compose pattern as the other LTX
two-stage pipelines, get ancestral Euler from an existing scheduler, and put the
temporal rounds in their own pipeline so the schedule is not switched mid-call.
`LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers
compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round,
and a documented recipe is the copy-paste 1080p path. The recipe knobs
(`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and
the required upsampler components are gone; `height`/`width` are this pass, not
the final output.
`ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a
re-application of the conditioning blend, which ancestral noise would otherwise
erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline`
owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on
`FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than
silently taking a deterministic step and returning a softer canvas.
Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to
`LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are
raw on both sides of every boundary, `output_type="latent"` returns the untrimmed
canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim
before decode. `LTX2DFRPipelineOutput` carries `keyframes` and
`keyframe_positions`, which cannot be re-derived after a round.
Verified against the pre-split implementation: bit-exact on dummy components in
fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds;
and within one bf16 ulp on the real checkpoint, where the only difference is that
the split normalizes upsampled latents in fp32 rather than bf16.
Also in this pass:
- `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a
resolved token plan. Resolving one needs the RoPE coordinates that only exist
once `prepare_latents` has run, so a caller could not build the plan at all.
- `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The
composed epilogue needs it, so it was public API in practice while named
private, and its normalized return forced callers into `_denormalize_latents`.
- Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline.
Enhancement belongs to stage 1 -- re-running it would denoise the canvas under
a different prompt than the one that generated it -- and the batch is set by
the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised.
- Fix the docs recipe: `requested_frames` counted latent frames where
`trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the
three pipelines share components, so place them together instead of offloading
one and leaving `temporal_latent_upsampler` off the device; and the detailing
IC-LoRA is applied at 0.5, the strength the reference hardcodes.
- Pass `crf=0` on the conditions in the refine-round test. The default CRF sends
the image through H.264 re-compression, which needs PyAV, so the test failed on
any environment without it while testing nothing about re-compression.
* Drop LTX2DFRCoreMixin in favor of # Copied from
Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the
classes stay self-contained. Shared helpers copy from LTX2Pipeline /
LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline.
dfr_core.py keeps only the constants and canvas helpers.
* inline _finalize_output
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* remove dfr_core.py
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* remove the dfr_layout test suite
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* Resolve the revision once at the beginning of from_pretrained
Use `huggingface_hub.resolve_revision` (new in huggingface_hub 1.26.0) at the
top of the loading entrypoints so that every file fetched afterwards is pinned
to the same commit and can be served from the cache without re-resolving the
revision on each call.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* less comments
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* refactor a* pipeline tests.
* [tests] fix expected slice in the sparsectrl dict/tuple test
The `test_dict_tuple_outputs_equivalent` override ported in e10f8d08e carried a
nine-value expected slice that is not the one the pre-migration file used, so CI
failed on CPU with `Shape mismatch: actual torch.Size([6]) vs expected
torch.Size([9])`.
Six is correct: the base test slices `output_tuple[0, -3:, -3:, -1, -1]` and this
pipeline's dummy inputs use `num_frames=2`, so the frame axis yields 2 rather
than 3. Restored the original values, unchanged since #14120.
The literal is only read when `torch_device == "cpu"`, which is why local GPU
runs did not catch it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* refactor lumina2 pipeline tests to the new mixin structure
* move lumina2 lora tests to the pipeline-level mixins
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
