Update SDLens/hooked_sd_pipeline.py

SDLens/hooked_sd_pipeline.py

@@ -3,6 +3,7 @@ from diffusers import StableDiffusionXLPipeline, IFPipeline
 from typing import List, Dict, Callable, Union
 import torch
 from .hooked_scheduler import HookedNoiseScheduler
+import spaces
 
 def retrieve(io):
     if isinstance(io, tuple):
@@ -30,7 +31,7 @@ class HookedDiffusionAbstractPipeline:
     def from_pretrained(cls, *args, **kwargs):
         return cls(cls.parent_cls.from_pretrained(*args, **kwargs))
 
-
+    @spaces.GPU
     def run_with_hooks(self, 
         *args,
         position_hook_dict: Dict[str, Union[Callable, List[Callable]]], 
@@ -69,6 +70,8 @@ class HookedDiffusionAbstractPipeline:
         
         return output
 
+    
+    @spaces.GPU
     def run_with_cache(self, 
         *args,
         positions_to_cache: List[str],
@@ -123,6 +126,8 @@ class HookedDiffusionAbstractPipeline:
             cache_dict['output'] = cache_output
         return output, cache_dict
 
+    
+    @spaces.GPU
     def run_with_hooks_and_cache(self,
         *args,
         position_hook_dict: Dict[str, Union[Callable, List[Callable]]],