MOSS-Image: extra sample args to window or disable the block-22 negative-branch skip

examples/common/common.cpp

@@ -1180,7 +1180,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
+         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; MOSS-Image supports moss_uncond_skip, moss_uncond_skip_start, moss_uncond_skip_end (sigma-percent window of its block-22 negative-branch skip); lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
          (int)',',
          &extra_sample_args},
         {"",

src/pipeline/diffusion_engine.cpp

@@ -2472,6 +2472,40 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
     float img_cfg_scale = guidance.img_cfg;
     float slg_scale     = guidance.slg.scale;
     bool slg_uncond     = sd::guidance::parse_skip_layer_guidance_uncond_arg(extra_sample_args);
+    const sd::guidance::MossUncondSkipParams moss_uncond_skip =
+        version == VERSION_MOSS_IMAGE ? sd::guidance::parse_moss_uncond_skip_args(extra_sample_args)
+                                      : sd::guidance::MossUncondSkipParams{};
+    // Sigma bounds of the MOSS window (ComfyUI ModelSamplingDiscreteFlow.percent_to_sigma, multiplier 1).
+    float moss_skip_sigma_high = INFINITY;
+    float moss_skip_sigma_low  = -INFINITY;
+    if (version == VERSION_MOSS_IMAGE && moss_uncond_skip.enabled && moss_uncond_skip.windowed()) {
+        auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
+        float shift        = flow_denoiser ? flow_denoiser->shift : 5.0f;
+        auto percent_to_sigma = [shift](float percent) {
+            if (percent <= 0.0f) {
+                return 1.0f;
+            }
+            if (percent >= 1.0f) {
+                return 0.0f;
+            }
+            return time_snr_shift(shift, 1.0f - percent);
+        };
+        moss_skip_sigma_high = percent_to_sigma(moss_uncond_skip.start);
+        moss_skip_sigma_low  = percent_to_sigma(moss_uncond_skip.end);
+        LOG_INFO("MOSS uncond skip window: sigma-percent %.4f..%.4f (sigma %.4f..%.4f, shift %.3f)",
+                 moss_uncond_skip.start, moss_uncond_skip.end, moss_skip_sigma_low, moss_skip_sigma_high, shift);
+    } else if (version == VERSION_MOSS_IMAGE && !moss_uncond_skip.enabled) {
+        LOG_INFO("MOSS uncond skip disabled");
+    }
+    int moss_skipped_steps = 0;
+    // Step caches (EasyCache / UCache / CacheDIT) key the negative prediction by &uncond only, so a windowed skip
+    // drops that entry whenever the effective skip set flips; -1 = no negative pass yet. Never touched by default.
+    int moss_prev_uncond_skip   = -1;
+    int moss_cache_invalidations = 0;
+    // Spectrum extrapolates the guided output from its own history, which that drop never reaches: its history is
+    // cleared at each window edge instead. -1 = no step seen yet.
+    int moss_spectrum_window = -1;
+    int moss_spectrum_resets = 0;
 
     std::vector<float> guidance_schedule = sd::guidance::parse_guidance_schedule(extra_sample_args);
     if (!guidance_schedule.empty() && guidance_schedule.size() != sigmas.size() - 1) {
@@ -2608,6 +2642,17 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
         }
 
+        if (cache_runtime.spectrum_enabled && version == VERSION_MOSS_IMAGE && moss_uncond_skip.enabled &&
+            moss_uncond_skip.windowed() && !uncond.empty()) {
+            const int moss_window_now = static_cast<int>(sigma >= moss_skip_sigma_low && sigma <= moss_skip_sigma_high);
+            if (moss_spectrum_window >= 0 && moss_spectrum_window != moss_window_now) {
+                cache_runtime.spectrum.reset_history();
+                ++moss_spectrum_resets;
+                LOG_VERBOSE("MOSS uncond skip set changed at step %d: Spectrum history dropped\n", step);
+            }
+            moss_spectrum_window = moss_window_now;
+        }
+
         if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
             cache_runtime.spectrum.predict(&denoised);
             if (apply_denoise_mask) {
@@ -2736,9 +2781,17 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             }
 
             sd::Tensor<float> cached_output;
+            const bool moss_trace_uncond = version == VERSION_MOSS_IMAGE && moss_uncond_skip.enabled &&
+                                           moss_uncond_skip.windowed() && &condition == &uncond;
             if (step_cache.before_condition(&condition, noised_input, &cached_output)) {
+                if (moss_trace_uncond) {
+                    LOG_VERBOSE("MOSS uncond served from step cache at step %d\n", step);
+                }
                 return std::move(cached_output);
             }
+            if (moss_trace_uncond) {
+                LOG_VERBOSE("MOSS uncond computed at step %d\n", step);
+            }
 
             // A re-enabled condition can miss the cache even when the positive pass was reused.
             if (!uncond_controls_ready && !uncond.empty() &&
@@ -2907,8 +2960,22 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                 if (is_skiplayer_step && slg_uncond) {
                     LOG_VERBOSE("Skipping layers at uncond step %d\n", step);
                     uncond_skip_layers = &skip_layer_guidance.layers();
-                } else if (version == VERSION_MOSS_IMAGE) {
-                    uncond_skip_layers = &moss_skip_layers;
+                } else if (version == VERSION_MOSS_IMAGE && moss_uncond_skip.enabled) {
+                    const bool moss_skip_now = !moss_uncond_skip.windowed() ||
+                                               (sigma >= moss_skip_sigma_low && sigma <= moss_skip_sigma_high);
+                    if (moss_uncond_skip.windowed()) {
+                        if (moss_prev_uncond_skip >= 0 && moss_prev_uncond_skip != static_cast<int>(moss_skip_now)) {
+                            step_cache.invalidate_condition(&uncond);
+                            ++moss_cache_invalidations;
+                            LOG_VERBOSE("MOSS uncond skip set changed at step %d: negative-branch step cache dropped\n", step);
+                        }
+                        moss_prev_uncond_skip = static_cast<int>(moss_skip_now);
+                    }
+                    if (moss_skip_now) {
+                        LOG_VERBOSE("MOSS uncond skip layers at step %d (sigma %.4f)\n", step, sigma);
+                        uncond_skip_layers = &moss_skip_layers;
+                        ++moss_skipped_steps;
+                    }
                 }
                 uncond_out = run_condition(uncond,
                                            uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
@@ -2999,6 +3066,15 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
 
     auto x0 = std::move(x0_opt);
     sd_sample::log_sample_cache_summary(cache_runtime, steps);
+    if (version == VERSION_MOSS_IMAGE) {
+        LOG_INFO("MOSS uncond skip: %d of %zu steps skipped block 22 on the negative branch", moss_skipped_steps, steps);
+        if (moss_cache_invalidations > 0) {
+            LOG_INFO("MOSS uncond skip: %d window edge(s) dropped the cached negative prediction", moss_cache_invalidations);
+        }
+        if (moss_spectrum_resets > 0) {
+            LOG_INFO("MOSS uncond skip: %d window edge(s) reset the Spectrum history", moss_spectrum_resets);
+        }
+    }
     if (inverse_noise_scaling) {
         x0 = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x0);
     }

src/runtime/guidance.cpp

@@ -64,6 +64,34 @@ namespace sd::guidance {
         return uncond;
     }
 
+    MossUncondSkipParams parse_moss_uncond_skip_args(const char* extra_sample_args) {
+        MossUncondSkipParams params;
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) {
+            if (key == "moss_uncond_skip") {
+                bool parsed = true;
+                if (!parse_strict_bool(value, parsed)) {
+                    LOG_WARN("ignoring invalid MOSS extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                    continue;
+                }
+                params.enabled = parsed;
+            } else if (key == "moss_uncond_skip_start" || key == "moss_uncond_skip_end") {
+                float parsed = 0.0f;
+                if (!parse_strict_float(value, parsed) || !(parsed >= 0.0f && parsed <= 1.0f)) {
+                    LOG_WARN("ignoring invalid MOSS extra sample arg '%s=%s' (expected 0..1)", key.c_str(), value.c_str());
+                    continue;
+                }
+                (key == "moss_uncond_skip_start" ? params.start : params.end) = parsed;
+            }
+        }
+        if (params.start > params.end) {
+            LOG_WARN("moss_uncond_skip_start %.4f > moss_uncond_skip_end %.4f: keeping the skip on every step",
+                     params.start, params.end);
+            params.start = 0.0f;
+            params.end   = 1.0f;
+        }
+        return params;
+    }
+
     std::vector<float> parse_guidance_schedule_from_spec(std::string spec) {
         std::vector<float> schedule;
 

src/runtime/guidance.h

@@ -28,6 +28,17 @@ namespace sd::guidance {
     AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args);
     bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params);
     bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args);
+
+    // MOSS-Image's built-in negative-branch skip (block 22 on the uncond pass). extra_sample_args
+    // moss_uncond_skip=false turns it off; moss_uncond_skip_start / moss_uncond_skip_end bound it in
+    // sigma-percent of the shifted flow schedule (ComfyUI percent_to_sigma). Defaults keep every step.
+    struct MossUncondSkipParams {
+        bool enabled = true;
+        float start  = 0.0f;
+        float end    = 1.0f;
+        bool windowed() const { return start > 0.0f || end < 1.0f; }
+    };
+    MossUncondSkipParams parse_moss_uncond_skip_args(const char* extra_sample_args);
     std::vector<float> parse_guidance_schedule(const char* extra_sample_args);
 
     struct GuidanceInput {

src/runtime/spectrum.hpp

@@ -74,6 +74,13 @@ struct SpectrumState {
         cnt++;
     }
 
+    // Forget the fitted outputs but keep the step count, so predictions restart from fresh model evaluations.
+    void reset_history() {
+        H_buf.clear();
+        T_buf.clear();
+        num_cached = 0;
+    }
+
     void predict(sd::Tensor<float>* denoised) {
         GGML_ASSERT(denoised != nullptr);
         int64_t F    = (int64_t)H_buf[0].size();