app.py

import gradio as gr 
from functions import CAT_ARTIST, CAT_CHARACTER, CAT_COPYRIGHT, CAT_GENERAL, CAT_LORE, CAT_META, CAT_SPECIES, PromptBuilder, parse_tag, parse_tags, related_tags 
 
def query_tag(tag: str, category: int): 
 if category == -1: 
 category = None 
 return related_tags(parse_tag(tag), category=category) 
 
def generate_prompt(include: str, focus: str, exclude: str, avoid: str, skip: str, rating: str, general: int, artist: int, species: int, copyright: int, character: int, meta: int) -> str: 
 try: 
 builder = PromptBuilder(skip=list(parse_tags(skip)), min_posts=50, rating=rating) 
 for tag in parse_tags(include): 
 builder = builder.include(tag) 
 for tag in parse_tags(focus): 
 builder = builder.focus(tag) 
 for tag in parse_tags(exclude): 
 builder = builder.exclude(tag) 
 for tag in parse_tags(avoid): 
 builder = builder.avoid(tag) 
 if artist > 0: 
 builder = builder.pick(CAT_ARTIST, artist, 10) 
 if species > 0: 
 builder = builder.pick(CAT_SPECIES, species, 10) 
 if copyright > 0: 
 builder = builder.pick(CAT_COPYRIGHT, copyright, 10) 
 if character > 0: 
 builder = builder.pick(CAT_CHARACTER, character, 10) 
 if meta > 0: 
 builder = builder.pick(CAT_META, meta, 10) 
 if general > 0: 
 builder = builder.pick(CAT_GENERAL, general, 50) 
 return builder.get_one() 
 except Exception as e: 
 return str(e) 
 
with gr.Blocks() as demo: 
 with gr.Tab("Tag Explorer"): 
 tag = gr.Textbox(label="Tag") 
 category = gr.Dropdown(label="Category", choices=[("All", -1), ("General", CAT_GENERAL), ("Artist", CAT_ARTIST), ("Copyright", CAT_COPYRIGHT), ("Character", CAT_CHARACTER), ("Species", CAT_SPECIES), ("Meta", CAT_META), ("Lore", CAT_LORE)], value=-1) 
 query = gr.Button("Query") 
 output = gr.Dataframe() 
 query.click(fn=query_tag, inputs=[tag, category], outputs=output) 
 with gr.Tab("Prompt Expander"): 
 include = gr.Textbox(label="Positive Prompt - Start with these tags and weight picks toward them") 
 focus = gr.Textbox(label="Focus - Used for picks but not necessarily added to the prompt") 
 exclude = gr.Textbox(label="Negative Prompt - Put these in the negative prompt and weight picks against them") 
 avoid = gr.Textbox(label="Other Negatives - Weighted against picks but not put in the negatives") 
 skip = gr.Textbox(label="Skip - Never pick these tags") 
 rating = gr.Dropdown(label="Rating Limit (not 100% reliable)", choices=[("Safe", "s"), ("Questionable", "q"), ("Explicit", "e")], value="s") 
 with gr.Accordion(label="Tag Counts"): 
 with gr.Row(): 
 general = gr.Number(5, label="General", precision=0, minimum=0, maximum=20) 
 artist = gr.Number(0, label="Artist", precision=0, minimum=0, maximum=5) 
 species = gr.Number(0, label="Species", precision=0, minimum=0, maximum=5) 
 copyright = gr.Number(0, label="Copyright", precision=0, minimum=0, maximum=5) 
 character = gr.Number(0, label="Character", precision=0, minimum=0, maximum=5) 
 meta = gr.Number(0, label="Meta", precision=0, minimum=0, maximum=5) 
 
 generate = gr.Button("Generate") 
 output = gr.Textbox(label="Prompt") 
 generate.click(fn=generate_prompt, inputs=[include, focus, exclude, avoid, skip, rating, general, artist, species, copyright, character, meta], outputs=output) 
 
demo.launch()