fixes

Author: a-r-r-o-w

src/diffusers/__init__.py

@@ -28,6 +28,7 @@
 
 _import_structure = {
     "configuration_utils": ["ConfigMixin"],
+    "hooks": [],
     "loaders": ["FromOriginalModelMixin"],
     "models": [],
     "pipelines": [],
@@ -77,6 +78,7 @@
 else:
     _import_structure["hooks"].extend(
         [
+            "HookRegistry",
             "PyramidAttentionBroadcastConfig",
             "apply_pyramid_attention_broadcast",
         ]
@@ -592,7 +594,7 @@
     except OptionalDependencyNotAvailable:
         from .utils.dummy_pt_objects import *  # noqa F403
     else:
-        from .hooks import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
+        from .hooks import HookRegistry, PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
         from .models import (
             AllegroTransformer3DModel,
             AsymmetricAutoencoderKL,

src/diffusers/hooks/__init__.py

@@ -2,4 +2,5 @@
 
 
 if is_torch_available():
+    from .hooks import HookRegistry
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast

src/diffusers/hooks/hooks.py

@@ -147,6 +147,12 @@ def remove_hook(self, name: str) -> None:
         del self.hooks[name]
         self._hook_order.remove(name)
 
+    def reset_stateful_hooks(self):
+        for hook_name in self._hook_order:
+            hook = self.hooks[hook_name]
+            if hook._is_stateful:
+                hook.reset_state(self._module_ref)
+
     @classmethod
     def check_if_exists_or_initialize(cls, module: torch.nn.Module) -> "HookRegistry":
         if not hasattr(module, "_diffusers_hook"):

src/diffusers/pipelines/allegro/pipeline_allegro.py

@@ -24,9 +24,9 @@
 from transformers import T5EncoderModel, T5Tokenizer
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import HookRegistry
 from ...models import AllegroTransformer3DModel, AutoencoderKLAllegro
 from ...models.embeddings import get_3d_rotary_pos_embed_allegro
-from ...models.hooks import reset_stateful_hooks
 from ...pipelines.pipeline_utils import DiffusionPipeline
 from ...schedulers import KarrasDiffusionSchedulers
 from ...utils import (
@@ -948,7 +948,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (video,)

src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py

@@ -21,10 +21,10 @@
 from transformers import T5EncoderModel, T5Tokenizer
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import HookRegistry
 from ...loaders import CogVideoXLoraLoaderMixin
 from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
 from ...models.embeddings import get_3d_rotary_pos_embed
-from ...models.hooks import reset_stateful_hooks
 from ...pipelines.pipeline_utils import DiffusionPipeline
 from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
@@ -778,7 +778,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (video,)

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py

@@ -22,10 +22,10 @@
 from transformers import T5EncoderModel, T5Tokenizer
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import HookRegistry
 from ...loaders import CogVideoXLoraLoaderMixin
 from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
 from ...models.embeddings import get_3d_rotary_pos_embed
-from ...models.hooks import reset_stateful_hooks
 from ...pipelines.pipeline_utils import DiffusionPipeline
 from ...schedulers import KarrasDiffusionSchedulers
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
@@ -831,7 +831,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (video,)

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py

@@ -22,11 +22,11 @@
 from transformers import T5EncoderModel, T5Tokenizer
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import HookRegistry
 from ...image_processor import PipelineImageInput
 from ...loaders import CogVideoXLoraLoaderMixin
 from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
 from ...models.embeddings import get_3d_rotary_pos_embed
-from ...models.hooks import reset_stateful_hooks
 from ...pipelines.pipeline_utils import DiffusionPipeline
 from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
 from ...utils import (
@@ -892,7 +892,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (video,)

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py

@@ -22,10 +22,10 @@
 from transformers import T5EncoderModel, T5Tokenizer
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import HookRegistry
 from ...loaders import CogVideoXLoraLoaderMixin
 from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
 from ...models.embeddings import get_3d_rotary_pos_embed
-from ...models.hooks import reset_stateful_hooks
 from ...pipelines.pipeline_utils import DiffusionPipeline
 from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
@@ -857,7 +857,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (video,)

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -26,10 +26,10 @@
     T5TokenizerFast,
 )
 
+from ...hooks import HookRegistry
 from ...image_processor import PipelineImageInput, VaeImageProcessor
 from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
 from ...models import AutoencoderKL, FluxTransformer2DModel
-from ...models.hooks import reset_stateful_hooks
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import (
     USE_PEFT_BACKEND,
@@ -971,7 +971,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (image,)

src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py

@@ -20,9 +20,9 @@
 from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import HookRegistry
 from ...loaders import HunyuanVideoLoraLoaderMixin
 from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel
-from ...models.hooks import reset_stateful_hooks
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
@@ -692,7 +692,7 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
-        reset_stateful_hooks(self.transformer, recurse=True)
+        HookRegistry.check_if_exists_or_initialize(self.transformer).reset_stateful_hooks()
 
         if not return_dict:
             return (video,)