scripts/lib.py

import datetime
import gzip
import hashlib
import os
import glob
import textwrap
import gradio as gr
import numpy as np
import requests
import json
import base64
import webbrowser
from io import BytesIO
from PIL import Image
from PIL import PngImagePlugin
import random
import argparse

# Language
LANG_EN = {
    "character1": "Character list 1",
    "character2": "Character list 2",
    "character3": "Character list 3",
    "action": "Action list",
    "original_character": "Original Character",    
}

LANG = LANG_EN

# JavaScript
JAVA_SCRIPT = """
function refresh() {
    const url = new URL(window.location);

    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""

# CSS
CSS_SCRIPT = """
#custom_prompt_text textarea {color: darkorange}
#positive_prompt_text textarea {color: greenyellow}
#negative_prompt_text textarea {color: red}
#ai_prompt_text textarea {color: hotpink}
#prompt_ban_text textarea {color: Khaki}
"""

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 = {}
original_character_list = ''
original_character_dict = {}
wai_image_dict = {}

wai_illustrious_character_select_files = [
    {'name': 'original_character', 'file_path': os.path.join(json_folder, 'original_character.json'), 'url': 'local'},
    {'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url': 'local'},
    {'name': 'wai_image', 'file_path': os.path.join(json_folder, 'wai_character_thumbs.json'), 'url': 'local'},
]
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 load_jsons():
    global character_list
    global character_dict
    global original_character_dict
    global original_character_list
    global wai_image_dict
    
    # download file
    for item in wai_illustrious_character_select_files:
        name = item['name']
        file_path = item['file_path']
        url = item['url']        
            
        with open(file_path, 'r', encoding='utf-8') as file:
            if 'original_character' == name:
                original_character_dict.update(json.load(file))                
            elif '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")
    
    original_character_list = list(original_character_dict.keys())    
    original_character_list.insert(0, "none")

def illustrious_character_select_ex(character = 'random', optimise_tags = True, random_action_seed = 1):
    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 ENGLISH_CHARACTER_NAME:
        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 original_character_select_ex(character = 'random', random_action_seed = 1):
    chara = ''
    rnd_character = ''
    
    if 'none' == character:
        return '', ''
    
    if 'random' == character:
        index = random_action_seed % len(original_character_list)
        rnd_character = original_character_list[index]
        if 'random' == rnd_character:
            rnd_character = original_character_list[index+2]
        elif 'none' == rnd_character:
            rnd_character = original_character_list[index+1]
    else:
        rnd_character = character
    chara = original_character_dict[rnd_character]                                    
    
    opt_chara = chara
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'   
            
    return rnd_character, opt_chara

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

def init():
    global ENGLISH_CHARACTER_NAME
    global LANG
        
    ENGLISH_CHARACTER_NAME = True
    if ENGLISH_CHARACTER_NAME:
        print(f'[{CAT}]:Use tags as Character Name')
        LANG = LANG_EN
        
    load_jsons()    
    print(f'[{CAT}]:Starting...')
    
    return character_list, original_character_list, LANG