Fix CLIP model device handling: pass device to StyleCLIPGlobalDirection and use .to(device) instead of .cuda()

editings/latent_editor.py

@@ -154,7 +154,7 @@ class LatentEditor:
             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)
+        global_direction_calculator = StyleCLIPGlobalDirection(delta_i_c, s_std, text_prompt_templates, device=self.device)
         return global_direction_calculator
 
 

editings/styleclip/mapper/gloabl_mapper.py

@@ -11,7 +11,7 @@ STYLESPACE_DIMENSIONS = [512 for _ in range(15)] + [256, 256, 256] + [128, 128,
 TORGB_INDICES = list(range(1, len(STYLESPACE_DIMENSIONS), 3))
 STYLESPACE_INDICES_WITHOUT_TORGB = [i for i in range(len(STYLESPACE_DIMENSIONS)) if i not in TORGB_INDICES][:11]
 
-def features_channels_to_s(s_without_torgb, s_std):
+def features_channels_to_s(s_without_torgb, s_std, device="cpu"):
     s = []
     start_index_features = 0
     for c in range(len(STYLESPACE_DIMENSIONS)):
@@ -20,19 +20,20 @@ def features_channels_to_s(s_without_torgb, s_std):
             s_i = s_without_torgb[start_index_features:end_index_features] * s_std[c]
             start_index_features = end_index_features
         else:
-            s_i = torch.zeros(STYLESPACE_DIMENSIONS[c]).cuda()
+            s_i = torch.zeros(STYLESPACE_DIMENSIONS[c]).to(device)
         s_i = s_i.view(1, 1, -1, 1, 1)
         s.append(s_i)
     return s
 
 class StyleCLIPGlobalDirection:
 
-    def __init__(self, delta_i_c, s_std, text_prompts_templates):
+    def __init__(self, delta_i_c, s_std, text_prompts_templates, device="cpu"):
         super(StyleCLIPGlobalDirection, self).__init__()
         self.delta_i_c = delta_i_c
         self.s_std = s_std
         self.text_prompts_templates = text_prompts_templates
-        self.clip_model, _ = clip.load("ViT-B/32", device="cuda")
+        self.device = device
+        self.clip_model, _ = clip.load("ViT-B/32", device=device)
 
     def get_delta_s(self, neutral_text, target_text, beta):
         delta_i = self.get_delta_i([target_text, neutral_text]).float()
@@ -43,7 +44,7 @@ class StyleCLIPGlobalDirection:
         max_channel_value = torch.abs(delta_s).max()
         if max_channel_value > 0:
             delta_s /= max_channel_value
-        direction = features_channels_to_s(delta_s, self.s_std)
+        direction = features_channels_to_s(delta_s, self.s_std, self.device)
         return direction
 
     def get_delta_i(self, text_prompts):
@@ -61,11 +62,11 @@ class StyleCLIPGlobalDirection:
             text_features_list = []
             for text_prompt in text_prompts:
                 formatted_text_prompts = [template.format(text_prompt) for template in self.text_prompts_templates]  # format with class
-                formatted_text_prompts = clip.tokenize(formatted_text_prompts).cuda()  # tokenize
+                formatted_text_prompts = clip.tokenize(formatted_text_prompts).to(self.device)  # tokenize
                 text_embeddings = self.clip_model.encode_text(formatted_text_prompts)  # embed with text encoder
                 text_embeddings /= text_embeddings.norm(dim=-1, keepdim=True)
                 text_embedding = text_embeddings.mean(dim=0)
                 text_embedding /= text_embedding.norm()
                 text_features_list.append(text_embedding)
-            text_features = torch.stack(text_features_list, dim=1).cuda()
+            text_features = torch.stack(text_features_list, dim=1).to(self.device)
         return text_features.t()