298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-25 15:55:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
comfyanonymous 2023-01-29 00:50:46 -05:00
parent 36ec5690a6
commit 50db297cf6
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/ldm/modules/attention.py
View File

@ -194,7 +194,14 @@ class CrossAttentionBirchSan(nn.Module):
kv_chunk_size_min = None
query_chunk_size_x = 1024 * 4
#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
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_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: