utils.py

import re
import base64
import html
from io import BytesIO


def image_to_base64(image):
    # Convert the image to bytes
    buffered = BytesIO()
    image.save(buffered, format="PNG")  # You can change format to PNG if needed

    # Encode the bytes to base64
    img_str = base64.b64encode(buffered.getvalue())

    return img_str.decode('utf-8')  # Convert bytes to string


def remove_id_and_ext(text):
    text = re.sub(r'\[.*\]$', '', text)
    extension = text[-12:].strip()
    if extension == "safetensors":
        text = text[:-13]
    elif extension == "ckpt":
        text = text[:-4]
    return text


def extract_data(text):
    results = {}
    patterns = {
        'prompt': r'(.*)',
        'negative_prompt': r'Negative prompt: (.*)',
        'steps': r'Steps: (\d+),',
        'seed': r'Seed: (\d+),',
        'sampler': r'Sampler:\s*([^\s,]+(?:\s+[^\s,]+)*)',
        'model': r'Model:\s*([^\s,]+)',
        'cfg_scale': r'CFG scale:\s*([\d\.]+)',
        'size': r'Size:\s*([0-9]+x[0-9]+)'
    }
    for key in ['prompt', 'negative_prompt', 'steps', 'seed', 'sampler', 'model', 'cfg_scale', 'size']:
        match = re.search(patterns[key], text)
        if match:
            results[key] = match.group(1)
        else:
            results[key] = None
    if results['size'] is not None:
        w, h = results['size'].split("x")
        results['w'] = w
        results['h'] = h
    else:
        results['w'] = None
        results['h'] = None
    return results


def place_lora(current_prompt, lora_name):
    pattern = r"<lora:" + lora_name + r":.*?>"

    if re.search(pattern, current_prompt):
        yield re.sub(pattern, "", current_prompt)
    else:
        yield current_prompt + " <lora:" + lora_name + ":1> "


def plaintext_to_html(text, classname=None):
    content = "<br>\n".join(html.escape(x) for x in text.split('\n'))

    return f"<p class='{classname}'>{content}</p>" if classname else f"<p>{content}</p>"


def get_exif_data(image):
    items = image.info

    info = ''
    for key, text in items.items():
        info += f"""
        <div>
        <p><b>{plaintext_to_html(str(key))}</b></p>
        <p>{plaintext_to_html(str(text))}</p>
        </div>
        """.strip() + "\n"

    if len(info) == 0:
        message = "Nothing found in the image."
        info = f"<div><p>{message}<p></div>"

    return info