Skip to content

Commit 8d4a2fb

Browse files
authored
Merge pull request #4 from agonzc34/better-loadingbar
Better Loading Bar
2 parents cfd77a3 + 5d7e9fa commit 8d4a2fb

File tree

1 file changed

+41
-16
lines changed

1 file changed

+41
-16
lines changed

src/huggingface_hub.cpp

Lines changed: 41 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
#include <iomanip>
3030
#include <regex>
3131
#include <sstream>
32+
#include <sys/ioctl.h>
33+
#include <unistd.h>
3234

3335
#include <curl/curl.h>
3436
#include <sys/stat.h>
@@ -47,7 +49,7 @@ void log_debug(const std::string &message) {
4749
if (!log_verbose) {
4850
return;
4951
}
50-
fprintf(stderr, "[DEBUG] %s\n", message.c_str());
52+
fprintf(stderr, "\036[31m[DEBUG] %s\033[0m\n", message.c_str());
5153
fflush(stderr);
5254
}
5355

@@ -63,7 +65,7 @@ void log_info_with_carriage_return(const std::string &message) {
6365
}
6466

6567
void log_error(const std::string &message) {
66-
fprintf(stderr, "[ERROR] %s\n", message.c_str());
68+
fprintf(stderr, "\033[31m[ERROR] %s\033[0m\n", message.c_str());
6769
fflush(stderr);
6870
}
6971

@@ -212,6 +214,18 @@ get_model_metadata_from_hf(const std::string &repo, const std::string &file) {
212214
return metadata;
213215
}
214216

217+
int get_terminal_width() {
218+
struct winsize w;
219+
if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &w) == 0) {
220+
return w.ws_col;
221+
} else {
222+
return 80; // Default terminal width
223+
}
224+
}
225+
226+
std::chrono::steady_clock::time_point last_print_time =
227+
std::chrono::steady_clock::now();
228+
215229
// Progress bar function
216230
int progress_callback(void *userdata, curl_off_t total, curl_off_t now,
217231
curl_off_t, curl_off_t) {
@@ -220,9 +234,15 @@ int progress_callback(void *userdata, curl_off_t total, curl_off_t now,
220234
uint64_t size = metadata->size;
221235
uint64_t byte_offset = total - size;
222236
uint64_t downloaded = now - byte_offset;
237+
int terminal_width = get_terminal_width();
238+
auto elapsed = std::chrono::steady_clock::now() - last_print_time;
239+
240+
if (total > 0 && (now == downloaded ||
241+
std::chrono::duration<double>(elapsed).count() > 0.08)) {
242+
last_print_time = std::chrono::steady_clock::now();
223243

224-
if (total > 0) {
225-
int width = 30; // Progress bar width
244+
bool show_speed = terminal_width > 65;
245+
int width = terminal_width - 65 + (show_speed ? 0 : 10);
226246
float percent = static_cast<float>(downloaded) / size;
227247
int filled = static_cast<int>(percent * width);
228248

@@ -231,22 +251,27 @@ int progress_callback(void *userdata, curl_off_t total, curl_off_t now,
231251
1e-6); // Avoid division by zero
232252
double remaining = (total - now) / speed;
233253

234-
// Print progress bar
235254
std::ostringstream progress;
236-
progress << "[";
237-
for (int i = 0; i < filled; ++i)
238-
progress << "#";
239-
for (int i = filled; i < width; ++i)
240-
progress << " ";
241-
progress << "] " << std::fixed << std::setprecision(2) << (percent * 100)
242-
<< "%";
255+
256+
// Print progress bar
257+
if (terminal_width > 50) {
258+
progress << "[";
259+
for (int i = 0; i < filled; ++i)
260+
progress << "#";
261+
for (int i = filled; i < width; ++i)
262+
progress << " ";
263+
progress << "] ";
264+
}
265+
progress << std::fixed << std::setprecision(2) << (percent * 100) << "%";
243266

244267
progress << " " << downloaded / 1024 / 1024 << " MB / "
245-
<< size / 1024 / 1024 << " MB";
268+
<< size / 1024 / 1024 << " MB ";
246269

247-
// Print speed and ETA
248-
progress << " " << (speed / 1024 / 1024)
249-
<< " MB/s | ETA: " << static_cast<int>(remaining) << "s ";
270+
if (show_speed) {
271+
progress << " " << (speed / 1024 / 1024) << " MB/s ";
272+
}
273+
progress << " | ETA: " << std::fixed << std::setprecision(1) << remaining
274+
<< "s";
250275
log_info_with_carriage_return(progress.str());
251276
}
252277
if (stop_download) {

0 commit comments

Comments
 (0)