Fix MPS prompt staging memory, failure recovery and pipeline provenance

- Encode prompts under no_grad so cached embeddings hold no autograd graph
  that would keep the released text encoder alive.
- Release and reload the pipeline when a prompt stage fails after the DiT
  was dropped, instead of failing every later request on transformer None.
- Hash the bundle's pipeline_moss_image.py that the MPS path executes, and
  record mps_runtime.py separately, instead of hashing only the wrapper.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

gradio_app.py

@@ -306,6 +306,11 @@ class ModelRunner:
                 raise ValueError(
                     "GPU memory exhausted. Try fewer/smaller images or a lower-memory offload mode."
                 ) from exc
+            except Exception:
+                # A failed MPS prompt stage leaves the pipeline without its DiT; reload it next time.
+                if self.pipe is not None and self.pipe.transformer is None:
+                    self.unload()
+                raise
             report(0, desc="Saving images and parameters")
             folder = self.output_dir / uuid4().hex
             folder.mkdir(parents=True, exist_ok=True)
@@ -324,6 +329,10 @@ class ModelRunner:
                         p for p in (root / "tokenizer").iterdir() if p.suffix in (".json", ".jinja", ".txt")
                     )
                 )
+            # On MPS the pipeline class is mps_runtime's wrapper around the bundle's pipeline source.
+            pipeline_file = getattr(self.pipe, "pipeline_file", None)
+            wrapper = Path(mps_runtime.__file__) if pipeline_file else None
+            pipeline_file = Path(pipeline_file or inspect.getfile(type(self.pipe)))
             record = {
                 "schema_version": 2,
                 "model": root.name if root.is_dir() else model,
@@ -333,9 +342,8 @@ class ModelRunner:
                 "sampler": "Euler",
                 "weight_dtype": str(self.pipe.transformer.dtype).removeprefix("torch."),
                 "vae_dtype": "float32",
-                "pipeline_sha256": hashlib.sha256(
-                    Path(inspect.getfile(type(self.pipe))).read_bytes()
-                ).hexdigest(),
+                "pipeline_sha256": hashlib.sha256(pipeline_file.read_bytes()).hexdigest(),
+                "mps_runtime_sha256": hashlib.sha256(wrapper.read_bytes()).hexdigest() if wrapper else None,
                 "model_config_sha256": {
                     str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
                     for p in configs

mps_runtime.py

@@ -117,7 +117,9 @@ def _staged_pipeline_class(module):
             key = ((prompt,) if isinstance(prompt, str) else tuple(prompt), max_sequence_length)
             if key not in self.embeddings:
                 self.use_text_encoder()
-                self.embeddings[key] = super().encode_prompt(prompt, max_sequence_length)
+                # An autograd graph would keep the released text encoder alive in the cache.
+                with torch.no_grad():
+                    self.embeddings[key] = super().encode_prompt(prompt, max_sequence_length)
                 while len(self.embeddings) > CACHED_PROMPTS:
                     self.embeddings.popitem(last=False)
             self.embeddings.move_to_end(key)
@@ -158,6 +160,8 @@ def load_pipeline(model, local_files_only=False):
         vae=AutoencoderKLFlux2.from_pretrained(root / "vae", torch_dtype=torch.float32).to(DEVICE).eval(),
     )
     pipe.root = root
+    # The pipeline class is defined here; provenance records the bundle source it patches and runs.
+    pipe.pipeline_file = root / "pipeline_moss_image.py"
     pipe.embeddings = OrderedDict()
     pipe.use_transformer()
     return pipe

tests/test_mps_runtime.py

@@ -1,4 +1,6 @@
 """Regression checks for Mac adaptation; no weights or GPU required."""
+import hashlib
+import json
 import tempfile
 import unittest
 from collections import OrderedDict
@@ -7,9 +9,21 @@ from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import torch
+from PIL import Image
 
 import mps_runtime as runtime
 
+MODEL = "OpenMOSS-Team/MOSS-Image-Base"
+
+
+def fake_pipe(**attributes):
+    transformer = SimpleNamespace(config=SimpleNamespace(num_layers=32), dtype=torch.bfloat16)
+    return SimpleNamespace(transformer=transformer, remove_all_hooks=Mock(), **attributes)
+
+
+def fake_output():
+    return SimpleNamespace(images=[Image.new("RGB", (16, 16))])
+
 
 class TimestepTests(unittest.TestCase):
     def test_cpu_patch_preserves_upstream_float64_math(self):
@@ -46,7 +60,8 @@ class StagingTests(unittest.TestCase):
                 if self.text_encoder is None:
                     raise AssertionError("Encoding requires the text encoder")
                 events.append("encode")
-                return torch.zeros(1, 2, 4), torch.ones(1, 2)
+                weight = torch.ones(4, requires_grad=True)  # stands in for text-encoder parameters
+                return torch.zeros(1, 2, 4) * weight, torch.ones(1, 2)
 
             def __call__(self, prompt, **kwargs):
                 if self.transformer is None or self.text_encoder is not None:
@@ -97,6 +112,75 @@ class StagingTests(unittest.TestCase):
         self.assertIn((("prompt-0",), 4096), self.pipe.embeddings)
         self.assertNotIn((("prompt-1",), 4096), self.pipe.embeddings)
 
+    def test_cached_embeddings_hold_no_autograd_graph(self):
+        # A graph would keep the released text encoder's weights and activations alive.
+        embeds, _ = self.pipe.encode_prompt("coffee")
+        self.assertFalse(embeds.requires_grad)
+        self.assertIsNone(embeds.grad_fn)
+
+    def test_load_pipeline_records_the_bundle_pipeline_it_executes(self):
+        class BasePipeline:
+            def __init__(self, **components):
+                self.text_encoder = None
+                self.__dict__.update(components)
+
+        module = SimpleNamespace(
+            MOSSImagePipeline=BasePipeline, MOSSImageTransformer=Mock(), MOSSImageScheduler=Mock()
+        )
+        root = Path("bundle")
+        with patch.object(runtime, "_bundle_root", return_value=root), \
+             patch.object(runtime, "_import_pipeline_module", return_value=module), \
+             patch.object(runtime.AutoTokenizer, "from_pretrained"), \
+             patch.object(runtime.AutoencoderKLFlux2, "from_pretrained"):
+            pipe = runtime.load_pipeline("model")
+        self.assertEqual(pipe.pipeline_file, root / "pipeline_moss_image.py")
+
+
+class RunnerTests(unittest.TestCase):
+    def setUp(self):
+        import gradio_app
+
+        self.app = gradio_app
+        directory = tempfile.TemporaryDirectory()
+        self.addCleanup(directory.cleanup)
+        self.runner = gradio_app.ModelRunner(directory.name, local_files_only=True)
+        self.config = gradio_app.GenerationConfig(prompt="coffee")
+        for target, name in ((gradio_app, "set_scheduler"), (runtime, "empty_cache")):
+            patcher = patch.object(target, name)
+            patcher.start()
+            self.addCleanup(patcher.stop)
+
+    def test_failed_prompt_stage_releases_pipeline_and_reloads_it(self):
+        broken, fresh = fake_pipe(), fake_pipe()
+
+        def run(pipe, *args, **kwargs):
+            if pipe is broken:
+                pipe.transformer = None  # released for the text encoder, which then failed
+                raise RuntimeError("MPS backend out of memory")
+            return fake_output()
+
+        with patch.object(self.app, "load_pipeline", side_effect=[broken, fresh]) as load, \
+             patch.object(self.app, "generate_with_progress", side_effect=run):
+            with self.assertRaisesRegex(RuntimeError, "out of memory"):
+                self.runner.generate(MODEL, self.config)
+            self.assertIsNone(self.runner.pipe)
+            broken.remove_all_hooks.assert_called_once()
+            self.runner.generate(MODEL, self.config)
+        self.assertEqual(load.call_count, 2)
+
+    def test_mps_provenance_hashes_bundle_pipeline_and_wrapper(self):
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory) / "pipeline_moss_image.py"
+            source.write_text("# bundle pipeline\n")
+            with patch.object(self.app, "load_pipeline", return_value=fake_pipe(pipeline_file=source)), \
+                 patch.object(self.app, "generate_with_progress", return_value=fake_output()):
+                _, metadata, _ = self.runner.generate(MODEL, self.config)
+        record = json.loads(Path(metadata).read_text())
+        self.assertEqual(record["pipeline_sha256"], hashlib.sha256(b"# bundle pipeline\n").hexdigest())
+        self.assertEqual(
+            record["mps_runtime_sha256"], hashlib.sha256(Path(runtime.__file__).read_bytes()).hexdigest()
+        )
+
 
 class IntegrationTests(unittest.TestCase):
     def test_mac_routing_and_unsupported_compile(self):