| import gradio as gr |
| import random |
| import spaces |
| import torch |
| from transformers import pipeline |
| import csv |
| import os |
| import uuid |
| import datetime |
|
|
| DATA_DIR = os.path.join(os.getcwd(), "session_data") |
| os.makedirs(DATA_DIR, exist_ok=True) |
| csv_path = None |
|
|
|
|
| pipe = pipeline("text-generation", model="mecoffey/NPC_brain") |
|
|
| DESC_SYSTEM_PROMPT = ( |
| "Write a description of a fantasy character based on the user prompt" |
| "response Must include: fantasy Race (Human, Dragonborn, Dwarf, Eladrin, Elf, Fairy, Firbolg, Gnome, Goblin, Half-elf, Half-orc, Halfling, Human, Orc, Tiefling, other as specified by prompt), height, build, eye color, hair color, clothing." |
| "response may include: noticeable features such as scars, tattoos, jewelry, or expression" |
| "Response must not include any extraneous details." |
| "response must be no longer than 150 words" |
| ) |
| BACK_SYSTEM_PROMPT = ( |
| "format the backstory as a timeline in a bulleted list with 1-3 sentences per point." |
| "Keep each point short and fitting for a background character in a fantasy world." |
| "include things like occupations they have had, family ties, livelihood, or mannerisms." |
| "response must include one line about what the character is thinking or doing right now." |
| "Do not include any descriptions of physical features." |
| ) |
|
|
|
|
| DEFAULT_SESSION_STATE = { |
| "csv_path": None, |
| "strength": "", |
| "dexterity": "", |
| "constitution": "", |
| "intelligence": "", |
| "wisdom": "", |
| "charisma": "", |
| "ac": "", |
| "hp": "", |
| "speed": "", |
| "charm": "", |
| "pref_traits": "", |
| "ick_traits": "", |
| } |
|
|
| def new_session_state(): |
| return dict(DEFAULT_SESSION_STATE) |
|
|
| def stats_gen(state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| state.update({ |
| "strength": random.randint(2, 12), |
| "dexterity": random.randint(2, 12), |
| "constitution": random.randint(2, 12), |
| "intelligence": random.randint(2, 12), |
| "wisdom": random.randint(2, 12), |
| "charisma": random.randint(2, 12), |
| }) |
| return state["strength"], state["dexterity"], state["constitution"], state["intelligence"], state["wisdom"], state["charisma"], state |
|
|
| def stats_gen_2(state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| state.update({ |
| "ac": random.randint(2, 14), |
| "hp": state.get("constitution", 0) + random.randint(1, 6), |
| "speed": "25 ft", |
| }) |
| return state["ac"], state["hp"], state["speed"], state |
|
|
| def traits_picker(filepath): |
| with open(filepath, 'r') as f: |
| items = [line.strip() for line in f if line.strip()] |
| count = random.randint(1, min(3, len(items))) |
| selected = random.sample(items,count) |
| return ', '.join(selected) |
|
|
| def rand_race_gen(filepath): |
| with open(filepath, 'r') as f: |
| items = [line.strip() for line in f if line.strip()] |
| return random.choice(items) |
|
|
| def bard_stats_gen(state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| state.update({ |
| "charm": random.randint(1, 30), |
| "pref_traits": traits_picker('traits.txt'), |
| "ick_traits": traits_picker('traits.txt'), |
| }) |
| return state["charm"], state["pref_traits"], state["ick_traits"], state |
|
|
|
|
| def generate_character(prompt, state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| state.update({ |
| "strength": random.randint(2, 12), |
| "dexterity": random.randint(2, 12), |
| "constitution": random.randint(2, 12), |
| "intelligence": random.randint(2, 12), |
| "wisdom": random.randint(2, 12), |
| "charisma": random.randint(2, 12), |
| }) |
| state.update({ |
| "ac": random.randint(2, 14), |
| "hp": state.get("constitution", 0) + random.randint(1, 6), |
| "speed": "25 ft", |
| }) |
| state.update({ |
| "charm": random.randint(1, 30), |
| "pref_traits": traits_picker('traits.txt'), |
| "ick_traits": traits_picker('traits.txt'), |
| }) |
| if not state.get("csv_path"): |
| state["csv_path"] = make_session_csv() |
| description, backstory = char_description(prompt) |
| append_to_csv( |
| state["csv_path"], |
| prompt, |
| description, |
| backstory, |
| state["strength"], |
| state["dexterity"], |
| state["constitution"], |
| state["intelligence"], |
| state["wisdom"], |
| state["charisma"], |
| state["ac"], |
| state["hp"], |
| state["speed"], |
| state["charm"], |
| state["pref_traits"], |
| state["ick_traits"], |
| ) |
| return ( |
| state["strength"], |
| state["dexterity"], |
| state["constitution"], |
| state["intelligence"], |
| state["wisdom"], |
| state["charisma"], |
| state["ac"], |
| state["hp"], |
| state["speed"], |
| state["charm"], |
| state["pref_traits"], |
| state["ick_traits"], |
| description, |
| backstory, |
| state, |
| ) |
|
|
|
|
| def generate_random_character(state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| state.update({ |
| "strength": random.randint(2, 12), |
| "dexterity": random.randint(2, 12), |
| "constitution": random.randint(2, 12), |
| "intelligence": random.randint(2, 12), |
| "wisdom": random.randint(2, 12), |
| "charisma": random.randint(2, 12), |
| }) |
| state.update({ |
| "ac": random.randint(2, 14), |
| "hp": state.get("constitution", 0) + random.randint(1, 6), |
| "speed": "25 ft", |
| }) |
| state.update({ |
| "charm": random.randint(1, 30), |
| "pref_traits": traits_picker('traits.txt'), |
| "ick_traits": traits_picker('traits.txt'), |
| }) |
| if not state.get("csv_path"): |
| state["csv_path"] = make_session_csv() |
| rand_race = rand_race_gen('core_races.txt') |
| description, backstory = rand_char_description(rand_race) |
| append_to_csv( |
| state["csv_path"], |
| rand_race, |
| description, |
| backstory, |
| state["strength"], |
| state["dexterity"], |
| state["constitution"], |
| state["intelligence"], |
| state["wisdom"], |
| state["charisma"], |
| state["ac"], |
| state["hp"], |
| state["speed"], |
| state["charm"], |
| state["pref_traits"], |
| state["ick_traits"], |
| ) |
| return ( |
| state["strength"], |
| state["dexterity"], |
| state["constitution"], |
| state["intelligence"], |
| state["wisdom"], |
| state["charisma"], |
| state["ac"], |
| state["hp"], |
| state["speed"], |
| state["charm"], |
| state["pref_traits"], |
| state["ick_traits"], |
| description, |
| backstory, |
| state, |
| ) |
|
|
|
|
| def make_session_csv(): |
| session_id = uuid.uuid4().hex |
| return os.path.join(DATA_DIR, f"story_session_{session_id}.csv") |
|
|
|
|
| def append_to_csv(csv_path, prompt, description, backstory, strength, dexterity, constitution, intelligence, wisdom, charisma, ac, hp, speed, charm, pref_traits, ick_traits): |
| is_new = not os.path.exists(csv_path) |
| with open(csv_path, "a", newline="", encoding="utf-8") as f: |
| writer = csv.writer(f) |
| if is_new: |
| writer.writerow(["timestamp", "prompt", "description", "backstory", "strength", "dexterity", "constitution", "intelligence", "wisdom", "charisma", "ac", "hp", "speed", "charm", "pref_traits", "ick_traits"]) |
| writer.writerow([ |
| datetime.datetime.utcnow().isoformat(), |
| prompt, |
| description, |
| backstory, |
| strength, |
| dexterity, |
| constitution, |
| intelligence, |
| wisdom, |
| charisma, |
| ac, |
| hp, |
| speed, |
| charm, |
| pref_traits, |
| ick_traits, |
| ]) |
|
|
| def char_description(prompt): |
|
|
| prompt1 = [{"role": "user", "content": prompt}] |
| output1 = pipe(prompt1, return_full_text=False, max_new_tokens=256) |
| output1_text = output1[0]["generated_text"].strip() |
| combine = str(prompt) + "\n" + str(output1_text) + "\nBackstory:" |
| prompt2 = [ |
| {"role": "user", "content": combine}, |
| ] |
| output2 = pipe(prompt2, return_full_text=False, max_new_tokens=256) |
| return output1_text, output2[0]["generated_text"].strip() |
|
|
| def rand_char_description(prompt): |
| prompt1 = [{"role": "user", "content": prompt}] |
| output1 = pipe(prompt1, return_full_text=False, max_new_tokens=256) |
| output1_text = output1[0]["generated_text"].strip() |
| combine = str(prompt) + "\n" + str(output1_text) + "\nBackstory:" |
| prompt2 = [ |
| {"role": "user", "content": combine}, |
| ] |
| output2 = pipe(prompt2, return_full_text=False, max_new_tokens=256) |
| return output1_text, output2[0]["generated_text"].strip() |
|
|
|
|
| def char_description_csv(prompt, state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| if not state.get("csv_path"): |
| state["csv_path"] = make_session_csv() |
| description, backstory = char_description(prompt) |
| append_to_csv( |
| state["csv_path"], |
| prompt, |
| description, |
| backstory, |
| state.get("strength", ""), |
| state.get("dexterity", ""), |
| state.get("constitution", ""), |
| state.get("intelligence", ""), |
| state.get("wisdom", ""), |
| state.get("charisma", ""), |
| state.get("ac", ""), |
| state.get("hp", ""), |
| state.get("speed", ""), |
| state.get("charm", ""), |
| state.get("pref_traits", ""), |
| state.get("ick_traits", ""), |
| ) |
| return description, backstory, state |
|
|
|
|
| def rand_char_description_csv(state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| if not state.get("csv_path"): |
| state["csv_path"] = make_session_csv() |
| rand_race = rand_race_gen('core_races.txt') |
| description, backstory = rand_char_description(rand_race) |
| append_to_csv( |
| state["csv_path"], |
| rand_race, |
| description, |
| backstory, |
| state.get("strength", ""), |
| state.get("dexterity", ""), |
| state.get("constitution", ""), |
| state.get("intelligence", ""), |
| state.get("wisdom", ""), |
| state.get("charisma", ""), |
| state.get("ac", ""), |
| state.get("hp", ""), |
| state.get("speed", ""), |
| state.get("charm", ""), |
| state.get("pref_traits", ""), |
| state.get("ick_traits", ""), |
| ) |
| return description, backstory, state |
|
|
|
|
| def get_csv_file(state): |
| state = dict(state or DEFAULT_SESSION_STATE) |
| if not state.get("csv_path"): |
| state["csv_path"] = make_session_csv() |
| return state["csv_path"], state |
|
|
|
|
| |
| with gr.Blocks() as demo: |
| title = gr.HTML("<center><h1>NPC Generator</h1></center>") |
| description = gr.HTML("<center>For tabletop GMs when players ask for a description of random people in the tavern.</center>") |
| tags = gr.HTML("<center>| Backyard AI | Open BMB | Tiny Titan | Fine-tuned | Modal |</center>") |
| demo_video = gr.HTML("<center>See a demo here!: https://youtu.be/b7M0ejRoxQU</center>") |
|
|
| with gr.Row(equal_height = True): |
| with gr.Column(scale = 2): |
| text_in = gr.Textbox(label="What npc do you need?", placeholder = "An Orc Barkeep") |
| with gr.Column(scale = 0): |
| btn = gr.Button("Generate") |
| rand = gr.Button("Random") |
|
|
| with gr.Row(): |
|
|
| session_state = gr.State(value=new_session_state()) |
|
|
| with gr.Column(scale = 0): |
| stre_box = gr.Textbox(value = "", label = "Strength", interactive = False) |
| dex_box = gr.Textbox(value = "", label = "Dexterity", interactive = False) |
| con_box = gr.Textbox(value = "", label = "Constitution", interactive = False) |
| intl_box = gr.Textbox(value = "", label = "Intelligence", interactive = False) |
| wis_box = gr.Textbox(value = "", label = "Wisdom", interactive = False) |
| char_box = gr.Textbox(value = "", label = "Charisma", interactive = False) |
|
|
| with gr.Column(scale = 2): |
| text_out1 = gr.Textbox(label = "Description") |
| text_out2 = gr.Textbox(label = "Backstory") |
|
|
| with gr.Column(scale = 0): |
| ac_box = gr.Textbox(value = "", label = "AC", interactive = False) |
| hp_box = gr.Textbox(value = "", label = "HP", interactive = False) |
| speed_box = gr.Textbox(value = "", label = "Speed", interactive = False) |
| bard_stats = gr.HTML("~~Bard Stats~~") |
| with gr.Accordion(open = False): |
| dcs = gr.Textbox(value = "", label = "DC to Charm") |
| prefs = gr.Textbox(value = "", label = "Advantage if:") |
| icks = gr.Textbox(value = "", label = "Disdvantage if:") |
|
|
| btn.click( |
| fn=generate_character, |
| inputs=[text_in, session_state], |
| outputs=[ |
| stre_box, |
| dex_box, |
| con_box, |
| intl_box, |
| wis_box, |
| char_box, |
| ac_box, |
| hp_box, |
| speed_box, |
| dcs, |
| prefs, |
| icks, |
| text_out1, |
| text_out2, |
| session_state, |
| ], |
| ) |
| rand.click( |
| fn=generate_random_character, |
| inputs=session_state, |
| outputs=[ |
| stre_box, |
| dex_box, |
| con_box, |
| intl_box, |
| wis_box, |
| char_box, |
| ac_box, |
| hp_box, |
| speed_box, |
| dcs, |
| prefs, |
| icks, |
| text_out1, |
| text_out2, |
| session_state, |
| ], |
| ) |
|
|
| with gr.Row(): |
| download_btn = gr.Button("Download CSV") |
| with gr.Row(): |
| download_file = gr.File(label="Generated CSV", interactive=False) |
| download_btn.click(fn=get_csv_file, inputs=session_state, outputs=[download_file, session_state]) |
|
|
|
|
| demo.launch(theme = gr.themes.Monochrome()) |
