unsloth/plans/diffusion-workflows-pr-plan.md

Stacked-PR plan: Studio diffusion workflows (Images redesign)

Branch tip: diffusion-eager-and-compile-cache (latest commit "Phase 16 review fixes"). Remote: oobabooga/unsloth. New PRs stack on top of the existing diffusion stack (ultimately on top of unslothai/unsloth#6658), treated as one logical change.

Nothing here is committed yet (commit/push only on explicit request).

CRITICAL: the working tree holds TWO uncommitted streams, and three core files INTERMINGLE them

A full git status / marker audit (branch diffusion-eager-and-compile-cache, tip "Phase 16 review fixes") shows the uncommitted tree is NOT a clean single feature. There are two streams:

A) The eager/compile-cache phase (the branch's own in-progress work, NOT this session's — zero of this feature's markers). Purely-its files, safe to NOT touch in the workflow PRs:

• new modules: diffusion_arch_patches.py, diffusion_compile_cache.py, diffusion_eager_patches.py, diffusion_gguf_compile.py, diffusion_patch_backend.py
• new tests: test_diffusion_arch_patches.py, test_diffusion_compile_cache.py, test_diffusion_eager_patches.py, test_diffusion_gguf_compile.py
• modified: diffusion_speed.py, test_diffusion_speed.py, conftest.py, scripts/diffusion_bench.py, and ~25 untracked scripts/*bench*/*probe*/*orchestrator*.

B) The Images workflows feature (this session): the workflow engine + frontend + installer.

The two streams INTERMINGLE inside three shared files and are NOT separable by file:

• studio/backend/core/inference/diffusion.py — this feature's workflow hunks are interleaved with the eager/compile wiring (imports at L67-75; install_arch_patches/compile_cache.begin/ .restore/.save and the eager_patched/compile_cache_ctx state throughout load_pipeline/generate/unload). A single diffusion.py cannot go into one PR without the other stream's hunks.
• studio/backend/models/inference.py — this feature's init_image/mask_image/ reference_images/upscale/model_kind fields sit next to the pre-existing speed_mode/ transformer_prequant_path fields in the same request models.
• studio/backend/tests/test_diffusion_backend.py — this feature's workflow tests sit next to the pre-existing test_failed_load_rolls_back_eager_patches (imports diffusion_eager_patches).

Implication / options (USER DECIDES — it is their branch + their eager/compile work):

• CLEANLY separable now (purely this feature, can be committed/PR'd on their own anytime): frontend images-page.tsx + api.ts + pickers.tsx, and the sd.cpp installer install_sd_cpp_prebuilt.py + test_sd_cpp_install.py. (These are PR 2 and PR 3 below.)
• The backend engine (PR 1) CANNOT be cleanly split from the eager/compile phase via files. Realistic paths: (a) finalize + commit the eager/compile phase first, then this feature's backend lands as a clean diff on top; or (b) commit both streams together as the branch's next chunk (consistent with treating the stack as one logical change); or (c) a manual git add -p hunk split of the three shared files (tedious, risks a non-compiling intermediate). NOT auto-doable safely without the owner's intent for the eager/compile work.

Proposed stack (3 PRs, bottom to top)

PR 1 - Backend: diffusion workflow engine (safetensors + image-conditioned + editing)

Files:

• studio/backend/core/inference/diffusion.py (the feature hunks: three load "kinds" gguf/single_file/pipeline; _workflow_pipe via from_pipe(torch_dtype=None); _align_vae_dtype; generate() routing for reference/img2img/inpaint/upscale/edit; image-conditioned width/height from the input image (but reference + txt2img use the slider size); upscale (hires fix) branch on the img2img pipe; reference (FLUX.2) branch that passes the image(s) to the loaded pipe directly (no from_pipe, no strength) incl. multi- reference (reference_images combined into a list, capped at 3 extra); branch ORDER inpaint/upscale before reference so a mask/upscale request on a reference family still routes right; _family_workflows (adds "upscale" wherever img2img is supported, "reference" for reference families, "outpaint" only for size-preserving inpaint families); kind on state + model_kind in status; load_progress double-count fix). NOTE: this file ALSO carries pre-existing speed hunks if any landed here - review per-hunk and exclude non-feature hunks.
• studio/backend/core/inference/diffusion_families.py (img2img/inpaint pipeline slots; edit flag + reference flag + inpaint_preserves_size flag; qwen-image-edit + flux.1-kontext families; flux.2-klein gains reference + inpaint (no outpaint: FLUX.2 normalizes to ~1MP); detect_family longest-match + leftover-reject; layered reject).
• studio/backend/core/inference/diffusion_engine_router.py (model_kind -> diffusers for non-gguf kinds).
• studio/backend/core/inference/diffusion_memory.py (estimate_safetensors_dense_mib).
• studio/backend/core/inference/sd_cpp_backend.py (model_kind passthrough; reject img2img/inpaint on the native engine).
• studio/backend/models/inference.py (load request: optional gguf_filename, model_kind, init/mask/strength, advanced knobs; status: workflows, model_kind).
• studio/backend/routes/inference.py (model_kind forwarding; ValueError -> 400; exc_info logging).
• Tests: test_diffusion_backend.py, test_diffusion_routes.py.

Title: Studio diffusion: safetensors + image-conditioned + instruction-editing workflows Summary: Adds non-GGUF safetensors loading (full bnb-4bit pipelines + single-file fp8, gated to unsloth/*), the image-conditioned workflows (img2img, inpaint, outpaint via the inpaint path) built with Pipeline.from_pipe for zero-extra-VRAM component reuse, and instruction editing as its own family kind (Qwen-Image-Edit-2511 + FLUX.1-Kontext-dev). Fixes two bugs: from_pipe defaulting to a float32 recast that crashed torchao-quantized transformers, and image-conditioned calls forcing the slider size onto the input image.

PR 2 - Frontend: redesigned Images page (workflow tabs + Advanced Options)

Files:

• studio/frontend/src/features/images/images-page.tsx (workflow tabs Create/Transform/ Inpaint/Extend/Upscale/Reference/Edit; capability gating + auto-switch; MaskCanvas; buildOutpaint; Upscale tab with Scale + Detail-strength sliders; Reference tab (FLUX.2, reference image + add/remove extra references, no strength); Advanced Options accordion gated to GGUF for transformer-quant; spinner-overlap fix).
• studio/frontend/src/features/images/api.ts (request/status types incl. model_kind, upscale, reference_images).
• studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx (curated safetensors + edit GGUF rows; SUPPORTED_EDIT_KEYWORDS un-hide; layered hide).

Title: Studio Images: workflow tabs (create/transform/inpaint/extend/upscale/reference/edit) + Advanced Options Summary: Redesigns the Images page around capability-gated workflow tabs with a brush mask editor, client-side outpaint, a hires-fix upscale tab, a FLUX.2 reference tab, an instruction-edit tab, and an Advanced Options panel (speed/quant/attention/memory/step-cache/offload), plus the number-input spinner fix.

PR 3 - sd.cpp prebuilt installer hardening

Files: studio/install_sd_cpp_prebuilt.py, studio/backend/tests/test_sd_cpp_install.py. Title: Studio sd.cpp: pin release + verify sha256 + mirror-ready source Summary: Pins the stable-diffusion.cpp release (was tracking latest), verifies each download's sha256 against GitHub's published asset digest before extract/execute, adds a download timeout, and makes the source repo configurable (UNSLOTH_SD_CPP_REPO) so a future unslothai mirror needs no code change. Cleanly separable from the rest.

Pre-PR review (done)

An independent 3-angle review (backend correctness, frontend/UX, security/robustness) ran over the full session diff. No High findings; the load-gating to unsloth/* and the multi-reference count caps were verified intact end-to-end. Fixes applied before the PRs:

• Frontend [Med]: multi-reference slots no longer renumber mid-edit (dropped the eager filter(Boolean) in the per-slot onChange; empties dropped only at send).
• Backend [Med]: upscale now caps the ABSOLUTE output (longest side <= 2048), not just the factor, so a large upload * 4x can't OOM.
• Backend/security [Med]: _decode_b64_image rejects images > 4096px/side (uniform guard for init/mask/reference vs decompression-bomb / OOM inputs); base64 image fields capped at 32 MiB.
• Security [Low]: the native sd.cpp engine guard also rejects reference_images. All covered by new tests (82 backend pass) and a post-fix e2e (all five workflows still pass).

Post-deploy user feedback fixes (done)

From live use of the deployed studio:

• Backend [Med]: image-conditioned workflows passed the raw upload size to pipelines that require multiples of 16 (Z-Image/Qwen/FLUX), so an odd upload (e.g. 186px) failed with "Height must be divisible by 16". Added _snap_to_multiple and auto-resize init (and the matched mask) to the nearest /16 for img2img/inpaint/extend/edit. Verified live: a 186x250 Transform and Inpaint now return 200 at 192x256. Tests added.
• Frontend [Med]: the Advanced options (FP8/INT8 quant, speed, attention, memory) were a collapsed, muted accordion at the bottom of the left rail that users missed (HF screenshots discussion #25). Moved them into a RIGHT-DOCKED panel mirroring Chat's settings panel. Per follow-up (discussion #26): CLOSED by default, toggled by a SINGLE fixed top-bar button using Chat's LayoutAlignRightIcon that stays in the exact same position in both states (verified x/y identical open vs closed) and highlights when open. Controls extracted to a render-local advancedControls; unused Accordion import removed.

Constraints for execution (when authorized)

• Write as the user; no AI/bot mentions, no emojis, no em dashes; well-formatted bodies.
• gh auth status first. Push to oobabooga/unsloth, stack on the current branch.
• Keep PR 3 independent; PR 2 depends on PR 1 (frontend needs the backend contract).
• Re-run pytest studio/backend/tests/test_diffusion_*.py test_sd_cpp_install.py + frontend tsc/build before each PR.

Out of scope / follow-ups

Scope decisions are backed by HF download data in plans/diffusion-popularity-findings.md. The seven shipped workflows are create, transform, inpaint, extend, upscale, reference, edit.

• Publish the unslothai/stable-diffusion.cpp mirror + macOS/Windows staging (#152).
• ControlNet / style-transfer: the goal's "most popular" set is covered by the seven shipped workflows. ControlNet on the modern diffusers/FLUX/Qwen stack is niche by downloads (~3K/30d), so deferring it is data-backed, not just an effort call.
• FLUX.2-klein inpaint and multi-reference are DONE (shipped). Outpaint is intentionally not offered for FLUX.2 (it scales >1MP inputs to ~1MP). No further FLUX.2 follow-ups outstanding.