holm/easydiffusion
1
0
Fork 0
forked from extern/easydiffusion
Code Issues Pull Requests 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

4
ui/server.py
View File

@ -208,6 +208,7 @@ def getModels():
os.makedirs(models_dir)
for file in os.listdir(models_dir):
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) )
@ -215,6 +216,9 @@ def getModels():
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)]