Upload modeling_isnet.py with huggingface_hub

modeling_isnet.py

@@ -7,7 +7,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
 from transformers import PreTrainedModel, PretrainedConfig
-from transformers.models.auto.modeling_auto import AutoModelForImageSegmentation
 
 # Import the ISNet model
 import sys
@@ -24,6 +23,7 @@ class ISNetConfig(PretrainedConfig):
         self.in_ch = in_ch
         self.out_ch = out_ch
         self.num_labels = out_ch  # Required for AutoModelForImageSegmentation
+        self.architectures = ["ISNetForImageSegmentation"]
 
 class ISNetForImageSegmentation(PreTrainedModel):
     """Transformers-compatible ISNet model for image segmentation"""
@@ -78,19 +78,49 @@ class ISNetForImageSegmentation(PreTrainedModel):
                 )
                 state_dict = torch.load(model_file, map_location="cpu")
             except:
-                # Fallback to the original model file
-                model_file = cached_file(
-                    pretrained_model_name_or_path,
-                    "supplyswap_isnet.pth",
-                    **kwargs
-                )
-                state_dict = torch.load(model_file, map_location="cpu")
+                try:
+                    # Try model.safetensors
+                    model_file = cached_file(
+                        pretrained_model_name_or_path,
+                        "model.safetensors",
+                        **kwargs
+                    )
+                    from safetensors import safe_open
+                    with safe_open(model_file, framework="pt", device="cpu") as f:
+                        state_dict = {key: f.get_tensor(key) for key in f.keys()}
+                except:
+                    # Fallback to the original model file
+                    model_file = cached_file(
+                        pretrained_model_name_or_path,
+                        "supplyswap_isnet.pth",
+                        **kwargs
+                    )
+                    state_dict = torch.load(model_file, map_location="cpu")
+        
+        # Handle different state dict formats
+        if isinstance(state_dict, dict):
+            # Check if the state dict has the expected keys
+            if any(key.startswith('isnet.') for key in state_dict.keys()):
+                # State dict already has the correct prefix
+                pass
+            elif any(key.startswith('conv_in.') or key.startswith('stage') for key in state_dict.keys()):
+                # State dict is from the original ISNet model, needs to be wrapped
+                wrapped_state_dict = {}
+                for key, value in state_dict.items():
+                    wrapped_state_dict[f"isnet.{key}"] = value
+                state_dict = wrapped_state_dict
+            else:
+                # Try to load directly
+                pass
         
         # Load the weights into the ISNet model
-        model.isnet.load_state_dict(state_dict)
+        try:
+            model.isnet.load_state_dict(state_dict)
+        except Exception as e:
+            print(f"Warning: Could not load state dict directly: {e}")
+            print("Attempting to load with strict=False...")
+            model.isnet.load_state_dict(state_dict, strict=False)
+        
         model.eval()
         
-        return model
-
-# Register the model with transformers
-AutoModelForImageSegmentation.register(ISNetConfig, ISNetForImageSegmentation) 
+        return model