Add error handling for missing checkpoint files in load_weights methods

models/methods.py

@@ -88,29 +88,44 @@ class FSEFull(nn.Module):
     def load_weights(self):
         if self.opts.checkpoint_path != "":
             print(f"Loading from checkpoint: {self.opts.checkpoint_path}")
-            ckpt = torch.load(self.opts.checkpoint_path, map_location="cpu")
-            self.load_disc_from_ckpt(ckpt)
-            self.encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
-            self.inverter.load_state_dict(get_keys(ckpt, "inverter"), strict=True)
+            try:
+                ckpt = torch.load(self.opts.checkpoint_path, map_location="cpu")
+                self.load_disc_from_ckpt(ckpt)
+                self.encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+                self.inverter.load_state_dict(get_keys(ckpt, "inverter"), strict=True)
+            except FileNotFoundError:
+                print(f"Warning: {self.opts.checkpoint_path} not found, using uninitialized weights")
         else:
             print(f"Loading Discriminator and Inverter from Inverter checkpoint: {self.inverter_pth}")
-            ckpt = torch.load(self.inverter_pth, map_location="cpu")
-            self.load_disc_from_ckpt(ckpt)
-            self.inverter.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+            try:
+                ckpt = torch.load(self.inverter_pth, map_location="cpu")
+                self.load_disc_from_ckpt(ckpt)
+                self.inverter.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+            except FileNotFoundError:
+                print(f"Warning: {self.inverter_pth} not found, using uninitialized weights")
 
         self.inverter = self.inverter.eval().to(self.device)
         toogle_grad(self.inverter, False)
 
         print("Loading Decoder from", self.opts.stylegan_weights)
-        ckpt = torch.load(self.opts.stylegan_weights)
-        self.decoder.load_state_dict(ckpt["g_ema"], strict=False)
-        self.latent_avg = ckpt['latent_avg'].to(self.device)
+        try:
+            ckpt = torch.load(self.opts.stylegan_weights)
+            self.decoder.load_state_dict(ckpt["g_ema"], strict=False)
+            self.latent_avg = ckpt['latent_avg'].to(self.device)
+        except FileNotFoundError:
+            print(f"Warning: {self.opts.stylegan_weights} not found, using uninitialized decoder")
+            self.latent_avg = torch.zeros(18, 512).to(self.device)
+        
         self.decoder = self.decoder.eval().to(self.device)
         toogle_grad(self.decoder, False)
 
         print("Loading E4E from", self.opts.e4e_path)
-        ckpt = torch.load(self.opts.e4e_path, map_location="cpu")
-        self.e4e_encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+        try:
+            ckpt = torch.load(self.opts.e4e_path, map_location="cpu")
+            self.e4e_encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+        except FileNotFoundError:
+            print(f"Warning: {self.opts.e4e_path} not found, using uninitialized E4E encoder")
+        
         self.e4e_encoder = self.e4e_encoder.eval().to(self.device)
         toogle_grad(self.e4e_encoder, False)
 
@@ -227,14 +242,21 @@ class FSEInverter(nn.Module):
     def load_weights(self):
         if self.opts.checkpoint_path != "":
             print("Loading  from checkpoint: {}".format(self.opts.checkpoint_path))
-            ckpt = torch.load(self.opts.checkpoint_path, map_location="cpu")
-            self.load_disc_from_ckpt(ckpt)
-            self.encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+            try:
+                ckpt = torch.load(self.opts.checkpoint_path, map_location="cpu")
+                self.load_disc_from_ckpt(ckpt)
+                self.encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
+            except FileNotFoundError:
+                print(f"Warning: {self.opts.checkpoint_path} not found, using uninitialized weights")
 
         print("Loading decoder from", self.opts.stylegan_weights)
-        ckpt = torch.load(self.opts.stylegan_weights)
-        self.decoder.load_state_dict(ckpt["g_ema"], strict=False)
-        self.latent_avg = ckpt['latent_avg'].to(self.device)
+        try:
+            ckpt = torch.load(self.opts.stylegan_weights)
+            self.decoder.load_state_dict(ckpt["g_ema"], strict=False)
+            self.latent_avg = ckpt['latent_avg'].to(self.device)
+        except FileNotFoundError:
+            print(f"Warning: {self.opts.stylegan_weights} not found, using uninitialized decoder")
+            self.latent_avg = torch.zeros(18, 512).to(self.device)
 
     def set_encoder(self):
         inverter = psp_encoders.Inverter(opts=self.opts, n_styles=18)