From 4e1c622d20cff957988858e1979f499c6084ecd6 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 25 May 2026 17:28:06 +0000
Subject: [PATCH] Fix/adjust diffusion: gate accelerate preflight on
 cpu_offload for PR #5754

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).
---
 studio/backend/core/inference/diffusion.py | 23 ++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 316f298c3d..efc8799ba0 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -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 "