Merge 45a2363e6a into 25757a53c9

Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>

.github/workflows/stable-release.yml

@@ -145,6 +145,8 @@ jobs:
           cp -r ComfyUI/.ci/windows_${{ inputs.rel_name }}_base_files/* ./
           cp ../update_comfyui_and_python_dependencies.bat ./update/
 
+          echo 'local-portable' > ComfyUI/.comfy_environment
+
           cd ..
 
           "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=9 -mfb=128 -md=768m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable

comfy/float.py

@@ -55,6 +55,11 @@ def stochastic_rounding(value, dtype, seed=0):
     if dtype == torch.bfloat16:
         return value.to(dtype=torch.bfloat16)
     if dtype == torch.float8_e4m3fn or dtype == torch.float8_e5m2:
+        # MPS does not support FP8 dtypes — perform rounding on CPU and return the result there.
+        on_mps = value.device.type == "mps"
+        if on_mps:
+            value = value.cpu()
+
         generator = torch.Generator(device=value.device)
         generator.manual_seed(seed)
         output = torch.empty_like(value, dtype=dtype)
@@ -159,6 +164,12 @@ def stochastic_round_quantize_nvfp4(x, per_tensor_scale, pad_16x, seed=0):
         """Round up x to the nearest multiple."""
         return ((x + multiple - 1) // multiple) * multiple
 
+    # MPS does not support FP8 dtypes used for block scales — perform on CPU.
+    on_mps = x.device.type == "mps"
+    if on_mps:
+        x = x.cpu()
+        per_tensor_scale = per_tensor_scale.cpu() if isinstance(per_tensor_scale, torch.Tensor) else per_tensor_scale
+
     generator = torch.Generator(device=x.device)
     generator.manual_seed(seed)
 
@@ -179,6 +190,12 @@ def stochastic_round_quantize_nvfp4_by_block(x, per_tensor_scale, pad_16x, seed=
         """Round up x to the nearest multiple."""
         return ((x + multiple - 1) // multiple) * multiple
 
+    # MPS does not support FP8 dtypes used for block scales — perform on CPU.
+    on_mps = x.device.type == "mps"
+    if on_mps:
+        x = x.cpu()
+        per_tensor_scale = per_tensor_scale.cpu() if isinstance(per_tensor_scale, torch.Tensor) else per_tensor_scale
+
     orig_shape = x.shape
 
     # Handle padding

comfy/model_patcher.py

@@ -26,6 +26,7 @@ import uuid
 from typing import Callable, Optional
 
 import torch
+import tqdm
 
 import comfy.float
 import comfy.hooks
@@ -1651,7 +1652,11 @@ class ModelPatcherDynamic(ModelPatcher):
                 self.model.model_loaded_weight_memory += casted_buf.numel() * casted_buf.element_size()
 
             force_load_stat = f" Force pre-loaded {len(self.backup)} weights: {self.model.model_loaded_weight_memory // 1024} KB." if len(self.backup) > 0 else ""
-            logging.info(f"Model {self.model.__class__.__name__} prepared for dynamic VRAM loading. {allocated_size // (1024 ** 2)}MB Staged. {num_patches} patches attached.{force_load_stat}")
+            log_key = (self.patches_uuid, allocated_size, num_patches, len(self.backup), self.model.model_loaded_weight_memory)
+            in_loop = bool(getattr(tqdm.tqdm, "_instances", None))
+            level = logging.DEBUG if in_loop and getattr(self, "_last_prepare_log_key", None) == log_key else logging.INFO
+            self._last_prepare_log_key = log_key
+            logging.log(level, f"Model {self.model.__class__.__name__} prepared for dynamic VRAM loading. {allocated_size // (1024 ** 2)}MB Staged. {num_patches} patches attached.{force_load_stat}")
 
             self.model.device = device_to
             self.model.current_weight_patches_uuid = self.patches_uuid

comfy/quant_ops.py

@@ -93,6 +93,12 @@ class _TensorCoreFP8LayoutBase(_CKFp8Layout):
         if not isinstance(scale, torch.Tensor):
             scale = torch.tensor(scale, device=tensor.device, dtype=torch.float32)
 
+        # MPS does not support FP8 dtypes — move to CPU for quantization.
+        on_mps = tensor.device.type == "mps"
+        if on_mps:
+            tensor = tensor.cpu()
+            scale = scale.cpu()
+
         if stochastic_rounding > 0:
             if inplace_ops:
                 tensor *= (1.0 / scale).to(tensor.dtype)

comfy_api_nodes/nodes_grok.py

@@ -54,7 +54,12 @@ class GrokImageNode(IO.ComfyNode):
             inputs=[
                 IO.Combo.Input(
                     "model",
-                    options=["grok-imagine-image-pro", "grok-imagine-image", "grok-imagine-image-beta"],
+                    options=[
+                        "grok-imagine-image-quality",
+                        "grok-imagine-image-pro",
+                        "grok-imagine-image",
+                        "grok-imagine-image-beta",
+                    ],
                 ),
                 IO.String.Input(
                     "prompt",
@@ -111,10 +116,12 @@ class GrokImageNode(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution"]),
                 expr="""
                 (
-                  $rate := $contains(widgets.model, "pro") ? 0.07 : 0.02;
+                  $rate := widgets.model = "grok-imagine-image-quality"
+                    ? (widgets.resolution = "1k" ? 0.05 : 0.07)
+                    : ($contains(widgets.model, "pro") ? 0.07 : 0.02);
                   {"type":"usd","usd": $rate * widgets.number_of_images}
                 )
                 """,
@@ -167,7 +174,12 @@ class GrokImageEditNode(IO.ComfyNode):
             inputs=[
                 IO.Combo.Input(
                     "model",
-                    options=["grok-imagine-image-pro", "grok-imagine-image", "grok-imagine-image-beta"],
+                    options=[
+                        "grok-imagine-image-quality",
+                        "grok-imagine-image-pro",
+                        "grok-imagine-image",
+                        "grok-imagine-image-beta",
+                    ],
                 ),
                 IO.Image.Input("image", display_name="images"),
                 IO.String.Input(
@@ -228,11 +240,19 @@ class GrokImageEditNode(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution"]),
                 expr="""
                 (
-                  $rate := $contains(widgets.model, "pro") ? 0.07 : 0.02;
-                  {"type":"usd","usd": 0.002 + $rate * widgets.number_of_images}
+                  $isQualityModel := widgets.model = "grok-imagine-image-quality";
+                  $isPro := $contains(widgets.model, "pro");
+                  $rate := $isQualityModel
+                    ? (widgets.resolution = "1k" ? 0.05 : 0.07)
+                    : ($isPro ? 0.07 : 0.02);
+                  $base := $isQualityModel ? 0.01 : 0.002;
+                  $output := $rate * widgets.number_of_images;
+                  $isPro
+                    ? {"type":"usd","usd": $base + $output}
+                    : {"type":"range_usd","min_usd": $base + $output, "max_usd": 3 * $base + $output}
                 )
                 """,
             ),

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.43.17
-comfyui-workflow-templates==0.9.69
+comfyui-workflow-templates==0.9.72
 comfyui-embedded-docs==0.4.4
 torch
 torchsde

server.py

@@ -560,7 +560,7 @@ class PromptServer():
                             buffer.seek(0)
 
                             return web.Response(body=buffer.read(), content_type=f'image/{image_format}',
-                                                headers={"Content-Disposition": f"attachment; filename=\"{filename}\""})
+                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
 
                     if 'channel' not in request.rel_url.query:
                         channel = 'rgba'
@@ -580,7 +580,7 @@ class PromptServer():
                             buffer.seek(0)
 
                             return web.Response(body=buffer.read(), content_type='image/png',
-                                                headers={"Content-Disposition": f"attachment; filename=\"{filename}\""})
+                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
 
                     elif channel == 'a':
                         with Image.open(file) as img:
@@ -597,7 +597,7 @@ class PromptServer():
                             alpha_buffer.seek(0)
 
                             return web.Response(body=alpha_buffer.read(), content_type='image/png',
-                                                headers={"Content-Disposition": f"attachment; filename=\"{filename}\""})
+                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
                     else:
                         # Use the content type from asset resolution if available,
                         # otherwise guess from the filename.
@@ -614,7 +614,7 @@ class PromptServer():
                         return web.FileResponse(
                             file,
                             headers={
-                                "Content-Disposition": f"attachment; filename=\"{filename}\"",
+                                "Content-Disposition": f"filename=\"{filename}\"",
                                 "Content-Type": content_type
                             }
                         )

tests/test_fp8_mps.py

@@ -0,0 +1,147 @@
+"""
+Tests for FP8 quantization on MPS (Apple Silicon) devices.
+
+MPS does not natively support float8_e4m3fn or float8_e5m2 dtypes.
+These tests verify that:
+  1. FP8 operations correctly fall back to CPU when on MPS.
+  2. The round-trip (quantize on CPU -> result on original device) is numerically sound.
+  3. No "Placeholder storage has not been allocated on MPS device!" errors occur.
+"""
+import pytest
+import torch
+
+import comfy.float
+from comfy.quant_ops import TensorCoreFP8E4M3Layout, TensorCoreFP8E5M2Layout
+
+# Skip the entire module if MPS is not available
+pytestmark = pytest.mark.skipif(
+    not torch.backends.mps.is_available(),
+    reason="MPS backend not available"
+)
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+def _make_mps_tensor(shape=(256, 256), dtype=torch.float32):
+    return torch.randn(shape, device="mps", dtype=dtype)
+
+
+# ── Tests for comfy.float ────────────────────────────────────────────────────
+
+class TestStochasticRoundingMPS:
+    """Tests for comfy.float.stochastic_rounding on MPS device."""
+
+    def test_stochastic_rounding_fp8_e4m3fn_on_mps(self):
+        """stochastic_rounding must not crash when input is on MPS and target dtype is float8_e4m3fn."""
+        x = _make_mps_tensor((64, 64), dtype=torch.float32)
+        result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e4m3fn, seed=42)
+
+        assert result.dtype == torch.float8_e4m3fn
+        assert result.shape == x.shape
+
+    def test_stochastic_rounding_fp8_e5m2_on_mps(self):
+        """stochastic_rounding must not crash when input is on MPS and target dtype is float8_e5m2."""
+        x = _make_mps_tensor((64, 64), dtype=torch.float32)
+        result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e5m2, seed=42)
+
+        assert result.dtype == torch.float8_e5m2
+        assert result.shape == x.shape
+
+    def test_stochastic_rounding_fp8_result_on_cpu(self):
+        """Result of FP8 rounding from MPS input should be on CPU (since MPS can't hold FP8)."""
+        x = _make_mps_tensor((32, 32), dtype=torch.float32)
+        result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e4m3fn, seed=42)
+
+        # FP8 tensors cannot live on MPS, so result must be on CPU
+        assert result.device.type == "cpu"
+
+    def test_stochastic_rounding_non_fp8_still_works(self):
+        """Non-FP8 dtypes on MPS must still work as before (no regression)."""
+        x = _make_mps_tensor((32, 32), dtype=torch.float32)
+
+        r16 = comfy.float.stochastic_rounding(x, dtype=torch.float16, seed=0)
+        assert r16.dtype == torch.float16
+        assert r16.device.type == "mps"
+
+        rbf16 = comfy.float.stochastic_rounding(x, dtype=torch.bfloat16, seed=0)
+        assert rbf16.dtype == torch.bfloat16
+        assert rbf16.device.type == "mps"
+
+    def test_stochastic_rounding_fp8_numerical_sanity(self):
+        """FP8 round-trip (float32 -> fp8 -> float32) should have bounded error."""
+        x = torch.randn(128, 128, device="mps", dtype=torch.float32)
+        x_clamped = torch.clamp(x, min=-448, max=448)  # FP8 e4m3fn range
+
+        fp8 = comfy.float.stochastic_rounding(x_clamped, dtype=torch.float8_e4m3fn, seed=123)
+        # Convert back to float32 for comparison
+        reconstructed = fp8.to(torch.float32)
+
+        # Max relative error should be bounded (FP8 e4m3fn has ~0.125 relative precision)
+        x_cpu = x_clamped.cpu()
+        max_abs_err = (reconstructed - x_cpu).abs().max().item()
+        # FP8 e4m3fn max value is 448, min subnormal ~0.001953
+        # For random normal data, error should be well under 1.0
+        assert max_abs_err < 2.0, f"FP8 round-trip error too large: {max_abs_err}"
+
+
+class TestManualStochasticRoundMPS:
+    """Tests for comfy.float.manual_stochastic_round_to_float8 on MPS device."""
+
+    def test_manual_round_fp8_on_mps_tensor(self):
+        """stochastic_rounding handles MPS generator internally without 'Placeholder storage' error."""
+        x = _make_mps_tensor((16, 16), dtype=torch.float32)
+        result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e4m3fn, seed=42)
+        assert result.dtype == torch.float8_e4m3fn
+
+
+class TestNVFP4StochasticRoundMPS:
+    """Tests for NVFP4 stochastic rounding on MPS - also creates FP8 tensors internally."""
+
+    def test_nvfp4_stochastic_round_on_mps(self):
+        """stochastic_round_quantize_nvfp4 creates FP8 block scales internally."""
+        # NVFP4 requires 2D input with dimensions divisible by 16
+        x = torch.randn(32, 32, device="mps", dtype=torch.float32)
+        scale = torch.tensor(1.0, device="mps", dtype=torch.float32)
+
+        # This should not crash - internally creates float8_e4m3fn block scales
+        qdata, block_scale = comfy.float.stochastic_round_quantize_nvfp4(
+            x, scale, pad_16x=False, seed=42
+        )
+        assert qdata.dtype == torch.uint8
+
+
+# ── Tests for comfy.quant_ops (integration) ──────────────────────────────────
+
+class TestQuantOpsMPS:
+    """Tests for the quantization ops layer that calls into comfy.float."""
+
+    def test_fp8_layout_quantize_on_mps(self):
+        """TensorCoreFP8E4M3Layout.quantize must work with MPS tensors."""
+        x = _make_mps_tensor((64, 64), dtype=torch.bfloat16)
+        qdata, params = TensorCoreFP8E4M3Layout.quantize(
+            x, scale="recalculate", stochastic_rounding=42
+        )
+
+        assert qdata.dtype == torch.float8_e4m3fn
+        assert params.orig_dtype == torch.bfloat16
+
+    def test_fp8_layout_quantize_without_stochastic_on_mps(self):
+        """TensorCoreFP8E4M3Layout.quantize with stochastic_rounding=0 uses ck.quantize_per_tensor_fp8."""
+        x = _make_mps_tensor((64, 64), dtype=torch.bfloat16)
+        qdata, params = TensorCoreFP8E4M3Layout.quantize(
+            x, scale="recalculate", stochastic_rounding=0
+        )
+
+        assert qdata.dtype == torch.float8_e4m3fn
+
+    def test_fp8_e5m2_layout_quantize_on_mps(self):
+        """TensorCoreFP8E5M2Layout.quantize must work with MPS tensors."""
+        x = _make_mps_tensor((64, 64), dtype=torch.float32)
+        qdata, params = TensorCoreFP8E5M2Layout.quantize(
+            x, scale="recalculate", stochastic_rounding=42
+        )
+
+        assert qdata.dtype == torch.float8_e5m2
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])