Spaces:
Running
Running
File size: 3,663 Bytes
379b837 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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()
|