From 08aec643126a88971cafcbab97dd0101185c2c3d Mon Sep 17 00:00:00 2001 From: Daniele <57776841+daniandtheweb@users.noreply.github.com> Date: Sun, 20 Jul 2025 21:54:34 +0200 Subject: [PATCH] Added prediction argument --- README.md | 1 + examples/cli/main.cpp | 19 ++++++ stable-diffusion.cpp | 134 +++++++++++++++++++++++++++++++----------- stable-diffusion.h | 13 ++++ 4 files changed, 132 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 8ce98137f..f74503e7d 100644 --- a/README.md +++ b/README.md @@ -333,6 +333,7 @@ arguments: -s SEED, --seed SEED RNG seed (default: 42, use random seed for < 0) -b, --batch-count COUNT number of images to generate --schedule {discrete, karras, exponential, ays, gits} Denoiser sigma schedule (default: discrete) + --prediction {eps, v, edm_v, sd3_flow, flux_flow} Prediction type override --clip-skip N ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1) <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x --vae-tiling process vae in tiles to reduce memory usage diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index b3ae569e6..0d6646fa1 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -85,6 +85,7 @@ struct SDParams { sample_method_t sample_method = EULER_A; schedule_t schedule = DEFAULT; + prediction_t prediction = DEFAULT_PRED; int sample_steps = 20; float strength = 0.75f; float control_strength = 0.9f; @@ -156,6 +157,7 @@ void print_params(SDParams params) { printf(" height: %d\n", params.height); printf(" sample_method: %s\n", sd_sample_method_name(params.sample_method)); printf(" schedule: %s\n", sd_schedule_name(params.schedule)); + printf(" prediction: %s\n", sd_prediction_name(params.prediction)); printf(" sample_steps: %d\n", params.sample_steps); printf(" strength(img2img): %.2f\n", params.strength); printf(" rng: %s\n", sd_rng_type_name(params.rng_type)); @@ -224,6 +226,7 @@ void print_usage(int argc, const char* argv[]) { printf(" -s SEED, --seed SEED RNG seed (default: 42, use random seed for < 0)\n"); printf(" -b, --batch-count COUNT number of images to generate\n"); printf(" --schedule {discrete, karras, exponential, ays, gits} Denoiser sigma schedule (default: discrete)\n"); + printf(" --prediction {eps, v, edm_v, sd3_flow, flux_flow} Prediction type override.\n"); printf(" --clip-skip N ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n"); printf(" <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n"); printf(" --vae-tiling process vae in tiles to reduce memory usage\n"); @@ -494,6 +497,20 @@ void parse_args(int argc, const char** argv, SDParams& params) { return 1; }; + auto on_prediction_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + const char* arg = argv[index]; + params.prediction = str_to_prediction(arg); + if (params.prediction == PREDICTION_COUNT) { + fprintf(stderr, "error: invalid prediction type %s\n", + arg); + return -1; + } + return 1; + }; + auto on_sample_method_arg = [&](int argc, const char** argv, int index) { if (++index >= argc) { return -1; @@ -564,6 +581,7 @@ void parse_args(int argc, const char** argv, SDParams& params) { {"-s", "--seed", "", on_seed_arg}, {"", "--sampling-method", "", on_sample_method_arg}, {"", "--schedule", "", on_schedule_arg}, + {"", "--prediction", "", on_prediction_arg}, {"", "--skip-layers", "", on_skip_layers_arg}, {"-r", "--ref-image", "", on_ref_image_arg}, {"-h", "--help", "", on_help_arg}, @@ -883,6 +901,7 @@ int main(int argc, const char* argv[]) { params.wtype, params.rng_type, params.schedule, + params.prediction, params.clip_on_cpu, params.control_net_cpu, params.vae_on_cpu, diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 402585f1c..f177a7c9e 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -561,47 +561,83 @@ class StableDiffusionGGML { int64_t t1 = ggml_time_ms(); LOG_INFO("loading model from '%s' completed, taking %.2fs", SAFE_STR(sd_ctx_params->model_path), (t1 - t0) * 1.0f / 1000); - // check is_using_v_parameterization_for_sd2 - - if (sd_version_is_sd2(version)) { - if (is_using_v_parameterization_for_sd2(ctx, sd_version_is_inpaint(version))) { - is_using_v_parameterization = true; - } - } else if (sd_version_is_sdxl(version)) { - if (model_loader.tensor_storages_types.find("edm_vpred.sigma_max") != model_loader.tensor_storages_types.end()) { - // CosXL models - // TODO: get sigma_min and sigma_max values from file - is_using_edm_v_parameterization = true; + if (sd_ctx_params->prediction != DEFAULT_PRED) { + switch (sd_ctx_params->prediction) { + case EPS_PRED: + LOG_INFO("running in eps-prediction mode"); + break; + case V_PRED: + LOG_INFO("running in v-prediction mode"); + denoiser = std::make_shared(); + break; + case EDM_V_PRED: + LOG_INFO("running in v-prediction EDM mode"); + denoiser = std::make_shared(); + break; + case SD3_FLOW_PRED: + LOG_INFO("running in FLOW mode"); + denoiser = std::make_shared(); + break; + case FLUX_FLOW_PRED: + { + LOG_INFO("running in Flux FLOW mode"); + float shift = 1.0f; // TODO: validate + for (auto pair : model_loader.tensor_storages_types) { + if (pair.first.find("model.diffusion_model.guidance_in.in_layer.weight") != std::string::npos) { + shift = 1.15f; + break; + } + } + denoiser = std::make_shared(shift); + break; + } + default: + LOG_ERROR("Unknown parametrization %i", sd_ctx_params->prediction); + abort(); } - if (model_loader.tensor_storages_types.find("v_pred") != model_loader.tensor_storages_types.end()) { + } else { + // check is_using_v_parameterization_for_sd2 + + if (sd_version_is_sd2(version)) { + if (is_using_v_parameterization_for_sd2(ctx, sd_version_is_inpaint(version))) { + is_using_v_parameterization = true; + } + } else if (sd_version_is_sdxl(version)) { + if (model_loader.tensor_storages_types.find("edm_vpred.sigma_max") != model_loader.tensor_storages_types.end()) { + // CosXL models + // TODO: get sigma_min and sigma_max values from file + is_using_edm_v_parameterization = true; + } + if (model_loader.tensor_storages_types.find("v_pred") != model_loader.tensor_storages_types.end()) { + is_using_v_parameterization = true; + } + } else if (version == VERSION_SVD) { + // TODO: V_PREDICTION_EDM is_using_v_parameterization = true; } - } else if (version == VERSION_SVD) { - // TODO: V_PREDICTION_EDM - is_using_v_parameterization = true; - } - if (sd_version_is_sd3(version)) { - LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(); - } else if (sd_version_is_flux(version)) { - LOG_INFO("running in Flux FLOW mode"); - float shift = 1.0f; // TODO: validate - for (auto pair : model_loader.tensor_storages_types) { - if (pair.first.find("model.diffusion_model.guidance_in.in_layer.weight") != std::string::npos) { - shift = 1.15f; - break; + if (sd_version_is_sd3(version)) { + LOG_INFO("running in FLOW mode"); + denoiser = std::make_shared(); + } else if (sd_version_is_flux(version)) { + LOG_INFO("running in Flux FLOW mode"); + float shift = 1.0f; // TODO: validate + for (auto pair : model_loader.tensor_storages_types) { + if (pair.first.find("model.diffusion_model.guidance_in.in_layer.weight") != std::string::npos) { + shift = 1.15f; + break; + } } + denoiser = std::make_shared(shift); + } else if (is_using_v_parameterization) { + LOG_INFO("running in v-prediction mode"); + denoiser = std::make_shared(); + } else if (is_using_edm_v_parameterization) { + LOG_INFO("running in v-prediction EDM mode"); + denoiser = std::make_shared(); + } else { + LOG_INFO("running in eps-prediction mode"); } - denoiser = std::make_shared(shift); - } else if (is_using_v_parameterization) { - LOG_INFO("running in v-prediction mode"); - denoiser = std::make_shared(); - } else if (is_using_edm_v_parameterization) { - LOG_INFO("running in v-prediction EDM mode"); - denoiser = std::make_shared(); - } else { - LOG_INFO("running in eps-prediction mode"); } if (sd_ctx_params->schedule != DEFAULT) { @@ -1290,6 +1326,31 @@ enum schedule_t str_to_schedule(const char* str) { return SCHEDULE_COUNT; } +const char* prediction_to_str[] = { + "default", + "eps", + "v", + "edm_v", + "sd3_flow", + "flux_flow", +}; + +const char* sd_prediction_name(enum prediction_t prediction) { + if (prediction < PREDICTION_COUNT) { + return prediction_to_str[prediction]; + } + return NONE_STR; +} + +enum prediction_t str_to_prediction(const char* str) { + for (int i = 0; i < PREDICTION_COUNT; i++) { + if (!strcmp(str, prediction_to_str[i])) { + return (enum prediction_t)i; + } + } + return PREDICTION_COUNT; +} + void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { memset((void*)sd_ctx_params, 0, sizeof(sd_ctx_params_t)); sd_ctx_params->vae_decode_only = true; @@ -1299,6 +1360,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->wtype = SD_TYPE_COUNT; sd_ctx_params->rng_type = CUDA_RNG; sd_ctx_params->schedule = DEFAULT; + sd_ctx_params->prediction = DEFAULT_PRED; sd_ctx_params->keep_clip_on_cpu = false; sd_ctx_params->keep_control_net_on_cpu = false; sd_ctx_params->keep_vae_on_cpu = false; @@ -1333,6 +1395,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "wtype: %s\n" "rng_type: %s\n" "schedule: %s\n" + "prediction: %s\n" "keep_clip_on_cpu: %s\n" "keep_control_net_on_cpu: %s\n" "keep_vae_on_cpu: %s\n" @@ -1358,6 +1421,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { sd_type_name(sd_ctx_params->wtype), sd_rng_type_name(sd_ctx_params->rng_type), sd_schedule_name(sd_ctx_params->schedule), + sd_prediction_name(sd_ctx_params->prediction), BOOL_STR(sd_ctx_params->keep_clip_on_cpu), BOOL_STR(sd_ctx_params->keep_control_net_on_cpu), BOOL_STR(sd_ctx_params->keep_vae_on_cpu), diff --git a/stable-diffusion.h b/stable-diffusion.h index a60325923..6f4be31f2 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -60,6 +60,16 @@ enum schedule_t { SCHEDULE_COUNT }; +enum prediction_t { + DEFAULT_PRED, + EPS_PRED, + V_PRED, + EDM_V_PRED, + SD3_FLOW_PRED, + FLUX_FLOW_PRED, + PREDICTION_COUNT +}; + // same as enum ggml_type enum sd_type_t { SD_TYPE_F32 = 0, @@ -130,6 +140,7 @@ typedef struct { enum sd_type_t wtype; enum rng_type_t rng_type; enum schedule_t schedule; + enum prediction_t prediction; bool keep_clip_on_cpu; bool keep_control_net_on_cpu; bool keep_vae_on_cpu; @@ -219,6 +230,8 @@ SD_API const char* sd_sample_method_name(enum sample_method_t sample_method); SD_API enum sample_method_t str_to_sample_method(const char* str); SD_API const char* sd_schedule_name(enum schedule_t schedule); SD_API enum schedule_t str_to_schedule(const char* str); +SD_API const char* sd_prediction_name(enum prediction_t prediction); +SD_API enum prediction_t str_to_prediction(const char* str); SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params); SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);