Fix DeltaEditor device handling: add device parameter, error handling for missing files, replace .cuda() with .to(device)

editings/deltaedit/editor.py

@@ -24,14 +24,25 @@ def improved_ds(ds, select):
 
 
 class DeltaEditor:
-    def __init__(self):
-        device = "cuda"
-        self.fs3 = np.load("pretrained_models/fs3.npy")
+    def __init__(self, device="cpu"):
+        self.device = device
+        try:
+            self.fs3 = np.load("pretrained_models/fs3.npy")
+        except FileNotFoundError:
+            # If fs3.npy is not available, create a dummy array
+            # This is a fallback for when the file is not downloaded yet
+            self.fs3 = np.zeros((512, 512))  # Dummy fs3 array
+            print("Warning: fs3.npy not found, using dummy array")
+        
         np.set_printoptions(suppress=True)
 
         self.net = DeltaMapper()
-        net_ckpt = torch.load("pretrained_models/delta_mapper.pt")
-        self.net.load_state_dict(net_ckpt)
+        try:
+            net_ckpt = torch.load("pretrained_models/delta_mapper.pt", map_location=device)
+            self.net.load_state_dict(net_ckpt)
+        except FileNotFoundError:
+            print("Warning: delta_mapper.pt not found, using uninitialized network")
+        
         self.net = self.net.to(device).eval()
 
         self.clip_model, self.preprocess = clip.load("ViT-B/32", device=device)
@@ -43,11 +54,11 @@ class DeltaEditor:
             classnames = [target, neutral]
             dt = map_tool.GetDt(classnames, self.clip_model)
             select = GetBoundary(self.fs3, dt, trash)
-            dt = torch.Tensor(dt).cuda()
+            dt = torch.Tensor(dt).to(self.device)
             dt = dt / dt.norm(dim=-1, keepdim=True).float().clamp(min=1e-5)
 
             img_gen_for_clip = self.avg_pool(orig_image)
-            c_latents = self.clip_model.encode_image(img_gen_for_clip.cuda())
+            c_latents = self.clip_model.encode_image(img_gen_for_clip.to(self.device))
             c_latents = c_latents / c_latents.norm(dim=-1, keepdim=True).float()
 
             delta_c = torch.cat((c_latents, dt.unsqueeze(0)), dim=1)

editings/latent_editor.py

@@ -120,7 +120,7 @@ class LatentEditor:
                 "fs_makeup": "editings/bound/Heavy_Makeup_boundary.npy"
              }
 
-            self.deltaedit_editor = DeltaEditor()
+            self.deltaedit_editor = DeltaEditor(device=self.device)
 
         elif self.domain == "car":