mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
e97a1ad2a287c1e8c0f54d63cdccdb0fdf810e1b
251 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
David Bertoin
|
a487e2f131 |
Add PRXPixelPipeline: pixel-space PRX text-to-image pipeline (#13928)
Adds a pixel-space variant of PRX that denoises raw RGB directly (no VAE), conditioned on a Qwen3-VL text encoder: - PRXTransformer2DModel: new optional config args `bottleneck_size` (two-layer img_in projection for large patch dims) and `resolution_embeds` (PRXResolutionEmbedder conditions the timestep modulation on the latent resolution) - PRXPipeline: support for subclass-tuned tokenizer max length, light text cleaning, x-prediction flow matching (x0 -> velocity conversion), and non-unit initial noise scale - PRXPixelPipeline: thin subclass wiring the above together (vae optional/None, vae_scale_factor=1, 1024px default) - conversion script support for the pixel checkpoint format - registration in __init__ files + dummy objects, docs autodoc entry, fast pipeline tests Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
MaciejBalaNV
|
212a98d4f2 |
Add Sound Encoder to Cosmos3 (#13911)
* Initial version of sound encoder Signed-off-by: Maciej Bala <mbala@nvidia.com> * Answered review comments Signed-off-by: Maciej Bala <mbala@nvidia.com> * Reshaped encoder Snake1d Signed-off-by: Maciej Bala <mbala@nvidia.com> * Updated test Signed-off-by: Maciej Bala <mbala@nvidia.com> * Fixed test Signed-off-by: Maciej Bala <mbala@nvidia.com> * Apply style fixes * Added docstring Signed-off-by: Maciej Bala <mbala@nvidia.com> --------- Signed-off-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
rmatif
|
c943837899 |
Add Anima modular pipeline (#13732)
* Add Anima pipeline * Fix empty Anima negative prompts * Fix Anima registration * Clean up Anima conditioner * Refactor Anima to modular * Use modular loader in Anima docs * Move Anima text conditioner * Apply style fixes * Address Anima review nits * Fix Anima rotary autocast * Clean up Anima LoRA checks * Update Anima model id * Document Anima LoRA loader * Document Anima conditioner forward --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
Guian Fang
|
d94ad81b87 |
[AnyFlow] FAR: standalone causal-mask builder + torch.compile follow-up (#13792)
* [AnyFlow] FAR: standalone causal-mask builder + torch.compile follow-up
Follow-up to #13745. Extracts FAR mask construction to a module-level
helper and adds an `attention_mask` forward kwarg so
AnyFlowFARTransformer3DModel can be wrapped in
`torch.compile(fullgraph=True)`. The pipeline pre-builds the mask during
KV-cache prefill so users get end-to-end fullgraph compile.
* Public method `AnyFlowFARTransformer3DModel.build_attention_mask(...)`
(modes: "train", "cache") plus private module-level helper
`_build_anyflow_far_causal_block_mask(...)`.
* `_build_freqs` cache lookup/write bypassed under
`torch.compiler.is_compiling()` to avoid a Dynamo guard recompile on
the second compiled call (applied in bidi source; synced to FAR via
`# Copied from`).
* `TestAnyFlowFARTransformer3DCompile(TorchCompileTesterMixin)` —
recompilation_and_graph_break, repeated_blocks, and group_offloading
pass on H200; AOT is `@pytest.mark.skip`'d (torch.export rejects
BlockMask as a pytree input).
* Base `get_dummy_inputs` omits `attention_mask` so every non-compile
test class exercises the in-forward fallback; the compile class
overrides to inject a pre-built mask.
* Bit-exact: pre-built path vs internal-build fallback max|Δ|=0.0e+00.
* [AnyFlow] docs: full author list, repo demo examples, slimmer pipeline page
* Full author list and NVIDIA → NUS → MIT institution order; TL;DR +
abstract + Available Models bullets.
* Rewritten pipeline-selection tip describing both pipelines symmetrically.
* T2V / I2V / V2V examples now use the canonical 81-frame setup and the
demo prompts / conditioning assets shipped under
`NVlabs/AnyFlow/assets/evaluation/` (linked via raw.githubusercontent.com).
* Drop the inline "Optimizing Memory" and "torch.compile" sections — those
notes will live in the NVlabs/AnyFlow repo's own performance guide rather
than the diffusers pipeline reference.
* Sync zh user guide and the two model-API stubs.
* [AnyFlow] FAR: move chunk_partition default into transformer config
- AnyFlowFARTransformer3DModel.__init__ now accepts chunk_partition via
@register_to_config (default (1, 3, 3, 3, 3, 3, 3, 2) for the released
81-frame checkpoints, matching the field on Hub).
- AnyFlowFARPipeline.__call__ no longer requires chunk_partition; defaults
to self.transformer.config.chunk_partition. Per-call override still
supported for V2V / non-default num_frames.
- Drop the AnyFlowFARPipeline.default_chunk_partition class attribute.
- Update docs (en pipelines/models, zh using-diffusers) and the conversion
script to match.
* [AnyFlow] FAR pipeline: fix `timesteps` shadowing across chunks
Inside the per-chunk rollout loop, the local variable `timesteps` was
reassigned to `self.scheduler.timesteps` after `set_timesteps()`. On the
next chunk iteration the same name was passed back into
`set_timesteps(timesteps=...)`, where a non-None value enters the
*custom-schedule* branch — `apply_shift` re-runs on already-shifted
values, double-shifting the schedule for every chunk after the first.
Concretely, with `shift=5` and `num_inference_steps=4`:
- chunk 0 timesteps: [1000, 937.5, 833.3, 625] (correct)
- chunk 1+ timesteps: [1000, 986.8, 961.3, 892.9] (double-shifted)
The later steps drift toward `t=1000` instead of toward `t=0`, the
flow-map model is conditioned on the wrong source sigma, and the chunk
KV cache accumulates errors that show up as artifacts in later video
frames.
Fix: rebind the cached schedule to a fresh local name
(`scheduler_timesteps`) so the outer-scope `timesteps` kwarg (the
user-provided custom schedule, when any) stays untouched across chunks.
Layer-by-layer verification against the NVlabs reference implementation
on H200 (elephant prompt, seed 0, 4 NFE, 81 frames):
- chunk 0 inference: bit-exact (0.0 mean diff)
- chunk 1 step 0: 0.194 → 0.014 (-93%)
- chunk 7 last step: 0.564 → 0.274 (-51%)
* [AnyFlow] FAR: doc-builder line wrap for chunk_partition docstrings
Pure rewrap to satisfy `doc-builder style --max_len 119`. Two docstrings
introduced in
|
||
Atharva Joshi
|
a1c7df4801 |
Adding Cosmos 3 to Diffusers (#13818)
* Adding Cosmos 3 * removed dead code * Change customer TimeEmbedding Layer to DIffusers Time Embedding * removed dependency on hugging face transformers * refactor 1 * Fixed Attention Pattern * Removed from Pretrain overrides * Removing normalization from the audio Tokenizer * fixed diffusers checkpoint * fixed video save uint conversion * added forward hook for cpu offload case * removed dead params for sound tokenizer * renaming audio encoder for readability * ruff format * Fix checkpoint conversion script for sound tokenizer * Audio Decoder trim and removing some dead code * removing dead sequence packing code * refactor pipeline to diffusers style formatting * removing use of cosmos3 audio encoder * Revert "removing use of cosmos3 audio encoder" This reverts commit 1b8b99a95e5acc749b5c100637ed382de138c606. * refactor audio encoder * inline remaining sequence packing functions and lint * Removed GenerationDataClean class and Action logic * inlined default args * removed dead code and refactoring * drop pipeline-helper @no_grad, inline derive helper, move guidance check to check_inputs Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop unused list bookkeeping from PackedSequence attn_modes was never read; sample_lens collapses into the existing sequence_length int (we only pack a single sample at a time); split_lens collapses into a single und_len int (only split_lens[0] was ever read). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * extract pack-time build state from PackedSequence dataclass curr, _use_mrope, _mrope_temporal_offset, _mrope_reset_spatial were transient counters used while building the joint sequence, not part of the finalized output. Thread them through _pack_*_tokens as positional args/returns so the dataclass only carries fields the pipeline actually reads back. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop private-API isinstance/shape asserts These were build-mode guards on PackedSequence internals and shape checks on tensors the pipeline itself constructs, both flagged as noise in private code per reviewer. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * build PackedSequence tensors on target device, drop to_cuda Thread device through pack_input_sequence and _pack_*_tokens helpers so all torch.tensor/zeros/arange calls land on the target device directly. Move CPU-side mRoPE tensors over with .to(device) at the append site. Pass device to finalize so list-to-tensor conversion lands on device too. Delete PackedSequence.to_cuda() and the helper _modality_to_cuda; drop the corresponding call sites in __call__. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop SequencePlan, skip_text_tokens, and bos_token_id branch SequencePlan's has_text and has_vision were True at every construction site and has_sound was derivable from x0_tokens_sound is not None. condition_frame_indexes_vision is now passed directly as a List[int] arg to pack_input_sequence. Removed the skip_text_tokens flag (never True) and the dead bos_token_id shift branch in _pack_text_tokens (special_tokens never carries one in this pipeline). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * use retrieve_latents helper in _encode_video Copied from stable_diffusion_img2img matching cosmos2_5 convention so make fix-copies keeps it synced. Functionally identical to the prior .latent_dist.mode() call but handles latent_dist/latents attribute variants uniformly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop get_data_and_condition + data_batch dict scaffolding normalize_video_databatch_inplace, augment_image_dim_inplace and remove_padding_from_latent were no-ops once is_preprocessed=True (always set by this pipeline) and the pipeline never pads. get_data_and_condition just orchestrated those plus a never-taken multi-vision branch. Replaced the whole chain with a few lines inline in prepare_latents: build vision_tensor on device, call _encode_video, set fps_vision. prepare_latents no longer needs input_caption_key, input_video_key, input_image_key, or the prompt kwarg. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop _load_image_as_tensor; use VideoProcessor for conditioning frame Image loading is the caller's responsibility (load_image from diffusers.utils), matching the cosmos2_5 example. The pipeline registers a VideoProcessor in __init__ and calls preprocess() to resize + normalize caller-supplied PIL / np / tensor inputs to [1, 3, H, W] in [-1, 1]. prepare_latents fills the temporal dim in two lines (single frame at t=0, repeat-pad the rest) preserving the prior i2v behavior. Inference script updated to call load_image() before passing to the pipeline. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * move encode/decode helpers + transformer forward into Cosmos3OmniTransformer Pull encode_text / encode_vision / encode_sound_tokens / decode_vision / decode_sound_tokens (and their pure-tensor helpers patchify_and_pack_latents, unpatchify_and_unpack_latents, apply_timestep_embeds_to_noisy_tokens, _pack_sound_latents, _unpack_sound_latents) from the pipeline into Cosmos3OmniTransformer as methods. The transformer's forward(packed_seq) now runs the full per-step pass: encode text/vision/sound, rotary + layer loop, decode vision/sound — returns (preds_vision, preds_sound). The pipeline's CFG loop drops the encode_*/decode_* method calls and the manual und/gen split/concat; each pass is now a single self.transformer(packed_seq) call. No self.transformer.{embed_tokens, vae2llm, llm2vae, sound2llm, llm2sound, time_embedder, time_proj, sound_modality_embed} access remains in the pipeline. Cosmos3VLTextModel is kept as a structural wrapper for now — flattening it would break the published checkpoint layout. Tracked separately. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * remove Cosmos3VLTextModel; flatten transformer layout embed_tokens / layers / norm / norm_moe_gen / rotary_emb are now direct attributes of Cosmos3OmniTransformer. The converter strips the leading `model.` prefix from the source language_model state-dict so new conversions land at the flat layout natively. Published Hub artifact (nvidia/Cosmos3-Nano) needs its transformer safetensors + index.json re-keyed with the same prefix strip before this code can load it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * move save_img_or_video/save_wav to cosmos/export_utils.py Per reviewer guidance, custom video/audio export helpers belong in a pipeline-local export_utils.py (mirroring pipelines/ltx2/export_utils.py) rather than living inside the pipeline file. Pipeline imports trim the now-unused pathlib/numpy/export_to_video; inference example updated to import from the new location. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop @torch.no_grad on pipeline __call__ Diffusers pipeline convention: __call__ does not wear a torch.no_grad decorator; the responsibility for grad context sits with the caller. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * follow standard transformer conventions in Cosmos3OmniTransformer - Declare _no_split_modules, _repeated_blocks, _skip_layerwise_casting_patterns, _keep_in_fp32_modules, _supports_gradient_checkpointing on the transformer. - Wire self.gradient_checkpointing + the _gradient_checkpointing_func branch in forward so the flag is honest (models.md gotcha #3). - Add PeftAdapterMixin and AttentionMixin to the mixin set so LoRA loading and the attention-backend setters work. - CosmosAttnProcessor3_0 now declares _attention_backend / _parallel_config and forwards them to dispatch_attention_fn, matching the pattern in models.md and transformer_wan.py. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * restore @torch.no_grad() on Cosmos3OmniDiffusersPipeline.__call__ Reverting bea4eecf9 — removing the decorator causes GPU OOM during inference because the autograd graph accumulates across the full denoising loop (35 steps × dual cond/uncond passes × full transformer). pipelines.md gotcha #2 documents this exact failure mode and the convention is upheld by every sibling pipeline (pipeline_flux.py:652, pipeline_qwenimage.py:462, pipeline_wan.py:381, pipeline_ltx.py:535). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop training-only dummy projection paths in transformer decode helpers _decode_vision and _decode_sound both guarded a "no noisy tokens" branch that ran zeroed projections to keep the autograd graph intact. Those branches only fire when a pure-conditioning step has no MSE-loss tokens, which never happens in the inference pipeline — every workflow has at least one noisy vision frame, and _decode_sound is gated on has_sound which itself requires noisy sound tokens. Deleting per CLAUDE.md: "delete training-time code paths… only keep the inference path." Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * inline modality encode/decode helpers into transformer forward models.md "Coding style": all layer calls should be visible directly in forward — avoid helper functions that hide nn.Module calls. Inlines _encode_text / _encode_vision / _encode_sound / _decode_vision / _decode_sound into forward so embed_tokens, vae2llm, sound2llm, llm2vae, llm2sound, and time_embedder are all visible at the call site. Pure- tensor helpers (_patchify_and_pack_latents, _unpatchify_and_unpack_latents, _pack_sound_latents, _unpack_sound_latents, _apply_timestep_embeds_to_noisy_tokens) stay as methods since they don't hide layer state. Also drops the inference-unreachable guards while collapsing the helpers: the "vision is None" / "sound is None" / "mse_loss_indexes.numel() > 0" branches never fire because the pipeline always packs vision, only routes to the sound branch when sound is present, and condition_frame_indexes never covers the entire stream. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * trim dead flags and idioms in Cosmos3 pipeline __call__ - Drop assert config.use_moe — use_moe is never read by the model and asserts vanish under python -O; if !use_moe is unsupported the place to surface it is check_inputs, not a stripped assert. - Delete the joint_attn_implementation == "flex" path entirely (the include_end_of_generation_token branch in pack_input_sequence, the include_eog hoist, and both call-site kwargs). The published config is "two_way"; the flex branch and the end_of_generation special token were dead under every shipped checkpoint. - Drop torch._inductor.cudagraph_mark_step_begin() from the step loop — cudagraph stepping belongs in the caller's torch.compile wrapper, not the inference pipeline. - Replace four int(torch.prod(torch.tensor(shape))) idioms with math.prod(shape) — no tensor allocation, no .item() sync, and math is already imported. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * introduce Cosmos3Condition for prepare_latents return shape Replace the 7-tuple return from prepare_latents with a Cosmos3Condition dataclass that carries the encoded conditioning latents (vision + optional sound), their fps tensors, the conditioning frame indices, and the num_vision_items count. The denoising loop and _postprocess_latents now read these as named attributes instead of positional tuple unpacking. Addresses reviewer thread huggingface/diffusers-new-model-addition-cosmos#1 comment 3278569263 ("create something like Cosmos3Condition class for condition input"). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * restructure pack helpers as data-returning pipeline methods Replace the four module-level pack functions (_pack_text_tokens, _pack_vision_tokens, _pack_sound_tokens, pack_input_sequence) with methods on Cosmos3OmniDiffusersPipeline. The three per-segment methods now build and return their own data (text_ids/mrope_ids tuple for text, a populated ModalityData for vision/sound) instead of mutating a shared PackedSequence builder; pack_input_sequence orchestrates them. Other cleanups along the way: - Drop dead branches that the published config never exercises: use_mrope=False (model is always unified_3d_mrope), has_generation=False (always True), multi-vision-items (num_vision_items always 1), and the curr_rope_id non-mrope path. - Move the bf16 cast of per-step noisy tokens to before pack_input_sequence so the build-then-mutate pattern on packed_seq.vision.tokens disappears. - Drop the latent_patch_size / config hoists in __call__ that are now read directly inside the pack methods. Addresses reviewer threads huggingface/diffusers-new-model-addition-cosmos#1 comments 3278871807, 3278908766, 3278918514. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * collapse builder-pattern PackedSequence/ModalityData into flat dataclasses ModalityData and PackedSequence carried a list-or-tensor union for every field so they could double as builders during packing, with finalize() converting the lists to tensors at the end. Now that the pack methods each build their segment in one shot, finalize() is just a list->tensor conversion the pack methods can do themselves. - Rename ModalityData to _ModalityData (internal) with all-tensor fields (lists only for per-item entries like tokens / condition_mask). - Rename PackedSequence to Cosmos3PackedSequence and drop its finalize() method; fields are direct tensors at construction time. - _pack_vision_tokens / _pack_sound_tokens now build finalized _ModalityData directly via torch.arange / torch.tensor / torch.full. - pack_input_sequence builds the final Cosmos3PackedSequence in one return statement; no more two-stage build-then-finalize. - Update the transformer's forward docstring to reference the new name. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * exclude dtype from saved transformer config via ignore_for_config ModelMixin.from_pretrained injects dtype into init_dict (configuration_utils.py:289) whenever it appears in the loader's unused_kwargs, so Cosmos3OmniTransformer.__init__ has to accept dtype. But the default @register_to_config decorator was also serializing it into config.json on every save — leaving a stray "dtype": "bfloat16" key that doesn't describe the architecture, just the load-time runtime preference. Adding ignore_for_config = ["dtype"] keeps the decorator from registering dtype while still accepting it on the init signature. New saves omit dtype; existing checkpoints that have it log a warning at load time but the value is re-injected and __init__ ignores it, so loading is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * add Cosmos3 pipeline and transformer docs pages; retire JSON example inputs - New docs/source/en/api/pipelines/cosmos3.md with copy-pasteable text-to-image, text-to-video, image-to-video, and text-to-video-with-sound example snippets. The docs page is now the canonical reference for application code instead of the JSON-driven inputs/ directory. - New docs/source/en/api/models/cosmos3_omni_transformer.md describing the MoT dual-pathway architecture and showing a from_pretrained snippet. - Wire both pages into docs/source/en/_toctree.yml. - Export Cosmos3OmniDiffusersPipeline from the top-level diffusers package (matching every sibling pipeline) and add the corresponding dummy class for torch/transformers-unavailable environments. - Delete examples/cosmos3/inputs/omni/{t2i,t2v,i2v}.json and rewrite inference_cosmos3.py to take --prompt / --vision-path / --num-frames directly as CLI args. The script stays as a development smoke-test runner; canonical usage now lives in the docs. - Refresh examples/cosmos3/README.md to point at the docs page and reflect the new CLI surface. Addresses reviewer thread huggingface/diffusers-new-model-addition-cosmos#1 comment 3278258567. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * restore full prompts in cosmos3 docs from original example JSONs The text-to-video and image-to-video example prompts were shortened when porting from examples/cosmos3/inputs/omni/*.json into the docs page. Restore them verbatim from the JSONs so the docs reflect the prompts the model was actually demonstrated against, and so users copying from the docs get the same conditioning the example was tuned for. Also align the text-to-video-with-sound example: it now reuses the exact same prompt as the text-to-video block with only enable_sound=True added, instead of a hand-written waterfall prompt. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Renamed Cosmos3 module attributes Signed-off-by: Maciej Bala <mbala@nvidia.com> * bugfix Signed-off-by: Maciej Bala <mbala@nvidia.com> * Removed unnecessary helper function; added extra comments Signed-off-by: Maciej Bala <mbala@nvidia.com> * Moved from encode_prompt to tokenize_prompt Signed-off-by: Maciej Bala <mbala@nvidia.com> * Bring back video system prompt Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove multi frame conditioning for now Signed-off-by: Maciej Bala <mbala@nvidia.com> * Removed default negative prompts from code Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove Cosmos3condition; simplify sequence pack Signed-off-by: Maciej Bala <mbala@nvidia.com> * Clean up multiple parameters Signed-off-by: Maciej Bala <mbala@nvidia.com> * Simplify decode video; remove remainings of batching Signed-off-by: Maciej Bala <mbala@nvidia.com> * simple renames Signed-off-by: Maciej Bala <mbala@nvidia.com> * Refactored schedulers for sound Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove unnecessary autocast Signed-off-by: Maciej Bala <mbala@nvidia.com> * Update sound example Signed-off-by: Maciej Bala <mbala@nvidia.com> * Simplify loops in transformer_cosmos3.py Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove unused config attributes Signed-off-by: Maciej Bala <mbala@nvidia.com> * Cleanup audio decoder Signed-off-by: Maciej Bala <mbala@nvidia.com> * Reuse encoder_video from LTX2 for Cosmos3 Signed-off-by: Maciej Bala <mbala@nvidia.com> * Fixed a few nits Signed-off-by: Maciej Bala <mbala@nvidia.com> * Moved to RMSNorm for Cosmos3 Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove meta_tensor usage Signed-off-by: Maciej Bala <mbala@nvidia.com> * Improved rope handling Signed-off-by: Maciej Bala <mbala@nvidia.com> * Improved prompt templates Signed-off-by: Maciej Bala <mbala@nvidia.com> * Added extra docs for templatete Signed-off-by: Maciej Bala <mbala@nvidia.com> * remove dataclasses * Cleanup after merging Signed-off-by: Maciej Bala <mbala@nvidia.com> * Added guardrails Signed-off-by: Maciej Bala <mbala@nvidia.com> * bugfixed guardrails Signed-off-by: Maciej Bala <mbala@nvidia.com> * Bugfix guardrails v2 Signed-off-by: Maciej Bala <mbala@nvidia.com> * Simplified input_timestep Signed-off-by: Maciej Bala <mbala@nvidia.com> * Add TODO Co-authored-by: YiYi Xu <yixu310@gmail.com> * simplify conditional mask generation * Inlined _postprocess_latents Signed-off-by: Maciej Bala <mbala@nvidia.com> * removed pack_input_sequence helper * restore export utils with deprecation warning * moved sampling rate to pipeline attribute * inlined sound and image condition mask * seperating static and timestep based sound and vision token packing * unpack transformer args * enabled selection of cosmos3 super * Update src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Apply suggestions from code review Co-authored-by: YiYi Xu <yixu310@gmail.com> * fixed sound and vision conditioning use from prepare latents * ruff format and doc builder * ran fix copies * move typing to python3.10 * move special token application from pack_text_tokens to tokenize_prompt * rename packing methods to process methods * remove guidance_scale check * fix nits * respect vae dtype in the pipeline * use vae dtype for vae normalization stats * skip CFG if guidance_scale is 1 * Remove unnecessary parameter Signed-off-by: Maciej Bala <mbala@nvidia.com> * fix CFG for sound * bugfix for sound CFG * ruff format * Fix apply_chat_template return dict arg to return BatchEncoding * added option to select attention processor * docs: refresh Cosmos 3 pipeline intro Replace the terse architectural lede with the launch-style positioning (unified WFM for Physical AI, consolidating Predict/Reason/Transfer into one omni-model) and split out "What's new" and "Available checkpoints" sections so the page leads with capability rather than repo IDs. * docs: document Cosmos3OmniPipeline.__call__ arguments Add the missing Args block to Cosmos3OmniPipeline.__call__ so utils/check_forward_call_docstrings.py passes — covers all 21 parameters from prompt through enable_safety_check, plus a Returns section describing the output dataclass. * style: doc-builder reflow on Cosmos3OmniPipeline.__call__ docstring --------- Signed-off-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: Yuliya Zhautouskaya <yzhautouskay@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: Dima Zhylko <dzhylko@nvidia.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
|
Guian Fang
|
e39aecff57 |
Add AnyFlow Any-Step Video Diffusion Pipelines (Bidirectional + FAR Causal) (#13745)
* [Pipelines] AnyFlow: scaffold pipelines/anyflow + register all top-level imports
This is the lazy-loader scaffolding only. Body files (pipeline_anyflow.py,
pipeline_anyflow_causal.py, transformer_anyflow.py,
scheduling_flow_map_euler_discrete.py) come in subsequent commits.
* [Schedulers] AnyFlow: add FlowMapEulerDiscreteScheduler
The flow-map scheduler advances samples from timestep t to caller-provided
target r in a single Euler step, supporting any-step sampling on flow-map-
distilled checkpoints. It is a general-purpose scheduler — not specific to the
AnyFlow checkpoints.
Tests: 12 standalone tests covering instantiation, set_timesteps endpoints,
shift identity/monotonicity, step shape preservation, zero-interval identity,
one-shot sampling, train weight schemes, scale_noise endpoints.
Docs: api/schedulers/flow_map_euler_discrete.md
* [Models] AnyFlow: add AnyFlowTransformer3DModel
A 3D DiT extending the v0.35.1 Wan2.1 backbone with two config-toggled modules:
* FAR causal blocks (init_far_model=True): block-sparse causal attention via
flex_attention + compressed-frame patch embedding for frame-level
autoregressive generation (Gu et al., 2025, arXiv:2503.19325).
* Dual-timestep flow-map embedding (init_flowmap_model=True): adds a delta
timestep embedder enabling flow-map sampling z_t -> z_r over arbitrary
intervals (AnyFlow).
With both flags off, the model reduces to stock Wan2.1.
The class is intentionally self-contained rather than annotated with
'# Copied from diffusers.models.transformers.transformer_wan' because upstream
Wan has been refactored extensively since v0.35.1 (new WanAttention class,
different processor architecture).
Tests: 9 unit tests covering construction in 3 modes, bidi forward shape and
determinism, return_dict variants, save/load round-trip with and without
init_far_model, gradient checkpointing toggle.
Docs: api/models/anyflow_transformer3d.md
* [Pipelines] AnyFlow: add AnyFlowPipeline and AnyFlowCausalPipeline
* AnyFlowPipeline (pipeline_anyflow.py, ~590 LOC): bidirectional T2V using
flow-map sampling. Loads checkpoints from nvidia/AnyFlow-Wan2.1-T2V-{1.3B,14B}.
* AnyFlowCausalPipeline (pipeline_anyflow_causal.py, ~700 LOC): FAR-based
causal pipeline supporting T2V/I2V/TV2V via task_type kwarg. Loads checkpoints
from nvidia/AnyFlow-FAR-Wan2.1-{1.3B,14B}-Diffusers.
Both pipelines reuse stock WanLoraLoaderMixin, AutoencoderKLWan, UMT5EncoderModel,
and AutoTokenizer from upstream. The transformer is the AnyFlowTransformer3DModel
introduced in the previous commit. The scheduler is FlowMapEulerDiscreteScheduler.
Tests:
* tests/pipelines/anyflow/test_anyflow.py: PipelineTesterMixin fast tests +
slow integration test against nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers.
* tests/pipelines/anyflow/test_anyflow_causal.py: same structure for FAR variant.
Reference slices for slow integration tests are deferred to Phase 7
(Final quality pass) where the user runs them on a real GPU.
* [Docs] AnyFlow: add main pipeline documentation page
Modeled on the Helios pipeline doc (PR #13208). Sections: paper link + abstract,
supported checkpoints table, memory/speed optimization tabs, T2V/I2V/TV2V
examples for both bidirectional and causal variants, autodoc trailers.
* [Auto/Scripts] AnyFlow: register AutoPipelineForText2Video + add conversion script
* Register AnyFlowPipeline in AUTO_TEXT2VIDEO_PIPELINES_MAPPING.
* AnyFlowCausalPipeline is intentionally NOT registered for AutoPipeline because
its task switch (t2v / i2v / tv2v) is too rich for a single auto-resolve key.
* scripts/convert_anyflow_to_diffusers.py: convert .pt training checkpoints
(with 'ema' state dict) into a diffusers save_pretrained layout. Supports all
4 released NVIDIA AnyFlow variants. Replaces the omegaconf-based config in the
upstream repo with argparse to match other diffusers conversion scripts.
* [Quality] AnyFlow: ruff-format + regenerated dummy stubs
* ruff format pass on all 5 source files (long lines + trailing comma fixes)
* check_dummies.py --fix_and_overwrite regenerated:
- dummy_pt_objects.py: AnyFlowTransformer3DModel + FlowMapEulerDiscreteScheduler
- dummy_torch_and_transformers_objects.py: AnyFlowPipeline + AnyFlowCausalPipeline
Local fast tests: 21/21 passed
- 12 scheduler tests (FlowMapEulerDiscreteScheduler)
- 9 transformer tests (AnyFlowTransformer3DModel construction + bidi forward + save/load)
The pipeline fast tests in tests/pipelines/anyflow/ require a local dev install
that matches the diffusers main branch's transformers >= compatibility floor.
The reference slices for slow integration tests (real GPU + 1.3B/14B
checkpoints) are intentionally left as TODO stubs to be captured by the user
on a real GPU machine before opening the PR.
* [AnyFlow] address review feedback: bug fixes + DMD wording + EN/ZH tutorials
Critical bug fixes (verified against precision-validation review):
* pipeline_anyflow.py / pipeline_anyflow_causal.py: replace hardcoded
transformer_dtype = torch.bfloat16 with self.transformer.dtype, so
pipe.to("cpu") and PipelineTesterMixin save/load tests do not crash on a
dtype mismatch in the patch_embedding conv3d.
* transformer_anyflow.py: drop the duplicate `base = base = ...` assignment in
_build_causal_mask (was a copy-paste typo carried over from FAR-Dev).
* transformer_anyflow.py: drop unused `q_is_context` / `k_is_context` locals
and the `# noqa: F841` markers that were silencing the dead-store warning.
* transformer_anyflow.py: remove `CacheMixin` from the inheritance list — the
pipeline manages KV cache directly, the mixin's interface is unused.
* transformer_anyflow.py: guard the module-level `torch.compile(flex_attention)`
with try/except so the file imports cleanly on CPU CI / no-Triton machines.
* convert_anyflow_to_diffusers.py: replace ad-hoc print warnings with the
stdlib logger (warning_once-style) and a module-level basicConfig.
Documentation accuracy:
* AnyFlowCausalPipeline class docstring + main pipeline doc + EN/ZH tutorial:
drop the fictitious `task_type` / `image` / `video` arguments and document
the real API: pass `context_sequence={"raw": tensor}` (or `{"latent": ...}`)
to switch between T2V (None) / I2V (1-frame) / TV2V (4n+1-frame) modes.
* Pipeline class docstrings + main doc: explicitly describe AnyFlow's
two-stage LoRA distillation including DMD reverse-divergence supervision
with Flow-Map backward simulation in stage 2 (was previously implicit).
* training_rollout: add detailed docstring explaining its role as the
3-segment Flow-Map backward simulation entry point used during DMD training.
* Long-form tutorial doc `using-diffusers/anyflow.md` (EN, 239 LOC) and
Chinese mirror `docs/source/zh/using-diffusers/anyflow.md` (224 LOC) added
and registered in both `_toctree.yml` files.
Tests:
* Skip `test_attention_slicing_forward_pass` in both pipeline test classes
with a clear rationale (custom attention processor does not support slicing).
* All 21 standalone tests still pass (12 scheduler + 9 transformer).
Quality gates:
* `ruff check` clean across all AnyFlow files.
* `ruff format --check` reports 6 files already formatted.
* `python utils/check_copies.py` reports no diff.
Out of scope for this commit (deferred until reviewer feedback):
* Splitting AnyFlowTransformer3DModel into bidi + causal subclasses
* Unifying _forward_inference / _forward_cache return types
* Migrating model tests from plain unittest to BaseModelTesterConfig + mixins
* HF model card / config.json metadata updates on the nvidia/* repos
(push to Hub manually before opening the PR)
* [AnyFlow] rename Causal->FAR + explicit forward signature + dataclass output
Round 2 of review feedback. Three groups of changes; transformer state-dict
keys, module hierarchy, and tensor flow are unchanged so the H200 bit-exact
validation remains valid.
A. Pipeline rename (mechanical, no behavior change):
* Class: AnyFlowCausalPipeline -> AnyFlowFARPipeline (Causal in diffusers
usually means an attention mask; AnyFlow's variant is FAR autoregressive,
so the FAR name is more specific and matches the paper).
* File: pipeline_anyflow_causal.py -> pipeline_anyflow_far.py (git mv).
* Test file: test_anyflow_causal.py -> test_anyflow_far.py (git mv).
* All references updated in src/, tests/, docs/, scripts/, plus stale
anyflowcausalpipeline anchor links in tutorial markdown.
B. Pipeline test bug fixes (closes 19 fast-test failures reported by
precision-validation reviewer):
* pipeline_anyflow.py / pipeline_anyflow_far.py: __call__ now sets
self._num_timesteps = num_inference_steps before the rollout, so the
PipelineTesterMixin callback tests can read pipe.num_timesteps.
* tests/pipelines/anyflow/test_anyflow_far.py: drop the fictitious
task_type="t2v" kwarg that crashed every causal fast test (the FAR
pipeline selects mode via context_sequence, not a task_type arg).
C. Transformer architecture cleanups (review-driven, no tensor changes):
* Replace forward(*args, **kwargs) dispatcher with an explicit signature
listing every supported kwarg (hidden_states, timestep, r_timestep,
encoder_hidden_states, encoder_hidden_states_image, chunk_partition,
clean_hidden_states, clean_timestep, kv_cache, kv_cache_flag, is_causal,
attention_kwargs, return_dict). Helps IDE / type-checker / torch.compile
tracing.
* Drop SimpleNamespace returns. Add AnyFlowFARTransformerOutput
(BaseOutput dataclass with sample + kv_cache fields) for the two causal
paths that need to also propagate kv_cache (_forward_inference and the
newly return_dict-aware _forward_cache). _forward_train and
_forward_bidirection now consistently return Transformer2DModelOutput.
Pipeline call sites already use return_dict=False with positional
unpacking, so the fix is transparent there.
Out of scope (deferred until canonical-org HF metadata sync):
* Splitting AnyFlowTransformer3DModel into a bidi class plus an
AnyFlowFARTransformer3DModel subclass — touches register_to_config keys
and would require updating model_index.json on every released checkpoint.
* Promoting chunk_partition from register_to_config to a forward-time
argument (same reason).
* Renaming training_rollout to _denoise — would break callers in the
FAR-Dev on-policy trainer that produced the released checkpoints.
Local fast tests: 21/21 still pass (12 scheduler + 9 transformer).
ruff check, ruff format, and check_copies.py are all clean.
* [AnyFlow] wire callback_on_step_end through inference_range + add chunk_partition to FAR fast-test fixture
Two root causes for the 19 remaining PipelineTesterMixin failures, identified
by the H200 reviewer:
1. callback_on_step_end was accepted by __call__ but never invoked. Both
pipelines pass it through to training_rollout (and FAR additionally through
inference()), and inference_range now fires it after scheduler.step in
the standard inference branch:
if callback_on_step_end is not None:
callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = ...
negative_prompt_embeds = ...
`nonlocal prompt_embeds, negative_prompt_embeds` lets the callback rewrite
the closure-captured embeddings, matching upstream WanPipeline semantics.
The 3-segment grad_timestep training rollout does not invoke the callback;
it is intentionally training-only.
2. tests/pipelines/anyflow/test_anyflow_far.py::get_dummy_components built
the dummy transformer without a `chunk_partition`, leaving it None on the
model config and crashing the pipeline at `sum(self.transformer.config.chunk_partition)`.
Set `chunk_partition=[1, 1, 1]` in the fixture (3 chunks of 1 latent frame
each, matching the test's num_frames=9 -> 3 latent frames).
Local fast tests: 21/21 still pass.
ruff check, ruff format, and check_copies.py are all clean.
* [AnyFlow] Phase 2: split transformer + drop chunk_partition from config + rename helpers
Major architectural refactor that aligns the integration with diffusers conventions
ahead of the canonical-org Hub upload. State-dict keys, module hierarchy, and
tensor flow are unchanged so the H200 bit-exact validation remains valid; only
the on-disk transformer/config.json fields move.
Changes:
1. **Sibling transformer classes** replace the flag-driven single class:
* AnyFlowTransformer3DModel — bidirectional only. Drops compressed_patch_size /
full_chunk_limit / init_far_model / init_flowmap_model / chunk_partition
kwargs (always-on for AnyFlow distilled checkpoints).
* AnyFlowFARTransformer3DModel — adds far_patch_embedding + the 3 FAR forward
paths (train / cache-prefill / autoregressive inference).
* AnyFlowTimeTextImageEmbedding (the legacy single-time embedder used only by
the old setup_flowmap_model bootstrap) is removed; both classes now build
AnyFlowDualTimestepTextImageEmbedding directly in __init__.
* setup_flowmap_model / setup_far_model methods are removed; weight warm-start
for far_patch_embedding (trilinear interpolation from patch_embedding) moves
into AnyFlowFARTransformer3DModel.__init__.
2. **chunk_partition** is no longer a model config field. The FAR pipeline owns
the schedule:
* AnyFlowFARPipeline.default_chunk_partition = [1, 3, 3, 3, 3, 3, 3, 2]
matches the released 81-frame NVIDIA checkpoints.
* AnyFlowFARPipeline.__call__ / _denoise_rollout accept a chunk_partition
argument that overrides the default for non-default num_frames.
3. **training_rollout -> _denoise_rollout** rename across both pipelines and all
English / Chinese docs that referenced it. Signals the method is internal to
the pipeline driver, not a public training API.
4. **Conversion script + tests + docs + registries**:
* scripts/convert_anyflow_to_diffusers.py: VARIANTS dict picks the right
transformer class per variant; init_far_model / init_flowmap_model /
chunk_partition kwargs are removed from the from_pretrained call.
* Transformer test file split into AnyFlowTransformer3DModelTest and
AnyFlowFARTransformer3DModelTest classes.
* Pipeline test fixtures use the right class and pass chunk_partition via
get_dummy_inputs (3-frame schedule [1, 1, 1] for the 9-frame test).
* New docs page docs/source/en/api/models/anyflow_far_transformer3d.md;
anyflow_transformer3d.md rewritten for the bidi-only class.
* AnyFlowFARTransformer3DModel registered in src/diffusers/__init__.py,
src/diffusers/models/__init__.py, models/transformers/__init__.py and the
dummy_pt_objects.py stubs.
* docs/source/en/_toctree.yml: new entry for the FAR transformer page.
5. **Cleanups**:
* Pipeline __call__ no longer passes is_causal=False to the bidi forward (the
bidi class doesn't accept it).
* Pipeline class docstrings drop stale references to init_*_model flags.
Local tests: 22/22 pass (12 scheduler + 10 transformer covering both classes).
ruff check / format / check_copies clean.
Hub artifacts (model_index.json, transformer/config.json, scheduler config) need
to be regenerated for the released checkpoints; the HF update guide will be
delivered separately.
* [AnyFlow] Phase 3: convention compliance against .ai/AGENTS.md + .ai/models.md
Hard violations (per official diffusers guidelines):
* drop einops dependency — replace 25+ rearrange() calls with native
permute/reshape/unflatten in transformer + both pipelines
* device-gate torch.float64 — apply_rotary_emb and AnyFlowRotaryPosEmbed now
fall back to float32 / complex64 on MPS / NPU; freqs are lazily rebuilt
per-device via _build_freqs (matches transformer_wan / transformer_flux
pattern)
* migrate attention to dispatch_attention_fn — replace direct
F.scaled_dot_product_attention calls with dispatch_attention_fn (works
with sage / flash / native backends); introduce AnyFlowAttention(
AttentionModuleMixin) with _default_processor_cls / _available_processors;
rename processors to AnyFlowAttnProcessor / AnyFlowCrossAttnProcessor and
declare _attention_backend / _parallel_config class attrs
* drop dead config fields — qk_norm and added_kv_proj_dim are pruned from
both transformer __init__ signatures and AnyFlowTransformerBlock;
AnyFlowAttention is hardcoded to rms-norm-across-heads (the only scheme
the released checkpoints use) and has no add_k_proj path (T2V only)
* add _repeated_blocks = ["AnyFlowTransformerBlock"] to both transformer
classes for compile_repeated_blocks() support (matches Wan)
* annotate prepare_latents with `# Copied from diffusers.pipelines.wan.
pipeline_wan.WanPipeline.prepare_latents`; the pipeline-side rearrange
to (B, T, C, H, W) layout is moved to the call site
State-dict keys are preserved (legacy Attention had identical to_q / to_k /
to_v / to_out / norm_q / norm_k naming), so existing AnyFlow checkpoints load
bit-exactly into the new AnyFlowAttention class.
The HF Hub config-update guide is updated correspondingly: transformer/
config.json now drops qk_norm and added_kv_proj_dim alongside the previous
init_far_model / init_flowmap_model / chunk_partition removals.
22 fast CPU tests still pass; ruff format / ruff check / check_copies all
clean.
* [AnyFlow] FAR fast-test compat: rope 0-dim guard + flex_attention CPU/head-dim fallbacks + KV-cache dtype + num_timesteps
Phase 3 migrated bidi + cross-attention to dispatch_attention_fn but the FAR
causal path still calls flex_attention directly, which has hard requirements
(CPU compile, head_dim >= 16) that fail on PipelineTesterMixin's tiny dummy
components. Real ckpts (head_dim=128, CUDA) never hit these branches; bit-exact
numerical equivalence with FAR-Dev preserved on all 4 released ckpts (forward
0.00e+00, backward kernel-nondet only, ratio 1.000).
Code fixes:
1. AnyFlowRotaryPosEmbed._forward_compressed_frame / _forward_full_frame now
short-circuit to an empty tensor when num_frames / height / width is 0.
PipelineTesterMixin's dummy VAE has scale_factor_spatial=8, so a 16x16 raw
spatial input becomes a 2x2 latent which then floors to 0 against
compressed_patch_size=(1, 4, 4); the original
`freqs[:0].view(0, k, 1, -1)` reshape was ambiguous in that regime.
2. flex_attention dispatch: split the module-load
`torch.compile(flex_attention, dynamic=True)` into `_flex_attention_eager`
(always available) plus `_flex_attention_compiled`, with a tiny wrapper
that picks compiled for CUDA tensors and eager for CPU. Avoids
torch._inductor C++ codegen failures that broke fast tests after
`pipe.to("cpu")`. CUDA performance unchanged (L10 benchmark: 0.0% delta on
bidi 1.3B fwd, 0.0% delta on FAR causal 1.3B fwd).
3. AnyFlowAttnProcessor (FAR causal branch): when head_dim < 16
(flex_attention's hard minimum) zero-pad q/k/v's last dim to 16 and pass
`scale=1/sqrt(original_head_dim)` to flex_attention. Padded value rows
contribute 0, so trimming the output back is mathematically equivalent.
Released ckpts use head_dim=128 so the branch is never taken in production.
4. pipeline_anyflow_far.encode_kv_cache: replace the hardcoded
`latents.to(torch.bfloat16)` with `self.transformer.dtype`. The hardcoded
bf16 crashed conv3d on dummy fp32 components ("Input type (BFloat16) and
bias type (float) should be the same"); real bf16 ckpts are unaffected.
5. pipeline_anyflow_far._denoise_rollout sets
`self._num_timesteps = (len(chunk_partition) - num_context_chunks) * num_inference_steps`
before the chunk loop, so PipelineTesterMixin.test_callback_cfg's
`pipe.num_timesteps`-based assertion matches the actual number of callback
fires (chunks * NFE) instead of the previous hardcoded num_inference_steps.
Tests:
* test_callback_inputs cannot pass without changing FAR's chunk-wise output
semantics — it zeroes latents on the final step and asserts the *entire*
output buffer is zero, but only the active chunk's slice is overwritten in
a chunk-wise rollout. Marked `@unittest.skip` with a detailed rationale;
callback functionality itself is still covered by test_callback_cfg.
* Full pytest run on tests/pipelines/anyflow/ +
tests/models/transformers/test_models_transformer_anyflow.py +
tests/schedulers/test_scheduler_flow_map_euler_discrete.py: 81 passed,
0 failed, 11 skipped.
Quality gates:
* `ruff check` and `ruff format --check` clean across all AnyFlow files.
* `python utils/check_copies.py` clean.
* `python utils/check_dummies.py` clean.
* [AnyFlow] docs/code: paper-release tidy-up
User-facing alignment with the official HF Hub model card and the day-of-announcement
materials at https://huggingface.co/collections/nvidia/anyflow.
* Fill in the arXiv identifier 2605.13724 (5 paper links + 2 BibTeX entries).
* Rename TV2V → V2V across docs + pipeline_anyflow{,_far}.py so the diffusers
copy uses the same Video-to-Video terminology as the official model card.
* Add the [nvidia/anyflow](https://huggingface.co/collections/nvidia/anyflow)
HF collection link to the three tutorial intros.
* Drop the temporary "guyuchao/* staging" tip from the EN tutorial / API page
/ ZH tutorial — the nvidia/AnyFlow-*-Diffusers repos are now live.
* Wire up NVlabs/AnyFlow (training code) and nvlabs.github.io/AnyFlow (project
page) in place of the prior <github-org> / <project-page-url> placeholders.
* Cite the authors (Yuchao Gu, Guian Fang et al.) and NUS ShowLab × NVIDIA
affiliation in the main tutorial, API pipeline page, and both transformer
model pages; BibTeX uses the standard `and others` to elide the full list
until the next pass.
Working tree, CI gates, and tests after the change:
ruff format --check ✓
ruff check ✓
python utils/check_copies.py ✓
python utils/check_dummies.py ✓
pytest tests/models + tests/schedulers (22 fast) ✓
No production code logic changes — only docstring wording inside pipeline
files (TV2V → V2V).
* [AnyFlow] docs: drop in official BibTeX (full author list)
Replace the placeholder ``@article{gu2026anyflow, author = {Gu, Yuchao and
Fang, Guian and others}, ...}`` block in both the English and Chinese
tutorials with the canonical ``@misc{gu2026anyflowanystepvideodiffusion,
...}`` form from arxiv.org/abs/2605.13724, which lists all seven authors:
Yuchao Gu, Guian Fang, Yuxin Jiang, Weijia Mao, Song Han, Han Cai,
Mike Zheng Shou.
Docs-only.
* [AnyFlow] align with diffusers conventions + drop training-only code
Scheduler
- FlowMapEulerDiscreteScheduler.step now returns a
FlowMapEulerDiscreteSchedulerOutput dataclass (or tuple with return_dict=False)
and uses the conventional positional order (model_output, timestep, sample,
r_timestep).
- Drop training-only helpers: adaptive_weighting, set_train_weight,
get_train_weight, linear_timesteps_weights, and the weight_type config field.
- Add scale_model_input no-op for API parity; raise ValueError on missing
r_timestep.
Transformer
- Remove gate_track debug write inside
AnyFlowDualTimestepTextImageEmbedding.forward_timestep.
- Compile flex_attention lazily on first CUDA call instead of at import time.
- Replace assert with ValueError in build_block_mask.
- Resolve <arxiv-id> placeholders to 2605.13724.
Pipelines (AnyFlowPipeline + AnyFlowFARPipeline)
- Add EXAMPLE_DOC_STRING + @replace_example_docstring and full __call__
docstrings covering every argument.
- Move use_mean_velocity from __init__ to __call__ so save/load round-trips.
- Drop _denoise_rollout's grad_timestep branch (DMD on-policy training rollout),
the inner inference_range closure, and the redundant negative-prompt concat.
- Replace asserts with ValueError; wire show_progress to tqdm; rename inference
-> _inference; remove dead current_timestep property.
- Update scheduler.step call sites to the new signature.
- Trim class docstrings to inference-only language.
Pipeline output
- Add Apache 2.0 license header; switch to relative import.
Auto pipeline / conversion script
- Register AnyFlowFARPipeline in AUTO_IMAGE2VIDEO_PIPELINES_MAPPING and
AUTO_VIDEO2VIDEO_PIPELINES_MAPPING.
- Document the weights_only=False requirement in the conversion script.
Tests
- Scheduler tests use the new step signature and verify the Output dataclass
contract.
- Drop the four obsolete training-weight tests; drop weight_type kwarg from
pipeline test fixtures; remove internal milestone names from TODO comments.
Docs
- Resolve <arxiv-id> in the scheduler docs page.
- Trim DMD / on-policy distillation language in EN/ZH tutorials and the
pipelines page; the paper abstract quote is preserved verbatim.
* [AnyFlow] split FAR causal transformer into transformer_anyflow_far.py
Per @dg845's review on #13745: extract FAR causal modules into a dedicated
sibling file so each transformer variant reads in isolation. Shared submodules
are duplicated via `# Copied from` so `make fix-copies` keeps both in sync.
- `transformer_anyflow.py`: bidi-only. `AnyFlowAttnProcessor` no longer carries
the flex/KV-cache branch (was: dispatch in one branch, bare flex_attention in
the other); `AnyFlowRotaryPosEmbed` drops the compressed-frame helpers and
the `is_causal` arg; `AnyFlowDualTimestepTextImageEmbedding` drops its causal
branch. `AnyFlowTransformerBlock` keeps a single class with a new
`is_causal: bool = False` ctor flag that selects the self-attn processor —
the forward path is identical in both modes, only the processor differs.
- `transformer_anyflow_far.py`: new. Contains `AnyFlowFARTransformerOutput`,
`AnyFlowCausalAttnProcessor` (routed through `dispatch_attention_fn(backend=
"flex")` with a clear ValueError when a non-flex backend is configured; the
BlockMask is consumed only by the flex backend in `_native_flex_attention`),
`AnyFlowDualTimestepTextImageEmbeddingCausal`, `AnyFlowCausalRotaryPosEmbed`,
`AnyFlowFARTransformer3DModel`, and `# Copied from` clones of the shared
shared `AnyFlowAttention`/`AnyFlowCrossAttnProcessor`/`AnyFlowImageEmbedding`/
`AnyFlowTransformerBlock`/`AnyFlowAttnProcessor` modules.
Verified bit-exact against the pre-refactor branch on H200 (float32):
- bidi: L2 = 0.000e+00, max|Δ| = 0.000e+00
- FAR : L2 = 4.772e-06, max|Δ| = 3.576e-07
The FAR delta is fp32 accumulation noise from the dispatch path permuting
(B,L,H,D) ↔ (B,H,L,D) around the same `flex_attention` kernel.
Addresses review comments at transformer_anyflow.py:215, :261, :450, :622,
:671, :958.
* [AnyFlow] pipeline cleanup: video_processor, encode_video, inline rollout, kwarg rename
Per @dg845's review on #13745, applied to both bidi `AnyFlowPipeline` and
causal `AnyFlowFARPipeline`:
- Use `self.video_processor.preprocess_video(...)` instead of the manual
`* 2 - 1` normalize.
- Merge `vae_encode` + `encode_latents` + `_normalize_latents` into a single
`encode_video` method, mirroring `WanImageToVideoPipeline.encode_image`'s
flat structure.
- Inline `_denoise_rollout` into `AnyFlowPipeline.__call__`. For the FAR
pipeline, inline both `_denoise_rollout` and `_inference` as a nested loop
(outer over chunks, inner over denoising steps), mirroring
`WanAnimatePipeline.__call__`. `encode_kv_cache` is intentionally kept as a
method — it is one transformer call with a different `kv_cache_flag` mode
(cache-write), and inlining it would interleave two distinct forward
semantics in the same loop body and lose readability.
- Rename `context_sequence` → `video` (pixel-space) + `video_latents`
(pre-encoded), matching `WanVideoToVideoPipeline`. For the FAR pipeline,
the old `{"raw"/"latent"}` dict form is replaced by the two kwargs.
Mutually-exclusive validation raises `ValueError`.
Addresses review comments at pipeline_anyflow.py:358, :372, :393, :473 and
pipeline_anyflow_far.py:395, :489, :675.
* [AnyFlow] scheduler: N-length timesteps + step defaults r_timestep
Per @dg845's review on #13745:
- `set_timesteps(N)` now produces `N` timesteps backed by an internal
`sigmas[N+1]` linspace, matching `FlowMatchEulerDiscreteScheduler.set_
timesteps`. The final sigma (== 0) is the implicit r-endpoint of the last
step; the pipeline rollouts iterate `for i, t in enumerate(timesteps)`
without the old `[:-1]` slicing.
- `step(r_timestep=None)` now defaults to the next timestep on the schedule
(resolved via fp-tolerant `argmin` over `sigmas[:-1]`), instead of raising.
Any-step sampling is preserved when `r_timestep` is explicit. The raise
stays only for the case where the caller passes a `timestep` value that
isn't on the schedule and provides no `r_timestep` — there's no sensible
default in that case.
- Build sigmas in float64 on CPU then move to the target device, with a
float32 downcast for MPS / NPU (float64 isn't supported on those backends).
Pipeline rollout loops updated to compute `r = sigmas[i + 1] * num_train_
timesteps` for the model's `r_timestep` input and pass `r_timestep=None` to
`scheduler.step` (which resolves it from the schedule internally).
Addresses review comments at scheduling_flow_map_euler_discrete.py:107 and
:148.
* [AnyFlow] tests: regenerate via generate_model_tests.py; split bidi/FAR files
Per @dg845's review on #13745: replaced the hand-rolled transformer tests
with the standard mixin-based suite produced by `utils/generate_model_tests
.py`, and split the FAR causal model tests into their own file to mirror the
transformer file split.
- `tests/models/transformers/test_models_transformer_anyflow.py`: regenerated
bidi suite. Pulls in `ModelTesterMixin`, `MemoryTesterMixin`,
`TrainingTesterMixin`, `AttentionTesterMixin`, `TorchCompileTesterMixin` via
`BaseModelTesterConfig`, with `get_init_dict()` / `get_dummy_inputs()`
filled in for the small bidi config used in CI.
- `tests/models/transformers/test_models_transformer_anyflow_far.py`: new.
Same mixin set (TorchCompile is intentionally skipped — FAR's
`_build_causal_mask` uses `flex_attention.create_block_mask(_compile=False)`
which conflicts with the standard compile tester's assumptions; the bidi
file covers compile, FAR is bit-exact-validated end-to-end on H200 via the
pipeline replay). Also carries an `AnyFlowCausalAttnProcessor` smoke test
that exercises the backend gate (non-flex backends must raise) and asserts
the `AnyFlowFARTransformerOutput` dataclass exposes the expected fields.
Addresses review comments at test_models_transformer_anyflow.py:71 and :128.
* [AnyFlow] docs: update for video / video_latents kwarg rename
Following the pipeline kwarg refactor in e9d50b2, sweep the user-facing docs
to reflect the new API:
- `docs/source/en/api/pipelines/anyflow.md`: T2V / I2V / V2V code examples now
use `video=` instead of `context_sequence={"raw": ...}`. The "Generation
with AnyFlow (FAR Causal)" intro describes the new mutually-exclusive
`video` / `video_latents` selector.
- `docs/source/en/using-diffusers/anyflow.md`: the scenario selector table,
the "Image-to-video and video-to-video" walkthrough, and the closing note
about pre-encoded latents are all updated. `vae_encode` references are
replaced with `encode_video`.
* [AnyFlow] tests: skip FAR training tests on CPU (flex backward); align scheduler tests with N-length timesteps
- TestAnyFlowFARTransformer3DTraining: skip test_training / test_training_with_ema /
test_gradient_checkpointing_equivalence on CPU. FAR causal self-attn uses
torch.nn.attention.flex_attention whose backward kernel is GPU-only.
- test_scheduler_flow_map_euler_discrete: assert timesteps is N-length (not N+1) and
the sigma=0 r-endpoint lives in self.sigmas[-1]; test_step_one_shot_sampling now
exercises r_timestep=None (resolved from sigmas) since N=1 has no timesteps[1].
* [AnyFlow] docs: complete forward() Args: sections for check_forward_call_docstrings
main #13758 added utils/check_forward_call_docstrings.py which requires every signature
arg to appear as its own `name (...):` entry under Args:. Expand the bidi and FAR
transformer forward docstrings to list each parameter individually.
* [AnyFlow] apply 5/21 review suggestions (A: 1-click)
FAR transformer:
- AnyFlowCausalAttnProcessor: default _attention_backend = 'flex' (was None);
remove None from _SUPPORTED_BACKENDS. None previously fell through to SDPA
which silently ignored the BlockMask; failing loudly is the right default.
- dispatch_attention_fn call: read self._attention_backend instead of hardcoded
'flex', so '_native_flex' selection works.
- _build_freqs / _forward_full_frame: add '# Copied from' to bidi RoPE.
Pipelines:
- bidi + FAR docstrings: video shape (B, C, T, H, W) -> (B, T, C, H, W) to
match VideoProcessor.preprocess_video.
- FAR EXAMPLE_DOC_STRING: single-frame I2V tensor wrap uses unsqueeze(1) for the
T axis instead of unsqueeze(2).
- FAR encode_video: drop duplicated @torch.no_grad() decorator.
Tests:
- test_anyflow / test_anyflow_far: lift the test_save_load_optional_components
skip (the test actually passes).
- FAR processor smoke test: assert default backend is 'flex' (was 'None').
* [AnyFlow] apply 5/21 review suggestions (B: refactors)
Pipelines:
- check_inputs accepts video / video_latents and raises early on:
(a) mutual exclusion (was checked late in __call__);
(b) FAR's (num_frames - 1) % 4 == 0 constraint.
__call__ no longer carries duplicate validation.
- FAR pipeline: drop the show_progress kwarg and replace the single tqdm with
nested progress bars in the LLaDA-2 pattern — outer 'Chunks' (position=0)
and per-chunk inner 'Inference Steps' (position=1, leave=False) — both
picking up DiffusionPipeline._progress_bar_config (so set_progress_bar_config
controls them, including disable=None).
Scheduler:
- step() resolves source and target sigmas by indexing self.sigmas via the new
index_for_timestep(), instead of dividing the input timesteps by
num_train_timesteps. This keeps the math correct for any future schedule
whose timestep/sigma relationship is non-linear. For an off-schedule
r_timestep the code falls back to r / num_train_timesteps, so explicit
any-step sampling outside the schedule still works (and t off-schedule with
r=None still raises a clear ValueError, as before).
Numerical equivalence: for the shipped linspace+shift schedule the two
formulations are bit-identical (verified: max abs diff = 0.0 over an N=8,
shift=5 schedule).
* [AnyFlow] apply Claude bot review (5/21): 8 findings beyond dg845's list
Finding #1 — attention_kwargs plumbing:
Both transformers now decorate forward() with @apply_lora_scale('attention_kwargs')
(matches Wan); pipelines forward attention_kwargs to the transformer + encode_kv_cache,
and the unused parameter is dropped from the inner _forward_train / _forward_cache /
_forward_inference signatures. Pipeline docstrings updated to the standard wording.
Finding #2 — naming:
Rename far_cfg -> layout_cfg in the bidi transformer (the bidi path is not FAR; the
FAR transformer keeps far_cfg, which is accurate there).
Finding #3 — scheduler state machine:
Add _step_index, _begin_index, step_index property, begin_index property,
set_begin_index(), _init_step_index(). step() lazily initializes and advances the
counter so downstream callbacks / composable schedulers can observe rollout progress.
Sigma resolution remains a pure function of (timestep, r_timestep) — calling step()
twice with identical args still returns identical prev_sample (idempotent).
Finding #4 — redundant @torch.no_grad():
Drop the redundant decorators on bidi pipeline's encode_video and FAR pipeline's
encode_kv_cache (callers are already in __call__'s no-grad scope).
Finding #5 — dead code:
Remove the unreachable temb.ndim == 2 else branch from the bidi transformer's
output-norm path (condition_embedder.forward always returns a 3D temb).
Finding #6 — private rename:
forward_far_patchify[_inference] -> _forward_far_patchify[_inference] (only called
internally by _forward_train / _forward_cache / _forward_inference).
Finding #7 — pipeline comment numbering:
Bidi + FAR pipelines renumber steps so the # 4. slot is no longer skipped.
Finding #8 — mask-mod comment numbering:
_build_causal_mask numbered comments now run 1) 2) 3) ... (was 1) 3) 4) ...).
Tests:
- New test_step_index_advances + test_set_begin_index_anchors_step_index in the
scheduler test file exercise the new state machine.
- All existing pipeline / transformer / scheduler tests still pass (85 passed,
85 skipped on CPU).
Bit-exact: 8-step rollout vs the previous formulation, max abs diff = 0.0 (the new
sigma-lookup is byte-identical to t/num_train_timesteps on this schedule).
* [AnyFlow] scheduler: honour off-schedule any-step in _init_step_index; drop dead _resolve_next_timestep
Audit caught two issues in the previous scheduler commit:
1. The new state machine raised in _init_step_index whenever the first timestep
wasn't on the active schedule, contradicting the documented contract that
step() falls back to t/num_train_timesteps for off-schedule any-step
sampling. The fall-back numerics were intact but they were unreachable —
the init check fired first.
Fix: _init_step_index now initializes _step_index to 0 when the timestep is
off-schedule (still a valid observable counter for callbacks). step()'s
sigma resolution is untouched, so on-schedule rollouts stay bit-exact and
off-schedule any-step sampling actually runs again. Regression test:
test_step_off_schedule_anystep_supported.
2. _resolve_next_timestep had no remaining callers after the step() rewrite
inlined the same lookup. Removed (private helper, no external API).
* [AnyFlow] docs: align user guides with video shape + kwarg fixes
- en api/pipelines/anyflow.md: video shape (B, C, T, H, W) -> (B, T, C, H, W);
example tensor wrap uses unsqueeze(0).unsqueeze(1) and permute(0, 3, 1, 2)
to match VideoProcessor.preprocess_video's 5D contract.
- zh using-diffusers/anyflow.md: same shape fixes; also flip the I2V / V2V
examples from the obsolete context_sequence={...} dict to the current
video= / video_latents= kwargs; helper to_video_tensor returns (1, T, C, H, W);
add a note about mutual exclusion.
* [AnyFlow] tests: drop @slow integration test scaffolds for initial PR
.ai/skills/model-integration/SKILL.md is explicit: 'No integration / slow
tests in the initial PR — don't add anything gated on @slow / RUN_SLOW=1
yet.' Our two integration test classes were shape-only assertions with TODOs
for a future numeric reference, so dropping them loses no actual coverage —
the relevant rollouts are covered by H200 bit-exact replay outside the
pytest suite. Can land a follow-up PR after merge with proper numeric
reference slices once the maintainer is comfortable enabling slow tests.
* Apply style fixes
* [AnyFlow] apply 5/22 dg845 review: comment cleanups + custom sigmas/timesteps schedule
dg845 third pass — 7 of 9 comments applied; the 8th (custom sigmas/timesteps support)
matches FlowMatchEulerDiscreteScheduler conventions; the 9th (_build_causal_mask
refactor) is explicitly marked non-blocking and deferred to a follow-up that also
re-enables TorchCompileTesterMixin.
Comment cleanups:
- transformer_anyflow.py:704 temb output-norm comment: drop redundant 'no ndim==2 branch'.
- pipeline_anyflow.py:550 denoise loop comment: '# 6. Denoising loop'.
- pipeline_anyflow_far.py:684 denoise loop comment: '# 8. Denoising loop (outer over
chunks, inner over timesteps).'.
- pipeline_anyflow_far.py:702 drop trailing inline comment on `timesteps = scheduler.timesteps`.
- scheduling_flow_map_euler_discrete.py: clearer wording on the off-schedule `r_timestep`
error.
Custom schedule support:
- FlowMapEulerDiscreteScheduler.set_timesteps gains `sigmas` and `timesteps` kwargs
mirroring FlowMatchEulerDiscreteScheduler. Default behaviour is unchanged
(linspace + shift); the validation + length-N → length-N+1 terminal-0 append are
shared with the default path so on-schedule rollouts stay bit-exact.
- AnyFlowPipeline.__call__ and AnyFlowFARPipeline.__call__ accept `sigmas` and
`timesteps` kwargs, override num_inference_steps from their length, and forward
to set_timesteps (matches LTX2Pipeline pattern).
- New scheduler tests: test_set_timesteps_custom_sigmas and
test_set_timesteps_custom_timesteps cover both override paths.
Dtype skip on save/load:
- TestAnyFlowTransformer3D and TestAnyFlowFARTransformer3D now skip
test_from_save_pretrained_dtype_inference (parametrized over fp16/bf16), mirroring
WanTransformer3DModel's skip — the test's tolerance requirements are too high for
meaningful signal under AnyFlow's flow-map mixed-precision sampling.
* [AnyFlow] docs: apply hf-doc-builder line wrap (max_len 119)
CI doc-builder style check flagged 3 files with docstring lines >119 chars.
Ran 'doc-builder style src/diffusers docs/source --max_len 119' to autoformat;
content unchanged, line wrapping only.
* [AnyFlow] apply 5/22 follow-up review: new_zeros terminal sigma + cleanup
dg845 blocking suggestion (r3287274209):
- scheduling_flow_map_euler_discrete.py:185 — use `working_sigmas.new_zeros(1)`
instead of `torch.zeros(1, dtype=...)` so the appended terminal sigma inherits
both device and dtype from working_sigmas. The current working_sigmas always
starts on CPU so the device mismatch is latent, but new_zeros is the correct
defensive pattern and matches how the published FAR test fixtures run on CUDA.
Claude bot final-review follow-ups:
- transformer_anyflow_far.py: drop three stale `# step 3: generate attention mask`
comments left over from the original numbered-step structure (bot #6).
- pipeline_anyflow_far.py: annotate `encode_video` with
`# Copied from diffusers.pipelines.anyflow.pipeline_anyflow.AnyFlowPipeline.encode_video`
and align docstring + inline comment so `make fix-copies` keeps them in sync (bot #3).
Skipped (not real / judgment-call):
- bot #2 (private rename of `_forward_far_patchify*`) — already done in 84605d5;
bot was looking at a stale snapshot.
- bot #4 (check_inputs `# Copied from`) — FAR's check_inputs has an extra
`(num_frames - 1) % 4 == 0` constraint that doesn't map onto the bidi version,
so a clean `# Copied from` link would require restructuring. Bot called it a
consistency nit; leaving as-is.
- bot #5 (`encode_kv_cache` → `_encode_kv_cache`) — bot itself flagged this as
judgment-call territory; the helper is a coherent operation that advanced
inference callers may want to invoke directly.
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
|
||
MQ
|
10302496a6 |
[feat] JoyAI-JoyImage-Edit support (#13444)
* [feat] JoyAI-JoyImage-Edit support * [fix] remove rearrange * [refactor] two pass when do cfg * [refactor] remove repa, use wantimetextembeding, refactor modulate code * [refactor] Joyimage Attention refactor * remove vae tiling and autocast * [fix] remove einops from setup.py * [refactor] Refactor JoyImageEditPipeline to use explicit arguments instead of namespace and remove _build_arg * [fix] remove deprecated method decode_latents * [refactor] refactor the image pre-processing logic into a separate VaeImageProcessor subclass * [refactor] add JoyImageAttention to align with Attention + AttnProcessor design and update conversion script for new weight key mapping (e.g. img_attn_qkv -> attn.img_attn_qkv) * [refactor] simplify bucket logic in JoyImageEditImageProcessor by replacing runtime generation with precomputed lookup tables * [fix] remove leftover training-only parameters * [fix] add layerwise casting and fp32 module patterns to JoyImageTransformer3DModel. Reference WanTransformer3DModel to fix layer casting errors during inference. * [test] add JoyImageEditPipeline fast tests and JoyImageEditTransformer3DModel model tests * [fix] fix some pipeline args to support batch inference * [fix] duplicate images to match batch size when fewer images than prompts in JoyImageEditPipeline * [fix] remove no longer used config parameters * Apply style fixes * [fix] remove unused dataclass and rewrite helpers as inline functions * [fix] make dummy objects for JoyImageEdit * [fix] allow test_torch_compile_repeated_blocks to pass * [fix] add examples on JoyImageEditPipeline * fix code style issues with ruff and black * Apply style fixes * [fix] change default num_inference_steps to 40 * [fix] use forward hook to extract pre-norm hidden states for transformers 5.x compatibility * [fix] change the assert to ValueError in pipeline * [fix] rename JoyImageTransformer3DModel to JoyImageEditTransformer3DModel, clean up anything about the alias * [fix] support gradient checkpointing * [refactor] simplify RoPE utilities, inline helpers, copy WanTimeTextImageEmbedding locally and remove unused parameters * [fix] remove _get_text_encoder_ckpt and qwen_processor * [fix] change nn.RMSNorm to FP32LayerNorm * [fix] small fixes for suggestions given by Claude * [refactor] build model using from _pretained instead of config * [refactor] auto-wrap prompt and support text-to-image in JoyImage Edit pipeline * make style, make quality and make fix-copies * [test] small fix to use vocab_size=1024 * [refactor] separate encode_prompt_multiple_images from encode_prompt, support prompt_embeds/prompt_embesd_mask/num_images_per_prompt in edit mode * [test] fix CI: use strict=False for xfail and add @require_torch_accelerator to group offloading test * [refactor] separate image_latents from latents in prepare_latents to align with flux2 * make style --------- Co-authored-by: zhangmaoquan.1 <zhangmaoquan.1@jd.com> Co-authored-by: huangfeice <huangfeice@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
Gong Junmin
|
1a8a17b71b |
Add ACE-Step pipeline for text-to-music generation (#13095)
* Add ACE-Step pipeline for text-to-music generation
Rebased on origin/main from the original pr-13095 branch (3 commits squashed).
- AceStepDiTModel: Diffusion Transformer with RoPE, GQA, sliding window,
AdaLN timestep conditioning, and cross-attention.
- AceStepConditionEncoder: fuses text / lyric / timbre into a single
cross-attention sequence.
- AceStepPipeline: text2music / cover / repaint / extract / lego / complete.
- Conversion script for the original checkpoint layout.
- Docs + tests.
* Fix ACE-Step pipeline audio quality and auto-detect turbo/base/sft variants
The PR's original inference produced low-quality audio on turbo because the
pipeline (a) mangled the SFT prompt format, (b) applied classifier-free guidance
with the wrong unconditional embedding (empty-string encoded vs. the learned
`null_condition_emb`), and (c) hardcoded turbo defaults even when loading a
base/SFT checkpoint.
Changes:
* Converter preserves `null_condition_emb` (stored under the condition encoder)
and propagates `is_turbo`/`model_version` into the transformer config so the
pipeline can route per-variant defaults.
* `AceStepConditionEncoder` registers `null_condition_emb` as a learned
parameter matching the original module.
* Pipeline auto-detects variant via `is_turbo`/`model_version` and picks
defaults that match `acestep/inference.py`:
* turbo: steps=8, shift=3.0, guidance_scale=1.0 (no CFG)
* base/SFT: steps=27, shift=1.0, guidance_scale=7.0
* Base/SFT timestep schedule uses the linear+shift transform from
`acestep/models/base/modeling_acestep_v15_base.py`; turbo still uses the
hardcoded 8-step `SHIFT_TIMESTEPS` table.
* CFG reuses the learned `null_condition_emb` and batches the
conditional+unconditional forwards into a single transformer call.
* `SFT_GEN_PROMPT` matches the newline layout in `acestep/constants.py` so the
text encoder sees the same prompt distribution it was trained on.
DiT parity vs. the original ACE-Step 1.5 turbo DiT is bit-identical
(max_abs=0.0 in fp32 eager/SDPA across 4 seed/shape cases) — see
scripts/dit_parity_test.py.
* Add ACE-Step parity test scripts
Two developer-facing parity harnesses live under scripts/:
* dit_parity_test.py — loads the same converted turbo weights into the
original AceStepDiTModel and the diffusers AceStepDiTModel, drives
identical (hidden_states, timestep, timestep_r, encoder_hidden_states,
context_latents) inputs, and asserts max-abs-diff ≤ 1e-5 in fp32
eager/SDPA. Currently passes bit-identical (max_abs=0) across four
shape/seed cases including batched + odd-length paths.
* audio_parity_jieyue.py — full end-to-end audio parity. Given the same
JSON example, runs both the original ACE-Step 1.5 pipeline and the
diffusers AceStepPipeline at matched seed/precision (bf16 + FA2 by
default) and saves side-by-side .wav files for listening verification.
Supports text2music / cover / repaint × turbo / base / sft via a
--matrix mode that writes 18 wavs named
{variant}_{task}_{official,diffusers}.wav.
* Route SFT parity to acestep-v15-sft checkpoint
On jieyue the release tree has a dedicated SFT checkpoint at
checkpoints/acestep-v15-sft with its own modeling_acestep_v15_base.py
shipped under acestep/models/sft/. Point the SFT row of the parity matrix
at that checkpoint / module so we're testing the actual SFT weights, not
the plain base ones.
* audio_parity_jieyue: fix doubled 'acestep-' in cache path; --converted-root flag
Previously the converted-pipeline cache dir was
`/tmp/acestep-<variant>-diffusers` but <variant> already starts with
"acestep-", giving `/tmp/acestep-acestep-v15-turbo-diffusers`. Drop the
prefix.
On jieyue the overlay rootfs (including /tmp) only has a few GB free; a
full turbo conversion needs ~5 GB per variant. Add --converted-root (env
ACESTEP_CONVERTED_ROOT) so the cache can live on vepfs.
* audio_parity_jieyue: two-phase matrix bootstraps cover/repaint from text2music
The ACE-Step release bundle on jieyue doesn't ship sample .wav/.mp3
files, so matrix mode had no default --src-audio and would skip
cover/repaint entirely. Run text2music first for every variant, then
reuse the TURBO official text2music output as the shared source for the
cover/repaint rows. Users can still override with --src-audio.
* audio_parity_jieyue: seed the diffusers generator on the pipeline device
The ORIGINAL ACE-Step pipeline seeds on the execution device
(`torch.Generator(device=device).manual_seed(seed)`), i.e. the CUDA RNG
stream when running on GPU. Previously the parity harness seeded the
diffusers side with a CPU generator, so even though the seed integer
matched, the two sides drew different noise from the outset and the
final outputs were essentially uncorrelated. Use the execution-device
generator on both sides for a fair comparison.
* Fix ACE-Step pipeline: switch to APG guidance + peak normalization
Two issues found after the first jieyue audio parity run:
1. The original base/SFT pipeline uses APG (Adaptive Projected Guidance,
acestep/models/common/apg_guidance.py) with a stateful momentum
buffer and norm/projection steps — NOT vanilla CFG. Using vanilla CFG
produced uncorrelated outputs vs. the reference (pearson ~0.0 on
20 s samples); this PR ports `_apg_forward` + `_APGMomentumBuffer`
and plugs them into the denoising loop when `guidance_scale > 1`.
Momentum is instantiated once per pipeline call (persists across
denoising steps) to match the reference semantics.
2. The post-VAE "anti-clipping normalization" in this pipeline was
`audio /= std * 5` with a `std<1 -> std=1` guard. The original
post-processing in
acestep/core/generation/handler/generate_music_decode.py is simple
peak normalization: `if audio.abs().max() > 1: audio /= peak`. The
std-based proxy both (a) let clips with peak < 1 leak through
unchanged (over-quiet) and (b) failed to bring clipping peaks to
exactly 1 in a bunch of base/SFT cases (observed max=1.000, std=0.200
repeatedly in the first parity run). Switch to peak normalization on
both sides.
Tested via scripts/audio_parity_jieyue.py on A800; re-run pending to
confirm the base/SFT correlation improvements.
* Fix ACE-Step chunk mask values to match the original pipeline
The DiT receives `context_latents = concat(src_latents, chunk_mask)` on the
channel dim, and was trained with chunk_mask values drawn from the three
sentinels documented in acestep/inference.py:
2.0 -> model-decided (default for text2music / cover / full-generation)
1.0 -> keep this latent frame from src_latents (repaint preserved region)
0.0 -> explicitly repaint this frame (only inside the repaint window)
Previously _build_chunk_mask returned all-1.0 for text2music (and cover /
lego), and an inverted 0/1 mask for repaint (1 inside the window, 0 outside).
Either case puts context_latents out of distribution. Switch text2music /
cover to the 2.0 sentinel and flip the repaint mask so it's 1.0 outside /
0.0 inside. Update the repaint src_latents zero-out to multiply by the new
mask (was `1 - chunk_mask`) so the zero region still lines up with the
repaint window.
* Add direct invoker for ACE-Step generate_music (ground truth)
Our earlier audio_parity_jieyue.py reconstructs the original pipeline by
calling AceStepConditionGenerationModel.generate_audio() directly, which
silently skips a lot of the real handler plumbing (conditioning masks,
silence-latent tiling, cover/repaint pre-processing, etc.). That made the
'official' wavs we saved sound wrong — flat, drone-like, not music.
This new script calls acestep.inference.generate_music end-to-end through
the real AceStepHandler, with LM + CoT explicitly disabled so we still have
a deterministic comparison. Use it to generate the ground-truth 'official'
wav for a given JSON example, then separately run the diffusers pipeline
with the same inputs and diff the two.
* run_official_generate_music: call initialize_service to bind a DiT variant
AceStepHandler() is a shell — you have to call handler.initialize_service(
project_root=..., config_path=..., device=..., use_flash_attention=..., ...)
before generate_music will work. Mirror what cli.py does at the equivalent
spot (around cli.py:1400).
* Fix silence-reference for ACE-Step timbre encoder
The root cause for the flat / drone-like outputs I was seeing (including
in my 'official' reconstruction): when no reference_audio is provided the
pipeline was feeding literal zeros to the timbre encoder. The real
handler feeds a slice of the learned `silence_latent` tensor.
The handler also transposes silence_latent on load (see
acestep/core/generation/handler/init_service_loader.py:214:
self.silence_latent = torch.load(...).transpose(1, 2)
) converting [1, 64, 15000] -> [1, 15000, 64] so that
`silence_latent[:, :750, :]` yields the expected [1, 750, 64] shape.
Changes:
* Converter: load silence_latent.pt, transpose to [1, T, C], bake into
the condition_encoder safetensors under key `silence_latent`.
(Also keeps the raw .pt file at the pipeline root for debugging.)
* AceStepConditionEncoder: register `silence_latent` as a persistent
buffer so from_pretrained loads it alongside the trained weights.
* Pipeline: when reference_audio is None, slice
`condition_encoder.silence_latent[:, :timbre_fix_frame, :]` and
broadcast across the batch instead of zeros. Emits a loud warning
(and falls back to zeros) if the buffer is all-zero — that means the
checkpoint was produced by an older converter and should be rebuilt.
* audio_parity_jieyue.py: the reference path now matches the handler's
silence-latent slicing.
Without this fix, every variant/task combo produced drone-like audio
even when my numeric DiT-forward parity claimed they were identical.
* Fix three more ACE-Step pipeline bugs I found by dumping real inputs
Instrumented the live generate_audio call in the real ACE-Step handler and
observed the exact tensors it sees — my diffusers pipeline was wrong in
three independent ways:
1. src_latents for text2music should be silence_latent tiled to
latent_length, NOT zeros. The handler fills no-target cases from
silence_latent_tiled (observed std=0.96). Zeros are OOD for the DiT
context_latents concat and produce drone-like outputs.
2. chunk_mask values cap at 1.0 (not 2.0). The handler starts with a
bool tensor (True inside the generate span, False outside); the
chunk_mask_modes=auto -> 2.0 override does NOT take effect because
the underlying tensor is bool, so setting entry = 2.0 casts to True.
After the later .to(dtype) float cast, the DiT sees 1.0/0.0 — exactly
what I observed in the captured tensor (unique values = [True]).
3. Default shift is 1.0 for ALL variants, including turbo. I was
defaulting turbo to shift=3.0 which picks a different SHIFT_TIMESTEPS
table (the 8-step schedule is keyed by shift, not variant).
Also:
* Added _silence_latent_tiled() helper that slices / tiles the learned
silence_latent (now loaded as a buffer on the condition encoder) to
the requested latent length.
* Repaint path now substitutes silence_latent (not raw zeros) inside
the repaint window — matches conditioning_masks.py.
* audio_parity_jieyue.py mirrors the same src/chunk/shift choices on
its 'original' leg for apples-to-apples parity once the buggy
reconstruction is removed from the picture.
* Add peak+loudness post-normalization to AceStepPipeline
The real pipeline normalizes audio in two stages (see
acestep/audio_utils.py:72 normalize_audio + generate_music_decode.py):
1. if peak > 1: audio /= peak (anti-clip)
2. audio *= target_amp / peak (target_amp = 10 ** (-1/20) ~ 0.891)
Step 2 is loudness normalization to -1 dBFS. Without it diffusers outputs
had peak=1.0 vs the real 0.891 — same music content (pearson was ~0.86
already), just 1.12x louder. Add step 2 after the existing anti-clip step.
* Match acestep/inference.py inference_steps=8 for ALL variants
GenerationParams.inference_steps default is 8 — turbo AND base/SFT. I had
base/SFT defaulting to 27 here, so every base/SFT parity run was comparing
a 27-step diffusers trajectory against an 8-step real trajectory. Different
number of denoising steps means different audio even at fixed seed.
This likely explains the lower base/SFT correlation in my earlier jieyue
runs (turbo was 0.86, base/SFT were 0.32-0.34). Aligning step counts
should bring base/SFT closer to turbo parity.
* Address PR #13095 review: rename classes + reuse diffusers primitives
Response to dg845's PR comments batch 1+2. DiT parity harness still bit-identical
(max_abs=0 on fp32 / SDPA across 4 shape cases).
Transformer file:
* Rename AceStepDiTModel -> AceStepTransformer1DModel (alias kept).
* Rename AceStepDiTLayer -> AceStepTransformerBlock (alias kept).
* Inherit AttentionMixin + CacheMixin on the DiT model.
* Swap in diffusers.models.normalization.RMSNorm for the hand-rolled
AceStepRMSNorm (weight-key-compatible).
* Swap the hand-rolled rotary embedding + apply_rotary for diffusers'
get_1d_rotary_pos_embed + apply_rotary_emb (use_real_unbind_dim=-2 to
match the cat-half convention ACE-Step inherits from Qwen3).
* Use get_timestep_embedding with flip_sin_to_cos=True — keeps the
(cos, sin) ordering of the original sinusoidal. State-dict-compatible.
* Drop max_position_embeddings arg from DiT config (RoPE computes freqs
per call based on seq_len); converter drops it.
* Gradient-checkpoint call now takes just the layer module (matches the
Flux2 idiom).
Pipeline modeling file (pipelines/ace_step/modeling_ace_step.py):
* Moved _pack_sequences + AceStepEncoderLayer here — they aren't used
by the DiT, so they shouldn't live in the transformer file.
* AceStepLyricEncoder + AceStepTimbreEncoder set
_supports_gradient_checkpointing = True and wrap encoder-layer calls
through the checkpointing func when enabled.
* Use diffusers RMSNorm + the RoPE helper from the transformer file
(shared single implementation).
Converter (scripts/convert_ace_step_to_diffusers.py):
* model_index.json now carries AceStepTransformer1DModel.
* Drop max_position_embeddings / use_sliding_window from the emitted
configs.
No numerical regressions: scripts/dit_parity_test.py PASSES with
max_abs=0.0 on fp32/SDPA across short, long, batched, and
padding-path shape variants.
* Address PR #13095 review: pipeline polish + converter HF-hub support
Response to dg845 review comments on the pipeline side. DiT parity still
bit-identical (max_abs=0 across 4 shape cases).
Pipeline (pipelines/ace_step/pipeline_ace_step.py):
* Add `sample_rate` + `latents_per_second` properties sourced from the
VAE config so the pipeline no longer hardcodes 48000 / 25 / 1920.
Propagates through prepare_latents, chunk_mask window math, and the
audio-duration round-trip.
* Add `do_classifier_free_guidance` property (matches LTX2 et al.).
* Add `check_inputs(...)` called from `__call__` before allocating noise.
Validates prompt type, lyrics type, task_type, step count, guidance
scale, shift, cfg interval bounds and repaint window ordering.
* Add `callback_on_step_end` + `callback_on_step_end_tensor_inputs` —
the modern callback form. The legacy `callback` / `callback_steps`
pair is kept for back-compat. Setting `pipe._interrupt = True` inside
the callback stops the loop early.
* Expose `encode_audio(audio)` as a public helper that wraps the tiled
VAE encode + (B, T, D) transpose the pipeline performs internally.
Converter (scripts/convert_ace_step_to_diffusers.py):
* Accept a Hugging Face Hub repo id for `--checkpoint_dir`; resolves it
via `huggingface_hub.snapshot_download` when the argument isn't a
local path.
Exports:
* Register `AceStepTransformer1DModel` in the top-level __init__,
models/__init__, models/transformers/__init__, and dummy_pt_objects so
`from diffusers import AceStepTransformer1DModel` works and the
pipeline loader resolves the new class name from model_index.json.
Deferred for a follow-up (commented inline in the PR): full
`Attention + AttnProcessor + dispatch_attention_fn` refactor and
`FlowMatchEulerDiscreteScheduler` migration — both would benefit from a
dedicated parity re-run and review.
* Fix stale ACE-Step 1.0-era docs / class names in the 1.5 integration
Docs and docstrings still carried a mix of 1.0 paper title, non-existent
`ACE-Step/ACE-Step-v1-5-turbo` hub id, `shift=3.0` turbo default, and
the old `AceStepDiTModel` class name. Cleaned up to match the actual
1.5 release:
* pipelines/ace_step.md: correct citation title ("ACE-Step 1.5: Pushing
the Boundaries of Open-Source Music Generation"), correct repo
(`ace-step/ACE-Step-1.5`), new variants table with real HF ids
(`Ace-Step1.5` / `acestep-v15-base` / `acestep-v15-sft`) and their
per-variant step/CFG defaults, drop the wrong `shift=3.0` tip.
* models/ace_step_transformer.md: page renamed to
`AceStepTransformer1DModel` with a short 1.5-specific description;
`AceStepDiTModel` noted as a backwards-compat alias.
* pipeline_ace_step.py: import, docstring, `Args`, and `__init__`
annotation reference `AceStepTransformer1DModel`; example model id
now `ACE-Step/Ace-Step1.5`; `_variant_defaults` docstring and the
`__call__` variant-fallback comment no longer claim `shift=3.0` /
`27 steps` — real defaults are 8 steps / shift=1.0 across all
variants, guidance=1.0 (turbo) vs 7.0 (base+sft).
* Address PR #13095 review: VAE tiling on AutoencoderOobleck + Timesteps class
Two more deferred review threads from dg845 addressed:
* Move tiled encode/decode onto AutoencoderOobleck
(https://github.com/huggingface/diffusers/pull/13095#discussion_r2785513647).
AutoencoderOobleck now carries `use_tiling` + `tile_sample_min_length` /
`tile_sample_overlap` / `tile_latent_min_length` / `tile_latent_overlap`
attributes and private `_tiled_encode` / `_tiled_decode` methods; the
existing `encode` / `_decode` dispatch to them when tiling is enabled and
the input exceeds the threshold. `AutoencoderMixin.enable_tiling()` is
already inherited.
AceStepPipeline's private `_tiled_encode` / `_tiled_decode` and the
`use_tiled_decode` `__call__` arg are gone; `__init__` now calls
`self.vae.enable_tiling()` so the long-audio memory behaviour is preserved
by default. Users can opt out with `pipe.vae.disable_tiling()`.
Note: the VAE-side tiling concatenates encoder features (h) and samples
the posterior once, instead of the old per-tile `.sample()` calls. This
is the standard diffusers pattern; numerically differs only in the
structure of the noise across tile boundaries.
* Use the Timesteps nn.Module for the sinusoid
(https://github.com/huggingface/diffusers/pull/13095#discussion_r2785420234).
`AceStepTimestepEmbedding` wraps `Timesteps(in_channels, flip_sin_to_cos=
True, downscale_freq_shift=0)` instead of calling `get_timestep_embedding`
directly — reviewer asked for the Module form.
* Address PR #13095 review: refactor AceStepAttention to Attention + AttnProcessor
Splits the monolithic AceStepAttention into the diffusers standard
Attention + AttnProcessor layout:
- AceStepAttention (torch.nn.Module, AttentionModuleMixin) holds the
to_q/to_k/to_v/to_out projections and norm_q/norm_k RMSNorms.
- AceStepAttnProcessor2_0 runs the attention dispatch through
dispatch_attention_fn so users can pick flash / sage / native backends
via model.set_attention_backend(...) or the attention_backend context
manager.
GQA (Q has 16 heads / K,V have 8) is preserved by passing enable_gqa=True
to dispatch_attention_fn instead of repeat_interleave; fusion is disabled
(_supports_qkv_fusion = False) because Q and K,V have different output
sizes.
The converter is updated to rename the six attention sub-keys
(q_proj -> to_q, k_proj -> to_k, v_proj -> to_v, o_proj -> to_out.0,
q_norm -> norm_q, k_norm -> norm_k) on both the DiT decoder path and the
condition encoder path, since AceStepLyricEncoder / AceStepTimbreEncoder
share the same AceStepAttention class.
Addresses review comments r2785433213 and r2785450463.
* Address PR #13095 review: migrate to FlowMatchEulerDiscreteScheduler
Replace the hand-rolled flow-matching Euler loop with
`FlowMatchEulerDiscreteScheduler`. ACE-Step still computes its own shifted /
turbo sigma schedule via `_get_timestep_schedule`, but now passes it to
`scheduler.set_timesteps(sigmas=...)` and delegates the ODE step to
`scheduler.step()`. The scheduler is configured with `num_train_timesteps=1`
and `shift=1.0` so `scheduler.timesteps` stays in `[0, 1]` (the convention the
DiT was trained on) and the scheduler doesn't re-shift already-shifted sigmas.
The scheduler's appended terminal `sigma=0` reproduces the old loop's
final-step "project to x0" case exactly: `prev = x + (0 - t_curr) * v`.
Parity on jieyue (seed=42, bf16 + flash-attn, turbo text2music, 8 steps):
waveform Pearson = 0.999999
spectral Pearson = 1.000000
max |diff| = 2.5e-3 (fp32 step-math vs previous bf16 step-math)
fp32 Euler-loop A/B against the hand-rolled path: max |diff| = 3.6e-7.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Address PR #13095 review: move DiT tests + drop stale test kwargs
- Move the DiT transformer tests out of the pipeline test file into a new
tests/models/transformers/test_models_transformer_ace_step.py that follows
the standard BaseModelTesterConfig + ModelTesterMixin scaffold (matches
test_models_transformer_longcat_audio_dit.py).
- Drop `max_position_embeddings` from the remaining AceStepDiTModel and
AceStepConditionEncoder test fixtures — neither constructor accepts that
argument anymore.
- Drop `use_sliding_window` from the same fixtures — also no longer a
constructor argument (the actual `sliding_window` int kwarg is kept).
- Wire `FlowMatchEulerDiscreteScheduler(num_train_timesteps=1, shift=1.0)`
into `get_dummy_components()` now that the pipeline requires it.
Resolves https://github.com/huggingface/diffusers/pull/13095#discussion_r3115653554,
r3115664850, r3115673059, r3115676580, r3115680700.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Address PR #13095 review from dg845 (2026-04-23)
Fixes 5 review threads + style:
1. Converter now builds `AceStepPipeline` in memory and calls
`save_pretrained`. Previously the hand-written `model_index.json` was
missing the `scheduler` entry — fresh converter output couldn't be loaded
by `AceStepPipeline.from_pretrained` (r3127767785). This also makes the
converter robust to future `__init__` signature changes.
2. `latent_length` uses `math.ceil(...)` instead of `int(...)` so non-integer
products (e.g. `latents_per_second=2.0, audio_duration=0.4 → 0.8`) round up
to `1` instead of truncating to `0` and crashing shape checks (r3127790939).
3. Add `_callback_tensor_inputs = ["latents"]` on `AceStepPipeline` so the
standard diffusers callback tests pick up the right tensor (r3127795954).
4. `AceStepConditionEncoder.silence_latent` no longer hard-codes the channel
dim to 64. The placeholder buffer now uses the `timbre_hidden_dim`
constructor argument, so smaller test configs with `timbre_hidden_dim != 64`
load without shape errors (r3127812932).
5. Revert `self.vae.enable_tiling()` from `AceStepPipeline.__init__`. Users can
call `pipe.vae.enable_tiling()` themselves for long-form generation; that
matches the opt-in convention used by the rest of diffusers (r3127777296).
6. `ruff check --fix` + `ruff format` over all ACE-Step sources (the style fix
dg845 asked for via `@bot /style`).
Also: converter now accepts sharded `model.safetensors.index.json` layouts
alongside the single-file `model.safetensors`, so the 5B XL turbo variant
converts without a pre-processing step.
Parity on jieyue (seed=42, bf16 + flash-attn, turbo text2music 160s, fresh
converter output loaded via `from_pretrained`):
waveform Pearson = 0.999954
spectral Pearson = 0.999977
max |a-b| bf16 = 4.3e-02 (dominated by the VAE tiling default flip)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Address PR #13095 review from yiyixuxu (2026-04-23)
Code-level (22 threads):
1. Delete 3 dev/parity scripts (`scripts/audio_parity_jieyue.py`,
`scripts/dit_parity_test.py`, `scripts/run_official_generate_music.py`)
that shouldn't have been committed.
2. Rename `AutoencoderOobleck._encode_one` → `_encode` to match the convention
used by other diffusers VAEs.
3. Delete the hard-coded `SHIFT_TIMESTEPS` / `VALID_SHIFTS` table in
`pipeline_ace_step.py`: the per-shift turbo schedules are recovered
exactly by `linspace(1, 0, N+1)[:-1]` plus the flow-match shift formula
that the non-turbo branch already uses, so a single code path covers both.
4. Drop the backwards-compat `AceStepDiTModel` / `AceStepDiTLayer` aliases
and every reference (top-level `__init__`, `models/__init__`,
`transformers/__init__`, dummy objects, tests, docs toctree, model card).
`AceStepTransformer1DModel` is the only exported name now.
5. Remove the unused `attention_mask` / `encoder_attention_mask` args from
`AceStepTransformer1DModel.forward`; the model rebuilds its masks from
the sequence shape and never consumed them.
6. In the DiT forward and both encoders, pass `None` instead of an all-zero
`full_attn_mask` / `encoder_4d_mask` to non-sliding attention layers — SDPA
dispatches to a faster kernel when the mask is None.
7. Inline the shared `_run_encoder_layers` helper directly into
`AceStepLyricEncoder.forward` / `AceStepTimbreEncoder.forward` so layer
calls are visible at the forward boundary (diffusers style).
8. Move `is_turbo` / `sample_rate` / `latents_per_second` from `@property`s
that re-read module configs each call to cached attributes populated in
`__init__` (Flux2-style), with a default-ACE-Step fallback when
`self.vae` is offloaded. Drop the now-unused `SAMPLE_RATE = 48000`
module-level constant and the three property definitions.
9. Warn + coerce `guidance_scale` to 1.0 on turbo (guidance-distilled)
checkpoints, following `pipeline_flux2_klein`. Prevents over-guided
audio when users forward their base/sft CFG settings to a turbo pipe.
10. Remove the `logger.warning(...)` paths that triggered on
`silence_latent` missing/zero — those only fired for author-side
unconverted checkpoints and tests; end users always load converted
weights where the buffer is baked in.
11. Drop the redundant `with torch.no_grad():` wrappers inside
`encode_prompt` — the pipeline's `__call__` runs under `torch.no_grad`
already.
12. Strip "reviewer comment on PR #13095" attribution comments from three
docstrings (here and everywhere).
Parity on jieyue (seed=42, bf16 + flash-attn, XL turbo 160s text2music):
waveform Pearson = 0.9747
spectral Pearson = 0.9895
The shift comes from full-attention layers switching `attn_mask=0_tensor` →
`attn_mask=None`, which dispatches to a different SDPA kernel on bf16. The
two outputs are algebraically equivalent for fp32 eager; on bf16+FA the
delta is dominated by kernel-level ULPs, well within the sampler-noise
band (ear-check on the 160s example confirms no audible regression).
Still open — AudioTokenizer/Detokenizer (deferred) + APG guider follow-up
(dims differ from `diffusers.guiders.adaptive_projected_guidance`, not a
drop-in; worth a separate PR).
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Address ACE-Step audio token and APG review
* Fix ACE-Step docs CI
* Address ACE-Step pipeline cleanup review
* Fix ACE-Step flash attention sliding windows
* Add ACE-Step callback properties
* Address ACE-Step final review comments
---------
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
|
||
Lancer
|
c507097764 |
[Bugfix] Fix shape mismatch in LongCatAudioDiTTransformer conversion (#13494)
Signed-off-by: Lancer <maruixiang6688@gmail.com> |
||
Lancer
|
c41a3c3ed8 |
[Feat] Adds LongCat-AudioDiT pipeline (#13390)
* Add LongCat-AudioDiT pipeline Signed-off-by: Lancer <maruixiang6688@gmail.com> * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * upd * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * Apply style fixes * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * Apply style fixes * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> * Apply style fixes * upd Signed-off-by: Lancer <maruixiang6688@gmail.com> --------- Signed-off-by: Lancer <maruixiang6688@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
dg845
|
072d15ee42 |
Add Support for LTX-2.3 Models (#13217)
* Initial implementation of perturbed attn processor for LTX 2.3 * Update DiT block for LTX 2.3 + add self_attention_mask * Add flag to control using perturbed attn processor for now * Add support for new video upsampling blocks used by LTX-2.3 * Support LTX-2.3 Big-VGAN V2-style vocoder * Initial implementation of LTX-2.3 vocoder with bandwidth extender * Initial support for LTX-2.3 per-modality feature extractor * Refactor so that text connectors own all text encoder hidden_states normalization logic * Fix some bugs for inference * Fix LTX-2.X DiT block forward pass * Support prompt timestep embeds and prompt cross attn modulation * Add LTX-2.3 configs to conversion script * Support converting LTX-2.3 DiT checkpoints * Support converting LTX-2.3 Video VAE checkpoints * Support converting LTX-2.3 Vocoder with bandwidth extender * Support converting LTX-2.3 text connectors * Don't convert any upsamplers for now * Support self attention mask for LTX2Pipeline * Fix some inference bugs * Support self attn mask and sigmas for LTX-2.3 I2V, Cond pipelines * Support STG and modality isolation guidance for LTX-2.3 * make style and make quality * Make audio guidance values default to video values by default * Update to LTX-2.3 style guidance rescaling * Support cross timesteps for LTX-2.3 cross attention modulation * Fix RMS norm bug for LTX-2.3 text connectors * Perform guidance rescale in sample (x0) space following original code * Support LTX-2.3 Latent Spatial Upsampler model * Support LTX-2.3 distilled LoRA * Support LTX-2.3 Distilled checkpoint * Support LTX-2.3 prompt enhancement * Make LTX-2.X processor non-required so that tests pass * Fix test_components_function tests for LTX2 T2V and I2V * Fix LTX-2.3 Video VAE configuration bug causing pixel jitter * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Refactor LTX-2.X Video VAE upsampler block init logic * Refactor LTX-2.X guidance rescaling to use rescale_noise_cfg * Use generator initial seed to control prompt enhancement if available * Remove self attention mask logic as it is not used in any current pipelines * Commit fixes suggested by claude code (guidance in sample (x0) space, denormalize after timestep conditioning) * Use constant shift following original code --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Junsong Chen
|
c6f72ad2f6 |
add ltx2 vae in sana-video; (#13229)
* add ltx2 vae in sana-video; * add ltx vae in conversion script; * Update src/diffusers/pipelines/sana_video/pipeline_sana_video.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * condition `vae_scale_factor_xxx` related settings on VAE types; * make the mean/std depends on vae class; --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
Ando
|
8ec0a5ccad |
feat: implement rae autoencoder. (#13046)
* feat: implement three RAE encoders(dinov2, siglip2, mae) * feat: finish first version of autoencoder_rae * fix formatting * make fix-copies * initial doc * fix latent_mean / latent_var init types to accept config-friendly inputs * use mean and std convention * cleanup * add rae to diffusers script * use imports * use attention * remove unneeded class * example traiing script * input and ground truth sizes have to be the same * fix argument * move loss to training script * cleanup * simplify mixins * fix training script * fix entrypoint for instantiating the AutoencoderRAE * added encoder_image_size config * undo last change * fixes from pretrained weights * cleanups * address reviews * fix train script to use pretrained * fix conversion script review * latebt normalization buffers are now always registered with no-op defaults * Update examples/research_projects/autoencoder_rae/README.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * use image url * Encoder is frozen * fix slow test * remove config * use ModelTesterMixin and AutoencoderTesterMixin * make quality * strip final layernorm when converting * _strip_final_layernorm_affine for training script * fix test * add dispatch forward and update conversion script * update training script * error out as soon as possible and add comments * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * use buffer * inline * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * remove optional * _noising takes a generator * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * fix api * rename * remove unittest * use randn_tensor * fix device map on multigpu * check if the key is missing in the original state dict and only then add to the allow_missing set * remove initialize_weights --------- Co-authored-by: wangyuqi <wangyuqi@MBP-FJDQNJTWYN-0208.local> Co-authored-by: Kashif Rasul <kashif.rasul@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> |
||
Dhruv Nair
|
47e8faf3b9 |
Clean up accidental files (#13202)
update |
||
|
Dhruv Nair
|
84ff061b1d |
[Modular] Save Modular Pipeline weights to Hub (#13168)
* update * update * update * update * update * update |
||
Miguel Martin
|
212db7b999 |
Cosmos Transfer2.5 Auto-Regressive Inference Pipeline (#13114)
* AR * address comments * address comments 2 |
||
Sayak Paul
|
2843b3d37a |
Sunset Python 3.8 & get rid of explicit typing exports where possible (#12524)
* drop python 3.8 * remove list, tuple, dict from typing * fold Unions into | * up * fix a bunch and please me. * up * up * up * up * up * up * enforce 3.10.0. * up * up * up * up * up * up * up * up * Update setup.py Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> * up. * python 3.10. * ifx * up * up * up * up * final * up * fix typing utils. * up * up * up * up * up * up * fix * up * up * up * up * up * up * handle modern types. * up * up * fix ip adapter type checking. * up * up * up * up * up * up * up * revert docstring changes. * keep deleted files deleted. * keep deleted files deleted. --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
|
Miguel Martin
|
a1816166a5 |
Cosmos Transfer2.5 inference pipeline: general/{seg, depth, blur, edge} (#13066)
* initial conversion script * cosmos control net block * CosmosAttention * base model conversion * wip * pipeline updates * convert controlnet * pipeline: working without controls * wip * debugging * Almost working * temp * control working * cleanup + detail on neg_encoder_hidden_states * convert edge * pos emb for control latents * convert all chkpts * resolve TODOs * remove prints * Docs * add siglip image reference encoder * Add unit tests * controlnet: add duplicate layers * Additional tests * skip less * skip less * remove image_ref * minor * docs * remove skipped test in transfer * Don't crash process * formatting * revert some changes * remove skipped test * make style * Address comment + fix example * CosmosAttnProcessor2_0 revert + CosmosAttnProcessor2_5 changes * make style * make fix-copies |
||
Vinh H. Pham
|
a2ea45a5da |
LTX2 distilled checkpoint support (#12934)
* add constants for distill sigmas values and allow ltx pipeline to pass in sigmas * add time conditioning conversion and token packing for latents * make style & quality * remove prenorm * add sigma param to ltx2 i2v * fix copies and add pack latents to i2v * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Infer latent dims if latents/audio_latents is supplied * add note for predefined sigmas * run make style and quality * revert distill timesteps & set original_state_dict_repo_idd to default None * add latent normalize * add create noised state, delete last sigmas * remove normalize step in latent upsample pipeline and move it to ltx2 pipeline * add create noise latent to i2v pipeline * fix copies * parse none value in weight conversion script * explicit shape handling * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * make style * add two stage inference tests * add ltx2 documentation * update i2v expected_audio_slice * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Apply suggestion from @dg845 Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update ltx2.md to remove one-stage example Removed one-stage generation example code and added comments for noise scale in two-stage generation. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Daniel Gu <dgu8957@gmail.com> |
||
YiYi Xu
|
61f175660a |
Flux2 klein (#12982)
* flux2-klein * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Klein tests (#2) * tests * up * tests * up * support step-distilled * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * doc string etc * style * more * copies * klein lora training scripts (#3) * initial commit * initial commit * remove remote text encoder * initial commit * initial commit * initial commit * revert * img2img fix * text encoder + tokenizer * text encoder + tokenizer * update readme * guidance * guidance * guidance * test * test * revert changes not needed for the non klein model * Update examples/dreambooth/train_dreambooth_lora_flux2_klein.py Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * fix guidance * fix validation * fix validation * fix validation * fix path * space --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * style * Update src/diffusers/pipelines/flux2/pipeline_flux2_klein.py * Apply style fixes * auto pipeline --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
dg845
|
c10bdd9b73 |
Add LTX 2.0 Video Pipelines (#12915)
* Initial LTX 2.0 transformer implementation * Add tests for LTX 2 transformer model * Get LTX 2 transformer tests working * Rename LTX 2 compile test class to have LTX2 * Remove RoPE debug print statements * Get LTX 2 transformer compile tests passing * Fix LTX 2 transformer shape errors * Initial script to convert LTX 2 transformer to diffusers * Add more LTX 2 transformer audio arguments * Allow LTX 2 transformer to be loaded from local path for conversion * Improve dummy inputs and add test for LTX 2 transformer consistency * Fix LTX 2 transformer bugs so consistency test passes * Initial implementation of LTX 2.0 video VAE * Explicitly specify temporal and spatial VAE scale factors when converting * Add initial LTX 2.0 video VAE tests * Add initial LTX 2.0 video VAE tests (part 2) * Get diffusers implementation on par with official LTX 2.0 video VAE implementation * Initial LTX 2.0 vocoder implementation * Use RMSNorm implementation closer to original for LTX 2.0 video VAE * start audio decoder. * init registration. * up * simplify and clean up * up * Initial LTX 2.0 text encoder implementation * Rough initial LTX 2.0 pipeline implementation * up * up * up * up * Add imports for LTX 2.0 Audio VAE * Conversion script for LTX 2.0 Audio VAE Decoder * Add Audio VAE logic to T2V pipeline * Duplicate scheduler for audio latents * Support num_videos_per_prompt for prompt embeddings * LTX 2.0 scheduler and full pipeline conversion * Add script to test full LTX2Pipeline T2V inference * Fix pipeline return bugs * Add LTX 2 text encoder and vocoder to ltx2 subdirectory __init__ * Fix more bugs in LTX2Pipeline.__call__ * Improve CPU offload support * Fix pipeline audio VAE decoding dtype bug * Fix video shape error in full pipeline test script * Get LTX 2 T2V pipeline to produce reasonable outputs * Make LTX 2.0 scheduler more consistent with original code * Fix typo when applying scheduler fix in T2V inference script * Refactor Audio VAE to be simpler and remove helpers (#7) * remove resolve causality axes stuff. * remove a bunch of helpers. * remove adjust output shape helper. * remove the use of audiolatentshape. * move normalization and patchify out of pipeline. * fix * up * up * Remove unpatchify and patchify ops before audio latents denormalization (#9) --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Add support for I2V (#8) * start i2v. * up * up * up * up * up * remove uniform strategy code. * remove unneeded code. * Denormalize audio latents in I2V pipeline (analogous to T2V change) (#11) * test i2v. * Move Video and Audio Text Encoder Connectors to Transformer (#12) * Denormalize audio latents in I2V pipeline (analogous to T2V change) * Initial refactor to put video and audio text encoder connectors in transformer * Get LTX 2 transformer tests working after connector refactor * precompute run_connectors,. * fixes * Address review comments * Calculate RoPE double precisions freqs using torch instead of np * Further simplify LTX 2 RoPE freq calc * Make connectors a separate module (#18) * remove text_encoder.py * address yiyi's comments. * up * up * up * up --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com> * up (#19) * address initial feedback from lightricks team (#16) * cross_attn_timestep_scale_multiplier to 1000 * implement split rope type. * up * propagate rope_type to rope embed classes as well. * up * When using split RoPE, make sure that the output dtype is same as input dtype * Fix apply split RoPE shape error when reshaping x to 4D * Add export_utils file for exporting LTX 2.0 videos with audio * Tests for T2V and I2V (#6) * add ltx2 pipeline tests. * up * up * up * up * remove content * style * Denormalize audio latents in I2V pipeline (analogous to T2V change) * Initial refactor to put video and audio text encoder connectors in transformer * Get LTX 2 transformer tests working after connector refactor * up * up * i2v tests. * up * Address review comments * Calculate RoPE double precisions freqs using torch instead of np * Further simplify LTX 2 RoPE freq calc * revert unneded changes. * up * up * update to split style rope. * up --------- Co-authored-by: Daniel Gu <dgu8957@gmail.com> * up * use export util funcs. * Point original checkpoint to LTX 2.0 official checkpoint * Allow the I2V pipeline to accept image URLs * make style and make quality * remove function map. * remove args. * update docs. * update doc entries. * disable ltx2_consistency test * Simplify LTX 2 RoPE forward by removing coords is None logic * make style and make quality * Support LTX 2.0 audio VAE encoder * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Remove print statement in audio VAE * up * Fix bug when calculating audio RoPE coords * Ltx 2 latent upsample pipeline (#12922) * Initial implementation of LTX 2.0 latent upsampling pipeline * Add new LTX 2.0 spatial latent upsampler logic * Add test script for LTX 2.0 latent upsampling * Add option to enable VAE tiling in upsampling test script * Get latent upsampler working with video latents * Fix typo in BlurDownsample * Add latent upsample pipeline docstring and example * Remove deprecated pipeline VAE slicing/tiling methods * make style and make quality * When returning latents, return unpacked and denormalized latents for T2V and I2V * Add model_cpu_offload_seq for latent upsampling pipeline --------- Co-authored-by: Daniel Gu <dgu8957@gmail.com> * Fix latent upsampler filename in LTX 2 conversion script * Add latent upsample pipeline to LTX 2 docs * Add dummy objects for LTX 2 latent upsample pipeline * Set default FPS to official LTX 2 ckpt default of 24.0 * Set default CFG scale to official LTX 2 ckpt default of 4.0 * Update LTX 2 pipeline example docstrings * make style and make quality * Remove LTX 2 test scripts * Fix LTX 2 upsample pipeline example docstring * Add logic to convert and save a LTX 2 upsampling pipeline * Document LTX2VideoTransformer3DModel forward pass --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com> |
||
|
Miguel Martin
|
973a077c6a |
Cosmos Predict2.5 14b Conversion (#12863)
14b conversion |
||
Miguel Martin
|
b5309683cb |
Cosmos Predict2.5 Base: inference pipeline, scheduler & chkpt conversion (#12852)
* cosmos predict2.5 base: convert chkpt & pipeline - New scheduler: scheduling_flow_unipc_multistep.py - Changes to TransformerCosmos for text embeddings via crossattn_proj * scheduler cleanup * simplify inference pipeline * cleanup scheduler + tests * Basic tests for flow unipc * working b2b inference * Rename everything * Tests for pipeline present, but not working (predict2 also not working) * docstring update * wrapper pipelines + make style * remove unnecessary files * UniPCMultistep: support use_karras_sigmas=True and use_flow_sigmas=True * use UniPCMultistepScheduler + fix tests for pipeline * Remove FlowUniPCMultistepScheduler * UniPCMultistepScheduler for use_flow_sigmas=True & use_karras_sigmas=True * num_inference_steps=36 due to bug in scheduler used by predict2.5 * Address comments * make style + make fix-copies * fix tests + remove references to old pipelines * address comments * add revision in from_pretrained call * fix tests |
||
|
YiYi Xu
|
671149e036 |
[HunyuanVideo1.5] support step-distilled (#12802)
* support step-distilled * style |
||
Guo-Hua Wang
|
4f136f842c |
Add support for Ovis-Image (#12740)
* add ovis_image * fix code quality * optimize pipeline_ovis_image.py according to the feedbacks * optimize imports * add docs * make style * make style * add ovis to toctree * oops --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
|
YiYi Xu
|
6156cf8f22 |
Hunyuanvideo15 (#12696)
* add --------- Co-authored-by: yiyi@huggingface.co <yiyi@ip-26-0-161-123.ec2.internal> Co-authored-by: yiyi@huggingface.co <yiyi@ip-26-0-160-103.ec2.internal> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
Sayak Paul
|
5ffb73d4ae |
let's go Flux2 🚀 (#12711)
* add vae * Initial commit for Flux 2 Transformer implementation * add pipeline part * small edits to the pipeline and conversion * update conversion script * fix * up up * finish pipeline * Remove Flux IP Adapter logic for now * Remove deprecated 3D id logic * Remove ControlNet logic for now * Add link to ViT-22B paper as reference for parallel transformer blocks such as the Flux 2 single stream block * update pipeline * Don't use biases for input projs and output AdaNorm * up * Remove bias for double stream block text QKV projections * Add script to convert Flux 2 transformer to diffusers * make style and make quality * fix a few things. * allow sft files to go. * fix image processor * fix batch * style a bit * Fix some bugs in Flux 2 transformer implementation * Fix dummy input preparation and fix some test bugs * fix dtype casting in timestep guidance module. * resolve conflicts., * remove ip adapter stuff. * Fix Flux 2 transformer consistency test * Fix bug in Flux2TransformerBlock (double stream block) * Get remaining Flux 2 transformer tests passing * make style; make quality; make fix-copies * remove stuff. * fix type annotaton. * remove unneeded stuff from tests * tests * up * up * add sf support * Remove unused IP Adapter and ControlNet logic from transformer (#9) * copied from * Apply suggestions from code review Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: apolinário <joaopaulo.passos@gmail.com> * up * up * up * up * up * Refactor Flux2Attention into separate classes for double stream and single stream attention * Add _supports_qkv_fusion to AttentionModuleMixin to allow subclasses to disable QKV fusion * Have Flux2ParallelSelfAttention inherit from AttentionModuleMixin with _supports_qkv_fusion=False * Log debug message when calling fuse_projections on a AttentionModuleMixin subclass that does not support QKV fusion * Address review comments * Update src/diffusers/pipelines/flux2/pipeline_flux2.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * up * Remove maybe_allow_in_graph decorators for Flux 2 transformer blocks (#12) * up * support ostris loras. (#13) * up * update schdule * up * up (#17) * add training scripts (#16) * add training scripts Co-authored-by: Linoy Tsaban <linoytsaban@gmail.com> * model cpu offload in validation. * add flux.2 readme * add img2img and tests * cpu offload in log validation * Apply suggestions from code review * fix * up * fixes * remove i2i training tests for now. --------- Co-authored-by: Linoy Tsaban <linoytsaban@gmail.com> Co-authored-by: linoytsaban <linoy@huggingface.co> * up --------- Co-authored-by: yiyixuxu <yixu310@gmail.com> Co-authored-by: Daniel Gu <dgu8957@gmail.com> Co-authored-by: yiyi@huggingface.co <yiyi@ip-10-53-87-203.ec2.internal> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: apolinário <joaopaulo.passos@gmail.com> Co-authored-by: yiyi@huggingface.co <yiyi@ip-26-0-160-103.ec2.internal> Co-authored-by: Linoy Tsaban <linoytsaban@gmail.com> Co-authored-by: linoytsaban <linoy@huggingface.co> |
||
|
Junsong Chen
|
1afc21855e |
SANA-Video Image to Video pipeline SanaImageToVideoPipeline support (#12634)
* move sana-video to a new dir and add `SanaImageToVideoPipeline` with no modify; * fix bug and run text/image-to-vidoe success; * make style; quality; fix-copies; * add sana image-to-video pipeline in markdown; * add test case for sana image-to-video; * make style; * add a init file in sana-video test dir; * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana_video/test_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana_video/test_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * minor update; * fix bug and skip fp16 save test; Co-authored-by: Yuyang Zhao <43061147+HeliosZhao@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * add copied from for `encode_prompt` * Apply style fixes --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Yuyang Zhao <43061147+HeliosZhao@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
dg845
|
d8e4805816 |
[WIP]Add Wan2.2 Animate Pipeline (Continuation of #12442 by tolgacangoz) (#12526)
--------- Co-authored-by: Tolga Cangöz <mtcangoz@gmail.com> Co-authored-by: Tolga Cangöz <46008593+tolgacangoz@users.noreply.github.com> |
||
Yashwant Bezawada
|
0fd58c7706 |
fix: correct import path for load_model_dict_into_meta in conversion scripts (#12616)
The function load_model_dict_into_meta was moved from modeling_utils.py to model_loading_utils.py but the imports in the conversion scripts were not updated, causing ImportError when running these scripts. This fixes the import in 6 conversion scripts: - scripts/convert_sd3_to_diffusers.py - scripts/convert_stable_cascade_lite.py - scripts/convert_stable_cascade.py - scripts/convert_stable_audio.py - scripts/convert_sana_to_diffusers.py - scripts/convert_sana_controlnet_to_diffusers.py Fixes #12606 |
||
|
Junsong Chen
|
b3e9dfced7 |
[SANA-Video] Adding 5s pre-trained 480p SANA-Video inference (#12584)
* 1. add `SanaVideoTransformer3DModel` in transformer_sana_video.py 2. add `SanaVideoPipeline` in pipeline_sana_video.py 3. add all code we need for import `SanaVideoPipeline` * add a sample about how to use sana-video; * code update; * update hf model path; * update code; * sana-video can run now; * 1. add aspect ratio in sana-video-pipeline; 2. add reshape function in sana-video-processor; 3. fix convert pth to safetensor bugs; * default to use `use_resolution_binning`; * make style; * remove unused code; * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * support `dispatch_attention_fn` * 1. add sana-video markdown; 2. fix typos; * add two test case for sana-video (need check) * fix text-encoder in test-sana-video; * Update tests/pipelines/sana/test_sana_video.py * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/video_processor.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * make style make quality make fix-copies * toctree yaml update; * add sana-video-transformer3d markdown; * Apply style fixes --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
YiYi Xu
|
a138d71ec1 |
HunyuanImage21 (#12333)
* add hunyuanimage2.1 --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
David Bertoin
|
dd07b19e27 |
Prx (#12525)
* rename photon to prx * rename photon into prx * Revert .gitignore to state before commit |
||
|
David Bertoin
|
cefc2cf82d |
Add Photon model and pipeline support (#12456)
* Add Photon model and pipeline support This commit adds support for the Photon image generation model: - PhotonTransformer2DModel: Core transformer architecture - PhotonPipeline: Text-to-image generation pipeline - Attention processor updates for Photon-specific attention mechanism - Conversion script for loading Photon checkpoints - Documentation and tests * just store the T5Gemma encoder * enhance_vae_properties if vae is provided only * remove autocast for text encoder forwad * BF16 example * conditioned CFG * remove enhance vae and use vae.config directly when possible * move PhotonAttnProcessor2_0 in transformer_photon * remove einops dependency and now inherits from AttentionMixin * unify the structure of the forward block * update doc * update doc * fix T5Gemma loading from hub * fix timestep shift * remove lora support from doc * Rename EmbedND for PhotoEmbedND * remove modulation dataclass * put _attn_forward and _ffn_forward logic in PhotonBlock's forward * renam LastLayer for FinalLayer * remove lora related code * rename vae_spatial_compression_ratio for vae_scale_factor * support prompt_embeds in call * move xattention conditionning out computation out of the denoising loop * add negative prompts * Use _import_structure for lazy loading * make quality + style * add pipeline test + corresponding fixes * utility function that determines the default resolution given the VAE * Refactor PhotonAttention to match Flux pattern * built-in RMSNorm * Revert accidental .gitignore change * parameter names match the standard diffusers conventions * renaming and remove unecessary attributes setting * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * quantization example * added doc to toctree * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * use dispatch_attention_fn for multiple attention backend support * naming changes * make fix copy * Update docs/source/en/api/pipelines/photon.md Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Add PhotonTransformer2DModel to TYPE_CHECKING imports * make fix-copies * Use Tuple instead of tuple Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * restrict the version of transformers Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/photon/test_pipeline_photon.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/photon/test_pipeline_photon.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * change | for Optional * fix nits. * use typing Dict --------- Co-authored-by: davidb <davidb@worker-10.soperator-worker-svc.soperator.svc.cluster.local> Co-authored-by: David Briand <david@photoroom.com> Co-authored-by: davidb <davidb@worker-8.soperator-worker-svc.soperator.svc.cluster.local> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: sayakpaul <spsayakpaul@gmail.com> |
||
Ali Imran
|
1b456bd5d5 |
docs: cleanup of runway model (#12503)
* cleanup of runway model * quality fixes |
||
|
YiYi Xu
|
f072c64bf2 |
ltx0.9.8 (without IC lora, autoregressive sampling) (#12493)
update Co-authored-by: Aryan <aryan@huggingface.co> |
||
Linoy Tsaban
|
b50014067d |
Add Wan2.2 VACE - Fun (#12324)
* support Wan2.2-VACE-Fun-A14B * support Wan2.2-VACE-Fun-A14B * support Wan2.2-VACE-Fun-A14B * Apply style fixes * test --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
YiYi Xu
|
a6d9f6a1a9 |
[WIP] Wan2.2 (#12004)
* support wan 2.2 i2v * add t2v + vae2.2 * add conversion script for vae 2.2 * add * add 5b t2v * conversion script * refactor out reearrange * remove a copied from in skyreels * Apply suggestions from code review Co-authored-by: bagheera <59658056+bghira@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_wan.py * fix fast tests * style --------- Co-authored-by: bagheera <59658056+bghira@users.noreply.github.com> |
||
Tolga Cangöz
|
7298bdd817 |
Add SkyReels V2: Infinite-Length Film Generative Model (#11518)
* style * Fix class name casing for SkyReelsV2 components in multiple files to ensure consistency and correct functionality. * cleaning * cleansing * Refactor `get_timestep_embedding` to move modifications into `SkyReelsV2TimeTextImageEmbedding`. * Remove unnecessary line break in `get_timestep_embedding` function for cleaner code. * Remove `skyreels_v2` entry from `_import_structure` and update its initialization to directly assign the list of SkyReelsV2 components. * cleansing * Refactor attention processing in `SkyReelsV2AttnProcessor2_0` to always convert query, key, and value to `torch.bfloat16`, simplifying the code and improving clarity. * Enhance example usage in `pipeline_skyreels_v2_diffusion_forcing.py` by adding VAE initialization and detailed prompt for video generation, improving clarity and usability of the documentation. * Refactor import structure in `__init__.py` for SkyReelsV2 components and improve formatting in `pipeline_skyreels_v2_diffusion_forcing.py` to enhance code readability and maintainability. * Update `guidance_scale` parameter in `SkyReelsV2DiffusionForcingPipeline` from 5.0 to 6.0 to enhance video generation quality. * Update `guidance_scale` parameter in example documentation and class definition of `SkyReelsV2DiffusionForcingPipeline` to ensure consistency and improve video generation quality. * Update `causal_block_size` parameter in `SkyReelsV2DiffusionForcingPipeline` to default to `None`. * up * Fix dtype conversion for `timestep_proj` in `SkyReelsV2Transformer3DModel` to *ensure* correct tensor operations. * Optimize causal mask generation by replacing repeated tensor with `repeat_interleave` for improved efficiency in `SkyReelsV2Transformer3DModel`. * style * Enhance example documentation in `SkyReelsV2DiffusionForcingPipeline` with guidance scale and shift parameters for T2V and I2V. Remove unused `retrieve_latents` function to streamline the code. * Refactor sample scheduler creation in `SkyReelsV2DiffusionForcingPipeline` to use `deepcopy` for improved state management during inference steps. * Enhance error handling and documentation in `SkyReelsV2DiffusionForcingPipeline` for `overlap_history` and `addnoise_condition` parameters to improve long video generation guidance. * Update documentation and progress bar handling in `SkyReelsV2DiffusionForcingPipeline` to clarify asynchronous inference settings and improve progress tracking during denoising steps. * Refine progress bar calculation in `SkyReelsV2DiffusionForcingPipeline` by rounding the step size to one decimal place for improved readability during denoising steps. * Update import statements in `SkyReelsV2DiffusionForcingPipeline` documentation for improved clarity and organization. * Refactor progress bar handling in `SkyReelsV2DiffusionForcingPipeline` to use total steps instead of calculated step size. * update templates for i2v, v2v * Add `retrieve_latents` function to streamline latent retrieval in `SkyReelsV2DiffusionForcingPipeline`. Update video latent processing to utilize this new function for improved clarity and maintainability. * Add `retrieve_latents` function to both i2v and v2v pipelines for consistent latent retrieval. Update video latent processing to utilize this function, enhancing clarity and maintainability across the SkyReelsV2DiffusionForcingPipeline implementations. * Remove redundant ValueError for `overlap_history` in `SkyReelsV2DiffusionForcingPipeline` to streamline error handling and improve user guidance for long video generation. * Update default video dimensions and flow matching scheduler parameter in `SkyReelsV2DiffusionForcingPipeline` to enhance video generation capabilities. * Refactor `SkyReelsV2DiffusionForcingPipeline` to support Image-to-Video (i2v) generation. Update class name, add image encoding functionality, and adjust parameters for improved video generation. Enhance error handling for image inputs and update documentation accordingly. * Improve organization for image-last_image condition. * Refactor `SkyReelsV2DiffusionForcingImageToVideoPipeline` to improve latent preparation and video condition handling integration. * style * style * Add example usage of PIL for image input in `SkyReelsV2DiffusionForcingImageToVideoPipeline` documentation. * Refactor `SkyReelsV2DiffusionForcingPipeline` to `SkyReelsV2DiffusionForcingVideoToVideoPipeline`, enhancing support for Video-to-Video (v2v) generation. Introduce video input handling, update latent preparation logic, and improve error handling for input parameters. * Refactor `SkyReelsV2DiffusionForcingImageToVideoPipeline` by removing the `image_encoder` and `image_processor` dependencies. Update the CPU offload sequence accordingly. * Refactor `SkyReelsV2DiffusionForcingImageToVideoPipeline` to enhance latent preparation logic and condition handling. Update image input type to `Optional`, streamline video condition processing, and improve handling of `last_image` during latent generation. * Enhance `SkyReelsV2DiffusionForcingPipeline` by refining latent preparation for long video generation. Introduce new parameters for video handling, overlap history, and causal block size. Update logic to accommodate both short and long video scenarios, ensuring compatibility and improved processing. * refactor * fix num_frames * fix prefix_video_latents * up * refactor * Fix typo in scheduler method call within `SkyReelsV2DiffusionForcingVideoToVideoPipeline` to ensure proper noise scaling during latent generation. * up * Enhance `SkyReelsV2DiffusionForcingImageToVideoPipeline` by adding support for `last_image` parameter and refining latent frame calculations. Update preprocessing logic. * add statistics * Refine latent frame handling in `SkyReelsV2DiffusionForcingImageToVideoPipeline` by correcting variable names and reintroducing latent mean and standard deviation calculations. Update logic for frame preparation and sampling to ensure accurate video generation. * up * refactor * up * Refactor `SkyReelsV2DiffusionForcingVideoToVideoPipeline` to improve latent handling by enforcing tensor input for video, updating frame preparation logic, and adjusting default frame count. Enhance preprocessing and postprocessing steps for better integration. * style * fix vae output indexing * upup * up * Fix tensor concatenation and repetition logic in `SkyReelsV2DiffusionForcingImageToVideoPipeline` to ensure correct dimensionality for video conditions and latent conditions. * Refactor latent retrieval logic in `SkyReelsV2DiffusionForcingVideoToVideoPipeline` to handle tensor dimensions more robustly, ensuring compatibility with both 3D and 4D video inputs. * Enhance logging in `SkyReelsV2DiffusionForcing` pipelines by adding iteration print statements for better debugging. Clean up unused code related to prefix video latents length calculation in `SkyReelsV2DiffusionForcingImageToVideoPipeline`. * Update latent handling in `SkyReelsV2DiffusionForcingImageToVideoPipeline` to conditionally set latents based on video iteration state, improving flexibility for video input processing. * Refactor `SkyReelsV2TimeTextImageEmbedding` to utilize `get_1d_sincos_pos_embed_from_grid` for timestep projection. * Enhance `get_1d_sincos_pos_embed_from_grid` function to include an optional parameter `flip_sin_to_cos` for flipping sine and cosine embeddings, improving flexibility in positional embedding generation. * Update timestep projection in `SkyReelsV2TimeTextImageEmbedding` to include `flip_sin_to_cos` parameter, enhancing the flexibility of time embedding generation. * Refactor tensor type handling in `SkyReelsV2AttnProcessor2_0` and `SkyReelsV2TransformerBlock` to ensure consistent use of `torch.float32` and `torch.bfloat16`, improving integration. * Update tensor type in `SkyReelsV2RotaryPosEmbed` to use `torch.float32` for frequency calculations, ensuring consistency in data types across the model. * Refactor `SkyReelsV2TimeTextImageEmbedding` to utilize automatic mixed precision for timestep projection. * down * down * style * Add debug tensor tracking to `SkyReelsV2Transformer3DModel` for enhanced debugging and output analysis; update `Transformer2DModelOutput` to include debug tensors. * up * Refactor indentation in `SkyReelsV2AttnProcessor2_0` to improve code readability and maintain consistency in style. * Convert query, key, and value tensors to bfloat16 in `SkyReelsV2AttnProcessor2_0` for improved performance. * Add debug print statements in `SkyReelsV2TransformerBlock` to track tensor shapes and values for improved debugging and analysis. * debug * debug * Remove commented-out debug tensor tracking from `SkyReelsV2TransformerBlock` * Add functionality to save processed video latents as a Safetensors file in `SkyReelsV2DiffusionForcingPipeline`. * up * Add functionality to save output latents as a Safetensors file in `SkyReelsV2DiffusionForcingPipeline`. * up * Remove additional commented-out debug tensor tracking from `SkyReelsV2TransformerBlock` and `SkyReelsV2Transformer3DModel` for cleaner code. * style * cleansing * Update example documentation and parameters in `SkyReelsV2Pipeline`. Adjusted example code for loading models, modified default values for height, width, num_frames, and guidance_scale, and improved output video quality settings. * Update shift parameter in example documentation and default values across SkyReels V2 pipelines. Adjusted shift values for I2V from 3.0 to 5.0 and updated related example code for consistency. * Update example documentation in SkyReels V2 pipelines to include available model options and update model references for loading. Adjusted model names to reflect the latest versions across I2V, V2V, and T2V pipelines. * Add test templates * style * Add docs template * Add SkyReels V2 Diffusion Forcing Video-to-Video Pipeline to imports * style * fix-copies * convert i2v 1.3b * Update transformer configuration to include `image_dim` for SkyReels V2 models and refactor imports to use `SkyReelsV2Transformer3DModel`. * Refactor transformer import in SkyReels V2 pipeline to use `SkyReelsV2Transformer3DModel` for consistency. * Update transformer configuration in SkyReels V2 to increase `in_channels` from 16 to 36 for i2v conf. * Update transformer configuration in SkyReels V2 to set `added_kv_proj_dim` values for different model types. * up * up * up * Add SkyReelsV2Pipeline support for T2V model type in conversion script * upp * Refactor model type checks in conversion script to use substring matching for improved flexibility * upp * Fix shard path formatting in conversion script to accommodate varying model types by dynamically adjusting zero padding. * Update sharded safetensors loading logic in conversion script to use substring matching for model directory checks * Update scheduler parameters in SkyReels V2 test files for consistency across image and video pipelines * Refactor conversion script to initialize text encoder, tokenizer, and scheduler for SkyReels pipelines, enhancing model integration * style * Update documentation for SkyReels-V2, introducing the Infinite-length Film Generative model, enhancing text-to-video generation examples, and updating model references throughout the API documentation. * Add SkyReelsV2Transformer3DModel and FlowMatchUniPCMultistepScheduler documentation, updating TOC and introducing new model and scheduler files. * style * Update documentation for SkyReelsV2DiffusionForcingPipeline to correct flow matching scheduler parameter for I2V from 3.0 to 5.0, ensuring clarity in usage examples. * Add documentation for causal_block_size parameter in SkyReelsV2DF pipelines, clarifying its role in asynchronous inference. * Simplify min_ar_step calculation in SkyReelsV2DiffusionForcingPipeline to improve clarity. * style and fix-copies * style * Add documentation for SkyReelsV2Transformer3DModel Introduced a new markdown file detailing the SkyReelsV2Transformer3DModel, including usage instructions and model output specifications. * Update test configurations for SkyReelsV2 pipelines - Adjusted `in_channels` from 36 to 16 in `test_skyreels_v2_df_image_to_video.py`. - Added new parameters: `overlap_history`, `num_frames`, and `base_num_frames` in `test_skyreels_v2_df_video_to_video.py`. - Updated expected output shape in video tests from (17, 3, 16, 16) to (41, 3, 16, 16). * Refines SkyReelsV2DF test parameters * Update src/diffusers/models/modeling_outputs.py Co-authored-by: Aryan <contact.aryanvs@gmail.com> * Refactor `grid_sizes` processing by using already-calculated post-patch parameters to simplify * Update docs/source/en/api/pipelines/skyreels_v2.md Co-authored-by: Aryan <contact.aryanvs@gmail.com> * Refactor parameter naming for diffusion forcing in SkyReelsV2 pipelines - Changed `flag_df` to `enable_diffusion_forcing` for clarity in the SkyReelsV2Transformer3DModel and associated pipelines. - Updated all relevant method calls to reflect the new parameter name. * Revert _toctree.yml to adjust section expansion states * style * Update docs/source/en/api/models/skyreels_v2_transformer_3d.md Co-authored-by: YiYi Xu <yixu310@gmail.com> * Add copying label to SkyReelsV2ImageEmbedding from WanImageEmbedding. * Refactor transformer block processing in SkyReelsV2Transformer3DModel - Ensured proper handling of hidden states during both gradient checkpointing and standard processing. * Update SkyReels V2 documentation to remove VRAM requirement and streamline imports - Removed the mention of ~13GB VRAM requirement for the SkyReels-V2 model. - Simplified import statements by removing unused `load_image` import. * Add SkyReelsV2LoraLoaderMixin for loading and managing LoRA layers in SkyReelsV2Transformer3DModel - Introduced SkyReelsV2LoraLoaderMixin class to handle loading, saving, and fusing of LoRA weights specific to the SkyReelsV2 model. - Implemented methods for state dict management, including compatibility checks for various LoRA formats. - Enhanced functionality for loading weights with options for low CPU memory usage and hotswapping. - Added detailed docstrings for clarity on parameters and usage. * Update SkyReelsV2 documentation and loader mixin references - Corrected the documentation to reference the new `SkyReelsV2LoraLoaderMixin` for loading LoRA weights. - Updated comments in the `SkyReelsV2LoraLoaderMixin` class to reflect changes in model references from `WanTransformer3DModel` to `SkyReelsV2Transformer3DModel`. * Enhance SkyReelsV2 integration by adding SkyReelsV2LoraLoaderMixin references - Added `SkyReelsV2LoraLoaderMixin` to the documentation and loader imports for improved LoRA weight management. - Updated multiple pipeline classes to inherit from `SkyReelsV2LoraLoaderMixin` instead of `WanLoraLoaderMixin`. * Update SkyReelsV2 model references in documentation - Replaced placeholder model paths with actual paths for SkyReels-V2 models in multiple pipeline files. - Ensured consistency across the documentation for loading models in the SkyReelsV2 pipelines. * style * fix-copies * Refactor `fps_projection` in `SkyReelsV2Transformer3DModel` - Replaced the sequential linear layers for `fps_projection` with a `FeedForward` layer using `SiLU` activation for better integration. * Update docs * Refactor video processing in SkyReelsV2DiffusionForcingPipeline - Renamed parameters for clarity: `video` to `video_latents` and `overlap_history` to `overlap_history_latent_frames`. - Updated logic for handling long video generation, including adjustments to latent frame calculations and accumulation. - Consolidated handling of latents for both long and short video generation scenarios. - Final decoding step now consistently converts latents to pixels, ensuring proper output format. * Update activation function in `fps_projection` of `SkyReelsV2Transformer3DModel` - Changed activation function from `silu` to `linear-silu` in the `fps_projection` layer for improved performance and integration. * Add fps_projection layer renaming in convert_skyreelsv2_to_diffusers.py - Updated key mappings for the `fps_projection` layer to align with new naming conventions, ensuring consistency in model integration. * Fix fps_projection assignment in SkyReelsV2Transformer3DModel - Corrected the assignment of the `fps_projection` layer to ensure it is properly cast to the appropriate data type, enhancing model functionality. * Update _keep_in_fp32_modules in SkyReelsV2Transformer3DModel - Added `fps_projection` to the list of modules that should remain in FP32 precision, ensuring proper handling of data types during model operations. * Remove integration test classes from SkyReelsV2 test files - Deleted the `SkyReelsV2DiffusionForcingPipelineIntegrationTests` and `SkyReelsV2PipelineIntegrationTests` classes along with their associated setup, teardown, and test methods, as they were not implemented and not needed for current testing. * style * Refactor: Remove hardcoded `torch.bfloat16` cast in attention * Refactor: Simplify data type handling in transformer model Removes unnecessary data type conversions for the FPS embedding and timestep projection. This change simplifies the forward pass by relying on the inherent data types of the tensors. * Refactor: Remove `fps_projection` from `_keep_in_fp32_modules` in `SkyReelsV2Transformer3DModel` * Update src/diffusers/models/transformers/transformer_skyreels_v2.py Co-authored-by: Aryan <contact.aryanvs@gmail.com> * Refactor: Remove unused flags and simplify attention mask handling in SkyReelsV2AttnProcessor2_0 and SkyReelsV2Transformer3DModel Refactor: Simplify causal attention logic in SkyReelsV2 Removes the `flag_causal_attention` and `_flag_ar_attention` flags to simplify the implementation. The decision to apply a causal attention mask is now based directly on the `num_frame_per_block` configuration, eliminating redundant flags and conditional checks. This streamlines the attention mechanism and simplifies the `set_ar_attention` methods. * Refactor: Clarify variable names for latent frames Renames `base_num_frames` to `base_latent_num_frames` to make it explicit that the variable refers to the number of frames in the latent space. This change improves code readability and reduces potential confusion between latent frames and decoded video frames. The `num_frames` parameter in `generate_timestep_matrix` is also renamed to `num_latent_frames` for consistency. * Enhance documentation: Add detailed docstring for timestep matrix generation in SkyReelsV2DiffusionForcingPipeline * Docs: Clarify long video chunking in pipeline docstring Improves the explanation of long video processing within the pipeline's docstring. The update replaces the abstract description with a concrete example, illustrating how the sliding window mechanism works with overlapping chunks. This makes the roles of `base_num_frames` and `overlap_history` clearer for users. * Docs: Move visual demonstration and processing details for SkyReelsV2DiffusionForcingPipeline to docs page from the code * Docs: Update asynchronous processing timeline and examples for long video handling in SkyReels-V2 documentation * Enhance timestep matrix generation documentation and logic for synchronous/asynchronous video processing * Update timestep matrix documentation and enhance analysis for clarity in SkyReelsV2DiffusionForcingPipeline * Docs: Update visual demonstration section and add detailed step matrix construction example for asynchronous processing in SkyReelsV2DiffusionForcingPipeline * style * fix-copies * Refactor parameter names for clarity in SkyReelsV2DiffusionForcingImageToVideoPipeline and SkyReelsV2DiffusionForcingVideoToVideoPipeline * Refactor: Avoid VAE roundtrip in long video generation Improves performance and quality for long video generation by operating entirely in latent space during the iterative generation process. Instead of decoding latents to video and then re-encoding the overlapping section for the next chunk, this change passes the generated latents directly between iterations. This avoids a computationally expensive and potentially lossy VAE decode/encode cycle within the loop. The full video is now decoded only once from the accumulated latents at the end of the process. * Refactor: Rename prefix_video_latents_length to prefix_video_latents_frames for clarity * Refactor: Rename num_latent_frames to current_num_latent_frames for clarity in SkyReelsV2DiffusionForcingImageToVideoPipeline * Refactor: Enhance long video generation logic and improve latent handling in SkyReelsV2DiffusionForcingImageToVideoPipeline Refactor: Unify video generation and pass latents directly Unifies the separate code paths for short and long video generation into a single, streamlined loop. This change eliminates the inefficient decode-encode cycle during long video generation. Instead of converting latents to pixel-space video between chunks, the pipeline now passes the generated latents directly to the next iteration. This improves performance, avoids potential quality loss from intermediate VAE steps, and enhances code maintainability by removing significant duplication. * style * Refactor: Remove overlap_history parameter and streamline long video generation logic in SkyReelsV2DiffusionForcingImageToVideoPipeline Refactor: Streamline long video generation logic Removes the `overlap_history` parameter and simplifies the conditioning process for long video generation. This change avoids a redundant VAE encoding step by directly using latent frames from the previous chunk for conditioning. It also moves image preprocessing outside the main generation loop to prevent repeated computations and clarifies the handling of prefix latents. * style * Refactor latent handling in i2v diffusion forcing pipeline Improves the latent conditioning and accumulation logic within the image-to-video diffusion forcing loop. - Corrects the splitting of the initial conditioning tensor to robustly handle both even and odd lengths. - Simplifies how latents are accumulated across iterations for long video generation. - Ensures the final latents are trimmed correctly before decoding only when a `last_image` is provided. * Refactor: Remove overlap_history parameter from SkyReelsV2DiffusionForcingImageToVideoPipeline * Refactor: Adjust video_latents parameter handling in prepare_latents method * style * Refactor: Update long video iteration print statements for clarity * Fix: Update transformer config with dynamic causal block size Updates the SkyReelsV2 pipelines to correctly set the `causal_block_size` in the transformer's configuration when it's provided during a pipeline call. This ensures the model configuration reflects the user's specified setting for the inference run. The `set_ar_attention` method is also renamed to `_set_ar_attention` to mark it as an internal helper. * style * Refactor: Adjust video input size and expected output shape in inference test * Refactor: Rename video variables for clarity in SkyReelsV2DiffusionForcingVideoToVideoPipeline * Docs: Clarify time embedding logic in SkyReelsV2 Adds comments to explain the handling of different time embedding tensor dimensions. A 2D tensor is used for standard models with a single time embedding per batch, while a 3D tensor is used for Diffusion Forcing models where each frame has its own time embedding. This clarifies the expected input for different model variations. * Docs: Update SkyReels V2 pipeline examples Updates the docstring examples for the SkyReels V2 pipelines to reflect current best practices and API changes. - Removes the `shift` parameter from pipeline call examples, as it is now configured directly on the scheduler. - Replaces the `set_ar_attention` method call with the `causal_block_size` argument in the pipeline call for diffusion forcing examples. - Adjusts recommended parameters for I2V and V2V examples, including inference steps, guidance scale, and `ar_step`. * Refactor: Remove `shift` parameter from SkyReelsV2 pipelines Removes the `shift` parameter from the call signature of all SkyReelsV2 pipelines. This parameter is a scheduler-specific configuration and should be set directly on the scheduler during its initialization, rather than being passed at runtime through the pipeline. This change simplifies the pipeline API. Usage examples are updated to reflect that the `shift` value should now be passed when creating the `FlowMatchUniPCMultistepScheduler`. * Refactors SkyReelsV2 image-to-video tests and adds last image case Simplifies the test suite by removing a duplicated test class and streamlining the dummy component and input generation. Adds a new test to verify the pipeline's behavior when a `last_image` is provided as input for conditioning. * test: Add image components to SkyReelsV2 pipeline test Adds the `image_encoder` and `image_processor` to the test components for the image-to-video pipeline. Also replaces a hardcoded value for the positional embedding sequence length with a more descriptive calculation, improving clarity. * test: Add callback configuration test for SkyReelsV2DiffusionForcingVideoToVideoPipeline test: Add callback test for SkyReelsV2DFV2V pipeline Adds a test to validate the callback functionality for the `SkyReelsV2DiffusionForcingVideoToVideoPipeline`. This test confirms that `callback_on_step_end` is invoked correctly and can modify the pipeline's state during inference. It uses a callback to dynamically increase the `guidance_scale` and asserts that the final value is as expected. The implementation correctly accounts for the nested denoising loops present in diffusion forcing pipelines. * style * fix: Update image_encoder type to CLIPVisionModelWithProjection in SkyReelsV2ImageToVideoPipeline * UP * Add conversion support for SkyReels-V2-FLF2V models Adds configurations for three new FLF2V model variants (1.3B-540P, 14B-540P, and 14B-720P) to the conversion script. This change also introduces specific handling to zero out the image positional embeddings for these models and updates the main script to correctly initialize the image-to-video pipeline. * Docs: Update and simplify SkyReels V2 usage examples Simplifies the text-to-video example by removing the manual group offloading configuration, making it more straightforward. Adds comments to pipeline parameters to clarify their purpose and provides guidance for different resolutions and long video generation. Introduces a new section with a code example for the video-to-video pipeline. * style * docs: Add SkyReels-V2 FLF2V 1.3B model to supported models list * docs: Update SkyReels-V2 documentation * Move the initialization of the `gradient_checkpointing` attribute to its suggested location. * Refactor: Use logger for long video progress messages Replaces `print()` calls with `logger.debug()` for reporting progress during long video generation in SkyReelsV2DF pipelines. This change reduces console output verbosity for standard runs while allowing developers to view progress by enabling debug-level logging. * Refactor SkyReelsV2 timestep embedding into a module Extract the sinusoidal timestep embedding logic into a new `SkyReelsV2Timesteps` `nn.Module`. This change encapsulates the embedding generation, which simplifies the `SkyReelsV2TimeTextImageEmbedding` class and improves code modularity. * Fix: Preserve original shape in timestep embeddings Reshapes the timestep embedding tensor to match the original input shape. This ensures that batched timestep inputs retain their batch dimension after embedding, preventing potential shape mismatches. * style * Refactor: Move SkyReelsV2Timesteps to model file Colocates the `SkyReelsV2Timesteps` class with the SkyReelsV2 transformer model. This change moves model-specific timestep embedding logic from the general embeddings module to the transformer's own file, improving modularity and making the model more self-contained. * Refactor parameter dtype retrieval to use utility function Replaces manual parameter iteration with the `get_parameter_dtype` helper to determine the time embedder's data type. This change improves code readability and centralizes the logic. * Add comments to track the tensor shape transformations * Add copied froms * style * fix-copies * up * Remove FlowMatchUniPCMultistepScheduler Deletes the `FlowMatchUniPCMultistepScheduler` as it is no longer being used. * Refactor: Replace FlowMatchUniPC scheduler with UniPC Removes the `FlowMatchUniPCMultistepScheduler` and integrates its functionality into the existing `UniPCMultistepScheduler`. This consolidation is achieved by using the `use_flow_sigmas=True` parameter in `UniPCMultistepScheduler`, simplifying the scheduler API and reducing code duplication. All usages, documentation, and tests are updated accordingly. * style * Remove text_encoder parameter from SkyReelsV2DiffusionForcingPipeline initialization * Docs: Rename `pipe` to `pipeline` in SkyReels examples Updates the variable name from `pipe` to `pipeline` across all SkyReels V2 documentation examples. This change improves clarity and consistency. * Fix: Rename shift parameter to flow_shift in SkyReels-V2 examples * Fix: Rename shift parameter to flow_shift in example documentation across SkyReels-V2 files * Fix: Rename shift parameter to flow_shift in UniPCMultistepScheduler initialization across SkyReels test files * Removes unused generator argument from scheduler step The `generator` parameter is not used by the scheduler's `step` method within the SkyReelsV2 diffusion forcing pipelines. This change removes the unnecessary argument from the method call for code clarity and consistency. * Fix: Update time_embedder_dtype assignment to use the first parameter's dtype in SkyReelsV2TimeTextImageEmbedding * style * Refactor: Use get_parameter_dtype utility function Replaces manual parameter iteration with the `get_parameter_dtype` helper. * Fix: Prevent (potential) error in parameter dtype check Adds a check to ensure the `_keep_in_fp32_modules` attribute exists on a parameter before it is accessed. This prevents a potential `AttributeError`, making the utility function more robust when used with models that do not define this attribute. --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Aryan <contact.aryanvs@gmail.com> |
||
Aryan
|
a79c3af6bb |
[single file] Cosmos (#11801)
* update * update * update docs |
||
|
Aryan
|
a4df8dbc40 |
Update more licenses to 2025 (#11746)
update |
||
Carl Thomé
|
81426b0f19 | Fix misleading comment (#11722) | ||
|
Aryan
|
9f91305f85 |
Cosmos Predict2 (#11695)
* support text-to-image * update example * make fix-copies * support use_flow_sigmas in EDM scheduler instead of maintain cosmos-specific scheduler * support video-to-world * update * rename text2image pipeline * make fix-copies * add t2i test * add test for v2w pipeline * support edm dpmsolver multistep * update * update * update * update tests * fix tests * safety checker * make conversion script work without guardrail |
||
|
Aryan
|
73a9d5856f |
Wan VACE (#11582)
* initial support * make fix-copies * fix no split modules * add conversion script * refactor * add pipeline test * refactor * fix bug with mask * fix for reference images * remove print * update docs * update slices * update * update * update example |
||
|
Aryan
|
06fee551e9 |
LTX Video 0.9.7 (#11516)
* add upsampling pipeline * ltx upsample pipeline conversion; pipeline fixes * make fix-copies * remove print * add vae convenience methods * update * add tests * support denoising strength for upscaling & video-to-video * update docs * update doc checkpoints * update docs * fix --------- Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> |
||
|
Aryan
|
7b904941bc |
Cosmos (#10660)
* begin transformer conversion * refactor * refactor * refactor * refactor * refactor * refactor * update * add conversion script * add pipeline * make fix-copies * remove einops * update docs * gradient checkpointing * add transformer test * update * debug * remove prints * match sigmas * add vae pt. 1 * finish CV* vae * update * update * update * update * update * update * make fix-copies * update * make fix-copies * fix * update * update * make fix-copies * update * update tests * handle device and dtype for safety checker; required in latest diffusers * remove enable_gqa and use repeat_interleave instead * enforce safety checker; use dummy checker in fast tests * add review suggestion for ONNX export Co-Authored-By: Asfiya Baig <asfiyab@nvidia.com> * fix safety_checker issues when not passed explicitly We could either do what's done in this commit, or update the Cosmos examples to explicitly pass the safety checker * use cosmos guardrail package * auto format docs * update conversion script to support 14B models * update name CosmosPipeline -> CosmosTextToWorldPipeline * update docs * fix docs * fix group offload test failing for vae --------- Co-authored-by: Asfiya Baig <asfiyab@nvidia.com> |
||
co63oc
|
86294d3c7f |
Fix typos in docs and comments (#11416)
* Fix typos in docs and comments * Apply style fixes --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
YiYi Xu
|
5873377a66 |
[Wan2.1-FLF2V] update conversion script (#11365)
update scheuler config in conversion sript |
||
|
YiYi Xu
|
0021bfa1e1 |
support Wan-FLF2V (#11353)
* update transformer --------- Co-authored-by: Aryan <aryan@huggingface.co> |
||
Ishan Modi
|
f1f38ffbee |
[ControlNet] Adds controlnet for SanaTransformer (#11040)
* added controlnet for sana transformer
* improve code quality
* addressed PR comments
* bug fixes
* added test cases
* update
* added dummy objects
* addressed PR comments
* update
* Forcing update
* add to docs
* code quality
* addressed PR comments
* addressed PR comments
* update
* addressed PR comments
* added proper styling
* update
* Revert "added proper styling"
This reverts commit
|