From 6089720c0c8fd0bf153c47644d1ab4d2aa420fd5 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 00:06:17 +0000 Subject: [PATCH] Fix/adjust diffusion: export unload + sd3.5 alias for PR #5754 - routes/export.py load_checkpoint now unloads the diffusion pipeline alongside the existing inference + training unloads, so an export load after Images does not OOM the export subprocess. - Remove the 'sd3.5' alias from the stable-diffusion-3 family. SD3.5 needs its own family + base_repo (and its own smoke test); pairing it with the SD3 Medium base produced a misleading load. --- studio/backend/core/inference/diffusion.py | 8 +++++++- studio/backend/routes/export.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 1291a60b45..50bebdd98c 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -113,7 +113,13 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "StableDiffusion3Pipeline", transformer_class = "SD3Transformer2DModel", base_repo = "stabilityai/stable-diffusion-3-medium-diffusers", - aliases = ("sd3-medium", "stable-diffusion-3-medium", "sd3.5"), + # Intentionally NOT including "sd3.5" / "stable-diffusion-3.5" + # here: the SD3.5 family uses a different transformer config and + # base repo than SD3 Medium, and silently pairing SD3.5 GGUFs + # with the Medium base produces a misleading load. Add a + # dedicated SD3.5 family with its own base_repo when we ship + # smoke coverage for it. + aliases = ("sd3-medium", "stable-diffusion-3-medium"), ), # SDXL: full diffusers path only (no GGUF). SDXL uses a UNet (not a # transformer) and wiring UNet2DConditionModel.from_single_file + diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 7dbc52dbed..b35ad56e4f 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -81,6 +81,20 @@ async def load_checkpoint( except Exception as e: logger.warning("Could not unload inference model: %s", e) + # Also unload any active diffusion pipeline (Images page); it + # competes for the same GPU and would survive the inference + # shutdown above. Best effort; silently skip if the module is + # absent. + try: + from core.inference.diffusion import get_diffusion_backend + + diff = get_diffusion_backend() + if diff.is_loaded: + logger.info("Unloading diffusion model to free GPU memory for export") + diff.unload_model() + except Exception as e: + logger.debug("diffusion unload skipped for export: %s", e) + try: from core.training import get_training_backend