Comment-only pass over the Python this PR touches: drop what the code already
says, collapse multi-line explanations that still read on one line, and keep
the reasoning that is not recoverable from the code. No code, docstring
semantics or behaviour changes; verified with an AST comparison against the
previous revision, and the backend suite is unchanged (same 37 environment
failures as before: the API integration tests that need a live keyed server,
the flash-attn install hooks, and the GPU memory fields).
Two correctness fixes:
- The cache keyed the checkpoint and its companion base by name only, so
a Hub repo advancing to a new commit, or a local directory updated in
place, kept returning embeddings from the previous text encoder. Pair
both with a revision marker: the locally resolved commit sha for a Hub
repo, config plus text-encoder file stats for a directory. Neither
loads the encoders, so a warm run still keeps them off the GPU.
- flow_shift only checked positivity, but JSON accepts 1e309, which
floats to inf, and inf <= 0 is False while NaN fails every comparison.
The sigma table then evaluates s * u / (1 + (s - 1) * u) as NaN, which
poisons every sampled sigma and saves a corrupted adapter while
progress looks normal. Require a finite value.
- start_training forwards resume_source_run_id to _start_training_impl, which
reads it. Without it every start raised NameError.
- Restore main's anchor in the load-marker order test: the file now has an
earlier `if config.is_gguf:`, so indexing the first one compared the wrong
branch.
- The two diffusion tests that reach diffusers now skip when it is absent,
matching the CPU repo-test env.
- The UI smoke finds nav rows that live in the sidebar's More flyout.
Qwen-Image's scheduler skips its static shift under use_dynamic_shifting,
so the DiT trainer was drawing UNSHIFTED uniform-schedule sigmas for it
(mean sigma 0.50) while inference always runs the exponential mu = log 3
shift plus the shift_terminal 0.02 stretch. Add a flow_shift config lever:
"auto" (the new qwen-image default) rebuilds the training sigma table
through the scheduler's own time_shift and stretch_shift_to_terminal so
the draw matches the inference distribution exactly (mean sigma 0.72);
a numeric value applies the standard linear shift s*u/(1+(s-1)*u); 1.0
keeps the historical identity behavior and stays the default for FLUX,
Z-Image and Krea 2. The model timestep conditioning follows the shifted
sigma, gathered in fp32 so bf16 rounding never skews it.
Also wire two opt-in levers with off defaults: cfg_dropout (per-sample
empty-prompt conditioning dropout, encoded alongside the captions before
the text encoders are freed) and weighting_scheme="bell" (bsmntw-style
mid-schedule Gaussian loss weighting normalized to mean 1).
Verified with two 80-step rank-8 bf16 LoRA runs on Qwen/Qwen-Image
(identity vs auto, same seed): both converge with finite decreasing loss
and produce coherent same-seed previews. Unit tests cover the exact
transform, the shifted sampling distribution, per-family defaults and
config plumbing.
