first commit

app.py

@@ -0,0 +1,152 @@
+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()

diff.py

@@ -0,0 +1,107 @@
+from functools import partial
+
+import torch
+from diffusers import StableDiffusionXLKDiffusionPipeline
+from k_diffusion.sampling import get_sigmas_polyexponential
+from k_diffusion.sampling import sample_dpmpp_2m_sde
+
+
+def set_timesteps_polyexponential(self, orig_sigmas, num_inference_steps, device=None):
+    self.num_inference_steps = num_inference_steps
+
+    self.sigmas = get_sigmas_polyexponential(
+        num_inference_steps + 1,
+        sigma_min=orig_sigmas[-2],
+        sigma_max=orig_sigmas[0],
+        rho=0.666666,
+        device=device or "cpu",
+    )
+    self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas.new_zeros([1])])
+
+
+def load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon", device="cuda"):
+    pipe: StableDiffusionXLKDiffusionPipeline
+    pipe = StableDiffusionXLKDiffusionPipeline.from_pretrained(
+        model_id, torch_dtype=torch.float16
+    ).to(device)
+    pipe.scheduler.set_timesteps = partial(
+        set_timesteps_polyexponential, pipe.scheduler, pipe.scheduler.sigmas
+    )
+    pipe.sampler = partial(sample_dpmpp_2m_sde, eta=0.35, solver_type="heun")
+    return pipe
+
+
+def encode_prompts(pipe: StableDiffusionXLKDiffusionPipeline, prompt, neg_prompt):
+    max_length = pipe.tokenizer.model_max_length
+
+    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
+    input_ids2 = pipe.tokenizer_2(prompt, return_tensors="pt").input_ids.to("cuda")
+
+    negative_ids = pipe.tokenizer(
+        neg_prompt,
+        truncation=False,
+        padding="max_length",
+        max_length=input_ids.shape[-1],
+        return_tensors="pt",
+    ).input_ids.to("cuda")
+    negative_ids2 = pipe.tokenizer_2(
+        neg_prompt,
+        truncation=False,
+        padding="max_length",
+        max_length=input_ids.shape[-1],
+        return_tensors="pt",
+    ).input_ids.to("cuda")
+
+    if negative_ids.size() > input_ids.size():
+        input_ids = pipe.tokenizer(
+            prompt,
+            truncation=False,
+            padding="max_length",
+            max_length=negative_ids.shape[-1],
+            return_tensors="pt",
+        ).input_ids.to("cuda")
+        input_ids2 = pipe.tokenizer_2(
+            prompt,
+            truncation=False,
+            padding="max_length",
+            max_length=negative_ids.shape[-1],
+            return_tensors="pt",
+        ).input_ids.to("cuda")
+
+    concat_embeds = []
+    neg_embeds = []
+    for i in range(0, input_ids.shape[-1], max_length):
+        concat_embeds.append(pipe.text_encoder(input_ids[:, i : i + max_length])[0])
+        neg_embeds.append(pipe.text_encoder(negative_ids[:, i : i + max_length])[0])
+
+    concat_embeds2 = []
+    neg_embeds2 = []
+    pooled_embeds2 = []
+    neg_pooled_embeds2 = []
+    for i in range(0, input_ids.shape[-1], max_length):
+        hidden_states = pipe.text_encoder_2(
+            input_ids2[:, i : i + max_length], output_hidden_states=True
+        )
+        concat_embeds2.append(hidden_states.hidden_states[-2])
+        pooled_embeds2.append(hidden_states[0])
+
+        hidden_states = pipe.text_encoder_2(
+            negative_ids2[:, i : i + max_length], output_hidden_states=True
+        )
+        neg_embeds2.append(hidden_states.hidden_states[-2])
+        neg_pooled_embeds2.append(hidden_states[0])
+
+    prompt_embeds = torch.cat(concat_embeds, dim=1)
+    negative_prompt_embeds = torch.cat(neg_embeds, dim=1)
+    prompt_embeds2 = torch.cat(concat_embeds2, dim=1)
+    negative_prompt_embeds2 = torch.cat(neg_embeds2, dim=1)
+    prompt_embeds = torch.cat([prompt_embeds, prompt_embeds2], dim=-1)
+    negative_prompt_embeds = torch.cat(
+        [negative_prompt_embeds, negative_prompt_embeds2], dim=-1
+    )
+
+    pooled_embeds2 = torch.mean(torch.stack(pooled_embeds2, dim=0), dim=0)
+    neg_pooled_embeds2 = torch.mean(torch.stack(neg_pooled_embeds2, dim=0), dim=0)
+
+    return prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2
+

dtg.py

@@ -0,0 +1,92 @@
+import time
+import pathlib
+
+import kgen.models as models
+from kgen.formatter import seperate_tags, apply_format, apply_dtg_prompt
+from kgen.metainfo import TARGET
+from kgen.generate import tag_gen
+from kgen.logging import logger
+
+
+SEED_MAX = 2**31 - 1
+DEFAULT_FORMAT = """<|special|>, 
+<|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>, 
+
+<|quality|>, <|meta|>, <|rating|>"""
+
+
+def process(
+    prompt: str,
+    aspect_ratio: float,
+    seed: int,
+    tag_length: str,
+    ban_tags: str,
+    format: str,
+    temperature: float,
+):
+    propmt_preview = prompt.replace("\n", " ")[:40]
+    logger.info(f"Processing propmt: {propmt_preview}...")
+    logger.info(f"Processing with seed: {seed}")
+    black_list = [tag.strip() for tag in ban_tags.split(",") if tag.strip()]
+    all_tags = [tag.strip() for tag in prompt.strip().split(",") if tag.strip()]
+
+    tag_length = tag_length.replace(" ", "_")
+    len_target = TARGET[tag_length]
+
+    tag_map = seperate_tags(all_tags)
+    dtg_prompt = apply_dtg_prompt(tag_map, tag_length, aspect_ratio)
+    for _, extra_tokens, iter_count in tag_gen(
+        models.text_model,
+        models.tokenizer,
+        dtg_prompt,
+        tag_map["special"] + tag_map["general"],
+        len_target,
+        black_list,
+        temperature=temperature,
+        top_p=0.8,
+        top_k=80,
+        max_new_tokens=512,
+        max_retry=10,
+        max_same_output=5,
+        seed=seed % SEED_MAX,
+    ):
+        pass
+    tag_map["general"] += extra_tokens
+    prompt_by_dtg = apply_format(tag_map, format)
+    logger.info(
+        "Prompt processing done. General Tags Count: "
+        f"{len(tag_map['general'] + tag_map['special'])}"
+        f" | Total iterations: {iter_count}"
+    )
+    return prompt_by_dtg
+
+
+if __name__ == "__main__":
+    models.model_dir = pathlib.Path(__file__).parent / "models"
+
+    file = models.download_gguf()
+    files = models.list_gguf()
+    file = files[-1]
+    logger.info(f"Use gguf model from local file: {file}")
+    models.load_model(file, gguf=True)
+
+    prompt = """
+1girl, ask (askzy), masterpiece
+"""
+
+    t0 = time.time_ns()
+    result = process(
+        prompt,
+        aspect_ratio=1.0,
+        seed=1,
+        tag_length="long",
+        ban_tags="",
+        format=DEFAULT_FORMAT,
+        temperature=1.35,
+    )
+    t1 = time.time_ns()
+    logger.info(f"Result:\n{result}")
+    logger.info(f"Time cost: {(t1 - t0) / 10**6:.1f}ms")

meta.py

@@ -0,0 +1,37 @@
+DEFAULT_STYLE_LIST = {
+    "style 1": "ask (askzy), torino aqua, migolu",
+    "style 2": "azuuru, torino aqua, kedama milk, fuzichoco, ask (askzy), chen bin, atdan, hito, mignon",
+    "style 3": "nou (nounknown), shikimi (yurakuru), namiki itsuki, lemon89h, satsuki (miicat), chon (chon33v), omutatsu, mochizuki kei",
+    "style 4": "ciloranko, maccha (mochancc), lobelia (saclia), migolu, ask (askzy), wanke, jiu ye sang, rumoon, mizumi zumi",
+    "style 5": "reoen, alchemaniac, rella, watercolor (medium)",
+    "no style": "",
+}
+
+MODEL_DEFAULT_QUALITY_LIST = {
+    "KBlueLeaf/Kohaku-XL-Epsilon": "masterpiece, newest, absurdres, safe",
+    "cagliostrolab/animagine-xl-3.1": "masterpiece, newest, very aesthetic, absurdres, safe",
+}
+
+MODEL_FORMAT_LIST = {
+    "KBlueLeaf/Kohaku-XL-Epsilon": """<|special|>, 
+<|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>, 
+
+<|quality|>, <|meta|>, <|rating|>""",
+    "cagliostrolab/animagine-xl-3.1": """<|special|>, 
+<|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>, 
+
+<|quality|>, <|meta|>, <|rating|>""",
+}
+
+
+DEFAULT_NEGATIVE_PROMPT = """
+low quality, worst quality, normal quality, text, signature, jpeg artifacts, 
+bad anatomy, old, early, mini skirt, nsfw, chibi, multiple girls, multiple boys, 
+multiple tails, multiple views, copyright name, watermark, artist name, signature
+"""

requirements.txt

@@ -0,0 +1,7 @@
+diffusers
+transformers
+k_diffusion
+requests
+sentencepiece
+tipo-kgen
+spaces