file corruption checking system

modules/patch.py

@@ -1,3 +1,5 @@
+import os.path
+
 import torch
 import time
 import fcbh.model_base
@@ -20,6 +22,7 @@ import fcbh.cli_args
 import args_manager
 import modules.advanced_parameters as advanced_parameters
 import warnings
+import safetensors.torch
 
 from fcbh.k_diffusion import utils
 from fcbh.k_diffusion.sampling import BrownianTreeNoiseSampler, trange
@@ -481,6 +484,40 @@ def patched_load_models_gpu(*args, **kwargs):
     return y
 
 
+def build_loaded(module, loader_name):
+    original_loader_name = loader_name + '_origin'
+
+    if not hasattr(module, original_loader_name):
+        setattr(module, original_loader_name, getattr(module, loader_name))
+
+    original_loader = getattr(module, original_loader_name)
+
+    def loader(*args, **kwargs):
+        result = None
+        try:
+            result = original_loader(*args, **kwargs)
+        except Exception as e:
+            result = None
+            exp = str(e) + '\n'
+            for path in list(args) + list(kwargs.values()):
+                if isinstance(path, str):
+                    if os.path.exists(path):
+                        exp += f'File corrupted: {path} \n'
+                        corrupted_backup_file = path + '.corrupted'
+                        if os.path.exists(corrupted_backup_file):
+                            os.remove(corrupted_backup_file)
+                        os.replace(path, corrupted_backup_file)
+                        if os.path.exists(path):
+                            os.remove(path)
+                        exp += f'Fooocus has tried to move the corrupted file to {corrupted_backup_file} \n'
+                        exp += f'You may try again now and Fooocus will download models again. \n'
+            raise ValueError(exp)
+        return result
+
+    setattr(module, loader_name, loader)
+    return
+
+
 def patch_all():
     if not fcbh.model_management.DISABLE_SMART_MEMORY:
         vram_inadequate = fcbh.model_management.total_vram < 20 * 1024
@@ -509,4 +546,7 @@ def patch_all():
 
     warnings.filterwarnings(action='ignore', module='torchsde')
 
+    build_loaded(safetensors.torch, 'load_file')
+    build_loaded(torch, 'load')
+
     return