Parse fooocus log files into DiffusionToolkit compatible text

parse_logs.py

from pathlib import Path
import re

switch_image = '<div id="(.+?)" class="image-container"><hr><table><tr>'
image_tag = '<td class=\'key\'>(.+?)</td><td class=\'value\'>(.+?)</td>'

tag_renames = {
    'Guidance Scale': 'CFG scale',
    'Resolution': 'Size',
    'Base Model': 'Model'
}

def process_log(logfile):
    with logfile.open() as fh:
        lines = fh.readlines()
    current_filename = None
    fh_dicts = {}
    for line in lines:
        match = re.match(switch_image, line)
        if match:
            current_filename = match.group(1)
            fh_dicts[current_filename] = {};
        if '<td class=\'key\'>' in line and '<td class=\'value\'>' in line and current_filename:
            fh_dicts[current_filename].update(dict(re.findall(image_tag, line)))
    return fh_dicts

def main():
    p = Path('.')
    logs = p.glob('Fooocus/outputs/**/log.html')
    for log in logs:
        log_results = process_log(log)
        for file_stub, file_values in log_results.items():
            file_stem, file_suffix = file_stub.rsplit('_', maxsplit = 1)
            image_file = log.with_name(file_stem).with_suffix(f'.{file_suffix}')
            metadata_file = log.with_name(file_stem).with_suffix('.txt')
            if not image_file.exists():
                if metadata_file.exists():
                    print(f'** removing metadata file for removed image {image_file}')
                    metadata_file.unlink()
                continue

            metadata_lines = [
                f'parameters: {file_values.get("Prompt", "")}\n',
                f'Negative prompt: {file_values.get("Negative Prompt", "")}\n',
                f'Steps: {file_values.get("Steps", 60)}'
            ]

            for key, value in file_values.items():
                if key in ['Prompt', 'Negative Prompt', 'Steps']:
                    continue
                key = tag_renames.get(key, key)
                if key == 'Size':
                    value = value[1:-1].replace(', ', 'x')
                metadata_lines[2] += f", {key}: {value}"

            with metadata_file.open('wt') as fh:
                fh.writelines(metadata_lines)

if __name__ == "__main__":
    main()