Spaces:

flagrantia
/

character_select_saa

Running

mirabarukaso

Update lib.py

f076f06 3 months ago

11.5 kB

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

from io import BytesIO
from PIL import Image

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

LANG = LANG_EN

TITLE = "WAI Character Select Preview"
CAT =       "CSP"
CAT_IL =    "IL "
CAT_ANIMA = "ANI"
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')

THUMB_PLACEHOLDER_SIZE = (307, 460)

character_list = ''
character_dict = {}
wai_image_dict = {}
character_list_cn = ''
il_tag_assist = {}

anima_character_list = ''
anima_character_dict = {}
wai_anima_image_dict = {}
anima_character_list_cn = ''
anima_tag_assist = {}

PROMPT_MANAGER = None

total_counter = 0
total_counter_il = 0
total_counter_anima = 0

wai_illustrious_character_select_files = [
    {'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url':'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/wai_characters_v160.csv?download=true'},
    {'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_v160.json?download=true'},
    {'name': 'wai_tag_assist', 'file_path': os.path.join(json_folder, 'wai_tag_assist.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiIllustriousSDXL_v160_tag_assist.json?download=true'},
    
    {'name': 'wai_anima_character', 'file_path': os.path.join(json_folder, 'wai_anima_characters.csv'), 'url':'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_characters.csv?download=true'},
    {'name': 'wai_anima_image', 'file_path': os.path.join(json_folder, 'wai_anima_character_thumbs.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_thumbs.json?download=true'},
    {'name': 'wai_anima_tag_assist', 'file_path': os.path.join(json_folder, 'wai_anima_tag_assist.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_tag_assist.json?download=true'},
    
    # {'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 make_thumb_placeholder(size=THUMB_PLACEHOLDER_SIZE):
    return Image.new('RGB', size, (39, 39, 42))

def ensure_thumb_size(image, size=THUMB_PLACEHOLDER_SIZE):
    if image.size != size:
        return image.resize(size, Image.LANCZOS)
    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 anima_character_list
    global anima_character_dict
    global wai_anima_image_dict
    global anima_character_list_cn
    
    global il_tag_assist
    global anima_tag_assist
    
    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)
            
        with open(file_path, 'r', encoding='utf-8') as file:        
            if 'wai_character' == name:
                print(f"[{CAT}] Loading wai_character...")
                lines = file.readlines()
                for line in lines:
                    key, value = line.split(',')
                    character_dict[key.strip()]=value.strip()
            elif 'wai_image' == name:
                print(f"[{CAT}] Loading wai_image...")
                wai_image_dict = json.load(file)
            elif 'wai_anima_character' == name:
                print(f"[{CAT}] Loading wai_anima_character...")
                lines = file.readlines()
                for line in lines:
                    key, value = line.split(',')
                    anima_character_dict[key.strip()]=value.strip()
            elif 'wai_anima_image' == name:
                print(f"[{CAT}] Loading wai_anima_image...")
                wai_anima_image_dict = json.load(file)
            elif 'wai_tag_assist' == name:
                print(f"[{CAT}] Loading wai_tag_assist...")
                il_tag_assist = json.load(file)
            elif 'wai_anima_tag_assist' == name:
                print(f"[{CAT}] Loading wai_anima_tag_assist...")
                anima_tag_assist = json.load(file)
                
                                        
    # Create IL 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")
    
    # Create Anima List
    anima_character_list = list(anima_character_dict.values())   
    anima_character_list.insert(0, "none")
    anima_character_list_cn = list(anima_character_dict.keys())   
    anima_character_list_cn.insert(0, "none")


def illustrious_character_select_ex(character = 'none', optimise_tags = True, use_cn=False):
    global total_counter
    global total_counter_il
    
    chara = ''
    tag_assist = ''
    
    if 'none' == character:
        return '', '', None
            
    if not use_cn:
        chara = character
    else:
        chara = character_dict[character]
        
    md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))                
    thumb_image = make_thumb_placeholder()
    if wai_image_dict.keys().__contains__(md5_chara):
        thumb_image = ensure_thumb_size(base64_to_image(wai_image_dict.get(md5_chara)))
    
    opt_chara = chara
    if optimise_tags:
        opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')  
        total_counter = total_counter + 1
        total_counter_il = total_counter_il + 1
        print(f'{CAT_IL}:{total_counter_il}/{total_counter}:[{chara}]->[{opt_chara}]')
    
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'  
    
    if il_tag_assist.keys().__contains__(chara):
        tag_assist = il_tag_assist[chara]
            
    return character, opt_chara, thumb_image, tag_assist

def anima_character_select_ex(character = 'none', optimise_tags = True, use_cn=False):
    global total_counter
    global total_counter_anima
        
    chara = ''
    tag_assist = ''
    
    if 'none' == character:
        return '', '', None
            
    if not use_cn:
        chara = character
    else:
        chara = anima_character_dict[character]
        
    md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))                
    thumb_image = make_thumb_placeholder()
    if wai_anima_image_dict.keys().__contains__(md5_chara):
        thumb_image = ensure_thumb_size(base64_to_image(wai_anima_image_dict.get(md5_chara)))
    
    opt_chara = chara
    if optimise_tags:
        opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')  
        total_counter = total_counter + 1
        total_counter_anima = total_counter_anima + 1
        print(f'{CAT_ANIMA}:{total_counter_anima}/{total_counter}:[{chara}]->[{opt_chara}]')
    
    if not opt_chara.endswith(','):
        opt_chara = f'{opt_chara},'   
        
    if anima_tag_assist.keys().__contains__(chara):
        tag_assist = anima_tag_assist[chara]
            
    return character, opt_chara, thumb_image, tag_assist

def create_prompt_info(rnd_character1, opt_chara1, tag_assist1,
                        rnd_character2, opt_chara2, tag_assist2,
                        rnd_character3, opt_chara3, tag_assist3,
                        rnd_character4, opt_chara4, tag_assist4):
    info = ''    
    if '' != opt_chara1:
        info += f'Anima EN: {rnd_character1}\nPrompt: {opt_chara1}\nTag Assist: {tag_assist1}\n\n'
        
    if '' != opt_chara2:
        info += f'ILLUS EN: {rnd_character2}\nPrompt: {opt_chara2}\nTag Assist: {tag_assist2}\n\n'
    
    if '' != opt_chara3:
        info += f'Anima CN: {rnd_character3}\n提示词: {opt_chara3}\n标签辅助: {tag_assist3}\n\n'
    
    if '' != opt_chara4:
        info += f'ILLUS CN: {rnd_character4}\n提示词: {opt_chara4}\n标签辅助: {tag_assist4}'

    prompt = f'{opt_chara1}{opt_chara2}{opt_chara3}{opt_chara4}'

    return prompt, info

def refresh_character_thumb_image(*args):
    thumb_image = [make_thumb_placeholder() for _ in range(4)]
    rnd_character = [''] *4
    opt_chara = [''] *4
    tag_assist = [''] * 4
    
    num_inputs = len(args)
    if num_inputs == 4:
        character1, character2, character3, character4 = args
    else:
        print(f'[refresh_character_thumb_image] Ignore old/incorrect inputs, num_inputs={num_inputs}')    
        return thumb_image, 'Refresh your page'
    
    if 'none' != character1:
        rnd_character[0], opt_chara[0], thumb_image1, tag_assist[0] = anima_character_select_ex(character = character1)        
        thumb_image[0] = thumb_image1
    if 'none' != character2:
        rnd_character[1], opt_chara[1], thumb_image2, tag_assist[1] = illustrious_character_select_ex(character = character2)
        thumb_image[1] = thumb_image2
    if 'none' != character3:
        rnd_character[2], opt_chara[2], thumb_image3, tag_assist[2] = anima_character_select_ex(character = character3, use_cn=True)
        thumb_image[2] = thumb_image3
    if 'none' != character4:
        rnd_character[3], opt_chara[3], thumb_image4, tag_assist[3] = illustrious_character_select_ex(character = character4, use_cn=True)        
        thumb_image[3] = thumb_image4
    
    _, character_info = create_prompt_info(
                        rnd_character[0], opt_chara[0], tag_assist[0],
                        rnd_character[1], opt_chara[1], tag_assist[1],
                        rnd_character[2], opt_chara[2], tag_assist[2],
                        rnd_character[3], opt_chara[3], tag_assist[3])
    
    return thumb_image, character_info

def get_prompt_manager():
    return PROMPT_MANAGER

def init():    
    load_jsons()   
    print(f'[{CAT}] Starting...')
    
    return character_list, character_list_cn, LANG, anima_character_list, anima_character_list_cn

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

	from io import BytesIO
	from PIL import Image

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

	LANG = LANG_EN

	TITLE = "WAI Character Select Preview"
	CAT = "CSP"
	CAT_IL = "IL "
	CAT_ANIMA = "ANI"
	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')

	THUMB_PLACEHOLDER_SIZE = (307, 460)

	character_list = ''
	character_dict = {}
	wai_image_dict = {}
	character_list_cn = ''
	il_tag_assist = {}

	anima_character_list = ''
	anima_character_dict = {}
	wai_anima_image_dict = {}
	anima_character_list_cn = ''
	anima_tag_assist = {}

	PROMPT_MANAGER = None

	total_counter = 0
	total_counter_il = 0
	total_counter_anima = 0

	wai_illustrious_character_select_files = [
	{'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url':'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/wai_characters_v160.csv?download=true'},
	{'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_v160.json?download=true'},
	{'name': 'wai_tag_assist', 'file_path': os.path.join(json_folder, 'wai_tag_assist.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiIllustriousSDXL_v160_tag_assist.json?download=true'},

	{'name': 'wai_anima_character', 'file_path': os.path.join(json_folder, 'wai_anima_characters.csv'), 'url':'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_characters.csv?download=true'},
	{'name': 'wai_anima_image', 'file_path': os.path.join(json_folder, 'wai_anima_character_thumbs.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_thumbs.json?download=true'},
	{'name': 'wai_anima_tag_assist', 'file_path': os.path.join(json_folder, 'wai_anima_tag_assist.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/waiANIMA_v10Base10_tag_assist.json?download=true'},

	# {'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 make_thumb_placeholder(size=THUMB_PLACEHOLDER_SIZE):
	return Image.new('RGB', size, (39, 39, 42))

	def ensure_thumb_size(image, size=THUMB_PLACEHOLDER_SIZE):
	if image.size != size:
	return image.resize(size, Image.LANCZOS)
	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 anima_character_list
	global anima_character_dict
	global wai_anima_image_dict
	global anima_character_list_cn

	global il_tag_assist
	global anima_tag_assist

	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)

	with open(file_path, 'r', encoding='utf-8') as file:
	if 'wai_character' == name:
	print(f"[{CAT}] Loading wai_character...")
	lines = file.readlines()
	for line in lines:
	key, value = line.split(',')
	character_dict[key.strip()]=value.strip()
	elif 'wai_image' == name:
	print(f"[{CAT}] Loading wai_image...")
	wai_image_dict = json.load(file)
	elif 'wai_anima_character' == name:
	print(f"[{CAT}] Loading wai_anima_character...")
	lines = file.readlines()
	for line in lines:
	key, value = line.split(',')
	anima_character_dict[key.strip()]=value.strip()
	elif 'wai_anima_image' == name:
	print(f"[{CAT}] Loading wai_anima_image...")
	wai_anima_image_dict = json.load(file)
	elif 'wai_tag_assist' == name:
	print(f"[{CAT}] Loading wai_tag_assist...")
	il_tag_assist = json.load(file)
	elif 'wai_anima_tag_assist' == name:
	print(f"[{CAT}] Loading wai_anima_tag_assist...")
	anima_tag_assist = json.load(file)


	# Create IL 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")

	# Create Anima List
	anima_character_list = list(anima_character_dict.values())
	anima_character_list.insert(0, "none")
	anima_character_list_cn = list(anima_character_dict.keys())
	anima_character_list_cn.insert(0, "none")


	def illustrious_character_select_ex(character = 'none', optimise_tags = True, use_cn=False):
	global total_counter
	global total_counter_il

	chara = ''
	tag_assist = ''

	if 'none' == character:
	return '', '', None

	if not use_cn:
	chara = character
	else:
	chara = character_dict[character]

	md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))
	thumb_image = make_thumb_placeholder()
	if wai_image_dict.keys().__contains__(md5_chara):
	thumb_image = ensure_thumb_size(base64_to_image(wai_image_dict.get(md5_chara)))

	opt_chara = chara
	if optimise_tags:
	opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')
	total_counter = total_counter + 1
	total_counter_il = total_counter_il + 1
	print(f'{CAT_IL}:{total_counter_il}/{total_counter}:[{chara}]->[{opt_chara}]')

	if not opt_chara.endswith(','):
	opt_chara = f'{opt_chara},'

	if il_tag_assist.keys().__contains__(chara):
	tag_assist = il_tag_assist[chara]

	return character, opt_chara, thumb_image, tag_assist

	def anima_character_select_ex(character = 'none', optimise_tags = True, use_cn=False):
	global total_counter
	global total_counter_anima

	chara = ''
	tag_assist = ''

	if 'none' == character:
	return '', '', None

	if not use_cn:
	chara = character
	else:
	chara = anima_character_dict[character]

	md5_chara = get_md5_hash(chara.replace('(','\\(').replace(')','\\)'))
	thumb_image = make_thumb_placeholder()
	if wai_anima_image_dict.keys().__contains__(md5_chara):
	thumb_image = ensure_thumb_size(base64_to_image(wai_anima_image_dict.get(md5_chara)))

	opt_chara = chara
	if optimise_tags:
	opt_chara = opt_chara.replace('(', '\\(').replace(')', '\\)')
	total_counter = total_counter + 1
	total_counter_anima = total_counter_anima + 1
	print(f'{CAT_ANIMA}:{total_counter_anima}/{total_counter}:[{chara}]->[{opt_chara}]')

	if not opt_chara.endswith(','):
	opt_chara = f'{opt_chara},'

	if anima_tag_assist.keys().__contains__(chara):
	tag_assist = anima_tag_assist[chara]

	return character, opt_chara, thumb_image, tag_assist

	def create_prompt_info(rnd_character1, opt_chara1, tag_assist1,
	rnd_character2, opt_chara2, tag_assist2,
	rnd_character3, opt_chara3, tag_assist3,
	rnd_character4, opt_chara4, tag_assist4):
	info = ''
	if '' != opt_chara1:
	info += f'Anima EN: {rnd_character1}\nPrompt: {opt_chara1}\nTag Assist: {tag_assist1}\n\n'

	if '' != opt_chara2:
	info += f'ILLUS EN: {rnd_character2}\nPrompt: {opt_chara2}\nTag Assist: {tag_assist2}\n\n'

	if '' != opt_chara3:
	info += f'Anima CN: {rnd_character3}\n提示词: {opt_chara3}\n标签辅助: {tag_assist3}\n\n'

	if '' != opt_chara4:
	info += f'ILLUS CN: {rnd_character4}\n提示词: {opt_chara4}\n标签辅助: {tag_assist4}'

	prompt = f'{opt_chara1}{opt_chara2}{opt_chara3}{opt_chara4}'

	return prompt, info

	def refresh_character_thumb_image(*args):
	thumb_image = [make_thumb_placeholder() for _ in range(4)]
	rnd_character = [''] *4
	opt_chara = [''] *4
	tag_assist = [''] * 4

	num_inputs = len(args)
	if num_inputs == 4:
	character1, character2, character3, character4 = args
	else:
	print(f'[refresh_character_thumb_image] Ignore old/incorrect inputs, num_inputs={num_inputs}')
	return thumb_image, 'Refresh your page'

	if 'none' != character1:
	rnd_character[0], opt_chara[0], thumb_image1, tag_assist[0] = anima_character_select_ex(character = character1)
	thumb_image[0] = thumb_image1
	if 'none' != character2:
	rnd_character[1], opt_chara[1], thumb_image2, tag_assist[1] = illustrious_character_select_ex(character = character2)
	thumb_image[1] = thumb_image2
	if 'none' != character3:
	rnd_character[2], opt_chara[2], thumb_image3, tag_assist[2] = anima_character_select_ex(character = character3, use_cn=True)
	thumb_image[2] = thumb_image3
	if 'none' != character4:
	rnd_character[3], opt_chara[3], thumb_image4, tag_assist[3] = illustrious_character_select_ex(character = character4, use_cn=True)
	thumb_image[3] = thumb_image4

	_, character_info = create_prompt_info(
	rnd_character[0], opt_chara[0], tag_assist[0],
	rnd_character[1], opt_chara[1], tag_assist[1],
	rnd_character[2], opt_chara[2], tag_assist[2],
	rnd_character[3], opt_chara[3], tag_assist[3])

	return thumb_image, character_info

	def get_prompt_manager():
	return PROMPT_MANAGER

	def init():
	load_jsons()
	print(f'[{CAT}] Starting...')

	return character_list, character_list_cn, LANG, anima_character_list, anima_character_list_cn