fixing

SAE/sae.py

@@ -105,7 +105,6 @@ class SparseAutoencoder(nn.Module):
     def n_dirs(self):
         return self.n_dirs_local
 
-    @spaces.GPU
     def encode(self, x):
         x = x.to('cuda') - self.pre_bias
         latents_pre_act = self.encoder(x) + self.latent_bias

SDLens/hooked_sd_pipeline.py

@@ -309,7 +309,6 @@ class HookedDiffusionAbstractPipeline:
     def __setattr__(self, name, value):
         return setattr(self.pipe, name, value)
 
-    @spaces.GPU
     def __call__(self, *args, **kwargs):
         return self.pipe(*args, **kwargs)
 

app.py

@@ -20,7 +20,7 @@ code_to_block = {
 }
 lock = threading.Lock()
 
-@spaces.GPU
+
 def process_cache(cache, saes_dict):
 
     top_features_dict = {}
@@ -70,6 +70,9 @@ def create_prompt_part(pipe, saes_dict, demo):
     @spaces.GPU
     def image_gen(prompt):
         lock.acquire()
+        pipe.to('cuda')
+        for sae in saes_dict.values:
+            sae.to('cuda')
         try:
             images, cache = pipe.run_with_cache(
                 prompt,
@@ -376,6 +379,7 @@ if __name__ == "__main__":
         variant=("fp16" if dtype==torch.float16 else None)
     )
     pipe.set_progress_bar_config(disable=True)
+    pipe.to('cuda')
 
     path_to_checkpoints = './checkpoints/'
 

utils/hooks.py

@@ -1,7 +1,6 @@
 import torch
 import spaces
 
-@spaces.GPU
 @torch.no_grad()
 def add_feature(sae, feature_idx, value, module, input, output):
     diff = (output[0] - input[0]).permute((0, 2, 3, 1)).to(sae.device)
@@ -12,7 +11,6 @@ def add_feature(sae, feature_idx, value, module, input, output):
     return (output[0] + to_add.permute(0, 3, 1, 2).to(output[0].device),)
 
 
-@spaces.GPU
 @torch.no_grad()
 def add_feature_on_area(sae, feature_idx, activation_map, module, input, output):
     diff = (output[0] - input[0]).permute((0, 2, 3, 1)).to(sae.device)
@@ -25,7 +23,6 @@ def add_feature_on_area(sae, feature_idx, activation_map, module, input, output)
     return (output[0] + to_add.permute(0, 3, 1, 2).to(output[0].device),)
 
 
-@spaces.GPU
 @torch.no_grad()
 def replace_with_feature(sae, feature_idx, value, module, input, output):
     diff = (output[0] - input[0]).permute((0, 2, 3, 1)).to(sae.device)
@@ -36,7 +33,6 @@ def replace_with_feature(sae, feature_idx, value, module, input, output):
     return (input[0] + to_add.permute(0, 3, 1, 2).to(output[0].device),)
 
 
-@spaces.GPU
 @torch.no_grad()
 def reconstruct_sae_hook(sae, module, input, output):
     diff = (output[0] - input[0]).permute((0, 2, 3, 1)).to(sae.device)
@@ -45,7 +41,6 @@ def reconstruct_sae_hook(sae, module, input, output):
     return (input[0] + reconstructed.permute(0, 3, 1, 2).to(output[0].device),)
 
 
-@spaces.GPU
 @torch.no_grad()
 def ablate_block(module, input, output):
     return input