Spaces:
Running
Running
| 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() | |