fixed prediction selection logic

leejet · Aug 7, 2024 · e995b3b · e995b3b
1 parent 1e9fd07
commit e995b3b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -49,6 +49,7 @@ const char* schedule_str[] = {
 };
 
 const char* prediction_str[] = {
+    "default",
     "eps",
     "v",
     "flow",
@@ -104,7 +105,7 @@ struct SDParams {
 
     sample_method_t sample_method = EULER_A;
     schedule_t schedule           = DEFAULT;
-    prediction_t prediction       = EPS_PRED;
+    prediction_t prediction       = DEFAULT_PRED;
     int sample_steps              = 20;
     float strength                = 0.75f;
     float control_strength        = 0.9f;

diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -423,7 +423,7 @@ class StableDiffusionGGML {
         LOG_INFO("loading model from '%s' completed, taking %.2fs", model_path.c_str(), (t1 - t0) * 1.0f / 1000);
 
 
-        if (prediction != EPS_PRED) {
+        if (prediction != DEFAULT_PRED) {
             switch (prediction) {
                 case EPS_PRED:
                     LOG_INFO("running in eps-prediction mode");

diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -54,6 +54,7 @@ enum schedule_t {
 };
 
 enum prediction_t {
+    DEFAULT_PRED,
     EPS_PRED,
     V_PRED,
     FLOW_PRED,