From cd79877e2a0e6afc98d2dd5b9de496e26e9a4601 Mon Sep 17 00:00:00 2001 From: "remyx-ai[bot]" <289541483+remyx-ai[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 22:00:48 +0000 Subject: [PATCH 1/2] Clockwork Diffusion: Efficient Generation With Model-Step Distillation --- src/diffusers/__init__.py | 4 + src/diffusers/hooks/__init__.py | 1 + src/diffusers/hooks/clockwork_cache.py | 310 ++++++++++++++++++++++++ src/diffusers/utils/dummy_pt_objects.py | 19 ++ tests/hooks/test_clockwork_cache.py | 115 +++++++++ 5 files changed, 449 insertions(+) create mode 100644 src/diffusers/hooks/clockwork_cache.py create mode 100644 tests/hooks/test_clockwork_cache.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index da77fa67df52..702f814df661 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -175,6 +175,7 @@ ) _import_structure["hooks"].extend( [ + "ClockworkCacheConfig", "FasterCacheConfig", "FirstBlockCacheConfig", "HookRegistry", @@ -184,6 +185,7 @@ "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "TextKVCacheConfig", + "apply_clockwork_cache", "apply_faster_cache", "apply_first_block_cache", "apply_layer_skip", @@ -1036,6 +1038,7 @@ TangentialClassifierFreeGuidance, ) from .hooks import ( + ClockworkCacheConfig, FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, @@ -1045,6 +1048,7 @@ SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, TextKVCacheConfig, + apply_clockwork_cache, apply_faster_cache, apply_first_block_cache, apply_layer_skip, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..832819ff201b 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -16,6 +16,7 @@ if is_torch_available(): + from .clockwork_cache import ClockworkCacheConfig, apply_clockwork_cache from .context_parallel import apply_context_parallel from .faster_cache import FasterCacheConfig, apply_faster_cache from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache diff --git a/src/diffusers/hooks/clockwork_cache.py b/src/diffusers/hooks/clockwork_cache.py new file mode 100644 index 000000000000..3ee607bff4c9 --- /dev/null +++ b/src/diffusers/hooks/clockwork_cache.py @@ -0,0 +1,310 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +import torch + +from ..utils import get_logger +from ..utils.torch_utils import unwrap_module +from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS +from ._helpers import TransformerBlockRegistry +from .hooks import BaseState, HookRegistry, ModelHook, StateManager + + +logger = get_logger(__name__) # pylint: disable=invalid-name + +_CLOCKWORK_LEADER_BLOCK_HOOK = "clockwork_leader_block_hook" +_CLOCKWORK_BLOCK_HOOK = "clockwork_block_hook" + + +@dataclass +class ClockworkCacheConfig: + r""" + Configuration for [Clockwork + Cache](https://huggingface.co/papers/2312.08128), an inference-time adaptation of the Clockwork Diffusion caching + insight. + + Clockwork Diffusion observes that the *low-resolution / semantically important* feature maps inside a denoiser + evolve slowly across diffusion steps, whereas the high-resolution feature maps are sensitive to small + perturbations. It therefore reuses the slow-changing features across a window of steps and recomputes them only + periodically, on a fixed "clock" schedule. This hook ports that cross-step reuse policy onto diffusers' + transformer block stack: every `clock_interval` denoising steps the full block stack is recomputed, and on the + steps in between the cached block-stack output is replayed instead. This is a *periodic* schedule, in contrast to + [`apply_first_block_cache`], which decides reuse adaptively from the residual magnitude at every step. + + Note: the original paper additionally *distills* a student model that is robust to the frozen low-resolution + features ("model-step distillation") to recover the quality lost by freezing them. That training step is out of + scope for this inference-only hook and is left to whoever trains a compatible checkpoint — the other caches in + this gallery ship the inference-time acceleration of their papers without the accompanying training procedures, + and this one follows the same convention. + + Args: + clock_interval (`int`, defaults to `3`): + The number of denoising steps between two full recomputations of the transformer block stack. Every + `clock_interval`-th step recomputes; the steps in between reuse the cached output. `1` disables caching + (every step recomputes). Must be `>= 1`. + warmup (`int`, defaults to `1`): + The number of leading denoising steps that always recompute, to populate and stabilize the cache before + any reuse happens. Must be `>= 0`. + """ + + clock_interval: int = 3 + warmup: int = 1 + + def __post_init__(self): + if self.clock_interval < 1: + raise ValueError(f"`clock_interval` must be >= 1, got {self.clock_interval}.") + if self.warmup < 0: + raise ValueError(f"`warmup` must be >= 0, got {self.warmup}.") + + +class ClockworkCacheState(BaseState): + def __init__(self) -> None: + super().__init__() + + # Cached leader output and residual, used to reconstruct the block-stack output on reuse steps. + self.head_block_output: torch.Tensor | tuple[torch.Tensor, ...] = None + self.head_block_residual: torch.Tensor = None + # Delta between the tail and leader outputs captured on the last compute step. + self.tail_block_residuals: torch.Tensor | tuple[torch.Tensor, ...] = None + # Per-step bookkeeping shared between the leader and the remaining blocks. + self.should_compute: bool = True + # Denoising-step counter, persisted across the whole `__call__` so the clock can reference it. + self.step: int = -1 + self.last_compute_step: int = -1 + + def reset(self): + # Mirrors FirstBlockCache: only the per-step bookkeeping is reset; the leader/tail residuals and the step + # counter persist across the generation (they are rebuilt from scratch by StateManager when the cache is + # reset before a new generation). + self.tail_block_residuals = None + self.should_compute = True + + +class ClockworkLeaderBlockHook(ModelHook): + _is_stateful = True + + def __init__(self, state_manager: StateManager, clock_interval: int, warmup: int): + self.state_manager = state_manager + self.clock_interval = clock_interval + self.warmup = warmup + self._metadata = None + + def initialize_hook(self, module): + unwrapped_module = unwrap_module(module) + self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__) + return module + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs) + + # The leader always recomputes itself; it is the reference against which the cached tail delta is applied. + output = self.fn_ref.original_forward(*args, **kwargs) + is_output_tuple = isinstance(output, tuple) + + if is_output_tuple: + hidden_states_residual = output[self._metadata.return_hidden_states_index] - original_hidden_states + else: + hidden_states_residual = output - original_hidden_states + + shared_state: ClockworkCacheState = self.state_manager.get_state() + shared_state.step += 1 + + should_compute = self._should_compute(shared_state) + shared_state.should_compute = should_compute + + if not should_compute: + # Reuse step: replay the cached block-stack output by adding the stale tail delta to the current leader + # output, then skip the remaining blocks (they pass this through unchanged). + if is_output_tuple: + hidden_states = ( + shared_state.tail_block_residuals[0] + output[self._metadata.return_hidden_states_index] + ) + else: + hidden_states = shared_state.tail_block_residuals[0] + output + + encoder_hidden_states = None + if self._metadata.return_encoder_hidden_states_index is not None: + assert is_output_tuple + encoder_hidden_states = ( + shared_state.tail_block_residuals[1] + output[self._metadata.return_encoder_hidden_states_index] + ) + + if is_output_tuple: + return_output = [None] * len(output) + return_output[self._metadata.return_hidden_states_index] = hidden_states + return_output[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states + return_output = tuple(return_output) + else: + return_output = hidden_states + output = return_output + else: + # Compute step: cache the leader output/residual so the tail can record the block-stack delta. + shared_state.last_compute_step = shared_state.step + if is_output_tuple: + head_block_output = [None] * len(output) + head_block_output[0] = output[self._metadata.return_hidden_states_index] + head_block_output[1] = output[self._metadata.return_encoder_hidden_states_index] + else: + head_block_output = output + shared_state.head_block_output = head_block_output + shared_state.head_block_residual = hidden_states_residual + + return output + + def reset_state(self, module): + self.state_manager.reset() + return module + + @torch.compiler.disable + def _should_compute(self, shared_state: ClockworkCacheState) -> bool: + # No cache populated yet -> the first step must always compute. + if shared_state.head_block_residual is None: + return True + # Warmup window always recomputes to stabilize the cache. + if shared_state.step < self.warmup: + return True + # Clock: recompute every `clock_interval` steps since the last compute. + if shared_state.step - shared_state.last_compute_step >= self.clock_interval: + return True + return False + + +class ClockworkBlockHook(ModelHook): + def __init__(self, state_manager: StateManager, is_tail: bool = False): + super().__init__() + self.state_manager = state_manager + self.is_tail = is_tail + self._metadata = None + + def initialize_hook(self, module): + unwrapped_module = unwrap_module(module) + self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__) + return module + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs) + original_encoder_hidden_states = None + if self._metadata.return_encoder_hidden_states_index is not None: + original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( + "encoder_hidden_states", args, kwargs + ) + + shared_state = self.state_manager.get_state() + + if shared_state.should_compute: + output = self.fn_ref.original_forward(*args, **kwargs) + if self.is_tail: + # Record the delta between the tail and leader outputs for reuse on later steps. + if isinstance(output, tuple): + hidden_states_residual = ( + output[self._metadata.return_hidden_states_index] - shared_state.head_block_output[0] + ) + encoder_hidden_states_residual = ( + output[self._metadata.return_encoder_hidden_states_index] - shared_state.head_block_output[1] + ) + else: + hidden_states_residual = output - shared_state.head_block_output + encoder_hidden_states_residual = None + shared_state.tail_block_residuals = (hidden_states_residual, encoder_hidden_states_residual) + return output + + # Reuse step: pass the input through unchanged; the cached delta is applied by the leader block. + if original_encoder_hidden_states is None: + return_output = original_hidden_states + else: + return_output = [None, None] + return_output[self._metadata.return_hidden_states_index] = original_hidden_states + return_output[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states + return_output = tuple(return_output) + return return_output + + +def apply_clockwork_cache(module: torch.nn.Module, config: ClockworkCacheConfig) -> None: + """ + Applies [Clockwork Cache](https://huggingface.co/papers/2312.08128) to a given module. + + Clockwork Cache reuses the slowly-evolving transformer block stack output across a window of denoising steps and + recomputes it on a fixed periodic schedule (the "clock"). It is built on the same `StateManager` mechanism as + [`apply_first_block_cache`], but selects reuse steps by a periodic schedule rather than by per-step residual + magnitude. + + Args: + module (`torch.nn.Module`): + The pytorch module to apply Clockwork Cache to. Typically this should be a transformer architecture + supported in Diffusers, such as `CogVideoXTransformer3DModel`, but external implementations may also + work. + config (`ClockworkCacheConfig`): + The configuration to use for applying the Clockwork Cache method. + + Example: + ```python + >>> import torch + >>> from diffusers import CogView4Pipeline + >>> from diffusers.hooks import ClockworkCacheConfig, apply_clockwork_cache + + >>> pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> apply_clockwork_cache(pipe.transformer, ClockworkCacheConfig(clock_interval=3, warmup=1)) + + >>> prompt = "A photo of an astronaut riding a horse on mars" + >>> image = pipe(prompt, generator=torch.Generator().manual_seed(42)).images[0] + >>> image.save("output.png") + ``` + """ + + state_manager = StateManager(ClockworkCacheState, (), {}) + remaining_blocks = [] + + for name, submodule in module.named_children(): + if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList): + continue + for index, block in enumerate(submodule): + remaining_blocks.append((f"{name}.{index}", block)) + + if len(remaining_blocks) < 2: + raise ValueError( + "Clockwork Cache requires at least two transformer blocks to cache (a leader block and a tail block). " + f"Found {len(remaining_blocks)} block(s) under known transformer-block identifiers " + f"({_ALL_TRANSFORMER_BLOCK_IDENTIFIERS})." + ) + + head_block_name, head_block = remaining_blocks.pop(0) + tail_block_name, tail_block = remaining_blocks.pop(-1) + + logger.debug(f"Applying ClockworkLeaderBlockHook to '{head_block_name}'") + _apply_clockwork_leader_block_hook(head_block, state_manager, config) + + for name, block in remaining_blocks: + logger.debug(f"Applying ClockworkBlockHook to '{name}'") + _apply_clockwork_block_hook(block, state_manager) + + logger.debug(f"Applying ClockworkBlockHook to tail block '{tail_block_name}'") + _apply_clockwork_block_hook(tail_block, state_manager, is_tail=True) + + +def _apply_clockwork_leader_block_hook( + block: torch.nn.Module, state_manager: StateManager, config: ClockworkCacheConfig +) -> None: + registry = HookRegistry.check_if_exists_or_initialize(block) + hook = ClockworkLeaderBlockHook(state_manager, config.clock_interval, config.warmup) + registry.register_hook(hook, _CLOCKWORK_LEADER_BLOCK_HOOK) + + +def _apply_clockwork_block_hook(block: torch.nn.Module, state_manager: StateManager, is_tail: bool = False) -> None: + registry = HookRegistry.check_if_exists_or_initialize(block) + hook = ClockworkBlockHook(state_manager, is_tail) + registry.register_hook(hook, _CLOCKWORK_BLOCK_HOOK) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 8439a2b93371..fdd1358c4681 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -167,6 +167,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ClockworkCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class FasterCacheConfig(metaclass=DummyObject): _backends = ["torch"] @@ -306,6 +321,10 @@ def apply_faster_cache(*args, **kwargs): requires_backends(apply_faster_cache, ["torch"]) +def apply_clockwork_cache(*args, **kwargs): + requires_backends(apply_clockwork_cache, ["torch"]) + + def apply_first_block_cache(*args, **kwargs): requires_backends(apply_first_block_cache, ["torch"]) diff --git a/tests/hooks/test_clockwork_cache.py b/tests/hooks/test_clockwork_cache.py new file mode 100644 index 000000000000..cee27adc76f0 --- /dev/null +++ b/tests/hooks/test_clockwork_cache.py @@ -0,0 +1,115 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import ClockworkCacheConfig, apply_clockwork_cache +from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry +from diffusers.models import ModelMixin + + +class CountingBlock(torch.nn.Module): + """A transformer block that doubles its input and records how often its forward runs.""" + + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + self.forward_calls += 1 + return hidden_states * 2.0 + + +class CountingTransformer(ModelMixin): + def __init__(self, num_blocks: int = 3): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([CountingBlock() for _ in range(num_blocks)]) + + def forward(self, hidden_states, encoder_hidden_states=None): + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states) + return hidden_states + + +@pytest.fixture(autouse=True) +def register_counting_blocks(): + TransformerBlockRegistry.register( + CountingBlock, + TransformerBlockMetadata(return_hidden_states_index=None, return_encoder_hidden_states_index=None), + ) + + +def _set_context(model, context_name): + """Helper to set context on all hooks in the model.""" + for module in model.modules(): + if hasattr(module, "_diffusers_hook"): + module._diffusers_hook._set_context(context_name) + + +def test_clockwork_cache_validation(): + """Invalid config is rejected.""" + with pytest.raises(ValueError): + ClockworkCacheConfig(clock_interval=0) + with pytest.raises(ValueError): + ClockworkCacheConfig(warmup=-1) + + +def test_clockwork_cache_clock_schedule(): + """ + Every `clock_interval` steps the block stack recomputes; the steps in between reuse the cached output, so the + non-leader blocks are skipped. With clock_interval=3 and warmup=1 over 4 steps: steps 0 and 3 compute, steps 1 + and 2 reuse. + """ + model = CountingTransformer(num_blocks=3) + apply_clockwork_cache(model, ClockworkCacheConfig(clock_interval=3, warmup=1)) + + _set_context(model, "default") + x = torch.randn(1, 4) + + for _ in range(4): + model(x) + + leader, middle, tail = model.transformer_blocks + # The leader recomputes itself every step (it is the reference for the cached delta). + assert leader.forward_calls == 4 + # Middle and tail only run on the compute steps (0 and 3). + assert middle.forward_calls == 2 + assert tail.forward_calls == 2 + + +def test_clockwork_cache_replays_correctly_for_static_input(): + """ + For a deterministic block, the cached replay on a reuse step must match a full recomputation of the same input. + """ + cached_model = CountingTransformer(num_blocks=3) + plain_model = CountingTransformer(num_blocks=3) + apply_clockwork_cache(cached_model, ClockworkCacheConfig(clock_interval=3, warmup=1)) + + _set_context(cached_model, "default") + x = torch.randn(1, 4) + + expected = plain_model(x) + + # Step 0 is a compute step (warmup) and must reproduce the full forward. + assert torch.allclose(cached_model(x), expected) + # Step 1 is a reuse step and must replay the cached output. + assert torch.allclose(cached_model(x), expected) + + +def test_clockwork_cache_requires_multiple_blocks(): + """A single-block transformer cannot be cached (needs a leader and a tail).""" + model = CountingTransformer(num_blocks=1) + with pytest.raises(ValueError, match="at least two transformer blocks"): + apply_clockwork_cache(model, ClockworkCacheConfig()) From 31c3bbf443f7560e0b02431d333b0c3b2a964b69 Mon Sep 17 00:00:00 2001 From: "remyx-ai[bot]" <289541483+remyx-ai[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 22:16:30 +0000 Subject: [PATCH 2/2] chore: align PR with target-repo conventions Convention-shape patches extracted from huggingface/diffusers's recent merged PRs. Algorithm logic is left untouched. Ruff auto-fixed lint-trivial issues on patched files. --- docs/source/en/api/cache.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/source/en/api/cache.md b/docs/source/en/api/cache.md index a5ed8751118d..89da0a3887d7 100644 --- a/docs/source/en/api/cache.md +++ b/docs/source/en/api/cache.md @@ -46,3 +46,9 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate [[autodoc]] MagCacheConfig [[autodoc]] apply_mag_cache + +## ClockworkCacheConfig + +[[autodoc]] ClockworkCacheConfig + +[[autodoc]] apply_clockwork_cache