From 1bfa440ef19b7b949c05ab4a78090149a2c46f82 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Fri, 24 Nov 2023 18:49:50 +0100
Subject: [PATCH] Add progress bar to model loader

---
 exllama.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/exllama.py b/exllama.py
index 5dcab75..67b7720 100644
--- a/exllama.py
+++ b/exllama.py
@@ -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)