extern/easydiffusion
1
0
Fork 1
mirror of https://github.com/easydiffusion/easydiffusion.git synced 2024-12-23 23:49:11 +01:00
Code Issues Packages Projects Releases Wiki Activity

Don't crash if an invalid model file is beign scanned

Browse Source
This commit is contained in:
cmdr2 2022-11-18 16:12:45 +05:30
parent 32dfb765dd
commit 921711a679
2 changed files with 12 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ui/media/js/main.js
View File

@ -1159,7 +1159,7 @@ async function getModels() {
previewPane.style.textAlign="center"
previewPane.innerHTML = '<H1>🔥Malware alert!🔥</H1><h2>The file <i>' + models['scan-error'] + '</i> in your <tt>models/stable-diffusion</tt> folder is probably malware infected.</h2><h2>Please delete this file from the folder before proceeding!</h2>After deleting the file, reload this page.<br><br><button onClick="window.location.reload();">Reload Page</button>'
makeImageBtn.disabled = true
}
}
let modelOptions = models['options']
let stableDiffusionOptions = modelOptions['stable-diffusion']
let vaeOptions = modelOptions['vae']

18
ui/server.py
View File

@ -208,13 +208,17 @@ def getModels():
os.makedirs(models_dir)
for file in os.listdir(models_dir):
scan_result = picklescan.scanner.scan_file_path( os.path.join(models_dir, file))
if ( scan_result.issues_count >0 or scan_result.infected_files >0):
rich.print(":warning: [bold red]Scan %s: %d scanned, %d issue, %d infected.[/bold red]" % ( file, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files) )
models['scan-error'] = file
return models
else:
rich.print("Scan %s: [green]%d scanned, %d issue, %d infected.[/green]" % ( file, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files ) )
try:
scan_result = picklescan.scanner.scan_file_path( os.path.join(models_dir, file))
if ( scan_result.issues_count >0 or scan_result.infected_files >0):
rich.print(":warning: [bold red]Scan %s: %d scanned, %d issue, %d infected.[/bold red]" % ( file, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files) )
models['scan-error'] = file
return models
else:
rich.print("Scan %s: [green]%d scanned, %d issue, %d infected.[/green]" % ( file, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files ) )
except Exception as e:
print('error while scanning', os.path.join(models_dir, file), 'error:', e)
for model_extension in model_extensions:
if file.endswith(model_extension):
model_name = file[:-len(model_extension)]