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