fix error with loading T5 in low-VRAM mode

seruva19 · Dec 11, 2024 · 6806394 · 6806394
1 parent 8401e4f
commit 6806394
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/src/models/model_30/kandinsky3/condition_encoders.py b/src/models/model_30/kandinsky3/condition_encoders.py
@@ -81,7 +81,8 @@ def __init__(
         if environment.kd30_low_vram:
             quantization_config = BitsAndBytesConfig(
                 load_in_4bit=True,
-                bnb_4bit_compute_type=torch.bfloat16,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                llm_int8_enable_fp32_cpu_offload=True,
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_use_double_quant=False,
             )

diff --git a/src/models/model_31/kandinsky31/condition_encoders.py b/src/models/model_31/kandinsky31/condition_encoders.py
@@ -30,7 +30,8 @@ def __init__(
         if environment.kd31_low_vram:
             quantization_config = BitsAndBytesConfig(
                 load_in_4bit=True,
-                bnb_4bit_compute_type=torch.bfloat16,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                llm_int8_enable_fp32_cpu_offload=True,
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_use_double_quant=False,
             )

diff --git a/src/utils/web.py b/src/utils/web.py