app.py

import os
import random
from functools import partial

if os.environ.get("IN_SPACES", None) is not None:
    in_spaces = True
    import spaces
else:
    in_spaces = False
import gradio as gr
import torch

try:
    # pre-import triton can avoid diffusers/transformers make import error
    import triton
except ImportError:
    print("Triton not found, skip pre import")

## HDM model dep
import xut.env
xut.env.TORCH_COMPILE = True
xut.env.USE_LIGER = False
xut.env.USE_VANILLA = False
xut.env.USE_XFORMERS = True
xut.env.USE_XFORMERS_LAYERS = True
from hdm.pipeline import HDMXUTPipeline

## TIPO
import kgen.models as kgen_models
import kgen.executor.tipo as tipo
from kgen.formatter import apply_format, seperate_tags


torch.set_float32_matmul_precision("high")


DEFAULT_FORMAT = """
<|special|>,
<|characters|>, <|copyrights|>,
<|artist|>,

<|general|>,

<|extended|>.

<|quality|>, <|meta|>, <|rating|>
""".strip()


def GPU(func=None, duration=None):
    if func is None:
        return partial(GPU, duration=duration)
    if in_spaces:
        if duration:
            return spaces.GPU(func, duration=duration)
        else:
            return spaces.GPU(func)
    else:
        return func


def prompt_opt(tags, nl_prompt, aspect_ratio, seed):
    meta, operations, general, nl_prompt = tipo.parse_tipo_request(
        seperate_tags(tags.split(",")),
        nl_prompt,
        tag_length_target="long",
        nl_length_target="short",
        generate_extra_nl_prompt=True,
    )
    meta["aspect_ratio"] = f"{aspect_ratio:.3f}"
    result, timing = tipo.tipo_runner(meta, operations, general, nl_prompt, seed=seed)
    return apply_format(result, DEFAULT_FORMAT).strip().strip(".").strip(",")


print("Loading models, please wait...")
device = torch.device("cuda")

model = (
    HDMXUTPipeline.from_pretrained(
        "KBlueLeaf/HDM-xut-340M-anime",
        trust_remote_code=True,
    )
    .to(torch.float16)
    .to(device)
)

tipo_model_name, gguf_list = kgen_models.tipo_model_list[0]
kgen_models.load_model(tipo_model_name, device="cuda")
print("Models loaded successfully. UI is ready.")


@GPU(duration=10)
@torch.no_grad()
def generate(
    nl_prompt: str,
    tag_prompt: str,
    negative_prompt: str,
    tipo_enable: bool,
    format_enable: bool,
    num_images: int,
    steps: int,
    cfg_scale: float,
    size: int,
    aspect_ratio: str,
    fixed_short_edge: bool,
    zoom: float,
    x_shift: float,
    y_shift: float,
    tread_gamma1: float,
    tread_gamma2: float,
    seed: int,
    progress=gr.Progress(),
):
    as_w, as_h = aspect_ratio.split(":")
    aspect_ratio = float(as_w) / float(as_h)
    # Set seed for reproducibility
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)

    # TIPO
    if tipo_enable:
        tipo.BAN_TAGS = [i.strip() for i in negative_prompt.split(",") if i.strip()]
        final_prompt = prompt_opt(tag_prompt, nl_prompt, aspect_ratio, seed)
    elif format_enable:
        final_prompt = apply_format(nl_prompt, DEFAULT_FORMAT)
    else:
        final_prompt = tag_prompt + "\n" + nl_prompt

    yield None, final_prompt

    prompts_to_generate = [final_prompt.replace("\n", " ")] * num_images
    negative_prompts_to_generate = [negative_prompt] * num_images

    if fixed_short_edge:
        if aspect_ratio > 1:
            h_factor = 1
            w_factor = aspect_ratio
        else:
            h_factor = 1 / aspect_ratio
            w_factor = 1
    else:
        w_factor = aspect_ratio**0.5
        h_factor = 1 / w_factor

    w = int(size * w_factor / 16) * 16
    h = int(size * h_factor / 16) * 16

    print("=" * 100)
    print(
        f"Generating {num_images} image(s) with seed: {seed} and resolution {w}x{h}"
    )
    print("-" * 80)
    print(f"Final prompt: {final_prompt}")
    print("-" * 80)
    print(f"Negative prompt: {negative_prompt}")
    print("-" * 80)

    prompts_batch = prompts_to_generate
    neg_prompts_batch = negative_prompts_to_generate

    images = model(
        prompts_batch,
        neg_prompts_batch,
        num_inference_steps=steps,
        cfg_scale=cfg_scale,
        width=w,
        height=h,
        camera_param={
            "zoom": zoom,
            "x_shift": x_shift,
            "y_shift": y_shift,
        },
        tread_gamma1=tread_gamma1,
        tread_gamma2=tread_gamma2,
    ).images

    yield images, final_prompt


# --- Gradio UI Definition ---
with gr.Blocks(title="HDM Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# HDM Demo")
    gr.Markdown(
        "### Enter a natural language prompt and/or specific tags to generate an image."
    )
    with gr.Accordion("Introduction", open=False):
        gr.Markdown("""
# HDM: HomeDiffusion Model Project
HDM is a project to implement a series of generative model that can be pretrained at home.

* Project Source code: https://github.com/KBlueLeaf/HDM
* Model: https://huggingface.co/KBlueLeaf/HDM-xut-340M-anime

## Usage
This early model used a model trained on anime image set only, 
so you should expect to see anime style images only in this demo.

For prompting, enter danbooru tag prompt to the box "Tag Prompt" with comma seperated and remove the underscore.
enter natural language prompt to the box "Natural Language Prompt" and enter negative prompt to the box "Negative Prompt".

If you don't want to spent so much effort on prompting, try to keep "Enable TIPO" selected.

If you don't want to apply any pre-defined format, unselect "Enable TIPO" and "Enable Format".

## Model Spec
- Backbone: 343M XUT(UViT modified) arch
- Text Encoder: Qwen3 0.6B (596M)
- VAE: EQ-SDXL-VAE, an EQ-VAE finetuned sdxl vae.

## Pretraining Dataset
- Danbooru 2023 (latest id around 8M)
- Pixiv famous artist set
- some pvc figure photos
""")

    with gr.Row():
        with gr.Column(scale=2):
            nl_prompt_box = gr.Textbox(
                label="Natural Language Prompt",
                placeholder="e.g., A beautiful anime girl standing in a blooming cherry blossom forest",
                lines=3,
            )
            tag_prompt_box = gr.Textbox(
                label="Tag Prompt (comma-separated)",
                placeholder="e.g., 1girl, solo, long hair, cherry blossoms, school uniform",
                lines=3,
            )
            neg_prompt_box = gr.Textbox(
                label="Negative Prompt",
                value=(
                    "llow quality, worst quality, text, signature, jpeg artifacts, bad anatomy, old, early, copyright name, watermark, artist name, signature, weibo username, realistic"
                ),
                lines=3,
            )
            with gr.Row():
                tipo_enable = gr.Checkbox(
                    label="Enable TIPO",
                    value=True,
                )
                format_enable = gr.Checkbox(
                    label="Enable Format",
                    value=True,
                )
            with gr.Row():
                zoom_slider = gr.Slider(
                    label="Zoom", minimum=0.5, maximum=2.0, value=1.0, step=0.01
                )
                x_shift_slider = gr.Slider(
                    label="X Shift", minimum=-0.5, maximum=0.5, value=0.0, step=0.01
                )
                y_shift_slider = gr.Slider(
                    label="Y Shift", minimum=-0.5, maximum=0.5, value=0.0, step=0.01
                )
        with gr.Column(scale=1):
            with gr.Row():
                num_images_slider = gr.Slider(
                    label="Number of Images", minimum=1, maximum=4, value=1, step=1
                )
                steps_slider = gr.Slider(
                    label="Inference Steps", minimum=1, maximum=50, value=24, step=1
                )

            with gr.Row():
                cfg_slider = gr.Slider(
                    label="CFG Scale", minimum=1.0, maximum=7.0, value=4.0, step=0.1
                )
                seed_input = gr.Number(
                    label="Seed",
                    value=-1,
                    precision=0,
                    info="Set to -1 for a random seed.",
                )

            with gr.Row():
                tread_gamma1_slider = gr.Slider(
                    label="Tread Gamma 1",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.0,
                    step=0.05,
                    interactive=True,
                )
                tread_gamma2_slider = gr.Slider(
                    label="Tread Gamma 2",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.0,
                    step=0.05,
                    interactive=True,
                )

            with gr.Row():
                size_slider = gr.Slider(
                    label="Base Image Size",
                    minimum=768,
                    maximum=1280,
                    value=1024,
                    step=16,
                )
            with gr.Row():
                aspect_ratio_box = gr.Textbox(
                    label="Ratio (W:H)",
                    value="1:1",
                )
                fixed_short_edge = gr.Checkbox(
                    label="Fixed Edge",
                    value=True,
                )

    with gr.Row():
        with gr.Column(scale=1):
            generate_button = gr.Button("Generate", variant="primary")
            output_prompt = gr.TextArea(
                label="Final Prompt",
                show_label=True,
                interactive=False,
                lines=32,
                max_lines=32,
            )
        with gr.Column(scale=2):
            output_gallery = gr.Gallery(
                label="Generated Images",
                show_label=True,
                elem_id="gallery",
                columns=2,
                rows=3,
                height="800px",
            )

    generate_button.click(
        fn=generate,
        inputs=[
            nl_prompt_box,
            tag_prompt_box,
            neg_prompt_box,
            tipo_enable,
            format_enable,
            num_images_slider,
            steps_slider,
            cfg_slider,
            size_slider,
            aspect_ratio_box,
            fixed_short_edge,
            zoom_slider,
            x_shift_slider,
            y_shift_slider,
            tread_gamma1_slider,
            tread_gamma2_slider,
            seed_input,
        ],
        outputs=[output_gallery, output_prompt],
        show_progress_on=output_gallery,
    )

if __name__ == "__main__":
    demo.launch()