Fix image-generation load, training, and inference bugs#6872
Fix image-generation load, training, and inference bugs#6872oobabooga wants to merge 11 commits into
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
for more information, see https://pre-commit.ci
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
…hai#6872 huggingface_hub 1.x makes HfHubHTTPError's response argument required and reads only .headers and .request from it, while 0.x defaults it to None. Pass a small stub response so the test constructs the error on both, verified against 0.36.2 and the 1.22.0 wheel.
|
Thanks for the batch, these are great catches. I verified the FBCache reset claim against diffusers 0.39 (CacheMixin only exposes _reset_stateful_cache, so the old call really was a no-op) and the companion size double count, both fixes look right. Pushed ee61c7b to your branch to clear the last CI failure: the new hub error test constructed RepositoryNotFoundError without a response, which huggingface_hub 1.x (what CI resolves) makes required. It now passes a small stub that satisfies both 0.x and 1.x, checked against 0.36.2 and the 1.22.0 wheel. All 258 tests in the touched suites plus the import hoist lint pass locally on your head. One heads up: we will port the _reset_stateful_cache fix to the video backend on the video PR stack as a follow up, since it reuses the same pattern. |
The video backend never cleared FBCache residuals between clips, so with the step cache engaged a second generation at a different resolution would hit stale state from the first. Mirrors the image backend fix from unslothai#6872: call the transformer level _reset_stateful_cache (reset_stateful_hooks only exists on the HookRegistry in diffusers 0.39), covering transformer_2 for the Wan dual expert, only when a cache is engaged.
# Conflicts: # studio/backend/core/inference/diffusion.py # studio/backend/core/training/diffusion_train_common.py
# Conflicts: # studio/backend/core/inference/diffusion.py # studio/backend/routes/training.py # studio/backend/tests/test_diffusion_dataset_api.py
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 14dc3b86a7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for p in written: | ||
| try: | ||
| dest.unlink(missing_ok = True) | ||
| p.unlink(missing_ok = True) |
There was a problem hiding this comment.
When an upload targets an existing dataset and reuses a filename, this rollback now blindly unlinks every destination path if a later file in the same batch exceeds the size limit or hits an I/O error. Because this endpoint explicitly allows repeat uploads into the same dataset to accumulate, a failed batch like existing.png followed by an over-limit file can delete the user's pre-existing existing.png; write to temporary files and only replace originals after the whole batch is committed, or track/restore pre-existing paths.
Useful? React with 👍 / 👎.
| # Multiple revisions may be cached; the active one is the fullest, so take the max. | ||
| return max( | ||
| ( | ||
| DiffusionBackend._local_dir_weight_bytes(rev, exclude_transformer = True) |
There was a problem hiding this comment.
When transformer_quant is requested for a GGUF load without a pre-quantized source, the fast path later materializes the base repo's dense bf16 transformer/ before quantizing it, but this estimate now always excludes those shards from the resident budget. On cards that fit the selected GGUF plus companions but not the dense transformer, plan.offload_policy can stay none, so the loader evicts the current model and then OOMs/falls back during finalization instead of skipping the dense path up front; include the dense transformer size for that specific preflight.
Useful? React with 👍 / 👎.
2 participants
This fixes a batch of bugs in the local image-generation feature, found by running it on NVIDIA GPUs (RTX 6000 Ada, RTX 3090) and an Apple Silicon Mac (8GB M1), and from a code review pass. Branched off
image-generationand kept in sync with it. There are two commits. The first covers the highest-impact bugs on the load-and-generate path. The second covers everything else. Each bullet names how to reproduce the problem on theimage-generationbranch.All the image-generation test suites pass (1067 tests). The only failures are 14 unrelated, pre-existing ones in a flash-attention test that fails the same way on the base branch. The frontend compiles clean. A few of these were fixed upstream while this was in progress, so they are not duplicated here.
Commit 1: loading and generating images
unsloth/Z-Image-Turbo-GGUF(Q8_0) withtransformer_quant=fp8on a CUDA card with plenty of free VRAM. It loads unquantized: the status reportstransformer_quant: nulland the output matches a plain GGUF load, with no error. The load was overestimating the memory the quantized path needed (it double-counted the base model's transformer), decided it would not fit, and fell back. Confirmed live on a 49GB Ada: the same request now reportsfp8and produces a correct image.transformer_cache=fbcache, generate one image (say 1024x1024), then generate another at a different size (768x768). The second crashes on leftover state from the first. Confirmed live: two images at different sizes with the cache on now both succeed.unsloth/Z-Image-Turbo-GGUFat Q2_K on an 8GB M1) with default settings. It fails at load with "MPS backend out of memory" even though it fits in system RAM. Confirmed live on an 8GB M1: it now loads. While testing I also caught that my first attempt at the fix had no effect, and corrected it.Commit 2: everything else
Training:
lr_scheduler=constnat). It passed the initial checks, unloaded your model to free the GPU, and only then failed inside the training process. It is now rejected immediately. Confirmed live: an invalid value is refused with a 4xx before anything is unloaded.model_family=z-image(orqwen-image) andmixed_precision=fp16. It returned "running", unloaded your model, and failed a moment later on the next status poll. It is now refused up front. Confirmed live: refused immediately with "requires bf16"./models/gguf-experiments/). It was mistaken for a GGUF file, which cannot be trained. Confirmed by test.Generation:
unsloth/Qwen-Image-2512-FP8). It expands to roughly twice its on-disk size once loaded, but the planner budgeted the on-disk size, so it could accept a load that then runs out of memory. The estimate now accounts for the expansion. Confirmed by test.transformer_quant=fp8andmemory_mode=balanced(orlow_vram). The quantization is dropped, because those modes stream the model instead of keeping it fully in VRAM. Before, there was no indication; now it logs why. Confirmed live: the explanation appears in the log.Errors and edge cases:
Frontend polish:
Scripts:
sd-cli/sd-server) is still running. The delete could be refused. It now stops those processes first. (A related auto-delete was left out on purpose, since it could remove a folder the user manages themselves.)sparse_accum_probe.py,nvfp4_t211_probe.py) crashed on any machine without the original dev box's hard-coded output path. Fixed to a relative path.