work

leejet · Sep 8, 2024 · f4e2f56 · f4e2f56
1 parent 14206fd
commit f4e2f56
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 13 deletions.
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -664,10 +664,12 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
             break;
     }
 
-    if (params->color == true) {
-        fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str);
-    } else {
-        fprintf(out_stream, "[%-5s] ", level_str);
+    if (level != SD_LOG_PROGRESS){
+        if (params->color) {
+            fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str);
+        } else {
+            fprintf(out_stream, "[%-5s] ", level_str);
+        }
     }
     fputs(log, out_stream);
     fflush(out_stream);

diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -102,7 +102,8 @@ enum sd_log_level_t {
     SD_LOG_DEBUG,
     SD_LOG_INFO,
     SD_LOG_WARN,
-    SD_LOG_ERROR
+    SD_LOG_ERROR,
+    SD_LOG_PROGRESS
 };
 
 typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);

diff --git a/util.cpp b/util.cpp
@@ -28,6 +28,8 @@
 #define STB_IMAGE_RESIZE_IMPLEMENTATION
 #include "stb_image_resize.h"
 
+#define LOG_BUFFER_SIZE 1024
+
 bool ends_with(const std::string& str, const std::string& ending) {
     if (str.length() >= ending.length()) {
         return (str.compare(str.length() - ending.length(), ending.length(), ending) == 0);
@@ -330,12 +332,17 @@ void pretty_progress(int step, int steps, float time) {
         }
     }
     progress += "|";
-    printf(time > 1.0f ? "\r%s %i/%i - %.2fs/it" : "\r%s %i/%i - %.2fit/s",
-           progress.c_str(), step, steps,
-           time > 1.0f || time == 0 ? time : (1.0f / time));
+
+    char log_text[LOG_BUFFER_SIZE];
+    snprintf(log_text,LOG_BUFFER_SIZE,time > 1.0f ? "\r%s %i/%i - %.2fs/it" : "\r%s %i/%i - %.2fit/s",
+             progress.c_str(), step, steps,
+             time > 1.0f || time == 0 ? time : (1.0f / time));
+
+    LOG_PROGRESS(log_text);
+
     fflush(stdout);  // for linux
     if (step == steps) {
-        printf("\n");
+        LOG_PROGRESS("\n");
     }
 }
 
@@ -360,19 +367,24 @@ std::string trim(const std::string& s) {
 static sd_log_cb_t sd_log_cb = NULL;
 void* sd_log_cb_data         = NULL;
 
-#define LOG_BUFFER_SIZE 1024
-
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
     va_list args;
     va_start(args, format);
 
     static char log_buffer[LOG_BUFFER_SIZE + 1];
-    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line);
+
+    int written = 0;
+    if (level != SD_LOG_PROGRESS) {
+        written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line);
+    }
 
     if (written >= 0 && written < LOG_BUFFER_SIZE) {
         vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
     }
-    strncat(log_buffer, "\n", LOG_BUFFER_SIZE - strlen(log_buffer));
+
+    if (level != SD_LOG_PROGRESS){
+        strncat(log_buffer, "\n", LOG_BUFFER_SIZE - strlen(log_buffer));
+    }
 
     if (sd_log_cb) {
         sd_log_cb(level, log_buffer, sd_log_cb_data);

diff --git a/util.h b/util.h
@@ -58,4 +58,5 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
 #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#define LOG_PROGRESS(format, ...) log_printf(SD_LOG_PROGRESS, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #endif  // __UTIL_H__