Spaces:
Runtime error
Runtime error
File size: 5,091 Bytes
7d4afe8 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import random
from time import time_ns
import torch
import spaces
import gradio as gr
from transformers import set_seed
from kgen import models
from diff import load_model, encode_prompts
from dtg import process
from meta import (
DEFAULT_STYLE_LIST,
MODEL_FORMAT_LIST,
MODEL_DEFAULT_QUALITY_LIST,
DEFAULT_NEGATIVE_PROMPT,
)
sdxl_pipe = load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon", device="cuda")
models.load_model(models.model_list[0])
models.text_model.cuda()
current_dtg_model = models.model_list[0]
current_sdxl_model = "KBlueLeaf/Kohaku-XL-Epsilon"
@spaces.GPU
def gen(
sdxl_model: str,
dtg_model: str,
style: str,
base_prompt: str,
addon_prompt: str = "",
):
global current_dtg_model, current_sdxl_model, sdxl_pipe
if sdxl_model != current_sdxl_model:
sdxl_pipe = load_model(model_id=sdxl_model, device="cuda")
current_sdxl_model = sdxl_model
if dtg_model != current_dtg_model:
models.load_model(dtg_model)
models.text_model.cuda()
current_dtg_model = dtg_model
t0 = time_ns()
seed = random.randint(0, 2**31 - 1)
prompt = (
f"{base_prompt}, {addon_prompt}, "
f"{DEFAULT_STYLE_LIST[style]}, "
f"{MODEL_DEFAULT_QUALITY_LIST[sdxl_model]}, "
)
full_prompt = process(
prompt,
aspect_ratio=1.0,
seed=seed,
tag_length="short",
ban_tags=".*alternate.*, character doll, multiple.*, .*cosplay.*, .*name, .*text.*",
format=MODEL_FORMAT_LIST[sdxl_model],
temperature=1.2,
)
torch.cuda.empty_cache()
prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
encode_prompts(sdxl_pipe, full_prompt, DEFAULT_NEGATIVE_PROMPT)
)
set_seed(seed)
with torch.autocast("cuda"):
result = sdxl_pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_embeds2,
negative_pooled_prompt_embeds=neg_pooled_embeds2,
num_inference_steps=24,
width=1024,
height=1024,
guidance_scale=6.0,
).images[0]
torch.cuda.empty_cache()
t1 = time_ns()
return result.convert("RGB"), full_prompt, f"Cost: {(t1 - t0) / 1e9:.2}sec"
if __name__ == "__main__":
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("""# This Cute Dragon Girl Doesn't Exist""")
with gr.Accordion("Introduction and Instructions", open=False):
gr.Markdown(
"""
### What is this:
"This Cute Dragon Girl Doesn't Exist" is a Demo for KGen System(DanTagGen) with SDXL anime models.
It is aimed to show how the DanTagGen can be used to "refine/upsample" simple prompt to help the T2I model.
Since I already have some application and demo on DanTagGen.
This demo is designed to be more "simple" than before.
Just one click, and get the result with high quality and high diversity.
### How to use it:
click "Next" button until you get the dragon girl you like.
### Resources:
- My anime model: [Kohaku XL Epsilon](https://huggingface.co/KBlueLeaf/Kohaku-XL-Epsilon)
- DanTagGen: [DanTagGen](https://huggingface.co/KBlueLeaf/DanTagGen-beta)
- DanTagGen extension: [z-a1111-sd-webui-dtg](https://github.com/KohakuBlueleaf/z-a1111-sd-webui-dtg)
"""
)
with gr.Row():
with gr.Column(scale=3):
with gr.Row():
sdxl_model = gr.Dropdown(
MODEL_FORMAT_LIST,
label="SDXL Model",
value=list(MODEL_FORMAT_LIST)[0],
)
dtg_model = gr.Dropdown(
models.model_list,
label="DTG Model",
value=models.model_list[0],
)
base_prompt = gr.Textbox(
label="Base prompt",
lines=1,
value="1girl, solo, dragon girl, dragon wings, dragon horns, dragon tail",
interactive=False,
)
with gr.Row():
addon_propmt = gr.Textbox(
label="Addon prompt",
lines=1,
value="cowboy shot, loli",
)
style = gr.Dropdown(
DEFAULT_STYLE_LIST,
label="Style",
value=list(DEFAULT_STYLE_LIST)[0],
)
submit = gr.Button("Next")
dtg_output = gr.TextArea(label="DTG output", lines=9, show_copy_button=True)
cost_time = gr.Markdown()
with gr.Column(scale=4):
result = gr.Image(label="Result", type="numpy", interactive=False)
submit.click(
fn=gen,
inputs=[sdxl_model, dtg_model, style, base_prompt, addon_propmt],
outputs=[result, dtg_output, cost_time],
)
demo.launch()
|