diff --git a/comfy/ldm/cosmos/predict2.py b/comfy/ldm/cosmos/predict2.py index d391d50b1..70a835cd2 100644 --- a/comfy/ldm/cosmos/predict2.py +++ b/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, diff --git a/comfy/ldm/lightricks/model.py b/comfy/ldm/lightricks/model.py index 9953b6679..9266c3c3c 100644 --- a/comfy/ldm/lightricks/model.py +++ b/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) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index e6500cff4..bd6a08525 100644 --- a/comfy/ldm/modules/attention.py +++ b/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) diff --git a/comfy/ldm/pixart/blocks.py b/comfy/ldm/pixart/blocks.py index 2225076e5..58d838613 100644 --- a/comfy/ldm/pixart/blocks.py +++ b/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) diff --git a/comfy/ldm/wan/ar_model.py b/comfy/ldm/wan/ar_model.py index d72f53602..06c9da784 100644 --- a/comfy/ldm/wan/ar_model.py +++ b/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, diff --git a/comfy/ldm/wan/model.py b/comfy/ldm/wan/model.py index 1c9782a38..8d7fa0bfa 100644 --- a/comfy/ldm/wan/model.py +++ b/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: diff --git a/comfy/model_management.py b/comfy/model_management.py index 222005b6f..e3b43636b 100644 --- a/comfy/model_management.py +++ b/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: diff --git a/comfy/ops.py b/comfy/ops.py index 13c2604fb..d48589f6a 100644 --- a/comfy/ops.py +++ b/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: diff --git a/comfy/sd.py b/comfy/sd.py index 9d7fa731f..ae97a7ff1 100644 --- a/comfy/sd.py +++ b/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: diff --git a/comfy/utils.py b/comfy/utils.py index 61c2a22dd..febbb3396 100644 --- a/comfy/utils.py +++ b/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) diff --git a/comfy_api/latest/__init__.py b/comfy_api/latest/__init__.py index 294ad425e..3d81652b5 100644 --- a/comfy_api/latest/__init__.py +++ b/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 diff --git a/comfy_api/latest/_ui.py b/comfy_api/latest/_ui.py index b48713d41..b13764948 100644 --- a/comfy_api/latest/_ui.py +++ b/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) diff --git a/comfy_api/torch_helpers/torch_compile.py b/comfy_api/torch_helpers/torch_compile.py index 9223f58db..a935a3e75 100644 --- a/comfy_api/torch_helpers/torch_compile.py +++ b/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: diff --git a/tests-unit/comfy_api_test/torch_compile_test.py b/tests-unit/comfy_api_test/torch_compile_test.py new file mode 100644 index 000000000..e9a17bb34 --- /dev/null +++ b/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 diff --git a/tests-unit/comfy_test/attention_context_test.py b/tests-unit/comfy_test/attention_context_test.py new file mode 100644 index 000000000..df568eed4 --- /dev/null +++ b/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 diff --git a/tests-unit/comfy_test/utils_image_conversion_test.py b/tests-unit/comfy_test/utils_image_conversion_test.py new file mode 100644 index 000000000..8b0072797 --- /dev/null +++ b/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