update code example

Author: a-r-r-o-w

src/diffusers/hooks/hooks.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 
@@ -141,9 +141,9 @@ def new_forward(module, *args, **kwargs):
         self.hooks[name] = hook
         self._hook_order.append(name)
 
-    def get_hook(self, name: str) -> ModelHook:
+    def get_hook(self, name: str) -> Optional[ModelHook]:
         if name not in self.hooks.keys():
-            raise ValueError(f"Hook with name {name} not found.")
+            return None
         return self.hooks[name]
 
     def remove_hook(self, name: str) -> None:

src/diffusers/hooks/pyramid_attention_broadcast.py

@@ -169,12 +169,15 @@ def apply_pyramid_attention_broadcast(
     ```python
     >>> import torch
     >>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
+    >>> from diffusers.utils import export_to_video
 
     >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
     >>> pipe.to("cuda")
 
     >>> config = PyramidAttentionBroadcastConfig(
-    ...     spatial_attention_block_skip_range=2, spatial_attention_timestep_skip_range=(100, 800)
+    ...     spatial_attention_block_skip_range=2,
+    ...     spatial_attention_timestep_skip_range=(100, 800),
+    ...     current_timestep_callback=lambda: pipe._current_timestep,
     ... )
     >>> apply_pyramid_attention_broadcast(pipe.transformer, config)
     ```