fix(studio): img2img/inpaint on GGUF image models — from_pipe fallback with no recast (#9186) - #9193
Conversation
![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: e704b33c6c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Do not force UNSLOTH_DISABLE_MLX_AUTOREPAIR=1 for every caller of this shared boot script. The no-mlx-repair scenario in tests/studio/mac_capability_verdict_smoke.py explicitly removes this variable and waits for the repair installer to start, and .github/workflows/mlx-ci.yml invokes that scenario through this script. This assignment reintroduces the variable in the server process, so the installer is never invoked and the scenario waits up to 600 seconds before failing.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed and pushed (6ac313f incl. pre-commit formatting): quantized residents are now detected BEFORE the fast path (quantization_config probe on every component) and go straight to the manual re-wire — from_pipe's internal .to(dtype) recasts components in place before reaching the quantized transformer and raising, so skipping it entirely is the only way to keep the resident txt2img pipeline unpolluted. New AST test proves the mutating path is never entered when a quantized component is resident.
The env-scoping P1 from the other thread is already fixed on #9191 (the boot script is back to neutral; the opt-out lives in the smoke job's env).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6ac313f5d2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| components manually: same resident modules, same references, no cast. | ||
| """ | ||
| try: | ||
| return pipe_cls.from_pipe(pipe, torch_dtype = None, **extra) |
There was a problem hiding this comment.
When the bundled from_pipe reaches the quantized transformer after visiting another shared component, those earlier components have already been recast to float32 before the exception is raised. The fallback then reuses these mutated references rather than restoring their original dtypes, so GGUF Qwen img2img/inpaint can still incur a large memory increase or fail from mixed dtypes, and the resident txt2img pipeline is modified as well. Detect the quantized path before calling the mutating fast path, or restore every component's dtype before constructing the fallback.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed and pushed (6ac313f incl. pre-commit formatting): quantized residents are now detected BEFORE the fast path (quantization_config probe on every component) and go straight to the manual re-wire — from_pipe's internal .to(dtype) recasts components in place before reaching the quantized transformer and raising, so skipping it entirely is the only way to keep the resident txt2img pipeline unpolluted. New AST test proves the mutating path is never entered when a quantized component is resident.
The env-scoping P1 from the other thread is already fixed on #9191 (the boot script is back to neutral; the opt-out lives in the smoke job's env).
| if (!prompt) { | ||
| return false; | ||
| } | ||
| return /{{\s*\$now\s*}}|{{\s*\$time\s*}}/i.test(prompt); |
There was a problem hiding this comment.
The /i flag makes prompts such as {{$NOW}} and {{$TIME}} trigger the warning, but resolveSystemPromptVariables performs a case-sensitive lookup against only $now and $time; uppercase names are therefore left unchanged (or can resolve as custom variables) and do not invalidate the cache each request. This produces a false performance warning for those inputs, including the uppercase case asserted by the new test.
Useful? React with 👍 / 👎.
|
Confirmed the call in studio/backend/core/inference/diffusion.py still passes torch_dtype=None into from_pipe, so the no-recast fallback is aimed at the right place. The branch also carries the mlx_repair import retry, the prompt time variable UI warning and the CI timeout changes, so could you narrow it to the diffusion fix and its test and split the rest out? |
|
@Lyxot could you take a look at this one? |
…k with no recast (unslothai#9186) Loading a GGUF-quantized image model and generating with an input photo failed with 'Casting a quantized model to a new dtype is unsupported': the app already passes torch_dtype=None to from_pipe intending 'reuse resident modules, no recast', but the bundled diffusers resolves None to float32 inside from_pipe and calls .to(dtype) on every component, which hard-crashes the GGUF-quantized transformer. Txt2img never calls from_pipe, which is why only photo workflows broke. _from_pipe_no_recast keeps from_pipe as the fast path and, when the quantized-cast error fires, re-wires the resident components directly (same module references, same quantized state, no cast) — matching what from_pipe does minus the dtype step. Strict constructor signatures get exactly the parameters they declare; **kwargs-style pipelines accept the full component set. Both the workflow pipe and the ControlNet pipe route through it. Four duck-typed tests (AST-extracted helper, the repo's established no-torch convention): the fast path still calls from_pipe with torch_dtype=None; the quantized crash falls back and preserves module identity; unrelated errors still raise; extra components (controlnet) forward through the fallback.
…lothai#9186) from_pipe hands the workflow pipeline the resident pipeline's own module objects and then casts them, so the cast rewrites the loaded model rather than a copy of it. It resolves its dtype argument to float32 whenever the caller names none, and an explicit torch_dtype=None stopped counting as naming one, so the argument both call sites passed to prevent the cast no longer prevents anything. On a quantized denoiser that raises "Casting a quantized model to a new `dtype` is unsupported", which is why img2img, inpaint, upscale and ControlNet failed on GGUF image models while txt2img, which never calls from_pipe, kept working. With nothing quantized it does not raise at all: every component is silently upcast to float32, including the resident text-to-image pipeline's, since the objects are shared. Catching the error is not enough. Components are cast in name order, so on Qwen-Image the text encoder is already float32 by the time the transformer refuses, and the resident pipeline is left holding it. Build the workflow and ControlNet pipelines through a class whose .to() drops the dtype and keeps the device instead. from_pipe still does its own assembly, its final cast becomes a no-op, and no dtype argument is passed, which behaves the same on the released diffusers and on the pinned revision.
6ac313f to
c3e223c
Compare
|
Force-pushed a rebase onto current Dropped six unrelated commits that were riding along on this branch — the MLX gate race (#9120), the prompt time-variable warning (#9177), and three CI self-heal fixes (#9183 and two follow-ups). The PR is now 2 commits over 3 files, and no longer conflicting. The original fix commit is kept as-authored. Reworked the fix. The fast-path-then-catch approach turns out to be insufficient, and it took a real GGUF load to see why. There is also a second failure the catch could never have reached: when nothing is quantized, So the recast is now prevented rather than caught: the workflow and ControlNet pipelines are built through a class whose One thing worth flagging for reviewers, because it changes how this reads: Diffusers 0.39.0 recasts here too. Its Validated on both revisions against a real |
|
@codex review |
4 participants
Closes #9186
Closes #9241
The bug
from_pipehands the new pipeline the resident pipeline's own module objects, then casts them. It resolves its dtype argument to float32 whenever the caller names none — and on the pinned Diffusers revision an explicittorch_dtype = Noneno longer counts as naming one, which is exactly what both call sites passed, under comments calling it load-bearing.So the cast lands on the loaded model rather than on a copy of it, and it fails in two different ways:
ModelMixin.to()refuses any dtype onceis_quantizedis set, so image-conditioned workflows die withCasting a quantized model to a new dtype is unsupported. This is the reported crash. Txt2img never callsfrom_pipe, which is why only photo workflows broke.Both reach
ControlNettoo, not just img2img/inpaint/upscale:_controlnet_pipemade the same call.The fix
Build the workflow and ControlNet pipelines through a class whose
.to()drops the dtype and keeps the device.from_pipestill performs its own assembly — optional components, type-hint filtering, config re-registration — and its terminal cast becomes a no-op. No dtype argument is passed at all.Two properties this buys over suppressing the cast per call site:
torch_dtypedefault istorch.float32, notNone— so the suppression is load-bearing on the released Diffusers as well as the pin. And if huggingface/diffusers#12762 or a successor ever lands,from_pipestops casting and the override simply never fires. Nothing needs to change here either way; the test suite pins both cases.Why not catch the error instead
An earlier revision of this PR kept
from_pipe(torch_dtype = None)as a fast path and rebuilt frompipe.componentswhen the quantized-cast error fired. That is not sufficient, and measurement against a real GGUF Qwen-Image load is what showed it:DiffusionPipeline.to()walks components insorted()order —scheduler, text_encoder, tokenizer, transformer, vae— so the Qwen2.5-VL text encoder is already float32 by the time the transformer refuses. Catching the error does not undo it, and the module is shared, so txt2img afterwards runs a 33 GB fp32 encoder. It also left the unquantized silent-upcast case entirely unaddressed, since nothing raises there.Validation
Reproduced and fixed end to end against the pinned Diffusers revision in an isolated environment, using a real
unsloth/Qwen-Image-2512-GGUFQ4_K_M transformer with its companion VAE, text encoder, tokenizer and scheduler. The helper is lifted out ofdiffusion.pyby AST at run time rather than retyped, so the validation exercises the source that ships.Checked on both Diffusers 0.39.0 and the pinned 0.40.0.dev0:
to(dtype = torch.float32)on a workflow pipeline is a no-op; a device move still worksTest-side: every added test was mutation-checked. Ten mutations of the production change — including bypassing the no-recast class at each call site, and removing each individual line of
to()— each fail a test that names the behaviour.test_from_pipe_no_recast_leaves_every_component_at_its_loaded_dtypeis parametrized over quantized/unquantized and over afrom_pipethat recasts versus one that has stopped.Note on the tests
The standalone
test_diffusion_from_pipe_no_recast.pyfrom the first revision is gone; its coverage lives intest_diffusion_backend.py, which already owns_workflow_pipe, and intest_diffusion_controlnet.pyfor the ControlNet path. One pre-existing assertion there changed: it asserted that Studio passestorch_dtype = Nonetofrom_pipe, which is the belief that caused this bug. It now asserts that the terminal cast never reaches the components.