Browse Source

Try to fix OOM issues with cards that have less vram than mine.

pull/3/head
comfyanonymous 2 years ago
parent
commit
50db297cf6
  1. 7
      comfy/ldm/modules/attention.py

7
comfy/ldm/modules/attention.py

@ -194,7 +194,14 @@ class CrossAttentionBirchSan(nn.Module):
kv_chunk_size_min = None kv_chunk_size_min = None
#not sure at all about the math here
#TODO: tweak this
if mem_free_total > 8192 * 1024 * 1024 * 1.3:
query_chunk_size_x = 1024 * 4 query_chunk_size_x = 1024 * 4
elif mem_free_total > 4096 * 1024 * 1024 * 1.3:
query_chunk_size_x = 1024 * 2
else:
query_chunk_size_x = 1024
kv_chunk_size_min_x = None kv_chunk_size_min_x = None
kv_chunk_size_x = (int((chunk_threshold_bytes // (batch_x_heads * bytes_per_token * query_chunk_size_x)) * 2.0) // 1024) * 1024 kv_chunk_size_x = (int((chunk_threshold_bytes // (batch_x_heads * bytes_per_token * query_chunk_size_x)) * 2.0) // 1024) * 1024
if kv_chunk_size_x < 1024: if kv_chunk_size_x < 1024:

Loading…
Cancel
Save