Merge eeccb878e0 into 1d1099bea0

comfy/ldm/cosmos/predict2.py

@@ -38,7 +38,13 @@ class GPT2FeedForward(nn.Module):
         return x
 
 
-def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
+def torch_attention_op(
+    q_B_S_H_D: torch.Tensor,
+    k_B_S_H_D: torch.Tensor,
+    v_B_S_H_D: torch.Tensor,
+    transformer_options: dict = {},
+    is_self_attention: bool = False,
+) -> torch.Tensor:
     """Computes multi-head attention using PyTorch's native implementation.
 
     This function provides a PyTorch backend alternative to Transformer Engine's attention operation.
@@ -65,7 +71,16 @@ def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H
     q_B_H_S_D = rearrange(q_B_S_H_D, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
     k_B_H_S_D = rearrange(k_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
     v_B_H_S_D = rearrange(v_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
-    return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True, transformer_options=transformer_options)
+    return optimized_attention(
+        q_B_H_S_D,
+        k_B_H_S_D,
+        v_B_H_S_D,
+        in_q_shape[-2],
+        skip_reshape=True,
+        transformer_options=transformer_options,
+        is_self_attention=is_self_attention,
+        attention_token_shape=transformer_options.get("attention_token_shape", tuple(in_q_shape[1:-2])),
+    )
 
 
 class Attention(nn.Module):
@@ -197,7 +212,13 @@ class Attention(nn.Module):
         return q, k, v
 
     def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
-        result = self.attn_op(q, k, v, transformer_options=transformer_options)  # [B, S, H, D]
+        result = self.attn_op(
+            q,
+            k,
+            v,
+            transformer_options=transformer_options,
+            is_self_attention=self.is_selfattn,
+        )  # [B, S, H * D]
         return self.output_dropout(self.output_proj(result))
 
     def forward(
@@ -516,6 +537,8 @@ class Block(nn.Module):
         gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d")
 
         B, T, H, W, D = x_B_T_H_W_D.shape
+        self_attn_options = transformer_options.copy()
+        self_attn_options["attention_token_shape"] = (T, H, W)
 
         def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D):
             return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D
@@ -532,7 +555,7 @@ class Block(nn.Module):
                 rearrange(normalized_x_B_T_H_W_D.to(compute_dtype), "b t h w d -> b (t h w) d"),
                 None,
                 rope_emb=rope_emb_L_1_1_D,
-                transformer_options=transformer_options,
+                transformer_options=self_attn_options,
             ),
             "b (t h w) d -> b t h w d",
             t=T,

comfy/ldm/lightricks/model.py

@@ -454,6 +454,7 @@ class CrossAttention(nn.Module):
 
     def forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
         q = self.to_q(x)
+        is_self_attention = context is None
         context = x if context is None else context
         k = self.to_k(context)
         v = self.to_v(context)
@@ -466,11 +467,11 @@ class CrossAttention(nn.Module):
             k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
 
         if mask is None:
-            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention)
         elif isinstance(mask, GuideAttentionMask):
             out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
         else:
-            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention)
 
         # Apply per-head gating if enabled
         if self.to_gate_logits is not None:
@@ -970,6 +971,12 @@ class LTXBaseModel(torch.nn.Module, ABC):
         merged_args = {**transformer_options, **kwargs}
         x, pixel_coords, additional_args = self._process_input(x, keyframe_idxs, denoise_mask, **merged_args)
         merged_args.update(additional_args)
+        transformer_options.pop("attention_token_grid", None)
+        orig_shape = additional_args.get("orig_shape")
+        if isinstance(x, torch.Tensor) and isinstance(orig_shape, (tuple, list)) and len(orig_shape) == 5:
+            token_grid = tuple(orig_shape[-3:])
+            if math.prod(token_grid) == x.shape[1]:
+                transformer_options["attention_token_grid"] = token_grid
 
         # Prepare timestep and context
         timestep, embedded_timestep, prompt_timestep = self._prepare_timestep(timestep, batch_size, input_dtype, **merged_args)

comfy/ldm/modules/attention.py

@@ -844,6 +844,7 @@ class CrossAttention(nn.Module):
 
     def forward(self, x, context=None, value=None, mask=None, transformer_options={}):
         q = self.to_q(x)
+        is_self_attention = context is None and value is None
         context = default(context, x)
         k = self.to_k(context)
         if value is not None:
@@ -853,9 +854,9 @@ class CrossAttention(nn.Module):
             v = self.to_v(context)
 
         if mask is None:
-            out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention)
         else:
-            out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention)
         return self.to_out(out)
 
 

comfy/ldm/pixart/blocks.py

@@ -296,7 +296,7 @@ class LabelEmbedder(nn.Module):
         Drops labels to enable classifier-free guidance.
         """
         if force_drop_ids is None:
-            drop_ids = torch.rand(labels.shape[0]).cuda() < self.dropout_prob
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
         else:
             drop_ids = force_drop_ids == 1
         labels = torch.where(drop_ids, self.num_classes, labels)
@@ -328,7 +328,7 @@ class CaptionEmbedder(nn.Module):
         Drops labels to enable classifier-free guidance.
         """
         if force_drop_ids is None:
-            drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob
+            drop_ids = torch.rand(caption.shape[0], device=caption.device) < self.uncond_prob
         else:
             drop_ids = force_drop_ids == 1
         caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption)
@@ -363,7 +363,7 @@ class CaptionEmbedderDoubleBr(nn.Module):
         Drops labels to enable classifier-free guidance.
         """
         if force_drop_ids is None:
-            drop_ids = torch.rand(global_caption.shape[0]).cuda() < self.uncond_prob
+            drop_ids = torch.rand(global_caption.shape[0], device=global_caption.device) < self.uncond_prob
         else:
             drop_ids = force_drop_ids == 1
         global_caption = torch.where(drop_ids[:, None], self.embedding, global_caption)

comfy/ldm/wan/ar_model.py

@@ -62,6 +62,7 @@ class CausalWanSelfAttention(nn.Module):
                 v.view(b, s, n * d),
                 heads=self.num_heads,
                 transformer_options=transformer_options,
+                is_self_attention=True,
             )
         else:
             end = kv_cache["end"]
@@ -78,6 +79,8 @@ class CausalWanSelfAttention(nn.Module):
                 kv_cache["v"][:, :new_end].view(b, new_end, n * d),
                 heads=self.num_heads,
                 transformer_options=transformer_options,
+                is_self_attention=True,
+                is_kv_cached_attention=True,
             )
 
         x = self.o(x)
@@ -170,7 +173,7 @@ class CausalWanModel(WanModel):
             device=device, dtype=dtype, operations=operations)
 
     def forward_block(self, x, timestep, context, start_frame,
-                      kv_caches, crossattn_caches, clip_fea=None):
+                      kv_caches, crossattn_caches, clip_fea=None, transformer_options={}):
         """
         Forward one temporal block for autoregressive inference.
 
@@ -191,6 +194,7 @@ class CausalWanModel(WanModel):
 
         x = self.patch_embedding(x.float()).to(x.dtype)
         grid_sizes = x.shape[2:]
+        transformer_options["grid_sizes"] = grid_sizes
         x = x.flatten(2).transpose(1, 2)
 
         # Per-frame time embedding
@@ -212,11 +216,15 @@ class CausalWanModel(WanModel):
         freqs = self.rope_encode(t, h, w, t_start=start_frame, device=x.device, dtype=x.dtype)
 
         # Transformer blocks
+        transformer_options["total_blocks"] = len(self.blocks)
+        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
+            transformer_options["block_index"] = i
             x = block(x, e=e0, freqs=freqs, context=context,
                       context_img_len=context_img_len,
                       kv_cache=kv_caches[i],
-                      crossattn_cache=crossattn_caches[i])
+                      crossattn_cache=crossattn_caches[i],
+                      transformer_options=transformer_options)
 
         # Head
         x = self.head(x, e)
@@ -269,6 +277,7 @@ class CausalWanModel(WanModel):
                 kv_caches=ar_state["kv_caches"],
                 crossattn_caches=ar_state["crossattn_caches"],
                 clip_fea=clip_fea,
+                transformer_options=transformer_options,
             )
 
         return super().forward(x, timestep, context, clip_fea=clip_fea,

comfy/ldm/wan/model.py

@@ -86,6 +86,7 @@ class WanSelfAttention(nn.Module):
             self.v(x).view(b, s, n * d),
             heads=self.num_heads,
             transformer_options=transformer_options,
+            is_self_attention=True,
         )
 
         if "attn1_patch" in patches:

comfy/model_management.py

@@ -395,6 +395,7 @@ def raise_non_oom(e):
 
 XFORMERS_VERSION = ""
 XFORMERS_ENABLED_VAE = True
+ENABLE_PYTORCH_VAE_ON_AMD = "COMFYUI_ENABLE_PYTORCH_VAE_ON_AMD"
 if args.disable_xformers:
     XFORMERS_IS_AVAILABLE = False
 else:
@@ -1628,9 +1629,14 @@ def pytorch_attention_enabled():
 
 def pytorch_attention_enabled_vae():
     if is_amd():
-        return False  # enabling pytorch attention on AMD currently causes crash when doing high res
+        if os.getenv(ENABLE_PYTORCH_VAE_ON_AMD) == "1":
+            return hasattr(torch.nn.functional, "scaled_dot_product_attention")
+        return False  # enabling pytorch attention on AMD can corrupt high-res VAE decode
     return pytorch_attention_enabled()
 
+def pytorch_attention_vae_single_batch():
+    return sys.platform == "win32" and is_amd() and pytorch_attention_enabled_vae()
+
 def pytorch_attention_flash_attention():
     global ENABLE_PYTORCH_ATTENTION
     if ENABLE_PYTORCH_ATTENTION:

comfy/ops.py

@@ -41,7 +41,7 @@ def scaled_dot_product_attention(q, k, v, *args, **kwargs):
 
 
 try:
-    if torch.cuda.is_available() and comfy.model_management.WINDOWS:
+    if torch.cuda.is_available() and comfy.model_management.WINDOWS and comfy.model_management.is_nvidia():
         from torch.nn.attention import SDPBackend, sdpa_kernel
         import inspect
         if "set_priority" in inspect.signature(sdpa_kernel).parameters:

comfy/sd.py

@@ -1106,6 +1106,8 @@ class VAE:
                 free_memory = self.patcher.get_free_memory(self.device)
                 batch_number = int(free_memory / memory_used)
                 batch_number = max(1, batch_number)
+                if model_management.pytorch_attention_vae_single_batch():
+                    batch_number = 1
 
                 # Pre-allocate output for VAEs that support direct buffer writes
                 preallocated = False
@@ -1964,10 +1966,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
     if unet_dtype is None:
         unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)
 
-    if model_config.quant_config is not None:
-        manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
-    else:
-        manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if model_config.clip_vision_prefix is not None:
@@ -2105,10 +2104,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable
     else:
         unet_dtype = dtype
 
-    if model_config.quant_config is not None:
-        manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
-    else:
-        manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if custom_operations is not None:

comfy/utils.py

@@ -1056,11 +1056,17 @@ def bislerp(samples, width, height):
     result = result.reshape(n, h_new, w_new, c).movedim(-1, 1)
     return result.to(orig_dtype)
 
+def image_to_uint8(image):
+    i = image.cpu().numpy()
+    i = np.nan_to_num(i, nan=0.0, posinf=1.0, neginf=0.0)
+    return (np.clip(i, 0, 1) * 255.).astype(np.uint8)
+
+
 def lanczos(samples, width, height):
     #the below API is strict and expects grayscale to be squeezed
     if samples.ndim == 4:
         samples = samples.squeeze(1) if samples.shape[1] == 1 else samples.movedim(1, -1)
-    images = [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in samples]
+    images = [Image.fromarray(image_to_uint8(image)) for image in samples]
     images = [image.resize((width, height), resample=Image.Resampling.LANCZOS) for image in images]
     images = [torch.from_numpy(t).movedim(-1, 0) if (t := np.array(image).astype(np.float32) / 255.0).ndim == 3 else torch.from_numpy(t) for image in images]
     result = torch.stack(images)

comfy_api/latest/__init__.py

@@ -12,7 +12,7 @@ from comfy_execution.utils import get_executing_context
 from comfy_execution.progress import get_progress_state, PreviewImageTuple
 from PIL import Image
 from comfy.cli_args import args
-import numpy as np
+import comfy.utils
 
 
 class ComfyAPI_latest(ComfyAPIBase):
@@ -65,9 +65,7 @@ class ComfyAPI_latest(ComfyAPIBase):
                     if len(tensor.shape) == 4:
                         tensor = tensor[0]
 
-                    # Convert to numpy array and scale to 0-255
-                    image_np = (tensor.cpu().numpy() * 255).astype(np.uint8)
-                    to_display = Image.fromarray(image_np)
+                    to_display = Image.fromarray(comfy.utils.image_to_uint8(tensor))
 
                 if isinstance(to_display, Image.Image):
                     # Detect image format from PIL Image

comfy_api/latest/_ui.py

@@ -7,7 +7,6 @@ import uuid
 from io import BytesIO
 
 import av
-import numpy as np
 import torch
 try:
     import torchaudio
@@ -18,6 +17,7 @@ from PIL import Image as PILImage
 from PIL.PngImagePlugin import PngInfo
 
 import folder_paths
+import comfy.utils
 
 # used for image preview
 from comfy.cli_args import args
@@ -79,7 +79,7 @@ class ImageSaveHelper:
     @staticmethod
     def _convert_tensor_to_pil(image_tensor: torch.Tensor) -> PILImage.Image:
         """Converts a single torch tensor to a PIL Image."""
-        return PILImage.fromarray(np.clip(255.0 * image_tensor.cpu().numpy(), 0, 255).astype(np.uint8))
+        return PILImage.fromarray(comfy.utils.image_to_uint8(image_tensor))
 
     @staticmethod
     def _create_png_metadata(cls: type[ComfyNode] | None) -> PngInfo | None:
@@ -440,8 +440,7 @@ class PreviewUI3D(_UIOutput):
         self.bg_image_path = None
         bg_image = kwargs.get("bg_image", None)
         if bg_image is not None:
-            img_array = (bg_image[0].cpu().numpy() * 255).astype(np.uint8)
-            img = PILImage.fromarray(img_array)
+            img = PILImage.fromarray(comfy.utils.image_to_uint8(bg_image[0]))
             temp_dir = folder_paths.get_temp_directory()
             filename = f"bg_{uuid.uuid4().hex}.png"
             bg_image_path = os.path.join(temp_dir, filename)

comfy_api/torch_helpers/torch_compile.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
+import logging
+import sys
 import torch
 
 import comfy.utils
 from comfy.patcher_extension import WrappersMP
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 if TYPE_CHECKING:
     from comfy.model_patcher import ModelPatcher
     from comfy.patcher_extension import WrapperExecutor
@@ -11,6 +13,35 @@ if TYPE_CHECKING:
 
 COMPILE_KEY = "torch.compile"
 TORCH_COMPILE_KWARGS = "torch_compile_kwargs"
+WINDOWS_ROCM_INDUCTOR_OPTIONS = {
+    "triton.cudagraphs": False,
+    "triton.cudagraph_trees": False,
+}
+
+
+def _is_windows_rocm_inductor(backend: Optional[str]) -> bool:
+    return backend == "inductor" and sys.platform == "win32" and getattr(torch.version, "hip", None) is not None
+
+
+def normalize_torch_compile_kwargs(compile_kwargs: dict[str, Any]) -> dict[str, Any]:
+    compile_kwargs = dict(compile_kwargs)
+    if _is_windows_rocm_inductor(compile_kwargs.get("backend")) and compile_kwargs.get("mode") in (None, "", "default"):
+        options = dict(compile_kwargs.get("options") or {})
+        if set(options) <= {"guard_filter_fn"}:
+            compile_kwargs["mode"] = None
+            compile_kwargs["options"] = None
+            logging.info("torch.compile: using default mode for Windows ROCm inductor.")
+        else:
+            changed = False
+            for key, value in WINDOWS_ROCM_INDUCTOR_OPTIONS.items():
+                if options.get(key) is not value:
+                    options[key] = value
+                    changed = True
+            compile_kwargs["options"] = options
+            compile_kwargs["mode"] = None
+            if changed:
+                logging.info("torch.compile: disabled inductor cudagraphs for Windows ROCm.")
+    return compile_kwargs
 
 
 def apply_torch_compile_factory(compiled_module_dict: dict[str, Callable]) -> Callable:
@@ -30,7 +61,7 @@ def apply_torch_compile_factory(compiled_module_dict: dict[str, Callable]) -> Ca
     return apply_torch_compile_wrapper
 
 
-def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Optional[dict[str,str]]=None,
+def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Optional[dict[str, Any]]=None,
                               mode: Optional[str]=None, fullgraph=False, dynamic: Optional[bool]=None,
                               keys: list[str]=["diffusion_model"], *args, **kwargs):
     '''
@@ -52,6 +83,7 @@ def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Option
         "fullgraph": fullgraph,
         "dynamic": dynamic,
     }
+    compile_kwargs = normalize_torch_compile_kwargs(compile_kwargs)
     # get a dict of compiled keys
     compiled_modules = {}
     for key in keys:

tests-unit/comfy_api_test/torch_compile_test.py

@@ -0,0 +1,50 @@
+import torch
+
+from comfy_api.torch_helpers import torch_compile
+
+
+def test_windows_rocm_default_mode_drops_injected_guard_options(monkeypatch):
+    monkeypatch.setattr(torch_compile.sys, "platform", "win32")
+    monkeypatch.setattr(torch.version, "hip", "7.15", raising=False)
+
+    result = torch_compile.normalize_torch_compile_kwargs(
+        {
+            "backend": "inductor",
+            "mode": "default",
+            "options": {"guard_filter_fn": object()},
+        }
+    )
+
+    assert result["mode"] is None
+    assert result["options"] is None
+
+
+def test_windows_rocm_custom_options_disable_cudagraphs(monkeypatch):
+    monkeypatch.setattr(torch_compile.sys, "platform", "win32")
+    monkeypatch.setattr(torch.version, "hip", "7.15", raising=False)
+
+    result = torch_compile.normalize_torch_compile_kwargs(
+        {
+            "backend": "inductor",
+            "mode": "default",
+            "options": {"max_autotune": True},
+        }
+    )
+
+    assert result["mode"] is None
+    assert result["options"] == {
+        "max_autotune": True,
+        "triton.cudagraphs": False,
+        "triton.cudagraph_trees": False,
+    }
+
+
+def test_non_rocm_compile_options_are_unchanged(monkeypatch):
+    monkeypatch.setattr(torch_compile.sys, "platform", "linux")
+    compile_kwargs = {
+        "backend": "inductor",
+        "mode": "default",
+        "options": {"max_autotune": True},
+    }
+
+    assert torch_compile.normalize_torch_compile_kwargs(compile_kwargs) == compile_kwargs

tests-unit/comfy_test/attention_context_test.py

@@ -0,0 +1,48 @@
+import torch
+
+from comfy.ldm.cosmos import predict2
+from comfy.ldm.wan.ar_model import CausalWanModel
+
+
+def test_cosmos_attention_passes_self_attention_context(monkeypatch):
+    captured = {}
+
+    def capture_attention(q, k, v, heads, **kwargs):
+        captured.update(kwargs)
+        return q
+
+    monkeypatch.setattr(predict2, "optimized_attention", capture_attention)
+    q = torch.zeros((1, 2, 3, 4))
+
+    predict2.torch_attention_op(q, q, q, is_self_attention=True)
+
+    assert captured["is_self_attention"] is True
+    assert captured["attention_token_shape"] == (2,)
+
+
+def test_causal_wan_forward_passes_transformer_options_to_ar_block():
+    transformer_options = {
+        "ar_state": {
+            "start_frame": 2,
+            "kv_caches": [object()],
+            "crossattn_caches": [object()],
+        },
+        "optimized_attention_override": object(),
+    }
+    captured = {}
+
+    class FakeCausalWan:
+        def forward_block(self, **kwargs):
+            captured.update(kwargs)
+            return "result"
+
+    result = CausalWanModel.forward(
+        FakeCausalWan(),
+        x=torch.zeros((1, 4, 3, 2, 2)),
+        timestep=torch.zeros(1),
+        context=torch.zeros((1, 1, 1)),
+        transformer_options=transformer_options,
+    )
+
+    assert result == "result"
+    assert captured["transformer_options"] is transformer_options

tests-unit/comfy_test/utils_image_conversion_test.py

@@ -0,0 +1,36 @@
+import warnings
+
+import numpy as np
+import torch
+
+import comfy.utils
+
+
+def test_image_to_uint8_sanitizes_nonfinite_values_without_runtime_warning():
+    image = torch.tensor(
+        [
+            [
+                [float("nan"), float("inf"), -float("inf")],
+                [-1.0, 0.5, 2.0],
+            ]
+        ],
+        dtype=torch.float32,
+    )
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", RuntimeWarning)
+        result = comfy.utils.image_to_uint8(image)
+
+    assert result.dtype == np.uint8
+    assert result.tolist() == [[[0, 255, 0], [0, 127, 255]]]
+
+
+def test_image_to_uint8_does_not_modify_source_tensor():
+    image = torch.tensor([float("nan"), float("inf"), -float("inf"), 0.5])
+
+    comfy.utils.image_to_uint8(image)
+
+    assert torch.isnan(image[0])
+    assert torch.isinf(image[1])
+    assert torch.isinf(image[2])
+    assert image[3] == 0.5