Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling

Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level,
use_stream) that keeps the transformer flowing through the GPU a few blocks at a
time while the text encoder / VAE stay resident, and fix VAE tiling to drive the
VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the
pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so
status never overstates either, and group falls back to whole-module offload when
the transformer can't be streamed.

Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group
cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 ->
2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names
now match that tradeoff: balanced = stream the transformer, low_vram = offload
every component. auto picks group when the companions fit resident, else model.

112 CPU tests pass.

studio/backend/core/inference/diffusion.py

@@ -458,9 +458,12 @@ class DiffusionBackend:
                 plan = self._plan_memory(
                     target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload
                 )
-                # apply_memory_plan returns the policy ACTUALLY engaged (it may fall
-                # back from sequential to whole-module offload), so status stays honest.
-                effective_policy = apply_memory_plan(pipe, plan, device = device, logger = logger)
+                # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it
+                # may fall back to whole-module offload, and tiling is a no-op on a
+                # pipeline with no tiling control), so status stays honest.
+                effective_policy, effective_tiling = apply_memory_plan(
+                    pipe, plan, device = device, logger = logger
+                )
 
                 self._state = _LoadState(
                     pipe = pipe,
@@ -471,13 +474,13 @@ class DiffusionBackend:
                     dtype = str(dtype).replace("torch.", ""),
                     cpu_offload = effective_policy != OFFLOAD_NONE,
                     offload_policy = effective_policy,
-                    vae_tiling = plan.vae_tiling,
+                    vae_tiling = effective_tiling,
                     memory_mode = plan.requested_mode,
                 )
 
         logger.info(
             "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",
-            repo_id, base, device, effective_policy, plan.vae_tiling, "; ".join(plan.reasons),
+            repo_id, base, device, effective_policy, effective_tiling, "; ".join(plan.reasons),
         )
         return self.status()
 
@@ -511,6 +514,7 @@ class DiffusionBackend:
             target = target,
             device_memory = device_memory,
             model_dense_mib = model_dense_mib,
+            companion_dense_mib = companion_mib,
             runtime_headroom_mib = runtime_headroom,
             requested_mode = memory_mode,
             explicit_offload = cpu_offload,

studio/backend/core/inference/diffusion_memory.py

@@ -36,15 +36,23 @@ MEMORY_MODES = (
 )
 
 # ── offload policies (what the loader actually does) ──────────────────────────
-# none       -> all weights resident on the device (fastest; fits only if there
-#               is room). model -> diffusers enable_model_cpu_offload(): one
-#               top-level module on the GPU at a time (modest VRAM cut, small
-#               speed cost). sequential -> enable_sequential_cpu_offload():
-#               submodule-level streaming (lowest VRAM, large speed cost).
+# none   -> all weights resident on the device (fastest; fits only if there is
+#           room). model -> diffusers enable_model_cpu_offload(): one top-level
+#           module on the GPU at a time (modest VRAM cut, small speed cost).
+# group  -> apply_group_offloading() on the transformer: stream it a few blocks at
+#           a time with a prefetch stream (lowest practical VRAM for the dominant
+#           module, less penalty than submodule sequential). sequential ->
+#           enable_sequential_cpu_offload(): submodule-level (broken for GGUF on
+#           diffusers 0.38, kept only as an explicit escape hatch).
 OFFLOAD_NONE = "none"
 OFFLOAD_MODEL = "model"
+OFFLOAD_GROUP = "group"
 OFFLOAD_SEQUENTIAL = "sequential"
 
+# Blocks of the transformer kept resident per group under group offloading: fewer
+# = lower peak VRAM, more host<->device traffic. One is the lowest-VRAM setting.
+DEFAULT_GROUP_BLOCKS = 1
+
 DEFAULT_IMAGE_WIDTH = 1024
 DEFAULT_IMAGE_HEIGHT = 1024
 # A flat allowance for the pipeline's fixed costs (scheduler, embeddings, the
@@ -312,29 +320,47 @@ def plan_diffusion_memory(
     device_memory: DeviceMemory,
     model_dense_mib: Optional[int],
     runtime_headroom_mib: int,
+    companion_dense_mib: Optional[int] = None,
     base_overhead_mib: int = DEFAULT_BASE_OVERHEAD_MIB,
     requested_mode: Optional[str] = None,
     explicit_offload: bool = False,
 ) -> MemoryPlan:
     """Pick an offload policy + VAE memory savers for the current load.
 
-    ``model_dense_mib`` is the estimated resident device size of the weights
-    (transformer + companion text-encoder / VAE). ``explicit_offload`` is the
-    back-compat ``cpu_offload=True`` request: it forces at least whole-module
-    offload wherever offload is meaningful."""
+    ``model_dense_mib`` is the estimated resident device size of all weights
+    (transformer + companion text-encoder / VAE); ``companion_dense_mib`` is just
+    the companions, which stay resident under streamed (group) offload while the
+    transformer is streamed block by block. ``explicit_offload`` is the back-compat
+    ``cpu_offload=True`` request: it forces whole-module offload.
+
+    Policy meanings, ordered by measured speed/VRAM tradeoff:
+      none   - everything resident: fastest, highest VRAM.
+      group  - stream the transformer, companions resident: near-resident speed,
+               moderate VRAM cut (the balanced tradeoff).
+      model  - offload every component incl. the text encoder: lowest VRAM, slow.
+    """
     mode = normalize_memory_mode(requested_mode) or MEMORY_MODE_AUTO
     can_offload = bool(getattr(target, "supports_model_cpu_offload", False))
     budget = _safe_device_budget_mib(device_memory)
     required = _sum_required(model_dense_mib, runtime_headroom_mib, base_overhead_mib)
+    # The resident floor under group offload: companions stay, the transformer streams.
+    group_floor = _sum_required(companion_dense_mib, runtime_headroom_mib, base_overhead_mib)
     reasons: list[str] = []
     estimates: dict[str, Optional[int]] = {
         "safe_device_budget_mib": budget,
         "model_dense_mib": model_dense_mib,
+        "companion_dense_mib": companion_dense_mib,
         "runtime_headroom_mib": runtime_headroom_mib,
         "base_overhead_mib": base_overhead_mib,
         "resident_required_mib": required,
+        "group_floor_mib": group_floor,
     }
 
+    def _group_fits() -> bool:
+        # Group offload only helps if the resident remainder (companions) fits; when
+        # the text encoder itself is too big, only whole-module offload will do.
+        return group_floor is not None and budget is not None and group_floor <= budget
+
     if not can_offload or device_memory.is_unified:
         # MPS / CPU can't stream to a separate device, and on unified / system
         # memory CPU offload just shuffles bytes within the same pool.
@@ -346,28 +372,29 @@ def plan_diffusion_memory(
     elif mode == MEMORY_MODE_FAST:
         policy = OFFLOAD_NONE
         if budget is not None and required is not None and required > budget:
-            policy = OFFLOAD_MODEL
-            reasons.append("fast requested but weights do not fit resident; whole-module offload")
+            # Doesn't fit resident: the fastest offload is the streamed transformer.
+            policy = OFFLOAD_GROUP if _group_fits() else OFFLOAD_MODEL
+            reasons.append("fast requested but weights do not fit resident; offloading")
         else:
             reasons.append("fast requested; weights resident on device")
     elif mode == MEMORY_MODE_BALANCED:
-        policy = OFFLOAD_MODEL
-        reasons.append("balanced requested; whole-module CPU offload")
+        policy = OFFLOAD_GROUP
+        reasons.append("balanced requested; streamed block-level transformer offload")
     elif mode == MEMORY_MODE_LOW_VRAM:
-        policy = OFFLOAD_SEQUENTIAL
-        reasons.append("low_vram requested; submodule sequential CPU offload")
+        policy = OFFLOAD_MODEL
+        reasons.append("low_vram requested; whole-module offload of every component")
     elif budget is None or required is None:
         policy = OFFLOAD_NONE
         reasons.append("device budget or model size unknown; staying resident")
     elif required <= int(budget * 0.85):
         policy = OFFLOAD_NONE
         reasons.append("weights fit resident with headroom")
-    elif required <= budget:
-        policy = OFFLOAD_MODEL
-        reasons.append("tight fit; whole-module CPU offload")
+    elif _group_fits():
+        policy = OFFLOAD_GROUP
+        reasons.append("tight fit; stream the transformer, companions resident")
     else:
-        policy = OFFLOAD_SEQUENTIAL
-        reasons.append("weights exceed device budget; submodule sequential CPU offload")
+        policy = OFFLOAD_MODEL
+        reasons.append("companions exceed budget; whole-module offload of every component")
 
     if explicit_offload and policy == OFFLOAD_NONE and can_offload and not device_memory.is_unified:
         policy = OFFLOAD_MODEL
@@ -393,24 +420,31 @@ def plan_diffusion_memory(
 # ── apply to a built pipeline ─────────────────────────────────────────────────
 
 
-def apply_memory_plan(pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = None) -> str:
+def apply_memory_plan(
+    pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = None,
+) -> tuple[str, bool]:
     """Apply ``plan`` to a freshly built diffusers pipeline: enable the VAE memory
-    savers (best-effort; not every pipeline exposes them) then place / offload the
-    weights. Exactly one placement call runs, so the pipeline ends up either fully
-    resident or wired for CPU offload, never both.
+    savers then place / offload the weights. Exactly one placement call runs, so the
+    pipeline ends up either fully resident or wired for offload, never both.
 
-    Returns the offload policy actually engaged, which can differ from the plan:
-    submodule sequential offload is unreliable for GGUF-dequantised transformers on
-    some diffusers versions, so it falls back to the robust whole-module offload
-    (the proper low-VRAM GGUF path is the CPU-resident cache, a later phase)."""
+    Returns ``(offload_policy, vae_tiling)`` ACTUALLY engaged, which can differ from
+    the plan: VAE tiling is a no-op on a pipeline that exposes no tiling control, and
+    block-level / sequential offload fall back to the robust whole-module offload if
+    the transformer doesn't support them (e.g. submodule sequential is broken for
+    GGUF on diffusers 0.38). Status then reflects what really happened."""
+    tiling_engaged = False
     if plan.vae_tiling:
-        _try_call(pipe, "enable_vae_tiling", logger)
+        tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger)
     if plan.vae_slicing:
-        _try_call(pipe, "enable_vae_slicing", logger)
+        _enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger)
 
     policy = plan.offload_policy
     if policy == OFFLOAD_MODEL:
         pipe.enable_model_cpu_offload()
+    elif policy == OFFLOAD_GROUP:
+        if not _apply_group_offload(pipe, device, logger):
+            pipe.enable_model_cpu_offload()
+            policy = OFFLOAD_MODEL
     elif policy == OFFLOAD_SEQUENTIAL:
         try:
             pipe.enable_sequential_cpu_offload()
@@ -424,15 +458,59 @@ def apply_memory_plan(pipe: Any, plan: MemoryPlan, *, device: str, logger: Any =
             policy = OFFLOAD_MODEL
     else:
         pipe.to(device)
-    return policy
+    return policy, tiling_engaged
 
 
-def _try_call(pipe: Any, method: str, logger: Any) -> None:
-    fn = getattr(pipe, method, None)
-    if not callable(fn):
-        return
+def _enable_vae_saver(pipe: Any, pipe_method: str, vae_method: str, logger: Any) -> bool:
+    """Turn on a VAE memory saver, trying the pipeline-level shortcut first and the
+    VAE submodule directly otherwise (some pipelines, e.g. Z-Image, only expose it
+    on ``pipe.vae``). Returns whether it actually engaged."""
+    for owner, method in ((pipe, pipe_method), (getattr(pipe, "vae", None), vae_method)):
+        fn = getattr(owner, method, None)
+        if not callable(fn):
+            continue
+        try:
+            fn()
+            return True
+        except Exception as exc:  # noqa: BLE001 — a VAE saver is an optimisation, never fatal
+            if logger is not None:
+                logger.warning("diffusion.memory: %s() failed: %s", method, exc)
+    return False
+
+
+def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
+    """Stream the transformer through the device a few blocks at a time via
+    diffusers group offloading, keeping the smaller components resident. Returns
+    False (so the caller falls back to whole-module offload) on any failure."""
+    transformer = getattr(pipe, "transformer", None)
+    if transformer is None:
+        return False
     try:
-        fn()
-    except Exception as exc:  # noqa: BLE001 — a VAE saver is an optimisation, never fatal
+        import torch
+        from diffusers.hooks import apply_group_offloading
+
+        onload = torch.device(device)
+        use_stream = onload.type == "cuda"  # overlap H2D copies with compute on CUDA
+        apply_group_offloading(
+            transformer,
+            onload_device = onload,
+            offload_device = torch.device("cpu"),
+            offload_type = "block_level",
+            num_blocks_per_group = DEFAULT_GROUP_BLOCKS,
+            use_stream = use_stream,
+        )
+        # Place the remaining (smaller) components resident; the streamed
+        # transformer manages its own placement via the offloading hooks.
+        for name, comp in getattr(pipe, "components", {}).items():
+            if name == "transformer":
+                continue
+            if isinstance(comp, torch.nn.Module):
+                comp.to(onload)
+        return True
+    except Exception as exc:  # noqa: BLE001 — fall back to whole-module offload
         if logger is not None:
-            logger.warning("diffusion.memory: %s() failed: %s", method, exc)
+            logger.warning(
+                "diffusion.memory: group offload failed (%s); falling back to "
+                "whole-module offload", exc,
+            )
+        return False

studio/backend/models/inference.py

@@ -1704,7 +1704,8 @@ class DiffusionLoadRequest(BaseModel):
     memory_mode: Optional[Literal["auto", "fast", "balanced", "low_vram"]] = Field(
         None,
         description = "Memory policy: auto (measured), fast (resident), balanced "
-                      "(model CPU offload), low_vram (sequential CPU offload). "
+                      "(stream the transformer, near-resident speed, moderate VRAM "
+                      "cut), low_vram (offload every component, lowest VRAM, slower). "
                       "Overrides cpu_offload when set.",
     )
 

studio/backend/tests/test_diffusion_backend.py

@@ -798,28 +798,34 @@ def _force_cuda_target(backend, monkeypatch):
     monkeypatch.setattr(backend, "_pick_device_and_dtype", lambda: ("cuda", torch.bfloat16))
 
 
-def test_load_memory_mode_balanced_engages_model_offload(fake_runtime, tmp_path, monkeypatch):
+def test_load_memory_mode_balanced_streams_or_falls_back(fake_runtime, tmp_path, monkeypatch):
+    # balanced requests streamed block-level (group) offload. Under the stub there is
+    # no real diffusers.hooks, so group can't engage and the applier falls back to
+    # whole-module offload, reporting the policy actually engaged (the real "group"
+    # path is GPU-verified in the bench).
     (tmp_path / "m.gguf").write_bytes(b"x")
     backend = DiffusionBackend()
     _force_cuda_target(backend, monkeypatch)
     status = backend.load_pipeline(
         str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", memory_mode = "balanced"
     )
-    assert status["offload_policy"] == "model" and status["cpu_offload"] is True
+    assert status["offload_policy"] in ("group", "model") and status["cpu_offload"] is True
     assert status["memory_mode"] == "balanced"
-    pipe = backend._state.pipe
-    assert pipe.offloaded is True and pipe.moved_to is None  # offload owns placement
+    assert backend._state.pipe.offloaded is True  # model-offload fallback engaged
 
 
-def test_load_memory_mode_low_vram_uses_sequential_offload(fake_runtime, tmp_path, monkeypatch):
+def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path, monkeypatch):
+    # low_vram offloads every component (lowest VRAM); whole-module offload is the
+    # robust path and engages directly (no streaming, so no diffusers.hooks needed).
     (tmp_path / "m.gguf").write_bytes(b"x")
     backend = DiffusionBackend()
     _force_cuda_target(backend, monkeypatch)
     status = backend.load_pipeline(
         str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", memory_mode = "low_vram"
     )
-    assert status["offload_policy"] == "sequential" and status["cpu_offload"] is True
-    assert backend._state.pipe.sequential_offloaded is True
+    assert status["offload_policy"] == "model" and status["cpu_offload"] is True
+    pipe = backend._state.pipe
+    assert pipe.offloaded is True and pipe.moved_to is None  # offload owns placement
 
 
 def test_load_explicit_cpu_offload_engages_model_offload_on_cuda(fake_runtime, tmp_path, monkeypatch):

studio/backend/tests/test_diffusion_memory.py

@@ -20,10 +20,12 @@ from core.inference.diffusion_memory import (
     MEMORY_MODE_BALANCED,
     MEMORY_MODE_FAST,
     MEMORY_MODE_LOW_VRAM,
+    OFFLOAD_GROUP,
     OFFLOAD_MODEL,
     OFFLOAD_NONE,
     OFFLOAD_SEQUENTIAL,
     DeviceMemory,
+    MemoryPlan,
     apply_memory_plan,
     estimate_gguf_dense_mib,
     estimate_image_runtime_mib,
@@ -166,14 +168,41 @@ def test_auto_model_offload_on_tight_fit():
     assert plan.vae_tiling is True  # offloading -> device is tight -> tile
 
 
-def test_auto_sequential_when_model_exceeds_budget():
+def test_auto_group_offload_when_transformer_overflows_but_companions_fit():
+    # Big transformer pushes the resident total over budget, but the companions
+    # (text encoder + VAE) still fit -> stream the transformer (fast, moderate cut).
+    plan = plan_diffusion_memory(
+        target = _target(),
+        device_memory = _discrete(8000, 8000),
+        model_dense_mib = 40000,
+        companion_dense_mib = 1500,
+        runtime_headroom_mib = 1000,
+        base_overhead_mib = 1000,
+    )
+    assert plan.offload_policy == OFFLOAD_GROUP
+
+
+def test_auto_model_offload_when_companions_exceed_budget():
+    # The text encoder itself is too big to stay resident -> offload everything.
+    plan = plan_diffusion_memory(
+        target = _target(),
+        device_memory = _discrete(8000, 8000),
+        model_dense_mib = 40000,
+        companion_dense_mib = 30000,
+        runtime_headroom_mib = 4000,
+    )
+    assert plan.offload_policy == OFFLOAD_MODEL
+
+
+def test_auto_model_offload_when_companion_size_unknown():
+    # Without a companion estimate the planner can't prove group fits -> safest cut.
     plan = plan_diffusion_memory(
         target = _target(),
         device_memory = _discrete(8000, 8000),
         model_dense_mib = 40000,
         runtime_headroom_mib = 4000,
     )
-    assert plan.offload_policy == OFFLOAD_SEQUENTIAL
+    assert plan.offload_policy == OFFLOAD_MODEL
 
 
 def test_auto_stays_resident_when_budget_unknown():
@@ -199,11 +228,11 @@ def test_explicit_modes_force_policy_regardless_of_budget():
     assert plan_diffusion_memory(
         target = _target(), device_memory = roomy, model_dense_mib = 1000,
         runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_BALANCED,
-    ).offload_policy == OFFLOAD_MODEL
+    ).offload_policy == OFFLOAD_GROUP
     assert plan_diffusion_memory(
         target = _target(), device_memory = roomy, model_dense_mib = 1000,
         runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_LOW_VRAM,
-    ).offload_policy == OFFLOAD_SEQUENTIAL
+    ).offload_policy == OFFLOAD_MODEL
 
 
 def test_fast_falls_back_to_model_offload_when_it_does_not_fit():
@@ -313,31 +342,74 @@ def _plan(policy, *, tiling):
         runtime_headroom_mib = 1000,
         requested_mode = {
             OFFLOAD_NONE: MEMORY_MODE_FAST,
-            OFFLOAD_MODEL: MEMORY_MODE_BALANCED,
-            OFFLOAD_SEQUENTIAL: MEMORY_MODE_LOW_VRAM,
+            OFFLOAD_GROUP: MEMORY_MODE_BALANCED,
+            OFFLOAD_MODEL: MEMORY_MODE_LOW_VRAM,
         }[policy],
     )
 
 
+def _manual_plan(policy, *, tiling):
+    """Build a plan for a policy the auto/explicit modes no longer emit (sequential)."""
+    return MemoryPlan(
+        requested_mode = "manual",
+        offload_policy = policy,
+        vae_tiling = tiling,
+        vae_slicing = tiling,
+        device_memory = _discrete(4000, 8000),
+        estimates = {},
+    )
+
+
 def test_apply_none_places_resident():
     pipe = _RecordingPipe()
-    effective = apply_memory_plan(pipe, _plan(OFFLOAD_NONE, tiling = False), device = "cuda")
+    effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_NONE, tiling = False), device = "cuda")
     assert pipe.calls == ["to:cuda"]  # no tiling on a roomy resident run
-    assert effective == OFFLOAD_NONE
+    assert effective == OFFLOAD_NONE and tiled is False
 
 
 def test_apply_model_offload_engages_offload_and_tiling():
     pipe = _RecordingPipe()
-    effective = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda")
+    effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda")
     assert "model_offload" in pipe.calls
     assert "to:cuda" not in pipe.calls  # offload owns placement; never both
     assert "vae_tiling" in pipe.calls and "vae_slicing" in pipe.calls
-    assert effective == OFFLOAD_MODEL
+    assert effective == OFFLOAD_MODEL and tiled is True
+
+
+def test_apply_vae_tiling_falls_back_to_vae_submodule():
+    # Z-Image-style pipeline: no pipeline-level enable_vae_tiling, only pipe.vae.
+    class _VaeOnly:
+        def __init__(self):
+            self.vae = types.SimpleNamespace(
+                tiled = False, sliced = False,
+                enable_tiling = self._tile, enable_slicing = self._slice,
+            )
+
+        def _tile(self):
+            self.vae.tiled = True
+
+        def _slice(self):
+            self.vae.sliced = True
+
+        def enable_model_cpu_offload(self):
+            self.offloaded = True
+
+    pipe = _VaeOnly()
+    effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda")
+    assert tiled is True and pipe.vae.tiled and pipe.vae.sliced
+
+
+def test_apply_group_falls_back_to_model_without_transformer():
+    # The recording pipe has no .transformer, so group offload can't engage and the
+    # applier falls back to whole-module offload, reporting the real policy.
+    pipe = _RecordingPipe()
+    effective, _ = apply_memory_plan(pipe, _plan(OFFLOAD_GROUP, tiling = True), device = "cuda")
+    assert effective == OFFLOAD_MODEL and "model_offload" in pipe.calls
 
 
 def test_apply_sequential_offload():
     pipe = _RecordingPipe()
-    effective = apply_memory_plan(pipe, _plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
+    effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
     assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls
     assert effective == OFFLOAD_SEQUENTIAL
 
@@ -350,7 +422,7 @@ def test_apply_sequential_falls_back_to_model_offload_when_unsupported():
             raise RuntimeError("sequential offload not supported for this transformer")
 
     pipe = _NoSeqPipe()
-    effective = apply_memory_plan(pipe, _plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
+    effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
     assert effective == OFFLOAD_MODEL
     assert "model_offload" in pipe.calls
 
@@ -365,5 +437,5 @@ def test_apply_tolerates_pipe_without_vae_savers():
             self.moved = device
 
     bare = _Bare()
-    apply_memory_plan(bare, _plan(OFFLOAD_NONE, tiling = False), device = "cpu")
-    assert bare.moved == "cpu"
+    _, tiled = apply_memory_plan(bare, _plan(OFFLOAD_NONE, tiling = False), device = "cpu")
+    assert bare.moved == "cpu" and tiled is False