refactor: make image schedule configurable

Signed-off-by: thxCode <thxcode0824@gmail.com>
gpustack · Nov 28, 2024 · 25ce758 · 25ce758
1 parent c64c47e
commit 25ce758
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -675,6 +675,7 @@ The available endpoints for the LLaMA Box server mode are:
         "size": "512x512",
         "prompt": "A lovely cat",
         "sampler": "euler",      // required, select from euler_a;euler;heun;dpm2;dpm++2s_a;dpm++2m;dpm++2mv2;ipndm;ipndm_v;lcm
+        "schedule": "default",   // optional, select from default;discrete;karras;exponential;ays;gits
         "seed": null,            // optional, random seed
         "cfg_scale": 4.5,        // optional, for sampler, the scale of classifier-free guidance in the output phase
         "sample_steps": 20,      // optional, number of sample steps
@@ -736,6 +737,7 @@ The available endpoints for the LLaMA Box server mode are:
       image=...                         // required
       mask=...                          // optional
       sampler=euler                     // required, select from euler_a;euler;heun;dpm2;dpm++2s_a;dpm++2m;dpm++2mv2;ipndm;ipndm_v;lcm
+      schedule=default                  // optional, select from default;discrete;karras;exponential;ays;gits
       seed=null                         // optional, random seed
       cfg_scale=4.5                     // optional, for sampler, the scale of classifier-free guidance in the output phase
       sample_steps=20                   // optional, number of sample steps

diff --git a/llama-box/server.cpp b/llama-box/server.cpp
@@ -1161,6 +1161,7 @@ struct server_context {
             slot.sdsparams.height          = json_value(data, "height", 512);
             slot.sdsparams.width           = json_value(data, "width", 512);
             slot.sdsparams.sampler         = json_value(data, "sampler", sdparams.sampler);
+            slot.sdsparams.schedule        = json_value(data, "schedule", sdparams.schedule);
             slot.sdsparams.cfg_scale       = json_value(data, "cfg_scale", sdparams.cfg_scale);
             slot.sdsparams.sample_steps    = json_value(data, "sample_steps", sdparams.sample_steps);
             slot.sdsparams.negative_prompt = json_value(data, "negative_prompt", std::string(""));

diff --git a/llama-box/stablediffusion.hpp b/llama-box/stablediffusion.hpp
@@ -57,6 +57,7 @@ struct stablediffusion_sampler_params {
     int height              = 512;
     int width               = 512;
     sample_method_t sampler = EULER_A;
+    schedule_t schedule     = DEFAULT;
     float cfg_scale         = 4.5f;
     int sample_steps        = 20;
     std::string negative_prompt;
@@ -213,6 +214,7 @@ stablediffusion_sampling_stream *stablediffusion_context::generate_stream(const
             sparams.width,
             sparams.height,
             sparams.sampler,
+            sparams.schedule,
             sparams.sample_steps,
             params.strength,
             sparams.seed,
@@ -234,6 +236,7 @@ stablediffusion_sampling_stream *stablediffusion_context::generate_stream(const
             sparams.width,
             sparams.height,
             sparams.sampler,
+            sparams.schedule,
             sparams.sample_steps,
             sparams.seed,
             control_img,

diff --git a/llama-box/tools/image_edit.sh b/llama-box/tools/image_edit.sh
@@ -28,6 +28,7 @@ QUALITY="${QUALITY:-"standard"}"
 IMAGE="${IMAGE:-""}"
 MASK="${MASK:-""}"
 SAMPLER="${SAMPLER:-"null"}"
+SCHEDULE="${SCHEDULE:-"default"}"
 SEED="${SEED:-"null"}"
 CFG_SCALE="${CFG_SCALE:-"9"}"
 SAMPLE_STEPS="${SAMPLE_STEPS:-"20"}"
@@ -110,6 +111,7 @@ image_edit() {
             --argjson response_format "\"${RESPONSE_FORMAT}\"" \
             --argjson size "\"${SIZE}\"" \
             --argjson sampler "\"${SAMPLER}\"" \
+            --argjson schedule "\"${SCHEDULE}\"" \
             --argjson seed "${SEED}" \
             --argjson cfg_scale "${CFG_SCALE}" \
             --argjson sample_steps "${SAMPLE_STEPS}" \
@@ -121,6 +123,7 @@ image_edit() {
                   response_format: $response_format,
                   size: $size,
                   sampler: $sampler,
+                  schedule: $schedule,
                   seed: $seed,
                   cfg_scale: $cfg_scale,
                   sample_steps: $sample_steps,
@@ -168,6 +171,7 @@ image_edit() {
                 --form "response_format=${RESPONSE_FORMAT}" \
                 --form "size=${SIZE}" \
                 --form "sampler=${SAMPLER}" \
+                --form "schedule=${SCHEDULE}" \
                 --form "seed=${SEED}" \
                 --form "cfg_scale=${CFG_SCALE}" \
                 --form "sample_steps=${SAMPLE_STEPS}" \
@@ -190,6 +194,7 @@ image_edit() {
                 --form "response_format=${RESPONSE_FORMAT}" \
                 --form "size=${SIZE}" \
                 --form "sampler=${SAMPLER}" \
+                --form "schedule=${SCHEDULE}" \
                 --form "seed=${SEED}" \
                 --form "cfg_scale=${CFG_SCALE}" \
                 --form "sample_steps=${SAMPLE_STEPS}" \
@@ -249,6 +254,7 @@ echo "QUALITY           : ${QUALITY} // ONE OF [standard, hd]"
 echo "IMAGE             : ${IMAGE}"
 echo "MASK              : ${MASK}"
 echo "SAMPLER           : ${SAMPLER} // OVERRIDE \"QUALITY\" and \"STYLE\" IF NOT NULL, ONE OF [euler_a, euler, heun, dpm2, dpm++2s_a, dpm++2mv2, ipndm, ipndm_v, lcm]"
+echo "SCHEDULE          : ${SCHEDULE} // AVAILABLE FOR SAMPLER, ONE OF [default, discrete, karras, exponential, ays, gits]"
 echo "SEED              : ${SEED} // AVAILABLE FOR SAMPLER"
 echo "CFG_SCALE         : ${CFG_SCALE} // AVAILABLE FOR SAMPLER"
 echo "SAMPLE_STEPS      : ${SAMPLE_STEPS} // AVAILABLE FOR SAMPLER"

diff --git a/llama-box/tools/image_generate.sh b/llama-box/tools/image_generate.sh
@@ -27,6 +27,7 @@ SIZE="${SIZE:-"512x512"}"
 QUALITY="${QUALITY:-"standard"}"
 STYLE="${STYLE:-"null"}"
 SAMPLER="${SAMPLER:-"null"}"
+SCHEDULE="${SCHEDULE:-"default"}"
 SEED="${SEED:-"null"}"
 CFG_SCALE="${CFG_SCALE:-"9"}"
 SAMPLE_STEPS="${SAMPLE_STEPS:-"20"}"
@@ -109,6 +110,7 @@ image_generate() {
             --argjson response_format "\"${RESPONSE_FORMAT}\"" \
             --argjson size "\"${SIZE}\"" \
             --argjson sampler "\"${SAMPLER}\"" \
+            --argjson schedule "\"${SCHEDULE}\"" \
             --argjson seed "${SEED}" \
             --argjson cfg_scale "${CFG_SCALE}" \
             --argjson sample_steps "${SAMPLE_STEPS}" \
@@ -118,14 +120,15 @@ image_generate() {
                   response_format: $response_format,
                   size: $size,
                   sampler: $sampler,
+                  schedule: $schedule,
                   seed: $seed,
                   cfg_scale: $cfg_scale,
                   sample_steps: $sample_steps,
                   negative_prompt: $negative_prompt,
                   stream: true,
                   stream_options: {
                     chunk_result: true,
-                    chunk_size: 32768
+                    chunk_size: 65536
                   }
                 } * .')"
     elif [[ "${STYLE}" != "null" ]]; then
@@ -144,7 +147,7 @@ image_generate() {
                   stream: true,
                   stream_options: {
                     chunk_result: true,
-                    chunk_size: 32768
+                    chunk_size: 65536
                   }
                 } * .')"
     else
@@ -161,7 +164,7 @@ image_generate() {
                   stream: true,
                   stream_options: {
                     chunk_result: true,
-                    chunk_size: 32768
+                    chunk_size: 65536
                   }
                 } * .')"
     fi
@@ -197,6 +200,7 @@ echo "SIZE              : ${SIZE}"
 echo "QUALITY           : ${QUALITY} // ONE OF [standard, hd]"
 echo "STYLE             : ${STYLE} // ONE OF [natural, vivid]"
 echo "SAMPLER           : ${SAMPLER} // OVERRIDE \"QUALITY\" and \"STYLE\" IF NOT NULL, ONE OF [euler_a, euler, heun, dpm2, dpm++2s_a, dpm++2mv2, ipndm, ipndm_v, lcm]"
+echo "SCHEDULE          : ${SCHEDULE} // AVAILABLE FOR SAMPLER, ONE OF [default, discrete, karras, exponential, ays, gits]"
 echo "SEED              : ${SEED} // AVAILABLE FOR SAMPLER"
 echo "CFG_SCALE         : ${CFG_SCALE} // AVAILABLE FOR SAMPLER"
 echo "SAMPLE_STEPS      : ${SAMPLE_STEPS} // AVAILABLE FOR SAMPLER"

diff --git a/llama-box/utils.hpp b/llama-box/utils.hpp
@@ -1014,6 +1014,7 @@ static json oaicompat_images_generations_request(const struct stablediffusion_pa
             throw std::runtime_error("Illegal param: quality must be one of 'hd' or 'standard'");
         }
         llama_params["sampler"]      = params.sampler;
+        llama_params["schedule"]     = params.schedule;
         llama_params["sample_steps"] = params.sample_steps;
         llama_params["cfg_scale"]    = params.cfg_scale;
         if (quality == "hd") {
@@ -1042,6 +1043,8 @@ static json oaicompat_images_generations_request(const struct stablediffusion_pa
     } else {
         std::string sampler_str         = json_value(body, "sampler", std::string("euler_a"));
         llama_params["sampler"]         = sd_argument_to_sample_method(sampler_str.c_str());
+        std::string schedule_str        = json_value(body, "schedule", std::string("default"));
+        llama_params["schedule"]        = sd_argument_to_schedule(schedule_str.c_str());
         llama_params["cfg_scale"]       = json_value(body, "cfg_scale", params.cfg_scale);
         llama_params["sample_steps"]    = json_value(body, "sample_steps", params.sample_steps);
         llama_params["negative_prompt"] = json_value(body, "negative_prompt", std::string(""));
@@ -1154,15 +1157,18 @@ static json oaicompat_images_edits_request(const struct stablediffusion_params &
             throw std::runtime_error("Illegal param: quality must be one of 'hd' or 'standard'");
         }
         llama_params["sampler"]      = params.sampler;
+        llama_params["schedule"]     = params.schedule;
         llama_params["sample_steps"] = params.sample_steps;
         llama_params["cfg_scale"]    = params.cfg_scale;
         if (quality == "hd") {
-            llama_params["sample_steps"]    = params.sample_steps + 10;
+            llama_params["sample_steps"]    = params.sample_steps + 2;
             llama_params["negative_prompt"] = "low quality";
         }
     } else {
         std::string sampler_str         = json_value(body, "sampler", std::string("euler_a"));
         llama_params["sampler"]         = sd_argument_to_sample_method(sampler_str.c_str());
+        std::string schedule_str        = json_value(body, "schedule", std::string("default"));
+        llama_params["schedule"]        = sd_argument_to_schedule(schedule_str.c_str());
         llama_params["cfg_scale"]       = json_value(body, "cfg_scale", params.cfg_scale);
         llama_params["sample_steps"]    = json_value(body, "sample_steps", params.sample_steps);
         llama_params["negative_prompt"] = json_value(body, "negative_prompt", std::string(""));