Improve weight loading verification and skip configs with too many missing keys

app/pytorch_colorizer.py

@@ -78,10 +78,11 @@ class ResNetGenerator(nn.Module):
         model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
         model += [nn.Tanh()]
         
-        self.model = nn.Sequential(*model)
+        # Wrap in Sequential with 'layers' attribute to match state_dict structure
+        self.layers = nn.Sequential(*model)
     
     def forward(self, input):
-        return self.model(input)
+        return self.layers(input)
 
 
 class UNetGenerator(nn.Module):
@@ -211,15 +212,18 @@ class PyTorchColorizer:
                 
                 # Log state dict keys to understand model structure
                 if isinstance(state_dict, dict):
-                    keys = list(state_dict.keys())[:20]  # First 20 keys
+                    keys = list(state_dict.keys())[:30]  # First 30 keys
                     logger.info(f"Model state_dict keys (sample): {keys}")
                     logger.info(f"Total state_dict keys: {len(state_dict.keys())}")
                     
                     # Try to infer architecture from key names
+                    if any('layers' in k.lower() for k in keys):
+                        logger.info("Detected sequential 'layers' structure")
                     if any('down' in k.lower() or 'up' in k.lower() for k in keys):
                         logger.info("Detected U-Net style architecture")
                     if any('resnet' in k.lower() for k in keys):
                         logger.info("Detected ResNet style architecture")
+                    
                 
             except Exception as e:
                 logger.error(f"Failed to load model file: {e}")
@@ -250,12 +254,17 @@ class PyTorchColorizer:
                     
                     # Try strict loading first
                     try:
-                        model.load_state_dict(state_dict, strict=True)
-                        logger.info(f"✅ Successfully loaded {model_type} model with strict matching: {config_copy}")
-                    except:
-                        # If strict fails, try non-strict
-                        model.load_state_dict(state_dict, strict=False)
-                        logger.info(f"✅ Successfully loaded {model_type} model with non-strict matching: {config_copy}")
+                        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+                        if not missing_keys and not unexpected_keys:
+                            logger.info(f"✅ Successfully loaded {model_type} model with perfect matching: {config_copy}")
+                        else:
+                            logger.warning(f"⚠️ Loaded {model_type} model with mismatches - Missing: {len(missing_keys)}, Unexpected: {len(unexpected_keys)}")
+                            if len(missing_keys) > len(state_dict) * 0.5:  # If more than 50% missing, skip
+                                logger.warning(f"Skipping this config - too many missing keys ({len(missing_keys)}/{len(state_dict)})")
+                                continue
+                    except Exception as e:
+                        logger.debug(f"Failed to load state_dict: {e}")
+                        continue
                     
                     model.eval()
                     model.to(self.device)