ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix/adjust diffusion: gate accelerate preflight on cpu_offload for PR #5754

Browse source 
Backend CI on Python 3.11 failed 15 diffusion tests after R30's
accelerate preflight because the CI test environment does not
install accelerate, but the tests mock from_pretrained and never
exercise the CPU-offload path that actually needs it.

Gate the find_spec("accelerate") check on enable_model_cpu_offload
so the dependency is only required for the path that uses it.
transformers preflight stays unconditional (it is always touched
by from_pretrained). Tests with offload=False (the default) pass
without accelerate; production loads with offload=True still get
the fail-fast unload-protection guard the original round-30 fix
added.

97 targeted backend tests pass (test_diffusion_routes,
test_diffusion_backend, test_inference_model_validation,
test_data_recipe_seed, test_training_raw_support,
test_export_log_cursor).
This commit is contained in:
Daniel Han-Chen 2026-05-25 17:28:06 +00:00
commit 4e1c622d20
1 changed files with 15 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

23
studio/backend/core/inference/diffusion.py
View file

@ -800,16 +800,23 @@ class DiffusionBackend:
"loading an image model."
) from exc
# Round 30 P1 #11: also preflight transformers + accelerate
# BEFORE any destructive unload. Diffusers can expose stub
# pipeline classes when transformers is missing or broken, so
# the load would otherwise tear down chat first and fail
# later inside from_pretrained. Use find_spec (no module
# execution) so test environments that stub these modules
# still pass the preflight without us actually importing them.
# Round 30 P1 #11: also preflight transformers BEFORE any
# destructive unload. Diffusers can expose stub pipeline
# classes when transformers is missing or broken, so the load
# would otherwise tear down chat first and fail later inside
# from_pretrained. Use find_spec (no module execution) so test
# environments that stub these modules still pass the preflight
# without us actually importing them.
# Round 34: accelerate is only needed for the CPU-offload path
# (``enable_model_cpu_offload`` / ``device_map="auto"`` /
# offload hooks); gate the preflight on the offload flag so
# tests and offload=False inference paths do not require it.
import importlib.util as _ilu
for _mod in ("transformers", "accelerate"):
_required = ["transformers"]
if enable_model_cpu_offload:
_required.append("accelerate")
for _mod in _required:
if _ilu.find_spec(_mod) is None:
raise RuntimeError(
"Diffusion image generation requires the Studio torch "