Staged sampling - #1688
Open
dxqb wants to merge 5 commits into
Open
Staged sampling#1688dxqb wants to merge 5 commits into
dxqb wants to merge 5 commits into
Conversation
Suppress a handful of specific, noisy-but-harmless messages emitted while launching the UI and starting training: - diffusers/transformers logger.warning() lines (Modular Diffusers experimental notice, unexpected-config-attributes, unrecognized loss_type) via filters on the exact emitting loggers - huggingface_hub local_dir_use_symlinks deprecation and the torch.compile inductor performance notes via warnings/logger filters - Qt gnome portal dbus errors via QT_LOGGING_RULES - tensorboard subprocess banner/notices by discarding its stdout/stderr Each filter targets one specific message, so other warnings from the same libraries still come through. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # modules/util/ui/pyside6_util.py
A cold torch.compile cache announces every frame it compiles, which scrolls the progress bar off the screen. There is no knowable total to build a real progress bar from, so the announcement goes into the postfix of the innermost running bar instead, and is cleared again by that bar's next redraw or by its close. tqdm keeps its bars in an unordered WeakSet, so which of the nested bars is the innermost one cannot be recovered from it. modules/util/tqdm_util.py subclasses tqdm to track that itself and adds show_status() next to tqdm.write(); every tqdm import in the repo now comes from there. Bars owned by mgds are outside this and still draw as before. Also gates the warning filters on OT_DEBUG_WARNINGS, so setting it brings every suppressed message back, and silences the diffusers attention-backend experimental notice. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sampling used to run each sample end to end, so with several sample prompts
every model part was brought on-device once per sample: text encoder, then
transformer, then VAE, then back to the text encoder for the next prompt.
This adds modules/util/staged_pipeline.py. run_staged_pipeline() takes a list of
(label, stage) pairs and a column-oriented dict of per-sample inputs, and runs
every sample through stage 0, then every sample through stage 1, and so on. Each
sample carries a context dict that accumulates the stages' outputs, so a value
produced early reaches any later stage without being threaded through the ones
between; a stage only receives the arguments its signature names.
The Flux2, Qwen, Krea2, ZImage and Anima samplers are split into __encode,
__denoise and __decode stages and drive them through that pipeline in a new
sample_all(). Each part is now materialized once per stage instead of once per
sample. BaseModelSampler.sample_all() has an item-major default that just loops
sample(), so samplers that are not split keep working unchanged.
GenericTrainer hands the whole enabled sample list to sample_all() in one call.
Custom samples ("sample now") all share the samples/custom directory, so their
filenames now get the sample index appended - without it a batch finishing
within the same second wrote every image to the same path. Custom sample
requests that arrive while sampling is already running are drained before
training resumes instead of waiting for the next step.
The sample window keeps the images of a batch and gets prev/next buttons with a
counter to page through them.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…uashed # Conflicts: # modules/trainer/GenericTrainer.py
Collaborator
Author
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.
Summary
Instead of cycling model parts for each sample, run all samples batched through the pipeline
Meaning: run text encoder for all samples, then denoising for all samples, then VAE for all samples
this does not mean the samples are send batched through the denoiser. only the sampling stages run batched
Intended for future large models and #1642 but some of the more recent and in-use samplers were ported to staged samplers anyway:
Others are left sample-major
Test plan
pre-commit run --all-filespassesAI assistance