fix

Author: a-r-r-o-w

src/diffusers/models/hooks.py

@@ -114,6 +114,7 @@ def reset_state(self, module):
         for hook in self.hooks:
             if hook._is_stateful:
                 hook.reset_state(module)
+        return module
 
 
 def add_hook_to_module(module: torch.nn.Module, hook: ModelHook, append: bool = False) -> torch.nn.Module:

src/diffusers/pipelines/fastercache_utils.py

@@ -526,7 +526,7 @@ def new_forward(self, module: nn.Module, *args, **kwargs) -> Any:
 
         state.iteration += 1
         output = (hidden_states, *output[1:]) if isinstance(output, tuple) else hidden_states
-        return output
+        return module._diffusers_hook.post_forward(module, output)
 
     def reset_state(self, module: nn.Module) -> nn.Module:
         module._fastercache_state.reset()