Merge remote-tracking branch 'origin/image-generation' into diffusion-image-workflows
# Conflicts: # studio/backend/core/inference/diffusion.py # studio/backend/core/inference/diffusion_families.py # studio/backend/core/inference/diffusion_speed.py # studio/backend/core/inference/sd_cpp_engine.py # studio/backend/tests/test_diffusion_speed.py
This commit is contained in:
Daniel Han
commit
a4197d24c0
120 changed files with 580 additions and 160 deletions
|
|
@ -0,0 +1,34 @@
|
|||
- Decision: Used plan mode because `workflows/benchmarking_workflow.md` explicitly requires planning before benchmark execution.
|
||||
- Decision: User selected full scope: `all 8 variants` and both `GGUF` and dense `safetensors` fp8/int8 paths.
|
||||
- Decision: Chose a focused-comprehensive benchmark matrix of `44 configs` instead of an exhaustive `549-630-config` Cartesian product because the latter would require excessive downloads/GPU time and produce an unreadable table.
|
||||
- Decision: Excluded `FLUX.2-klein-9B` from actual sweep after `snapshot_download` returned `GatedRepoError: 403 Client Error`; needs manual HF license acceptance.
|
||||
- Decision: Treated failures as benchmark findings, not harness bugs: dense int8 small-M failures, `mxfp8` unsupported, Flux T5 fp8 text-encoder failure, and FBCache GGUF slowdown.
|
||||
- Created task `#66`: `Benchmark all supported diffusion models across all optimization levers`.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/diffusion_bench.py`: added/plumbed benchmark flags including `--transformer-cache {off,fbcache}` and related metrics/config output.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/benchall_orchestrator.py`: encodes the model/config matrix, runs `diffusion_bench.py` across GPUs, supports `--list`, `--gpu`, `--variants`, and `--aggregate`.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/benchall_predownload.py` via shell heredoc during predownload setup.
|
||||
- Overwrote `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/wobbly-jumping-narwhal.md` with approved benchmark plan.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/diffusion_benchall/SUMMARY.md` with benchmark findings and headline tables.
|
||||
- Generated `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/diffusion_benchall/results.csv` from aggregation; contains `40 rows`.
|
||||
- Commands run: `ls workflows/benchmarking_workflow.md unsloth/workflows/benchmarking_workflow.md`, `ls parallel_planner.py gpu_pool_runner.sh`, `ls unsloth/scripts/diffusion_bench.py unsloth/scripts`; found workflow, planner, runner, `diffusion_bench.py`, `diffusion_quality.py`.
|
||||
- Command run: `python3 -u parallel_planner.py --output_file extra_plan.md --prompt ...`; background task completed and wrote `/mnt/disks/unslothai/ubuntu/workspace_81/async_task_outputs/extra_plan.md`.
|
||||
- Command run: `python -c "import ast; ast.parse(open('scripts/diffusion_bench.py').read()); print('bench OK')"`; exit `0`, output `bench OK`.
|
||||
- Command run: `python scripts/diffusion_bench.py --help ... grep ...`; verified `--transformer-cache {off,fbcache}` appears.
|
||||
- Command run: HF access resolver using `huggingface_hub.HfApi`; found all 8 GGUF `Q4_K_M` filenames and initially `list_repo_files` access appeared OK.
|
||||
- Resolved GGUF files include `z-image-turbo-Q4_K_M.gguf`, `z-image-Q4_K_M.gguf`, `qwen-image-2512-Q4_K_M.gguf`, `qwen-image-Q4_K_M.gguf`, `flux1-schnell-Q4_K_M.gguf`.
|
||||
- Command run: `python scripts/benchall_orchestrator.py --list`; output showed `44 configs`, including `z-image-turbo (5)`, `z-image (5)`, `qwen-image-2512 (5)`, `qwen-image (6)`, `flux.1-schnell (4)`, `flux.1-dev (11)`, plus klein configs.
|
||||
- Command run: dry-run `CUDA_VISIBLE_DEVICES=5 python -u scripts/benchall_orchestrator.py ...`; exit `0`; all `z-image-turbo` configs passed: `G0_eager`, `G1_default`, `D1_fp8`, `D2_int8`, `G1_max`.
|
||||
- Command run: predownload all assets; exit `0` overall, but logged `[base ERR] black-forest-labs/FLUX.2-klein-9B: GatedRepoError: 403 Client Error`; other 7 models downloaded.
|
||||
- Command run: metrics schema check on `outputs/diffusion_benchall/_dryrun/z-image-turbo/G1_default/metrics.json`; verified keys `env`, `load`, `generate`, `config`, `accuracy`.
|
||||
- Command run: launched full sweep across GPUs `4`, `5`, `6`, `7` using `scripts/benchall_orchestrator.py`; background task completed exit `0`.
|
||||
- Command run: first progress check showed failures for `flux.1-dev/D2_int8`, `qwen-image/D2_int8`, `qwen-image-2512/D2_int8`; later aggregation found `33` completed metrics and `7` failed configs.
|
||||
- Failed configs: `flux.1-dev/D2_int8`, `flux.1-dev/D_mxfp8`, `flux.1-dev/G1_teq_fp8`, `qwen-image/D2_int8`, `flux.1-schnell/D2_int8`, `qwen-image-2512/D2_int8`, `z-image/G2_fbcache`.
|
||||
- Error root cause: dense int8 on Flux/Qwen hits `torch._int_mm: self.size(0) needs to be greater than 16, but got 1`; small-M conditioning embedders should be excluded from int8.
|
||||
- Error root cause: `mxfp8` failed with `CUBLAS_STATUS_NOT_SUPPORTED` on B200/cuBLAS stack.
|
||||
- Error root cause: Flux text encoder fp8 failed with `normal_kernel_cuda not implemented for Float8_e4m3fn`.
|
||||
- Finding: FBCache on GGUF is slower despite compile engaging; log showed `diffusion.cache: fbcache engaged (threshold=0.08)` and `speed_optims` included `compiled`; likely due to `fullgraph=False` plus GGUF per-op dequant graph breaks.
|
||||
- Final measured headline: GGUF compile+cuDNN stack gives `1.9x-3.55x`; dense fp8 gives `2.3x-5.0x`; GGUF lowest VRAM, dense fp8 fastest.
|
||||
- Example final table values: `Z-Image-Turbo` best `3.58x`; `Z-Image` best `3.87x`; `Qwen-Image-2512` best `4.89x`; `Qwen-Image` best `4.97x`; `FLUX.1-schnell` best `3.32x`; `FLUX.1-dev` best `4.47x`; `FLUX.2-klein-4B` best `2.31x`; `FLUX.2-klein-9B` blocked.
|
||||
- Task `#66` marked completed.
|
||||
- COMPLETED: plan approved, benchmark harness patched, orchestrator created, dry run passed, assets predownloaded where accessible, 7-model full sweep completed, results aggregated, summary written.
|
||||
- PENDING: optional follow-up fix recommended: gate FBCache to dense-only/off for GGUF; optional fix for dense int8 should exclude small-M embedders; accept HF license for `black-forest-labs/FLUX.2-klein-9B` before rerunning that model.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- PENDING: fp8/int8 UI; auto-resize 16
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
- Decision: Launched 6 parallel research agents to avoid duplicated work and cover shipped optimizations plus five lever clusters: caching, attention/token merging, peak-memory reduction, `torch.compile`/Inductor, and quant/kernel/cross-platform.
|
||||
- Decision: Prioritized attention-backend selection after research and B200 probing showed it was the only immediately validated speed win: `_native_cudnn` gave `1.18x` end-to-end with `LPIPS=0.004448794759809971`; Sage/Flash/FBCache were not usable on the available Z-Image setup.
|
||||
- Decision: Split Inductor flags out of the main implementation because measured gain was negligible on B200/bf16: `inductor_flags 0.681s (1.01x vs base) peak=40.9G LPIPS=0.0047416952438652515`.
|
||||
- Decision: Rejected `flash_4_hub` for now because installing `kernels` upgraded `huggingface-hub` to `1.21.0`, breaking diffusers `0.38.0` which needs `<1.0`; env was restored to `huggingface-hub==0.36.2`.
|
||||
- Decision: Rejected SageAttention for B200 because `sageattention` v1.0.6 lacked usable Blackwell/`sm_100` kernels: `Sage Attention backend 'sage' is not usable`.
|
||||
- Decision: Implemented Phase 11 consumer-aware quant ladder because consumer/workstation GPUs have nerfed FP8 FP32-accumulate while INT8 is full-rate; kept FP8 first on data-center GPUs.
|
||||
- Decision: Rejected VAE tiling as Phase 12 for Z-Image inference peak memory after measuring worse/negligible peaks; denoise transformer activations dominate, not VAE decode.
|
||||
- Decision: Deferred FBCache shipping because it fails for Z-Image even eager with `ValueError: Parameter 'hidden_states' not found in function signature but was requested.` and `KeyError: 'hidden_states'`; it needs validation on many-step Flux/Qwen.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/quant_research/03_NEXT_LEVERS.md`: synthesis of next levers and later updated with negative VAE/FBCache findings.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/perf_levers_probe.py`: B200 probe for baseline, Inductor flags, attention backends, and FBCache.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_attention.py`: normalizes attention backend requests, selects `auto -> _native_cudnn` on NVIDIA CUDA, supports opt-in `sage`/`flash`/etc., applies via `set_attention_backend`, gracefully falls back to native.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion.py`: threaded `attention_backend` through load flow, applied attention backend before speed compile, added `_LoadState.attention_backend`, and exposed it in status.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/models/inference.py`: added request and status model fields for `attention_backend`.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/routes/inference.py`: forwards `attention_backend` from API request to backend load call.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_attention.py`: hermetic tests for backend normalization/selection/application/fallback.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_routes.py`: added route coverage for threading and invalid enum handling.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_transformer_quant.py`: added consumer-aware `auto` reorder so consumer GPUs prefer `int8`; B200/data-center keeps `fp8` first.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_transformer_quant.py`: updated torch stub with device names, added 3 consumer ladder tests, fixed data-center names.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/temp/phase10_pr_body.md` and `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/temp/phase11_pr_body.md` for PR bodies.
|
||||
- Command: checked diffusers APIs; key output: `diffusers 0.38.0 torch 2.9.1+cu128`, `ZImageTransformer2DModel has set_attention_backend: True`, cache configs present, Inductor flags present.
|
||||
- Command: listed attention backends; key output included `flash`, `flash_4_hub`, `aiter`, `flex`, `native`, `_native_cudnn`, `sage`, `sage_hub`, `sage_varlen`, `xformers`.
|
||||
- Command: `uv pip install sageattention kernels`; caused `huggingface-hub` incompatibility.
|
||||
- Command: restored env with `uv pip uninstall kernels` and `uv pip install "huggingface-hub>=0.34.0,<1.0"`; verified `hub 0.36.2 diffusers 0.38.0`, `ZImage import OK`, `sageattention OK`.
|
||||
- Command: ran B200 perf probe; final key output: `baseline 0.686s peak=22.6G`, `attn_cudnn 0.584s (1.18x vs base) peak=40.6G LPIPS=0.004448794759809971`.
|
||||
- Command: created branch `diffusion-phase10-attention`; verified `has set_attention_backend: True`, `has reset_attention_backend: True`, `hip: None`.
|
||||
- Error resolved: initial tests failed because test expected `"flash-3"` alias; actual normalization did not map it. Fixed test to use valid case.
|
||||
- Command: `python -m pytest tests/ -q -k "diffusion"` after Phase 10; result `213 passed, 4993 deselected, 1 warning`.
|
||||
- Command: B200 smoke for product attention functions; key output: `select(auto, speed_active=True) -> _native_cudnn`, `apply -> _native_cudnn`, `image finite: True shape: (1024, 1024, 3)`, `ATTN-SMOKE-OK`, `=== EXIT 0 ===`.
|
||||
- Commit: `b92367554 Studio diffusion (Phase 10): attention-backend selection`.
|
||||
- Command: pushed `diffusion-phase10-attention` and created PR `https://github.com/unslothai/unsloth/pull/6701` stacked on `diffusion-phase9-prequant`.
|
||||
- Command: pre-commit.ci moved Phase 10 remote; rebased successfully; sync output `0 0`.
|
||||
- Command: re-ran tests for PR #6701; result `213 passed, 4993 deselected, 1 warning in 15.14s`; PR state `{"base":"diffusion-phase9-prequant","mergeable":"MERGEABLE","number":6701,"state":"OPEN"}`.
|
||||
- Command: created branch `diffusion-phase11-consumer-int8`.
|
||||
- Command: Phase 11 tests: `python -m pytest tests/ -q -k "diffusion"`; result `216 passed` (tail showed `216 p`).
|
||||
- Command: B200 non-regression for quant ladder; key output: `device: NVIDIA B200`, `_is_consumer_gpu('cuda'): False`, `reorder (B200): ('fp8', 'nvfp4', 'mxfp8', 'int8')`, `auto scheme on B200: fp8`.
|
||||
- Commit: `764a3e1dc Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder`.
|
||||
- Command: pushed `diffusion-phase11-consumer-int8` and created PR `https://github.com/unslothai/unsloth/pull/6702` stacked on `diffusion-phase10-attention`.
|
||||
- Command: Phase 11 pre-commit wait showed `no movement`, sync `0 0`.
|
||||
- Command: PR state check: #6701 `MERGEABLE`/`OPEN`, #6702 `MERGEABLE`/`OPEN`.
|
||||
- Command: measured VAE decode peak with pipeline-level tiling; errors: `'ZImagePipeline' object has no attribute 'enable_vae_tiling'` and `'ZImagePipeline' object has no attribute 'enable_vae_slicing'`.
|
||||
- Command: checked Z-Image VAE methods; key output: `VAE class: AutoencoderKL`, `vae.enable_tiling: True`, `vae.enable_slicing: True`, `pipe.enable_vae_tiling: False`, `pipe.vae_scale_factor: 8`.
|
||||
- Command: measured VAE-level tiling; key output: `res 1024: no-tile peak=23.3G | vae-tile peak=30.6G | saved=-7.3G (-31%) | LPIPS=0.0000`; `res 1536: no-tile peak=33.8G | vae-tile peak=37.9G | saved=-4.1G (-12%) | LPIPS=0.0037`; `res 2048: no-tile peak=45.6G | vae-tile peak=45.3G | saved=0.3G (1%) | LPIPS=0.0037`; `=== EXIT 0 ===`.
|
||||
- Command: probed FBCache compatibility; key output: `[cache_eager] FAIL ValueError: Parameter 'hidden_states' not found in function signature but was requested.`, `[cache+regional_nofullgraph] FAIL KeyError: 'hidden_states'`, `[cache+fullcompile_nofullgraph] FAIL KeyError: 'hidden_states'`, `FBCACHE-COMPAT-DONE`, `=== EXIT 0 ===`.
|
||||
- Command still in flight at end: prefetch Flux.1-dev for FBCache validation, background ID `bbujpwspc`, log path pattern `logs/flux_prefetch_$(date +%Y%m%d_%H%M%S).log`.
|
||||
- Completed: Phase 10 PR #6701 is open/mergeable, tests passing, B200 smoke passing.
|
||||
- Completed: Phase 11 PR #6702 is open/mergeable, tests passing, B200 non-regression verified.
|
||||
- Completed: VAE tiling investigated and rejected for Z-Image inference peak.
|
||||
- Completed: FBCache investigated on Z-Image and found incompatible due to `hidden_states` signature.
|
||||
- Pending: Wait for Flux.1-dev prefetch, then validate FBCache on compatible many-step model and potentially ship gated Phase 12 for Flux/Qwen only.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- #156 done: FLUX.2 reference; 99 pass
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
- User directed to expand scope to every architecture: `Yes do all arches.`
|
||||
- User also requested validation benchmarks across all architectures: `Then also do a benchmark to check all.`
|
||||
- User specifically requested benchmarking the alternative buffer-based implementation: `Also benchmark the buffers approach`
|
||||
- COMPLETED: captured new requirements only.
|
||||
- PENDING: implement/test all architectures; run comparative benchmarks including buffer approach.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- `upscale` added+verified; hold for push/mirror
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
- Decision: implement faithful `addcmul` fusions for all four Flux block classes (`FluxTransformerBlock`, `FluxSingleTransformerBlock`, `Flux2TransformerBlock`, `Flux2SingleTransformerBlock`) using exact `diffusers` 0.38 forward bodies to avoid semantic drift.
|
||||
- Decision: run two follow-up tasks: `Task #112` for Flux.1/Flux.2 patch implementation with body-drift guards/specs/allclose tests, and `Task #113` for isolation benchmarks across all four arch families plus focused global buffer on/off benchmark.
|
||||
- Commands run from `/mnt/disks/unslothai/ubuntu/workspace_81`: `CUDA_VISIBLE_DEVICES="" UNSLOTH_ALLOW_CPU=1 python - <<'PY' 2>/dev/null` to inspect `transformer_flux2` and `transformer_flux` sources.
|
||||
- Key outputs: fetched `Flux2TransformerBlock.forward`, `Flux2SingleTransformerBlock.forward`, `FluxTransformerBlock.forward`, `FluxSingleTransformerBlock.forward`, relevant `__init__` signatures, and `Flux2Modulation` source.
|
||||
- Finding: `Flux2Modulation.split` returns modulation tensors shaped `[B,1,dim]`, so test inputs can rely on clean broadcasting.
|
||||
- Error: `API Error: Connection closed mid-response. The response above may be incomplete.` Work was interrupted before edits.
|
||||
- Completed: source inspection and task creation.
|
||||
- Pending: edit `diffusion_arch_patches.py`, add specs/guards/tests, run allclose checks and benchmarks.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Backend PR split blocked; plan updated.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Done: default=compile-dequant; max=compile
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- Done: right-sidebar move #25, Advanced discoverability, `/16` auto-resize verified.
|
||||
- Pending: push, backend path, mirror, cancel loop.
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the user’s directive to make a PR or push, scoped as “Complete PR, new branch” per the answered scoping question. The last completed tool was the full diffusion pre-commit pytest run, logging to `/mnt/disks/unslothai/ubuntu/workspace_81/logs/pretest_230339.log`; do not rerun it. Next, inspect the test result/log as needed, then proceed with the planned commit/push flow: push the Phase-16 base branch to `oobabooga`, create the new complete PR branch, commit only the intended frontend/installer/backend changes while excluding logs/temp/plans/junk, push it, and open the PR against that base.
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
- User requested benchmarking performance gains across all PRs for every supported model format: `GGUFs` and `safetensors`.
|
||||
- Verbatim directive to preserve: `And for all odels that are supported - GGUFs or just safetensors - across all PRs - what are the performance gains - benchmark all`
|
||||
- Decision implied: scope expands from isolated PR review to cross-PR, cross-model benchmark comparison.
|
||||
- Files created/edited: none in this span.
|
||||
- Commands run: none in this span.
|
||||
- Errors encountered: none; note typo `odels` appears in user directive.
|
||||
- COMPLETED: request captured.
|
||||
- PENDING: identify supported models/formats, enumerate relevant PRs, run benchmarks, report performance gains.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Done: PR synced; no pending
|
||||
|
|
@ -0,0 +1 @@
|
|||
- PR `#6694` opened; tests pass.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked check `https://github.com/city96/ComfyUI-GGUF`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Faster inference; `5 Opus subagents`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked: `non GGUFs`, `fp8, int8, bf16`
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- User asked: `Also for https://github.com/unslothai/unsloth/pull/6658 - which models are supported?`
|
||||
- PENDING: inspect PR `6658` and determine supported models.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- `FLUX.2-klein` inpaint done; extend excluded
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- Done: HF data
|
||||
- Pending: PR push
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Phase12 done: PR `#6703`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Rebased, pushed PR `#6675`; 127 pass
|
||||
|
|
@ -0,0 +1 @@
|
|||
- quant built; 146 pass; GPU pending
|
||||
|
|
@ -0,0 +1 @@
|
|||
- No actions; context-summary request only
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
- Decision: implemented all remaining diffusion arch addcmul patches because tests/benchmarks required all four families covered and bit-correct: `qwen-image`, `z-image`, `flux.1`, `flux.2-klein`.
|
||||
- Decision: kept source/signature drift guards around arch patches because patch safety depends on matching the expected diffusers forward bodies.
|
||||
- Decision: treated initial `flux.1-schnell` benchmark `-20.5%` as invalid because logs showed `offload=model` and latency was ~10x normal; reran contaminated cases sequentially on free GPU 5 resident.
|
||||
- Decision: concluded global weight buffer is neutral because qwen was `+0.2%`, flux.1 repeated from `+11.0%` to `-2.4%`, and allocator profile had no allocation retries/steady reserved VRAM.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_arch_patches.py`: added all family/block patches and updated module docstring to state all arches are implemented.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_arch_patches.py`: added allclose tests for four flux blocks and updated install-count assertion.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/buffer_bench.py`: dedicated global-buffer on/off benchmark.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/arch_patch/SUMMARY.md`: final benchmark summary and interpretation.
|
||||
- Ran `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && CUDA_VISIBLE_DEVICES="" python -m pytest tests/test_diffusion_arch_patches.py -q 2>&1 | tail -16`; result: `.......... [100%]`, 10 arch tests passed with warning.
|
||||
- Ran `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && CUDA_VISIBLE_DEVICES="" python -m pytest tests/ -q -k "diffusion" 2>&1 | tail -4`; result: `300 passed, 1 skipped, 5012 deselected, 1 warning in 13.25s`.
|
||||
- Launched benchmarks: `scripts/arch_patch_bench.py --model flux.1-schnell --gpu 4`, `--model flux.2-klein-4b --gpu 5`, `scripts/buffer_bench.py --model qwen-image --gpu 6`, `--model flux.1-schnell --gpu 7`.
|
||||
- Error: result collection shell test produced `/bin/bash: line 23: [: 0 0: integer expression expected`; non-blocking, later results were collected another way.
|
||||
- Error: GPUs 4/6/7 were occupied by external `python3` processes using ~158GB each; caused offload contamination. Did not kill external PIDs.
|
||||
- Ran `kill 325007` to stop own offloaded qwen buffer bench; output `killed my qwen buffer bench (offloaded)`.
|
||||
- Ran clean sequential re-runs on GPU 5: `scripts/arch_patch_bench.py --model flux.1-schnell --gpu 5`, `--model qwen-image --gpu 5`, `scripts/buffer_bench.py --model flux.1-schnell --gpu 5`, `--model qwen-image --gpu 5`; all 8 loads reported `offload=none`.
|
||||
- Final resident arch benchmark results: `z-image-turbo off=1.554s on=1.557s -0.2% PSNR=32.1`; `flux.2-klein-4b off=0.626s on=0.623s +0.4% PSNR=33.6`; `flux.1-schnell off=1.417s on=1.405s +0.8% PSNR=30.2`; `qwen-image off=9.508s on=9.667s -1.6% PSNR=45.2`.
|
||||
- Final resident buffer benchmark results: `qwen-image off=9.477s on=9.462s +0.2%`; `flux.1-schnell off=1.569/on=1.413 +11.0%` then repeat `off=1.415s on=1.450s -2.4% VRAM 18.8->18.9 PSNR=inf`.
|
||||
- Error: repeat flux.1 buffer pretty-print wrapper had quote-nesting bug and printed inline Python source; underlying `buffer_bench` still wrote `outputs/buffer_bench/flux.1-schnell/result.json`, so resolved by reading result JSON.
|
||||
- Ran final verification: `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && CUDA_VISIBLE_DEVICES="" python -m pytest tests/ -q -k "diffusion" 2>&1 | tail -3`; result: `300 passed, 1 skipped, 5012 deselected, 1 warning in 14.20s`.
|
||||
- Task status: completed task `#112` and `#113`.
|
||||
- COMPLETED: all arch patches implemented, flux allclose tests added, install-count assertion updated, buffer benchmark script added, clean benchmarks collected, summary written, final diffusion test suite green.
|
||||
- PENDING: none noted.
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
- Decision: continue active Stop-hook work on diffusion efficiency/memory without pausing, per directive `Keep working on making diffusion more efficent, use less memory and not reduce accuracy...`.
|
||||
- Must use cwd only: `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/temp/temporary_qui4gujp`; never exit it.
|
||||
- Goal constraints: optimize diffusion load/inference peak memory and speed across Mac/Windows/Linux, NVIDIA/AMD/Intel/CPU; support GGUF and safetensors; avoid large accuracy loss, `25% reduction is fine`.
|
||||
- Research targets preserved: `torch._int_mm`, `torchao/sparsity`, SDNQ, int8 fused, FBGEMM, TorchInductor config, Diffusers optimization docs, FP8/NVFP4/MXFP8, TensorRT, FlashAttention diffusion.
|
||||
- Existing PR context to consider: `#6658 #6670 #6675 #6679 #6680 #6690 #6694 #6700 #5872`.
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the latest benchmark/fix task from the completed state: the int8 dense quant crash investigation, fix, verification, and stacked PR are already done. Do not rerun the completed background benchmarks, pytest runs, branch push, or PR creation; PR #6716 exists, includes the pre-commit.ci autofix, and quant tests passed after syncing. Next tool action is to proceed only with the next separate benchmark failure if the broader effort continues, starting by inspecting the relevant code/logs for the T5 text-encoder fp8 issue rather than touching the completed int8 PR work.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- NVFP4 Phase2C done; pushed `dbb029256`
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
- User requested testing safetensors inplace ops and `global` buffers.
|
||||
- Desired defaults: `channels_last`, compiled dequant only, cuDNN if supported, inplace/eager opts.
|
||||
- Max adds compile all + tf32.
|
||||
- Check `https://github.com/unslothai/unsloth/blob/main/unsloth/models/llama.py#L1236`-`L1354`.
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- Advanced moved to right card, open default; Chat-like toggle.
|
||||
- `images-page.tsx`; `tsc`, build ok.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Found/fixed `is_gguf` bug in `core/inference/diffusion.py`; non-GGUF fp8/int8 now route to regional `compiled`.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Check NVFP4; test FP8/NVFP4 lowmem
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the latest directive to keep improving diffusion efficiency, memory use, and speed without unacceptable accuracy loss, specifically building on the validated Flux.1-dev FBCache Phase 12 work. The next concrete action is to continue wiring FBCache into the backend after the completed edit to `studio/backend/models/inference.py`: update the route/request plumbing as needed, finish status model exposure, then validate the new cache parameters and compile `fullgraph=False` behavior without re-running already completed Flux validation tools.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked `NXFP4` via PyTorch/TorchAO
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
- Decision/request: make `Advanced` closed by default, because user wants it not open initially.
|
||||
- Decision/request: when `Advanced` is expanded, keep the Advanced icon stationary; it should look like `Chat`, so the icon should not move.
|
||||
- Reference URL provided verbatim: `https://huggingface.co/datasets/danielhanchen/screenshots/discussions/26`
|
||||
- COMPLETED: no implementation in this span.
|
||||
- PENDING: apply the UI behavior/icon alignment changes.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- `https://huggingface.co/datasets/danielhanchen/screenshots/discussions/25`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Load Studio `--secure`; ask sd.cpp
|
||||
|
|
@ -0,0 +1,6 @@
|
|||
- Decision: next lever should target inference-time peak memory/speed, because Phases 7-9 already covered compile, dense+torchao quant, consumer fast-accum gating, and pre-quantized loading/load VRAM.
|
||||
- Goal preserved: push diffusion speed and `peak memory (both load and inference)` under ~25% accuracy cost across consumer/data-center GPUs and ideally Mac/AMD/Intel/CPU, for `GGUF + safetensors`.
|
||||
- Mentioned candidate levers: step caching, attention backends, VAE tiling for decode peak memory.
|
||||
- No files edited/created.
|
||||
- No commands completed.
|
||||
- Error: `API Error: Connection closed mid-response.` unresolved; assistant response may be incomplete.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- PENDING: Z-Image, LPIPS, CPU cmp
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Waiting user decision.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Studio launched: `--secure`, `:8890`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- COMPLETED; pending user go-ahead
|
||||
|
|
@ -0,0 +1 @@
|
|||
- PENDING: install `torch 2.11` `NVFP4`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Pending: run `gh pr checks 6703`; use `CronList`/`CronDelete` only if checks done.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Studio up: `HTTP 200`; URL `https://courtesy-complicated-newman-citizens.trycloudflare.com`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- fp8 auto; PR `6694` merged-ready.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Bench: 2.56x faster `1.83 s -> 0.71 s`; stack: compile+cudnn attention+channels_last+benchmark. FBCache extra for many-step only.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Complete; pending `push the PRs`/`set up the mirror`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Decided fp8 fastest; GGUF smallest/default.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- PR-ready; pending `push the PRs`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- No API errors; only benign ONNX warning
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the latest directive: remove all global buffers since they looked useless, then ensure nothing else broke. The last completed action read the affected section of tests/test_diffusion_speed.py, so the next action is to edit that file to remove the weight_buffer stub/assertions without re-reading or rerunning the completed Read.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- PENDING: user decide commit scope
|
||||
|
|
@ -0,0 +1 @@
|
|||
- compile gives most win; conv opts ~noop
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- Phase13 TE offload rejected: `+686%` latency for `-11%` peak; code reverted.
|
||||
- `#6703` green: `42 pass`, `1 skipping`.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Diffusion: fp8 fastest/`auto`; caveats VRAM/shape
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the latest directive: remove all global buffer machinery because the benchmarks showed it was useless, and ensure nothing else broke. The production module, diffusion comments, speed docstrings, and `test_diffusion_gguf_compile.py` were already edited; the last completed tool read `studio/backend/tests/test_diffusion_speed.py` and showed the remaining `weight_buffer` test stub/assertions. Next action is to edit `test_diffusion_speed.py` to remove `weight_buffer` from the stub and expectations, then grep for remaining buffer references and run the relevant diffusion tests.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Pending: needs user go-ahead.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- NVFP4 tried; Phase4 started
|
||||
|
|
@ -0,0 +1 @@
|
|||
- User asked: `So FP8 is the fastest?`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Decision: keep `default dynamic=True`.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Multi-ref FLUX.2 done; verified live.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Use 5 Opus subagents; recheck links
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked `How about memory usage and size`
|
||||
|
|
@ -0,0 +1 @@
|
|||
- `dynamic=True` default; global buffers
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Phase 7 pushed: `ede94176f`, green.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- User asks benchmark/settings recap
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked: `Search HuggingFace for any NVFP4 diffusion models and see if we can try them`
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
- Decision: discussion `#26` changes are complete because the UI is now closed by default and the Chat-style toggle is fixed.
|
||||
- Completed: verified the discussion `#26` behavior; no further autonomous action was identified in this span.
|
||||
- Pending: waiting for user go-ahead on remaining choices: `push the frontend + installer`, `backend path 1/2/3`, `set up the mirror`, `cancel the loop`, or continued UI feedback iteration.
|
||||
- Live studio remains available at `https://incentive-topics-patrick-commissioners.trycloudflare.com` with credentials `unsloth` / `diffusion-Thrg7s3tr1-UoGYq`.
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- fixed UI+resize
|
||||
- pending push/backend
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue making diffusion more efficient with lower loading and inference peak memory, faster execution, and acceptable accuracy loss only within the user’s stated tolerance. The latest edit to inference.py is done; do not re-run it. Next, proceed from the completed status-response-model edit by inspecting the updated inference.py context only as needed, then continue implementing or validating the remaining quantization/memory/speed changes implied by the stop-hook feedback.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked diffusers `int8 / float8`
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
- Decision: Used cached/local context first per `AGENTS.md`; no cache existed for PR `#6658`, so fetched live PR metadata with `gh`.
|
||||
- Decision: Summarized PR `#6658` as foundational diffusion/image-generation support; detection is architecture-based so one family entry covers quants/variants.
|
||||
- Command run: `cd /mnt/disks/unslothai/ubuntu/workspace_81 && echo "=== cached summary? ===" && ls -la data/pr_summaries/unslothai__unsloth/pr_6658.json 2>/dev/null || echo "(no cache)" && echo "=== PR meta ===" && gh pr view 6658 --repo unslothai/unsloth --json number,title,state,author,body ...`; key output: `(no cache)`, `#6658 [OPEN] Studio: add image generation with Unsloth GGUFs`, author `oobabooga`.
|
||||
- Completed: Captured PR body summary: new Image tab with download menu, generation inputs, output image/gallery.
|
||||
- Completed: Identified integrated families: `z-image`, `qwen-image`, `flux.1`, `flux.2-klein`; all Unsloth GGUF transformer-only with base diffusers components resolved via `base_model`.
|
||||
- Completed: Noted excluded models/reasons: `ERNIE-Image`, `FLUX.2-dev`, editing models.
|
||||
- Pending: none in this span.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Done; pending user choice
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
- Decision: discussion #26 fixed by using one always-present top-bar `LayoutAlignRightIcon`; guarantees no icon swap or movement and matches Chat style.
|
||||
- Decision: Advanced panel now closed by default so preview canvas loads full width.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend/src/features/images/images-page.tsx`: default Advanced state closed; replaced conditional gear/header close icons with a single fixed top-right `LayoutAlignRightIcon` toggle; removed redundant panel header close button.
|
||||
- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/diffusion-workflows-pr-plan.md`: updated plan note for discussion #26.
|
||||
- Ran `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend && npx tsc --noEmit ... && npm run build`; exit 0, key output `=== TSC OK ===`, `build exit 0 -> /mnt/disks/unslothai/ubuntu/workspace_81/logs/fe_build_222352.log`.
|
||||
- Ran Playwright script `/mnt/disks/unslothai/ubuntu/workspace_81/scripts/shoot_toggle.py`; exit 0, saved `toggle_closed.png` and `toggle_open.png`; output `ICON MOVED: False`, both boxes `{x:1558, y:11, width:34, height:34}`.
|
||||
- Verified visually via `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/toggle_open.png` and `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/toggle_closed.png`.
|
||||
- Completed: discussion #26 changes and verification. Pending: local frontend changes for PR 2 remain uncommitted.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Asked compile flags/CUDAGraphs
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Done: adv open; `/16` resize
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Done/no errors; pending go-ahead
|
||||
|
|
@ -0,0 +1 @@
|
|||
- FP8/INT8 PASS; all `/goal` done
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
- Decided pre-quantized torchao loading fixes fast-mode load VRAM; offline quant avoids dense bf16 GPU materialization.
|
||||
- Created PR `https://github.com/unslothai/unsloth/pull/6700`, base `diffusion-phase8-quant`, branch `diffusion-phase9-prequant`, commit `b90f83346`.
|
||||
- Added `diffusion_prequant.py`, builder/verify scripts, request field/tests.
|
||||
- Verified: load peak `21.2 GB` -> `14.6 GB`, disk `6.28 GB`, `LPIPS 0.0`; `200` diffusion tests passed.
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- User asked: `Is there a way to fix the loading VRAM issue - pre-quantized?`
|
||||
- Pending: answer likely needs repo/PR context; no commands or edits in this span.
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
- Approved plan saved at `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/wobbly-jumping-narwhal.md`.
|
||||
- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_patch_backend.py`; lazy wraps `unsloth_zoo.temporary_patches.utils` `patch_function`/`can_safely_patch`/`restore_original`.
|
||||
- Edited `diffusion_eager_patches.py`, `diffusion_gguf_compile.py`, `diffusion.py`, `tests/conftest.py`, `test_diffusion_gguf_compile.py`; CPU tests require `UNSLOTH_ALLOW_CPU=1`.
|
||||
- Created `diffusion_arch_patches.py` + `test_diffusion_arch_patches.py`; qwen `_modulate` and z-image gated residuals use compile-safe `torch.addcmul`.
|
||||
- Created `scripts/arch_patch_bench.py`.
|
||||
- Tests: `77 passed, 1 skipped`; full diffusion: `296 passed, 1 skipped`.
|
||||
- Bench: qwen parity default `5.952s`, max `2.741s`; arch isolation z-image `-0.2%`, qwen `+0.3%`.
|
||||
- Alloc profile: `alloc_retries delta == 0`, reserved VRAM flat; global scratch buffers not useful.
|
||||
- Pending: user asked whether to extend same addcmul recipe to flux.1/flux.2.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Done: summarized `outputs/` benches
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
- DONE: build verified
|
||||
- PENDING: choice
|
||||
|
|
@ -0,0 +1 @@
|
|||
- User asked: `So overall how much did increase the speed by?`
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the latest directive: examine why int8 failed for some cases and fix if possible. The latest assistant state says this was completed: int8 failures were root-caused, fixed, verified on B200, tests passed, and PR #6716 was shipped stacked on #6703. No tool is in flight; the next action is to carry forward that completed status and, if the conversation continues, answer only any new user follow-up such as tackling the separate T5 fp8 issue.
|
||||
|
|
@ -0,0 +1 @@
|
|||
Continue the latest directive: “make a PR or push to another,” meaning push the completed diffusion image workflow work to the user’s fork and open the appropriate stacked PR. The full diffusion test run already completed, so do not re-run it just because of compaction; inspect its result/log if needed, then proceed with committing the work, pushing the base/work branches as described, and opening the PR against the fork branch.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- Explained diffusion models/int8 UI gap.
|
||||
|
|
@ -0,0 +1 @@
|
|||
- NVFP4 slow; added `scripts/nvfp4_probe.py`
|
||||
|
|
@ -29,8 +29,10 @@ from loggers import get_logger
|
|||
from utils.hardware import clear_gpu_cache
|
||||
|
||||
from .diffusion_families import (
|
||||
DIFFUSION_CANCELLED_MSG,
|
||||
DIFFUSION_NOT_LOADED_MSG,
|
||||
DiffusionFamily,
|
||||
detect_family,
|
||||
detect_family_for_pick,
|
||||
resolve_base_repo,
|
||||
resolve_local_gguf_child,
|
||||
supported_family_names,
|
||||
|
|
@ -43,11 +45,10 @@ from .diffusion_device import (
|
|||
from .diffusion_memory import (
|
||||
OFFLOAD_NONE,
|
||||
apply_memory_plan,
|
||||
estimate_gguf_dense_mib,
|
||||
estimate_gguf_resident_mib,
|
||||
estimate_image_runtime_mib,
|
||||
estimate_safetensors_dense_mib,
|
||||
file_size_mib,
|
||||
infer_gguf_quant_label,
|
||||
plan_diffusion_memory,
|
||||
snapshot_device_memory,
|
||||
)
|
||||
|
|
@ -383,22 +384,6 @@ class DiffusionBackend:
|
|||
base, rfilename, hf_token, cancel_event = self._cancel_event
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _detect_family_for_pick(
|
||||
repo_id: str, gguf_filename: Optional[str], family_override: Optional[str]
|
||||
) -> Optional[DiffusionFamily]:
|
||||
"""Detect the family from the repo id, falling back to the combined
|
||||
path/filename for a direct local .gguf pick. The frontend splits such a
|
||||
pick into (parent dir, basename), so the family keyword can live only in
|
||||
the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent
|
||||
directory carries none; scan it too when the directory alone is
|
||||
undetectable. Only used as a fallback, so remote 'org/name' picks and
|
||||
explicit overrides behave exactly as before."""
|
||||
fam = detect_family(repo_id, family_override)
|
||||
if fam is None and gguf_filename and not family_override:
|
||||
fam = detect_family(f"{repo_id}/{gguf_filename}", family_override)
|
||||
return fam
|
||||
|
||||
def validate_load_request(
|
||||
self,
|
||||
repo_id: str,
|
||||
|
|
@ -414,7 +399,7 @@ class DiffusionBackend:
|
|||
undetectable family, and ValueError/FileNotFoundError for a bad local path.
|
||||
Touches no GPU, network, or state."""
|
||||
kind = resolve_model_kind(gguf_filename, model_kind)
|
||||
fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
if fam is None:
|
||||
raise ValueError(
|
||||
f"'{repo_id}' is not a supported diffusion image model. Supported families: "
|
||||
|
|
@ -537,7 +522,7 @@ class DiffusionBackend:
|
|||
# Resolve the base repo and estimate sizes on this thread (both network
|
||||
# calls) so begin_load returns instantly; the bar shows raw bytes until
|
||||
# the total lands. This is the only writer of _loading's fields here.
|
||||
fam = self._detect_family_for_pick(
|
||||
fam = detect_family_for_pick(
|
||||
kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
|
||||
)
|
||||
kind = resolve_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind"))
|
||||
|
|
@ -753,7 +738,6 @@ class DiffusionBackend:
|
|||
plan = self._plan_memory(
|
||||
target,
|
||||
single_file_path,
|
||||
gguf_filename,
|
||||
base,
|
||||
fam,
|
||||
memory_mode,
|
||||
|
|
@ -936,7 +920,11 @@ class DiffusionBackend:
|
|||
quant = transformer_quant_engaged,
|
||||
attention_backend = attention_engaged,
|
||||
compile_kwargs = {
|
||||
"fullgraph": cache_engaged is None,
|
||||
# Mirrors apply_speed_optims' fullgraph decision: an active
|
||||
# step cache OR a planned offload graph-breaks, so the cached
|
||||
# bundle must be keyed on the same fullgraph setting.
|
||||
"fullgraph": cache_engaged is None
|
||||
and plan.offload_policy == OFFLOAD_NONE,
|
||||
"dynamic": effective_speed != SPEED_MAX,
|
||||
"mode": "max-autotune-no-cudagraphs"
|
||||
if effective_speed == SPEED_MAX
|
||||
|
|
@ -952,8 +940,21 @@ class DiffusionBackend:
|
|||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
cache_active = cache_engaged is not None,
|
||||
# The planned offload policy: group/model/sequential offload installs
|
||||
# compiler-disabled onload hooks, so compile must drop fullgraph.
|
||||
offload_active = plan.offload_policy != OFFLOAD_NONE,
|
||||
logger = logger,
|
||||
)
|
||||
if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
|
||||
# Promotion above could not engage compile (e.g. the family is not
|
||||
# compile-friendly, or compile_repeated_blocks failed): the quantized
|
||||
# transformer is now running eager, which is far slower than the GGUF
|
||||
# path it replaced. Surface it loudly rather than hiding the regression.
|
||||
logger.warning(
|
||||
"diffusion.transformer_quant: %s engaged but the transformer is NOT "
|
||||
"compiled; eager torchao quant is ~30x slower than GGUF here",
|
||||
transformer_quant_engaged,
|
||||
)
|
||||
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
|
||||
# also before placement so the offload hooks move the smaller weights.
|
||||
te_quant = quantize_text_encoders(
|
||||
|
|
@ -1008,6 +1009,9 @@ class DiffusionBackend:
|
|||
if eager_patched:
|
||||
uninstall_patches()
|
||||
uninstall_arch_patches()
|
||||
# Also free the half-built pipe's VRAM: the failed load never
|
||||
# commits _state, so nothing else reclaims it until the next unload.
|
||||
clear_gpu_cache()
|
||||
|
||||
logger.info(
|
||||
"diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",
|
||||
|
|
@ -1106,7 +1110,6 @@ class DiffusionBackend:
|
|||
self,
|
||||
target: DiffusionDeviceTarget,
|
||||
single_file_path: Optional[str],
|
||||
gguf_filename: Optional[str],
|
||||
base: str,
|
||||
fam: DiffusionFamily,
|
||||
memory_mode: Optional[str],
|
||||
|
|
@ -1120,9 +1123,10 @@ class DiffusionBackend:
|
|||
offload policy + VAE memory savers. Kept on the backend so the cached base
|
||||
repo (companion text-encoder / VAE) feeds the size estimate.
|
||||
|
||||
The size estimate is per-kind: a GGUF dequantises (a 4-bit file ~4x), a
|
||||
safetensors single-file loads near its on-disk size, and a full pipeline is
|
||||
one cached download (transformer + companions) that is already compressed."""
|
||||
The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul
|
||||
transient dequant), so a GGUF loads near its on-disk size; a safetensors
|
||||
single-file loads near its on-disk size (it carries its dtype); and a full
|
||||
pipeline is one cached download (transformer + companions), already compressed."""
|
||||
device_memory = snapshot_device_memory(target)
|
||||
if kind == "pipeline":
|
||||
# The whole repo (transformer + companions) is one cached download; the
|
||||
|
|
@ -1134,18 +1138,18 @@ class DiffusionBackend:
|
|||
else:
|
||||
if kind == "single_file":
|
||||
# Safetensors single-file: no dequant expansion (it carries its dtype).
|
||||
transformer_dense = estimate_safetensors_dense_mib(file_size_mib(single_file_path))
|
||||
else:
|
||||
transformer_dense = estimate_gguf_dense_mib(
|
||||
file_size_mib(single_file_path), infer_gguf_quant_label(gguf_filename)
|
||||
transformer_resident = estimate_safetensors_dense_mib(
|
||||
file_size_mib(single_file_path)
|
||||
)
|
||||
else:
|
||||
transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path))
|
||||
# The companion components (VAE + text encoders) load near their on-disk
|
||||
# size; sum whatever the prefetch already placed in the base-repo cache.
|
||||
companion = self._cache_bytes(base)
|
||||
companion_mib = int(companion // (1024 * 1024)) if companion else None
|
||||
model_dense_mib = None
|
||||
if transformer_dense is not None:
|
||||
model_dense_mib = transformer_dense + (companion_mib or 0)
|
||||
if transformer_resident is not None:
|
||||
model_dense_mib = transformer_resident + (companion_mib or 0)
|
||||
runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name)
|
||||
return plan_diffusion_memory(
|
||||
target = target,
|
||||
|
|
@ -1244,7 +1248,7 @@ class DiffusionBackend:
|
|||
with self._lock:
|
||||
state = self._state
|
||||
if state is None:
|
||||
raise RuntimeError("No diffusion model is loaded.")
|
||||
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
|
||||
# Register under _lock so unload()/a load can signal THIS generation.
|
||||
# A cancel that arrived before now either nulled _state (we raised
|
||||
# above) or targets an older generation, so nothing is lost.
|
||||
|
|
@ -1422,7 +1426,7 @@ class DiffusionBackend:
|
|||
# A cancelled denoise returns early with a partial/garbage image;
|
||||
# don't hand it back to be persisted.
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# The first compiled generation just paid the compile cost; persist the
|
||||
# warm torch.compile cache bundle when saving is enabled (distributor /
|
||||
# first-run warm). Idempotent + best-effort -- never fails a generation.
|
||||
|
|
|
|||
|
|
@ -20,6 +20,14 @@ from pathlib import Path, PurePosixPath
|
|||
from typing import Optional
|
||||
|
||||
|
||||
# Runtime->route contract: the RuntimeError messages a backend raises for
|
||||
# client-recoverable generate states. The /images/generate route matches these
|
||||
# EXACTLY to return 409 (vs a sanitized 500 for real failures), so both engines
|
||||
# must raise them verbatim -- keep them named here, not as scattered literals.
|
||||
DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded."
|
||||
DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled."
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class DiffusionFamily:
|
||||
name: str
|
||||
|
|
@ -291,6 +299,24 @@ def supported_family_names() -> tuple[str, ...]:
|
|||
return tuple(fam.name for fam in _FAMILIES)
|
||||
|
||||
|
||||
def detect_family_for_pick(
|
||||
repo_id: str,
|
||||
gguf_filename: Optional[str] = None,
|
||||
override: Optional[str] = None,
|
||||
) -> Optional[DiffusionFamily]:
|
||||
"""``detect_family``, falling back to the combined path/filename for a direct
|
||||
local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename),
|
||||
so the family keyword can live only in the filename (e.g.
|
||||
``/models/z-image-turbo-Q4_K_M.gguf``) while the parent directory carries none;
|
||||
scan the combined string too when the directory alone is undetectable. Only a
|
||||
fallback, so remote ``org/name`` picks and explicit overrides behave exactly as
|
||||
``detect_family``. Shared by both engines so validation and load can't diverge."""
|
||||
fam = detect_family(repo_id, override)
|
||||
if fam is None and gguf_filename and not override:
|
||||
fam = detect_family(f"{repo_id}/{gguf_filename}", override)
|
||||
return fam
|
||||
|
||||
|
||||
def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
|
||||
"""The companion diffusers repo: caller-supplied if given, else the family fallback."""
|
||||
base = (base_repo or "").strip()
|
||||
|
|
|
|||
|
|
@ -218,51 +218,22 @@ def file_size_mib(path: Any) -> Optional[int]:
|
|||
return None
|
||||
|
||||
|
||||
def infer_gguf_quant_label(filename: Optional[str]) -> Optional[str]:
|
||||
"""Pull a quant tag (Q4_K_M, Q8_0, BF16, ...) out of a GGUF filename."""
|
||||
if not filename:
|
||||
return None
|
||||
from pathlib import Path
|
||||
def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]:
|
||||
"""Approximate the RESIDENT device size of a GGUF transformer loaded through
|
||||
diffusers' ``GGUFQuantizationConfig``.
|
||||
|
||||
stem = Path(filename).name
|
||||
if stem.lower().endswith(".gguf"):
|
||||
stem = stem[:-5]
|
||||
parts = [p.upper() for p in stem.replace("-", "_").split("_") if p]
|
||||
for index, part in enumerate(parts):
|
||||
if part in ("BF16", "F16", "FP16", "FP8", "Q8", "Q6", "Q5", "Q4", "Q3", "Q2"):
|
||||
suffix = parts[index + 1 :]
|
||||
# Quant names carry either a K-family suffix (Q4_K_M) or a legacy
|
||||
# numeric one (Q8_0, Q5_1); keep up to two suffix tokens.
|
||||
if suffix and suffix[0] in ("K", "M", "S", "L", "XS", "XXS", "0", "1"):
|
||||
return "_".join([part] + suffix[:2])
|
||||
return part
|
||||
if part.startswith("IQ") or part.startswith("UD"):
|
||||
return "_".join(parts[index : index + 3])
|
||||
return None
|
||||
The weights stay PACKED on the device as quantised bytes (``GGUFParameter`` /
|
||||
uint8); ``GGUFLinear.forward`` dequantises each weight to the bf16 compute dtype
|
||||
transiently for its matmul and frees it immediately, so the persistent footprint
|
||||
is ~= the on-disk tensor size, NOT the unpacked bf16 size. Measured on
|
||||
Z-Image-Turbo: Q2_K 3.64 GiB -> 3.68 GiB, Q8_0 7.22 GiB -> 7.25 GiB resident.
|
||||
The transient per-op dequant is covered by the separate runtime headroom.
|
||||
|
||||
|
||||
def estimate_gguf_dense_mib(storage_mib: Optional[int], quant: Optional[str]) -> Optional[int]:
|
||||
"""Approximate the dequantised (device) size of a GGUF from its on-disk size
|
||||
and quant label. The compute dtype is bf16/fp16, so a 4-bit file roughly
|
||||
quadruples once unpacked; higher-bit quants expand less."""
|
||||
(The prior per-quant expansion assumed a full unpack that never happens on this
|
||||
path; it over-estimated e.g. Q2 ~7.6x, forcing needless offload.)"""
|
||||
if storage_mib is None:
|
||||
return None
|
||||
q = (quant or "").upper()
|
||||
if any(t in q for t in ("BF16", "F16", "FP16")):
|
||||
return storage_mib
|
||||
if "FP8" in q or "Q8" in q:
|
||||
return int(storage_mib * 2.0)
|
||||
if "Q6" in q:
|
||||
return int(storage_mib * 2.8)
|
||||
if "Q5" in q:
|
||||
return int(storage_mib * 3.3)
|
||||
if "Q4" in q or "IQ4" in q or "UD" in q:
|
||||
return int(storage_mib * 4.0)
|
||||
if "Q3" in q or "IQ3" in q:
|
||||
return int(storage_mib * 5.3)
|
||||
if "Q2" in q or "Q1" in q or "IQ2" in q or "IQ1" in q:
|
||||
return int(storage_mib * 8.0)
|
||||
return int(storage_mib * 4.0) # unknown: assume 4-bit-ish
|
||||
return int(storage_mib * 1.05) # small margin for allocator + bf16 norms/biases
|
||||
|
||||
|
||||
def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]:
|
||||
|
|
|
|||
|
|
@ -164,11 +164,17 @@ def apply_speed_optims(
|
|||
family: Any,
|
||||
speed_mode: str = SPEED_OFF,
|
||||
cache_active: bool = False,
|
||||
offload_active: bool = False,
|
||||
logger: Any = None,
|
||||
) -> dict[str, bool]:
|
||||
"""Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline,
|
||||
BEFORE placement / offload. Returns which optimisations actually engaged. Every
|
||||
step is best-effort: a pipeline that doesn't support one is simply skipped."""
|
||||
step is best-effort: a pipeline that doesn't support one is simply skipped.
|
||||
|
||||
``offload_active`` is the planned offload policy != none: group/model/sequential
|
||||
offloading installs ``@torch.compiler.disable``d onload hooks, so the compile must
|
||||
drop ``fullgraph`` (same reason as an active step cache) or it crashes at the first
|
||||
denoise step."""
|
||||
applied = {
|
||||
"channels_last": False,
|
||||
"cudnn_benchmark": False,
|
||||
|
|
@ -178,12 +184,11 @@ def apply_speed_optims(
|
|||
"compiled_dequant": False,
|
||||
}
|
||||
mode = normalize_speed_mode(speed_mode)
|
||||
# TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this
|
||||
# load isn't max, so a later default/off diffusion load -- or chat inference in the
|
||||
# same long-lived process -- doesn't silently inherit a prior max load's TF32 and
|
||||
# lose the bit-identical default the regression harness checks.
|
||||
if mode != SPEED_MAX:
|
||||
_restore_tf32(logger)
|
||||
# TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max,
|
||||
# cudnn.benchmark on any non-off CUDA load). The caller snapshots them before this
|
||||
# call and restores on unload / failed load via snapshot_backend_flags /
|
||||
# restore_backend_flags, so a later `off` load -- or chat inference in the same
|
||||
# process -- never inherits them. We keep no separate bookkeeping here.
|
||||
if mode == SPEED_OFF:
|
||||
return applied
|
||||
|
||||
|
|
@ -212,11 +217,19 @@ def apply_speed_optims(
|
|||
applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger)
|
||||
elif compile_eligible(target, is_gguf = is_gguf, family = family):
|
||||
applied["compiled"] = _compile_repeated_blocks(
|
||||
pipe, logger, max_autotune = False, cache_active = cache_active
|
||||
pipe,
|
||||
logger,
|
||||
max_autotune = False,
|
||||
cache_active = cache_active,
|
||||
offload_active = offload_active,
|
||||
)
|
||||
elif mode == SPEED_MAX and compile_eligible(target, is_gguf = is_gguf, family = family):
|
||||
applied["compiled"] = _compile_repeated_blocks(
|
||||
pipe, logger, max_autotune = True, cache_active = cache_active
|
||||
pipe,
|
||||
logger,
|
||||
max_autotune = True,
|
||||
cache_active = cache_active,
|
||||
offload_active = offload_active,
|
||||
)
|
||||
|
||||
if mode == SPEED_MAX:
|
||||
|
|
@ -247,6 +260,7 @@ def _compile_repeated_blocks(
|
|||
*,
|
||||
max_autotune: bool = False,
|
||||
cache_active: bool = False,
|
||||
offload_active: bool = False,
|
||||
) -> bool:
|
||||
transformer = getattr(pipe, "transformer", None)
|
||||
fn = getattr(transformer, "compile_repeated_blocks", None)
|
||||
|
|
@ -259,14 +273,34 @@ def _compile_repeated_blocks(
|
|||
# / max-autotune) are deliberately NOT used: they crash on the regionally-compiled
|
||||
# block because its static output buffer is overwritten across denoise steps.
|
||||
#
|
||||
# fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is
|
||||
# ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip
|
||||
# inlining torch.compiler.disable()d function"). The break is cheap and the rest of the
|
||||
# block still compiles.
|
||||
kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune}
|
||||
# fullgraph drops to False when a step cache OR CPU offloading is engaged: both insert
|
||||
# an ``@torch.compiler.disable``d function into the forward -- FBCache's per-step
|
||||
# decision, and group/model/sequential offload's ``ModuleGroup.onload_`` streaming hook
|
||||
# -- i.e. a graph break, which fullgraph=True rejects ("Skip inlining
|
||||
# torch.compiler.disable()d function"). The break is cheap and the rest of the block
|
||||
# still compiles.
|
||||
kwargs: dict[str, Any] = {
|
||||
"fullgraph": not (cache_active or offload_active),
|
||||
"dynamic": not max_autotune,
|
||||
}
|
||||
if max_autotune:
|
||||
kwargs["mode"] = "max-autotune-no-cudagraphs"
|
||||
try:
|
||||
import torch
|
||||
|
||||
# Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block
|
||||
# shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default
|
||||
# recompile_limit of 8. Once the limit is hit a resident load hard-errors under
|
||||
# fullgraph (and an offload/cache load silently drops the overflow blocks to eager),
|
||||
# so raise it well past that (64) for headroom on larger heterogeneous DiTs. This is
|
||||
# diffusers' own documented fix for regional-compile recompilation (their guide bumps
|
||||
# cache_size_limit). Deliberately NOT force_parameter_static_shapes=False: it doesn't
|
||||
# cut the variant count here and makes each compile ~6x slower (24s -> 143s cold).
|
||||
dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
|
||||
if dynamo_cfg is not None:
|
||||
for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver
|
||||
if hasattr(dynamo_cfg, _limit_attr):
|
||||
setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
|
||||
fn(**kwargs)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
|
|
@ -284,22 +318,10 @@ def _enable_cudnn_benchmark(logger: Any) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
# The TF32 flag values from before the first max load flipped them, so a later
|
||||
# non-max load / unload can put the process back exactly as it found it (rather than
|
||||
# forcing a hardcoded default that might clobber another component's choice).
|
||||
_tf32_prev: Optional[tuple[bool, bool]] = None
|
||||
|
||||
|
||||
def _enable_tf32(logger: Any) -> bool:
|
||||
global _tf32_prev
|
||||
try:
|
||||
import torch
|
||||
|
||||
if _tf32_prev is None:
|
||||
_tf32_prev = (
|
||||
torch.backends.cuda.matmul.allow_tf32,
|
||||
torch.backends.cudnn.allow_tf32,
|
||||
)
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
return True
|
||||
|
|
@ -308,26 +330,6 @@ def _enable_tf32(logger: Any) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def restore_tf32(logger: Any = None) -> None:
|
||||
"""Put the process-global TF32 flags back to their pre-max-load values. No-op if
|
||||
a max load never set them. Called on a non-max load and on unload."""
|
||||
_restore_tf32(logger)
|
||||
|
||||
|
||||
def _restore_tf32(logger: Any) -> None:
|
||||
global _tf32_prev
|
||||
if _tf32_prev is None:
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
torch.backends.cuda.matmul.allow_tf32 = _tf32_prev[0]
|
||||
torch.backends.cudnn.allow_tf32 = _tf32_prev[1]
|
||||
except Exception as exc: # noqa: BLE001 — best-effort restore
|
||||
_warn(logger, "tf32_restore", exc)
|
||||
finally:
|
||||
_tf32_prev = None
|
||||
|
||||
|
||||
def _fuse_qkv(pipe: Any, logger: Any) -> bool:
|
||||
for owner in (pipe, getattr(pipe, "transformer", None)):
|
||||
fn = getattr(owner, "fuse_qkv_projections", None)
|
||||
|
|
|
|||
|
|
@ -35,8 +35,10 @@ from typing import Any, Optional
|
|||
|
||||
from core.inference.diffusion_device import resolve_diffusion_device_target
|
||||
from core.inference.diffusion_families import (
|
||||
DIFFUSION_CANCELLED_MSG,
|
||||
DIFFUSION_NOT_LOADED_MSG,
|
||||
DiffusionFamily,
|
||||
detect_family,
|
||||
detect_family_for_pick,
|
||||
family_sd_cpp_supported,
|
||||
resolve_base_repo,
|
||||
resolve_local_gguf_child,
|
||||
|
|
@ -246,7 +248,10 @@ class SdCppDiffusionBackend:
|
|||
raise ValueError(
|
||||
"gguf_filename is required: the native engine loads single-file GGUF checkpoints only."
|
||||
)
|
||||
fam = detect_family(repo_id, family_override)
|
||||
# Use the filename-fallback detector the route validated with, so a local
|
||||
# .gguf pick whose family keyword lives only in the basename doesn't pass
|
||||
# validation and then dead-end here on a no-GPU (native-routed) host.
|
||||
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
if fam is None:
|
||||
raise ValueError(
|
||||
f"'{repo_id}' is not a supported diffusion image model. Supported families: "
|
||||
|
|
@ -489,7 +494,7 @@ class SdCppDiffusionBackend:
|
|||
with self._lock:
|
||||
state = self._state
|
||||
if state is None:
|
||||
raise RuntimeError("No diffusion model is loaded.")
|
||||
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
|
||||
self._active_generate_cancel = cancel
|
||||
engine = self._resolve_engine()
|
||||
try:
|
||||
|
|
@ -510,7 +515,7 @@ class SdCppDiffusionBackend:
|
|||
with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir:
|
||||
for index in range(max(1, int(batch_size))):
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# Distinct seed per batch image (sd-cli is one image/run here),
|
||||
# so a batch is reproducible image-by-image from the base seed.
|
||||
# Mask to sd-cli's int64 range, NOT 53 bits: the request model and
|
||||
|
|
@ -547,7 +552,7 @@ class SdCppDiffusionBackend:
|
|||
images.append(im.copy())
|
||||
seeds.append(seed_i)
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# ``seeds`` is the per-image seed (each sd-cli run used seed+index), so
|
||||
# the route can persist the real seed for every image in the batch.
|
||||
return {
|
||||
|
|
@ -557,7 +562,7 @@ class SdCppDiffusionBackend:
|
|||
"repo_id": state.repo_id,
|
||||
}
|
||||
except SdCppCancelled as exc:
|
||||
raise RuntimeError("Diffusion generation was cancelled.") from exc
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc
|
||||
finally:
|
||||
self._gen = None
|
||||
with self._lock:
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ import time
|
|||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
from utils.process_lifetime import child_popen_kwargs
|
||||
from core.inference.sd_cpp_args import (
|
||||
SdCppGenParams,
|
||||
SdCppModelFiles,
|
||||
|
|
@ -352,6 +353,10 @@ class SdCppEngine:
|
|||
# Own session/process group so cancellation/timeout can kill the whole
|
||||
# tree, not just the parent (POSIX only; harmless flag elsewhere).
|
||||
start_new_session = (os.name == "posix"),
|
||||
# Bind the child to the parent's lifetime (Linux PR_SET_PDEATHSIG), so a
|
||||
# hard parent crash mid-generation can't orphan sd-cli holding VRAM/RAM --
|
||||
# matching every llama.cpp Popen site. Composes with start_new_session.
|
||||
**child_popen_kwargs(),
|
||||
)
|
||||
# Drain stdout on a reader thread so the timeout is enforced even when the
|
||||
# child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain
|
||||
|
|
|
|||
|
|
@ -10377,6 +10377,10 @@ async def generate_diffusion_image(
|
|||
):
|
||||
from core.inference import image_gallery
|
||||
from core.inference.diffusion_engine_router import get_active_diffusion_engine
|
||||
from core.inference.diffusion_families import (
|
||||
DIFFUSION_CANCELLED_MSG,
|
||||
DIFFUSION_NOT_LOADED_MSG,
|
||||
)
|
||||
|
||||
backend = get_active_diffusion_engine()
|
||||
try:
|
||||
|
|
@ -10401,11 +10405,15 @@ async def generate_diffusion_image(
|
|||
# doesn't support) — a 400 with the reason, not a generic 500.
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
except RuntimeError as exc:
|
||||
# Only "no model loaded" / cancelled are client-state (409). The native
|
||||
# sd.cpp engine also raises RuntimeError for execution failures (nonzero
|
||||
# exit, timeout, missing output), which are server errors (500).
|
||||
# Only "no model loaded" / user-cancelled are client-state (409); both engines
|
||||
# raise these two EXACT messages. The native sd.cpp engine also raises
|
||||
# RuntimeError for execution failures (nonzero exit, timeout, missing output)
|
||||
# whose text can embed the raw sd-cli tail (local paths / argv) -- those are
|
||||
# server errors (500) returned as a fixed literal, never echoed. Match the
|
||||
# sentinels exactly, not as a substring, so an sd-cli failure that merely
|
||||
# contains "cancelled" can't misroute to 409 and leak that output.
|
||||
msg = str(exc)
|
||||
if "No diffusion model is loaded" in msg or "cancelled" in msg.lower():
|
||||
if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG):
|
||||
raise HTTPException(status_code = 409, detail = msg)
|
||||
logger.error("diffusion.generate_failed: %s", exc, exc_info = True)
|
||||
raise HTTPException(status_code = 500, detail = "Image generation failed.")
|
||||
|
|
|
|||
|
|
@ -457,7 +457,8 @@ def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monk
|
|||
base_repo = "base/repo",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.detect_family", lambda repo_id, override = None: plain
|
||||
"core.inference.diffusion.detect_family_for_pick",
|
||||
lambda repo_id, gguf_filename = None, override = None: plain,
|
||||
)
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
|
|
@ -922,6 +923,38 @@ def test_generate_without_load_raises(fake_runtime):
|
|||
backend.generate(prompt = "x")
|
||||
|
||||
|
||||
def test_failed_load_restores_backend_flags(fake_runtime, tmp_path, monkeypatch):
|
||||
# A failure AFTER apply_speed_optims (here an OOM in apply_memory_plan) must go
|
||||
# through the load's try/finally and restore the process-global TF32 / cudnn flags,
|
||||
# so a later `off` load is still bit-identical, and must not commit a partial state.
|
||||
# Regression: a refactor dropped this guard, leaking the flags on a failed load.
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
|
||||
restored: list = []
|
||||
cleared: list = []
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.restore_backend_flags", lambda snap: restored.append(snap)
|
||||
)
|
||||
monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: cleared.append(True))
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.apply_memory_plan",
|
||||
lambda *a, **k: (_ for _ in ()).throw(RuntimeError("CUDA out of memory")),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match = "out of memory"):
|
||||
backend.load_pipeline(
|
||||
str(tmp_path),
|
||||
gguf_filename = "model.gguf",
|
||||
family_override = "z-image",
|
||||
base_repo = "base/repo",
|
||||
speed_mode = "max",
|
||||
)
|
||||
assert restored, "restore_backend_flags was not called on the failed-load path"
|
||||
assert cleared, "clear_gpu_cache was not called on the failed-load path (VRAM leak)"
|
||||
assert backend._state is None and backend.is_loaded is False
|
||||
|
||||
|
||||
def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch):
|
||||
from core.inference import diffusion
|
||||
from core.inference.diffusion_families import detect_family
|
||||
|
|
|
|||
|
|
@ -27,9 +27,8 @@ from core.inference.diffusion_memory import (
|
|||
DeviceMemory,
|
||||
MemoryPlan,
|
||||
apply_memory_plan,
|
||||
estimate_gguf_dense_mib,
|
||||
estimate_gguf_resident_mib,
|
||||
estimate_image_runtime_mib,
|
||||
infer_gguf_quant_label,
|
||||
normalize_memory_mode,
|
||||
plan_diffusion_memory,
|
||||
snapshot_device_memory,
|
||||
|
|
@ -70,29 +69,15 @@ def test_normalize_memory_mode_accepts_and_rejects():
|
|||
# ── filename / size estimates ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename,expected",
|
||||
[
|
||||
("z-image-turbo-Q4_K_M.gguf", "Q4_K_M"),
|
||||
("flux1-dev-Q8_0.gguf", "Q8_0"),
|
||||
("model-BF16.gguf", "BF16"),
|
||||
("qwen-image-IQ4_XS.gguf", "IQ4_XS"),
|
||||
("no-quant-here.gguf", None),
|
||||
(None, None),
|
||||
],
|
||||
)
|
||||
def test_infer_gguf_quant_label(filename, expected):
|
||||
assert infer_gguf_quant_label(filename) == expected
|
||||
|
||||
|
||||
def test_estimate_gguf_dense_mib_expansion():
|
||||
# 4-bit roughly quadruples once dequantised to bf16; F16 is already dense.
|
||||
assert estimate_gguf_dense_mib(1000, "Q4_K_M") == 4000
|
||||
assert estimate_gguf_dense_mib(1000, "Q8_0") == 2000
|
||||
assert estimate_gguf_dense_mib(1000, "BF16") == 1000
|
||||
assert estimate_gguf_dense_mib(None, "Q4_K_M") is None
|
||||
# Unknown quant falls back to the conservative 4-bit-ish factor.
|
||||
assert estimate_gguf_dense_mib(1000, None) == 4000
|
||||
def test_estimate_gguf_resident_mib_matches_packed_size():
|
||||
# GGUF weights stay packed (uint8) on-device; diffusers dequantises per-matmul
|
||||
# transiently, so the resident footprint ~= the on-disk size regardless of quant
|
||||
# level (measured on Z-Image-Turbo: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). A
|
||||
# small margin covers allocator overhead. The prior per-quant expansion over-
|
||||
# estimated (Q2 ~7.6x) and forced needless offload on a roomy card.
|
||||
assert estimate_gguf_resident_mib(1000) == 1050
|
||||
assert estimate_gguf_resident_mib(7220) == 7581
|
||||
assert estimate_gguf_resident_mib(None) is None
|
||||
|
||||
|
||||
def test_estimate_image_runtime_scales_with_pixels_and_family():
|
||||
|
|
|
|||
|
|
@ -303,6 +303,37 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
|
|||
assert "CUDA" not in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_generate_execution_error_with_cancelled_substring_is_sanitized_500(client, monkeypatch):
|
||||
# A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled"
|
||||
# must stay a sanitized 500, not misroute to 409 and echo that output (path/arg
|
||||
# leak). Regression: the handler matched "cancelled" as a substring.
|
||||
backend = diffusion_module.get_diffusion_backend()
|
||||
backend.loaded = True
|
||||
|
||||
def _fail(**kwargs):
|
||||
raise RuntimeError("sd-cli exited 1. Last output:\nop cancelled at /home/u/models/x.gguf")
|
||||
|
||||
monkeypatch.setattr(backend, "generate", _fail)
|
||||
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
||||
assert resp.status_code == 500
|
||||
assert resp.json()["detail"] == "Image generation failed."
|
||||
assert "cancelled" not in resp.json()["detail"] and "models" not in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_generate_user_cancellation_returns_409(client, monkeypatch):
|
||||
# The exact cancellation sentinel both engines raise is client-state (409).
|
||||
backend = diffusion_module.get_diffusion_backend()
|
||||
backend.loaded = True
|
||||
|
||||
def _cancel(**kwargs):
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
|
||||
monkeypatch.setattr(backend, "generate", _cancel)
|
||||
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
||||
assert resp.status_code == 409
|
||||
assert resp.json()["detail"] == "Diffusion generation was cancelled."
|
||||
|
||||
|
||||
def test_load_unknown_family_returns_400(client, monkeypatch):
|
||||
def _raise(*a, **k):
|
||||
raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.")
|
||||
|
|
|
|||
|
|
@ -239,6 +239,25 @@ def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch):
|
|||
assert called == {"compiled_dequant": 0}
|
||||
|
||||
|
||||
def test_offload_active_drops_fullgraph(monkeypatch):
|
||||
# Group/model/sequential offload installs a torch.compiler.disable'd onload hook;
|
||||
# compiling with fullgraph=True then crashes at the first denoise step. Same reason
|
||||
# as an active step cache -> fullgraph must drop to False when offload is planned.
|
||||
# (Dense model: on this branch GGUF `default` takes the compiled-dequant path.)
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe,
|
||||
_target(),
|
||||
is_gguf = False,
|
||||
family = _family(),
|
||||
speed_mode = SPEED_DEFAULT,
|
||||
offload_active = True,
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert pipe.compile_kwargs["fullgraph"] is False
|
||||
|
||||
|
||||
def test_speed_default_gguf_compiles_only_dequant(monkeypatch):
|
||||
# GGUF `default` is the LIGHT path: compile ONLY the dequant op chain, NOT the
|
||||
# regional block compile.
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue