[Enhancement] Add force_refresh support for GLM-Image for cache-dit 1.3.0 upgrade (#1858)
Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
This commit is contained in:
1 parent
b2b061a9b3
commit
9ead0d890d
3 files changed
+72
No files matched your search
+33
@@ -61,6 +61,8 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig:
|
||||
max_cached_steps=cache_config.max_cached_steps,
|
||||
max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
|
||||
residual_diff_threshold=cache_config.residual_diff_threshold,
|
||||
force_refresh_step_hint=cache_config.force_refresh_step_hint,
|
||||
force_refresh_step_policy=cache_config.force_refresh_step_policy,
|
||||
)
|
||||
|
||||
|
||||
@@ -1091,6 +1093,36 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
|
||||
return refresh_cache_context
|
||||
|
||||
|
||||
def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
|
||||
"""Enable cache-dit for GLM-Image pipeline.
|
||||
|
||||
GLM-Image processes prompt and image by calling the transformer before the
|
||||
denoising loop. When an input image is provided (editing mode), the cache must
|
||||
be force-refreshed after the preprocessing step so stale hidden states are
|
||||
discarded. Set force_refresh_step_hint = 1 for editing, None for text-to-image.
|
||||
"""
|
||||
db_cache_config = _build_db_cache_config(cache_config)
|
||||
|
||||
calibrator_config = None
|
||||
if cache_config.enable_taylorseer:
|
||||
calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=cache_config.taylorseer_order)
|
||||
logger.info(f"TaylorSeer enabled with order={cache_config.taylorseer_order}")
|
||||
|
||||
logger.info(
|
||||
f"Enabling cache-dit on GLM-Image transformer: "
|
||||
f"Fn={db_cache_config.Fn_compute_blocks}, "
|
||||
f"Bn={db_cache_config.Bn_compute_blocks}, "
|
||||
f"W={db_cache_config.max_warmup_steps}, "
|
||||
f"force_refresh_step_hint={db_cache_config.force_refresh_step_hint}, "
|
||||
)
|
||||
|
||||
cache_dit.enable_cache(
|
||||
pipeline.transformer,
|
||||
cache_config=db_cache_config,
|
||||
calibrator_config=calibrator_config,
|
||||
)
|
||||
|
||||
|
||||
def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
|
||||
"""Enable cache-dit for Flux.2-dev pipeline.
|
||||
|
||||
@@ -1180,6 +1212,7 @@ CUSTOM_DIT_ENABLERS.update(
|
||||
"LTX2Pipeline": enable_cache_for_ltx2,
|
||||
"LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
|
||||
"BagelPipeline": enable_cache_for_bagel,
|
||||
"GlmImagePipeline": enable_cache_for_glm_image,
|
||||
"Flux2Pipeline": enable_cache_for_flux2,
|
||||
}
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user