feat: added prediction argument (#334)

This commit is contained in:
Daniele
2025-10-15 17:00:10 +02:00
committed by GitHub
parent a7d6d296c7
commit e3702585cb
4 changed files with 155 additions and 54 deletions

View File

@@ -84,6 +84,7 @@ struct SDParams {
std::string prompt;
std::string negative_prompt;
int clip_skip = -1; // <= 0 represents unspecified
int width = 512;
int height = 512;
@@ -127,6 +128,8 @@ struct SDParams {
int chroma_t5_mask_pad = 1;
float flow_shift = INFINITY;
prediction_t prediction = DEFAULT_PRED;
sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
SDParams() {
@@ -188,6 +191,7 @@ void print_params(SDParams params) {
printf(" sample_params: %s\n", SAFE_STR(sample_params_str));
printf(" high_noise_sample_params: %s\n", SAFE_STR(high_noise_sample_params_str));
printf(" moe_boundary: %.3f\n", params.moe_boundary);
printf(" prediction: %s\n", sd_prediction_name(params.prediction));
printf(" flow_shift: %.2f\n", params.flow_shift);
printf(" strength(img2img): %.2f\n", params.strength);
printf(" rng: %s\n", sd_rng_type_name(params.rng_type));
@@ -281,6 +285,7 @@ void print_usage(int argc, const char* argv[]) {
printf(" --rng {std_default, cuda} RNG (default: cuda)\n");
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(" --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");
@@ -651,6 +656,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;
@@ -807,6 +826,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
{"", "--rng", "", on_rng_arg},
{"-s", "--seed", "", on_seed_arg},
{"", "--sampling-method", "", on_sample_method_arg},
{"", "--prediction", "", on_prediction_arg},
{"", "--scheduler", "", on_schedule_arg},
{"", "--skip-layers", "", on_skip_layers_arg},
{"", "--high-noise-sampling-method", "", on_high_noise_sample_method_arg},
@@ -1354,6 +1374,7 @@ int main(int argc, const char* argv[]) {
params.n_threads,
params.wtype,
params.rng_type,
params.prediction,
params.offload_params_to_cpu,
params.clip_on_cpu,
params.control_net_cpu,