bring the text encoder back before using it again

Production, 2026-10-28: Krea-2-Turbo answered one request and then 500ed.

    RuntimeError: Expected all tensors to be on the same device, but got
    index is on cuda:0, different from other tensors on cpu

The offload moved the encoder to the CPU and the pipeline is CACHED, so the
next request met an encoder that was no longer where encode_prompt indexes.
_free's docstring even said the encoder "waits in host RAM ... the next
request needs it again" -- and nothing ever brought it back.

What turned that into a 500 rather than a skipped optimisation: the side
effect outlives this function. The bail-out path returns the original kwargs
and lets the pipeline generate "the ordinary way", but the ordinary way calls
the SAME encoder, so it hit the same mismatch and the request failed anyway.
A fallback that leaves the thing it fell back to broken is not a fallback.

So the offload is a per-request cycle -- bring the encoder on, encode, move it
off -- and every bail-out restores the encoder first, because each one hands
the request to a path that needs it on the GPU.

The regression test drives a pipeline that really refuses to encode off-device
and asserts the SECOND request still works; without the restore it reproduces
the production message exactly. A second test pins the weaker promise that any
bail-out leaves the encoder where the ordinary path can use it.

Found by the operator mid-render, not by the suite: every check so far ran one
request per model against a freshly started engine, which is the one case that
cannot fail.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014T3rZozi1Ehk96NG1WesgC

codai/models/text_encoder_offload.py

@@ -219,6 +219,52 @@ def _assign_outputs(out, params: set, mask_arg: Optional[str]) -> Dict[str, Any]
     return assigned
 
 
+def _ensure_on_device(pipe, comp_names=("text_encoder", "text_encoder_2")) -> list:
+    """Bring the named components back to the pipeline's execution device.
+
+    REQUIRED, not an optimisation. The pipeline is cached and served again, so
+    a previous request's offload has left the encoder on the CPU. Encoding then
+    indexes cuda input ids against cpu weights:
+
+        RuntimeError: Expected all tensors to be on the same device, but got
+        index is on cuda:0, different from other tensors on cpu
+
+    and -- the part that made this a 500 rather than a slow path -- the damage
+    OUTLIVES this function. Falling back to "generate the ordinary way" hands
+    the request to a pipeline whose own encode_prompt hits the same mismatch,
+    so the request fails anyway. Measured in production 2026-10-28: the first
+    request succeeded and every one after it returned 500 until the model was
+    evicted.
+
+    So the offload is a per-request CYCLE -- bring the encoder on, encode, move
+    it off -- not a one-way move.
+    """
+    notes = []
+    try:
+        import torch
+    except Exception:
+        return notes
+    dev = getattr(pipe, "_execution_device", None)
+    if dev is None:
+        try:
+            dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        except Exception:
+            return notes
+    for name in comp_names:
+        comp = getattr(pipe, name, None)
+        if comp is None or not hasattr(comp, "to"):
+            continue
+        if _resident_on_gpu(comp):
+            continue
+        try:
+            comp.to(dev)
+            notes.append(f"{name}: brought back to {dev}")
+        except Exception as exc:
+            notes.append(f"{name}: could NOT be brought back to {dev} "
+                         f"({type(exc).__name__}: {str(exc)[:120]})")
+    return notes
+
+
 def encode_then_free(pipe, kwargs: Dict[str, Any], *,
                      log=None) -> Tuple[Dict[str, Any], float]:
     """Encode the prompt in ``kwargs``, free the encoder, return new kwargs.
@@ -277,11 +323,21 @@ def encode_then_free(pipe, kwargs: Dict[str, Any], *,
             enc_kwargs["do_classifier_free_guidance"] = bool(
                 gs is not None and float(gs) > 1.0)
 
+    # Put the encoder back before using it: a previous request on this cached
+    # pipeline may have left it on the CPU.
+    back = _ensure_on_device(pipe)
+    if back:
+        _say(f"  [text-encoder-offload] {'; '.join(back)}")
+
     try:
         import torch
         with torch.no_grad():
             out = pipe.encode_prompt(**enc_kwargs)
     except Exception as exc:
+        # Leave the pipeline able to generate on its own: the ordinary path
+        # calls this same encoder, so handing it back on the CPU would turn a
+        # skipped optimisation into a failed request.
+        _ensure_on_device(pipe)
         _say(f"  [text-encoder-offload] encoding failed ({exc}); "
              f"generating the ordinary way")
         return kwargs, 0.0
@@ -295,10 +351,12 @@ def encode_then_free(pipe, kwargs: Dict[str, Any], *,
     try:
         assigned = _assign_outputs(out, params, mask_arg)
     except Exception as exc:
+        _ensure_on_device(pipe)
         _say(f"  [text-encoder-offload] could not assemble the embeddings "
              f"({exc}); generating the ordinary way")
         return kwargs, 0.0
     if "prompt_embeds" not in assigned:
+        _ensure_on_device(pipe)
         _say("  [text-encoder-offload] encode_prompt returned no embeddings; "
              "generating the ordinary way")
         return kwargs, 0.0
@@ -308,6 +366,7 @@ def encode_then_free(pipe, kwargs: Dict[str, Any], *,
         # is exactly what happened on the first live LTX-2.3 run: the mask was
         # read from the wrong tuple slot and the call raised "Must provide
         # `prompt_attention_mask` when specifying `prompt_embeds`".
+        _ensure_on_device(pipe)
         _say(f"  [text-encoder-offload] encode_prompt returned no {mask_arg}; "
              f"generating the ordinary way")
         return kwargs, 0.0
@@ -343,6 +402,7 @@ def encode_then_free(pipe, kwargs: Dict[str, Any], *,
                 if "negative_prompt_embeds" not in new:
                     raise ValueError("no negative embeddings came back")
         except Exception as exc:
+            _ensure_on_device(pipe)
             _say(f"  [text-encoder-offload] could not encode the negative "
                  f"prompt ({exc}); generating the ordinary way")
             return kwargs, 0.0

tests/test_text_encoder_offload_outputs.py

@@ -233,3 +233,86 @@ def test_both_in_one_family_is_not_called_twice():
                                                "guidance_scale": 3.0})
     assert calls == ["a cat"], f"encoded more than once: {calls}"
     assert out["negative_prompt_attention_mask"] is not None
+
+
+class _StatefulPipe:
+    """A cached pipeline that really refuses to encode off-device.
+
+    Models the production failure: the pipeline is cached between requests, so
+    the previous request's offload has left the encoder on the CPU, and
+    encode_prompt indexes cuda ids against cpu weights.
+    """
+
+    def __init__(self):
+        self.device = "cuda"          # where the encoder currently is
+        self._execution_device = "cuda"
+        self.encodes = 0
+
+        class _Enc:
+            def __init__(self, outer):
+                self._outer = outer
+
+            def parameters(self):
+                class _P:
+                    pass
+                p = _P()
+                p.device = type("D", (), {"type": self._outer.device})()
+                return iter([p])
+
+            def to(self, dev):
+                self._outer.device = "cuda" if "cuda" in str(dev) else "cpu"
+                return self
+
+        self.text_encoder = _Enc(self)
+
+    def __call__(self, prompt=None, prompt_embeds=None,
+                 prompt_attention_mask=None, negative_prompt_embeds=None,
+                 negative_prompt_attention_mask=None):
+        raise AssertionError("must not be called by the helper")
+
+    def encode_prompt(self, prompt=None, negative_prompt=None,
+                      do_classifier_free_guidance=True, device=None):
+        if self.device != "cuda":
+            raise RuntimeError(
+                "Expected all tensors to be on the same device, but got index "
+                "is on cuda:0, different from other tensors on cpu")
+        self.encodes += 1
+        return _embeds(), _mask(), _embeds(5), _mask(5)
+
+
+def test_a_second_request_on_a_cached_pipeline_still_works():
+    """The offload is a per-request CYCLE. Without bringing the encoder back,
+    request 2 raised a device mismatch, the helper fell back to "the ordinary
+    way", and the pipeline's OWN encode_prompt then hit the same mismatch --
+    so the fallback produced a 500 rather than a slow path."""
+    pipe = _StatefulPipe()
+
+    out1, freed1 = teo.encode_then_free(pipe, {"prompt": "first"})
+    assert "prompt_embeds" in out1, "request 1 should use the offload"
+    assert pipe.device == "cpu", "the encoder should have been moved off"
+
+    out2, _ = teo.encode_then_free(pipe, {"prompt": "second"})
+    assert "prompt_embeds" in out2, (
+        "request 2 fell back instead of bringing the encoder back")
+    assert pipe.encodes == 2, "the second request never encoded"
+
+
+def test_a_bail_out_leaves_the_encoder_where_the_pipeline_can_use_it():
+    """Whatever goes wrong, the ordinary path must still be able to generate:
+    it calls the same encoder."""
+
+    class _BadShape(_StatefulPipe):
+        def encode_prompt(self, prompt=None, negative_prompt=None,
+                          do_classifier_free_guidance=True, device=None):
+            if self.device != "cuda":
+                raise RuntimeError("Expected all tensors to be on the same device")
+            # More tensors than can be named -> the helper abandons the attempt.
+            return tuple(_embeds() for _ in range(7))
+
+    pipe = _BadShape()
+    pipe.device = "cpu"               # as a previous request left it
+    out, freed = teo.encode_then_free(pipe, {"prompt": "x"})
+    assert out == {"prompt": "x"}, "the ordinary path needs its prompt back"
+    assert freed == 0.0
+    assert pipe.device == "cuda", (
+        "the encoder was left on the CPU, so the ordinary path will 500 too")