Skip to content

refactor(args): fold --diffusion-model into --hf-checkpoint - #142

Merged
Rockdu merged 2 commits into
refactor/args-recipe-groupsfrom
refactor/args-fold-diffusion-model
Aug 10, 2026
Merged

refactor(args): fold --diffusion-model into --hf-checkpoint#142
Rockdu merged 2 commits into
refactor/args-recipe-groupsfrom
refactor/args-fold-diffusion-model

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

This PR is stacked on #141.

What

  • Delete --diffusion-model. --hf-checkpoint becomes the single model reference, and is now required.
  • Point its three readers at the survivor: the sglang-d engine's model_path, LTX's DiT resolution, and resolve_diffusion_model_family.
  • Add --diffusion-model-family, so a checkpoint whose path lacks the family name can still name it.
  • Point the LTX recipe's --hf-checkpoint at Lightricks/LTX-2.3.

Why

Two flags named the same repo id from opposite ends, and five of the six recipes passed the identical value to both. The field the family lookup matches against has always been called hf_ckpt_name_patterns; miles has only --hf-checkpoint. One value means the engine and the trainer can no longer disagree about which pipeline is being trained.

Two things worth a reviewer's attention:

  1. A sentinel, not a read. The three if getattr(args, "diffusion_model", None): guards in miles_validate_args used that flag's non-empty default as an "is this a diffusion run" test, inherited from miles where the attribute does not exist. They are unconditional here — one caller, parse_args, whose parser always registers these arguments. Renaming underneath them would have flipped them off whenever --hf-checkpoint was absent, and silently: no TrainPipelineConfig resolution, no validate_sp_args, and sp_size pinned to 1 so dp_size and the global_batch_size divisibility assertion both came out wrong.

  2. Your own weights. One value also means a path lacking ltx/sd3/wan2.2/qwen-image no longer resolves — the normal case for a directory of your own weights, and previously servable by pointing the two flags at different places. That hatch existed only as MILES_DIFFUSION_MODEL_FAMILY, invisible in --help; it is now a flag, with precedence flag > env > name match. --train-pipeline-config-path still covers families that are not registered at all; passing both now raises. miles solves this from the other side — it derives identity from the checkpoint's own config with --model-name as the override — but LTX's official repo ships no model_index.json, so the pipeline-class route cannot serve the one family that needs it most.

Validation

Captured train_args from all six recipes plus the smoke and debug-alignment variants (7 total), before and after. The only differences: --diffusion-model and its duplicate value gone, and the LTX recipe's --hf-checkpoint now carrying Lightricks/LTX-2.3. num_gpus_per_node and the extra_env_vars passed to execute_train are unchanged for all seven, and the --diffusion-model-family commit leaves all seven byte-identical. Family resolution re-checked from --hf-checkpoint for every model the recipes use: LTX-2.3 → ltx, Wan2.2 → wan2_2, SD3.5 → sd3, Qwen-Image → qwen_image. The new flag's six branches were replayed standalone, including both error cases. pre-commit clean; py_compile clean.

Not verified here: pytest and train.py --help need torch/sglang/ray, which this environment lacks — CI covers both. No new tests: the added logic lives in miles_validate_args, reachable only through parse_args, which has no fixture in the repo.

Two flags named the same repo id from opposite ends. --hf-checkpoint fed the
diffusers train-side loader (components, scheduler, model_index.json) and the
eval dataset cache key; --diffusion-model fed the sglang-d engine's model_path,
LTX's checkpoint resolution, and the model-family lookup. Five of the six
recipes passed the identical value to both, and the config field the family
lookup reads has always been called hf_ckpt_name_patterns -- the split was
accretion, not design. miles has only --hf-checkpoint.

Keep --hf-checkpoint, delete --diffusion-model, and point its three readers at
the surviving flag. One value for one model means the engine and the trainer can
no longer disagree about which pipeline is being trained, and the family can no
longer be resolved from a different name than the weights.

--hf-checkpoint is now required. It defaulted to None while --diffusion-model
defaulted to stable-diffusion-3.5-medium, so a run that omitted both got an SD3.5
engine and family with a None training checkpoint -- it crashed in
DiffusersModelBackend anyway, just later and less clearly.

The three `if getattr(args, "diffusion_model", None):` guards in
miles_validate_args were using that non-empty default as an "is this a diffusion
run" sentinel. They are unconditional here: miles_validate_args has exactly one
caller, parse_args, whose parser always registers these arguments. Renaming
underneath them would have silently flipped them off whenever --hf-checkpoint was
absent, skipping the TrainPipelineConfig resolution and validate_sp_args, and
pinning sp_size to 1 so dp_size and the global-batch-size divisibility check came
out wrong. They are now plain statements.

The LTX recipe's `--hf-checkpoint gpt2` is dropped rather than carried over. LTX
loads through MilesModelBackend, which resolves the DiT from the model reference
and never touches hf_checkpoint; the only read on that path is the eval dataset
cache key, and the recipe sets no eval flags at all. gpt2 also matches no
family's hf_ckpt_name_patterns, which is precisely why the family lookup had to
read the other flag. The recipe now passes --hf-checkpoint Lightricks/LTX-2.3.

Verified by capturing the generated train_args of all six recipes plus the smoke
and debug-alignment variants (7 in total): the only change is the removal of
--diffusion-model and its duplicate value, and for LTX the removal of the dead
gpt2, with Lightricks/LTX-2.3 surviving on --hf-checkpoint.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… their family

Folding --diffusion-model into --hf-checkpoint left one path narrower than it was.
Before, you could point --hf-checkpoint at a local directory and --diffusion-model
at the hub id, and the family lookup read the hub id. Now there is one value, so a
checkpoint whose path does not happen to contain "ltx", "sd3", "wan2.2" or
"qwen-image" cannot be resolved -- which is the normal case for a directory of your
own weights.

The escape hatch for that already existed but only as MILES_DIFFUSION_MODEL_FAMILY,
an environment variable, so it never appears in --help. Promote it to a flag:

  --diffusion-model-family ltx --hf-checkpoint /root/models/my-own-dit

Precedence is flag, then the env var, then name matching. --train-pipeline-config-path
remains the answer for a family that is not registered at all; passing both now
raises instead of silently letting the config path win.

The value is normalized to stripped lowercase on args, because downstream readers
compare it exactly -- encoder_hub.get_encoder tests `family == "wan2_2"`, so
"WAN2_2" would have missed.

This mirrors what miles does for the same problem from the other side. miles never
pattern-matches the checkpoint name: it derives the model identity from the
checkpoint's own config, `type(load_hf_config(args.hf_checkpoint)).__name__.lower()`,
with --model-name as the explicit CLI override. Deriving from the pipeline class in
model_index.json would be the closer analogue here, but LTX's official repo ships no
model_index.json -- it loads a single-file DiT through MilesModelBackend -- so the
one family that most needs the hook is the one that path cannot serve. Name matching
with an explicit override is what works for all four families.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@Rockdu
Rockdu merged commit 44bca75 into main Aug 10, 2026
17 checks passed
Rockdu added a commit to Rockdu/miles_diffusion that referenced this pull request Aug 10, 2026
radixark#138); regroup train flags by concern (ckpt/rollout/eval/grpo/optimizer/lora/reward/wandb/sglang/train_backend/perf/misc); drop folded --diffusion-model
niehen6174 added a commit to niehen6174/miles_diffusion that referenced this pull request Aug 11, 2026
Resolve SKILL.md in favor of upstream's install-skill rewrite (radixark#146). The
H3 diffusers pin stays documented in requirements.txt. Also drop the recipe's
--diffusion-model flag, which radixark#142 folded into --hf-checkpoint.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant