Add progress bar to model loader

2023-11-24 18:49:50 +01:00
parent 09de8351b9
commit 1bfa440ef1
1 changed files with 6 additions and 1 deletions
@@ -67,7 +67,12 @@ class Loader:
            return

        self.ckpt = ExLlamaV2(self.config)
-        self.ckpt.load(gpu_split=self.gpu_split)
+        progress = ProgressBar(len(self.ckpt.modules))
+
+        self.ckpt.load(
+            gpu_split=self.gpu_split,
+            callback=lambda s, _: progress.update_absolute(s),
+        )

        self.cache = (
            ExLlamaV2Cache_8bit(self.ckpt)