Add native Metal acceleration for Apple Silicon inference SD1.5 First

README.md

@@ -3,6 +3,28 @@ A web interface for Stable Diffusion, implemented using Gradio library.
 
 ![](screenshot.png)
 
+## Apple Silicon Metal enhancements
+
+This branch adds targeted Apple Silicon inference optimizations while retaining the Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, extensions, and PyTorch fallbacks. Stable Diffusion 1.x inference—especially short DPM++ SDE runs—is the primary measured workload.
+
+- Shape-selective Metal Flash Attention for measured SD 1.x FP16 inference shapes, encoded on PyTorch's current MPS command buffer.
+- Unified-memory-aware attention routing with dynamically sized sub-quadratic chunks and streaming online softmax.
+- Native fused FP16 GroupNorm + SiLU and exact-parity GEGLU operations for compatible inference tensors.
+- Modern-PyTorch removal of obsolete MPS clones and unconditional FP32 LayerNorm workarounds.
+- An M1-validated FP16 VAE default with Automatic1111's FP32 NaN retry retained.
+- An opt-in coarse profiler enabled with `A1111_MPS_PROFILE=1 ./webui.sh`.
+- Local benchmarks and focused correctness tests for the Metal paths.
+
+Native paths perform isolated startup checks and fall back to PyTorch for unsupported shapes, dtypes, training, masks, runtime failures, and incompatible configurations. The regular safetensors format remains supported; no checkpoint conversion is required.
+
+On Apple M1, the macOS launcher uses the FP16 VAE path by default. A controlled 512x512, five-step DPM++ SDE/Karras comparison measured approximately 9.16 seconds with FP16 VAE versus 10.37 seconds with `--no-half-vae`, an approximately 1.21 second or 11.7% median improvement. The measured comparison had no NaN fallback; decoded RGB pixels had a mean absolute difference of 0.0227, a maximum difference of 2, and 6.548% changed pixels. Small FP16 rounding differences are expected.
+
+Automatic1111 retains its existing VAE NaN detection and retry behavior: if FP16 decoding produces non-finite output, the VAE is converted to FP32 and decoding is retried. Use `--no-half-vae` to force the conservative FP32 path. The FP16 VAE path was smoke-tested with txt2img, img2img, inpainting, and Hires Fix; external VAE coverage requires an external VAE asset.
+
+Changed Apple Silicon defaults include NGMS 1.0/all steps, Clip skip 2, FP16 sampling without the upstream sampling-upcast default, and the validated FP16 VAE route on M1-family Macs. These settings can change same-seed output compared with upstream defaults. They can be changed through the existing settings or local launch overrides.
+
+A recorded 16 GB M1 comparison at the same checkpoint hash and compute shape improved a five-step 384×640 DPM++ SDE/Karras request from 12.8 seconds to 8.7 seconds. The runs used different seeds, so this is a throughput observation rather than an image-parity comparison.
+
 ## Features
 [Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features):
 - Original txt2img and img2img modes

modules/launch_utils.py

@@ -275,39 +275,35 @@ def run_extensions_installers(settings_file):
                 startup_timer.record(dirname_extension)
 
 
-re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")
-
-
 def requirements_met(requirements_file):
     """
     Does a simple parse of a requirements.txt file to determine if all rerqirements in it
     are already installed. Returns True if so, False if not installed or parsing fails.
     """
 
     import importlib.metadata
-    import packaging.version
+    import packaging.requirements
 
     with open(requirements_file, "r", encoding="utf8") as file:
         for line in file:
-            if line.strip() == "":
+            line = line.split("#", 1)[0].strip()
+            if not line:
                 continue
 
-            m = re.match(re_requirement, line)
-            if m is None:
+            try:
+                requirement = packaging.requirements.Requirement(line)
+            except packaging.requirements.InvalidRequirement:
                 return False
 
-            package = m.group(1).strip()
-            version_required = (m.group(2) or "").strip()
-
-            if version_required == "":
+            if requirement.marker is not None and not requirement.marker.evaluate():
                 continue
 
             try:
-                version_installed = importlib.metadata.version(package)
-            except Exception:
+                version_installed = importlib.metadata.version(requirement.name)
+            except importlib.metadata.PackageNotFoundError:
                 return False
 
-            if packaging.version.parse(version_required) != packaging.version.parse(version_installed):
+            if requirement.specifier and version_installed not in requirement.specifier:
                 return False
 
     return True
@@ -427,17 +423,39 @@ def ensure_build_dependencies():
         """Ensure essential build tools are available"""
         if not is_installed("wheel"):
             run_pip("install wheel", "wheel")
-        # Check setuptools version compatibility
+        # Keep the historical A1111 pin for compatibility with old packages.
         try:
-            setuptools_version = run(f'"{python}" -c "import setuptools; print(setuptools.__version__)"', None, None).strip()
-            if setuptools_version >= "70":
+            setuptools_version = importlib.metadata.version("setuptools")
+            setuptools_match = re.match(r"(\d+)\.(\d+)", setuptools_version)
+            setuptools_tuple = tuple(map(int, setuptools_match.groups())) if setuptools_match else (0, 0)
+            if setuptools_tuple >= (70, 0):
                 run_pip("install setuptools==69.5.1", "setuptools")
         except Exception:
-            # If setuptools check fails, install compatible version
-            run_pip("install setuptools==69.5.1", "setuptools")
+            print("Could not validate setuptools compatibility; leaving the installed version unchanged.")
     # Install build dependencies early
     ensure_build_dependencies()
 
+    if platform.system() == "Darwin" and platform.machine() == "arm64":
+        mps_flash_installer = os.environ.get("MPS_FLASH_ATTENTION_INSTALLER", "")
+        torch_version_match = re.match(r"(\d+)\.(\d+)", importlib.metadata.version("torch"))
+        torch_version = tuple(map(int, torch_version_match.groups())) if torch_version_match else (0, 0)
+        has_current_mps_flash = check_run_python(
+            "import metal_flash_sdpa; "
+            "assert getattr(metal_flash_sdpa, 'A1111_MPS_DEFERRED_COMMIT', False); "
+            "assert getattr(metal_flash_sdpa, 'A1111_MPS_FUSED_GROUP_NORM_SILU', False)"
+        )
+        if mps_flash_installer and torch_version >= (2, 3) and not has_current_mps_flash:
+            try:
+                run(
+                    f'"{python}" "{mps_flash_installer}"',
+                    "Installing stream-safe Metal Flash Attention",
+                    "Couldn't install Metal Flash Attention",
+                    live=True,
+                )
+                startup_timer.record("install Metal Flash Attention")
+            except RuntimeError as exc:
+                print(f"Metal Flash Attention installation failed; continuing with native MPS attention: {exc}")
+
     if not is_installed("clip"):
         run_pip(f"install --no-build-isolation {clip_package}", "clip")
         startup_timer.record("install clip")

modules/mac_specific.py

@@ -1,4 +1,5 @@
 import logging
+import os
 
 import torch
 from torch import Tensor
@@ -30,6 +31,12 @@ def check_for_mps() -> bool:
 has_mps = check_for_mps()
 
 
+def legacy_mps_workaround_required(fixed_in: str) -> bool:
+    """Keep old MPS safety copies off runtimes where the underlying bug is fixed."""
+    force_legacy = os.environ.get("A1111_MPS_FORCE_LEGACY_OPS") == "1"
+    return force_legacy or version.parse(torch.__version__) < version.parse(fixed_in)
+
+
 def torch_mps_gc() -> None:
     try:
         if shared.state.current_latent is not None:
@@ -84,10 +91,15 @@ def interpolate_with_fp32_fallback(orig_func, *args, **kwargs) -> Tensor:
         cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs)
         CondFunc('torch.cumsum', cumsum_fix_func, None)
         CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None)
-        CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)
+
+        # Early MPS builds could crash when a large narrow() view was consumed.
+        # Current runtimes handle the view correctly, so avoid cloning every slice.
+        if legacy_mps_workaround_required("2.1"):
+            CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)
 
         # MPS workaround for https://github.com/pytorch/pytorch/issues/96113
-        CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps')
+        if legacy_mps_workaround_required("2.0.1"):
+            CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps')
 
         # MPS workaround for https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14046
         CondFunc('torch.nn.functional.interpolate', interpolate_with_fp32_fallback, None)

modules/mps_flash_attention.py

@@ -0,0 +1,221 @@
+"""Selective Draw Things-style Metal Flash Attention for Apple Silicon."""
+
+from __future__ import annotations
+
+import importlib
+import os
+import platform
+import subprocess
+import sys
+
+import torch
+import torch.nn.functional as F
+
+
+MIN_TORCH_VERSION = (2, 3)
+
+_extension = None
+_availability = None
+_availability_error = None
+_dispatch_count = 0
+_fallback_count = 0
+_runtime_failure_warned = False
+_first_dispatch_logged = False
+_native_sdpa_supports_gqa = "enable_gqa" in (F.scaled_dot_product_attention.__doc__ or "")
+
+
+def _version_tuple(version):
+    components = version.split("+", 1)[0].split(".")
+    try:
+        return tuple(int(component) for component in components[:2])
+    except ValueError:
+        return (0, 0)
+
+
+def should_use_mfa_shape(query_tokens, key_tokens, head_dim):
+    """Return whether the measured M1 routing table favors MFA for this shape."""
+    if query_tokens < 192:
+        return False
+
+    # These are the SD 1.x UNet head dimensions. With MFA encoded on the
+    # current command buffer, measured self- and cross-attention shapes all
+    # beat PyTorch SDPA without introducing a submission per attention call.
+    return head_dim in (40, 80, 160)
+
+
+def _run_isolated_self_test():
+    code = """
+import torch
+import torch.nn.functional as F
+from metal_flash_sdpa import MetalFlashAttentionForward, fused_geglu_forward, fused_group_norm_silu_add_embedding_forward, fused_group_norm_silu_forward
+
+torch.manual_seed(1)
+source = torch.randn((1, 256, 320), device='mps', dtype=torch.float16)
+q = source.view(1, 256, 8, 40).transpose(1, 2)
+k = q.clone()
+v = q.clone()
+projection = torch.randn((320, 320), device='mps', dtype=torch.float16)
+expected = F.scaled_dot_product_attention(q, k, v)
+expected = F.linear(expected.transpose(1, 2).reshape(1, 256, 320), projection)
+torch.mps.synchronize()
+actual = MetalFlashAttentionForward.apply(q, k, v, 40 ** -0.5, False)
+actual = F.linear(actual.transpose(1, 2).reshape(1, 256, 320), projection)
+torch.mps.synchronize()
+assert torch.isfinite(actual).all().item()
+assert (actual.float() - expected.float()).abs().max().item() < 0.05
+
+norm_source = torch.randn((1, 320, 48, 80), device='mps', dtype=torch.float16)
+norm_weight = torch.randn((320,), device='mps', dtype=torch.float16)
+norm_bias = torch.randn((320,), device='mps', dtype=torch.float16)
+expected_norm = F.silu(F.group_norm(norm_source, 32, norm_weight, norm_bias, 1e-5))
+actual_norm = fused_group_norm_silu_forward(norm_source, norm_weight, norm_bias, 32, 1e-5)
+actual_norm = actual_norm + 0
+torch.mps.synchronize()
+norm_difference = (actual_norm.float() - expected_norm.float()).abs()
+assert torch.isfinite(actual_norm).all().item()
+assert norm_difference.max().item() < 0.02
+assert norm_difference.mean().item() < 0.001
+
+embedding = torch.randn((1, 320), device='mps', dtype=torch.float16)
+expected_embedding_norm = F.silu(F.group_norm(norm_source + embedding[:, :, None, None], 32, norm_weight, norm_bias, 1e-5))
+actual_embedding_norm = fused_group_norm_silu_add_embedding_forward(norm_source, embedding, norm_weight, norm_bias, 32, 1e-5) + 0
+torch.mps.synchronize()
+embedding_difference = (actual_embedding_norm.float() - expected_embedding_norm.float()).abs()
+assert torch.isfinite(actual_embedding_norm).all().item()
+assert embedding_difference.max().item() < 0.02
+assert embedding_difference.mean().item() < 0.001
+
+import numpy as np
+geglu_source = torch.randn((1, 256, 2560), device='mps', dtype=torch.float16)
+half_values = np.arange(65536, dtype=np.uint16).view(np.float16).copy()
+geglu_lut = F.gelu(torch.from_numpy(half_values).to('mps')).contiguous()
+value, gate = geglu_source.chunk(2, dim=-1)
+expected_geglu = value * F.gelu(gate)
+actual_geglu = fused_geglu_forward(geglu_source, geglu_lut) + 0
+torch.mps.synchronize()
+assert torch.isfinite(actual_geglu).all().item()
+assert torch.equal(actual_geglu, expected_geglu)
+"""
+    environment = os.environ.copy()
+    environment["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+    result = subprocess.run(
+        [sys.executable, "-c", code],
+        capture_output=True,
+        text=True,
+        timeout=90,
+        env=environment,
+    )
+    if result.returncode != 0:
+        detail = result.stderr.strip().splitlines()
+        raise RuntimeError(detail[-1] if detail else f"native self-test exited {result.returncode}")
+
+
+def is_available():
+    """Load and crash-test the optional native extension once."""
+    global _extension, _availability, _availability_error
+    if _availability is not None:
+        return _availability
+
+    if platform.system() != "Darwin" or platform.machine() != "arm64":
+        _availability = False
+        return False
+    if not torch.backends.mps.is_available():
+        _availability = False
+        return False
+    if _version_tuple(torch.__version__) < MIN_TORCH_VERSION:
+        _availability_error = f"requires PyTorch {MIN_TORCH_VERSION[0]}.{MIN_TORCH_VERSION[1]} or newer"
+        _availability = False
+        return False
+
+    try:
+        _extension = importlib.import_module("metal_flash_sdpa")
+        if not hasattr(_extension, "MetalFlashAttentionForward"):
+            raise RuntimeError("extension does not expose MetalFlashAttentionForward")
+        if _version_tuple(torch.__version__) < (2, 11) and not getattr(_extension, "A1111_MPS_DEFERRED_COMMIT", False):
+            raise RuntimeError("native extension is missing the A1111 deferred MPS commit patch")
+        if not getattr(_extension, "A1111_MPS_FUSED_GROUP_NORM_SILU", False):
+            raise RuntimeError("native extension is missing fused GroupNorm+SiLU")
+        if not getattr(_extension, "A1111_MPS_FUSED_GROUP_NORM_SILU_EMBEDDING", False):
+            raise RuntimeError("native extension is missing fused GroupNorm+SiLU+embedding")
+        if not getattr(_extension, "A1111_MPS_FUSED_GEGLU", False):
+            raise RuntimeError("native extension is missing fused GEGLU")
+        _run_isolated_self_test()
+    except (ImportError, OSError, RuntimeError, subprocess.SubprocessError) as exc:
+        _availability_error = str(exc)
+        _availability = False
+        print(f"Metal Flash Attention unavailable: {_availability_error}")
+        return False
+
+    _availability = True
+    print("Metal self-test passed; deferred MFA, fused GroupNorm+SiLU, and fused GEGLU routing enabled.")
+    return True
+
+
+def _can_dispatch(query, key, value, attn_mask, dropout_p, enable_gqa, training):
+    if not is_available() or training or enable_gqa or dropout_p != 0.0 or attn_mask is not None:
+        return False
+    if query.device.type != "mps" or key.device != query.device or value.device != query.device:
+        return False
+    if query.dtype != torch.float16 or key.dtype != query.dtype or value.dtype != query.dtype:
+        return False
+    if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
+        return False
+    if query.shape[0] != key.shape[0] or key.shape != value.shape:
+        return False
+    if query.shape[1] != key.shape[1] or query.shape[-1] != key.shape[-1]:
+        return False
+    if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
+        return False
+    return should_use_mfa_shape(query.shape[-2], key.shape[-2], query.shape[-1])
+
+
+def scaled_dot_product_attention(
+    query,
+    key,
+    value,
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False,
+    scale=None,
+    enable_gqa=False,
+    *,
+    training=False,
+):
+    """Use native MFA for measured winning shapes, otherwise use PyTorch SDPA."""
+    global _dispatch_count, _fallback_count, _runtime_failure_warned, _first_dispatch_logged
+    if _can_dispatch(query, key, value, attn_mask, dropout_p, enable_gqa, training):
+        try:
+            _dispatch_count += 1
+            if not _first_dispatch_logged:
+                print(f"Metal Flash Attention first dispatch: Q={tuple(query.shape)}, K={tuple(key.shape)}")
+                _first_dispatch_logged = True
+            attention_scale = scale if scale is not None else query.shape[-1] ** -0.5
+            return _extension.MetalFlashAttentionForward.apply(query, key, value, attention_scale, is_causal)
+        except RuntimeError as exc:
+            if not _runtime_failure_warned:
+                print(f"Metal Flash Attention failed; using PyTorch SDPA: {exc}")
+                _runtime_failure_warned = True
+            torch.mps.empty_cache()
+
+    _fallback_count += 1
+    native_kwargs = {
+        "attn_mask": attn_mask,
+        "dropout_p": dropout_p,
+        "is_causal": is_causal,
+        "scale": scale,
+    }
+    if enable_gqa:
+        if not _native_sdpa_supports_gqa:
+            raise RuntimeError("grouped-query attention requires a newer PyTorch SDPA runtime")
+        native_kwargs["enable_gqa"] = True
+
+    return F.scaled_dot_product_attention(query, key, value, **native_kwargs)
+
+
+def diagnostics():
+    return {
+        "available": bool(_availability),
+        "error": _availability_error,
+        "dispatches": _dispatch_count,
+        "fallbacks": _fallback_count,
+    }