From 2bc2e8d73e6686e7ee54e84c7a0416708c4df1da Mon Sep 17 00:00:00 2001
From: Dan Sedano <47065799+thedandano@users.noreply.github.com>
Date: Thu, 13 Aug 2026 12:49:01 -0700
Subject: [PATCH] Add unit tests for the attention-dependent memory estimate

Covers which attention checks select the efficient formula in
BaseModel.memory_required, including the new flash_attention_enabled()
path, and pins the 7.5x bf16 ratio between the conservative and efficient
estimates so a retune of either formula fails the test and gets updated
consciously.
---
 tests-unit/comfy_test/memory_required_test.py | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)
 create mode 100644 tests-unit/comfy_test/memory_required_test.py

diff --git a/tests-unit/comfy_test/memory_required_test.py b/tests-unit/comfy_test/memory_required_test.py
new file mode 100644
index 000000000..76e2d45a3
--- /dev/null
+++ b/tests-unit/comfy_test/memory_required_test.py
@@ -0,0 +1,63 @@
+import math
+
+import torch
+
+import comfy.model_base
+import comfy.model_management
+
+
+class _StubModel:
+    """Minimal stand-in for BaseModel: just the attributes memory_required reads."""
+    memory_usage_factor_conds = ()
+    memory_usage_shape_process = {}
+    memory_usage_factor = 2.0
+
+    def get_dtype_inference(self):
+        return torch.bfloat16
+
+
+INPUT_SHAPE = (1, 16, 1, 180, 320)
+AREA = INPUT_SHAPE[0] * math.prod(INPUT_SHAPE[2:])
+DTYPE_SIZE = 2  # bf16
+EFFICIENT = AREA * DTYPE_SIZE * 0.01 * _StubModel.memory_usage_factor * (1024 * 1024)
+CONSERVATIVE = AREA * 0.15 * _StubModel.memory_usage_factor * (1024 * 1024)
+
+
+def _patch_attention(monkeypatch, xformers=False, pytorch_flash=False, flash=False):
+    monkeypatch.setattr(comfy.model_management, "xformers_enabled", lambda: xformers)
+    monkeypatch.setattr(comfy.model_management, "pytorch_attention_flash_attention", lambda: pytorch_flash)
+    monkeypatch.setattr(comfy.model_management, "flash_attention_enabled", lambda: flash)
+
+
+def _estimate():
+    return comfy.model_base.BaseModel.memory_required(_StubModel(), INPUT_SHAPE)
+
+
+def test_no_efficient_attention_uses_conservative_estimate(monkeypatch):
+    _patch_attention(monkeypatch)
+    assert _estimate() == CONSERVATIVE
+
+
+def test_pytorch_flash_attention_uses_efficient_estimate(monkeypatch):
+    _patch_attention(monkeypatch, pytorch_flash=True)
+    assert _estimate() == EFFICIENT
+
+
+def test_flash_attention_flag_uses_efficient_estimate(monkeypatch):
+    # --use-flash-attention must select the efficient estimate even when
+    # pytorch attention was not auto enabled (e.g. torch builds without
+    # working aotriton), otherwise the estimate is 7.5x too large.
+    _patch_attention(monkeypatch, flash=True)
+    assert _estimate() == EFFICIENT
+
+
+def test_conservative_estimate_is_7_5x_efficient_at_bf16(monkeypatch):
+    # Documents the size of the gap between the two formulas: at bf16 the
+    # conservative path asks for 7.5x more working memory than the efficient
+    # one for the same shapes. If either formula is retuned, this ratio (and
+    # the impact of picking the wrong branch) changes; update it consciously.
+    _patch_attention(monkeypatch, flash=True)
+    efficient = _estimate()
+    _patch_attention(monkeypatch)
+    conservative = _estimate()
+    assert conservative == 7.5 * efficient