diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index 9c9373807a9..977a7a033b8 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -1,4 +1,4 @@ -# DMD2 distillation for Qwen-Image +# DMD2 distillation for Qwen-Image and Qwen-Image-Edit Distill [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) into a **few-step generator** with DMD2 (Distribution Matching Distillation). The distilled student @@ -6,6 +6,10 @@ produces images in as few as **1–4 sampling steps** while matching the base mo output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's [`TrainDiffusionRecipe`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py). +The paired image-edit path supports +[`Qwen/Qwen-Image-Edit-2511`](https://huggingface.co/Qwen/Qwen-Image-Edit-2511), +including its one-or-more reference-image input contract. + > [!NOTE] > Qwen-Image is a third-party model with its own license terms. Review the > [Qwen-Image model card](https://huggingface.co/Qwen/Qwen-Image) before downloading or @@ -115,14 +119,135 @@ torchrun --nproc-per-node=8 \ Any `DMDConfig` field can be overridden on the CLI (e.g. `--dmd2.guidance_scale=3.5`). +## Qwen-Image-Edit-2511 training + +Image editing uses `configs/dmd2_qwen_image_edit_2511.yaml`. It is not a text-to-image +cache with an extra tensor: Edit-2511 conditions every transformer call with packed +reference-image latents, and its Qwen2.5-VL prompt embedding jointly encodes the edit +instruction and those same ordered references. Stable `diffusers>=0.37.0` is required so +the transformer's `zero_cond_t` path applies `t=0` modulation to the reference tokens. + +Preprocess native SpatialEdit WebDataset shards directly, without extracting the image +corpus: + +```bash +python examples/diffusers/fastgen/preprocess_qwen_image_edit.py \ + --input-dir /path/to/SpatialEdit-500K \ + --output-dir /path/to/qwen_image_edit_2511_cache \ + --model-name /path/to/Qwen-Image-Edit-2511 \ + --gpu-id 0 +``` + +The preprocessor also accepts `--manifest pairs.jsonl`. A row supplies a target, an edit +instruction, and one or more ordered references; image values may be local paths or +`{"archive": "/path/shard.tar", "member": "sample.0.jpg"}` descriptors: + +```json +{"id":"sample-1","target":"target.png","conditioning":["source.png"],"prompt":"Move the red cube left."} +``` + +Each cached sample contains the target latent, a list of deterministic reference latents, +and image-aware positive **and negative** embeddings. Consequently the edit dataloader +does not take `negative_prompt_embedding_path`; one global text-only negative embedding +would omit the reference-image visual tokens. + +Keep the edit dataloader at `batch_size: 1` with the current sampler. It buckets target +resolution only; batching multiple samples also requires matching reference count and every +reference-slot shape. The collate rejects incompatible batches instead of padding image tokens. + +```bash +torchrun --nproc-per-node=8 \ + examples/diffusers/fastgen/dmd2_finetune.py \ + --config examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml \ + --model.pretrained_model_name_or_path=/path/to/Qwen-Image-Edit-2511 \ + --data.dataloader.cache_dir=/path/to/qwen_image_edit_2511_cache \ + --fsdp.dp_size=8 --step_scheduler.global_batch_size=8 +``` + +The `qwen_image_edit` plugin packs the noisy target first, appends every clean reference, +builds `img_shapes=[target, reference_1, ...]`, and slices the transformer prediction back +to the target-token prefix. The same references are forwarded through student, +teacher, fake-score, CFG, backward-simulation, and GAN paths. + ### Checkpoints & resuming Checkpoints land under `checkpoint.checkpoint_dir`. Alongside the student, the recipe saves the DMD2 sidecars needed to resume exactly: the fake-score model + optimizer, the -student EMA (`ema_shadow.pt`), and the DMD iteration counter (`dmd_state.pt`). With +DMD iteration counter (`dmd_state.pt`), and, when EMA is enabled, the student EMA +(`ema_shadow.pt`). With `restore_from: LATEST` a re-launch auto-resumes from the newest checkpoint; pin a specific one with `--checkpoint.restore_from=epoch_0_step_500`. +## Quantization-aware training (QAT) + +Continue a full-precision DMD2 run with the **student quantized**, so the few-step model +stays accurate at FP8/NVFP4. QAT here is **restore-only**: the trainer loads a ModelOpt +quantizer state (recipe + frozen `amax`) from disk and **never calibrates**. Only the +student is quantized; the frozen teacher and trainable fake-score stay full precision so +the distribution-matching gradient is exact, and `amax` stays frozen for the whole run. + +QAT is driven by a `dmd2.quant` block — there's no dedicated config file. The cleanest +way to launch is to **reuse the exact config + overrides of the full-precision run you're +continuing** and add only the three `dmd2.quant.*` keys (plus a reduced LR), so the QAT +run is provably identical to the FP run except for quantization and learning rate. The +CLI parser creates the `dmd2.quant` subtree even when it's absent from the YAML. + +| Key | Role | +| --- | --- | +| `dmd2.quant.enabled` | Turn QAT on (restore-only student quantization). | +| `dmd2.quant.quant_state_path` | The `transformer.pt` from step 1 below (recipe + frozen `amax`). | +| `dmd2.quant.init_weights_from` | FP DMD2 checkpoint to warm-start student / fake-score / optimizers from on the first launch (the run `amax` was calibrated against). | + +1. **Calibrate once** with the quantization example to produce the quantizer state + (`amax`, no weights) for a trained student checkpoint: + + ```bash + python examples/diffusers/quantization/quantize.py \ + --model qwen-image-dmd2 --format fp8 \ + --extra-param student_path=<.../epoch_4_step_15999/model/consolidated> \ + --quantized-torch-ckpt-save-path <.../epoch_4_step_15999/quant> + # -> writes <.../epoch_4_step_15999/quant/transformer.pt> + ``` + +2. **Launch QAT** by re-running the FP run's command with a new output dir, a reduced + student LR, and the three quant keys appended: + + ```bash + torchrun --nproc-per-node= \ + examples/diffusers/fastgen/dmd2_finetune.py \ + --config examples/diffusers/fastgen/configs/.yaml \ + --checkpoint.checkpoint_dir= \ + <... the FP run's other overrides, unchanged ...> \ + --optim.learning_rate= --lr_scheduler.min_lr= \ + --dmd2.quant.enabled=true \ + --dmd2.quant.quant_state_path=<.../epoch_4_step_15999/quant/transformer.pt> \ + --dmd2.quant.init_weights_from=<.../epoch_4_step_15999> + ``` + +On the first launch (empty `checkpoint_dir`) the student / fake-score / discriminator / +optimizers warm-start from `init_weights_from`, then the student is quantized from +`quant_state_path`. `restore_from: LATEST` auto-resumes the new `checkpoint_dir` +thereafter. Because QAT is restore-only — amax never recalibrates — the recipe re-applies +`quant_state_path` on every resume rather than persisting a per-checkpoint copy, so keep +that file accessible for the whole run (it's the only quantization dependency). The saved +student weights are clean full precision (`model/consolidated` is a normal +`QwenImageTransformer2DModel`); re-apply `quant_state_path` to deploy or evaluate the +quantized QAT student via the quantization example. + +> The `quant_state_path` `amax` must have been calibrated against the student in +> `init_weights_from`, with the same few-step schedule (`dmd2.sample_t_cfg.t_list`) the +> student trains/infers with. Pass `dmd2.quant.enabled=true` on every resume too (it is +> what tells the recipe to quantize). Reduce only the student LR by keeping +> `--dmd2.fake_score_lr` / `--dmd2.discriminator_lr` at the FP value. + +For Qwen-Image-Edit, the restore-only QAT path itself is unchanged because the student is +still a `QwenImageTransformer2DModel`. The quantizer state must, however, be calibrated on +the exact edit student using target + reference tokens, multimodal image/instruction +embeddings, and the same `t_list`. The existing `--model qwen-image-dmd2` text-only +calibrator does **not** provide representative activation ranges for Edit-2511. Also +calibrate the non-EMA weights restored by `dmd2.quant.init_weights_from`; do not calibrate +an EMA overlay for that warm start. + ## Inference After training, sample from the distilled student. The pipeline loads your consolidated @@ -162,6 +287,23 @@ Set `num_inference_steps` to the number of steps the student was trained for (`dmd2.student_sample_steps` — e.g. 4 for the canonical config, or 1 for a single-step student). +For an edit student, use the companion pipeline and pass one or more ordered references: + +```python +from inference_dmd2_qwen_image_edit import QwenImageEditDMDInferencePipeline +from diffusers.utils import load_image + +pipe = QwenImageEditDMDInferencePipeline.from_pretrained( + student_path="/path/to/checkpoint/model/consolidated", + base_pipeline_path="Qwen/Qwen-Image-Edit-2511", +).to("cuda") +image = pipe( + [load_image("source.png")], + "Move the red cube left.", + num_inference_steps=4, +).images[0] +``` + ## Config reference | Section | Key | Role | @@ -170,7 +312,7 @@ student). | `model` | `mode` | `finetune` — loads the pretrained weights. | | `step_scheduler` | `global_batch_size`, `local_batch_size`, `max_steps`, `ckpt_every_steps`, `log_every` | Standard AutoModel scheduling knobs. | | `dmd2` | `recipe_path` | Built-in fastgen recipe to hydrate `DMDConfig` from (`general/distillation/dmd2_qwen_image`). | -| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / img_shapes). | +| `dmd2` | `pipeline_plugin` | `qwen_image` for T2I or `qwen_image_edit` for target + reference token packing. | | `dmd2` | `student_sample_steps` | Number of student sampling steps (e.g. 4). | | `dmd2` | `guidance_scale` | CFG strength on the teacher (`null` disables CFG; requires a negative-prompt embedding when set). | | `dmd2` | `gan_loss_weight_gen`, `gan_r1_reg_weight`, `gan_feature_indices`, … | GAN branch (set `gan_loss_weight_gen: 0` to disable). | @@ -178,7 +320,7 @@ student). | `dmd2` | `sample_t_cfg`, `ema` | Timestep sampling + student EMA settings. | | `optim` | `learning_rate`, `optimizer.*` | Student AdamW knobs. | | `fsdp` | `dp_size`, `tp_size`, `activation_checkpointing`, … | FSDP2 parallelism (set `dp_size` to your GPU count). | -| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Latent cache dir + optional CFG negative-prompt embedding. | +| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Latent cache. The static negative path is T2I-only; edit negatives are cached per sample. | | `checkpoint` | `checkpoint_dir`, `model_save_format`, `restore_from` | Output dir, save format, resume behavior. | ## Troubleshooting diff --git a/examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml b/examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml new file mode 100644 index 00000000000..e515fa2e911 --- /dev/null +++ b/examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml @@ -0,0 +1,124 @@ +# Qwen-Image-Edit-2511 DMD2 — paired image-edit training. +# +# Preprocess a paired dataset first (SpatialEdit-500K is supported directly): +# +# python examples/diffusers/fastgen/preprocess_qwen_image_edit.py \ +# --input-dir /path/to/SpatialEdit-500K \ +# --output-dir /path/to/qwen_image_edit_2511_cache \ +# --model-name Qwen/Qwen-Image-Edit-2511 +# +# Then launch with torchrun and override the cache/checkpoint paths as needed. +# A local model snapshot can be selected with: +# --model.pretrained_model_name_or_path=/path/to/Qwen-Image-Edit-2511 + +seed: 42 + +wandb: + project: fastgen-dmd2-qwen-image-edit + mode: online + name: qwen_image_edit_2511_dmd2 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + pretrained_model_name_or_path: Qwen/Qwen-Image-Edit-2511 + mode: finetune + +step_scheduler: + global_batch_size: 128 + local_batch_size: 1 + ckpt_every_steps: 500 + num_epochs: 4 + log_every: 1 + max_steps: 5000 + +dmd2: + recipe_path: general/distillation/dmd2_qwen_image + # Unlike text-to-image, this plugin appends the cached reference-image tokens + # to the noisy target tokens and trains on the target prefix only. + pipeline_plugin: qwen_image_edit + qwen_image_guidance: + + pred_type: flow + num_train_timesteps: + guidance_scale: 4.0 + student_sample_steps: 4 + student_sample_type: ode + backward_simulation: false + student_update_freq: 5 + fake_score_pred_type: x0 + + gan_loss_weight_gen: 0.03 + gan_use_same_t_noise: true + gan_r1_reg_weight: 0.1 + gan_r1_reg_alpha: 0.1 + + fake_score_lr: 2.0e-6 + discriminator_lr: 2.0e-6 + gan_feature_indices: [30] + gan_num_blocks: 60 + gan_inner_dim: 3072 + + sample_t_cfg: + time_dist_type: uniform + min_t: 0.001 + max_t: 0.999 + p_mean: 0.0 + p_std: 1.0 + t_list: [0.999, 0.74925, 0.4995, 0.24975, 0.0] + + # Full-tensor EMA materializes a complete FP32 copy of this 20B-class student on + # every rank, which is not viable on 80-GiB workers alongside the three DMD2 models. + # Keep it disabled until the EMA checkpoint path supports sharded shadows end to end. + ema: null + + # Restore-only student QAT uses the same recipe. Enable after calibrating the + # exact FP edit student with reference-token + multimodal conditioning and this + # t_list; the text-only qwen-image-dmd2 calibration path is not representative. + quant: + enabled: false + quant_state_path: + init_weights_from: + +optim: + learning_rate: 2.0e-6 + optimizer: + weight_decay: 0.01 + betas: [0.9, 0.999] + +lr_scheduler: + lr_decay_style: constant + lr_warmup_steps: 0 + min_lr: 2.0e-6 + +fsdp: + tp_size: 1 + cp_size: 1 + pp_size: 1 + dp_replicate_size: 1 + dp_size: 128 + activation_checkpointing: true + +# Each cache item contains: +# target latent, 1..N reference latents, multimodal positive prompt embedding, +# and a per-sample multimodal negative embedding built from the same references. +data: + dataloader: + _target_: fastgen_data.build_image_to_image_multiresolution_dataloader + cache_dir: /path/to/preprocessed/qwen_image_edit_2511 + base_resolution: [1024, 1024] + batch_size: 1 + drop_last: false + shuffle: true + num_workers: 0 + +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_edit_2511_dmd2/checkpoints + model_save_format: safetensors + save_consolidated: true + v4_compatible: true + diffusers_compatible: true + restore_from: LATEST diff --git a/examples/diffusers/fastgen/dmd2_finetune.py b/examples/diffusers/fastgen/dmd2_finetune.py index 6d91db94acd..22bddfb4aa1 100644 --- a/examples/diffusers/fastgen/dmd2_finetune.py +++ b/examples/diffusers/fastgen/dmd2_finetune.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Entrypoint for the DMD2 Qwen-Image AutoModel example. +"""Entrypoint for the DMD2 Qwen-Image / Qwen-Image-Edit AutoModel examples. Parses the YAML config + CLI overrides with AutoModel's argument parser, then hands control to :class:`DMD2DiffusionRecipe`. diff --git a/examples/diffusers/fastgen/dmd2_recipe.py b/examples/diffusers/fastgen/dmd2_recipe.py index 7934a07cf13..3ab18771755 100644 --- a/examples/diffusers/fastgen/dmd2_recipe.py +++ b/examples/diffusers/fastgen/dmd2_recipe.py @@ -20,10 +20,10 @@ drives ``modelopt.torch.fastgen.DMDPipeline`` (or a plugin subclass) through the three-phase DMD2 alternation (student update / fake-score update / EMA step). -Backbone: **Qwen-Image** (``Qwen/Qwen-Image``) — 4D ``image_latents``, -:class:`QwenImageDMDPipeline` handles 2x2 patch packing / img_shapes / -unpacking. Config: ``configs/dmd2_qwen_image.yaml`` — the canonical -real-data run (4-step + CFG + GAN). +Backbones: **Qwen-Image** and **Qwen-Image-Edit-2511** — both train 4D target +``image_latents``. :class:`QwenImageDMDPipeline` handles T2I patch packing; +:class:`QwenImageEditDMDPipeline` additionally appends clean reference-image tokens and +crops predictions back to the target prefix. Canonical configs live under ``configs/``. Launch:: @@ -37,6 +37,8 @@ from __future__ import annotations +import contextlib +import dataclasses import json import logging import os @@ -54,6 +56,11 @@ # and surfaced as a downstream ``TypeError: takes no arguments``. try: from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.distributed.parallelizer import ( + PARALLELIZATION_STRATEGIES, + DefaultParallelizationStrategy, + register_parallel_strategy, + ) from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe, is_main_process except ImportError as exc: raise ImportError( @@ -68,10 +75,67 @@ from torch import nn import modelopt.torch.fastgen as mtf +import modelopt.torch.opt as mto from modelopt.torch.fastgen.config import DMDConfig from modelopt.torch.fastgen.discriminators import Discriminator_ImageDiT from modelopt.torch.fastgen.methods.dmd import DMDPipeline from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin +from modelopt.torch.quantization.utils.core_utils import set_quantizer_state_dict + + +class _QwenImageParallelizationStrategy(DefaultParallelizationStrategy): + """Add full-block activation checkpointing to AutoModel's native FSDP flow. + + AutoModel's default decoder-layer checkpointing recognizes ``self_attn`` / ``mlp`` + attributes. Diffusers' Qwen image blocks instead contain joint ``attn``, ``img_mlp``, + and ``txt_mlp`` paths, so that generic logic wraps nothing. Diffusers checkpoints each + complete block; mirror that boundary, then delegate TP/FSDP behavior unchanged. + """ + + def parallelize( + self, + model, + device_mesh, + activation_checkpointing: bool = False, + **kwargs, + ): + if activation_checkpointing: + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + checkpoint_wrapper, + ) + + blocks = getattr(model, "transformer_blocks", None) + if blocks is None: + raise AttributeError( + "QwenImageTransformer2DModel does not expose `transformer_blocks`" + ) + for index, block in enumerate(blocks): + blocks[index] = checkpoint_wrapper( + block, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ) + logging.info( + "[DMD2] Qwen-Image activation checkpointing enabled for %d full blocks", + len(blocks), + ) + + return super().parallelize( + model, + device_mesh, + activation_checkpointing=False, + **kwargs, + ) + + +def _register_qwen_image_parallelization_strategy() -> None: + """Install the Qwen strategy unless AutoModel already provides a native one.""" + model_class_name = "QwenImageTransformer2DModel" + if model_class_name not in PARALLELIZATION_STRATEGIES: + register_parallel_strategy(name=model_class_name)(_QwenImageParallelizationStrategy) + + +_register_qwen_image_parallelization_strategy() # Keys under the ``dmd2:`` YAML block that shadow fields on :class:`DMDConfig`. The # recipe deep-merges these on top of the loaded built-in recipe so users can tweak DMD2 @@ -98,12 +162,41 @@ def _deep_merge_dicts(base: dict, override: dict) -> dict: return merged +def restore_quantizer_state(model: nn.Module, path: str) -> nn.Module: + """Re-insert the quantizer modules + load the frozen amax onto ``model`` from ``path``. + + ``path`` is the weight-free ModelOpt quantizer state written by the + ``examples/diffusers/quantization`` calibration example + (``--quantized-torch-ckpt-save-path`` → ``transformer.pt``): ``mto.modelopt_state`` + (which layers are quantized, FP8/NVFP4, axes) bundled with the per-quantizer buffers + (``amax``, ...) under ``modelopt_state_weights``. It carries NO model weights; ``model`` + must already hold the weights the amax was calibrated against (here: the DMD2 student + warm-started from the FP checkpoint). + + This re-applies the quantization recipe (module conversion -- ``nn.Linear`` -> + ``QuantLinear``, in place, preserving the existing ``weight``/``bias`` Parameter objects + so any pre-built FSDP2 optimizer's references stay valid) and loads the saved amax. It + is RESTORE-ONLY: no calibration forward pass is run, so amax stays exactly as it was on + disk and remains frozen for the whole training run. + """ + modelopt_state = mto.load_modelopt_state(str(path)) + quantizer_state = modelopt_state.pop("modelopt_state_weights", None) + mto.restore_from_modelopt_state(model, modelopt_state) + if quantizer_state is not None: + set_quantizer_state_dict(model, quantizer_state) + return model + + # Auto-detect substrings (matched case-insensitively against ``model_id``) that map to # DMDPipeline plugin subclasses. Keep this list small — adding a new entry is only the # right move when the model has a non-diffusers transformer signature that requires a # pack/unpack wrapper. Models with the standard ``(hidden_states, timestep, # encoder_hidden_states)`` signature work with the base :class:`DMDPipeline`. _PIPELINE_PLUGIN_BY_MODEL_SUBSTR = ( + # Edit must precede the generic Qwen-Image match: the edit transformer consumes + # target tokens followed by one or more reference-image token sequences. + ("qwen-image-edit", "qwen_image_edit"), + ("qwen_image_edit", "qwen_image_edit"), ("qwen-image", "qwen_image"), ("qwen_image", "qwen_image"), ) @@ -111,6 +204,29 @@ def _deep_merge_dicts(base: dict, override: dict) -> dict: _DMD_COMPLETE_MARKER = "dmd2_complete.marker" +@dataclasses.dataclass(frozen=True) +class _QuantSettings: + """Resolved ``dmd2.quant`` block — restore-only QAT of the student (no calibration). + + Attributes: + enabled: When ``True`` the student is quantized by RESTORING a ModelOpt quantizer + state from disk (recipe + frozen amax). The trainer never calibrates. + quant_state_path: Path to the quantizer-state file produced by the + ``examples/diffusers/quantization`` calibration example + (``--quantized-torch-ckpt-save-path`` → ``transformer.pt``). Used on the FIRST + launch (warm-start) to quantize the FP student. Required when ``enabled``. + init_weights_from: Optional path to the full-precision DMD2 checkpoint to + warm-start the student / fake_score / discriminator / EMA / optimizers from + when the run's own ``checkpoint_dir`` has no QAT checkpoint yet. The amax in + ``quant_state_path`` must have been calibrated against this checkpoint's + student weights. + """ + + enabled: bool + quant_state_path: str | None + init_weights_from: str | None + + class DMD2DiffusionRecipe(TrainDiffusionRecipe): """DMD2 recipe that reuses ``TrainDiffusionRecipe`` for the student path. @@ -379,6 +495,7 @@ def run_train_validation_loop(self) -> None: neg_text_embeds, neg_text_mask, ) = self._prepare_micro_batch(micro_batch) + model_kwargs = self._prepare_model_kwargs(micro_batch) if is_student_phase: # ``compute_student_loss`` reads ``guidance_scale`` from the @@ -394,6 +511,7 @@ def run_train_validation_loop(self) -> None: negative_encoder_hidden_states=neg_text_embeds, negative_encoder_hidden_states_mask=neg_text_mask, guidance_scale=None, + **model_kwargs, ) micro_vsd_losses.append(float(losses["vsd"].item())) else: @@ -402,6 +520,7 @@ def run_train_validation_loop(self) -> None: noise, encoder_hidden_states=text_embeds, encoder_hidden_states_mask=text_mask, + **model_kwargs, ) (losses["total"] / len(batch_group)).backward() @@ -421,6 +540,7 @@ def run_train_validation_loop(self) -> None: noise, encoder_hidden_states=text_embeds, encoder_hidden_states_mask=text_mask, + **model_kwargs, ) (disc_losses["total"] / len(batch_group)).backward() # Manual gradient all-reduce across DP ranks (the @@ -538,10 +658,33 @@ def load_checkpoint(self, restore_from: str | None = None): # ``nemo_automodel`` can be used unmodified. make_optimizer_partial_load_tolerant(self.checkpointer) + quant = self._resolve_quant_settings() + resolved = self._resolve_complete_dmd_checkpoint(restore_from) + + # QAT first launch: the run's own checkpoint_dir has no QAT checkpoint yet, so + # warm-start the FP student / fake_score / EMA / optimizers from the full-precision + # checkpoint named by ``dmd2.quant.init_weights_from`` (the FP run the amax was + # calibrated against). On later resumes ``resolved`` already points at a QAT + # checkpoint in checkpoint_dir, so this branch is skipped. + if quant.enabled and resolved is None and quant.init_weights_from: + resolved = self._resolve_complete_dmd_checkpoint(quant.init_weights_from) + if is_main_process(): + logging.info( + "[DMD2][qat] no QAT checkpoint in checkpoint_dir; warm-starting FP " + "state from dmd2.quant.init_weights_from=%s", + quant.init_weights_from, + ) + self.__dict__["_dmd2_resolved_restore_from"] = resolved if resolved is None: + if quant.enabled and is_main_process(): + logging.warning( + "[DMD2][qat] QAT enabled but no checkpoint resolved (checkpoint_dir empty " + "and dmd2.quant.init_weights_from unset). Quantizing a fresh base model — " + "its weights will NOT match the calibrated amax." + ) if ( restore_from is not None and str(restore_from).upper() == "LATEST" @@ -552,10 +695,137 @@ def load_checkpoint(self, restore_from: str | None = None): "Starting fresh.", self.checkpointer.config.checkpoint_dir, ) + # Even with no weights to restore, honor QAT by quantizing from the recipe file. + if quant.enabled: + self._quantize_student(quant.quant_state_path) return + # Single restore path. The student-weight checkpoints are always clean FP (the QAT + # save hides the quantizer buffers — see ``_student_quant_buffers_hidden``), so the + # strict DCP load matches an unquantized student whether this is a warm-start from + # the FP checkpoint or a resume from a QAT checkpoint. Quantization always happens + # AFTER the weights are in place. super().load_checkpoint(resolved) + if quant.enabled: + # Restore-only QAT never recalibrates, so the recipe + amax are invariant for the + # whole run — re-apply the same configured quantizer-state file on every (re)start + # rather than persisting an unchanging copy per checkpoint. + self._quantize_student(quant.quant_state_path) + + def _resolve_quant_settings(self) -> _QuantSettings: + """Parse and cache the ``dmd2.quant`` block (restore-only student QAT).""" + cached = self.__dict__.get("_quant_settings") + if cached is not None: + return cached + + node = self.cfg.get("dmd2.quant", None) + if node is None: + settings = _QuantSettings(enabled=False, quant_state_path=None, init_weights_from=None) + self.__dict__["_quant_settings"] = settings + return settings + + d = node.to_dict() if hasattr(node, "to_dict") else dict(node) + enabled = bool(d.get("enabled", False)) + quant_state_path = d.get("quant_state_path") + init_weights_from = d.get("init_weights_from") + if enabled and not quant_state_path: + raise ValueError( + "dmd2.quant.enabled is true but dmd2.quant.quant_state_path is not set. Point it " + "at the quantizer-state file produced by examples/diffusers/quantization " + "(its --quantized-torch-ckpt-save-path, e.g. .../quant/transformer.pt)." + ) + settings = _QuantSettings( + enabled=enabled, + quant_state_path=quant_state_path, + init_weights_from=init_weights_from, + ) + self.__dict__["_quant_settings"] = settings + return settings + + def _quantize_student(self, quant_state_path: str | None) -> None: + """Quantize the student by RESTORING a ModelOpt quantizer state from disk. + + Restore-only: re-inserts the quantizer modules from the saved recipe and loads the + precomputed, frozen amax. No ``mtq.quantize`` / calibration forward pass is ever + run, so amax stays exactly as it was on disk for the whole training run. Only the + student (``self.model``) is quantized — the frozen teacher and the trainable + fake_score stay full precision so the distribution-matching gradient is exact. + + Safe after FSDP2 wrapping: the conversion is an in-place ``__class__`` swap that + preserves the existing ``weight``/``bias`` Parameter objects, so the already-built + student optimizer's references stay valid (verified by ModelOpt's + ``tests/gpu/torch/quantization/test_fsdp2.py``). + """ + if not quant_state_path: + raise ValueError( + "[DMD2][qat] _quantize_student called without a quant_state_path. Set " + "dmd2.quant.quant_state_path." + ) + if not os.path.isfile(quant_state_path): + raise FileNotFoundError( + f"[DMD2][qat] quantizer-state file not found: {quant_state_path}. QAT here is " + "restore-only (no on-the-fly calibration); produce it with " + "examples/diffusers/quantization first." + ) + + if is_main_process(): + logging.info( + "[DMD2][qat] restoring student quantizer state (recipe + frozen amax) <- %s", + quant_state_path, + ) + restore_quantizer_state(self.model, quant_state_path) + + # amax (and any other quantizer buffers) come off disk on CPU. Move just the + # TensorQuantizer buffers onto the student device — these modules carry no + # parameters, so this never touches the FSDP2 DTensor weights. + from modelopt.torch.quantization.nn import TensorQuantizer + + for module in self.model.modules(): + if isinstance(module, TensorQuantizer): + module.to(self.device) + + if is_main_process(): + import modelopt.torch.quantization as mtq + + logging.info("[DMD2][qat] student quantized (restore-only). Quantizer summary:") + mtq.print_quant_summary(self.model) + + @contextlib.contextmanager + def _student_quant_buffers_hidden(self): + """Temporarily mark the student's quantizer buffers non-persistent. + + Wraps the parent ``save_checkpoint`` so the student's DCP shards AND the + consolidated/diffusers export stay clean full-precision (``ModelState.state_dict()`` + — and hence the consolidated index, which re-adds every state_dict key + (checkpointing.py:941-948) — excludes the ``amax`` buffers). The frozen amax is not + persisted per checkpoint at all: it is re-applied from ``dmd2.quant.quant_state_path`` + on every (re)start. Keeping the saved student weights amax-free both yields a clean + ``model/consolidated`` (a normal ``QwenImageTransformer2DModel``, re-quantizable for + deploy/eval) and makes every restore a clean ``load FP weights -> quantize`` path + (the strict DCP load matches an unquantized student). Mirrors ModelOpt's own trick in + ``quantization/plugins/transformers_trainer.py`` (``_modelopt_prepare``). + + No-op when QAT is disabled. + """ + if not self._resolve_quant_settings().enabled: + yield + return + + from modelopt.torch.quantization.nn import TensorQuantizer + + saved: list[tuple[TensorQuantizer, set[str]]] = [] + for module in self.model.modules(): + if isinstance(module, TensorQuantizer): + saved.append((module, set(module._non_persistent_buffers_set))) + module._non_persistent_buffers_set.update(module._buffers.keys()) + try: + yield + finally: + for module, original in saved: + module._non_persistent_buffers_set.clear() + module._non_persistent_buffers_set.update(original) + def save_checkpoint( self, epoch: int, @@ -588,7 +858,11 @@ def save_checkpoint( self.checkpointer.config.checkpoint_dir ) - super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) + # Hide the student's quantizer buffers during the parent save so the DCP shards and + # consolidated/diffusers export stay clean FP. Frozen amax remains external in + # ``dmd2.quant.quant_state_path`` and is re-applied on restore. No-op when QAT is disabled. + with self._student_quant_buffers_hidden(): + super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) if not self.checkpointer.config.enabled: return @@ -861,6 +1135,9 @@ def _is_dmd_checkpoint_complete(self, path: str) -> bool: if not complete: return False + # QAT adds no per-checkpoint artifact (the quantizer recipe + frozen amax are + # re-applied from dmd2.quant.quant_state_path on every (re)start), so QAT + # checkpoints use the same completeness criteria as full-precision ones. if self._cfg_gan_enabled(): return os.path.isfile(os.path.join(path, "discriminator.pt")) and os.path.isfile( os.path.join(path, "discriminator_optimizer.pt") @@ -976,10 +1253,10 @@ def _build_discriminator_optimizer(self) -> torch.optim.Optimizer | None: def _attach_gan_feature_capture(self) -> None: """Install Qwen-Image feature-capture hooks on the teacher when GAN is enabled. - Reads the latent resolution from the dataloader so the hook can reshape + Reads an initial latent resolution from the dataloader so the hook can reshape ``[B, num_image_patches, 3072]`` into ``[B, 3072, H_lat//2, W_lat//2]``. - Mock dataloader → spatial_h/spatial_w from the YAML. Real dataloader → - base_resolution / vae_scale. + The Qwen plugin refreshes that shape before every teacher forward, so real + multiresolution batches are captured using their actual target dimensions. """ feature_indices = list(self.cfg.get("dmd2.gan_feature_indices", [30])) @@ -1020,6 +1297,11 @@ def _attach_gan_feature_capture(self) -> None: feature_indices=feature_indices, h_lat=h_lat, w_lat=w_lat, + # Edit models append reference-image tokens after the target tokens. The GAN + # discriminator is defined on the generated target only, so capture its prefix. + target_prefix_only=( + self._resolve_pipeline_cls().__name__ == "QwenImageEditDMDPipeline" + ), ) if is_main_process(): logging.info( @@ -1123,13 +1405,18 @@ def _resolve_pipeline_cls(self) -> type[DMDPipeline]: from modelopt.torch.fastgen.plugins.qwen_image import QwenImageDMDPipeline return QwenImageDMDPipeline + if explicit == "qwen_image_edit": + from modelopt.torch.fastgen.plugins.qwen_image_edit import QwenImageEditDMDPipeline + + return QwenImageEditDMDPipeline raise ValueError( - f"Unknown dmd2.pipeline_plugin={explicit!r}. Supported: null/'base', 'qwen_image'." + f"Unknown dmd2.pipeline_plugin={explicit!r}. Supported: null/'base', " + "'qwen_image', 'qwen_image_edit'." ) def _resolve_pipeline_kwargs(self, pipeline_cls: type[DMDPipeline]) -> dict[str, Any]: """Extra kwargs to forward to the pipeline subclass constructor (plugin-specific).""" - if pipeline_cls.__name__ == "QwenImageDMDPipeline": + if pipeline_cls.__name__ in {"QwenImageDMDPipeline", "QwenImageEditDMDPipeline"}: # Optional ``guidance`` value passed to the transformer's guidance kwarg every # call. Independent of DMDConfig.guidance_scale (which drives the negative- # prompt CFG path on the teacher). Leave ``None`` to skip the embedding when @@ -1289,6 +1576,59 @@ def _prepare_micro_batch( noise = torch.randn_like(latents) return latents, noise, text_embeds, text_mask, negative_text_embeds, negative_text_mask + def _prepare_model_kwargs(self, micro_batch: dict[str, Any]) -> dict[str, Any]: + """Move optional model-specific conditioning to the training device. + + Qwen-Image-Edit caches one or more VAE-encoded reference images separately from + the clean target latent. The edit plugin packs these tensors and appends them to + every student / teacher / fake-score forward. Keeping the data as a list allows + references with different aspect ratios while retaining a regular batch dimension + for each reference slot. + """ + conditioning = micro_batch.get("conditioning_latents") + if conditioning is None: + return {} + + if torch.is_tensor(conditioning): + if conditioning.ndim == 4: + conditioning = [conditioning] + elif conditioning.ndim == 5: + # Collates may represent refs as [B, N, C, H, W]. Convert to the + # plugin's list-of-[B,C,H,W] contract. + conditioning = list(conditioning.unbind(dim=1)) + else: + raise ValueError( + "conditioning_latents tensor must be [B,C,H,W] or [B,N,C,H,W], " + f"got shape {tuple(conditioning.shape)}." + ) + elif isinstance(conditioning, (list, tuple)): + conditioning = list(conditioning) + else: + raise TypeError( + "conditioning_latents must be a tensor or a list/tuple of tensors; " + f"got {type(conditioning).__name__}." + ) + + if not conditioning: + raise ValueError("conditioning_latents must contain at least one reference image.") + + expected_batch = int(micro_batch["image_latents"].shape[0]) + prepared: list[torch.Tensor] = [] + for ref_index, ref in enumerate(conditioning): + if not torch.is_tensor(ref) or ref.ndim != 4: + shape = tuple(ref.shape) if torch.is_tensor(ref) else None + raise ValueError( + f"conditioning_latents[{ref_index}] must be [B,C,H,W], got {shape}." + ) + if int(ref.shape[0]) != expected_batch: + raise ValueError( + f"conditioning_latents[{ref_index}] batch={ref.shape[0]} does not match " + f"target batch={expected_batch}." + ) + prepared.append(ref.to(self.device, dtype=self.bf16, non_blocking=True)) + + return {"conditioning_latents": tuple(prepared)} + def _log_step( self, *, diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 771b93b1c0b..52d7cb6675a 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -38,9 +38,12 @@ # Convert a missing-helper ImportError into an actionable message naming the supported range. try: from .collate_fns import ( + build_image_to_image_multiresolution_dataloader, build_text_to_image_multiresolution_dataloader, + collate_fn_image_to_image, collate_fn_text_to_image, ) + from .image_to_image_dataset import ImageToImageDataset from .text_to_image_dataset import TextToImageDataset except ImportError as exc: # pragma: no cover - environment guard raise ImportError( @@ -53,8 +56,11 @@ ) from exc __all__ = [ + "ImageToImageDataset", "TextToImageDataset", + "build_image_to_image_multiresolution_dataloader", "build_text_to_image_multiresolution_dataloader", + "collate_fn_image_to_image", "collate_fn_text_to_image", ] diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index d669d2a7c4a..4049b646ded 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -38,11 +38,116 @@ from nemo_automodel.components.datasets.diffusion.sampler import SequentialBucketSampler from torchdata.stateful_dataloader import StatefulDataLoader +from .image_to_image_dataset import ImageToImageDataset from .text_to_image_dataset import TextToImageDataset logger = logging.getLogger(__name__) +def _pad_text_conditioning( + batch: list[dict], + embedding_key: str, + mask_key: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Right-pad variable-length embeddings and masks from cached edit samples.""" + + embeddings = [item[embedding_key] for item in batch] + masks = [item[mask_key] for item in batch] + if any(not torch.is_tensor(value) or value.ndim != 2 for value in embeddings): + shapes = [getattr(value, "shape", None) for value in embeddings] + raise ValueError(f"{embedding_key} values must have shape [seq,dim], got {shapes}") + hidden_dims = {value.shape[1] for value in embeddings} + dtypes = {value.dtype for value in embeddings} + if len(hidden_dims) != 1 or len(dtypes) != 1: + raise ValueError( + f"{embedding_key} hidden dimensions/dtypes must match across a batch: " + f"dims={hidden_dims}, dtypes={dtypes}" + ) + + max_length = max(value.shape[0] for value in embeddings) + hidden_dim = embeddings[0].shape[1] + padded = embeddings[0].new_zeros((len(batch), max_length, hidden_dim)) + padded_mask = torch.zeros((len(batch), max_length), dtype=torch.long) + for index, (embedding, mask) in enumerate(zip(embeddings, masks)): + if not torch.is_tensor(mask) or mask.ndim != 1 or mask.shape[0] != embedding.shape[0]: + raise ValueError( + f"{mask_key} for sample {index} must have shape [{embedding.shape[0]}], " + f"got {getattr(mask, 'shape', None)}" + ) + length = embedding.shape[0] + padded[index, :length] = embedding + padded_mask[index, :length] = mask.to(dtype=torch.long) + return padded, padded_mask + + +def collate_fn_image_to_image(batch: list[dict]) -> dict: + """Build a Qwen-Image-Edit batch and validate reference compatibility. + + ``conditioning_latents`` remains a list, one entry per reference image. Each entry is a + stacked ``[B,C,H,W]`` tensor, which preserves support for references with different aspect + ratios while ensuring a given reference slot is stackable across the batch. + """ + + if not batch: + raise ValueError("Cannot collate an empty image-to-image batch") + resolutions = {tuple(item["crop_resolution"].tolist()) for item in batch} + if len(resolutions) != 1: + raise ValueError(f"Mixed target resolutions in batch: {resolutions}") + + reference_counts = {len(item["conditioning_latents"]) for item in batch} + if len(reference_counts) != 1: + raise ValueError(f"Mixed conditioning-image counts in batch: {reference_counts}") + reference_count = reference_counts.pop() + if reference_count < 1: + raise ValueError("Every image-to-image sample must contain at least one reference") + + conditioning_latents = [] + for reference_index in range(reference_count): + shapes = {tuple(item["conditioning_latents"][reference_index].shape) for item in batch} + if len(shapes) != 1: + raise ValueError( + f"Reference {reference_index} has mixed latent shapes in batch: {shapes}" + ) + conditioning_latents.append( + torch.stack([item["conditioning_latents"][reference_index] for item in batch]) + ) + + text_embeddings, text_mask = _pad_text_conditioning( + batch, "prompt_embeds", "prompt_embeds_mask" + ) + negative_embeddings, negative_mask = _pad_text_conditioning( + batch, "negative_prompt_embeds", "negative_prompt_embeds_mask" + ) + + image_batch = { + "image_latents": torch.stack([item["latent"] for item in batch]), + "conditioning_latents": conditioning_latents, + "data_type": "image_edit", + "text_embeddings": text_embeddings, + "text_embeddings_mask": text_mask, + "negative_text_embeddings": negative_embeddings, + "negative_text_embeddings_mask": negative_mask, + "metadata": { + "sample_ids": [item["sample_id"] for item in batch], + "prompts": [item["prompt"] for item in batch], + "negative_prompts": [item["negative_prompt"] for item in batch], + "image_paths": [item["image_path"] for item in batch], + "conditioning_image_paths": [item["conditioning_image_paths"] for item in batch], + "conditioning_resolutions": [item["conditioning_resolutions"] for item in batch], + "target_latent_shapes": [item["target_latent_shape"] for item in batch], + "conditioning_latent_shapes": [item["conditioning_latent_shapes"] for item in batch], + "bucket_ids": [item["bucket_id"] for item in batch], + "aspect_ratios": [item["aspect_ratio"] for item in batch], + "crop_resolution": torch.stack([item["crop_resolution"] for item in batch]), + "original_resolution": torch.stack([item["original_resolution"] for item in batch]), + "crop_offset": torch.stack([item["crop_offset"] for item in batch]), + }, + } + if "source_metadata" in batch[0]: + image_batch["metadata"]["source_metadata"] = [item.get("source_metadata") for item in batch] + return image_batch + + def collate_fn_text_to_image( batch: list[dict], negative_text_embeddings: torch.Tensor | None = None, @@ -240,3 +345,69 @@ def build_text_to_image_multiresolution_dataloader( dp_world_size, ) return dataloader, sampler + + +def build_image_to_image_multiresolution_dataloader( + *, + cache_dir: str, + train_text_encoder: bool = False, + batch_size: int = 1, + dp_rank: int = 0, + dp_world_size: int = 1, + base_resolution: tuple[int, int] = (256, 256), + drop_last: bool = True, + shuffle: bool = True, + dynamic_batch_size: bool = False, + num_workers: int = 4, + pin_memory: bool = True, + prefetch_factor: int = 2, +) -> tuple[StatefulDataLoader, SequentialBucketSampler]: + """Build the cached Qwen-Image-Edit multiresolution dataloader. + + Positive and negative embeddings are both sample-specific because each contains visual + tokens from that sample's references, so this builder intentionally has no static negative + prompt embedding argument. + """ + + dataset = ImageToImageDataset( + cache_dir=cache_dir, + train_text_encoder=train_text_encoder, + ) + sampler = SequentialBucketSampler( + dataset, + base_batch_size=batch_size, + base_resolution=base_resolution, + drop_last=drop_last, + shuffle_buckets=shuffle, + shuffle_within_bucket=shuffle, + dynamic_batch_size=dynamic_batch_size, + num_replicas=dp_world_size, + rank=dp_rank, + ) + dataloader = StatefulDataLoader( + dataset, + batch_sampler=sampler, + collate_fn=collate_fn_image_to_image, + num_workers=num_workers, + pin_memory=pin_memory, + prefetch_factor=prefetch_factor if num_workers > 0 else None, + persistent_workers=num_workers > 0, + ) + logger.info( + "image-to-image dataloader | cache_dir=%s size=%d batches/epoch=%d batch_size=%d dp=%d/%d", + cache_dir, + len(dataset), + len(sampler), + batch_size, + dp_rank, + dp_world_size, + ) + return dataloader, sampler + + +__all__ = [ + "build_image_to_image_multiresolution_dataloader", + "build_text_to_image_multiresolution_dataloader", + "collate_fn_image_to_image", + "collate_fn_text_to_image", +] diff --git a/examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py b/examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py new file mode 100644 index 00000000000..e12d7d7289c --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cached image-to-image dataset for Qwen-Image-Edit DMD2 training.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +from nemo_automodel.components.datasets.diffusion.base_dataset import BaseMultiresolutionDataset + + +def _remove_cache_batch_dim( + tensor: torch.Tensor, + name: str, + unbatched_ndim: int, +) -> torch.Tensor: + """Remove the singleton encoder batch dimension while rejecting malformed caches.""" + + if not torch.is_tensor(tensor): + raise TypeError(f"Cached {name!r} must be a tensor, got {type(tensor).__name__}") + if tensor.ndim == unbatched_ndim + 1 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + if tensor.ndim != unbatched_ndim: + raise ValueError( + f"Cached {name!r} must be {unbatched_ndim}D after removing an optional " + f"singleton batch dimension, got shape {tuple(tensor.shape)}" + ) + return tensor + + +class ImageToImageDataset(BaseMultiresolutionDataset): + """Read target/reference latents and per-sample multimodal prompt embeddings.""" + + def __init__(self, cache_dir: str, train_text_encoder: bool = False): + if train_text_encoder: + raise NotImplementedError( + "Qwen-Image-Edit requires cached multimodal embeddings; on-the-fly text encoder " + "training is not supported by ImageToImageDataset." + ) + self.train_text_encoder = False + super().__init__(cache_dir, quantization=64) + + def _validated_cache_file(self, item: dict[str, Any]) -> Path: + cache_file = Path(item["cache_file"]).resolve() + cache_dir = Path(self.cache_dir).resolve() + try: + cache_file.relative_to(cache_dir) + except ValueError as exc: + raise ValueError( + f"Cache file {cache_file} is outside cache directory {cache_dir}" + ) from exc + return cache_file + + def __getitem__(self, idx: int) -> dict[str, Any]: + item = self.metadata[idx] + data = torch.load( + self._validated_cache_file(item), + map_location="cpu", + weights_only=True, + ) + target_latent = data.get("latent") + if not torch.is_tensor(target_latent) or target_latent.ndim != 3: + raise ValueError(f"Cache item {idx} target latent must have shape [C,H,W]") + + conditioning_latents = data.get("conditioning_latents") + if not isinstance(conditioning_latents, list) or not conditioning_latents: + raise ValueError( + f"Cache item {idx} must contain a non-empty `conditioning_latents` list" + ) + if not all(torch.is_tensor(latent) and latent.ndim == 3 for latent in conditioning_latents): + raise ValueError(f"Cache item {idx} conditioning latents must all have shape [C,H,W]") + + resolution_key = "bucket_resolution" if "bucket_resolution" in item else "crop_resolution" + prompt_embeds = _remove_cache_batch_dim( + data["prompt_embeds"], "prompt_embeds", unbatched_ndim=2 + ) + negative_prompt_embeds = _remove_cache_batch_dim( + data["negative_prompt_embeds"], "negative_prompt_embeds", unbatched_ndim=2 + ) + prompt_mask = data.get("prompt_embeds_mask") + if prompt_mask is None: + prompt_mask = torch.ones(prompt_embeds.shape[0], dtype=torch.long) + else: + prompt_mask = _remove_cache_batch_dim( + prompt_mask, "prompt_embeds_mask", unbatched_ndim=1 + ).long() + negative_mask = data.get("negative_prompt_embeds_mask") + if negative_mask is None: + negative_mask = torch.ones(negative_prompt_embeds.shape[0], dtype=torch.long) + else: + negative_mask = _remove_cache_batch_dim( + negative_mask, + "negative_prompt_embeds_mask", + unbatched_ndim=1, + ).long() + + output = { + "latent": target_latent, + "conditioning_latents": conditioning_latents, + "prompt_embeds": prompt_embeds, + "prompt_embeds_mask": prompt_mask, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_prompt_embeds_mask": negative_mask, + "crop_resolution": torch.tensor(item[resolution_key]), + "original_resolution": torch.tensor(item["original_resolution"]), + "crop_offset": torch.tensor(data["crop_offset"]), + "prompt": data["prompt"], + "negative_prompt": data.get("negative_prompt", " "), + "image_path": data["image_path"], + "conditioning_image_paths": data["conditioning_image_paths"], + "conditioning_resolutions": data.get( + "conditioning_resolutions", + [None] * len(conditioning_latents), + ), + "target_latent_shape": data.get("target_latent_shape", tuple(target_latent.shape)), + "conditioning_latent_shapes": data.get( + "conditioning_latent_shapes", + [tuple(value.shape) for value in conditioning_latents], + ), + "sample_id": data.get("sample_id", str(idx)), + "bucket_id": item["bucket_id"], + "aspect_ratio": item.get("aspect_ratio", 1.0), + } + if "source_metadata" in data: + output["source_metadata"] = data["source_metadata"] + return output + + +__all__ = ["ImageToImageDataset"] diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py index 5907d0f1b86..7297679db8d 100644 --- a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py +++ b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py @@ -294,15 +294,6 @@ def __call__( num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) - txt_seq_lens = ( - prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None - ) - neg_txt_seq_lens = ( - neg_prompt_embeds_mask.sum(dim=1).int().tolist() - if neg_prompt_embeds_mask is not None - else None - ) - # ---- 3. Build initial noisy latents at t = schedule[0] --------------- if isinstance(prompt, str): batch_size = 1 @@ -334,7 +325,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, guidance=None, return_dict=False, )[0] @@ -352,7 +342,6 @@ def __call__( encoder_hidden_states_mask=neg_prompt_embeds_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=neg_txt_seq_lens, guidance=None, return_dict=False, )[0] diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py b/examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py new file mode 100644 index 00000000000..25bdf687142 --- /dev/null +++ b/examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Few-step inference for a DMD2-trained Qwen-Image-Edit-2511 student. + +The stock EditPlus pipeline is reused for multimodal prompt encoding, reference-image +preprocessing, VAE encode/decode, and output postprocessing. Only its denoising loop is +replaced with the exact rectified-flow schedule used by DMD2 training. Target tokens are +followed by the fixed reference-image tokens on every transformer call; only the target +prediction prefix is stepped. +""" + +from __future__ import annotations + +import argparse +import itertools +import logging +import math +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from diffusers import QwenImageEditPlusPipeline, QwenImageTransformer2DModel +from diffusers.utils import load_image +from diffusers.utils.torch_utils import randn_tensor + +logger = logging.getLogger(__name__) + +_CONDITION_IMAGE_AREA = 384 * 384 +_VAE_IMAGE_AREA = 1024 * 1024 + + +def _calculate_dimensions(target_area: int, ratio: float) -> tuple[int, int]: + """Match Diffusers EditPlus' area-preserving, 32-pixel-quantized resize.""" + raw_width = math.sqrt(target_area * ratio) + raw_height = raw_width / ratio + return round(raw_width / 32) * 32, round(raw_height / 32) * 32 + + +def _overlay_ema(student: torch.nn.Module, ema_path: str | os.PathLike[str]) -> None: + payload = torch.load(str(ema_path), map_location="cpu", weights_only=True) + shadow = payload.get("shadow", payload) if isinstance(payload, dict) else payload + if not isinstance(shadow, dict): + raise TypeError(f"EMA payload must be a state dict, got {type(shadow).__name__}.") + missing, unexpected = student.load_state_dict(shadow, strict=False) + if missing or unexpected: + logger.warning("EMA overlay: %d missing, %d unexpected keys", len(missing), len(unexpected)) + + +@dataclass +class QwenImageEditDMDOutput: + images: list[Any] + + +class QwenImageEditDMDInferencePipeline: + """DMD sampler around a stock :class:`QwenImageEditPlusPipeline`.""" + + def __init__(self, pipeline: QwenImageEditPlusPipeline, max_t: float = 0.999) -> None: + self._pipe = pipeline + self.max_t = float(max_t) + + @classmethod + def from_pretrained( + cls, + student_path: str | os.PathLike[str], + base_pipeline_path: str | os.PathLike[str] = "Qwen/Qwen-Image-Edit-2511", + *, + ema_path: str | os.PathLike[str] | None = None, + torch_dtype: torch.dtype = torch.bfloat16, + max_t: float = 0.999, + ) -> QwenImageEditDMDInferencePipeline: + student_path = str(student_path) + if not os.path.isdir(student_path): + raise FileNotFoundError(f"student_path is not a directory: {student_path}") + student = QwenImageTransformer2DModel.from_pretrained(student_path, torch_dtype=torch_dtype) + if ema_path is not None: + _overlay_ema(student, ema_path) + student.eval() + pipeline = QwenImageEditPlusPipeline.from_pretrained( + str(base_pipeline_path), transformer=student, torch_dtype=torch_dtype + ) + return cls(pipeline, max_t=max_t) + + def to(self, device: str | torch.device) -> QwenImageEditDMDInferencePipeline: + self._pipe.to(device) + return self + + @property + def device(self) -> torch.device: + return self._pipe.transformer.device + + @property + def dtype(self) -> torch.dtype: + return next(self._pipe.transformer.parameters()).dtype + + @staticmethod + def _resolve_schedule( + num_inference_steps: int, + max_t: float, + t_list: list[float] | None, + ) -> list[float]: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + if t_list is None: + return torch.linspace(max_t, 0.0, num_inference_steps + 1).tolist() + if len(t_list) != num_inference_steps + 1: + raise ValueError("t_list must contain num_inference_steps + 1 entries.") + schedule = [float(value) for value in t_list] + if abs(schedule[-1]) > 1e-6: + raise ValueError("t_list must end at 0.0.") + if any(left <= right for left, right in itertools.pairwise(schedule)): + raise ValueError("t_list must be strictly decreasing.") + return schedule + + @torch.no_grad() + def __call__( + self, + image: Any | list[Any], + prompt: str, + *, + negative_prompt: str | None = None, + num_inference_steps: int = 4, + guidance_scale: float = 1.0, + height: int | None = None, + width: int | None = None, + generator: torch.Generator | None = None, + max_t: float | None = None, + t_list: list[float] | None = None, + sample_type: str = "ode", + output_type: str = "pil", + ) -> QwenImageEditDMDOutput: + """Edit one image from one or more ordered references. + + A CFG-trained DMD2 student has already internalized teacher guidance, so the + default ``guidance_scale=1`` performs a single transformer call per step. + """ + if sample_type not in {"ode", "sde"}: + raise ValueError("sample_type must be 'ode' or 'sde'.") + references = list(image) if isinstance(image, (list, tuple)) else [image] + if not references: + raise ValueError("At least one reference image is required.") + + pipe = self._pipe + device, dtype = self.device, self.dtype + max_t = self.max_t if max_t is None else float(max_t) + schedule = self._resolve_schedule(num_inference_steps, max_t, t_list) + + # Match QwenImageEditPlusPipeline.__call__: the last reference determines the + # default target aspect ratio; each reference gets separate vision/VAE resolutions. + last_width, last_height = references[-1].size + default_width, default_height = _calculate_dimensions( + _VAE_IMAGE_AREA, last_width / last_height + ) + width = int(width or default_width) + height = int(height or default_height) + multiple = pipe.vae_scale_factor * 2 + width, height = width // multiple * multiple, height // multiple * multiple + + condition_images: list[Any] = [] + vae_images: list[torch.Tensor] = [] + vae_sizes: list[tuple[int, int]] = [] + for reference in references: + ref_width, ref_height = reference.size + ratio = ref_width / ref_height + cond_width, cond_height = _calculate_dimensions(_CONDITION_IMAGE_AREA, ratio) + vae_width, vae_height = _calculate_dimensions(_VAE_IMAGE_AREA, ratio) + condition_images.append(pipe.image_processor.resize(reference, cond_height, cond_width)) + vae_images.append( + pipe.image_processor.preprocess(reference, vae_height, vae_width).unsqueeze(2) + ) + vae_sizes.append((vae_width, vae_height)) + + prompt_embeds, prompt_mask = pipe.encode_prompt( + image=condition_images, prompt=prompt, device=device, num_images_per_prompt=1 + ) + do_cfg = guidance_scale != 1.0 + negative_embeds = negative_mask = None + if do_cfg: + negative_prompt = " " if negative_prompt is None else negative_prompt + negative_embeds, negative_mask = pipe.encode_prompt( + image=condition_images, + prompt=negative_prompt, + device=device, + num_images_per_prompt=1, + ) + + channels = pipe.transformer.config.in_channels // 4 + h_lat = 2 * (height // (pipe.vae_scale_factor * 2)) + w_lat = 2 * (width // (pipe.vae_scale_factor * 2)) + noise = randn_tensor( + (1, 1, channels, h_lat, w_lat), generator=generator, device=device, dtype=dtype + ) + target = pipe._pack_latents(noise * schedule[0], 1, channels, h_lat, w_lat) + target_tokens = target.shape[1] + + packed_references: list[torch.Tensor] = [] + for vae_image in vae_images: + ref_latent = pipe._encode_vae_image( + vae_image.to(device=device, dtype=dtype), generator=generator + ) + ref_h, ref_w = ref_latent.shape[3:] + packed_references.append(pipe._pack_latents(ref_latent, 1, channels, ref_h, ref_w)) + img_shapes = [ + [ + (1, h_lat // 2, w_lat // 2), + *[ + ( + 1, + vae_height // pipe.vae_scale_factor // 2, + vae_width // pipe.vae_scale_factor // 2, + ) + for vae_width, vae_height in vae_sizes + ], + ] + ] + + x = target + fixed_references = torch.cat(packed_references, dim=1) + for t_cur, t_next in itertools.pairwise(schedule): + model_input = torch.cat([x, fixed_references], dim=1) + timestep = torch.full((1,), float(t_cur), device=device, dtype=dtype) + flow = pipe.transformer( + hidden_states=model_input, + timestep=timestep, + guidance=None, + encoder_hidden_states_mask=prompt_mask, + encoder_hidden_states=prompt_embeds, + img_shapes=img_shapes, + return_dict=False, + )[0][:, :target_tokens] + if do_cfg: + negative_flow = pipe.transformer( + hidden_states=model_input, + timestep=timestep, + guidance=None, + encoder_hidden_states_mask=negative_mask, + encoder_hidden_states=negative_embeds, + img_shapes=img_shapes, + return_dict=False, + )[0][:, :target_tokens] + flow = ( + negative_flow.double() + + float(guidance_scale) * (flow.double() - negative_flow.double()) + ).to(dtype) + + x0 = (x.double() - float(t_cur) * flow.double()).to(dtype) + if t_next <= 1e-6: + x = x0 + continue + if sample_type == "ode": + eps = ( + (x.double() - (1.0 - float(t_cur)) * x0.double()) / max(float(t_cur), 1e-6) + ).to(dtype) + else: + eps = torch.randn(x.shape, generator=generator, device=device, dtype=dtype) + x = ((1.0 - float(t_next)) * x0.double() + float(t_next) * eps.double()).to(dtype) + + decoded_latents = pipe._unpack_latents(x, height, width, pipe.vae_scale_factor) + latents_mean = ( + torch.tensor(pipe.vae.config.latents_mean) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device, dtype) + ) + latents_std = ( + torch.tensor(pipe.vae.config.latents_std) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device, dtype) + ) + decoded_latents = decoded_latents * latents_std + latents_mean + decoded = pipe.vae.decode(decoded_latents, return_dict=False)[0][:, :, 0] + return QwenImageEditDMDOutput( + images=pipe.image_processor.postprocess(decoded, output_type=output_type) + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--student-path", required=True) + parser.add_argument("--base-pipeline-path", default="Qwen/Qwen-Image-Edit-2511") + parser.add_argument("--image", nargs="+", required=True, help="Ordered reference image(s).") + parser.add_argument("--prompt", required=True) + parser.add_argument("--negative-prompt") + parser.add_argument("--num-inference-steps", type=int, default=4) + parser.add_argument("--guidance-scale", type=float, default=1.0) + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--ema-path") + parser.add_argument("--output", default="qwen_image_edit_dmd2.png") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(args.seed) + pipeline = QwenImageEditDMDInferencePipeline.from_pretrained( + args.student_path, + args.base_pipeline_path, + ema_path=args.ema_path, + ).to(device) + references = [load_image(path).convert("RGB") for path in args.image] + output = pipeline( + references, + args.prompt, + negative_prompt=args.negative_prompt, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + height=args.height, + width=args.width, + generator=generator, + ) + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + output.images[0].save(args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/preprocess/processors/__init__.py b/examples/diffusers/fastgen/preprocess/processors/__init__.py index 2660d27b105..a9d3e67448e 100644 --- a/examples/diffusers/fastgen/preprocess/processors/__init__.py +++ b/examples/diffusers/fastgen/preprocess/processors/__init__.py @@ -23,6 +23,7 @@ get_caption_loader, ) from .qwen_image import QwenImageProcessor +from .qwen_image_edit import QwenImageEditProcessor from .registry import ProcessorRegistry __all__ = [ @@ -33,6 +34,7 @@ "JSONSidecarCaptionLoader", "MetaJSONCaptionLoader", "ProcessorRegistry", + "QwenImageEditProcessor", "QwenImageProcessor", "get_caption_loader", ] diff --git a/examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py b/examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py new file mode 100644 index 00000000000..f1af7f8c5d4 --- /dev/null +++ b/examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image-Edit preprocessing support. + +Qwen-Image-Edit-2511 conditions the denoiser through two independent paths: + +* every reference image is encoded by the Qwen2.5-VL prompt encoder together with the edit + instruction; and +* every reference image is encoded by the Qwen-Image VAE and appended to the noisy target + tokens by the denoiser adapter. + +This processor intentionally keeps those two representations separate in the cache. Target +latents use the same sampled-VAE convention as :class:`QwenImageProcessor`, while reference +latents use the deterministic posterior mode used by ``QwenImageEditPlusPipeline`` at inference. +""" + +from __future__ import annotations + +import logging +import math +from typing import TYPE_CHECKING, Any + +import torch + +from .qwen_image import QwenImageProcessor +from .registry import ProcessorRegistry + +if TYPE_CHECKING: + from PIL import Image + +logger = logging.getLogger(__name__) + +_CONDITION_IMAGE_AREA = 384 * 384 +_VAE_IMAGE_AREA = 1024 * 1024 + + +def _dimensions_for_area(image: Image.Image, area: int) -> tuple[int, int]: + """Return ``(width, height)`` preserving aspect ratio, rounded to multiples of 32.""" + + width, height = image.size + if width <= 0 or height <= 0: + raise ValueError(f"Invalid image dimensions: {image.size}") + ratio = width / height + resized_width = max(32, round(math.sqrt(area * ratio) / 32) * 32) + resized_height = max(32, round(math.sqrt(area / ratio) / 32) * 32) + return resized_width, resized_height + + +def _posterior_mode(encoder_output: Any) -> torch.Tensor: + """Extract the deterministic VAE posterior mode across diffusers return variants.""" + + if hasattr(encoder_output, "latent_dist"): + return encoder_output.latent_dist.mode() + if hasattr(encoder_output, "latents"): + return encoder_output.latents + raise AttributeError("Could not access VAE latents from encoder output") + + +@ProcessorRegistry.register("qwen_image_edit") +class QwenImageEditProcessor(QwenImageProcessor): + """Precompute the full Qwen-Image-Edit-2511 conditioning contract.""" + + @property + def model_type(self) -> str: + return "qwen_image_edit" + + @property + def default_model_name(self) -> str: + return "Qwen/Qwen-Image-Edit-2511" + + def load_models(self, model_name: str, device: str) -> dict[str, Any]: + """Load only the VAE and multimodal prompt encoder needed for caching.""" + + try: + from diffusers import QwenImageEditPlusPipeline + except ImportError as exc: # pragma: no cover - depends on the runtime environment + raise ImportError( + "Qwen-Image-Edit-2511 preprocessing requires a diffusers release that provides " + "QwenImageEditPlusPipeline (introduced in diffusers 0.36.0). Full 2511 " + "denoising also requires `zero_cond_t` support (stable diffusers>=0.37)." + ) from exc + + logger.info("[Qwen-Image-Edit] Loading preprocessing models from %s", model_name) + pipeline = QwenImageEditPlusPipeline.from_pretrained( + model_name, + transformer=None, + torch_dtype=torch.bfloat16, + ) + pipeline.vae.to(device=device, dtype=torch.bfloat16).eval() + pipeline.text_encoder.to(device).eval() + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return {"vae": pipeline.vae, "pipeline": pipeline} + + def _vision_images(self, images: list[Image.Image], pipeline: Any) -> list[Image.Image]: + """Resize references exactly as the Edit Plus prompt-encoding path does.""" + + prepared = [] + for image in images: + width, height = _dimensions_for_area(image, _CONDITION_IMAGE_AREA) + prepared.append(pipeline.image_processor.resize(image, height, width)) + return prepared + + def encode_conditioning_images( + self, + images: list[Image.Image], + models: dict[str, Any], + device: str, + *, + max_pixels: int = _VAE_IMAGE_AREA, + ) -> list[torch.Tensor]: + """Encode one or more references with deterministic VAE posterior modes. + + Each returned tensor is ``[C, H/8, W/8]``. References are kept as a list because the + Edit Plus model permits different aspect ratios for different references. + """ + + if not images: + raise ValueError("Qwen-Image-Edit requires at least one conditioning image") + vae = models["vae"] + pipeline = models["pipeline"] + latents = [] + for image in images: + width, height = _dimensions_for_area(image, max_pixels) + image_tensor = pipeline.image_processor.preprocess(image, height, width).unsqueeze(2) + image_tensor = image_tensor.to(device=device, dtype=torch.bfloat16) + with torch.no_grad(): + latent = _posterior_mode(vae.encode(image_tensor)) + + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latent = (latent - latents_mean) / latents_std + latents.append(latent.detach().cpu().to(torch.float16).squeeze(2).squeeze(0)) + return latents + + def encode_multimodal_text( + self, + prompt: str, + images: list[Image.Image], + models: dict[str, Any], + device: str, + ) -> dict[str, torch.Tensor]: + """Encode an instruction and its reference images with Qwen2.5-VL.""" + + if not images: + raise ValueError("Qwen-Image-Edit prompt encoding requires at least one image") + pipeline = models["pipeline"] + vision_images = self._vision_images(images, pipeline) + with torch.no_grad(): + prompt_embeds, prompt_embeds_mask = pipeline.encode_prompt( + prompt=prompt, + image=vision_images, + device=device, + ) + result = {"prompt_embeds": prompt_embeds.detach().cpu().to(torch.bfloat16)} + if prompt_embeds_mask is not None: + result["prompt_embeds_mask"] = prompt_embeds_mask.detach().cpu().to(torch.long) + return result + + def encode_edit_prompts( + self, + prompt: str, + negative_prompt: str, + images: list[Image.Image], + models: dict[str, Any], + device: str, + ) -> dict[str, torch.Tensor]: + """Encode positive and per-sample negative multimodal conditioning.""" + + positive = self.encode_multimodal_text(prompt, images, models, device) + negative = self.encode_multimodal_text(negative_prompt, images, models, device) + result = dict(positive) + result["negative_prompt_embeds"] = negative["prompt_embeds"] + if "prompt_embeds_mask" in negative: + result["negative_prompt_embeds_mask"] = negative["prompt_embeds_mask"] + return result + + def encode_text( + self, + prompt: str, + models: dict[str, Any], + device: str, + ) -> dict[str, torch.Tensor]: + """Reject text-only use, which would silently omit image tokens from the cache.""" + + raise ValueError( + "QwenImageEditProcessor.encode_text cannot encode a text-only prompt. Use " + "encode_multimodal_text/encode_edit_prompts with the conditioning images." + ) + + def verify_latent( + self, + latent: torch.Tensor, + models: dict[str, Any], + device: str, + ) -> bool: + """Decode a target latent with the VAE's actual dtype and validate finiteness.""" + + try: + vae = models["vae"] + vae_dtype = next(vae.parameters()).dtype + value = latent.unsqueeze(0).unsqueeze(2).to(device=device, dtype=vae_dtype) + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(device=device, dtype=vae_dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(device=device, dtype=vae_dtype) + ) + with torch.no_grad(): + decoded = vae.decode(value * latents_std + latents_mean).sample[:, :, 0] + return ( + decoded.ndim == 4 and decoded.shape[1] == 3 and torch.isfinite(decoded).all().item() + ) + except Exception as exc: + logger.warning("[Qwen-Image-Edit] Latent verification failed: %s", exc) + return False + + def get_cache_data( + self, + latent: torch.Tensor, + text_encodings: dict[str, torch.Tensor], + metadata: dict[str, Any], + ) -> dict[str, Any]: + """Construct an image-edit cache record consumed by ``ImageToImageDataset``.""" + + conditioning_latents = metadata.get("conditioning_latents") + if not isinstance(conditioning_latents, list) or not conditioning_latents: + raise ValueError("metadata['conditioning_latents'] must be a non-empty tensor list") + required_text = ("prompt_embeds", "negative_prompt_embeds") + missing = [key for key in required_text if key not in text_encodings] + if missing: + raise KeyError(f"Missing edit text encodings: {missing}") + + cache = { + "latent": latent, + "conditioning_latents": conditioning_latents, + "prompt_embeds": text_encodings["prompt_embeds"], + "negative_prompt_embeds": text_encodings["negative_prompt_embeds"], + "original_resolution": metadata["original_resolution"], + "bucket_resolution": metadata["bucket_resolution"], + "crop_offset": metadata["crop_offset"], + "prompt": metadata["prompt"], + "negative_prompt": metadata["negative_prompt"], + "image_path": metadata["image_path"], + "conditioning_image_paths": metadata["conditioning_image_paths"], + "conditioning_resolutions": metadata["conditioning_resolutions"], + "target_latent_shape": tuple(latent.shape), + "conditioning_latent_shapes": [tuple(value.shape) for value in conditioning_latents], + "bucket_id": metadata["bucket_id"], + "aspect_ratio": metadata["aspect_ratio"], + "sample_id": metadata["sample_id"], + "model_type": self.model_type, + } + for key in ("prompt_embeds_mask", "negative_prompt_embeds_mask"): + if key in text_encodings: + cache[key] = text_encodings[key] + if metadata.get("source_metadata") is not None: + cache["source_metadata"] = metadata["source_metadata"] + return cache + + +__all__ = ["QwenImageEditProcessor"] diff --git a/examples/diffusers/fastgen/preprocess_qwen_image_edit.py b/examples/diffusers/fastgen/preprocess_qwen_image_edit.py new file mode 100644 index 00000000000..e844e27b0e1 --- /dev/null +++ b/examples/diffusers/fastgen/preprocess_qwen_image_edit.py @@ -0,0 +1,792 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Precompute Qwen-Image-Edit-2511 DMD2 training caches. + +Two local input layouts are supported: + +* SpatialEdit-style WebDataset roots containing ``*.tar`` shards. A sample contains + ``.json`` and indexed images such as ``.0.jpg`` / ``.1.jpg``. Images are + ordered by index; the last image is the target and every preceding image is a reference. +* A JSONL manifest with ``target``, ``conditioning`` (a path or path list), and ``prompt``. + The data-tooling aliases ``generated_image`` / ``reference_image`` (targets) and + ``conditioning_images`` (sources) are also accepted, including ``{archive, member}`` + descriptors. Relative paths are resolved against the manifest directory. ``id``, + ``negative_prompt``, and ``metadata`` are optional. + +The output follows ``BaseMultiresolutionDataset``'s sharded ``metadata.json`` layout. Each +``.pt`` record contains a sampled target ``latent``, a list of deterministic +``conditioning_latents``, and positive/negative multimodal embeddings and masks. +""" + +from __future__ import annotations + +import argparse +import hashlib +import io +import json +import logging +import os +import re +import sys +import tarfile +import traceback +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from PIL import Image + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +logger = logging.getLogger(__name__) + +_IMAGE_MEMBER_RE = re.compile( + r"^(?P.+)\.(?P\d+)\.(?:jpe?g|png|webp)$", + flags=re.IGNORECASE, +) +_IMAGE_MARKER_RE = re.compile(r"(?:\s*)+", flags=re.IGNORECASE) + + +@dataclass +class EditSample: + """One loaded edit pair. Only one instance is retained while streaming input shards.""" + + sample_id: str + target_image: Image.Image + conditioning_images: list[Image.Image] + prompt: str + negative_prompt: str + target_path: str + conditioning_paths: list[str] + source_metadata: dict[str, Any] | None = None + + +def _load_rgb(path: Path) -> Image.Image: + with Image.open(path) as image: + return image.convert("RGB") + + +def _load_tar_rgb(archive: tarfile.TarFile, member: tarfile.TarInfo) -> Image.Image: + fileobj = archive.extractfile(member) + if fileobj is None: + raise OSError(f"Could not read {member.name!r} from {archive.name!r}") + with Image.open(io.BytesIO(fileobj.read())) as image: + return image.convert("RGB") + + +def _strip_image_markers(text: str) -> str: + """Remove dataset placeholders because EditPlus inserts its own visual tokens.""" + + return _IMAGE_MARKER_RE.sub("", text).strip() + + +def _conversation_text(payload: dict[str, Any], key: str) -> str | None: + conversation = payload.get(key) + if not isinstance(conversation, list): + return None + for message in conversation: + if isinstance(message, dict) and message.get("from") in {"human", "user"}: + value = message.get("value") or message.get("text") + if isinstance(value, str) and value.strip(): + return value + return None + + +def extract_spatialedit_prompt(payload: dict[str, Any], variant: str = "human") -> str: + """Extract a usable instruction across SpatialEdit's three shard families.""" + + metadata = payload.get("metadata") + metadata = metadata if isinstance(metadata, dict) else {} + if variant == "human": + candidates = ( + metadata.get("instruction_human"), + _conversation_text(payload, "conversations_human"), + metadata.get("instruction"), + _conversation_text(payload, "conversations"), + ) + elif variant == "raw": + candidates = ( + metadata.get("instruction"), + _conversation_text(payload, "conversations"), + metadata.get("instruction_human"), + _conversation_text(payload, "conversations_human"), + ) + else: + raise ValueError(f"Unknown prompt variant: {variant!r}") + for candidate in candidates: + if isinstance(candidate, str) and candidate.strip(): + prompt = _strip_image_markers(candidate) + if prompt: + return prompt + raise ValueError("SpatialEdit sample has no non-empty edit instruction") + + +def _spatialedit_sample_id(payload: dict[str, Any], fallback: str) -> str: + metadata = payload.get("metadata") + candidates = ( + payload.get("SAMPLE_ID"), + payload.get("id"), + metadata.get("id") if isinstance(metadata, dict) else None, + fallback, + ) + return str(next(value for value in candidates if value is not None and str(value))) + + +def _spatialedit_metadata( + payload: dict[str, Any], + tar_path: Path, + sample_key: str, +) -> dict[str, Any]: + result: dict[str, Any] = { + "source": "SpatialEdit-500K", + "tar_path": str(tar_path.resolve()), + "webdataset_key": sample_key, + } + for key in ("metadata", "meta", "data_type", "multi_image", "only_text"): + if key in payload: + result[key] = payload[key] + return result + + +def iter_spatialedit_samples( + root: Path, + *, + negative_prompt: str = " ", + prompt_variant: str = "human", + shard_rank: int = 0, + shard_world: int = 1, +) -> Iterator[EditSample]: + """Stream native SpatialEdit WebDataset pairs without extracting shards to disk.""" + + tar_paths = sorted(path for path in root.rglob("*.tar") if path.is_file()) + if not tar_paths: + raise FileNotFoundError(f"No .tar shards found under {root}") + selected = tar_paths[shard_rank::shard_world] + logger.info( + "SpatialEdit input: %d/%d tar shards assigned to rank %d", + len(selected), + len(tar_paths), + shard_rank, + ) + + for tar_path in selected: + try: + with tarfile.open(tar_path, mode="r:*") as archive: + grouped: dict[str, dict[str, Any]] = {} + for member in archive.getmembers(): + if not member.isfile(): + continue + if member.name.lower().endswith(".json"): + key = member.name[: -len(".json")] + grouped.setdefault(key, {})["json"] = member + continue + match = _IMAGE_MEMBER_RE.match(member.name) + if match: + group = grouped.setdefault(match.group("key"), {}) + group.setdefault("images", {})[int(match.group("index"))] = member + + for sample_key in sorted(grouped): + members = grouped[sample_key] + image_members = members.get("images", {}) + if "json" not in members or len(image_members) < 2: + logger.warning( + "Skipping incomplete WebDataset sample %s::%s (json=%s, images=%d)", + tar_path, + sample_key, + "json" in members, + len(image_members), + ) + continue + json_file = archive.extractfile(members["json"]) + if json_file is None: + raise OSError(f"Could not read metadata for {tar_path}::{sample_key}") + payload = json.loads(json_file.read()) + if not isinstance(payload, dict): + raise ValueError(f"Metadata for {tar_path}::{sample_key} is not an object") + + ordered = sorted(image_members.items()) + conditioning = [_load_tar_rgb(archive, member) for _, member in ordered[:-1]] + _, target_member = ordered[-1] + target = _load_tar_rgb(archive, target_member) + display_prefix = f"{tar_path.resolve()}::" + yield EditSample( + sample_id=_spatialedit_sample_id(payload, sample_key), + target_image=target, + conditioning_images=conditioning, + prompt=extract_spatialedit_prompt(payload, prompt_variant), + negative_prompt=negative_prompt, + target_path=f"{display_prefix}{target_member.name}", + conditioning_paths=[ + f"{display_prefix}{member.name}" for _, member in ordered[:-1] + ], + source_metadata=_spatialedit_metadata(payload, tar_path, sample_key), + ) + except Exception: + logger.error("Failed while reading WebDataset shard %s", tar_path) + logger.debug(traceback.format_exc()) + raise + + +def _manifest_value(record: dict[str, Any], names: tuple[str, ...]) -> Any: + for name in names: + if name in record and record[name] is not None: + return record[name] + return None + + +def _resolve_manifest_path(value: Any, base_dir: Path, field: str) -> Path: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Manifest field {field!r} must be a non-empty local path") + path = Path(value).expanduser() + if not path.is_absolute(): + path = base_dir / path + path = path.resolve() + if not path.is_file(): + raise FileNotFoundError(f"Manifest {field} image does not exist: {path}") + return path + + +class _ArchiveImageLoader: + """Small LRU of open tar files for archive/member JSONL descriptors.""" + + def __init__(self, max_open_archives: int = 8) -> None: + self.max_open_archives = max_open_archives + self._archives: OrderedDict[Path, tarfile.TarFile] = OrderedDict() + + def _archive(self, path: Path) -> tarfile.TarFile: + archive = self._archives.pop(path, None) + if archive is None: + archive = tarfile.open(path, mode="r:*") + self._archives[path] = archive + while len(self._archives) > self.max_open_archives: + _, stale = self._archives.popitem(last=False) + stale.close() + return archive + + def load(self, archive_path: Path, member_name: str) -> Image.Image: + archive = self._archive(archive_path) + try: + member = archive.getmember(member_name) + except KeyError as exc: + raise FileNotFoundError( + f"Archive member does not exist: {archive_path}::{member_name}" + ) from exc + return _load_tar_rgb(archive, member) + + def close(self) -> None: + for archive in self._archives.values(): + archive.close() + self._archives.clear() + + +def _load_manifest_image( + value: Any, + base_dir: Path, + field: str, + archive_loader: _ArchiveImageLoader, +) -> tuple[Image.Image, str]: + """Load a local path or ``{archive, member}`` image descriptor.""" + + if isinstance(value, str): + path = _resolve_manifest_path(value, base_dir, field) + return _load_rgb(path), str(path) + if not isinstance(value, dict): + raise ValueError( + f"Manifest field {field!r} must be a path or an {{archive, member}} object" + ) + if "path" in value: + path = _resolve_manifest_path(value["path"], base_dir, field) + return _load_rgb(path), str(path) + + archive_path = _resolve_manifest_path(value.get("archive"), base_dir, f"{field}.archive") + member = value.get("member") + if not isinstance(member, str) or not member: + raise ValueError(f"Manifest field {field!r}.member must be a non-empty string") + return archive_loader.load(archive_path, member), f"{archive_path}::{member}" + + +def iter_jsonl_samples( + manifest: Path, + *, + negative_prompt: str = " ", + shard_rank: int = 0, + shard_world: int = 1, +) -> Iterator[EditSample]: + """Stream generic local edit records from a JSONL manifest.""" + + base_dir = manifest.resolve().parent + archive_loader = _ArchiveImageLoader() + try: + with manifest.open("r", encoding="utf-8") as handle: + for line_index, line in enumerate(handle): + if not line.strip() or line.lstrip().startswith("#"): + continue + record = json.loads(line) + if not isinstance(record, dict): + raise ValueError(f"Manifest line {line_index + 1} is not a JSON object") + shard_value = record.get("archive_index") + if shard_value is None: + shard_value = line_index + try: + assigned_rank = int(shard_value) % shard_world + except (TypeError, ValueError) as exc: + raise ValueError( + f"Manifest line {line_index + 1} archive_index must be an integer" + ) from exc + if assigned_rank != shard_rank: + continue + + target_value = _manifest_value( + record, + ( + "target", + "target_image", + "generated_image", + "output_image", + "output", + "reference_image", + ), + ) + conditioning_value = _manifest_value( + record, + ( + "conditioning", + "conditioning_images", + "reference_images", + "source_images", + "source", + "input", + ), + ) + prompt_value = _manifest_value( + record, ("prompt", "instruction", "edit_instruction") + ) + if isinstance(conditioning_value, (str, dict)): + conditioning_value = [conditioning_value] + if not isinstance(conditioning_value, list) or not conditioning_value: + raise ValueError( + f"Manifest line {line_index + 1} must provide one or more " + "conditioning images" + ) + if not isinstance(prompt_value, str) or not prompt_value.strip(): + raise ValueError(f"Manifest line {line_index + 1} has no edit prompt") + + target_image, target_path = _load_manifest_image( + target_value, + base_dir, + "target", + archive_loader, + ) + loaded_conditioning = [ + _load_manifest_image(value, base_dir, "conditioning", archive_loader) + for value in conditioning_value + ] + per_sample_negative = record.get("negative_prompt", negative_prompt) + if not isinstance(per_sample_negative, str): + raise ValueError( + f"Manifest line {line_index + 1} negative_prompt must be a string" + ) + sample_id = next( + str(value) + for value in ( + record.get("id"), + record.get("sample_id"), + record.get("source_id"), + record.get("key"), + line_index, + ) + if value is not None and str(value) + ) + yield EditSample( + sample_id=sample_id, + target_image=target_image, + conditioning_images=[image for image, _ in loaded_conditioning], + prompt=_strip_image_markers(prompt_value), + negative_prompt=per_sample_negative, + target_path=target_path, + conditioning_paths=[path for _, path in loaded_conditioning], + source_metadata=record.get("metadata"), + ) + finally: + archive_loader.close() + + +def _write_json_atomic(path: Path, payload: Any) -> None: + temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, ensure_ascii=False) + os.replace(temporary, path) + + +class MetadataShardWriter: + """Incrementally write metadata shards so 500K samples need not stay in memory.""" + + def __init__( + self, + output_dir: Path, + shard_size: int, + shard_rank: int, + shard_world: int, + ) -> None: + self.output_dir = output_dir + self.shard_size = shard_size + self.shard_rank = shard_rank + self.shard_world = shard_world + self.buffer: list[dict[str, Any]] = [] + self.shards: list[str] = [] + self.total_items = 0 + + def add(self, item: dict[str, Any]) -> None: + self.buffer.append(item) + self.total_items += 1 + if len(self.buffer) >= self.shard_size: + self.flush() + + def flush(self) -> None: + if not self.buffer: + return + rank_prefix = f"r{self.shard_rank:02d}_" if self.shard_world > 1 else "" + filename = f"metadata_shard_{rank_prefix}s{len(self.shards):04d}.json" + _write_json_atomic(self.output_dir / filename, self.buffer) + self.shards.append(filename) + self.buffer = [] + + def finish(self, **config: Any) -> Path: + self.flush() + if not self.shards: + raise RuntimeError("No valid samples were preprocessed; metadata was not written") + index_name = ( + f"metadata_r{self.shard_rank:02d}.json" if self.shard_world > 1 else "metadata.json" + ) + payload = { + "processor": "qwen_image_edit", + "model_type": "qwen_image_edit", + "total_items": self.total_items, + "num_shards": len(self.shards), + "shard_size": self.shard_size, + "shards": self.shards, + **config, + } + if self.shard_world > 1: + payload.update(shard_rank=self.shard_rank, shard_world=self.shard_world) + index_path = self.output_dir / index_name + _write_json_atomic(index_path, payload) + return index_path + + +def _cache_identity( + sample: EditSample, + model_name: str, + resolution: tuple[int, int], + conditioning_max_pixels: int, +) -> str: + fields = ( + model_name, + sample.sample_id, + sample.target_path, + *sample.conditioning_paths, + sample.prompt, + sample.negative_prompt, + f"{resolution[0]}x{resolution[1]}", + str(conditioning_max_pixels), + ) + return hashlib.sha256("\0".join(fields).encode("utf-8")).hexdigest() + + +def _stable_sample_seed(sample_id: str, seed: int) -> int: + digest = hashlib.sha256(f"{seed}\0{sample_id}".encode()).digest() + return int.from_bytes(digest[:8], byteorder="big") % (2**31) + + +def preprocess_samples( + samples: Iterable[EditSample], + *, + output_dir: Path, + model_name: str, + device: str, + max_pixels: int, + conditioning_max_pixels: int, + metadata_shard_size: int, + shard_rank: int, + shard_world: int, + verify: bool, + overwrite: bool, + fail_fast: bool, + limit: int | None, + seed: int, + log_every: int, +) -> Path: + """Encode a stream of edit samples and return the generated metadata index path.""" + + import torch + from nemo_automodel.components.datasets.diffusion.multi_tier_bucketing import ( + MultiTierBucketCalculator, + ) + from preprocess.processors import QwenImageEditProcessor + + output_dir.mkdir(parents=True, exist_ok=True) + processor = QwenImageEditProcessor() + models = processor.load_models(model_name, device) + calculator = MultiTierBucketCalculator(quantization=64, max_pixels=max_pixels) + writer = MetadataShardWriter( + output_dir, + metadata_shard_size, + shard_rank, + shard_world, + ) + + attempted = 0 + failures = 0 + for sample in samples: + if limit is not None and attempted >= limit: + break + attempted += 1 + try: + original_width, original_height = sample.target_image.size + bucket = calculator.get_bucket_for_image(original_width, original_height) + target_width, target_height = bucket["resolution"] + resolution = (target_width, target_height) + cache_hash = _cache_identity( + sample, + model_name, + resolution, + conditioning_max_pixels, + ) + cache_subdir = output_dir / f"{target_width}x{target_height}" + cache_subdir.mkdir(parents=True, exist_ok=True) + cache_file = cache_subdir / f"{cache_hash}.pt" + + if overwrite or not cache_file.is_file(): + resized_target, crop_offset = calculator.resize_and_crop( + sample.target_image, + target_width, + target_height, + crop_mode="center", + ) + sample_seed = _stable_sample_seed(sample.sample_id, seed) + torch.manual_seed(sample_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(sample_seed) + target_tensor = processor.preprocess_image(resized_target) + latent = processor.encode_image(target_tensor, models, device) + if verify and not processor.verify_latent(latent, models, device): + raise ValueError("target latent verification failed") + + conditioning_latents = processor.encode_conditioning_images( + sample.conditioning_images, + models, + device, + max_pixels=conditioning_max_pixels, + ) + text_encodings = processor.encode_edit_prompts( + sample.prompt, + sample.negative_prompt, + sample.conditioning_images, + models, + device, + ) + cache_metadata = { + "conditioning_latents": conditioning_latents, + "original_resolution": (original_width, original_height), + "bucket_resolution": resolution, + "crop_offset": crop_offset, + "prompt": sample.prompt, + "negative_prompt": sample.negative_prompt, + "image_path": sample.target_path, + "conditioning_image_paths": sample.conditioning_paths, + "conditioning_resolutions": [ + tuple(image.size) for image in sample.conditioning_images + ], + "bucket_id": bucket["id"], + "aspect_ratio": bucket["aspect_ratio"], + "sample_id": sample.sample_id, + "source_metadata": sample.source_metadata, + } + cache = processor.get_cache_data(latent, text_encodings, cache_metadata) + temporary = cache_file.with_name(f".{cache_file.name}.tmp-{os.getpid()}") + torch.save(cache, temporary) + os.replace(temporary, cache_file) + else: + crop_offset = (0, 0) + + writer.add( + { + "cache_file": str(cache_file.resolve()), + "image_path": sample.target_path, + "conditioning_image_paths": sample.conditioning_paths, + "conditioning_resolutions": [ + list(image.size) for image in sample.conditioning_images + ], + "sample_id": sample.sample_id, + "bucket_resolution": [target_width, target_height], + "original_resolution": [original_width, original_height], + "prompt": sample.prompt, + "bucket_id": bucket["id"], + "aspect_ratio": bucket["aspect_ratio"], + "pixels": target_width * target_height, + "model_type": processor.model_type, + } + ) + if log_every > 0 and attempted % log_every == 0: + logger.info( + "Processed %d samples (%d failures, %d cached records)", + attempted, + failures, + writer.total_items, + ) + except Exception as exc: + failures += 1 + logger.error("Failed sample %s: %s", sample.sample_id, exc) + logger.debug(traceback.format_exc()) + if fail_fast: + raise + + index_path = writer.finish( + model_name=model_name, + max_pixels=max_pixels, + conditioning_max_pixels=conditioning_max_pixels, + attempted_items=attempted, + failed_items=failures, + negative_prompt_is_per_sample=True, + ) + logger.info( + "Finished preprocessing: %d records, %d failures; metadata=%s", + writer.total_items, + failures, + index_path, + ) + return index_path + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( + "--webdataset-root", + "--input-dir", + dest="webdataset_root", + type=Path, + help="SpatialEdit-style root recursively containing WebDataset .tar shards", + ) + source.add_argument( + "--manifest", + type=Path, + help="Local JSONL with target, conditioning, and prompt fields", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model-name", default="Qwen/Qwen-Image-Edit-2511") + parser.add_argument( + "--device", + default=None, + help="Torch device. Defaults to cuda or cuda: when --gpu-id is provided.", + ) + parser.add_argument("--gpu-id", type=int, help="GPU index used when --device is omitted") + parser.add_argument("--max-pixels", type=_positive_int, default=1024 * 1024) + parser.add_argument( + "--conditioning-max-pixels", + type=_positive_int, + default=1024 * 1024, + help="Per-reference VAE pixel budget (the official EditPlus default is 1024^2)", + ) + parser.add_argument("--negative-prompt", default=" ") + parser.add_argument( + "--prompt-variant", + choices=("human", "raw"), + default="human", + help="SpatialEdit instruction variant; ignored for JSONL manifests", + ) + parser.add_argument("--metadata-shard-size", type=_positive_int, default=10_000) + parser.add_argument("--shard-rank", "--shard-idx", dest="shard_rank", type=int, default=0) + parser.add_argument( + "--shard-world", + "--shard-count", + dest="shard_world", + type=_positive_int, + default=1, + ) + parser.add_argument("--limit", "--max-samples", dest="limit", type=_positive_int) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--log-every", type=int, default=100) + parser.add_argument("--verify", action="store_true") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--fail-fast", action="store_true") + return parser + + +def main(argv: list[str] | None = None) -> Path: + args = build_parser().parse_args(argv) + if not 0 <= args.shard_rank < args.shard_world: + raise ValueError( + f"shard_rank must satisfy 0 <= rank < world; got {args.shard_rank}/{args.shard_world}" + ) + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", + ) + device = args.device + if device is None: + device = f"cuda:{args.gpu_id}" if args.gpu_id is not None else "cuda" + + if args.webdataset_root is not None: + samples = iter_spatialedit_samples( + args.webdataset_root, + negative_prompt=args.negative_prompt, + prompt_variant=args.prompt_variant, + shard_rank=args.shard_rank, + shard_world=args.shard_world, + ) + else: + samples = iter_jsonl_samples( + args.manifest, + negative_prompt=args.negative_prompt, + shard_rank=args.shard_rank, + shard_world=args.shard_world, + ) + + return preprocess_samples( + samples, + output_dir=args.output_dir, + model_name=args.model_name, + device=device, + max_pixels=args.max_pixels, + conditioning_max_pixels=args.conditioning_max_pixels, + metadata_shard_size=args.metadata_shard_size, + shard_rank=args.shard_rank, + shard_world=args.shard_world, + verify=args.verify, + overwrite=args.overwrite, + fail_fast=args.fail_fast, + limit=args.limit, + seed=args.seed, + log_every=args.log_every, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index 5e5e79f0d12..6c543061e42 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -9,5 +9,9 @@ # fastgen_data/__init__.py adds a runtime guard with an actionable message if the helpers move. nemo_automodel[diffusion]>=0.4.0,<1.0 +# Qwen-Image-Edit-2511 needs EditPlus plus the transformer's ``zero_cond_t`` handling, +# which landed after the stable 0.36 release (the checkpoint reports 0.36.0.dev0). +diffusers>=0.37.0 + # Optional but recommended for the smoke logs. wandb diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index 27b1ec22436..bebc61970a3 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -21,6 +21,7 @@ from models_utils import MODEL_DEFAULTS, ModelType from pipeline_manager import PipelineManager from quantize_config import CalibrationConfig +from qwen_image_dmd2_sampler import dmd2_sample from tqdm import tqdm from utils import load_calib_prompts @@ -95,6 +96,9 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: elif self.model_type in [ModelType.WAN22_T2V_14b, ModelType.WAN22_T2V_5b]: # Special handling for WAN video models self._run_wan_video_calibration(prompt_batch, extra_args) + elif self.model_type == ModelType.QWEN_IMAGE_DMD2: + # DMD2 students use a custom few-step sampler, not the standard loop. + self._run_qwen_image_dmd2_calibration(prompt_batch) else: common_args = { "prompt": prompt_batch, @@ -105,6 +109,22 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: self.logger.debug(f"Completed calibration batch {i + 1}/{self.config.num_batches}") self.logger.info("Calibration completed successfully") + def _run_qwen_image_dmd2_calibration(self, prompt_batch: list[str]) -> None: + """Calibrate a DMD2 Qwen-Image student via its few-step sampler. + + Drives the same few-step DMD unroll the student was trained/served with + (NOT the standard denoising loop) so the collected activation statistics + are representative of inference. The VAE decode is skipped — calibration + only needs the transformer forwards. + """ + cfg = self.pipeline_manager.dmd_sampler_cfg + if cfg is None: + raise RuntimeError( + "DMD2 sampler config is not set; the qwen-image-dmd2 pipeline must be created " + "via PipelineManager.create_pipeline() before calibration." + ) + dmd2_sample(self.pipe, prompt_batch, decode=False, **cfg) + def _run_wan_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] ) -> None: diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 4d1bd803305..3dc41138797 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -64,6 +64,10 @@ class ModelType(str, Enum): WAN22_T2V_14b = "wan2.2-t2v-14b" WAN22_T2V_5b = "wan2.2-t2v-5b" QWEN_IMAGE = "qwen-image" + # DMD2-distilled few-step Qwen-Image student (from examples/diffusers/fastgen). + # Same architecture as QWEN_IMAGE, but loaded from a consolidated student dir + # and calibrated with the few-step DMD sampler instead of the standard loop. + QWEN_IMAGE_DMD2 = "qwen-image-dmd2" _FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { @@ -74,6 +78,7 @@ class ModelType(str, Enum): ModelType.WAN22_T2V_14b: filter_func_wan_video, ModelType.WAN22_T2V_5b: filter_func_wan_video, ModelType.QWEN_IMAGE: filter_func_qwen_image, + ModelType.QWEN_IMAGE_DMD2: filter_func_qwen_image, } _VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { @@ -107,6 +112,11 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: "Wan-AI/Wan2.2-T2V-A14B-Diffusers", ModelType.WAN22_T2V_5b: "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ModelType.QWEN_IMAGE: "Qwen/Qwen-Image", + # Base pipeline (VAE / text-encoder / tokenizer / scheduler) for DMD2 students; + # the trained transformer is loaded separately from a consolidated dir via the + # ``student_path`` extra-param. Override with ``--override-model-path`` or + # ``--extra-param base_pipeline_path=...``. + ModelType.QWEN_IMAGE_DMD2: "Qwen/Qwen-Image", } MODEL_PIPELINE: dict[ModelType, type[DiffusionPipeline] | None] = { @@ -122,6 +132,7 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: WanPipeline, ModelType.WAN22_T2V_5b: WanPipeline, ModelType.QWEN_IMAGE: QwenImagePipeline, + ModelType.QWEN_IMAGE_DMD2: QwenImagePipeline, } # Shared dataset configurations @@ -273,6 +284,14 @@ def get_model_filter_func( }, } +# DMD2 students share Qwen-Image's architecture, so they reuse the same block-range +# recipe, high-precision filter, base pipeline, and calibration dataset. They differ +# only in (a) loading -- a consolidated student dir swapped into the base pipeline +# (PipelineManager._create_qwen_image_dmd2_pipeline) -- and (b) calibration, which +# drives the few-step DMD sampler instead of the standard denoising loop +# (Calibrator._run_qwen_image_dmd2_calibration). Inherit so the recipe stays in sync. +MODEL_DEFAULTS[ModelType.QWEN_IMAGE_DMD2] = {**MODEL_DEFAULTS[ModelType.QWEN_IMAGE]} + def _coerce_extra_param_value(value: str) -> Any: lowered = value.lower() diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index af89ed568ff..4e76ee0d661 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -43,6 +43,9 @@ def __init__(self, config: ModelConfig, logger: logging.Logger): self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling self._transformer: torch.nn.Module | None = None self._video_decoder: torch.nn.Module | None = None + # Few-step sampler config for DMD2 students (populated when loading a + # qwen-image-dmd2 pipeline); consumed by the calibrator / sanity check. + self.dmd_sampler_cfg: dict[str, Any] | None = None @staticmethod def create_pipeline_from( @@ -100,6 +103,11 @@ def create_pipeline(self) -> Any: self.logger.info("LTX-2 pipeline created successfully") return self.pipe + if self.config.model_type == ModelType.QWEN_IMAGE_DMD2: + self.pipe = self._create_qwen_image_dmd2_pipeline() + self.logger.info("Qwen-Image DMD2 pipeline created successfully") + return self.pipe + pipeline_cls = MODEL_PIPELINE[self.config.model_type] if pipeline_cls is None: raise ValueError( @@ -266,6 +274,113 @@ def _create_ltx2_pipeline(self) -> Any: pipeline_kwargs.update(params) return TI2VidTwoStagesPipeline(**pipeline_kwargs) + def _create_qwen_image_dmd2_pipeline(self) -> Any: + """Build a QwenImagePipeline whose transformer is a DMD2-trained student. + + Loads the consolidated student transformer (the ``model/consolidated`` dir + produced by ``examples/diffusers/fastgen`` training), optionally overlays an + EMA shadow, and swaps it into the base Qwen-Image pipeline so the VAE / + text-encoder / tokenizer / scheduler come from the base checkpoint. + + Reads from ``extra_params``: + student_path (required): consolidated student dir. + base_pipeline_path: base Qwen-Image dir/HF id (defaults to the + registry id or ``--override-model-path``). + ema_path: optional ``ema_shadow.pt`` to overlay onto the student. + sample_steps / t_list / sample_type / guidance_scale / max_t: + few-step sampler schedule (defaults match the canonical 4-step + shift=3 student); stashed in ``self.dmd_sampler_cfg``. + """ + from qwen_image_dmd2_sampler import DEFAULT_MAX_T, resolve_schedule + + try: + from diffusers import QwenImagePipeline, QwenImageTransformer2DModel + except ImportError as e: + raise ImportError( + "qwen-image-dmd2 requires a diffusers version providing QwenImagePipeline " + "and QwenImageTransformer2DModel; upgrade diffusers." + ) from e + + params = dict(self.config.extra_params) + student_path = params.get("student_path") + if not student_path: + raise ValueError( + "Missing required extra_param: student_path (the consolidated DMD2 student " + "dir, e.g. .../epoch_4_step_17999/model/consolidated)." + ) + base_pipeline_path = params.get("base_pipeline_path") or self.config.model_path + ema_path = params.get("ema_path") + + default_dtype = self.config.model_dtype["default"] + transformer_dtype = self.config.model_dtype.get("transformer", default_dtype) + if torch.float16 in (default_dtype, transformer_dtype): + self.logger.warning( + "Qwen-Image is trained/served in bfloat16; float16 (Half) can overflow the " + "VAE and produce NaNs. Consider --model-dtype BFloat16." + ) + + self.logger.info("Loading DMD2 student transformer from %s", student_path) + transformer = QwenImageTransformer2DModel.from_pretrained( + student_path, torch_dtype=transformer_dtype + ) + + if ema_path: + self.logger.info("Overlaying EMA shadow from %s", ema_path) + ema_state = torch.load(str(ema_path), map_location="cpu") + shadow = ( + ema_state.get("shadow", ema_state) if isinstance(ema_state, dict) else ema_state + ) + if not isinstance(shadow, dict): + raise ValueError( + f"ema_path content has unexpected type {type(shadow).__name__}; " + "expected dict[str, Tensor]." + ) + missing, unexpected = transformer.load_state_dict(shadow, strict=False) + if unexpected: + self.logger.warning("EMA overlay had %d unexpected key(s)", len(unexpected)) + if missing: + self.logger.warning("EMA overlay missed %d student key(s)", len(missing)) + + transformer.eval() + + self.logger.info( + "Loading base Qwen-Image pipeline from %s (transformer replaced by student)", + base_pipeline_path, + ) + pipe = QwenImagePipeline.from_pretrained( + base_pipeline_path, transformer=transformer, torch_dtype=default_dtype + ) + pipe.set_progress_bar_config(disable=True) + + # Resolve and stash the few-step sampler config. Defaults match the + # canonical 4-step shift=3 student; the schedule MUST match training. + sample_steps = params.get("sample_steps") + sample_steps = int(sample_steps) if sample_steps is not None else 4 + t_list = params.get("t_list") + if isinstance(t_list, str): + t_list = [float(x) for x in t_list.split(",") if x.strip()] + max_t = float(params.get("max_t", DEFAULT_MAX_T)) + schedule = resolve_schedule(t_list, sample_steps, max_t) + defaults = MODEL_DEFAULTS[self.config.model_type].get("inference_extra_args", {}) + self.dmd_sampler_cfg = { + "schedule": schedule, + "sample_type": str(params.get("sample_type", "ode")), + "guidance_scale": float(params.get("guidance_scale", 1.0)), + "negative_prompt": params.get("negative_prompt"), + "height": int(params.get("height", defaults.get("height", 1024))), + "width": int(params.get("width", defaults.get("width", 1024))), + "max_sequence_length": int(params.get("max_sequence_length", 512)), + } + self.logger.info( + "DMD2 few-step sampler: steps=%d schedule=%s sample_type=%s guidance_scale=%s " + "(schedule must match the student's training t_list)", + len(schedule) - 1, + schedule, + self.dmd_sampler_cfg["sample_type"], + self.dmd_sampler_cfg["guidance_scale"], + ) + return pipe + def print_quant_summary(self): for name, backbone in self.iter_backbones(): self.logger.info(f"{name} quantization info:") diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 1d71c088652..26aac8ae07f 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -50,9 +50,8 @@ QuantFormat, QuantizationConfig, ) -from utils import check_conv_and_mha, check_lora +from utils import check_conv_and_mha, check_lora, restore_quantizer_state, save_quantizer_state -import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint @@ -309,8 +308,11 @@ def save_checkpoint( filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" target_path = ckpt_path / filename - self.logger.info(f"Saving backbone to {target_path}") - mto.save(backbone, str(target_path)) + # Save ONLY the quantization state (recipe + quantizer buffers incl. amax), + # not the model weights. The weights live in the base HF/diffusers checkpoint + # and are reloaded there on restore; this keeps the artifact tiny. + self.logger.info(f"Saving quantizer state (amax + recipe, no weights) to {target_path}") + save_quantizer_state(backbone, str(target_path)) self.logger.info("Checkpoint saved successfully") @@ -380,7 +382,9 @@ def restore_checkpoint(self) -> None: f"Checkpoint not found for '{backbone_name}' in {restore_path}" ) self.logger.info(f"Restoring {backbone_name} from {source_path}") - mto.restore(backbone, str(source_path)) + # The pipeline was just created with the base (unquantized) weights, so + # this re-applies the quantization recipe + amax on top of them. + restore_quantizer_state(backbone, str(source_path)) self.logger.info("Checkpoints restored successfully") diff --git a/examples/diffusers/quantization/qwen_image_dmd2_sampler.py b/examples/diffusers/quantization/qwen_image_dmd2_sampler.py new file mode 100644 index 00000000000..518bc6b47d5 --- /dev/null +++ b/examples/diffusers/quantization/qwen_image_dmd2_sampler.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compact DMD2 few-step sampler for Qwen-Image students. + +This is a vendored, calibration-friendly version of the few-step unroll in +``examples/diffusers/fastgen/inference_dmd2_qwen_image.py``. It is kept here so +the quantization example is self-contained (no cross-example ``sys.path`` +imports) and so calibration can run the **same forward logic the student was +trained/served with** — which is what makes the collected ``amax`` statistics +representative. + +The single :func:`dmd2_sample` entry point serves two callers: + +* **Calibration** (``decode=False``): runs only the transformer forwards of the + DMD unroll and returns ``None``. The VAE / image post-processing is skipped + because quantization only needs the transformer's activation statistics, and + skipping the VAE saves substantial time and memory on the 60-layer student. +* **Sanity inference** (``decode=True``): additionally runs the VAE decode and + returns a list of images, used to confirm a restored (quantized) student + still produces a finite image. + +The math is bit-aligned with the training-time ``_build_student_input`` in +``modelopt/torch/fastgen/methods/dmd.py`` and with the inference reference: + + for (t_cur, t_next) in pairwise(t_list): + v = student(x, t=t_cur, text_emb) # flow at t_cur + x_0 = x - t_cur * v # RF identity -> x_0 estimate + if t_next > 0: + eps = (x - (1 - t_cur) * x_0) / t_cur # ODE: invert RF forward + x = (1 - t_next) * x_0 + t_next * eps # re-noise to t_next + else: + x = x_0 # final step + +``t_list`` MUST match the student's training schedule (e.g. the LightX2V +"shift=3" 4-step shape ``[1.0, 0.9, 0.75, 0.5, 0.0]``); a mismatch produces a +train/inference gap and therefore misleading calibration statistics. +""" + +from __future__ import annotations + +import itertools + +import torch +from diffusers.utils.torch_utils import randn_tensor + +# Canonical 4-step "shift=3" student schedule (LightX2V-Qwen-Image-Lightning +# shape). t_list has student_sample_steps + 1 entries: the first N are the +# timesteps the student is evaluated at, the trailing 0.0 is the terminal the +# final Euler step lands on (NOT an extra evaluation). +DEFAULT_T_LIST: tuple[float, ...] = (1.0, 0.9, 0.75, 0.5, 0.0) +DEFAULT_MAX_T: float = 0.999 + + +def resolve_schedule( + t_list: list[float] | tuple[float, ...] | None, + sample_steps: int | None, + max_t: float = DEFAULT_MAX_T, +) -> list[float]: + """Resolve the sampling schedule (timesteps + terminal 0.0). + + Priority: + 1. An explicit ``t_list`` (must end at 0.0 and have ``sample_steps + 1`` + entries when ``sample_steps`` is given). + 2. ``sample_steps == 1`` -> ``[max_t, 0.0]`` (canonical single-step). + 3. ``sample_steps == 4`` (or None) with no ``t_list`` -> ``DEFAULT_T_LIST``. + 4. Otherwise a linear ``linspace(max_t, 0, sample_steps + 1)`` fallback. + """ + if t_list is not None: + schedule = [float(t) for t in t_list] + if abs(schedule[-1]) > 1e-6: + raise ValueError( + f"t_list must end at 0.0 (got {schedule[-1]}); the final step lands on x_0." + ) + if sample_steps is not None and len(schedule) != sample_steps + 1: + raise ValueError( + f"t_list must have sample_steps+1 entries " + f"(got {len(schedule)} for sample_steps={sample_steps})." + ) + return schedule + + if sample_steps == 1: + return [float(max_t), 0.0] + if sample_steps in (None, 4): + return list(DEFAULT_T_LIST) + return torch.linspace(float(max_t), 0.0, sample_steps + 1).tolist() + + +@torch.no_grad() +def dmd2_sample( + pipe, + prompt: str | list[str], + *, + schedule: list[float], + sample_type: str = "ode", + guidance_scale: float = 1.0, + negative_prompt: str | list[str] | None = None, + height: int = 1024, + width: int = 1024, + num_images_per_prompt: int = 1, + generator: torch.Generator | None = None, + max_sequence_length: int = 512, + decode: bool = False, + output_type: str = "pil", +) -> list | None: + """Run the DMD few-step unroll on ``pipe.transformer``. + + Args: + pipe: A ``QwenImagePipeline`` whose ``transformer`` is the DMD2 student. + prompt: A prompt or list of prompts (one calibration batch). + schedule: Full timestep schedule incl. trailing 0.0 (see + :func:`resolve_schedule`). + sample_type: ``"ode"`` (deterministic, recover eps via RF identity) or + ``"sde"`` (fresh Gaussian noise between steps). Must match training. + guidance_scale: Inference-time CFG. Leave at ``1.0`` for students trained + with an internalised (non-null) ``dmd2.guidance_scale`` — passing + ``> 1.0`` there would double-apply CFG. + negative_prompt: Negative prompt for CFG; defaults to ``""`` when CFG is + engaged and none is given. + height/width: Output spatial size (must be VAE-compatible). + num_images_per_prompt: Images per prompt. + generator: Optional RNG for reproducible noise. + max_sequence_length: Text-encoder max sequence length. + decode: If ``True`` run VAE decode + post-process and return images. If + ``False`` (calibration) skip the VAE and return ``None``. + output_type: Passed to the image processor when ``decode=True``. + + Returns: + A list of images when ``decode=True``, else ``None``. + """ + if sample_type not in ("ode", "sde"): + raise ValueError(f"sample_type must be 'ode' or 'sde', got {sample_type!r}") + + do_cfg = guidance_scale != 1.0 + if do_cfg and negative_prompt is None: + negative_prompt = "" + + device = pipe.transformer.device + dtype = next(pipe.transformer.parameters()).dtype + + # ---- Encode prompt(s) ------------------------------------------------ + prompt_embeds, prompt_embeds_mask = pipe.encode_prompt( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + neg_prompt_embeds = neg_prompt_embeds_mask = None + if do_cfg: + neg_prompt_embeds, neg_prompt_embeds_mask = pipe.encode_prompt( + prompt=negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + txt_seq_lens = ( + prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None + ) + neg_txt_seq_lens = ( + neg_prompt_embeds_mask.sum(dim=1).int().tolist() + if neg_prompt_embeds_mask is not None + else None + ) + + # ---- Build initial noisy latents at t = schedule[0] ------------------ + batch_size = (1 if isinstance(prompt, str) else len(prompt)) * num_images_per_prompt + num_channels_latents = pipe.transformer.config.in_channels // 4 # 64 // 4 = 16 + h_lat = 2 * (height // (pipe.vae_scale_factor * 2)) + w_lat = 2 * (width // (pipe.vae_scale_factor * 2)) + latent_shape = (batch_size, 1, num_channels_latents, h_lat, w_lat) + + noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) + latents_5d = noise * schedule[0] # RF: sigma(t0) = t0 + x_packed = pipe._pack_latents(latents_5d, batch_size, num_channels_latents, h_lat, w_lat) + img_shapes = [[(1, h_lat // 2, w_lat // 2)]] * batch_size + + # ---- DMD few-step unroll (transformer forwards) ---------------------- + for t_cur, t_next in itertools.pairwise(schedule): + timestep = torch.tensor([t_cur], device=device, dtype=dtype).expand(batch_size) + flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + if do_cfg: + neg_flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=neg_prompt_embeds, + encoder_hidden_states_mask=neg_prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=neg_txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + flow_packed = ( + neg_flow_packed.to(torch.float64) + + float(guidance_scale) + * (flow_packed.to(torch.float64) - neg_flow_packed.to(torch.float64)) + ).to(dtype) + + # RF identity: x_0 = x_t - t_cur * v (fp64 for stability). + x0_packed = (x_packed.to(torch.float64) - float(t_cur) * flow_packed.to(torch.float64)).to( + dtype + ) + + if t_next > 1e-6: + if sample_type == "ode": + alpha_cur = 1.0 - float(t_cur) + eps_packed = ( + (x_packed.to(torch.float64) - alpha_cur * x0_packed.to(torch.float64)) + / max(float(t_cur), 1e-6) + ).to(dtype) + else: + eps_packed = torch.randn( + x_packed.shape, generator=generator, device=device, dtype=dtype + ) + alpha_next = 1.0 - float(t_next) + x_packed = ( + alpha_next * x0_packed.to(torch.float64) + + float(t_next) * eps_packed.to(torch.float64) + ).to(dtype) + else: + x_packed = x0_packed + + if not decode: + # Calibration path: transformer forwards already ran; nothing to decode. + return None + + # ---- VAE decode (sanity-inference path only) ------------------------- + x0_5d = pipe._unpack_latents(x_packed, height, width, pipe.vae_scale_factor) + latents_mean = ( + torch.tensor(pipe.vae.config.latents_mean) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device=device, dtype=dtype) + ) + latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view( + 1, pipe.vae.config.z_dim, 1, 1, 1 + ).to(device=device, dtype=dtype) + x0_scaled = x0_5d / latents_std + latents_mean + image_5d = pipe.vae.decode(x0_scaled, return_dict=False)[0] + image_4d = image_5d[:, :, 0] # Qwen-Image treats images as 1-frame videos + return pipe.image_processor.postprocess(image_4d, output_type=output_type) diff --git a/examples/diffusers/quantization/sanity_check_dmd2.py b/examples/diffusers/quantization/sanity_check_dmd2.py new file mode 100644 index 00000000000..9f78901f20c --- /dev/null +++ b/examples/diffusers/quantization/sanity_check_dmd2.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Restore a quantized DMD2 Qwen-Image student and run one few-step inference. + +Confirms the round trip of the new ``qwen-image-dmd2`` quantization flow: + + 1. Load the base Qwen-Image pipeline with the consolidated student swapped in + (via the same :class:`PipelineManager` path quantize.py uses) -- this brings + the original (unquantized) weights. + 2. Reapply the weight-free quantization checkpoint saved by ``quantize.py`` + (``save_quantizer_state`` -> ``transformer.pt``) via + ``restore_quantizer_state``, which re-applies the quantizer recipe **and the + calibrated amax** buffers on top of the loaded weights. + 3. Run a single few-step DMD inference (with VAE decode) and assert the image + is finite and non-constant. + +This deliberately reuses :class:`PipelineManager` and +:func:`qwen_image_dmd2_sampler.dmd2_sample` so the inference path is identical to +calibration's (minus the VAE decode, which is enabled here). + +Usage:: + + python sanity_check_dmd2.py \\ + --quantized-ckpt ./qwen_dmd2_fp8/transformer.pt \\ + --student-path /.../epoch_4_step_17999/model/consolidated \\ + --base-pipeline-path /.../models/Qwen-Image \\ + --output-png ./qwen_dmd2_fp8/sanity.png +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + +import torch +from models_utils import ModelType +from pipeline_manager import PipelineManager +from quantize_config import ModelConfig +from qwen_image_dmd2_sampler import dmd2_sample +from utils import restore_quantizer_state + +import modelopt.torch.quantization as mtq + +logger = logging.getLogger("sanity_check_dmd2") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--quantized-ckpt", + required=True, + help="Path to the quantized checkpoint saved by quantize.py (e.g. .../transformer.pt).", + ) + parser.add_argument( + "--student-path", + required=True, + help="Consolidated DMD2 student dir (provides architecture + base weights to restore into).", + ) + parser.add_argument( + "--base-pipeline-path", + default="Qwen/Qwen-Image", + help="Base Qwen-Image dir/HF id for the VAE / text-encoder / tokenizer / scheduler.", + ) + parser.add_argument("--ema-path", default=None, help="Optional EMA shadow overlaid on load.") + parser.add_argument("--output-png", default="./qwen_dmd2_sanity.png") + parser.add_argument("--prompt", default="a small red cube on a white table") + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--seed", type=int, default=42) + # Few-step sampler knobs (defaults match the canonical 4-step shift=3 student). + parser.add_argument("--sample-steps", type=int, default=4) + parser.add_argument( + "--t-list", + default=None, + help="Comma-separated schedule incl. trailing 0.0, e.g. '1.0,0.9,0.75,0.5,0.0'.", + ) + parser.add_argument("--sample-type", default="ode", choices=["ode", "sde"]) + parser.add_argument("--guidance-scale", type=float, default=1.0) + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s" + ) + + # 1. Build the base pipeline with the student swapped in (unquantized). + extra_params: dict[str, object] = { + "student_path": args.student_path, + "base_pipeline_path": args.base_pipeline_path, + "sample_steps": args.sample_steps, + "sample_type": args.sample_type, + "guidance_scale": args.guidance_scale, + "height": args.height, + "width": args.width, + } + if args.ema_path: + extra_params["ema_path"] = args.ema_path + if args.t_list: + extra_params["t_list"] = args.t_list + + model_config = ModelConfig( + model_type=ModelType.QWEN_IMAGE_DMD2, + model_dtype={"default": torch.bfloat16}, + backbone=["transformer"], + extra_params=extra_params, + ) + pm = PipelineManager(model_config, logger) + pipe = pm.create_pipeline() + + # 2. Restore the quantized architecture + calibrated amax into the student. + logger.info( + "Restoring quantizer state (amax + recipe) from %s onto the loaded student", + args.quantized_ckpt, + ) + restore_quantizer_state(pipe.transformer, args.quantized_ckpt) + mtq.print_quant_summary(pipe.transformer) + pm.setup_device() + + # 3. One few-step inference (with VAE decode). + gen = torch.Generator(device=pipe.transformer.device).manual_seed(args.seed) + images = dmd2_sample(pipe, [args.prompt], decode=True, generator=gen, **pm.dmd_sampler_cfg) + image = images[0] + + import numpy as np + + arr = np.asarray(image) + stats = { + "prompt": args.prompt, + "quantized_ckpt": args.quantized_ckpt, + "schedule": pm.dmd_sampler_cfg["schedule"], + "image_shape": list(arr.shape), + "image_dtype": str(arr.dtype), + "image_min": float(arr.min()), + "image_max": float(arr.max()), + "image_mean": float(arr.mean()), + "image_std": float(arr.std()), + "is_finite": bool(np.isfinite(arr).all()), + "is_not_constant": bool(arr.std() > 0), + } + + os.makedirs(os.path.dirname(os.path.abspath(args.output_png)), exist_ok=True) + image.save(args.output_png) + with open(args.output_png.replace(".png", "_stats.json"), "w") as f: + json.dump(stats, f, indent=2) + print(json.dumps(stats, indent=2)) + + if not stats["is_finite"]: + logger.error("Sanity check FAILED: image contains non-finite values.") + sys.exit(1) + if not stats["is_not_constant"]: + logger.error("Sanity check FAILED: image is constant (std == 0).") + sys.exit(1) + logger.info( + "Sanity check PASSED: restored quantized student produced a finite image -> %s", + args.output_png, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..9a6b841a52e 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -24,8 +24,13 @@ from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear from diffusers.utils import load_image +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.plugins.diffusion.diffusers import AttentionModuleMixin +from modelopt.torch.quantization.utils.core_utils import ( + get_quantizer_state_dict, + set_quantizer_state_dict, +) USE_PEFT = True try: @@ -193,3 +198,36 @@ def mha_filter_func(name): if hasattr(F, "scaled_dot_product_attention"): mtq.disable_quantizer(backbone, mha_filter_func) + + +def save_quantizer_state(model: torch.nn.Module, path: str) -> None: + """Save ONLY ModelOpt's quantization state -- the recipe plus the quantizer + buffers (amax, pre_quant_scale, ...) -- and NOT the model weights. + + This is the same idiom ModelOpt uses internally (see + ``modelopt.torch.quantization.plugins.transformers_trainer``): the + ``modelopt_state`` (architecture/recipe from :func:`mto.modelopt_state`) is + bundled with the per-quantizer state from + :func:`get_quantizer_state_dict` under the ``modelopt_state_weights`` key. + The resulting checkpoint is tiny (KBs-MBs) and is reloaded on top of the + original (unquantized) model via :func:`restore_quantizer_state`. + """ + modelopt_state = mto.modelopt_state(model) + modelopt_state["modelopt_state_weights"] = get_quantizer_state_dict(model) + torch.save(modelopt_state, str(path)) + + +def restore_quantizer_state(model: torch.nn.Module, path: str) -> torch.nn.Module: + """Reload a checkpoint written by :func:`save_quantizer_state` onto ``model``. + + ``model`` must already hold its original (unquantized) weights (e.g. freshly + loaded from the base HF/diffusers checkpoint); this re-applies the + quantization recipe and loads the calibrated amax/quantizer buffers on top. + Mirrors ModelOpt's ``_restore_modelopt_state_with_weights``. + """ + modelopt_state = mto.load_modelopt_state(str(path)) + quantizer_state = modelopt_state.pop("modelopt_state_weights", None) + mto.restore_from_modelopt_state(model, modelopt_state) + if quantizer_state is not None: + set_quantizer_state_dict(model, quantizer_state) + return model diff --git a/modelopt/torch/fastgen/plugins/__init__.py b/modelopt/torch/fastgen/plugins/__init__.py index 8810470b26f..c33f29ec9d1 100644 --- a/modelopt/torch/fastgen/plugins/__init__.py +++ b/modelopt/torch/fastgen/plugins/__init__.py @@ -25,3 +25,4 @@ with import_plugin("qwen_image"): from .qwen_image import * + from .qwen_image_edit import * diff --git a/modelopt/torch/fastgen/plugins/qwen_image.py b/modelopt/torch/fastgen/plugins/qwen_image.py index 08a32b09301..0a4d3f8fc64 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image.py +++ b/modelopt/torch/fastgen/plugins/qwen_image.py @@ -48,9 +48,11 @@ from __future__ import annotations import contextlib +from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING, Any import torch +from packaging.version import Version from torch import nn from ..methods.dmd import DMDPipeline @@ -58,6 +60,16 @@ if TYPE_CHECKING: from ..config import DMDConfig + +try: + # Diffusers 0.35/0.36 requires explicit Python sequence lengths in Qwen's + # positional-embedding path. Starting in 0.37 the mask is authoritative and + # passing txt_seq_lens is deprecated. Keep the shared T2I plugin compatible + # with ModelOpt's broader diffusers extra while Edit-2511 pins the newer API. + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = Version(version("diffusers")) < Version("0.37.0") +except PackageNotFoundError: # pragma: no cover - optional plugin import guard + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = False + __all__ = [ "QwenImageDMDPipeline", "attach_feature_capture", @@ -65,6 +77,7 @@ "pack_latents", "remove_feature_capture", "unpack_latents", + "update_feature_capture_shape", ] @@ -206,6 +219,7 @@ def _call_model( packed = pack_latents(hidden_states) img_shapes = build_img_shapes(b, h, w) + update_feature_capture_shape(model, h, w) call_kwargs: dict[str, Any] = dict(model_kwargs) call_kwargs.pop("hidden_states", None) @@ -214,8 +228,10 @@ def _call_model( call_kwargs.pop("guidance", None) call_kwargs.pop("return_dict", None) txt_seq_lens = call_kwargs.pop("txt_seq_lens", None) - if txt_seq_lens is None and encoder_hidden_states_mask is not None: - txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + if _DIFFUSERS_NEEDS_TXT_SEQ_LENS: + if txt_seq_lens is None and encoder_hidden_states_mask is not None: + txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + call_kwargs["txt_seq_lens"] = txt_seq_lens guidance = None if self._guidance_value is not None: @@ -232,7 +248,6 @@ def _call_model( encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=encoder_hidden_states_mask, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, guidance=guidance, return_dict=False, **call_kwargs, @@ -266,6 +281,16 @@ def _call_model( _SHAPE_ATTR = "_fastgen_capture_shape" +def update_feature_capture_shape(model: nn.Module, h_lat: int, w_lat: int) -> None: + """Refresh a hooked teacher's target shape for the current multiresolution batch.""" + if h_lat % 2 or w_lat % 2: + raise ValueError( + f"feature capture requires even latent dims, got h_lat={h_lat}, w_lat={w_lat}." + ) + if hasattr(model, _HANDLES_ATTR): + setattr(model, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) + + def attach_feature_capture( teacher: nn.Module, feature_indices: list[int], @@ -273,6 +298,7 @@ def attach_feature_capture( w_lat: int, *, blocks_attr: str = "transformer_blocks", + target_prefix_only: bool = False, ) -> None: """Install forward hooks on ``teacher.transformer_blocks[i]`` for each ``i`` in ``feature_indices``. @@ -297,6 +323,11 @@ def attach_feature_capture( blocks_attr: Attribute under which the teacher exposes its block stack. Default ``"transformer_blocks"`` matches diffusers' ``QwenImageTransformer2DModel``. + target_prefix_only: When ``True``, allow the captured image-token sequence to + contain extra tokens after the target image and retain only the leading + ``(h_lat // 2) * (w_lat // 2)`` target tokens. Qwen-Image-Edit concatenates + packed reference-image tokens after the noisy target tokens. The default is + ``False`` so the text-to-image path keeps its strict sequence-length check. Raises: AttributeError: ``teacher`` does not expose ``blocks_attr``. @@ -336,8 +367,6 @@ def attach_feature_capture( setattr(teacher, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) handles: list[Any] = [] - h_half = h_lat // 2 - w_half = w_lat // 2 for idx in sorted_indices: block = blocks[idx] @@ -354,13 +383,20 @@ def _hook(_module: nn.Module, _inputs: Any, output: Any) -> None: ) # hidden: [B, num_image_patches, C] -> [B, C, H_half, W_half]. b, s, c = hidden.shape + h_half, w_half = getattr(teacher, _SHAPE_ATTR) expected_s = h_half * w_half - if s != expected_s: + if s < expected_s or (not target_prefix_only and s != expected_s): + expected_description = ( + f"at least {expected_s}" if target_prefix_only else str(expected_s) + ) raise RuntimeError( f"QwenImage feature-capture got hidden_states seq_len={s} but expected " - f"{expected_s} = (h_lat // 2) * (w_lat // 2). Did the input resolution " - f"drift from the attach_feature_capture-time setting?" + f"{expected_description} target tokens, where {expected_s} = " + "(h_lat // 2) * (w_lat // 2). Did the input resolution drift from " + "the attach_feature_capture-time setting?" ) + if target_prefix_only: + hidden = hidden[:, :expected_s] feat = hidden.permute(0, 2, 1).reshape(b, c, h_half, w_half) captured.append(feat) diff --git a/modelopt/torch/fastgen/plugins/qwen_image_edit.py b/modelopt/torch/fastgen/plugins/qwen_image_edit.py new file mode 100644 index 00000000000..4d56801c738 --- /dev/null +++ b/modelopt/torch/fastgen/plugins/qwen_image_edit.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image-Edit plumbing for DMD2. + +``QwenImageEditPlusPipeline`` conditions the transformer in two complementary ways: + +* the Qwen2.5-VL prompt embedding contains the edit instruction and visual context; and +* one or more clean VAE reference-image latents are packed and appended after the noisy + target-image tokens. + +Only the target image is diffused. DMD2 therefore keeps its external latent contract as +``[B, C, H, W]`` and forwards reference latents through the model kwargs under +``conditioning_latents``. This plugin packs ``[target, reference_1, ...]`` for every model +call, constructs the matching ``img_shapes``, and discards the reference-token suffix from +the model prediction before returning to the shared DMD math. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from .qwen_image import ( + QwenImageDMDPipeline, + pack_latents, + unpack_latents, + update_feature_capture_shape, +) +from .qwen_image import attach_feature_capture as _attach_qwen_feature_capture +from .qwen_image import remove_feature_capture as _remove_qwen_feature_capture + +if TYPE_CHECKING: + from ..config import DMDConfig + +__all__ = [ + "QwenImageEditDMDPipeline", +] + + +class QwenImageEditDMDPipeline(QwenImageDMDPipeline): + """DMD2 pipeline for Qwen-Image-Edit's target-plus-reference token layout. + + ``conditioning_latents`` must be supplied to each ``compute_*_loss`` call as a + non-empty list or tuple. Each entry is one clean reference image with shape + ``[B, C, H_ref, W_ref]``. Reference images may have different spatial shapes from + each other and from the target, but their batch/channel/device/dtype must match the + target latent. The reference order must match the order used when constructing the + multimodal prompt embedding. + """ + + def __init__( + self, + student: nn.Module, + teacher: nn.Module, + fake_score: nn.Module, + config: DMDConfig, + *, + discriminator: nn.Module | None = None, + guidance: float | None = None, + ) -> None: + """Initialize the shared Qwen pipeline and retain its timestep/guidance checks.""" + super().__init__( + student=student, + teacher=teacher, + fake_score=fake_score, + config=config, + discriminator=discriminator, + guidance=guidance, + ) + + @staticmethod + def _validate_conditioning_latents( + conditioning_latents: Any, + target: torch.Tensor, + ) -> list[torch.Tensor]: + """Validate and normalize the reference-latent sequence for one model call.""" + if not isinstance(conditioning_latents, (list, tuple)) or not conditioning_latents: + raise ValueError( + "QwenImageEditDMDPipeline requires non-empty `conditioning_latents` as a " + "list or tuple of [B, C, H, W] tensors." + ) + + b, c, _h, _w = target.shape + normalized: list[torch.Tensor] = [] + for index, latent in enumerate(conditioning_latents): + if not torch.is_tensor(latent): + raise TypeError( + f"conditioning_latents[{index}] must be a Tensor, got {type(latent).__name__}." + ) + if latent.ndim != 4: + raise ValueError( + f"conditioning_latents[{index}] must have shape [B, C, H, W], got " + f"{latent.ndim}D tensor with shape {tuple(latent.shape)}." + ) + if latent.shape[0] != b or latent.shape[1] != c: + raise ValueError( + f"conditioning_latents[{index}] batch/channels {tuple(latent.shape[:2])} " + f"must match target {(b, c)}." + ) + if latent.shape[2] % 2 or latent.shape[3] % 2: + raise ValueError( + f"conditioning_latents[{index}] requires even spatial dims for Qwen " + f"packing, got H={latent.shape[2]}, W={latent.shape[3]}." + ) + if latent.device != target.device: + raise ValueError( + f"conditioning_latents[{index}] is on {latent.device}, but target is on " + f"{target.device}." + ) + if latent.dtype != target.dtype: + raise ValueError( + f"conditioning_latents[{index}] has dtype {latent.dtype}, but target has " + f"dtype {target.dtype}." + ) + normalized.append(latent) + return normalized + + def _call_model( + self, + model: nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Pack target + references, call the transformer, and return target prediction only.""" + if hidden_states.ndim != 4: + raise ValueError( + "QwenImageEditDMDPipeline._call_model expects 4D hidden_states " + f"[B, C, H, W] (got {hidden_states.ndim}D)." + ) + b, _c, h, w = hidden_states.shape + + call_kwargs: dict[str, Any] = dict(model_kwargs) + conditioning_latents = self._validate_conditioning_latents( + call_kwargs.pop("conditioning_latents", None), hidden_states + ) + + target_packed = pack_latents(hidden_states) + update_feature_capture_shape(model, h, w) + conditioning_packed = [pack_latents(latent) for latent in conditioning_latents] + packed = torch.cat([target_packed, *conditioning_packed], dim=1) + target_num_patches = target_packed.shape[1] + + per_sample_shapes = [(1, h // 2, w // 2)] + [ + (1, latent.shape[2] // 2, latent.shape[3] // 2) for latent in conditioning_latents + ] + img_shapes = [list(per_sample_shapes) for _ in range(b)] + + # These values are owned by this wrapper. Drop caller copies so duplicate kwargs + # cannot leak through to the diffusers transformer. + call_kwargs.pop("hidden_states", None) + encoder_hidden_states_mask = call_kwargs.pop("encoder_hidden_states_mask", None) + call_kwargs.pop("img_shapes", None) + call_kwargs.pop("guidance", None) + call_kwargs.pop("return_dict", None) + # Stable Diffusers derives text lengths from encoder_hidden_states_mask. + call_kwargs.pop("txt_seq_lens", None) + + guidance = None + if self._guidance_value is not None: + guidance = torch.full( + (b,), + float(self._guidance_value), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + out = model( + hidden_states=packed, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + img_shapes=img_shapes, + guidance=guidance, + return_dict=False, + **call_kwargs, + ) + + if isinstance(out, tuple): + raw_packed = out[0] + elif isinstance(out, torch.Tensor): + raw_packed = out + elif hasattr(out, "sample"): + raw_packed = out.sample + else: + raise TypeError( + "QwenImageEditDMDPipeline._call_model could not extract a tensor from " + f"output of type {type(out).__name__!r}." + ) + + if raw_packed.ndim != 3: + raise ValueError( + "QwenImageEditDMDPipeline expected packed model output [B, tokens, C*4], " + f"got shape {tuple(raw_packed.shape)}." + ) + if raw_packed.shape[0] != b: + raise ValueError( + f"Packed model output batch {raw_packed.shape[0]} does not match target batch {b}." + ) + if raw_packed.shape[1] < target_num_patches: + raise ValueError( + f"Packed model output has {raw_packed.shape[1]} tokens but the target prefix " + f"requires {target_num_patches}." + ) + + # QwenImageEditPlusPipeline treats only the leading target tokens as the denoising + # prediction. The appended reference-token outputs are conditioning-only. + target_prediction = raw_packed[:, :target_num_patches] + return unpack_latents(target_prediction, h, w) + + +def attach_feature_capture( + teacher: nn.Module, + feature_indices: list[int], + h_lat: int, + w_lat: int, + *, + blocks_attr: str = "transformer_blocks", +) -> None: + """Capture only the target-token prefix from Qwen-Image-Edit teacher blocks.""" + _attach_qwen_feature_capture( + teacher, + feature_indices, + h_lat, + w_lat, + blocks_attr=blocks_attr, + target_prefix_only=True, + ) + + +def remove_feature_capture(teacher: nn.Module) -> None: + """Remove feature hooks installed through :func:`attach_feature_capture`.""" + _remove_qwen_feature_capture(teacher) diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..fdb5e3443c7 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -142,9 +142,16 @@ def _quantized_sdpa(self, *args, **kwargs): k_quantized_scale = self.k_bmm_quantizer._get_amax(key) v_quantized_scale = self.v_bmm_quantizer._get_amax(value) - # We don't need to calibrate the output of softmax - return self.bmm2_output_quantizer( - fp8_sdpa( + # We don't need to calibrate the output of softmax. + # ``FP8SDPA`` is an export-only autograd Function: it exists solely to attach the ONNX + # ``symbolic`` (export_fp8_mha), and its forward is just + # ``original_scaled_dot_product_attention``. It implements no ``backward``, so routing + # through it at runtime makes quantized attention non-differentiable and breaks training + # (QAT) -- ``loss.backward()`` raises "must implement either the backward or vjp method". + # Use it only during ONNX export; at runtime call SDPA directly (identical forward math, + # with q/k/v already fake-quantized above) so autograd works. + if torch.onnx.is_in_onnx_export(): + attn_output = fp8_sdpa( query, key, value, @@ -157,7 +164,19 @@ def _quantized_sdpa(self, *args, **kwargs): else "Half", self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, ) - ) + else: + # Pass attn_mask/dropout_p/is_causal/scale as keywords (``scale`` is keyword-only in + # recent torch), mirroring FP8SDPA.forward's own call to SDPA. + attn_output = original_scaled_dot_product_attention( + query, + key, + value, + attn_mask=param_dict["attn_mask"], + dropout_p=param_dict["dropout_p"], + is_causal=param_dict["is_causal"], + scale=param_dict["scale"], + ) + return self.bmm2_output_quantizer(attn_output) class _QuantAttention(_QuantFunctionalMixin): diff --git a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py new file mode 100644 index 00000000000..2747012ec59 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression test for the DMD2 QAT (restore-only) quantizer-state restore. + +The fastgen QAT path NEVER calibrates during training -- it RESTORES a ModelOpt quantizer +state (recipe + frozen amax) saved by the calibration example and re-applies it on every +(re)start. This test pins the guarantee the recipe depends on, on a tiny CPU model (no +GPU, milliseconds): ``dmd2_recipe.restore_quantizer_state`` onto a *fresh* model with +DIFFERENT weights reproduces the amax bit-identically and leaves that model's weights +untouched -- i.e. amax stays exactly as calibrated and the warm-started student weights +are preserved. + +The on-disk state is built here with ModelOpt's own idiom (the same one the calibration +example's ``--quantized-torch-ckpt-save-path`` uses), so the test also pins format +compatibility with that file. + +Dependency-guarded with ``importorskip`` so it skips where torch / modelopt are absent. +""" + +from __future__ import annotations + +import pathlib +import sys + +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +torch = pytest.importorskip("torch") +mtq = pytest.importorskip("modelopt.torch.quantization") +mto = pytest.importorskip("modelopt.torch.opt") +dmd2_recipe = pytest.importorskip("dmd2_recipe") + +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import get_quantizer_state_dict + + +def _tiny_model(seed: int) -> torch.nn.Module: + torch.manual_seed(seed) + return torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.ReLU(), torch.nn.Linear(8, 4)) + + +def _amax_by_name(model: torch.nn.Module) -> dict[str, torch.Tensor]: + return { + name: module.amax.detach().clone() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.amax is not None + } + + +def test_restore_quantizer_state_is_bit_identical_and_weight_free(tmp_path): + # Calibrate a tiny model (this is the ONLY place quantize/calibration happens -- the + # calibration example; the trainer never does this) and save its quantizer state in the + # weight-free format the calibration example writes (mto.modelopt_state + amax). + model = _tiny_model(seed=0) + calib = torch.randn(16, 8) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(calib)) + src_amax = _amax_by_name(model) + assert src_amax, "expected at least one calibrated TensorQuantizer amax" + + state = mto.modelopt_state(model) + state["modelopt_state_weights"] = get_quantizer_state_dict(model) + path = tmp_path / "transformer.pt" + torch.save(state, str(path)) + + # Restore onto a FRESH model with DIFFERENT weights; amax must come back + # bit-identically and the fresh model's weights must be untouched. + fresh = _tiny_model(seed=999) + before = {n: p.detach().clone() for n, p in fresh.named_parameters()} + dmd2_recipe.restore_quantizer_state(fresh, str(path)) + + restored_amax = _amax_by_name(fresh) + assert set(restored_amax) == set(src_amax) + for name, amax in src_amax.items(): + assert torch.equal(restored_amax[name], amax), f"amax mismatch at {name}" + + for n, p in fresh.named_parameters(): + if n in before: + assert torch.equal(p.detach(), before[n]), f"restore changed weight {n}" diff --git a/tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py b/tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py new file mode 100644 index 00000000000..e4f322cc0c7 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused CPU tests for the Qwen-Image-Edit preprocessing/data contracts.""" + +from __future__ import annotations + +import io +import json +import sys +import tarfile +from pathlib import Path + +import pytest +from PIL import Image + +_REPO_ROOT = Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +import preprocess_qwen_image_edit as edit_preprocess + + +def _jpeg_bytes(color: tuple[int, int, int]) -> bytes: + buffer = io.BytesIO() + Image.new("RGB", (32, 24), color=color).save(buffer, format="JPEG") + return buffer.getvalue() + + +def _add_tar_bytes(archive: tarfile.TarFile, name: str, payload: bytes) -> None: + member = tarfile.TarInfo(name) + member.size = len(payload) + archive.addfile(member, io.BytesIO(payload)) + + +def test_spatialedit_tar_parser_maps_source_target_and_instruction(tmp_path): + shard = tmp_path / "object_rotation" / "worker0-000000.tar" + shard.parent.mkdir() + metadata = { + "conversations": [ + {"from": "human", "value": "\nRotate the object to the right."}, + {"from": "gpt", "value": "\n"}, + ], + "meta": {"sample_id": "sample-7"}, + } + with tarfile.open(shard, "w") as archive: + _add_tar_bytes(archive, "key.0.jpg", _jpeg_bytes((255, 0, 0))) + _add_tar_bytes(archive, "key.1.jpg", _jpeg_bytes((0, 0, 255))) + _add_tar_bytes(archive, "key.json", json.dumps(metadata).encode()) + + sample = next(edit_preprocess.iter_spatialedit_samples(tmp_path)) + + # The WebDataset key is pair-unique; legacy ``meta.sample_id`` is only asset-unique. + assert sample.sample_id == "key" + assert sample.prompt == "Rotate the object to the right." + assert len(sample.conditioning_images) == 1 + assert sample.conditioning_paths[0].endswith("::key.0.jpg") + assert sample.target_path.endswith("::key.1.jpg") + # JPEG is lossy, so compare dominant channels rather than exact values. + assert sample.conditioning_images[0].getpixel((0, 0))[0] > 200 + assert sample.target_image.getpixel((0, 0))[2] > 200 + + +def test_jsonl_parser_supports_archive_descriptors_and_normalized_field_names(tmp_path): + archive_path = tmp_path / "images.tar" + with tarfile.open(archive_path, "w") as archive: + _add_tar_bytes(archive, "ref1.jpg", _jpeg_bytes((255, 0, 0))) + _add_tar_bytes(archive, "ref2.jpg", _jpeg_bytes((0, 255, 0))) + _add_tar_bytes(archive, "target.jpg", _jpeg_bytes((0, 0, 255))) + manifest = tmp_path / "data.jsonl" + manifest.write_text( + json.dumps( + { + "id": "multi-ref", + "conditioning_images": [ + {"archive": "images.tar", "member": "ref1.jpg"}, + {"archive": "images.tar", "member": "ref2.jpg"}, + ], + "reference_image": {"archive": "images.tar", "member": "target.jpg"}, + "prompt": "\nCombine both references.", + } + ) + + "\n" + ) + + sample = next(edit_preprocess.iter_jsonl_samples(manifest)) + + assert sample.sample_id == "multi-ref" + assert sample.prompt == "Combine both references." + assert sample.negative_prompt == " " + assert len(sample.conditioning_images) == 2 + assert sample.target_path.endswith("images.tar::target.jpg") + + +def test_launcher_cli_aliases_map_to_canonical_arguments(): + args = edit_preprocess.build_parser().parse_args( + [ + "--input-dir", + "raw", + "--output-dir", + "cache", + "--gpu-id", + "3", + "--shard-idx", + "1", + "--shard-count", + "4", + "--max-samples", + "25", + ] + ) + + assert args.webdataset_root == Path("raw") + assert args.gpu_id == 3 + assert (args.shard_rank, args.shard_world, args.limit) == (1, 4, 25) + + +def test_edit_collate_pads_multimodal_positive_and_negative_sequences(): + torch = pytest.importorskip("torch") + pytest.importorskip("nemo_automodel") + from fastgen_data import collate_fn_image_to_image + + def sample(pos_length: int, neg_length: int, sample_id: str): + return { + "latent": torch.randn(16, 8, 8), + "conditioning_latents": [torch.randn(16, 8, 8)], + "prompt_embeds": torch.randn(pos_length, 32), + "prompt_embeds_mask": torch.ones(pos_length, dtype=torch.long), + "negative_prompt_embeds": torch.randn(neg_length, 32), + "negative_prompt_embeds_mask": torch.ones(neg_length, dtype=torch.long), + "crop_resolution": torch.tensor([64, 64]), + "original_resolution": torch.tensor([64, 64]), + "crop_offset": torch.tensor([0, 0]), + "prompt": "edit", + "negative_prompt": " ", + "image_path": f"{sample_id}.jpg", + "conditioning_image_paths": [f"{sample_id}-ref.jpg"], + "conditioning_resolutions": [(64, 64)], + "target_latent_shape": (16, 8, 8), + "conditioning_latent_shapes": [(16, 8, 8)], + "sample_id": sample_id, + "bucket_id": 0, + "aspect_ratio": 1.0, + } + + output = collate_fn_image_to_image([sample(3, 2, "a"), sample(5, 4, "b")]) + + assert output["image_latents"].shape == (2, 16, 8, 8) + assert len(output["conditioning_latents"]) == 1 + assert output["conditioning_latents"][0].shape == (2, 16, 8, 8) + assert output["text_embeddings"].shape == (2, 5, 32) + assert output["text_embeddings_mask"].tolist() == [ + [1, 1, 1, 0, 0], + [1, 1, 1, 1, 1], + ] + assert output["negative_text_embeddings"].shape == (2, 4, 32) + assert output["negative_text_embeddings_mask"].tolist() == [ + [1, 1, 0, 0], + [1, 1, 1, 1], + ] + + +def test_qwen_image_edit_processor_is_registered(): + pytest.importorskip("torch") + from preprocess.processors import ProcessorRegistry + + assert ProcessorRegistry.is_registered("qwen_image_edit") diff --git a/tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py new file mode 100644 index 00000000000..2cb379d8f66 --- /dev/null +++ b/tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Qwen-Image-Edit DMD2 target/reference-token wrapper.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import DMDConfig +from modelopt.torch.fastgen.config import SampleTimestepConfig +from modelopt.torch.fastgen.plugins.qwen_image import ( + attach_feature_capture as attach_t2i_feature_capture, +) +from modelopt.torch.fastgen.plugins.qwen_image import ( + pack_latents, + remove_feature_capture, + update_feature_capture_shape, +) +from modelopt.torch.fastgen.plugins.qwen_image_edit import QwenImageEditDMDPipeline +from modelopt.torch.fastgen.plugins.qwen_image_edit import ( + attach_feature_capture as attach_edit_feature_capture, +) + + +class _CapturingModel(nn.Module): + """Record Qwen kwargs and return the packed input in a requested output style.""" + + def __init__(self, style: str = "tensor") -> None: + super().__init__() + self.style = style + self.last_kwargs: dict[str, object] = {} + + def forward(self, **kwargs): + self.last_kwargs = dict(kwargs) + output = kwargs["hidden_states"] + if self.style == "tensor": + return output + if self.style == "tuple": + return (output,) + if self.style == "sample": + return SimpleNamespace(sample=output) + raise ValueError(self.style) + + +class _CurrentDiffusersSignatureModel(nn.Module): + """Approximate current Diffusers Qwen forward, which removed ``txt_seq_lens``.""" + + def __init__(self) -> None: + super().__init__() + self.called = False + + def forward( + self, + hidden_states, + encoder_hidden_states, + encoder_hidden_states_mask, + timestep, + img_shapes, + guidance=None, + attention_kwargs=None, + return_dict=True, + ): + self.called = True + return hidden_states + + +class _GenericForwardWrapper(_CurrentDiffusersSignatureModel): + """Mimic a distributed wrapper that exposes ``**kwargs`` around an explicit model API.""" + + def forward(self, *args, **kwargs): + return super().forward(*args, **kwargs) + + +def _make_pipeline(student: nn.Module, *, guidance: float | None = None): + return QwenImageEditDMDPipeline( + student=student, + teacher=nn.Identity(), + fake_score=nn.Identity(), + config=DMDConfig(num_train_timesteps=None), + discriminator=None, + guidance=guidance, + ) + + +@pytest.mark.parametrize("style", ["tensor", "tuple", "sample"]) +def test_call_model_appends_multiple_references_and_crops_target_prefix(style): + """Packed references follow the target, while only target output tokens are unpacked.""" + b, c, h, w = 2, 4, 8, 10 + target = torch.arange(b * c * h * w, dtype=torch.float32).reshape(b, c, h, w) + references = [ + torch.full((b, c, 6, 8), 10_000.0), + torch.full((b, c, 4, 12), 20_000.0), + ] + model = _CapturingModel(style) + pipe = _make_pipeline(model, guidance=1.25) + timestep = torch.tensor([0.25, 0.75]) + text = torch.randn(b, 7, 16) + text_mask = torch.tensor([[1, 1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 1, 0, 0]]) + + output = pipe._call_model( + model, + target, + timestep, + encoder_hidden_states=text, + encoder_hidden_states_mask=text_mask, + conditioning_latents=references, + ) + + expected_packed = torch.cat( + [pack_latents(target), *(pack_latents(x) for x in references)], dim=1 + ) + kwargs = model.last_kwargs + assert torch.equal(kwargs["hidden_states"], expected_packed) + assert kwargs["img_shapes"] == [ + [(1, 4, 5), (1, 3, 4), (1, 2, 6)], + [(1, 4, 5), (1, 3, 4), (1, 2, 6)], + ] + assert "txt_seq_lens" not in kwargs + assert torch.equal(kwargs["timestep"], timestep) + assert torch.equal(kwargs["guidance"], torch.full((b,), 1.25)) + assert kwargs["return_dict"] is False + assert "conditioning_latents" not in kwargs + + # The capturing model echoes target+reference tokens. Cropping the prefix before + # unpacking must recover the target bit-exactly rather than trying to unpack all tokens. + assert torch.equal(output, target) + + +def test_conditioning_latent_validation_errors_are_actionable(): + b, c, h, w = 1, 4, 8, 8 + target = torch.randn(b, c, h, w) + text = torch.randn(b, 2, 8) + timestep = torch.tensor([0.5]) + pipe = _make_pipeline(_CapturingModel()) + + def call(conditioning_latents): + return pipe._call_model( + pipe.student, + target, + timestep, + encoder_hidden_states=text, + conditioning_latents=conditioning_latents, + ) + + with pytest.raises(ValueError, match="non-empty.*conditioning_latents"): + call(None) + with pytest.raises(ValueError, match="list or tuple"): + call(torch.randn_like(target)) + with pytest.raises(TypeError, match=r"conditioning_latents\[0\].*Tensor"): + call(["not-a-tensor"]) + with pytest.raises(ValueError, match=r"conditioning_latents\[0\].*\[B, C, H, W\]"): + call([torch.randn(c, h, w)]) + with pytest.raises(ValueError, match="batch/channels"): + call([torch.randn(2, c, h, w)]) + with pytest.raises(ValueError, match="even spatial"): + call([torch.randn(b, c, h - 1, w)]) + with pytest.raises(ValueError, match="dtype"): + call([torch.randn(b, c, h, w, dtype=torch.bfloat16)]) + + +def test_current_diffusers_signature_does_not_receive_removed_txt_seq_lens(): + model = _GenericForwardWrapper() + pipe = _make_pipeline(model) + target = torch.randn(1, 4, 8, 8) + reference = torch.randn(1, 4, 8, 8) + text = torch.randn(1, 3, 8) + mask = torch.ones(1, 3, dtype=torch.long) + + output = pipe._call_model( + model, + target, + torch.tensor([0.5]), + encoder_hidden_states=text, + encoder_hidden_states_mask=mask, + conditioning_latents=[reference], + ) + + assert model.called + assert output.shape == target.shape + + +class _TinyEditTransformer(nn.Module): + """Grad-capable packed-token transformer used for an end-to-end DMD loss call.""" + + def __init__(self, packed_dim: int = 16) -> None: + super().__init__() + self.proj = nn.Linear(packed_dim, packed_dim) + self.seen_token_counts: list[int] = [] + + def forward(self, hidden_states, **_kwargs): + self.seen_token_counts.append(hidden_states.shape[1]) + return self.proj(hidden_states) + + +def test_shared_dmd_losses_forward_references_to_student_teacher_and_fake_score(): + """All DMD branches receive the fixed reference suffix through ``model_kwargs``.""" + torch.manual_seed(0) + student = _TinyEditTransformer() + teacher = _TinyEditTransformer() + fake_score = _TinyEditTransformer() + config = DMDConfig( + pred_type="flow", + num_train_timesteps=None, + student_sample_steps=1, + guidance_scale=None, + gan_loss_weight_gen=0.0, + sample_t_cfg=SampleTimestepConfig(time_dist_type="uniform", min_t=0.001, max_t=0.999), + ema=None, + ) + pipe = QwenImageEditDMDPipeline(student, teacher, fake_score, config) + target = torch.randn(1, 4, 8, 8) # 16 target patches + noise = torch.randn_like(target) + reference = torch.randn(1, 4, 4, 8) # 8 reference patches + text = torch.randn(1, 3, 8) + + student_losses = pipe.compute_student_loss( + target, + noise, + encoder_hidden_states=text, + conditioning_latents=[reference], + ) + assert torch.isfinite(student_losses["total"]) + student_losses["total"].backward() + assert any(p.grad is not None for p in student.parameters()) + assert student.seen_token_counts == [24] + assert teacher.seen_token_counts == [24] + assert fake_score.seen_token_counts == [24] + + fake_score.zero_grad(set_to_none=True) + fake_losses = pipe.compute_fake_score_loss( + target, + noise, + encoder_hidden_states=text, + conditioning_latents=(reference,), + ) + assert torch.isfinite(fake_losses["total"]) + fake_losses["total"].backward() + assert any(p.grad is not None for p in fake_score.parameters()) + assert student.seen_token_counts[-1] == 24 + assert fake_score.seen_token_counts[-1] == 24 + + +class _TupleBlock(nn.Module): + def forward(self, hidden_states): + return torch.empty(0), hidden_states + + +class _TeacherWithBlocks(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = nn.ModuleList([_TupleBlock()]) + + +def test_edit_feature_capture_keeps_target_prefix_and_t2i_remains_strict(): + b, target_h, target_w, channels = 1, 8, 6, 5 + target_tokens = (target_h // 2) * (target_w // 2) + hidden = torch.arange(b * (target_tokens + 7) * channels, dtype=torch.float32).reshape( + b, target_tokens + 7, channels + ) + + edit_teacher = _TeacherWithBlocks() + attach_edit_feature_capture(edit_teacher, [0], target_h, target_w) + edit_teacher.transformer_blocks[0](hidden) + captured = edit_teacher._fastgen_captured + assert len(captured) == 1 + expected = ( + hidden[:, :target_tokens] + .permute(0, 2, 1) + .reshape(b, channels, target_h // 2, target_w // 2) + ) + assert torch.equal(captured[0], expected) + + # The installed hooks must follow later multiresolution batches instead of retaining + # the base-resolution shape present at hook registration time. + captured.clear() + dynamic_h, dynamic_w = 4, 8 + dynamic_tokens = (dynamic_h // 2) * (dynamic_w // 2) + dynamic_hidden = hidden[:, : dynamic_tokens + 3] + update_feature_capture_shape(edit_teacher, dynamic_h, dynamic_w) + edit_teacher.transformer_blocks[0](dynamic_hidden) + dynamic_expected = ( + dynamic_hidden[:, :dynamic_tokens] + .permute(0, 2, 1) + .reshape(b, channels, dynamic_h // 2, dynamic_w // 2) + ) + assert torch.equal(captured[0], dynamic_expected) + remove_feature_capture(edit_teacher) + + # The text-to-image helper keeps exact-length validation by default, preventing + # accidental resolution drift from being silently interpreted as edit references. + t2i_teacher = _TeacherWithBlocks() + attach_t2i_feature_capture(t2i_teacher, [0], target_h, target_w) + with pytest.raises(RuntimeError, match="seq_len"): + t2i_teacher.transformer_blocks[0](hidden) + remove_feature_capture(t2i_teacher) diff --git a/tests/unit/torch/fastgen/test_qwen_image_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_plugin.py index 498b6ce5f9f..07bd75fdce4 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_plugin.py @@ -33,6 +33,7 @@ from torch import nn from modelopt.torch.fastgen import DMDConfig +from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin from modelopt.torch.fastgen.plugins.qwen_image import ( QwenImageDMDPipeline, build_img_shapes, @@ -162,12 +163,13 @@ def _make_pipeline(student: nn.Module) -> QwenImageDMDPipeline: ) -def test_call_model_forwards_qwen_kwargs(): +def test_call_model_forwards_qwen_kwargs(monkeypatch): """``_call_model`` must forward the exact Qwen signature (hidden_states packed to ``[B, num_patches, 64]``, encoder_hidden_states verbatim, - encoder_hidden_states_mask verbatim, txt_seq_lens derived from the mask, + encoder_hidden_states_mask verbatim (Diffusers derives sequence lengths from it), img_shapes as ``[[(1, h//2, w//2)]] * B``, guidance=None, return_dict=False, timestep verbatim with no /1000 rescale).""" + monkeypatch.setattr(qwen_image_plugin, "_DIFFUSERS_NEEDS_TXT_SEQ_LENS", False) b, c, h, w = 2, 16, 32, 32 student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4), style="tensor") pipe = _make_pipeline(student) @@ -191,7 +193,7 @@ def test_call_model_forwards_qwen_kwargs(): assert tuple(kw["hidden_states"].shape) == (b, (h // 2) * (w // 2), c * 4) assert tuple(kw["encoder_hidden_states"].shape) == (b, 512, 3584) assert torch.equal(kw["encoder_hidden_states_mask"], mask) - assert kw["txt_seq_lens"] == [37, 42] + assert "txt_seq_lens" not in kw assert kw["img_shapes"] == [[(1, h // 2, w // 2)]] * b assert kw["guidance"] is None assert kw["return_dict"] is False @@ -199,6 +201,27 @@ def test_call_model_forwards_qwen_kwargs(): assert tuple(out.shape) == (b, c, h, w) +def test_call_model_forwards_legacy_txt_seq_lens(monkeypatch): + """Diffusers 0.35/0.36 still needs lengths derived from the attention mask.""" + monkeypatch.setattr(qwen_image_plugin, "_DIFFUSERS_NEEDS_TXT_SEQ_LENS", True) + b, c, h, w = 2, 16, 8, 8 + student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4)) + pipe = _make_pipeline(student) + mask = torch.zeros(b, 9, dtype=torch.long) + mask[0, :4] = 1 + mask[1, :7] = 1 + + pipe._call_model( + student, + torch.randn(b, c, h, w), + torch.tensor([0.25, 0.5]), + encoder_hidden_states=torch.randn(b, 9, 32), + encoder_hidden_states_mask=mask, + ) + + assert student.last_kwargs["txt_seq_lens"] == [4, 7] + + @pytest.mark.parametrize("style", ["tensor", "tuple", "sample"]) def test_call_model_unpacks_return_styles(style): """``_call_model`` must unpack ``tensor`` / ``tuple`` / ``.sample`` return