benchmark.py

import time
import logging
from modules import shared
from modules.processing import StableDiffusionProcessingTxt2Img, Processed, process_images


log = logging.getLogger('sd')

default_args = {
    # 'sd_model': None,
    'prompt': 'postapocalyptic steampunk city, exploration, cinematic, realistic, hyper detailed, photorealistic maximum detail, volumetric light, (((focus))), wide-angle, (((brightly lit))), (((vegetation))), lightning, vines, destruction, devastation, wartorn, ruins',
    'sampler_name': 'Euler a',
    'batch_size': 1,
    'n_iter': 1,
    'steps': 10,
    'cfg_scale': 7.0,
    'width': 512,
    'height': 512,
    'do_not_save_samples': True,
    'do_not_save_grid': True,
    'negative_prompt': '(((blurry))), ((foggy)), (((dark))), ((monochrome)), sun, (((depth of field)))',
    'do_not_reload_embeddings': True
}
args = {}


def run_benchmark(batch: int, steps: str, width: int, height: int):
    global args # pylint: disable=global-statement
    shared.state.begin('Benchmark')
    args = default_args.copy()
    # if args['sd_model'] is None:
    #     args['sd_model'] = shared.sd_model
    args['batch_size'] = batch
    if steps == 'turbo':
        args['steps'] = 10
    elif steps == 'long':
        args['steps'] = 50
    else:
        args['steps'] = 20
    if width and width > 0:
        args['width'] = int(width)
    if height and height > 0:
        args['height'] = int(height)
    mp = 0
    p = StableDiffusionProcessingTxt2Img(**args)
    t0 = time.time()
    try:
        processed: Processed = process_images(p)
        if processed is None or processed.images is None:
            log.error(f'SD-System-Info benchmark error: {batch} no results')
            return 'error'
        if len(processed.images) != batch:
            log.error(f'SD-System-Info benchmark error: {batch} results mismatch')
            return 'error'
        for image in processed.images:
            if image.width * image.height < 65536:
                log.error(f'SD-System-Info benchmark error: {batch} image too small')
                return 'error'
            mp += image.width * image.height
    except Exception as e:
        log.error(f'SD-System-Info benchmark error: {batch} {e}')
        from modules import errors
        errors.display(e, 'Benchmark')
        return 'error'
    t1 = time.time()
    shared.state.end()
    its = args['steps'] * batch / (t1 - t0)
    mp = mp / 1024 / 1024
    mps = mp / (t1 - t0)
    log.debug(f'SD-System-Info benchmark: batch={batch} time={t1-t0:.2f} steps={args["steps"]} its={its:.2f} mps={mps:.2f}')
    if its > 300:
        log.error(f'SD-System-Info benchmark: its={its:.2f} too high')
        return 'error'
    return round(its, 2)


class LogFilter(logging.Filter):
    import socket
    hostname = socket.gethostname()
    def filter(self, record):
        record.hostname = LogFilter.hostname
        return True


def submit_benchmark(data, username):
    from logging.handlers import SysLogHandler
    from hashlib import sha256

    syslog = SysLogHandler(address=('logs3.papertrailapp.com', 32554))
    syslog.addFilter(LogFilter())
    formatter = logging.Formatter(f'%(asctime)s %(hostname)s SDBENCHMARK: {username} %(message)s', datefmt='%b %d %H:%M:%S')
    syslog.setFormatter(formatter)
    remote = logging.getLogger('SDBENCHMARK')
    remote.setLevel(logging.INFO)
    for h in remote.handlers: # remove local handlers
        remote.removeHandler(h)
    remote.addHandler(syslog)
    for line in data:
        message = '|'.join(line).replace('  ', ' ').replace('"', '').strip()
        hash256 = sha256(message.encode('utf-8')).hexdigest()[:6]
        message = message + '|' + hash256
        remote.info(message)