init

.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

app.py

@@ -0,0 +1,302 @@
+import sys
+import gradio as gr
+
+import re
+import random
+from time import time
+
+import torch
+from transformers import set_seed
+if sys.platform == "win32":
+    #dev env in windows, @spaces.GPU will cause problem
+    def GPU(func):
+        return func
+else:
+    from spaces import GPU
+
+import kgen.models as models
+import kgen.executor.titpop as titpop
+from kgen.formatter import seperate_tags, apply_format
+from kgen.generate import generate
+
+from diff import load_model, encode_prompts
+from meta import DEFAULT_NEGATIVE_PROMPT
+
+
+sdxl_pipe = load_model()
+
+models.load_model(
+    "KBlueLeaf/TITPOP-200M-dev",
+    device="cuda",
+    subfolder="dan-cc-coyo_epoch2",
+)
+generate(max_new_tokens=4)
+
+
+DEFAULT_FORMAT = """<|special|>, <|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>,
+
+<|extended|>.
+
+<|quality|>, <|meta|>, <|rating|>
+""".strip()
+DEFAULT_TAGS = """
+1girl, 
+ningen mame, ciloranko, 
+solo, dragon girl,
+masterpiece, absurdres, safe, newest
+""".strip()
+DEFAULT_NL = """
+An illustration of a girl
+""".strip()
+
+
+def format_time(timing):
+    total = timing["total"]
+    generate_pass = timing["generate_pass"]
+
+    result = ""
+
+    result += f"""
+### Process Time
+| Total    | {total:5.2f} sec / {generate_pass:5} Passes | {generate_pass/total:7.2f} Passes Per Second|
+|-|-|-|
+"""
+    if "generated_tokens" in timing:
+        total_generated_tokens = timing["generated_tokens"]
+        total_input_tokens = timing["input_tokens"]
+    if "generated_tokens" in timing and "total_sampling" in timing:
+        sampling_time = timing["total_sampling"] / 1000
+        process_time = timing["prompt_process"] / 1000
+        model_time = timing["total_eval"] / 1000
+
+        result += f"""| Process  | {process_time:5.2f} sec / {total_input_tokens:5} Tokens | {total_input_tokens/process_time:7.2f} Tokens Per Second|
+| Sampling | {sampling_time:5.2f} sec / {total_generated_tokens:5} Tokens | {total_generated_tokens/sampling_time:7.2f} Tokens Per Second|
+| Eval     | {model_time:5.2f} sec / {total_generated_tokens:5} Tokens | {total_generated_tokens/model_time:7.2f} Tokens Per Second|
+"""
+
+    if "generated_tokens" in timing:
+        result += f"""
+### Processed Tokens:
+* {total_input_tokens:} Input Tokens
+* {total_generated_tokens:} Output Tokens
+"""
+    return result
+
+
+@GPU
+@torch.no_grad()
+def generate(
+    tags,
+    nl_prompt,
+    black_list,
+    temp,
+    target_length,
+    top_p,
+    min_p,
+    top_k,
+    seed,
+    escape_brackets,
+):
+    titpop.BAN_TAGS = [t.strip() for t in black_list.split(",") if t.strip()]
+    generation_setting = {
+        "seed": seed,
+        "temperature": temp,
+        "top_p": top_p,
+        "min_p": min_p,
+        "top_k": top_k,
+    }
+    inputs = seperate_tags(tags.split(","))
+    if nl_prompt:
+        if "<|extended|>" in DEFAULT_FORMAT:
+            inputs["extended"] = nl_prompt
+        elif "<|generated|>" in DEFAULT_FORMAT:
+            inputs["generated"] = nl_prompt
+    input_prompt = apply_format(inputs, DEFAULT_FORMAT)
+    if escape_brackets:
+        input_prompt = re.sub(r"([()\[\]])", r"\\\1", input_prompt)
+
+    meta, operations, general, nl_prompt = titpop.parse_titpop_request(
+        seperate_tags(tags.split(",")),
+        nl_prompt,
+        tag_length_target=target_length,
+        generate_extra_nl_prompt="<|generated|>" in DEFAULT_FORMAT or not nl_prompt,
+    )
+    t0 = time()
+    for result, timing in titpop.titpop_runner_generator(
+        meta, operations, general, nl_prompt, **generation_setting
+    ):
+        result = apply_format(result, DEFAULT_FORMAT)
+        if escape_brackets:
+            result = re.sub(r"([()\[\]])", r"\\\1", result)
+        timing["total"] = time() - t0
+        yield result, input_prompt, format_time(timing)
+
+
+@GPU
+@torch.no_grad()
+def generate_image(
+    seed,
+    prompt,
+    prompt2,
+):
+    torch.cuda.empty_cache()
+    prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
+        encode_prompts(sdxl_pipe, prompt, DEFAULT_NEGATIVE_PROMPT)
+    )
+    set_seed(seed)
+    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]
+    prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
+        encode_prompts(sdxl_pipe, prompt2, DEFAULT_NEGATIVE_PROMPT)
+    )
+    set_seed(seed)
+    result2 = 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()
+    return result2, result
+
+
+if __name__ == "__main__":
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""# TITPOP DEMO""")
+        with gr.Accordion("Introduction and Instructions", open=False):
+            gr.Markdown(
+                """
+### What is this:
+TITPOP
+
+**The implementation is a little bit inefficient, image gen may be a little bit slower than expected.**
+"""
+            )
+        with gr.Row():
+            with gr.Column(scale=5):
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        tags_input = gr.TextArea(
+                            label="Danbooru Tags",
+                            lines=6,
+                            show_copy_button=True,
+                            interactive=True,
+                            value=DEFAULT_TAGS,
+                            placeholder="Enter danbooru tags here",
+                        )
+                        nl_prompt_input = gr.Textbox(
+                            label="Natural Language Prompt",
+                            lines=6,
+                            show_copy_button=True,
+                            interactive=True,
+                            value=DEFAULT_NL,
+                            placeholder="Enter Natural Language Prompt here",
+                        )
+                        black_list = gr.TextArea(
+                            label="Black List (seperated by comma)",
+                            lines=4,
+                            interactive=True,
+                            value="monochrome",
+                            placeholder="Enter tag/nl black list here",
+                        )
+                    with gr.Column(scale=2):
+                        target_length = gr.Dropdown(
+                            label="Target Length",
+                            choices=["very_short", "short", "long", "very_long"],
+                            value="short",
+                        )
+                        temp = gr.Slider(
+                            label="Temp",
+                            minimum=0.0,
+                            maximum=1.5,
+                            value=0.5,
+                            step=0.05,
+                        )
+                        top_p = gr.Slider(
+                            label="Top P",
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=0.95,
+                            step=0.05,
+                        )
+                        min_p = gr.Slider(
+                            label="Min P",
+                            minimum=0.0,
+                            maximum=0.2,
+                            value=0.05,
+                            step=0.01,
+                        )
+                        top_k = gr.Slider(
+                            label="Top K", minimum=0, maximum=120, value=60, step=1
+                        )
+                        with gr.Row():
+                            seed = gr.Number(
+                                label="Seed",
+                                minimum=0,
+                                maximum=2147483647,
+                                value=20090220,
+                                step=1,
+                            )
+                            escape_brackets = gr.Checkbox(
+                                label="Escape Brackets", value=False
+                            )
+                        submit = gr.Button("TITPOP!", variant="primary")
+                with gr.Accordion("Speed statstics", open=False):
+                    cost_time = gr.Markdown()
+            with gr.Column(scale=5):
+                result = gr.TextArea(
+                    label="Result", lines=8, show_copy_button=True, interactive=False
+                )
+                input_prompt = gr.Textbox(
+                    label="Input Prompt", lines=1, interactive=False, visible=False
+                )
+                gen_img = gr.Button("Generate Image from Result", variant="primary")
+                with gr.Row():
+                    with gr.Column():
+                        img1 = gr.Image(label="Original Propmt", interactive=False)
+                    with gr.Column():
+                        img2 = gr.Image(label="Generated Prompt", interactive=False)
+        submit.click(
+            generate,
+            [
+                tags_input,
+                nl_prompt_input,
+                black_list,
+                temp,
+                target_length,
+                top_p,
+                min_p,
+                top_k,
+                seed,
+                escape_brackets,
+            ],
+            [
+                result,
+                input_prompt,
+                cost_time,
+            ],
+            queue=True,
+        )
+        gen_img.click(
+            generate_image,
+            [seed, result, input_prompt],
+            [img1, img2],
+            queue=True,
+        )
+
+    demo.launch()

diff.py

@@ -0,0 +1,120 @@
+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
+
+torch.set_float32_matmul_precision("medium")
+
+
+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 model_forward(k_diffusion_model: torch.nn.Module):
+    orig_forward = k_diffusion_model.forward
+
+    def forward(*args, **kwargs):
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            result = orig_forward(*args, **kwargs)
+        return result.float()
+
+    return forward
+
+
+def load_model(model_id="KBlueLeaf/Kohaku-XL-Zeta", 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")
+    pipe.k_diffusion_model.forward = model_forward(pipe.k_diffusion_model)
+    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

meta.py

@@ -0,0 +1,54 @@
+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)",
+    "style 6": "ogipote, misu kasumi, fuzichoco, ciloranko, ninjin nouka, ningen mame, ask (askzy), kita (kitairoha), maccha (mochancc)",
+    "no style": "",
+}
+
+MODEL_DEFAULT_QUALITY_LIST = {
+    "KBlueLeaf/Kohaku-XL-Zeta": "masterpiece, newest, absurdres",
+    "KBlueLeaf/Kohaku-XL-Epsilon-rev2": "masterpiece, newest, absurdres",
+    "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-Zeta": """<|special|>, 
+<|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>, 
+
+<|quality|>, <|meta|>, <|rating|>""",
+    "KBlueLeaf/Kohaku-XL-Epsilon-rev2": """<|special|>, 
+<|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>, 
+
+<|quality|>, <|meta|>, <|rating|>""",
+    "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,3 @@
+git+https://${GITHUB_TOKEN}@github.com/KohakuBlueleaf/TITPOP-KGen@titpop
+gradio
+spaces