Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 9 additions & 11 deletions examples/cli/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@ struct SDParams {
preview_t preview_method = PREVIEW_NONE;
int preview_interval = 1;
std::string preview_path = "preview.png";
float preview_fps = 16;
bool taesd_preview = false;
bool preview_noisy = false;

Expand Down Expand Up @@ -1630,25 +1631,22 @@ bool load_images_from_dir(const std::string dir,
return true;
}

std::string preview_path;
float preview_fps;

void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy) {
void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
(void)step;
(void)is_noisy;
SDParams* params = (SDParams*)data;
// is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
// unused in this app, it will either be always noisy or always denoised here
if (frame_count == 1) {
stbi_write_png(preview_path.c_str(), image->width, image->height, image->channel, image->data, 0);
stbi_write_png(params->preview_path.c_str(), image->width, image->height, image->channel, image->data, 0);
} else {
create_mjpg_avi_from_sd_images(preview_path.c_str(), image, frame_count, preview_fps);
create_mjpg_avi_from_sd_images(params->preview_path.c_str(), image, frame_count, params->preview_fps);
}
}

int main(int argc, const char* argv[]) {
SDParams params;
parse_args(argc, argv, params);
preview_path = params.preview_path;
if (params.video_frames > 4) {
size_t last_dot_pos = params.preview_path.find_last_of(".");
std::string base_path = params.preview_path;
Expand All @@ -1659,20 +1657,20 @@ int main(int argc, const char* argv[]) {
std::transform(file_ext.begin(), file_ext.end(), file_ext.begin(), ::tolower);
}
if (file_ext == ".png") {
preview_path = base_path + ".avi";
params.preview_path = base_path + ".avi";
}
}
preview_fps = params.fps;
params.preview_fps = params.fps;
if (params.preview_method == PREVIEW_PROJ)
preview_fps /= 4.0f;
params.preview_fps /= 4.0f;

params.sample_params.guidance.slg.layers = params.skip_layers.data();
params.sample_params.guidance.slg.layer_count = params.skip_layers.size();
params.high_noise_sample_params.guidance.slg.layers = params.high_noise_skip_layers.data();
params.high_noise_sample_params.guidance.slg.layer_count = params.high_noise_skip_layers.size();

sd_set_log_callback(sd_log_cb, (void*)&params);
sd_set_preview_callback((sd_preview_cb_t)step_callback, params.preview_method, params.preview_interval, !params.preview_noisy, params.preview_noisy);
sd_set_preview_callback(step_callback, params.preview_method, params.preview_interval, !params.preview_noisy, params.preview_noisy, (void*)&params);

if (params.verbose) {
print_params(params);
Expand Down
16 changes: 9 additions & 7 deletions stable-diffusion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1283,7 +1283,8 @@ class StableDiffusionGGML {
enum SDVersion version,
preview_t preview_mode,
ggml_tensor* result,
std::function<void(int, int, sd_image_t*, bool)> step_callback,
std::function<void(int, int, sd_image_t*, bool, void*)> step_callback,
void* step_callback_data,
bool is_noisy) {
const uint32_t channel = 3;
uint32_t width = latents->ne[0];
Expand Down Expand Up @@ -1354,7 +1355,7 @@ class StableDiffusionGGML {
for (int i = 0; i < frames; i++) {
images[i] = {width, height, channel, data + i * width * height * channel};
}
step_callback(step, frames, images, is_noisy);
step_callback(step, frames, images, is_noisy, step_callback_data);
free(data);
free(images);
} else {
Expand Down Expand Up @@ -1408,7 +1409,7 @@ class StableDiffusionGGML {
images[i].data = ggml_tensor_to_sd_image(result, i, ggml_n_dims(latents) == 4);
}

step_callback(step, frames, images, is_noisy);
step_callback(step, frames, images, is_noisy, step_callback_data);

ggml_ext_tensor_scale_inplace(result, 0);
for (int i = 0; i < frames; i++) {
Expand Down Expand Up @@ -1557,8 +1558,9 @@ class StableDiffusionGGML {
}

auto denoise = [&](ggml_tensor* input, float sigma, int step) -> ggml_tensor* {
auto sd_preview_cb = sd_get_preview_callback();
auto sd_preview_mode = sd_get_preview_mode();
auto sd_preview_cb = sd_get_preview_callback();
auto sd_preview_cb_data = sd_get_preview_callback_data();
auto sd_preview_mode = sd_get_preview_mode();
if (step == 1 || step == -1) {
pretty_progress(0, (int)steps, 0);
}
Expand Down Expand Up @@ -1627,7 +1629,7 @@ class StableDiffusionGGML {
}
if (sd_preview_cb != nullptr && sd_should_preview_noisy()) {
if (step % sd_get_preview_interval() == 0) {
preview_image(work_ctx, step, noised_input, version, sd_preview_mode, preview_tensor, sd_preview_cb, true);
preview_image(work_ctx, step, noised_input, version, sd_preview_mode, preview_tensor, sd_preview_cb, sd_preview_cb_data, true);
}
}

Expand Down Expand Up @@ -1775,7 +1777,7 @@ class StableDiffusionGGML {

if (sd_preview_cb != nullptr && sd_should_preview_denoised()) {
if (step % sd_get_preview_interval() == 0) {
preview_image(work_ctx, step, denoised, version, sd_preview_mode, preview_tensor, sd_preview_cb, false);
preview_image(work_ctx, step, denoised, version, sd_preview_mode, preview_tensor, sd_preview_cb, sd_preview_cb_data, false);
}
}

Expand Down
4 changes: 2 additions & 2 deletions stable-diffusion.h
Original file line number Diff line number Diff line change
Expand Up @@ -282,11 +282,11 @@ typedef struct sd_ctx_t sd_ctx_t;

typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy);
typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);

SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy);
SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
SD_API int32_t get_num_physical_cores();
SD_API const char* sd_get_system_info();

Expand Down
7 changes: 6 additions & 1 deletion util.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,7 @@ static sd_progress_cb_t sd_progress_cb = nullptr;
void* sd_progress_cb_data = nullptr;

static sd_preview_cb_t sd_preview_cb = nullptr;
static void* sd_preview_cb_data = nullptr;
preview_t sd_preview_mode = PREVIEW_NONE;
int sd_preview_interval = 1;
bool sd_preview_denoised = true;
Expand Down Expand Up @@ -335,8 +336,9 @@ void sd_set_progress_callback(sd_progress_cb_t cb, void* data) {
sd_progress_cb = cb;
sd_progress_cb_data = data;
}
void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode = PREVIEW_PROJ, int interval = 1, bool denoised = true, bool noisy = false) {
void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, bool denoised, bool noisy, void* data) {
sd_preview_cb = cb;
sd_preview_cb_data = data;
sd_preview_mode = mode;
sd_preview_interval = interval;
sd_preview_denoised = denoised;
Expand All @@ -346,6 +348,9 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode = PREVIEW_PROJ,
sd_preview_cb_t sd_get_preview_callback() {
return sd_preview_cb;
}
void* sd_get_preview_callback_data() {
return sd_preview_cb_data;
}

preview_t sd_get_preview_mode() {
return sd_preview_mode;
Expand Down
1 change: 1 addition & 0 deletions util.h
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ sd_progress_cb_t sd_get_progress_callback();
void* sd_get_progress_callback_data();

sd_preview_cb_t sd_get_preview_callback();
void* sd_get_preview_callback_data();
preview_t sd_get_preview_mode();
int sd_get_preview_interval();
bool sd_should_preview_denoised();
Expand Down
Loading