app.py

import os
import time
import spaces

import torch
from transformers import (
    AutoModelForPreTraining,
    AutoProcessor,
    AutoConfig,
    PreTrainedTokenizerFast,
)
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import gradio as gr


MODEL_NAME = os.environ.get("MODEL_NAME", None)
assert MODEL_NAME is not None
MODEL_PATH = hf_hub_download(repo_id=MODEL_NAME, filename="model.safetensors")
DEVICE = (
    torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda")
)

BAD_WORD_KEYWORDS = ["(medium)"]


def fix_compiled_state_dict(state_dict: dict):
    return {k.replace("._orig_mod.", "."): v for k, v in state_dict.items()}


def get_bad_words_ids(tokenizer: PreTrainedTokenizerFast):
    ids = [
        [id]
        for token, id in tokenizer.vocab.items()
        if any(word in token for word in BAD_WORD_KEYWORDS)
    ]
    return ids


def prepare_models():
    config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
    model = AutoModelForPreTraining.from_config(
        config, torch_dtype=torch.bfloat16, trust_remote_code=True
    )
    model.decoder_model.use_cache = True
    processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)

    state_dict = load_file(MODEL_PATH)
    state_dict = {k.replace("._orig_mod.", "."): v for k, v in state_dict.items()}
    model.load_state_dict(state_dict)

    model.eval()
    model = model.to(DEVICE)
    # model = torch.compile(model)

    return model, processor


def demo():
    model, processor = prepare_models()
    ban_ids = get_bad_words_ids(processor.decoder_tokenizer)

    @spaces.GPU(duration=5)
    @torch.inference_mode()
    def generate_tags(
        text: str,
        auto_detect: bool,
        copyright_tags: str = "",
        length: str = "short",
        max_new_tokens: int = 128,
        do_sample: bool = False,
        temperature: float = 0.1,
        top_k: int = 10,
        top_p: float = 0.1,
    ):
        tag_text = (
            "<|bos|>"
            f"<|aspect_ratio:tall|><|rating:general|><|length:{length}|>"
            "<|reserved_2|><|reserved_3|><|reserved_4|>"
            "<|translate:exact|><|input_end|>"
            "<copyright>" + copyright_tags.strip()
        )
        if not auto_detect:
            tag_text += "</copyright><character></character><general>"
        inputs = processor(
            encoder_text=text, decoder_text=tag_text, return_tensors="pt"
        )

        start_time = time.time()
        outputs = model.generate(
            input_ids=inputs["input_ids"].to(model.device),
            attention_mask=inputs["attention_mask"].to(model.device),
            encoder_input_ids=inputs["encoder_input_ids"].to(model.device),
            encoder_attention_mask=inputs["encoder_attention_mask"].to(model.device),
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            no_repeat_ngram_size=1,
            eos_token_id=processor.decoder_tokenizer.eos_token_id,
            pad_token_id=processor.decoder_tokenizer.pad_token_id,
            bad_words_ids=ban_ids,
        )
        elapsed = time.time() - start_time

        deocded = ", ".join(
            [
                tag
                for tag in processor.batch_decode(outputs[0], skip_special_tokens=True)
                if tag.strip() != ""
            ]
        )
        return [deocded, f"Time elapsed: {elapsed:.2f} seconds"]

    # warmup
    print("warming up...")
    print(generate_tags("Hatsune Miku is looking at viewer.", True))
    print("done.")

    with gr.Blocks() as ui:
        with gr.Column():
            with gr.Row():
                with gr.Column():
                    text = gr.Text(
                        label="Text",
                        info="Enter a prompt in natural language (currently only English is supported). But maybe danbooru tags are also supported.",
                        lines=4,
                        placeholder="A girl with fox ears and tail in maid costume is looking at viewer.",
                    )
                    auto_detect = gr.Checkbox(
                        label="Auto detect copyright tags.", value=False
                    )
                    copyright_tags = gr.Textbox(
                        label="Copyright tags",
                        info="You can specify copyright tags manually. This must be valid danbooru tags.",
                        placeholder="e.g.) vocaloid, blue archive",
                    )
                    length = gr.Dropdown(
                        label="Length",
                        choices=[
                            "very_short",
                            "short",
                            "long",
                            "very_long",
                        ],
                        value="short",
                    )
                    translate_btn = gr.Button(value="Translate", variant="primary")

                    with gr.Accordion(label="Advanced", open=False):
                        max_new_tokens = gr.Number(label="Max new tokens", value=128)
                        do_sample = gr.Checkbox(label="Do sample", value=False)
                        temperature = gr.Slider(
                            label="Temperature",
                            minimum=0.1,
                            maximum=1.0,
                            value=0.3,
                            step=0.1,
                        )
                        top_k = gr.Slider(
                            label="Top k",
                            minimum=1,
                            maximum=100,
                            value=10,
                            step=10,
                        )
                        top_p = gr.Slider(
                            label="Top p",
                            minimum=0.1,
                            maximum=1.0,
                            value=0.5,
                            step=0.1,
                        )

                with gr.Column():
                    output = gr.Textbox(label="Output", lines=4, interactive=False)
                    time_elapsed = gr.Markdown(value="")

            gr.Examples(
                examples=[
                    [
                        "A girl with fox ears and tail in maid costume is looking at viewer.",
                        False,
                        "",
                        "very_short",
                    ],
                    [
                        "Hatsune Miku is looking at viewer.",
                        True,
                        "",
                        "very_short",
                    ],
                    [
                        "Fujita Kotone, Tsukimura Temari, Hanami Saki from Gakuen Idolmaster. They are in the hole, there are some tables and chairs. One's face is shaded, one is crying, and one is 😊.",
                        True,
                        "",
                        "short",
                    ],
                    [
                        "A single girl wearing red hood is sleeping in the forest. View angle from above. grass field. many colorful flowers. Bright atmosphere.",
                        False,
                        "",
                        "short",
                    ],
                    [
                        "Arona and Plana are hugging each other.",
                        True,
                        "blue archive",
                        "short",
                    ],
                    [
                        "Arona and Plana are hugging each other.",
                        True,
                        "blue archive",
                        "very_long",
                    ],
                    [
                        "There are two girls. A vivacious blonde gyaru leans against a classroom desk, her flashy accessories jingling as she gestures animatedly. Across from her stands the prim and proper class representative, her long black hair neatly framing her face as she listens attentively, occasionally adjusting her glasses with a delicate touch.",
                        False,
                        "",
                        "long",
                    ],
                    [
                        "1girl, solo, white and blue medium hair with side braid, dark blue parka with hoodie, looking at somewhere else viewer, cowboy shot, with many cats. the girl has cat ears and tail.",
                        False,
                        "",
                        "short",
                    ],
                ],
                inputs=[text, auto_detect, copyright_tags, length],
            )

        gr.on(
            triggers=[
                translate_btn.click,
            ],
            fn=generate_tags,
            inputs=[
                text,
                auto_detect,
                copyright_tags,
                length,
                max_new_tokens,
                do_sample,
                temperature,
                top_k,
                top_p,
            ],
            outputs=[output, time_elapsed],
        )

    ui.launch()


if __name__ == "__main__":
    demo()