Skip to content

vllm_omni.diffusion.cache.seacache.hook

logger module-attribute 

logger = init_logger(__name__)

SeaCacheRootHook

Bases: ModelHook

Drive SeaCache gating and transformer forward control.

config instance-attribute 

config = config

current_sigma_callback instance-attribute 

current_sigma_callback = current_sigma_callback

current_step_callback instance-attribute 

current_step_callback = current_step_callback

extractor_fn instance-attribute 

extractor_fn = extractor_fn

full_count instance-attribute 

full_count = 0

num_inference_steps_callback instance-attribute 

num_inference_steps_callback = num_inference_steps_callback

skip_count instance-attribute 

skip_count = 0

state_manager instance-attribute 

state_manager = StateManager(SeaCacheState)

cache_context 

cache_context(name: str) -> Iterator[None]

initialize_hook 

initialize_hook(module: Module) -> Module

new_forward 

new_forward(
    module: Module, *args: Any, **kwargs: Any
) -> Any

refresh 

refresh(module: Module) -> None

reset_state 

reset_state(module: Module) -> Module

apply_sea_cache_hook 

apply_sea_cache_hook(
    module: Module,
    config: SeaCacheConfig,
    *,
    current_step_callback: Callable[[], int | Tensor | None]
    | None = None,
    current_sigma_callback: Callable[
        [], float | Tensor | None
    ]
    | None = None,
    num_inference_steps_callback: Callable[
        [], int | Tensor | None
    ]
    | None = None,
    extractor_fn: Callable[..., CacheContext] | None = None,
) -> SeaCacheRootHook