Model load progress bar fixed for multi-checkpoint models

This commit is contained in:
ebolam
2023-10-13 14:17:03 -04:00
parent 8cb956fd6e
commit 2319d1d87b

View File

@@ -707,7 +707,7 @@ class model_settings(settings):
self.loaded_layers = 0 # Used in UI 2 to show model loading progress
self.total_layers = 0 # Same as above
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.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
@@ -832,13 +832,22 @@ class model_settings(settings):
#Setup TQDP for model loading
elif name == "loaded_layers" and '_tqdm' in self.__dict__:
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
else:
if self.total_checkpoints == 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:
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
elif name == "total_download_chunks" and '_tqdm' in self.__dict__:
self._tqdm.reset(total=value)