scripts/lib.py

import gzip
import hashlib
import os
import requests
import json
import base64

from io import BytesIO
from PIL import Image
from tag_autocomplete import PromptManager

# Language
LANG_EN = {
    "character1": "Character list EN",
    "character2": "Character list EN",
    "character3": "Character list CN",
    "action": "Action list",
    "original_character": "Original Character",    
}

LANG = LANG_EN

TITLE = "WAI Character Select Preview"
CAT = "WAI_Character_Select"
ENGLISH_CHARACTER_NAME = True

current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
json_folder = os.path.join(parent_dir, 'json')

character_list = ''
character_dict = {}
wai_image_dict = {}
character_list_cn = ''
PROMPT_MANAGER = None

wai_illustrious_character_select_files = [
    {'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url':'https://raw.githubusercontent.com/mirabarukaso/character_select_stand_alone_app/refs/heads/main/json/wai_characters.csv'},
    {'name': 'wai_image', 'file_path': os.path.join(json_folder, 'wai_character_thumbs.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/wai_character_thumbs.json'},
    {'name': 'e621_sfw', 'file_path': os.path.join(json_folder, 'e621_sfw.csv'), 'url': 'https://raw.githubusercontent.com/DominikDoom/a1111-sd-webui-tagcomplete/refs/heads/main/tags/e621_sfw.csv'},
]
def get_md5_hash(input_str):
    md5_hash = hashlib.md5()
    md5_hash.update(input_str.encode('utf-8'))
    return md5_hash.hexdigest()

def base64_to_image(base64_data):
    compressed_data = base64.b64decode(base64_data)
    webp_data = gzip.decompress(compressed_data)
    image = Image.open(BytesIO(webp_data))  
    return image

def download_file(url, file_path):   
    response = requests.get(url)
    response.raise_for_status() 
    print(f'[{CAT}]:Downloading... {url}')
    with open(file_path, 'wb') as file:
        file.write(response.content)        

def load_text_file(file_path):
    raw_text = ''
    if os.path.exists(file_path):
        print(f'[{CAT}]:Loading {file_path}')
        with open(file_path, 'r', encoding='utf-8') as js_file:
            raw_text = js_file.read()
    else:
        print(f"[{CAT}] ERROR: {file_path} file missing!!!")
        
    return raw_text
            
def load_jsons():
    global character_list
    global character_dict
    global wai_image_dict
    global character_list_cn
    global PROMPT_MANAGER
    
    # download file
    for item in wai_illustrious_character_select_files:
        name = item['name']
        file_path = item['file_path']
        url = item['url']        
        
        if not os.path.exists(file_path):
            download_file(url, file_path)
            
        if 'e621_sfw' == name:
            PROMPT_MANAGER = PromptManager(file_path)
        else:
            with open(file_path, 'r', encoding='utf-8') as file:        
                if 'wai_character' == name:
                    lines = file.readlines()
                    for line in lines:
                        key, value = line.split(',')
                        character_dict[key.strip()]=value.strip()
                elif 'wai_image' == name:
                    wai_image_dict = json.load(file)                
                                        
    # Create list        
    character_list = list(character_dict.values())   
    character_list.insert(0, "none")
    character_list_cn = list(character_dict.keys())   
    character_list_cn.insert(0, "none")

def illustrious_character_select_ex(character = 'random', optimise_tags = True, random_action_seed = 1, use_cn=False):
    chara = ''
    rnd_character = ''
    
    if 'none' == character:
        return '', '', None
    
    if 'random' == character:
        index = random_action_seed % len(character_list)
        rnd_character = character_list[index]
        if 'random' == rnd_character:
            rnd_character = character_list[index+2]
        elif 'none' == rnd_character:
            rnd_character = character_list[index+1]
    else:
        rnd_character = character
    if not use_cn:
        chara = rnd_character
    else:
        chara = character_dict[rnd_character]
        
    md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))                
    thumb_image = Image.new('RGB', (128, 128), (128, 128, 128))
    if wai_image_dict.keys().__contains__(md5_chara):
        thumb_image = base64_to_image(wai_image_dict.get(md5_chara))
    
    opt_chara = chara
    if optimise_tags:
        opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')  
        #print(f'{CAT}:Optimise Tags:[{chara}]->[{opt_chara}]')
    
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'   
            
    return rnd_character, opt_chara, thumb_image

def refresh_character_thumb_image(character1, character2, character3):
    thumb_image = []
    if 'none' != character1 and 'random' != character1:
        _, _, thumb_image1 = illustrious_character_select_ex(character = character1)        
        thumb_image.append(thumb_image1)
    
    if 'none' != character2 and 'random' != character2:
        _, _, thumb_image2 = illustrious_character_select_ex(character = character2)        
        thumb_image.append(thumb_image2)
    if 'none' != character3 and 'random' != character3:
        _, _, thumb_image3 = illustrious_character_select_ex(character = character3, use_cn=True)        
        thumb_image.append(thumb_image3)
    return thumb_image

def get_prompt_manager():
    return PROMPT_MANAGER

def init():            
    lib_js_path = os.path.join(current_dir, 'lib.js')
    lib_css_path = os.path.join(current_dir, 'lib.css')            
    
    load_jsons()   
    js_script = load_text_file(lib_js_path)
    css_script = load_text_file(lib_css_path) 
    print(f'[{CAT}]:Starting...')
    
    return character_list, character_list_cn, LANG, js_script, css_script