Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add in capability to repeatedly run the upscaler in a row #174

Merged
merged 4 commits into from
Feb 24, 2024
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
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ cmake --build . --config Release
### Run

```
usage: ./bin/sd [arguments]
usage: ./build/bin/sd [arguments]

arguments:
-h, --help show this help message and exit
Expand All @@ -161,6 +161,7 @@ arguments:
--control-net [CONTROL_PATH] path to control net model
--embd-dir [EMBEDDING_PATH] path to embeddings.
--upscale-model [ESRGAN_PATH] path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.
--upscale-repeats Run the ESRGAN upscaler this many times (default 1)
--type [TYPE] weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
If not specified, the default is the type of the weight file.
--lora-model-dir [DIR] lora model directory
Expand All @@ -186,6 +187,7 @@ arguments:
<= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
--vae-tiling process vae in tiles to reduce memory usage
--control-net-cpu keep controlnet in cpu (for low vram)
--canny apply canny preprocessor (edge detection)
-v, --verbose print extra info
```

Expand Down
31 changes: 24 additions & 7 deletions examples/cli/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ struct SDParams {
bool vae_tiling = false;
bool control_net_cpu = false;
bool canny_preprocess = false;
int upscale_repeats = 1;
};

void print_params(SDParams params) {
Expand Down Expand Up @@ -129,6 +130,7 @@ void print_params(SDParams params) {
printf(" seed: %ld\n", params.seed);
printf(" batch_count: %d\n", params.batch_count);
printf(" vae_tiling: %s\n", params.vae_tiling ? "true" : "false");
printf(" upscale_repeats: %d\n", params.upscale_repeats);
}

void print_usage(int argc, const char* argv[]) {
Expand All @@ -145,6 +147,7 @@ void print_usage(int argc, const char* argv[]) {
printf(" --control-net [CONTROL_PATH] path to control net model\n");
printf(" --embd-dir [EMBEDDING_PATH] path to embeddings.\n");
printf(" --upscale-model [ESRGAN_PATH] path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n");
printf(" --upscale-repeats Run the ESRGAN upscaler this many times (default 1)\n");
printf(" --type [TYPE] weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n");
printf(" If not specified, the default is the type of the weight file.\n");
printf(" --lora-model-dir [DIR] lora model directory\n");
Expand Down Expand Up @@ -296,6 +299,16 @@ void parse_args(int argc, const char** argv, SDParams& params) {
break;
}
params.prompt = argv[i];
} else if (arg == "--upscale-repeats") {
if (++i >= argc) {
invalid_arg = true;
break;
}
params.upscale_repeats = std::stoi(argv[i]);
if (params.upscale_repeats < 1) {
fprintf(stderr, "error: upscale multiplier must be at least 1\n");
exit(1);
}
} else if (arg == "-n" || arg == "--negative-prompt") {
if (++i >= argc) {
invalid_arg = true;
Expand Down Expand Up @@ -700,7 +713,7 @@ int main(int argc, const char* argv[]) {
}

int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth
if (params.esrgan_path.size() > 0) {
if (params.esrgan_path.size() > 0 && params.upscale_repeats > 0) {
upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(params.esrgan_path.c_str(),
params.n_threads,
params.wtype);
Expand All @@ -712,13 +725,17 @@ int main(int argc, const char* argv[]) {
if (results[i].data == NULL) {
continue;
}
sd_image_t upscaled_image = upscale(upscaler_ctx, results[i], upscale_factor);
if (upscaled_image.data == NULL) {
printf("upscale failed\n");
continue;
sd_image_t current_image = results[i];
for (int u = 0; u < params.upscale_repeats; ++u) {
sd_image_t upscaled_image = upscale(upscaler_ctx, current_image, upscale_factor);
if (upscaled_image.data == NULL) {
printf("upscale failed\n");
break;
}
free(current_image.data);
current_image = upscaled_image;
}
free(results[i].data);
results[i] = upscaled_image;
results[i] = current_image; // Set the final upscaled image as the result
}
}
}
Expand Down
9 changes: 5 additions & 4 deletions stable-diffusion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -173,10 +173,11 @@ class StableDiffusionGGML {
if (version == VERSION_XL) {
scale_factor = 0.13025f;
if (vae_path.size() == 0 && taesd_path.size() == 0) {
LOG_WARN("!!!It looks like you are using SDXL model. "
"If you find that the generated images are completely black, "
"try specifying SDXL VAE FP16 Fix with the --vae parameter. "
"You can find it here: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors");
LOG_WARN(
"!!!It looks like you are using SDXL model. "
"If you find that the generated images are completely black, "
"try specifying SDXL VAE FP16 Fix with the --vae parameter. "
"You can find it here: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors");
}
}

Expand Down
Loading