Upload upload_model.py with huggingface_hub

upload_model.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+"""
+Upload ISNet model to Hugging Face Hub for AutoModelForImageSegmentation
+"""
+import os
+import torch
+from huggingface_hub import HfApi, login
+from isnet_transformers import ISNetForImageSegmentation, ISNetConfig
+
+def prepare_model_for_upload():
+    """Prepare the model for Hugging Face upload"""
+    
+    # Find the model file
+    model_paths = [
+        "supplyswap_isnet.pth",
+        "isnet-transformers/supplyswap_isnet.pth",
+        "models/supplyswap_isnet.pth"
+    ]
+    
+    model_path = None
+    for path in model_paths:
+        if os.path.exists(path):
+            model_path = path
+            break
+    
+    if not model_path:
+        print("❌ No model file found! Please place your trained model file in the current directory.")
+        return None
+    
+    print(f"✅ Found model file: {model_path}")
+    
+    # Create a temporary directory for upload
+    upload_dir = "model_for_upload"
+    os.makedirs(upload_dir, exist_ok=True)
+    
+    # Create config
+    config = ISNetConfig()
+    config.save_pretrained(upload_dir)
+    
+    # Load the model
+    print("Loading model...")
+    state_dict = torch.load(model_path, map_location="cpu")
+    
+    # Create transformers model
+    model = ISNetForImageSegmentation(config)
+    model.isnet.load_state_dict(state_dict)
+    model.eval()
+    
+    # Save in transformers format
+    print(f"Saving model to {upload_dir}...")
+    model.save_pretrained(upload_dir)
+    
+    # Also save the original weights for compatibility
+    torch.save(state_dict, os.path.join(upload_dir, "supplyswap_isnet.pth"))
+    
+    # Create a model card
+    model_card_content = """---
+language: en
+tags:
+- image-segmentation
+- background-removal
+- isnet
+license: mit
+model_type: isnet
+---
+
+# ISNet Background Remover
+
+This model removes backgrounds from images using the ISNet architecture.
+
+## Usage
+
+```python
+from transformers import AutoModelForImageSegmentation
+from isnet_transformers import ISNetForImageSegmentation
+
+# Download and load the model
+model = AutoModelForImageSegmentation.from_pretrained("YOUR_USERNAME/isnet-background-remover")
+# OR
+model = ISNetForImageSegmentation.from_pretrained("YOUR_USERNAME/isnet-background-remover")
+```
+
+## Features
+- Removes backgrounds from images
+- High-quality segmentation
+- Fast inference
+- Compatible with transformers library
+
+## Model Architecture
+- Based on ISNet (Interactive Image Segmentation Network)
+- Uses U-Net style encoder-decoder architecture
+- Outputs binary masks for background removal
+"""
+    
+    with open(os.path.join(upload_dir, "README.md"), "w") as f:
+        f.write(model_card_content)
+    
+    print(f"✅ Model prepared in {upload_dir}")
+    return upload_dir
+
+def upload_to_huggingface(repo_name, upload_dir):
+    """Upload the model to Hugging Face Hub"""
+    
+    # Login to Hugging Face
+    print("Please enter your Hugging Face token:")
+    token = input().strip()
+    login(token)
+    
+    api = HfApi()
+    
+    # Create the repository
+    print(f"Creating repository: {repo_name}")
+    api.create_repo(repo_name, exist_ok=True)
+    
+    # Upload all files
+    print("Uploading model files...")
+    for root, dirs, files in os.walk(upload_dir):
+        for file in files:
+            file_path = os.path.join(root, file)
+            relative_path = os.path.relpath(file_path, upload_dir)
+            
+            with open(file_path, "rb") as f:
+                api.upload_file(
+                    path_or_fileobj=f,
+                    path_in_repo=relative_path,
+                    repo_id=repo_name
+                )
+            print(f"✅ Uploaded: {relative_path}")
+    
+    print(f"🎉 Model successfully uploaded to: https://huggingface.co/{repo_name}")
+
+if __name__ == "__main__":
+    print("🚀 Preparing ISNet model for Hugging Face upload...")
+    
+    # Prepare the model
+    upload_dir = prepare_model_for_upload()
+    
+    if upload_dir:
+        # Get repository name
+        print("\nEnter your Hugging Face username:")
+        username = input().strip()
+        repo_name = f"{username}/isnet-background-remover"
+        
+        # Upload to Hugging Face
+        upload_to_huggingface(repo_name, upload_dir)
+        
+        print(f"\n🎯 After upload, users can download your model with:")
+        print(f"from transformers import AutoModelForImageSegmentation")
+        print(f"model = AutoModelForImageSegmentation.from_pretrained('{repo_name}')")
+        print(f"\nOR")
+        print(f"from isnet_transformers import ISNetForImageSegmentation")
+        print(f"model = ISNetForImageSegmentation.from_pretrained('{repo_name}')")