Update app.py

app.py

@@ -1,7 +1,8 @@
-import torch
+from safetensors.torch import load_file
 import os
 import requests
 
+# URLs de los archivos
 MODEL_URL = "https://huggingface.co/datasets/uruguayai/fooocus/resolve/main/juggernautXL_v8Rundiffusion.safetensors"
 MODEL_PATH = "/home/user/app/juggernautXL_v8Rundiffusion.safetensors"
 
@@ -9,19 +10,23 @@ def download_file(url, path):
     """Descarga un archivo desde una URL si no está presente en la ruta especificada."""
     if not os.path.exists(path):
         response = requests.get(url)
-        with open(path, 'wb') as f:
-            f.write(response.content)
-        print(f"Downloaded {path}")
+        if response.status_code == 200:
+            with open(path, 'wb') as f:
+                f.write(response.content)
+            print(f"Downloaded {path}")
+        else:
+            raise Exception(f"Failed to download {url}, status code: {response.status_code}")
 
 def load_model():
+    # Descargar el archivo si no existe
     download_file(MODEL_URL, MODEL_PATH)
-    # Cargar con weights_only=False si confías en la fuente del modelo
-    model = torch.load(MODEL_PATH, weights_only=False)
+    # Cargar los pesos utilizando safetensors
+    model = load_file(MODEL_PATH)  # Carga el modelo desde el archivo safetensors
     return model
 
 def main():
     model = load_model()
-    print("Model loaded successfully!")
+    print("Model loaded successfully with safetensors!")
 
 if __name__ == "__main__":
     main()