darask0

/

MOSS-Image_4-step_adapter_lora

Like 33

Hi, thanks for the node — it works well on CUDA. I hit one hard blocker when running it on
Intel Arc (torch.xpu) and wanted to offer a small patch.

Problem

_timestep_embedding (nodes.py L105–106) hardcodes torch.float64:

freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float64, device=timestep.device) / half)
phase = (timestep.double() * 1000.0)[:, None] * freqs

Devices without an fp64 unit cannot execute this. On Intel XPU and Apple MPS the graph dies at
the first sampler iteration:

[ERROR] !!! Exception during processing !!! Required aspect fp64 is not supported on the device

  File "custom_nodes/rapid-moss/nodes.py", line 231, in sample_dataspace
    denoised = model(x_in, s * s_in, **extra_args).float()
  ...
  File "custom_nodes/rapid-moss/nodes.py", line 105, in _timestep_embedding
    freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float64, device=timestep.device) / half)
RuntimeError: Required aspect fp64 is not supported on the device

Everything else in the node works fine on XPU — model loading, LoRA patching, the fp32 time path,
and the data-space sampler all execute. This single line is the only blocker.

Environment: Intel Arc A770 16 GB (DG2-512) · torch 2.14.0+xpu · ComfyUI 0.42.1 (63be2f9b)
· Windows 11 · base Comfy-Org/MOSS-Image → moss_image_sft_bf16.safetensors · adapter
lodestones/MOSS-Image-SFT-8-steps-dmd (lora_latest @ 0.7 + time_embedding + dataspace)

Proposed fix

Probe each device once, cache the result, and only take the fp64 path where it actually works.
CUDA / CPU behaviour is bit-for-bit unchanged — the probe returns True there.

_FP64_SUPPORT: dict = {}


def _supports_fp64(device) -> bool:
    """True if the device can actually run float64 kernels."""
    key = str(device)
    if key not in _FP64_SUPPORT:
        try:
            torch.zeros(1, dtype=torch.float64, device=device) * 1.0
            _FP64_SUPPORT[key] = True
        except Exception:
            _FP64_SUPPORT[key] = False
    return _FP64_SUPPORT[key]


def _timestep_embedding(timestep, dim=256):
    half = dim // 2
    device = timestep.device
    if _supports_fp64(device):
        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float64, device=device) / half)
        phase = (timestep.double() * 1000.0)[:, None] * freqs
    else:
        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=device) / half)
        phase = (timestep.float() * 1000.0)[:, None] * freqs
    return torch.cat((phase.cos(), phase.sin()), dim=-1).float()

torch.arange(half, ...) only produces exact integers, so the input downcast is bit-exact; the only
difference is in the exp and the subsequent multiply.

Numerical impact (measured, not estimated)

Embedding output compared for this adapter's actual step codes, cat(cos, sin):

Max absolute difference ~4e-5. For reference, bf16 weight representation precision is ~1e-2
— two and a half orders of magnitude coarser. (The relative-error column inflates where the cosine
crosses zero, which is why absolute difference is the meaningful figure.)

For precedent, the same fp64 → fp32 substitution is already the established workaround for the
official moss-image/model.py on Intel Arc — its freq / phase construction is identical.

Validation

Probe behaviour on the reporting machine:

cpu  : fp64 SUPPORTED
xpu  : fp64 NOT SUPPORTED - Required aspect fp64 is not supported on the device

With the patch, the full 8-step data-space DMD path runs end-to-end on the A770:

adapter=lodestones-8step  steps=8  lora_strength=0.7  batch=4  1024x1024
success  170.1 s  (42.5 s/image)
peak VRAM 12.05 GiB of 15.56 GiB  (lowest free 3.51 GiB)

Output is identical to the pre-fork fp32 build (PNG payloads differ only by the 3 extra characters
of a longer node class name in the embedded workflow metadata).

Patch

Against comfyui/rapid-moss/nodes.py @ c88f33bc (14,251 B, sha256 b2496288…), applies with patch -p1:

--- a/comfyui/rapid-moss/nodes.py
+++ b/comfyui/rapid-moss/nodes.py
@@ -99,11 +99,41 @@ def _patch_hf(patched, lora, time_path, strength):
 
 # ---- native ComfyUI MOSS-Image: the whole forward as a DIFFUSION_MODEL wrapper ----
 
+_FP64_SUPPORT: dict = {}
+
+
+def _supports_fp64(device) -> bool:
+    """True if the device can actually run float64 kernels.
+
+    CUDA and CPU can. Intel XPU (Arc A/B-series) and Apple MPS cannot: they raise
+    "Required aspect fp64 is not supported on the device". Probe once per device and cache.
+    """
+    key = str(device)
+    if key not in _FP64_SUPPORT:
+        try:
+            torch.zeros(1, dtype=torch.float64, device=device) * 1.0
+            _FP64_SUPPORT[key] = True
+        except Exception:
+            _FP64_SUPPORT[key] = False
+    return _FP64_SUPPORT[key]
+
+
 def _timestep_embedding(timestep, dim=256):
-    """The native sinusoidal embedding (timestep * 1000, cos | sin) in float64, never rounded to bf16."""
+    """The native sinusoidal embedding (timestep * 1000, cos | sin), never rounded to bf16.
+
+    Computed in float64 where the device supports it. On devices without an fp64 unit
+    (Intel XPU, Apple MPS) it falls back to float32: phase spans [0, 1000], where fp32
+    resolution is ~1.2e-4 -- two orders of magnitude finer than bf16 weight precision
+    (~1e-2) -- so the fallback is absorbed numerically.
+    """
     half = dim // 2
-    freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float64, device=timestep.device) / half)
-    phase = (timestep.double() * 1000.0)[:, None] * freqs
+    device = timestep.device
+    if _supports_fp64(device):
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float64, device=device) / half)
+        phase = (timestep.double() * 1000.0)[:, None] * freqs
+    else:
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=device) / half)
+        phase = (timestep.float() * 1000.0)[:, None] * freqs
     return torch.cat((phase.cos(), phase.sin()), dim=-1).float()

Happy to rework it if you'd prefer a different detection strategy (e.g. an explicit env var
override, or keying off torch.device.type). Thanks!