Repository navigation
fix(z-image): run Fun ControlNets through the shared hooks, as trained - #9943
Open
Pfannkuchensack wants to merge 12 commits into
Open
Pfannkuchensack wants to merge 12 commits into
Pfannkuchensack wants to merge 12 commits into
Conversation
…trol layers Add alibaba-pai's VACE-style control branch for Qwen-Image: identification, loader, the qwen_image_control node and control/inpaint conditioning in the denoise, injected through forward hooks on the diffusers transformer. The canvas gets a Qwen Image ControlNet layer kind with inpaint/outpaint feeding and a grayscale filter; the denoise now reserves working memory for its activations.
…ontrol layers Loads alibaba-pai's FLUX.2-dev union as a FLUX.2 ControlNet, runs its control blocks on the shared hook helper and adds a Canvas control kind with inpaint wiring. The control encodes inside Denoise - FLUX.2 retry tiled after running out of memory, like Image to Latents - FLUX.2.
…feat/flux2-fun-controlnet
…trings The Qwen-Image and FLUX.2 [dev] ControlNet strings add 0.5 KB to the locale file every route loads; request counts are unchanged.
The control blocks now hook into the diffusers forward, with each build's own layer map and the 2.x refiner hints, so regional prompting also works on control steps. Lite, 2.0 and the first Union now follow the control image instead of producing noise or ignoring it; the old custom control forward is removed.
Member
Author
This branch has not been deployed
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
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
Z-Image's Fun ControlNets ran through Invoke's own copy of the transformer forward, which did not match how alibaba-pai trained them: the hints sat after the wrong layers for the Lite and first Union builds, the 2.x refiner hints were missing (their weights were dropped as unexpected keys), 2.0 refined its control through the wrong blocks, the missing
x_pad_tokengot a random value, and regional prompting was lost on control steps. In practice Lite and 2.0 produced noisy images and the first Union ignored the control image.This PR runs the Z-Image ControlNets through the same hook helper as the Qwen-Image (#9914) and FLUX.2 (#9924) unions, on the diffusers forward:
[control | 0 | 0]); pad rows come from the base model'sx_pad_token.The node, its fields and the UI are unchanged; the control is now stronger at the same weight for the builds that were wrong before.
Related Issues / Discussions
Builds on #9924 (and #9914). GGUF Z-Image bases need #9942 (every GGUF Z-Image generation with 2-D pad tokens failed before that fix, with or without control).
QA Instructions
RTX 4090, Z-Image Turbo int8 + Qwen3-4B + FLUX VAE, 8 steps, scale 0.9 (Tile 0.85), same seeds, old code vs this PR on the same server root:
Working memory, measured on the hooked forward at 1024 px (hint-sized tensors above no control vs the
n + 4charge): 1.0 1.78 / 10, Lite 0.89 / 7, 2.0 11.61 / 19, 2.1 11.56 / 19 — the coefficient stays, the measured points in the tests are updated.CPU: per-build parity against an independent restatement of VideoX-Fun's forward on a tiny randomly initialised
ZImageTransformer2DModel(torch.equal; layer maps restated from the VideoX configs; wrong maps, dropped refiner hints, wrong pad rows and a missing caption join each fail it); lifetime tests that the control stream is released before the last control block; loader tests on the real headers (1.0, 2.1, Lite). Focusedpytestover the Fun-ControlNet, Z-Image, Qwen-Image and FLUX.2 control/denoise/working-memory/loader tests: passed. Ruff clean.Review
Resolved: the shared chain and the stage hand-over kept control tensors alive longer than needed (now owned and released, with lifetime tests; this also trims one tensor per forward for Qwen-Image and FLUX.2); the parity tests read the layer maps from the implementation (now restated independently); the working-memory points described the removed forward (re-measured); docs no longer claim inpainting through a Z-Image ControlNet. Remaining: regional prompting with several regions is not bit-reproducible run to run (the collected region order varies; unrelated to this change); inpaint inputs on the Z-Image control node are not wired (future work).
Compatibility / Rollout
No schema or persisted-data changes; existing workflows and canvases run unchanged but produce different (corrected) images for the affected builds. The calibration script's control option is now
--control-variant.Checklist
What's Newcopy (if doing a release after this PR)Need help on this PR? Tag
@codesmith-botwith what you need.