From 47b02d2defa65dbb538a2a01c5dbd725805cc624 Mon Sep 17 00:00:00 2001 From: Evanlin <166350005+EdwinjJ1@users.noreply.github.com> Date: Sun, 18 Oct 2026 19:44:18 +0800 Subject: [PATCH] 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 --- gradio_app.py | 14 +++++-- mps_runtime.py | 6 ++- tests/test_mps_runtime.py | 86 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 101 insertions(+), 5 deletions(-) diff --git a/gradio_app.py b/gradio_app.py index b7aaf1d..3bdf3f1 100644 --- a/gradio_app.py +++ b/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 diff --git a/mps_runtime.py b/mps_runtime.py index cf60528..795d9cd 100644 --- a/mps_runtime.py +++ b/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 diff --git a/tests/test_mps_runtime.py b/tests/test_mps_runtime.py index 2fe7370..87f0d9c 100644 --- a/tests/test_mps_runtime.py +++ b/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):