Spaces:
Starting
Starting
| 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 | |
| ##UI is fine as-is | |
| 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()) |