Upload folder using huggingface_hub

.gitignore

@@ -173,3 +173,6 @@ cython_debug/
 
 # Tensorboard
 /runs
+
+# Saved models
+saved_models

app.py

@@ -38,7 +38,7 @@ load_model_without_module(sod_model, 'u2net/saved_models/u2net-duts-msra.safeten
 
 style_files = os.listdir('./style_images')
 style_options = {' '.join(style_file.split('.')[0].split('_')): f'./style_images/{style_file}' for style_file in style_files}
-lrs = np.logspace(np.log10(0.001), np.log10(0.1), 10).tolist()
+lrs = np.logspace(np.log10(0.0025), np.log10(0.25), 10).tolist()
 img_size = 512
 
 cached_style_features = {}

u2net/inference.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from torchvision import transforms
+from safetensors.torch import load_file
 import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
@@ -46,12 +47,12 @@ def overlay_segmentation(original_image, binary_mask, alpha=0.5):
 
 if __name__ == '__main__':
     # ---
-    model_path = 'results/inter-u2net-duts.pt'
+    model_path = 'results/u2net-duts-msra.safetensors'
     image_path = 'images/ladies.jpg'
     # ---
     model = U2Net().to(device)
     model = nn.DataParallel(model)
-    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
+    model.load_state_dict(load_file(model_path, device=device.type))
     model.eval()
 
     mask = run_inference(model, image_path, threshold=None)