diff --git a/ui/media/js/main.js b/ui/media/js/main.js index d22d3169..27da7b2a 100644 --- a/ui/media/js/main.js +++ b/ui/media/js/main.js @@ -1159,7 +1159,7 @@ async function getModels() { previewPane.style.textAlign="center" previewPane.innerHTML = '

🔥Malware alert!🔥

The file ' + models['scan-error'] + ' in your models/stable-diffusion folder is probably malware infected.

Please delete this file from the folder before proceeding!

After deleting the file, reload this page.

' makeImageBtn.disabled = true - } + } let modelOptions = models['options'] let stableDiffusionOptions = modelOptions['stable-diffusion'] let vaeOptions = modelOptions['vae'] diff --git a/ui/server.py b/ui/server.py index 154fc6cd..07464242 100644 --- a/ui/server.py +++ b/ui/server.py @@ -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)]