fluxdev/stable-diffusion-webui-forge

1import torch2import os3import safetensors4 5 6def build_loaded(module, loader_name):7    original_loader_name = loader_name + '_origin'8 9    if not hasattr(module, original_loader_name):10        setattr(module, original_loader_name, getattr(module, loader_name))11 12    original_loader = getattr(module, original_loader_name)13 14    def loader(*args, **kwargs):15        result = None16        try:17            result = original_loader(*args, **kwargs)18        except Exception as e:19            result = None20            exp = str(e) + '\n'21            for path in list(args) + list(kwargs.values()):22                if isinstance(path, str):23                    if os.path.exists(path):24                        exp += f'File corrupted: {path} \n'25                        corrupted_backup_file = path + '.corrupted'26                        if os.path.exists(corrupted_backup_file):27                            os.remove(corrupted_backup_file)28                        os.replace(path, corrupted_backup_file)29                        if os.path.exists(path):30                            os.remove(path)31                        exp += f'Forge has tried to move the corrupted file to {corrupted_backup_file} \n'32                        exp += f'You may try again now and Forge will download models again. \n'33            raise ValueError(exp)34        return result35 36    setattr(module, loader_name, loader)37    return38 39 40def patch_all_basics():41    build_loaded(safetensors.torch, 'load_file')42    build_loaded(torch, 'load')43    return44