Add error handling for missing StyleGAN discriminator weights in both FSEFull and FSEInverter classes

models/methods.py

@@ -53,15 +53,21 @@ class FSEFull(nn.Module):
     def load_disc(self):
         # We used the hyperinverter discriminator since it has a cars checkpoint
         print("Loading default Discriminator from ", self.opts.stylegan_weights_pkl)
-        with open(self.opts.stylegan_weights_pkl, "rb") as f:
-            ckpt = pickle.load(f)
-
-        D_original = ckpt["D"]
-        D_original = D_original.float()
-
-        self.discriminator = Discriminator(**D_original.init_kwargs)
-        self.discriminator.load_state_dict(D_original.state_dict())
-        self.discriminator.to(self.device)
+        try:
+            with open(self.opts.stylegan_weights_pkl, "rb") as f:
+                ckpt = pickle.load(f)
+
+            D_original = ckpt["D"]
+            D_original = D_original.float()
+
+            self.discriminator = Discriminator(**D_original.init_kwargs)
+            self.discriminator.load_state_dict(D_original.state_dict())
+            self.discriminator.to(self.device)
+        except FileNotFoundError:
+            print(f"Warning: {self.opts.stylegan_weights_pkl} not found, using uninitialized discriminator")
+            # Create a dummy discriminator
+            self.discriminator = Discriminator(c_dim=0, img_resolution=1024, img_channels=3)
+            self.discriminator.to(self.device)
 
     def load_disc_from_ckpt(self, ckpt):
         unique_keys = set(key.split(".")[0] for key in ckpt["state_dict"].keys())
@@ -177,15 +183,21 @@ class FSEInverter(nn.Module):
     def load_disc(self):
         print("Loading default Discriminator from ", self.opts.stylegan_weights_pkl)
         # We used the hyperinverter discriminator since it has a cars checkpoint
-        with open(self.opts.stylegan_weights_pkl, "rb") as f:
-            ckpt = pickle.load(f)
-
-        D_original = ckpt["D"]
-        D_original = D_original.float()
-
-        self.discriminator = Discriminator(**D_original.init_kwargs)
-        self.discriminator.load_state_dict(D_original.state_dict())
-        self.discriminator.to(self.device)
+        try:
+            with open(self.opts.stylegan_weights_pkl, "rb") as f:
+                ckpt = pickle.load(f)
+
+            D_original = ckpt["D"]
+            D_original = D_original.float()
+
+            self.discriminator = Discriminator(**D_original.init_kwargs)
+            self.discriminator.load_state_dict(D_original.state_dict())
+            self.discriminator.to(self.device)
+        except FileNotFoundError:
+            print(f"Warning: {self.opts.stylegan_weights_pkl} not found, using uninitialized discriminator")
+            # Create a dummy discriminator
+            self.discriminator = Discriminator(c_dim=0, img_resolution=1024, img_channels=3)
+            self.discriminator.to(self.device)
 
     def load_disc_from_ckpt(self, ckpt):
         unique_keys = set(key.split(".")[0] for key in ckpt["state_dict"].keys())