animatediff/utils/wild_card.py

import glob
import os
import random
import re

wild_card_regex = r'(\A|\W)__([\w-]+)__(\W|\Z)'


def create_wild_card_map(wild_card_dir):
    result = {}
    if os.path.isdir(wild_card_dir):
        txt_list = glob.glob( os.path.join(wild_card_dir ,"**/*.txt"), recursive=True)
        for txt in txt_list:
            basename_without_ext = os.path.splitext(os.path.basename(txt))[0]
            with open(txt, encoding='utf-8') as f:
                try:
                    result[basename_without_ext] = [s.rstrip() for s in f.readlines()]
                except Exception as e:
                    print(e)
                    print("can not read ", txt)
    return result

def replace_wild_card_token(match_obj, wild_card_map):
    m1 = match_obj.group(1)
    m3 = match_obj.group(3)

    dict_name = match_obj.group(2)

    if dict_name in wild_card_map:
        token_list = wild_card_map[dict_name]
        token = token_list[random.randint(0,len(token_list)-1)]
        return m1+token+m3
    else:
        return match_obj.group(0)

def replace_wild_card(prompt, wild_card_dir):
    wild_card_map = create_wild_card_map(wild_card_dir)
    prompt = re.sub(wild_card_regex, lambda x: replace_wild_card_token(x, wild_card_map ), prompt)
    return prompt