Upload folder using huggingface_hub

app.py

@@ -9,6 +9,7 @@ import torchvision.models as models
 import numpy as np
 import gradio as gr
 from gradio_imageslider import ImageSlider
+from safetensors.torch import load_file
 
 from utils import preprocess_img, preprocess_img_from_path, postprocess_img
 from vgg.vgg19 import VGG_19
@@ -22,8 +23,7 @@ print('DEVICE:', device)
 if device == 'cuda': print('CUDA DEVICE:', torch.cuda.get_device_name())
 
 def load_model_without_module(model, model_path):
-    state_dict = torch.load(model_path, map_location=device, weights_only=False)
-
+    state_dict = load_file(model_path, device=device)
     new_state_dict = {}
     for k, v in state_dict.items():
         name = k[7:] if k.startswith('module.') else k

u2net/train.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, ConcatDataset
 from torch.amp import autocast, GradScaler
+from safetensors.torch import save_file
 
 from data_loader import DUTSDataset, MSRADataset
 from model import U2Net
@@ -78,11 +79,11 @@ if __name__ == '__main__':
 
         if val_loss < best_val_loss:
             best_val_loss = val_loss
-            torch.save(model.state_dict(), f'results/best-{model_name}.pt')
+            save_file(model.state_dict(), f'results/best-{model_name}.safetensors')
             print('Best model saved.')
             
         print(f'Epoch [{epoch+1}/{epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f} (Best: {best_val_loss:.4f})')
 
-    torch.save(model.state_dict(), f'results/{model_name}.pt')
+    save_file(model.state_dict(), f'results/{model_name}.safetensors')
     with open('results/loss.txt', 'wb') as f:
         pickle.dump(losses, f)