feat: add RegionE support for Qwen-Image-Edit.

xllm/core/common/global_flags.h

@@ -360,6 +360,12 @@ DECLARE_int64(dit_cache_start_blocks);
 
 DECLARE_int64(dit_cache_end_blocks);
 
+DECLARE_string(dit_regione_refresh_steps);
+DECLARE_double(dit_regione_region_threshold);
+DECLARE_double(dit_regione_cache_threshold);
+DECLARE_bool(dit_regione_use_avd_gamma);
+DECLARE_bool(dit_regione_erosion_dilation);
+
 DECLARE_bool(dit_sp_communication_overlap);
 
 DECLARE_int64(dit_generation_image_area_max);

xllm/core/framework/config/dit_config.cpp

@@ -23,7 +23,7 @@ DEFINE_int32(max_requests_per_batch, 1, "Max number of request per batch.");
 DEFINE_string(dit_cache_policy,
               "TaylorSeer",
               "The policy of dit cache(e.g. None, FBCache, TaylorSeer, "
-              "FBCacheTaylorSeer, ResidualCache).");
+              "FBCacheTaylorSeer, ResidualCache, RegionE).");
 
 DEFINE_int64(dit_cache_warmup_steps, 0, "The number of warmup steps.");
 
@@ -53,6 +53,31 @@ DEFINE_int64(dit_cache_end_blocks,
              5,
              "The number of blocks to skip at the end.");
 
+DEFINE_string(dit_regione_refresh_steps,
+              "16",
+              "RegionE: comma-separated full-image refresh steps in RAGS.");
+
+DEFINE_double(dit_regione_region_threshold,
+              0.80,
+              "RegionE: cosine threshold for adaptive region partition.");
+
+DEFINE_double(dit_regione_cache_threshold,
+              0.02,
+              "RegionE: AVDCache error threshold δ (paper Eq.8). "
+              "Reuse velocity while 1-accumulate <= threshold. "
+              "Default is 0.02.");
+
+DEFINE_bool(dit_regione_use_avd_gamma,
+            true,
+            "RegionE: use AVDCache with gamma (paper/inplace.py method). "
+            "Uses the original diffusers 28-step gamma curve, linearly "
+            "upsampled/downsampled to the actual inference step count. "
+            "Set false to use fixed skip_interval instead.");
+
+DEFINE_bool(dit_regione_erosion_dilation,
+            true,
+            "RegionE: enable erosion/dilation for region mask cleanup.");
+
 DEFINE_bool(dit_sp_communication_overlap,
             true,
             "Communication & Computation overlap for sequence parallel");
@@ -131,6 +156,11 @@ void DiTConfig::from_flags() {
   XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_steps);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_start_blocks);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_blocks);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_refresh_steps);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_region_threshold);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_cache_threshold);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_use_avd_gamma);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_erosion_dilation);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_laser_attention_enabled);
@@ -157,6 +187,11 @@ void DiTConfig::from_json(const JsonReader& json) {
   XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_steps);
   XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_start_blocks);
   XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_blocks);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_refresh_steps);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_region_threshold);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_cache_threshold);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_use_avd_gamma);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_erosion_dilation);
   XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap);
   XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print);
   XLLM_CONFIG_ASSIGN_FROM_JSON(dit_laser_attention_enabled);
@@ -194,6 +229,16 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const {
       config_json, default_config, dit_cache_start_blocks);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, dit_cache_end_blocks);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, dit_regione_refresh_steps);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, dit_regione_region_threshold);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, dit_regione_cache_threshold);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, dit_regione_use_avd_gamma);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, dit_regione_erosion_dilation);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, dit_sp_communication_overlap);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(

xllm/core/framework/config/dit_config.h

@@ -51,6 +51,11 @@ class DiTConfig final {
          "dit_cache_end_steps",
          "dit_cache_start_blocks",
          "dit_cache_end_blocks",
+         "dit_regione_refresh_steps",
+         "dit_regione_region_threshold",
+         "dit_regione_cache_threshold",
+         "dit_regione_use_avd_gamma",
+         "dit_regione_erosion_dilation",
          "dit_sp_communication_overlap",
          "dit_debug_print",
          "dit_laser_attention_enabled",
@@ -87,6 +92,16 @@ class DiTConfig final {
 
   PROPERTY(int64_t, dit_cache_end_blocks) = 5;
 
+  PROPERTY(std::string, dit_regione_refresh_steps) = "16";
+
+  PROPERTY(double, dit_regione_region_threshold) = 0.80;
+
+  PROPERTY(double, dit_regione_cache_threshold) = 0.02;
+
+  PROPERTY(bool, dit_regione_use_avd_gamma) = true;
+
+  PROPERTY(bool, dit_regione_erosion_dilation) = true;
+
   PROPERTY(bool, dit_sp_communication_overlap) = true;
 
   PROPERTY(bool, dit_debug_print) = false;

xllm/core/framework/dit_cache/CMakeLists.txt

@@ -13,6 +13,7 @@ cc_library(
   fbcache_taylorseer.h
   taylorseer.h
   residual_cache.h
+  regione.h
  SRCS
   dit_cache_impl.cpp
   dit_cache.cpp
@@ -21,8 +22,10 @@ cc_library(
   fbcache_taylorseer.cpp
   taylorseer.cpp
   residual_cache.cpp
+  regione.cpp
  DEPS
   torch
+  $<$<BOOL:${USE_NPU}>:torch_npu>
   glog::glog
   Folly::folly
   parallel_state

xllm/core/framework/dit_cache/dit_cache.cpp

@@ -19,6 +19,12 @@ namespace xllm {
 
 bool DiTCache::init(const DiTCacheConfig& cfg,
                     const ParallelArgs& parallel_args) {
+  regione_cache_.reset();
+  if (cfg.selected_policy == PolicyType::RegionE) {
+    regione_cache_ = std::make_unique<RegionECache>();
+    regione_cache_->init(cfg);
+  }
+
   active_cache_ = create_dit_cache(cfg);
   active_cond_cache_ = create_dit_cache(cfg);
   if (!active_cache_ || !active_cond_cache_) {

xllm/core/framework/dit_cache/dit_cache.h

@@ -14,7 +14,12 @@ limitations under the License.
 ==============================================================================*/
 
 #pragma once
+
+#include <memory>
+#include <utility>
+
 #include "dit_cache_impl.h"
+#include "regione.h"
 
 namespace xllm {
 
@@ -36,24 +41,30 @@ class DiTCache {
   bool init(const DiTCacheConfig& cfg, const ParallelArgs& parallel_args);
 
   bool on_before_block(const CacheBlockIn& blockin, bool use_cfg = false);
-
   CacheBlockOut on_after_block(const CacheBlockIn& blockin,
                                bool use_cfg = false);
-
   bool on_before_step(const CacheStepIn& stepin, bool use_cfg = false);
-
   CacheStepOut on_after_step(const CacheStepIn& stepin, bool use_cfg = false);
 
   void set_context(const CacheContext& context) {
+    if (regione_cache_) {
+      regione_cache_->set_infer_steps(context.infer_steps);
+      regione_cache_->set_num_blocks(context.num_blocks);
+    }
     active_cache_->set_context(context);
     active_cond_cache_->set_context(context);
   }
 
+  RegionECache* regione() { return regione_cache_.get(); }
+  const RegionECache* regione() const { return regione_cache_.get(); }
+
  private:
-  torch::Tensor get_tensor_or_empty(const TensorMap& m, const std::string& k);
+  static torch::Tensor get_tensor_or_empty(const TensorMap& m,
+                                           const std::string& k);
 
   std::unique_ptr<DitCacheImpl> active_cache_;
   std::unique_ptr<DitCacheImpl> active_cond_cache_;
+  std::unique_ptr<RegionECache> regione_cache_;
 };
 
 }  // namespace xllm

xllm/core/framework/dit_cache/dit_cache_config.h

@@ -15,14 +15,19 @@ limitations under the License.
 
 #pragma once
 
+#include <cstdint>
+#include <string>
+#include <vector>
+
 namespace xllm {
 
 enum class PolicyType {
   None,
   FBCache,
   TaylorSeer,
   FBCacheTaylorSeer,
-  ResidualCache
+  ResidualCache,
+  RegionE
 };
 
 struct DiTBaseCacheOptions {
@@ -51,6 +56,21 @@ struct FBCacheTaylorSeerOptions : public DiTBaseCacheOptions {
   int n_derivatives = 3;
 };
 
+struct RegionEOptions : public DiTBaseCacheOptions {
+  // Fallback fixed-interval AVD when gamma is disabled or step count
+  // mismatches.
+  int64_t skip_interval_steps = 3;
+  int64_t tail_steps = 1;
+  std::vector<int64_t> refresh_steps = {16};
+  float region_threshold = 0.80f;
+  // AVDCache δ in paper Eq.8/9.
+  float cache_threshold = 0.02f;
+  // Use fitted γ_t AVDCache (paper) instead of fixed skip_interval.
+  bool use_avd_gamma = true;
+  // Enable erosion/dilation morphological cleanup after ARP mask selection.
+  bool erosion_dilation = true;
+};
+
 struct ResidualCacheOptions {
   // The number of steps to skip at the start.
   int64_t dit_cache_start_steps = 5;
@@ -85,6 +105,9 @@ struct DiTCacheConfig {
 
   // the configuration for ResidualCache policy.
   ResidualCacheOptions residual_cache;
+
+  // the configuration for RegionE policy.
+  RegionEOptions regione;
 };
 
 }  // namespace xllm

xllm/core/framework/dit_cache/dit_cache_impl.cpp

@@ -82,6 +82,8 @@ std::unique_ptr<DitCacheImpl> create_dit_cache(const DiTCacheConfig& cfg) {
       return std::make_unique<FBCacheTaylorSeer>();
     case PolicyType::ResidualCache:
       return std::make_unique<ResidualCache>();
+    case PolicyType::RegionE:
+      return std::make_unique<DiTNonCache>();
     default:
       return std::make_unique<DiTNonCache>();
   }