Add llama.cpp GPU offload option (#2060)

2023-05-14 21:58:11 -04:00 · 2023-05-14 21:58:11 -04:00 · 071f0776ad
commit 071f0776ad
parent eee986348c
4 changed files with 23 additions and 5 deletions
--- a/modules/llamacpp_model.py
+++ b/modules/llamacpp_model.py
@ -27,7 +27,8 @@ class LlamaCppModel:
            'n_threads': shared.args.threads or None,
            'n_batch': shared.args.n_batch,
            'use_mmap': not shared.args.no_mmap,
-            'use_mlock': shared.args.mlock
+            'use_mlock': shared.args.mlock,
+            'n_gpu_layers': shared.args.n_gpu_layers
        }
        self.model = Llama(**params)
        self.model.set_cache(LlamaCache)