downloader: better progress bar

Touch-Night · Jul 29, 2024 · f4d95f3 · f4d95f3
1 parent 9dcff21
commit f4d95f3
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/download-model.py b/download-model.py
@@ -212,11 +212,15 @@ def get_single_file(self, url, output_folder, start_from_scratch=False):
                     total_size = int(r.headers.get('content-length', 0))
                     block_size = 1024 * 1024  # 1MB
 
+                    filename_str = str(filename)  # Convert PosixPath to string if necessary
+
                     tqdm_kwargs = {
                         'total': total_size,
-                        'unit': 'iB',
+                        'unit': 'B',
                         'unit_scale': True,
-                        'bar_format': '{l_bar}{bar}| {n_fmt}/{total_fmt} {rate_fmt}'
+                        'unit_divisor': 1024,
+                        'bar_format': '{desc}{percentage:3.0f}%|{bar:50}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]',
+                        'desc': f"{filename_str}: "
                     }
 
                     if 'COLAB_GPU' in os.environ:
@@ -233,7 +237,7 @@ def get_single_file(self, url, output_folder, start_from_scratch=False):
                                 t.update(len(data))
                                 if total_size != 0 and self.progress_bar is not None:
                                     count += len(data)
-                                    self.progress_bar(float(count) / float(total_size), f"{filename}")
+                                    self.progress_bar(float(count) / float(total_size), f"{filename_str}")
 
                     break  # Exit loop if successful
             except (RequestException, ConnectionError, Timeout) as e: