Avoid CUDA init on Spaces: map tensors to device, no .cuda() in LatentEditor

editings/latent_editor.py

@@ -42,9 +42,10 @@ STYLESPACE_IDX = [
 
 
 class LatentEditor:
-    def __init__(self, domain="human_faces"):
+    def __init__(self, domain="human_faces", device="cpu"):
 
         self.domain = domain
+        self.device = torch.device(device)
 
         if self.domain == "human_faces":
             self.interfacegan_directions = {
@@ -53,11 +54,11 @@ class LatentEditor:
                 "rotation": "editings/interfacegan_directions/rotation.pt",
             }
             self.interfacegan_tensors = {
-                name: torch.load(path).cuda()
+                name: torch.load(path, map_location=self.device)
                 for name, path in self.interfacegan_directions.items()
             }
 
-            self.ganspace_pca = torch.load("editings/ganspace_pca/ffhq_pca.pt")
+            self.ganspace_pca = torch.load("editings/ganspace_pca/ffhq_pca.pt", map_location=self.device)
             self.ganspace_directions = {
                 "eye_openness": (54, 7, 8, 5),
                 "trimmed_beard": (58, 7, 9, 7),
@@ -147,10 +148,10 @@ class LatentEditor:
 
 
     def load_styleclip_global(self):
-        delta_i_c = torch.from_numpy(np.load("editings/styleclip/global_mapper_data/delta_i_c.npy")).float().cuda()
+        delta_i_c = torch.from_numpy(np.load("editings/styleclip/global_mapper_data/delta_i_c.npy")).float().to(self.device)
         with open("editings/styleclip/global_mapper_data/S_mean_std", "rb") as channels_statistics:
             _, s_std = pickle.load(channels_statistics)
-            s_std = [torch.from_numpy(s_i).float().cuda() for s_i in s_std]
+            s_std = [torch.from_numpy(s_i).float().to(self.device) for s_i in s_std]
         with open("editings/styleclip/global_mapper_data/templates.txt", "r") as templates:
             text_prompt_templates = templates.readlines()
         global_direction_calculator = StyleCLIPGlobalDirection(delta_i_c, s_std, text_prompt_templates)
@@ -179,7 +180,7 @@ class LatentEditor:
         opts = argparse.Namespace(**opts)
         style_clip_net = StyleCLIPMapper(opts)
         style_clip_net.eval()
-        style_clip_net.cuda()
+        style_clip_net.to(self.device)
         direction = style_clip_net.mapper(start_w)
         for factor in factors:
             edited_latent = start_w + factor * direction
@@ -195,7 +196,7 @@ class LatentEditor:
         disentanglement = float(disentanglement)
 
         directions = self.styleclip_global_editor.get_delta_s(neutral_text, target_text, disentanglement)
-        factors = torch.tensor(factors).cuda().view(-1, 1)
+        factors = torch.tensor(factors).to(self.device).view(-1, 1)
         srart_ss, start_rgb = start_s
 
         edits_rgb = []
@@ -222,7 +223,7 @@ class LatentEditor:
         neutral_text, target_text, disentanglement = direction.split("_")
         disentanglement = float(disentanglement)
 
-        factors = torch.tensor(factors).cuda().view(-1, 1)
+        factors = torch.tensor(factors).to(self.device).view(-1, 1)
         srart_ss, edited_rgb = start_s
         edits_ss = self.deltaedit_editor.get_delta_s(neutral_text, target_text, disentanglement, original_image, srart_ss)
 

runners/base_runner.py

@@ -74,7 +74,8 @@ class BaseRunner:
         return edited_latents
 
     def _setup_latent_editor(self):
-        self.latent_editor = LatentEditor(self.config.exp.domain)
+        # Pass device to avoid unintended CUDA initialization on Spaces
+        self.latent_editor = LatentEditor(self.config.exp.domain, device=self.device)
 
     def _setup_device(self):
         config_device = self.config.model["device"].lower()