* feat(cosmos3): multi-GPU inference — context + tensor parallelism
Cosmos 3 cannot use diffusers' declarative `_cp_plan` CP path: it is grouped-query
attention (the shared Ulysses kernel assumes K/V share the query head count), its
understanding (causal) and generation (full) streams are separate packed sequences
(gen attends to cat(und, gen)), and per-pathway lengths are ragged. The model carries
no parallelism logic -- it exposes only small, CP-agnostic seams; all sharding lives
outside it, in a reusable example module.
Model (transformer_cosmos3.py): adds two default-None `forward` seams -- `_cp_shard_fn`
(shards und/gen + rotary before the decoder layers) and `_cp_gather_fn` (gathers/unpads
after the final norm) -- and extracts `Cosmos3AttnProcessor._run_attention` as an
override point. The non-parallel path is unchanged.
Helpers (examples/cosmos3/cosmos_parallel.py): one importable module, two orthogonal
and composable axes:
* Context parallelism (Ulysses) -- `enable_cosmos3_context_parallel`. Shards the
sequence; brackets the two attention pathways with all-to-all (DTensor redistribute),
repeats GQA KV heads, pads ragged lengths and masks padded generation keys.
* Tensor parallelism (Megatron) -- `enable_cosmos3_tensor_parallel`. Column/row-shards
the attention + MLP weights so a checkpoint that does not fit one GPU (Super, ~120 GB)
loads across several; weights load to CPU then shard layer by layer.
Both expand KV heads to the query-head count and call SDPA with enable_gqa=False so it
dispatches to the flash kernel; enable_gqa=True forces the math path, which materializes
the full [S, S] score matrix and OOMs on long videos. A dense `Cosmos3FlashAttnProcessor`
(`enable_cosmos3_flash_attention`) provides the same for TP without CP.
CLI (examples/cosmos3/inference_cosmos3.py): imports these helpers, so any modality
(text-to-image/video, image-to-video, sound, action) runs single- or multi-GPU via
`--tp-degree` / `--cp-degree` (their product must equal --nproc_per_node). Single-GPU
behavior is unchanged.
Docs + example README updated. Verified: CP attention core is bit-exact vs non-CP in
fp32 (max|d|=0), and a full 36-layer forward matches CP-on vs CP-off to ~1e-6 in fp32
(bf16 differs only by floating-point rounding).
* refactor(cosmos3): make CP/TP attention processors standalone
Address review feedback on PR #14054: the parallel attention processors no
longer subclass Cosmos3AttnProcessor, so the model file needs no override seam.
- transformer_cosmos3.py: revert Cosmos3AttnProcessor to inline the attention
in __call__ (remove the _run_attention seam); restores it to its base version.
- cosmos_parallel.py: Cosmos3CPAttnProcessor and Cosmos3FlashAttnProcessor each
get their own full __call__, sharing a _project_qkv_with_rope prologue helper.
Verified behavior-preserving on 4x RTX PRO 6000: cp_unit_test (fp32) passes at
1e-4; cp_numeric_check is byte-identical to the pre-refactor code; the end-to-end
CLI passes in CP-only, TP-only, and TP+CP modes.
* fix(cosmos3): support modular multi-GPU inference
* Remove JAX/Flax support
Deletes all Flax models, schedulers, pipelines, and loading utilities,
the from_flax loading path, flax/jax packaging extras, Flax CI filters,
Flax training examples, and Flax content in translated docs.
Hub compatibility is preserved: Flax class names in old model_index.json
files are still stripped to their PyTorch equivalents, and *.msgpack
files in repos are still ignored during download.
* Drop Flax class-name remapping from model_index.json loading
Flax-saved checkpoint metadata is no longer recognized: class names like
FlaxStableDiffusionPipeline in model_index.json now fail to resolve
instead of being remapped to their PyTorch equivalents. *.msgpack files
in hub repos are still ignored during download.
* Restore Flax training examples
The example scripts are standalone and still work with diffusers<=0.39.x.
A warning note is added to each Flax README section pointing out the
version requirement now that the library itself no longer ships Flax.
* update
* Repoint scheduler default-config tests off the Flax fixture
* Repoint example training tests off the Flax fixture
* Ideogram4 DreamBooth LoRA training
Add examples/dreambooth/train_dreambooth_lora_ideogram4.py + README + requirements:
- DreamBooth LoRA training for Ideogram 4 (flow-matching, dual transformer, Qwen3-VL TE)
- nf4 QLoRA and SDNQ fp8 bases: --do_fp8_training trains the fp8 checkpoint in place
(scaled matmul), or omit it to dequantize the SDNQ base to bf16
- --disable_training_autocast (Ideogram4's forward is corrupted by autocast)
- structured JSON caption support + --upsample_prompt
- model card with a validation-image gallery
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Apply style fixes
* Fix base model repo reference in Ideogram4 LoRA README
* Align Ideogram4 script deps with requirements file
---------
Co-authored-by: linoytsaban <linoy.tsaban@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.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: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
* Add Krea 2 LoRA DreamBooth trainer + Krea2 LoRA loader
Adds examples/dreambooth/train_dreambooth_lora_krea2.py and the
Krea2LoraLoaderMixin (Krea2Pipeline now inherits it) so LoRA adapters
can be saved/loaded for the Krea 2 transformer. The transformer/pipeline
forward now accept attention_kwargs for runtime LoRA scaling.
Co-Authored-By: linoytsaban <linoy.tsaban@gmail.com>
* inline _pack/_unpack_latents in trainer (pipeline methods are instance methods)
* add Krea2LoraLoaderMixin to LoRA loaders docs
* remove mu from validation call to align with pipeline on main
---------
Co-authored-by: linoytsaban <linoy.tsaban@gmail.com>
Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
`train_dreambooth_lora_flux2_klein_img2img.py` passed the whole batch to
`Flux2KleinPipeline._prepare_image_ids` as a list of reference images, so each
independent training sample got a different temporal id (T=10, 20, 30, ...) as if
they were multiple references of one instance. Each batch element is an independent
sample with a single conditional image, so they must share the same temporal id.
Generate ids for one sample and expand across the batch. This is the Flux2 Klein
counterpart of the same bug/fix in `train_dreambooth_lora_flux2_img2img.py` (#13811).
Fixes#13866
When training with `--mixed_precision="fp16"` and `--validation_prompt`,
the first optimizer step after a validation run fails with
`ValueError: Attempting to unscale FP16 gradients`.
Under fp16, `cast_training_params` keeps the trainable LoRA params in
fp32. The in-loop validation pipeline is built with the same live `unet`
object, and `log_validation` then calls `pipeline.to(device, dtype=torch_dtype)`,
which downcasts those fp32 LoRA params back to fp16. The next backward
therefore produces fp16 grads and `GradScaler.unscale_` raises.
Drop the dtype cast from that `.to(...)` so the shared `unet` keeps its
fp32 LoRA params. This matches train_dreambooth_lora_sdxl.py, which moves
the validation pipeline with `.to(accelerator.device)` only.
Fixes#13124
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Fix AttributeError in onnxruntime train_unconditional
The ORT example `examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py`
defines `--logger` (argparse line 196) but references `args.report_to`
in two places:
* Line 280: `if args.report_to == "wandb" and args.hub_token is not None:`
* Line 294: `log_with=args.report_to,`
Since there is no `--report_to` flag on the parser, both references crash
with `AttributeError: 'Namespace' object has no attribute 'report_to'`
as soon as `main()` starts — the script cannot run at all.
The sibling non-ORT script
`examples/unconditional_image_generation/train_unconditional.py` uses
`args.logger` consistently; this PR aligns the ORT variant with that
behavior, and updates the error message to refer to `--logger=wandb`
instead of the non-existent `--report_to=wandb`.
* 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>
* examples/dreambooth: chunk weighting tensor alongside model_pred and target when using prior preservation (flux LoRA)
* examples/dreambooth: chunk weighting tensor alongside model_pred and target when using prior preservation (SD3 LoRA)
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* support lora for cosmos 2.5
* Fix inconsistencies with cosmos official repo in VAE encoding, text encoder attention implementation, and timestep scaling
* Support f_min and f_max in linear_scheduler warmup
* Add requirements and dataset preprocessing scripts to run examples
* Add LoRA training scripts
* Add LoRA eval scripts
* add assets for blogpost
* Fix(scheduler): device mismatch from upstream b114620 - move rk and b to device before torch.stack
* Always upcast to fp32
* Directly inhrit from LoraBaseMixin
* remove flash-attn2
* Use _keep_in_fp32_modules instead of autocast
* remove the get_latent_shape_cthw method and fix style
* simplifiy the eval script to make it more user-friendly
* overwrite scheduling_unipc_multistep.py with main's version
* remove network_alphas and add # Copied from
* remove figures and assets
* revert scheduler
* revert fp32 upcast and support bs > 1
---------
Co-authored-by: Ting-Yun Chang <tingyunc@nvidia.com>
* Handle prompt embedding concat in Qwen dreambooth example
* remove wandb config
* Apply style fixes
* add a comment on how this is only relevant during prior preservation.
---------
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* add a profiling worflow.
* fix
* fix
* more clarification
* add points.
* up
* cache hooks
* improve readme.
* propagate deletion.
* up
* up
* wan fixes.
* more
* up
* add more traces.
* up
* better title
* cuda graphs.
* up
* Apply suggestions from code review
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* add torch.compile link.
* approach -> How the tooling works
* table
* unavoidable gaps.
* make important
* note on regional compilation
* Apply suggestions from code review
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* make regional compilation note clearer.
* Apply suggestions from code review
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* clarify scheduler related changes.
* Apply suggestions from code review
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* Update examples/profiling/README.md
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* up
* formatting
* benchmarking runtime
* up
* up
* up
* up
* Update examples/profiling/README.md
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
---------
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* 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>
* switch to transformers main again./
* more
* up
* up
* fix group offloading.
* attributes
* up
* up
* tie embedding issue.
* fix t5 stuff for more.
* matrix configuration to see differences between 4.57.3 and main failures.
* change qwen expected slice because of how init is handled in v5.
* same stuff.
* up
* up
* Revert "up"
This reverts commit 515dd06db5.
* Revert "up"
This reverts commit 5274ffdd7f.
* up
* up
* fix with peft_format.
* just keep main for easier debugging.
* remove torchvision.
* empty
* up
* up with skyreelsv2 fixes.
* fix skyreels type annotation.
* up
* up
* fix variant loading issues.
* more fixes.
* fix dduf
* fix
* fix
* fix
* more fixes
* fixes
* up
* up
* fix dduf test
* up
* more
* update
* hopefully ,final?
* one last breath
* always install from main
* up
* audioldm tests
* up
* fix PRX tests.
* up
* kandinsky fixes
* qwen fixes.
* prx
* hidream
* 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>
