40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
"""Regression test for init_kv_cache crashing when the text encoder runs on CPU."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from comfy.cli_args import args
|
|
|
|
if not torch.cuda.is_available():
|
|
args.cpu = True
|
|
|
|
import comfy.ops as ops # noqa: E402
|
|
import comfy.text_encoders.llama as llama # noqa: E402
|
|
|
|
|
|
def _make_model():
|
|
config = llama.Llama2Config(
|
|
vocab_size=16, hidden_size=8, intermediate_size=8,
|
|
num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2,
|
|
)
|
|
config.head_dim = 4
|
|
config.fixed_kv = True
|
|
return llama.Llama2_(config, device="cpu", dtype=torch.float32, ops=ops.manual_cast)
|
|
|
|
|
|
def test_init_kv_cache_on_cpu_does_not_probe_cuda_capability(monkeypatch):
|
|
def _raise_if_not_cuda(device):
|
|
if torch.device(device).type != "cuda":
|
|
raise ValueError(f"Expected a cuda device, but got: {device}")
|
|
return True
|
|
|
|
monkeypatch.setattr(
|
|
llama.comfy_kitchen, "flash_attention_decode_is_available", _raise_if_not_cuda
|
|
)
|
|
|
|
model = _make_model()
|
|
past = model.init_kv_cache(1, 4, torch.device("cpu"), torch.float32)
|
|
|
|
assert not isinstance(past[0], llama.FixedKV)
|