Fix bitsandbytes imports in unit tests on MacOS.

invoke-ai · Dec 30, 2024 · 9a0a226 · 9a0a226
1 parent 477d87e
commit 9a0a226
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/...er/load/model_cache/torch_module_autocast/custom_modules/test_custom_invoke_linear_nf4.py b/...er/load/model_cache/torch_module_autocast/custom_modules/test_custom_invoke_linear_nf4.py
@@ -1,13 +1,17 @@
 import pytest
 import torch
 
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_nf4 import (
-    CustomInvokeLinearNF4,
-)
 from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
     wrap_custom_layer,
 )
-from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA is not available", allow_module_level=True)
+else:
+    from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_nf4 import (
+        CustomInvokeLinearNF4,
+    )
+    from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
 
 
 def build_linear_nf4_layer(orig_layer: torch.nn.Linear | None = None):