feat(qwen image): dmd2 preprocess fixes - #3593
Open
linnanwang wants to merge 3 commits into
Open
Conversation
The Qwen-Image processor discarded the attention mask returned by encode_prompt, so downstream collation had to synthesize one from the raw embedding length -- wrong whenever encode_prompt pads to a fixed max_sequence_length, since padded positions would be marked as real tokens. Thread the real mask through processor -> dataset -> collate instead, with the old length-derived fallback kept for caches preprocessed before this fix. Signed-off-by: linnan wang <linnanw@nvidia.com>
The preprocessing CLI had no way to split a large image set across multiple jobs/nodes without them clobbering each other's metadata files in a shared output_dir. Add --shard_idx/--shard_count so each rank processes image_files[shard_idx::shard_count] and writes rank-namespaced metadata shard/index files; default single-job output is unchanged. Signed-off-by: linnan wang <linnanw@nvidia.com>
The DMD2 recipe could train a student but Automodel had no way to sample from or deploy the result: a DMD2 student is trained on the rectified-flow identity, not the noise-prediction loss a stock diffusers scheduler expects, so it can't be sampled with standard scheduler stepping. Add QwenImageDMDInferencePipeline (bit-aligned single/multi-step RF sampler, optional CFG and EMA overlay) and export_diffusers_pipeline (assembles a consolidated transformer checkpoint + base pipeline components into a full QwenImagePipeline.from_pretrained-loadable dir). Both depend only on diffusers/torch, not on ModelOpt, so they work independent of the pinned fastgen branch's availability. Signed-off-by: linnan wang <linnanw@nvidia.com>
akoumpa
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes a masking bug in Qwen-Image preprocessing, adds multi-node sharded
preprocessing, and adds a DMD2 inference sampler + diffusers export tool —
closing the data-preprocessing-through-inference gaps around this PR's DMD2
Qwen-Image distillation recipe.
Changelog
prompt_embeds_maskfrom Qwen-Image'sencode_promptthrough the processor,TextToImageDataset, andcollate_fn_production, instead of synthesizing a mask from the raw(possibly already-padded) embedding length — the synthesized fallback is
wrong whenever
encode_promptpads to a fixedmax_sequence_length.--shard_idx/--shard_counttotools/diffusion/preprocessing_multiprocess.py'simagesubcommand formulti-node/multi-job preprocessing sweeps, with rank-namespaced metadata
files so concurrent jobs sharing one
output_dirdon't collide. Defaultsingle-job output is unchanged.
tools/diffusion/inference_dmd2_qwen_image.py(
QwenImageDMDInferencePipeline) — a bit-aligned single/multi-steprectified-flow sampler for DMD2-trained Qwen-Image students, matching the
training-time math in ModelOpt's
fastgen, with optional inference-timeCFG and EMA-shadow overlay.
tools/diffusion/export_diffusers_pipeline.pytoassemble a consolidated transformer checkpoint + base pipeline components
into a full
QwenImagePipeline.from_pretrained-loadable directory(symlink by default,
--copyfor portability).sampler's schedule-resolution logic, and the export tool — all CPU-only,
no GPU/ModelOpt/diffusers required to run.
Before your PR is "Ready for review"
Pre checks:
Additional Information
main, since it extendsthe DMD2 recipe that PR introduces.
diffusers/torch, not onModelOpt, so they work independent of the pinned
fastgenbranch.pre-commit/ruffcould not be run in the environment that authoredthese commits (no
pre-commit/uvreachable there) — please runpre-commit runlocally before merging to confirm formatting/lint.