rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Model load progress bar fixed for multi-checkpoint models

Browse Source
This commit is contained in:
ebolam
2023-10-13 14:17:03 -04:00
parent 8cb956fd6e
commit 2319d1d87b
1 changed files with 14 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
koboldai_settings.py
View File

@@ -707,7 +707,7 @@ class model_settings(settings):
self.loaded_layers = 0 # Used in UI 2 to show model loading progress self.loaded_layers = 0 # Used in UI 2 to show model loading progress
self.total_layers = 0 # Same as above self.total_layers = 0 # Same as above
self.loaded_checkpoints = 0 self.loaded_checkpoints = 0
self.total_checkpoints = 0 self.total_checkpoints = 1
self.total_download_chunks = 0 # tracks how much of the model has downloaded for the UI 2 self.total_download_chunks = 0 # tracks how much of the model has downloaded for the UI 2
self.downloaded_chunks = 0 #as above self.downloaded_chunks = 0 #as above
self._tqdm = tqdm.tqdm(total=self.genamt, file=self.ignore_tqdm()) # tqdm agent for generating tokens. This will allow us to calculate the remaining time self._tqdm = tqdm.tqdm(total=self.genamt, file=self.ignore_tqdm()) # tqdm agent for generating tokens. This will allow us to calculate the remaining time
@@ -832,13 +832,22 @@ class model_settings(settings):
#Setup TQDP for model loading #Setup TQDP for model loading
elif name == "loaded_layers" and '_tqdm' in self.__dict__: elif name == "loaded_layers" and '_tqdm' in self.__dict__:
if value == 0: if value == 0:
self._tqdm.reset(total=self.total_layers) self._tqdm.reset(total=self.total_layers if self.total_checkpoints == 1 else 1000)
self.tqdm_progress = 0 self.tqdm_progress = 0
else: else:
if self.total_checkpoints == 1:
self._tqdm.update(1) self._tqdm.update(1)
self.tqdm_progress = int(float(self.loaded_layers)/float(self.total_layers)*100) elif self.total_layers != 0 and self.total_checkpoints != 0:
proper_progress = (self.loaded_checkpoints + value/self.total_layers)/self.total_checkpoints*1000
self._tqdm.update(proper_progress - self._tqdm.n)
self.tqdm_progress = int(float(self._tqdm.n)/float(self._tqdm.total)*100)
if self._tqdm.format_dict['rate'] is not None: if self._tqdm.format_dict['rate'] is not None:
self.tqdm_rem_time = str(datetime.timedelta(seconds=int(float(self.total_layers-self.loaded_layers)/self._tqdm.format_dict['rate']))) elapsed = self._tqdm.format_dict["elapsed"]
rate = self._tqdm.format_dict["rate"]
remaining = (self._tqdm.total - self._tqdm.n) / rate if rate and self._tqdm.total else 0
self.tqdm_rem_time = str(datetime.timedelta(seconds=remaining))
#Setup TQDP for model downloading #Setup TQDP for model downloading
elif name == "total_download_chunks" and '_tqdm' in self.__dict__: elif name == "total_download_chunks" and '_tqdm' in self.__dict__:
self._tqdm.reset(total=value) self._tqdm.reset(total=value)