Added missing prediction type overrides

Author: daniandtheweb

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

examples/cli/main.cpp

@@ -226,7 +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, flow}        Prediction type.\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");

stable-diffusion.cpp

@@ -570,10 +570,27 @@ class StableDiffusionGGML {
                     LOG_INFO("running in v-prediction mode");
                     denoiser = std::make_shared<CompVisVDenoiser>();
                     break;
-                case FLOW_PRED:
+                case EDM_V_PRED:
+                    LOG_INFO("running in v-prediction EDM mode");
+                    denoiser = std::make_shared<EDMVDenoiser>();
+                    break;
+                case SD3_FLOW_PRED:
                     LOG_INFO("running in FLOW mode");
                     denoiser = std::make_shared<DiscreteFlowDenoiser>();
                     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<FluxFlowDenoiser>(shift); 
+                    break;
+                    }
                 default:
                     LOG_ERROR("Unknown parametrization %i", sd_ctx_params->prediction);
                     abort();
@@ -1313,7 +1330,9 @@ const char* prediction_to_str[] = {
     "default",
     "eps",
     "v",
-    "flow",
+    "edm_v",
+    "sd3_flow",
+    "flux_flow",
 };
 
 const char* sd_prediction_name(enum prediction_t prediction) {

stable-diffusion.h

@@ -64,7 +64,9 @@ enum prediction_t {
     DEFAULT_PRED,
     EPS_PRED,
     V_PRED,
-    FLOW_PRED,
+    EDM_V_PRED,
+    SD3_FLOW_PRED,
+    FLUX_FLOW_PRED,
     PREDICTION_COUNT
 };