Update app.py

app.py

@@ -10,6 +10,7 @@ import torch.nn as nn
 from transformers import pipeline
 from PIL import Image
 import inspect
+import safetensors.torch
 
 # =============================================================================
 # Aesthetic-Shadow (using Hugging Face transformers pipeline)
@@ -73,8 +74,12 @@ def load_clip_models(name: str = "ViT-L/14", device='cuda'):
 
 def load_model(model_path: str, input_size=768, device: str = 'cuda', dtype=None):
     model = MLP(input_size=input_size)
-    state = torch.load(model_path, map_location=device, weights_only=False)
-    model.load_state_dict(state)
+    if model_path.endswith(".safetensors"):
+        state_dict = safetensors.torch.load_file(model_path, device=device)
+    else:
+        state = torch.load(model_path, map_location=device, weights_only=False)
+        state_dict = state
+    model.load_state_dict(state_dict)
     model.to(device)
     if dtype:
         model = model.to(dtype=dtype)