From 3d433bfba631b78379750369f3a030bed1fcd518 Mon Sep 17 00:00:00 2001
From: Evanlin <166350005+EdwinjJ1@users.noreply.github.com>
Date: Sun, 18 Oct 2026 17:35:04 +0800
Subject: [PATCH 1/2] Add experimental Apple silicon MPS inference support

---
 .gitignore                  |   2 -
 README.md                   |   4 +
 README_ZH.md                |   4 +
 comfyui/moss-image/model.py |   8 +-
 constraints-mac.txt         |   6 ++
 docs/macos.md               |  83 ++++++++++++++++++
 docs/macos_zh.md            |  69 +++++++++++++++
 gradio_app.py               |  21 +++--
 mps_runtime.py              | 163 ++++++++++++++++++++++++++++++++++++
 pyproject.toml              |   2 +-
 scripts/gradio-mac.sh       |  17 ++++
 scripts/setup-mac.sh        |  15 ++++
 tests/test_mps_runtime.py   | 137 ++++++++++++++++++++++++++++++
 13 files changed, 520 insertions(+), 11 deletions(-)
 create mode 100644 constraints-mac.txt
 create mode 100644 docs/macos.md
 create mode 100644 docs/macos_zh.md
 create mode 100644 mps_runtime.py
 create mode 100755 scripts/gradio-mac.sh
 create mode 100644 scripts/setup-mac.sh
 create mode 100644 tests/test_mps_runtime.py

diff --git a/.gitignore b/.gitignore
index 01031e2..4f2cee3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,10 +16,8 @@ outputs/
 .env.*
 *credential*
 github_ssh_public_key.txt
-/tests/
 /.github/workflows/tests.yml
 /examples/
-/docs/
 /export_diffusers.py
 .ruff_cache/
 /ckpt/
diff --git a/README.md b/README.md
index 9d28027..c019c89 100644
--- a/README.md
+++ b/README.md
@@ -139,6 +139,10 @@ source .venv/bin/activate
 pip install -r requirements.txt
 ```
 
+### Apple Silicon / MPS
+
+[Apple silicon Mac setup and limitations](docs/macos.md) · M1 Max / 32GB
+
 <a id="python"></a>
 
 ### 🐍 Python
diff --git a/README_ZH.md b/README_ZH.md
index 75090f9..033368a 100644
--- a/README_ZH.md
+++ b/README_ZH.md
@@ -134,6 +134,10 @@ source .venv/bin/activate
 pip install -r requirements.txt
 ```
 
+### Apple Silicon / MPS
+
+[Apple Silicon Mac 安装与限制](docs/macos_zh.md) · M1 Max / 32GB
+
 <a id="python"></a>
 
 ### 🐍 Python
diff --git a/comfyui/moss-image/model.py b/comfyui/moss-image/model.py
index 03f24a9..8be4322 100644
--- a/comfyui/moss-image/model.py
+++ b/comfyui/moss-image/model.py
@@ -199,9 +199,11 @@ class MOSSImageTransformer(nn.Module):
         times = times.expand(batch)
         # Training computes the sinusoid in float64 and feeds float32 into its MLP.
         half = self.freq_dim // 2
-        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
-        phase = times.double()[:, None] * freq
-        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        # MPS has no float64; keep the training sinusoid precision on CPU.
+        sinusoid_device = torch.device("cpu") if times.device.type == "mps" else times.device
+        freq = 10000 ** (-torch.arange(half, device=sinusoid_device, dtype=torch.float64) / half)
+        phase = times.to(sinusoid_device).double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float().to(times.device)
         time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
         mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(
             batch, 6, -1
diff --git a/constraints-mac.txt b/constraints-mac.txt
new file mode 100644
index 0000000..54196aa
--- /dev/null
+++ b/constraints-mac.txt
@@ -0,0 +1,6 @@
+# Versions used on the maintainer's M1 Max / 32 GB Mac.
+torch==2.14.1
+diffusers==0.41.0
+transformers==5.20.1
+accelerate==1.15.0
+gradio==6.30.0
diff --git a/docs/macos.md b/docs/macos.md
new file mode 100644
index 0000000..e98844d
--- /dev/null
+++ b/docs/macos.md
@@ -0,0 +1,83 @@
+# Experimental Apple silicon inference (MPS)
+
+[简体中文](macos_zh.md)
+
+This optional path runs MOSS-Image through PyTorch Metal/MPS. Recorded real-weight validation
+is on an **M1 Max / 32 GB unified-memory Mac**, with SFT. Other chips, lower-memory systems,
+Base and older dependency versions need testing. No quantization is applied.
+
+## Install and launch
+
+Use native arm64 Python 3.12 and a compatible macOS/PyTorch installation. Allow about 40 GB
+free disk for one full model bundle (~25 GB), the environment and download files.
+From the cloned repository root:
+
+```bash
+MOSS_PYTHON=python3.12 bash scripts/setup-mac.sh
+.venv/bin/hf download OpenMOSS-Team/MOSS-Image-SFT --local-dir models/MOSS-Image-SFT
+bash scripts/gradio-mac.sh
+```
+
+Open http://127.0.0.1:7860. Start with 512×512, 10 steps and one image. Keep the terminal
+running; loading and inference can take minutes. Original PNGs and parameters.json are saved
+under `outputs/gradio/`. Stop with Ctrl+C. Subsequent runs only need the launch command.
+The Mac launcher uses local files only and disables Gradio analytics.
+
+`constraints-mac.txt` lists the versions used for the recorded local runs: torch 2.14.1,
+diffusers 0.41.0, transformers 5.20.1, accelerate 1.15.0 and Gradio 6.30.0. These constraints
+apply to the optional Mac setup script, not to the existing installation path.
+
+The launcher defaults the DiT to FP16. To compare BF16:
+
+```bash
+MOSS_MPS_DIT_DTYPE=bfloat16 bash scripts/gradio-mac.sh
+```
+
+The text encoder remains BF16 and the VAE FP32. Direct `mps_runtime.load_pipeline` defaults
+to BF16 unless `MOSS_MPS_DIT_DTYPE` is set before import. CUDA offload modes and Compile DiT
+do not apply to this staged MPS path; the UI disables those controls on MPS.
+
+## Implementation and review points
+
+The Qwen encoder (~17 GB) and DiT (~10 GB) are staged so only one is resident on MPS at a
+time. Positive and negative prompts are encoded before the DiT is loaded. An eight-entry
+LRU prompt cache permits reuse when only seed, dimensions or sampling settings change.
+
+The downloaded custom pipeline computes a timestep sinusoid in FP64, unsupported by MPS.
+The small sinusoid operation is moved to CPU at FP64 and its result copied back. Because
+the standalone pipeline is distributed in the model bundle, this currently uses a narrowly
+matched source replacement and stops explicitly if that source changes. The ComfyUI model
+has the equivalent conditional CPU calculation; the full ComfyUI path was not validated here.
+
+A locally observed loading deadlock involved pending Metal copies and cache/synchronization
+calls holding the GIL. The current workaround drains copies through a CPU transfer before
+calling those methods. It wraps `torch.mps.empty_cache` and `torch.mps.synchronize` when MPS
+is available, including other callers in the process; that process-wide effect is an explicit
+review point for this experimental implementation.
+
+## Evidence and limits
+
+Recorded SFT calls on M1 Max / 32 GB:
+
+| Output | Settings | Full recorded generation call |
+| --- | --- | --- |
+| Chinese chalkboard | 1024×1024, 20 steps, BF16, seed 42 | 320.362 s |
+| English character posters | 768×1024, 28 steps, FP16 | 395.150 / 367.596 / 344.072 s |
+
+[Images, original metadata and prompts](https://github.com/EdwinjJ1/moss-image-mac/tree/main/assets/samples)
+are available as validation evidence. These are end-to-end records, not controlled throughput,
+precision-equivalence or peak-memory measurements. CUDA execution is retained but was not
+hardware-regression-tested by this Mac contribution.
+
+CPU regression checks, without downloading weights:
+
+```bash
+python -m pip install -e '.[ui]' -c constraints-mac.txt
+python -m unittest discover -s tests -p test_mps_runtime.py -v
+```
+
+The checks cover staging order, prompt cache eviction, CPU sinusoid arithmetic, source-change
+failure, MPS routing and UI construction. Real MPS generation remains a separate hardware check.
+If MPS is unavailable, verify native arm64 Python and macOS compatibility. For memory pressure,
+reduce resolution/count and close heavy applications. For a source-patch mismatch, report the
+model revision and error rather than bypassing the check.
diff --git a/docs/macos_zh.md b/docs/macos_zh.md
new file mode 100644
index 0000000..62faebf
--- /dev/null
+++ b/docs/macos_zh.md
@@ -0,0 +1,69 @@
+# Apple Silicon / MPS 实验性支持
+
+[English](macos.md)
+
+通过 PyTorch Metal/MPS 在苹果芯片 Mac 上运行。本次有真实权重出图记录的配置为
+**M1 Max / 32GB 统一内存 / SFT**。其他芯片、小内存配置、Base 和更早的依赖组合尚需实测。
+没有做权重量化，32GB 也不是已证明的最低内存门槛。
+
+## 安装与启动
+
+使用原生 arm64 Python 3.12 和兼容的 macOS/PyTorch，建议为一个完整模型、Python 环境和
+下载文件预留约 40GB 磁盘。从仓库根目录执行：
+
+```bash
+MOSS_PYTHON=python3.12 bash scripts/setup-mac.sh
+.venv/bin/hf download OpenMOSS-Team/MOSS-Image-SFT --local-dir models/MOSS-Image-SFT
+bash scripts/gradio-mac.sh
+```
+
+打开 http://127.0.0.1:7860，先用 512×512、10 步、1 张图试跑。保持终端运行，加载和推理可能
+需要数分钟。原图和 `parameters.json` 保存到 `outputs/gradio/`，按 Ctrl+C 停止服务。
+再次使用只需启动脚本，不必重复安装下载。脚本只读本地模型并关闭 Gradio 分析上报。
+
+可选安装脚本使用 `constraints-mac.txt` 中的本地实测版本：torch 2.14.1、diffusers 0.41.0、
+transformers 5.20.1、accelerate 1.15.0、Gradio 6.30.0，不影响现有安装入口的依赖范围。
+
+启动脚本默认 DiT FP16；比较 BF16 时：
+
+```bash
+MOSS_MPS_DIT_DTYPE=bfloat16 bash scripts/gradio-mac.sh
+```
+
+文本编码器使用 BF16，VAE 使用 FP32。直接调用 `mps_runtime.load_pipeline` 时，若导入前未设置
+环境变量，默认 DiT BF16。MPS 页面自动采用分阶段加载，并禁用 CUDA offload 选择和 Compile DiT。
+
+## 实现与待评审项
+
+约 17GB 的 Qwen 编码器和约 10GB 的 DiT 轮流驻留 MPS。先处理正负提示词，释放编码器后再加载
+DiT；最近八组提示词编码采用 LRU 缓存，仅改变 seed、尺寸或步数时可复用。
+
+MPS 不支持时间步正弦编码中的 FP64，因此将这一小段计算放到 CPU，保留精度后传回 MPS。
+独立 pipeline 位于下载的模型包中，当前用严格匹配的源码替换适配；源码变化会明确报错。
+ComfyUI 模型包含相同的条件 CPU 计算修复，但本次未验证完整 ComfyUI 出图链路。
+
+针对本地观察到的模型加载死锁，在清理缓存/同步前，通过回传 CPU 等待 Metal 复制完成。
+当前方案会在 MPS 可用时包装 `torch.mps.empty_cache` 和 `torch.mps.synchronize`，影响进程内其他
+调用者；这是实验方案中明确需要评审的部分。
+
+## 实测与检查
+
+| 样例 | 参数 | 完整生成调用耗时 |
+| --- | --- | --- |
+| 中文黑板 | 1024×1024、20 步、BF16、seed 42 | 320.362 秒 |
+| 三张英文角色海报 | 768×1024、28 步、FP16 | 395.150 / 367.596 / 344.072 秒 |
+
+[真实图片、提示词及原始参数](https://github.com/EdwinjJ1/moss-image-mac/tree/main/assets/samples)
+仅作为验证证据，不代表同条件速度、精度对照或峰值内存测量。本贡献保留 CUDA 路径，但未进行
+CUDA 硬件回归验证。
+
+不下载模型的 CPU 回归检查：
+
+```bash
+python -m pip install -e '.[ui]' -c constraints-mac.txt
+python -m unittest discover -s tests -p test_mps_runtime.py -v
+```
+
+覆盖模型切换、缓存淘汰、时间步运算、源码变更检查、设备路由及页面构建。实际 MPS 出图需另行
+硬件验证。MPS 不可用时检查 arm64 Python 和系统兼容性；内存不足先降尺寸和张数；补丁不匹配时
+请提供模型 revision 和报错，不要绕过检查。
diff --git a/gradio_app.py b/gradio_app.py
index fbcd666..b7aaf1d 100644
--- a/gradio_app.py
+++ b/gradio_app.py
@@ -21,6 +21,7 @@ from uuid import uuid4
 
 import torch
 
+import mps_runtime
 from gradio_presets import (
     PIXELS,
     fit_size,
@@ -217,8 +218,12 @@ def parse_parameters(parameters, version):
 
 
 def load_pipeline(model, offload, compile_transformer=False, local_files_only=False):
+    if mps_runtime.is_available():
+        if compile_transformer:
+            raise ValueError("Compile DiT is unavailable on Apple silicon. Disable it to use staged MPS inference.")
+        return mps_runtime.load_pipeline(model, local_files_only=local_files_only)
     if not torch.cuda.is_available():
-        raise ValueError("A CUDA GPU is required for generation. The interface can be previewed without one.")
+        raise ValueError("A CUDA GPU or Apple silicon is required for generation. The interface can be previewed without one.")
     pipe = load_hf_pipeline(model, local_files_only=local_files_only)
     return configure_pipeline(pipe, offload=offload, compile_transformer=compile_transformer)
 
@@ -240,6 +245,7 @@ class ModelRunner:
             gc.collect()
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
+            mps_runtime.empty_cache()
         return "Model released. It will load again on the next generation."
 
     def generate(self, model, config, progress=None):
@@ -325,7 +331,7 @@ class ModelRunner:
                 "pipeline_module": type(self.pipe).__module__,
                 "parameters": asdict(config),
                 "sampler": "Euler",
-                "weight_dtype": "bfloat16",
+                "weight_dtype": str(self.pipe.transformer.dtype).removeprefix("torch."),
                 "vae_dtype": "float32",
                 "pipeline_sha256": hashlib.sha256(
                     Path(inspect.getfile(type(self.pipe))).read_bytes()
@@ -339,6 +345,7 @@ class ModelRunner:
                     name: importlib.metadata.version(name)
                     for name in ("torch", "diffusers", "transformers", "accelerate")
                 },
+                "device": "cuda" if torch.cuda.is_available() else "mps" if mps_runtime.is_available() else "cpu",
                 "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
                 "elapsed_seconds": round(time.monotonic() - start, 3),
                 "images": [Path(p).name for p in paths],
@@ -376,6 +383,9 @@ def build_app(
             choices.append(model)
     initial = model or (choices[0] if choices else None)
     offload_default, height_default, width_default = profile_defaults(profile)
+    on_mps = mps_runtime.is_available()
+    if on_mps:
+        offload_default = "model"
     pixels_default, ratio_default = size_selection(width_default, height_default)
     min_budget, max_budget = PIXELS[0] / 1048576, PIXELS[-1] / 1048576
     defaults = GenerationConfig()
@@ -528,19 +538,20 @@ def build_app(
                             shift = gr.Number(value=defaults.shift, label="Shift", minimum=0.01)
                         with gr.Row():
                             offload = gr.Dropdown(
-                                choices=[
+                                choices=[("Staged · Apple silicon", "model")] if on_mps else [
                                     ("GPU resident", "none"),
                                     ("Component", "model"),
                                     ("Sequential · fallback", "sequential"),
                                 ],
                                 value=offload_default,
+                                interactive=not on_mps,
                                 label="Memory mode",
                                 scale=2,
                                 min_width=140,
                             )
                             compile_transformer = gr.Checkbox(
                                 value=False,
-                                interactive=offload_default != "sequential",
+                                interactive=not on_mps and offload_default != "sequential",
                                 label="Compile DiT",
                                 scale=0,
                                 min_width=0,
@@ -733,7 +744,7 @@ def build_app(
             )
 
         def update_compile(mode):
-            if mode == "sequential":
+            if on_mps or mode == "sequential":
                 return gr.update(value=False, interactive=False)
             return gr.update(interactive=True)
 
diff --git a/mps_runtime.py b/mps_runtime.py
new file mode 100644
index 0000000..cf60528
--- /dev/null
+++ b/mps_runtime.py
@@ -0,0 +1,163 @@
+"""Apple-silicon (MPS) loading that keeps one large model on the GPU at a time.
+
+The Qwen text encoder (~17 GB) and the DiT (~10 GB) do not fit in a 32 GB Mac's
+GPU working set together. The text encoder is loaded from disk only to encode a
+new prompt and is released before sampling; the DiT then comes back. Prompt
+embeddings are cached, so a new seed, size or step count skips the text encoder.
+"""
+from __future__ import annotations
+
+import gc
+import importlib.util
+import os
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+from diffusers import AutoencoderKLFlux2
+from huggingface_hub import snapshot_download
+from transformers import AutoModel, AutoTokenizer
+
+DEVICE = "mps"
+CACHED_PROMPTS = 8
+# MOSS_MPS_DIT_DTYPE=float16 runs the DiT in fp16, which M1/M2 GPUs compute natively.
+DIT_DTYPE = getattr(torch, os.environ.get("MOSS_MPS_DIT_DTYPE", "bfloat16"), None)
+if DIT_DTYPE not in (torch.float16, torch.bfloat16):
+    raise ValueError("MOSS_MPS_DIT_DTYPE must be float16 or bfloat16.")
+
+# MPS has no float64. The DiT computes its timestep sinusoid in float64, so
+# that one calculation runs on the CPU with the same precision.
+_FLOAT64_TIMESTEP = (
+    "        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)\n"
+    "        phase = times.double()[:, None] * freq\n"
+    "        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()\n"
+)
+_CPU_TIMESTEP = (
+    "        freq = 10000 ** (-torch.arange(half, dtype=torch.float64) / half)\n"
+    "        phase = times.cpu().double()[:, None] * freq\n"
+    "        embedding = torch.cat((phase.cos(), phase.sin()), -1).float().to(times.device)\n"
+)
+
+
+def _wait_without_gil(original):
+    # torch.mps.empty_cache() and synchronize() wait on the GPU while holding the GIL.
+    # After a non_blocking copy from a CPU tensor whose storage belongs to a Python object
+    # (safetensors, frombuffer), Metal's completion handler needs the GIL to release it, so
+    # the two deadlock; Diffusers does exactly this when loading with device_map.
+    # Copying a value back to the CPU waits with the GIL released and drains them first.
+    def wait(*args, **kwargs):
+        torch.zeros(1, device=DEVICE).cpu()
+        return original(*args, **kwargs)
+    return wait
+
+
+if torch.backends.mps.is_available():
+    torch.mps.empty_cache = _wait_without_gil(torch.mps.empty_cache)
+    torch.mps.synchronize = _wait_without_gil(torch.mps.synchronize)
+
+
+def is_available():
+    return not torch.cuda.is_available() and torch.backends.mps.is_available()
+
+
+def empty_cache():
+    gc.collect()
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+
+
+def _bundle_root(model, local_files_only):
+    path = Path(model)
+    if path.is_dir():
+        return path
+    return Path(snapshot_download(model, local_files_only=local_files_only))
+
+
+def _import_pipeline_module(root):
+    path = root / "pipeline_moss_image.py"
+    source = path.read_text()
+    if _FLOAT64_TIMESTEP not in source:
+        raise ValueError(f"{path} changed; update the float64 patch in mps_runtime.py.")
+    name = "pipeline_moss_image_mps"
+    module = importlib.util.module_from_spec(importlib.util.spec_from_loader(name, loader=None, origin=str(path)))
+    sys.modules[name] = module
+    exec(compile(source.replace(_FLOAT64_TIMESTEP, _CPU_TIMESTEP), str(path), "exec"), module.__dict__)
+    return module
+
+
+def _staged_pipeline_class(module):
+    class StagedMOSSImagePipeline(module.MOSSImagePipeline):
+        """Swaps the text encoder and the DiT on the GPU instead of holding both."""
+
+        @property
+        def _execution_device(self):
+            return torch.device(DEVICE)
+
+        def _release(self, name):
+            if getattr(self, name) is not None:
+                setattr(self, name, None)
+                empty_cache()
+
+        def use_text_encoder(self):
+            if self.text_encoder is None:
+                self._release("transformer")
+                self.text_encoder = AutoModel.from_pretrained(
+                    self.root / "text_encoder", dtype=torch.bfloat16, device_map=DEVICE
+                ).eval()
+
+        def use_transformer(self):
+            if self.transformer is None:
+                self._release("text_encoder")
+                self.transformer = module.MOSSImageTransformer.from_pretrained(
+                    self.root / "transformer", torch_dtype=DIT_DTYPE, device_map=DEVICE
+                ).eval()
+
+        def encode_prompt(self, prompt, max_sequence_length=4096):
+            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)
+                while len(self.embeddings) > CACHED_PROMPTS:
+                    self.embeddings.popitem(last=False)
+            self.embeddings.move_to_end(key)
+            return self.embeddings[key]
+
+        def __call__(self, prompt=None, *, negative_prompt="", guidance_scale=3.0, max_sequence_length=4096,
+                     prompt_embeds=None, prompt_attention_mask=None,
+                     negative_prompt_embeds=None, negative_prompt_attention_mask=None, **kwargs):
+            # Encode both prompts before the DiT returns to the GPU.
+            if prompt is not None and prompt_embeds is None:
+                prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+                prompt = None
+            if guidance_scale != 1.0 and negative_prompt_embeds is None and prompt_embeds is not None:
+                batch = prompt_embeds.shape[0]
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(
+                    negatives, max_sequence_length
+                )
+            self.use_transformer()
+            return super().__call__(
+                prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale,
+                max_sequence_length=max_sequence_length,
+                prompt_embeds=prompt_embeds, prompt_attention_mask=prompt_attention_mask,
+                negative_prompt_embeds=negative_prompt_embeds,
+                negative_prompt_attention_mask=negative_prompt_attention_mask, **kwargs,
+            )
+
+    return StagedMOSSImagePipeline
+
+
+def load_pipeline(model, local_files_only=False):
+    root = _bundle_root(model, local_files_only)
+    module = _import_pipeline_module(root)
+    pipe = _staged_pipeline_class(module)(
+        transformer=None,
+        scheduler=module.MOSSImageScheduler.from_pretrained(root / "scheduler"),
+        tokenizer=AutoTokenizer.from_pretrained(root / "tokenizer"),
+        vae=AutoencoderKLFlux2.from_pretrained(root / "vae", torch_dtype=torch.float32).to(DEVICE).eval(),
+    )
+    pipe.root = root
+    pipe.embeddings = OrderedDict()
+    pipe.use_transformer()
+    return pipe
diff --git a/pyproject.toml b/pyproject.toml
index 2d3f4c6..cfce7a9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ ui = ["gradio>=6.30"]
 
 [tool.setuptools]
 license-files = ["LICENSE", "LICENSES/*.txt"]
-py-modules = ["inference_runtime", "gradio_app", "gradio_presets"]
+py-modules = ["inference_runtime", "mps_runtime", "gradio_app", "gradio_presets"]
 packages = ["assets"]
 include-package-data = false
 
diff --git a/scripts/gradio-mac.sh b/scripts/gradio-mac.sh
new file mode 100755
index 0000000..617b398
--- /dev/null
+++ b/scripts/gradio-mac.sh
@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+# Gradio UI on an Apple-silicon Mac, using the bundles downloaded into models/.
+set -euo pipefail
+
+repo_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+export HF_HUB_OFFLINE=1 GRADIO_ANALYTICS_ENABLED=False
+# Use native fp16 DiT arithmetic on M1/M2; set bfloat16 to compare.
+export MOSS_MPS_DIT_DTYPE="${MOSS_MPS_DIT_DTYPE:-float16}"
+if [[ ! -x "${repo_dir}/.venv/bin/python" ]]; then
+  echo "Run bash scripts/setup-mac.sh first." >&2
+  exit 1
+fi
+cd "$repo_dir"
+exec "${repo_dir}/.venv/bin/python" "${repo_dir}/gradio_app.py" \
+  --model-root "${repo_dir}/models" \
+  --model "${repo_dir}/models/MOSS-Image-SFT" \
+  --local-files-only "$@"
diff --git a/scripts/setup-mac.sh b/scripts/setup-mac.sh
new file mode 100644
index 0000000..13719a4
--- /dev/null
+++ b/scripts/setup-mac.sh
@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+# Create an isolated Python environment; model download is a separate step.
+set -euo pipefail
+repo_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$repo_dir"
+if [[ $(uname -s) != Darwin || $(uname -m) != arm64 ]]; then
+  echo "This installer requires an Apple silicon Mac and a native arm64 terminal." >&2
+  exit 1
+fi
+python_bin=${MOSS_PYTHON:-python3}
+"$python_bin" -c 'import sys; assert sys.version_info >= (3, 10), "Python 3.10+ is required; 3.12 is recommended"'
+"$python_bin" -m venv .venv
+.venv/bin/python -m pip install --upgrade pip
+.venv/bin/python -m pip install -e '.[ui]' -c constraints-mac.txt
+.venv/bin/python -c 'import torch; assert torch.backends.mps.is_available(), "PyTorch cannot access MPS. Check macOS and the Python architecture."; print("MPS available. Next: download the SFT bundle using the macOS guide.")'
diff --git a/tests/test_mps_runtime.py b/tests/test_mps_runtime.py
new file mode 100644
index 0000000..2fe7370
--- /dev/null
+++ b/tests/test_mps_runtime.py
@@ -0,0 +1,137 @@
+"""Regression checks for Mac adaptation; no weights or GPU required."""
+import tempfile
+import unittest
+from collections import OrderedDict
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import torch
+
+import mps_runtime as runtime
+
+
+class TimestepTests(unittest.TestCase):
+    def test_cpu_patch_preserves_upstream_float64_math(self):
+        source = (
+            "import torch\nclass MOSSImagePipeline:\n"
+            "    def embedding(self, times, half):\n"
+            + runtime._FLOAT64_TIMESTEP + "        return embedding\n"
+        )
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "pipeline_moss_image.py").write_text(source)
+            module = runtime._import_pipeline_module(root)
+            times = torch.tensor([0.0, 0.125, 999.0], dtype=torch.float32)
+            actual = module.MOSSImagePipeline().embedding(times, 128)
+        phase = times.double()[:, None] * (10000 ** (-torch.arange(128, dtype=torch.float64) / 128))
+        expected = torch.cat((phase.cos(), phase.sin()), -1).float()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    def test_changed_upstream_source_fails_explicitly(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "pipeline_moss_image.py").write_text("# incompatible upstream pipeline\n")
+            with self.assertRaisesRegex(ValueError, "changed; update"):
+                runtime._import_pipeline_module(root)
+
+
+class StagingTests(unittest.TestCase):
+    def setUp(self):
+        self.events = []
+        events = self.events
+
+        class BasePipeline:
+            def encode_prompt(self, prompt, max_sequence_length):
+                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)
+
+            def __call__(self, prompt, **kwargs):
+                if self.transformer is None or self.text_encoder is not None:
+                    raise AssertionError("Sampling must retain only the DiT")
+                events.append("sample")
+                return kwargs
+
+        self.loader = Mock()
+        module = SimpleNamespace(MOSSImagePipeline=BasePipeline, MOSSImageTransformer=self.loader)
+        self.pipe = runtime._staged_pipeline_class(module)()
+        self.pipe.root = Path("unused-model")
+        self.pipe.text_encoder = None
+        self.pipe.transformer = object()
+        self.pipe.embeddings = OrderedDict()
+
+        def load_text(*args, **kwargs):
+            self.assertIsNone(self.pipe.transformer)
+            self.events.append("load-text")
+            return Mock()
+
+        def load_dit(*args, **kwargs):
+            self.assertIsNone(self.pipe.text_encoder)
+            self.events.append("load-dit")
+            return Mock()
+
+        self.loader.from_pretrained.side_effect = load_dit
+        self.text_patch = patch.object(runtime.AutoModel, "from_pretrained", side_effect=load_text)
+        self.cache_patch = patch.object(runtime, "empty_cache")
+        self.text_patch.start()
+        self.cache_patch.start()
+        self.addCleanup(self.text_patch.stop)
+        self.addCleanup(self.cache_patch.stop)
+
+    def test_both_prompts_encoded_before_sampling_and_cached(self):
+        first = self.pipe("coffee", negative_prompt="blur")
+        self.assertEqual(self.events, ["load-text", "encode", "encode", "load-dit", "sample"])
+        self.events.clear()
+        second = self.pipe("coffee", negative_prompt="blur", num_inference_steps=10, width=512)
+        self.assertEqual(self.events, ["sample"])
+        self.assertIs(first["prompt_embeds"], second["prompt_embeds"])
+
+    def test_prompt_cache_evicts_least_recently_used(self):
+        for index in range(runtime.CACHED_PROMPTS):
+            self.pipe.encode_prompt(f"prompt-{index}")
+        self.pipe.encode_prompt("prompt-0")
+        self.pipe.encode_prompt("new")
+        self.assertEqual(len(self.pipe.embeddings), runtime.CACHED_PROMPTS)
+        self.assertIn((("prompt-0",), 4096), self.pipe.embeddings)
+        self.assertNotIn((("prompt-1",), 4096), self.pipe.embeddings)
+
+
+class IntegrationTests(unittest.TestCase):
+    def test_mac_routing_and_unsupported_compile(self):
+        import gradio_app
+
+        with patch.object(runtime, "is_available", return_value=True), patch.object(runtime, "load_pipeline") as load:
+            gradio_app.load_pipeline("model", "model", local_files_only=True)
+            load.assert_called_once_with("model", local_files_only=True)
+            with self.assertRaisesRegex(ValueError, "Compile DiT"):
+                gradio_app.load_pipeline("model", "model", compile_transformer=True)
+
+    def test_cuda_routing_retains_existing_loader_and_configuration(self):
+        import gradio_app
+
+        with patch.object(runtime, "is_available", return_value=False), \
+             patch.object(torch.cuda, "is_available", return_value=True), \
+             patch.object(gradio_app, "load_hf_pipeline") as load, \
+             patch.object(gradio_app, "configure_pipeline") as configure:
+            result = gradio_app.load_pipeline("model", "sequential", local_files_only=True)
+            load.assert_called_once_with("model", local_files_only=True)
+            configure.assert_called_once_with(load.return_value, offload="sequential", compile_transformer=False)
+            self.assertIs(result, configure.return_value)
+
+    def test_gradio_ui_builds_for_mac_without_loading_weights(self):
+        import gradio_app
+
+        with patch.object(runtime, "is_available", return_value=True), patch.object(runtime, "load_pipeline") as load:
+            app = gradio_app.build_app("missing-model-directory")
+            self.assertEqual(app.title, "MOSS-Image")
+            controls = {item["props"].get("label"): item["props"] for item in app.config["components"]}
+            self.assertFalse(controls["Memory mode"]["interactive"])
+            self.assertFalse(controls["Compile DiT"]["interactive"])
+            load.assert_not_called()
+            app.close()
+
+
+if __name__ == "__main__":
+    unittest.main()

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 2/2] 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             | 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):