app.py

"""
DiffusionPen Ukrainian Handwriting Demo — HuggingFace Spaces
"""

import json
import os
import random
import sys
import time
from datetime import datetime
from types import SimpleNamespace

import gradio as gr
import numpy as np
import spaces
import torch
import torchvision.transforms
from diffusers import AutoencoderKL, DDIMScheduler
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import CanineTokenizer, CanineModel

_DEMO_ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _DEMO_ROOT)

from unet import UNetModel
from feature_extractor import ImageEncoder
from generate_sentence import (
    PUNCTUATION, strip_dp_prefix, detect_num_classes,
    prepare_style_reference_image, generate_single_word, crop_whitespace,
    compose_sentence_geometry,
    normalize_ink_brightness, sample_punctuation, split_word_for_generation,
)
from utils.word_dataset import char_classes as WORD_CHAR_CLASSES

# ---------------------------------------------------------------------------
# Config — update HF_MODEL_REPO to your model repo before deploying
# ---------------------------------------------------------------------------
HF_MODEL_REPO  = "kdonitz/diffusionpen-ukrainian"
SD_REPO        = "runwayml/stable-diffusion-v1-5"
STYLE_REFS_DIR = os.path.join(_DEMO_ROOT, "style_refs")
PUNCT_BANK_DIR = os.path.join(_DEMO_ROOT, "punct_bank")

IMG_HEIGHT     = 64
IMG_WIDTH      = 256
TEXT_MAX_LEN   = 40
CANVAS_HEIGHT  = 104
NUM_RES_BLOCKS = 2

MODEL: dict = {}
WRITER_LIST: list = []
WRITER_ID_MAP: dict = {}
_LUCKY_WRITERS: list = []  # top-20 pool — populated after load_models()

_LUCKY_PHRASES = [
    "дивовижний", "неперевершений", "неприродний", "кривавий", "жіночний",
    "бульбашка", "спати", "клуб", "гладити", "починати",
    "натякати", "родина", "квіти", "шум", "ніч",
    "відро", "торгівля", "забруднення", "палиця", "вітрило",
    "бомба", "книги", "автомобіль", "подія", "помилка",
    "край", "поїзд", "дерева", "вага", "колесо",
    "рік", "цинк", "батат", "письмо", "робота",
    "бажання", "крило", "вино", "відпустка", "теорія",
    "суспільство", "шарф", "сірник", "вимір", "розум",
    "нервовий", "зубчастий", "важливий", "загальний", "крихкий",
    "ламкий", "щасливий", "жорстокий", "жвавий", "іронічний",
    "різкий", "позбавляти", "направляти", "планувати", "презентувати",
    "бігти", "поспішати", "походити", "заявляти", "лякати",
    "плавати", "тривати", "запитувати", "спонукати", "терпіти",
    "ділити", "радити", "збирати", "впасти", "здобувати",
    "обіймати", "передавати", "обмінювати", "стерти", "фарбувати",
    "благати", "накладати", "посилатися", "гриміти", "шити",
    "метати", "сіяти", "вгадувати", "надавати", "приносити",
    "приймати", "гавань", "частина", "штовхати", "багатство",
    "пустеля", "система", "площа", "повага", "офіс",
    "новини", "чоловіки", "адвокат", "озеро", "залізо",
    "ґрунт", "губернатор", "парта", "розвиток", "контроль",
    "ланцюг", "кактус", "торт", "шанс", "печера",
    "брат", "сестри", "змія", "схил", "ріка",
    "мета", "економічний", "прямий", "захоплений", "вигідний",
    "великий", "некерований", "переможний", "блукаючий", "практичний",
    "мудрий", "смішний", "знайомий", "чорний", "покинутий",
    "здатний", "чарівний", "яскравий", "виникати", "бути",
    "дихати", "транслювати", "горіти", "заряджати", "плутати",
    "закривати", "крокувати", "готувати", "платити", "вказувати",
]

LOG_DIR     = os.path.join(_DEMO_ROOT, "demo_logs")
IMG_LOG_DIR = os.path.join(LOG_DIR, "images")
JSONL_LOG   = os.path.join(LOG_DIR, "generations.jsonl")

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
def load_models():
    global WRITER_LIST, WRITER_ID_MAP

    print("[demo] Downloading checkpoint from HF Hub...")
    ckpt_path      = hf_hub_download(repo_id=HF_MODEL_REPO, filename="ema_ckpt.pt")
    style_enc_path = hf_hub_download(repo_id=HF_MODEL_REPO, filename="mixed_ukr_mobilenetv2_100.pth")

    print("[demo] Loading models on CPU...")
    state_dict = torch.load(ckpt_path, map_location="cpu")
    state_dict = strip_dp_prefix(state_dict)
    num_classes = detect_num_classes(state_dict)
    print(f"[demo]   {num_classes} writer classes")

    tokenizer    = CanineTokenizer.from_pretrained("google/canine-c")
    canine_model = CanineModel.from_pretrained("google/canine-c")

    unet = UNetModel(
        image_size=(IMG_HEIGHT, IMG_WIDTH),
        in_channels=4, model_channels=320, out_channels=4,
        num_res_blocks=NUM_RES_BLOCKS,
        attention_resolutions=(1, 1), channel_mult=(1, 1), num_heads=4,
        num_classes=num_classes, context_dim=320,
        vocab_size=WORD_CHAR_CLASSES, text_encoder=canine_model,
        args=SimpleNamespace(interpolation=False, mix_rate=None),
    )
    unet.load_state_dict(state_dict)
    unet.eval()

    vae = AutoencoderKL.from_pretrained(SD_REPO, subfolder="vae")
    vae.requires_grad_(False)

    noise_scheduler = DDIMScheduler.from_pretrained(SD_REPO, subfolder="scheduler")

    style_extractor = ImageEncoder(model_name="mobilenetv2_100", num_classes=0,
                                   pretrained=False, trainable=False)
    style_sd    = torch.load(style_enc_path, map_location="cpu")
    model_dict  = style_extractor.state_dict()
    style_sd    = {k: v for k, v in style_sd.items()
                   if k in model_dict and model_dict[k].shape == v.shape}
    model_dict.update(style_sd)
    style_extractor.load_state_dict(model_dict)
    style_extractor.eval()

    # Build writer map from the bundled style_refs/ directory
    WRITER_ID_MAP = {d: i for i, d in enumerate(sorted(
        d for d in os.listdir(STYLE_REFS_DIR)
        if os.path.isdir(os.path.join(STYLE_REFS_DIR, d))
    ))}
    WRITER_LIST = sorted(WRITER_ID_MAP.keys())
    global _LUCKY_WRITERS
    _LUCKY_WRITERS = WRITER_LIST[:20]  # TODO: replace with curated top-20
    print(f"[demo]   {len(WRITER_LIST)} writers indexed from style_refs/")

    MODEL.update({
        "unet": unet, "vae": vae,
        "style_extractor": style_extractor,
        "tokenizer": tokenizer,
        "noise_scheduler": noise_scheduler,
        "num_classes": num_classes,
    })
    print("[demo] Ready.")


# ---------------------------------------------------------------------------
# Style ref loading from bundled style_refs/
# ---------------------------------------------------------------------------
_STYLE_TRANSFORM = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])


def _load_writer_style_ref(writer_str: str, device: torch.device) -> torch.Tensor:
    wdir   = os.path.join(STYLE_REFS_DIR, writer_str)
    fnames = sorted(f for f in os.listdir(wdir) if f.endswith(".png"))[:5]
    imgs   = []
    for fname in fnames:
        img = Image.open(os.path.join(wdir, fname)).convert("RGB")
        img = prepare_style_reference_image(img, IMG_HEIGHT, IMG_WIDTH)
        imgs.append(_STYLE_TRANSFORM(img))
    while len(imgs) < 5:
        imgs.append(imgs[0])
    return torch.stack(imgs).to(device)  # [5, 3, H, W]


# ---------------------------------------------------------------------------
# Logging helpers
# ---------------------------------------------------------------------------
def _log_generation(text, writer_str, writer_idx, cfg_scale, seed, duration_s, img_path, status):
    record = {
        "timestamp": datetime.now().isoformat(),
        "text": text, "writer_str": writer_str,
        "writer_idx": int(writer_idx), "cfg_scale": float(cfg_scale),
        "seed": int(seed), "duration_s": round(float(duration_s), 2),
        "output": img_path, "status": status,
    }
    with open(JSONL_LOG, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


def _load_recent_table(n: int = 10):
    if not os.path.exists(JSONL_LOG):
        return []
    rows = []
    with open(JSONL_LOG, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    rows.append(json.loads(line))
                except Exception:
                    pass
    table = []
    for r in reversed(rows[-n:]):
        ts = r.get("timestamp", "")[:19].replace("T", " ")
        table.append([
            ts, r.get("text", "")[:40], r.get("writer_str", ""),
            r.get("seed", ""), f"{r.get('duration_s', 0):.1f}s", r.get("status", ""),
        ])
    return table


def _make_diffusion_gif(frames, out_path, frame_ms=280, hold_ms=3000):
    if not frames:
        return None
    rgb = []
    for f in frames:
        pil = f if isinstance(f, Image.Image) else Image.fromarray(f)
        pil = pil.convert("RGB")
        pil = pil.resize((pil.width * 2, pil.height * 2), Image.NEAREST)
        rgb.append(pil)
    durations = [frame_ms] * len(rgb)
    durations[-1] = frame_ms + hold_ms
    rgb[0].save(out_path, save_all=True, append_images=rgb[1:],
                duration=durations, loop=0, optimize=False)
    return out_path


def _random_writer(lucky: bool = False):
    pool = _LUCKY_WRITERS if (lucky and _LUCKY_WRITERS) else WRITER_LIST
    return random.choice(pool) if pool else "Random"


def _lucky_fill():
    return random.choice(_LUCKY_PHRASES) if _LUCKY_PHRASES else ""


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------
@spaces.GPU
def generate(text_raw: str, seed_val: int, lucky: bool = False):
    if not MODEL:
        return None, "Models not loaded.", _load_recent_table(), None

    text = text_raw.strip()
    if not text:
        return None, "Please enter some text.", _load_recent_table(), None

    cfg_scale = 5.0
    writer_str = _random_writer(lucky=lucky)
    writer_idx = WRITER_ID_MAP[writer_str]

    seed = int(seed_val)
    if seed < 0:
        seed = random.randint(0, 2 ** 31 - 1)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    np.random.seed(seed % (2 ** 32))
    random.seed(seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    MODEL["unet"].to(device)
    MODEL["vae"].to(device)
    MODEL["style_extractor"].to(device)

    style_ref = _load_writer_style_ref(writer_str, device)

    # Pre-identify the longest alphabetic word for GIF capture
    gif_target = None
    best_len = 0
    for word in text.split():
        w = word
        while w and w[-1] in PUNCTUATION:
            w = w[:-1]
        if len(w) > best_len:
            best_len = len(w)
            gif_target = w
    gif_frames = []

    t0 = time.time()
    word_images, expanded_words, punct_flags, punct_standalone_flags = [], [], [], []

    for word in text.split():
        punct_suffix = []
        w = word
        while w and w[-1] in PUNCTUATION:
            punct_suffix.insert(0, w[-1])
            w = w[:-1]

        if w:
            for part, is_mark in split_word_for_generation(w, PUNCT_BANK_DIR, IMG_HEIGHT):
                if is_mark:
                    ch_arr = sample_punctuation(part, IMG_HEIGHT, PUNCT_BANK_DIR, writer_str)
                    word_images.append(ch_arr if ch_arr is not None else np.full((IMG_HEIGHT, 6), 255, dtype=np.uint8))
                    expanded_words.append(part)
                    punct_flags.append(True)
                    punct_standalone_flags.append(False)
                else:
                    capture = gif_frames if (part == gif_target and not gif_frames) else None
                    img_pil = generate_single_word(
                        word=part, unet=MODEL["unet"], vae=MODEL["vae"],
                        style_extractor=MODEL["style_extractor"],
                        tokenizer=MODEL["tokenizer"],
                        noise_scheduler=MODEL["noise_scheduler"],
                        style_ref=style_ref, writer_idx=writer_idx,
                        device=device, cfg_scale=cfg_scale,
                        img_height=IMG_HEIGHT, img_width=IMG_WIDTH,
                        text_max_len=TEXT_MAX_LEN,
                        intermediate_frames=capture,
                    )
                    word_images.append(crop_whitespace(img_pil))
                    expanded_words.append(part)
                    punct_flags.append(False)
                    punct_standalone_flags.append(False)

        for ch in punct_suffix:
            ch_arr = sample_punctuation(ch, IMG_HEIGHT, PUNCT_BANK_DIR, writer_str)
            word_images.append(ch_arr if ch_arr is not None else np.full((IMG_HEIGHT, 6), 255, dtype=np.uint8))
            expanded_words.append(ch)
            punct_flags.append(True)
            punct_standalone_flags.append(ch == '-' and not w)

    if not word_images:
        return None, "No words generated.", _load_recent_table(), None

    word_images = normalize_ink_brightness(word_images)

    paragraph, _ = compose_sentence_geometry(
        word_images=word_images,
        expanded_words=expanded_words,
        punct_flags=punct_flags,
        punct_standalone_flags=punct_standalone_flags,
        punct_bank=PUNCT_BANK_DIR,
        writer_str=writer_str,
        canvas_height=CANVAS_HEIGHT,
        anchor_long_words=True,
    )
    duration = time.time() - t0

    ts        = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_text = "".join(c if c.isalnum() or c in " _-" else "_" for c in text)[:40].strip().replace(" ", "_")
    img_fname = f"{ts}_{safe_text}_w{writer_str}_cfg{int(cfg_scale)}.png"
    img_path  = os.path.join(IMG_LOG_DIR, img_fname)
    paragraph.save(img_path)
    _log_generation(text, writer_str, writer_idx, cfg_scale, seed, duration, img_path, "ok")

    gif_path = None
    if gif_frames:
        gif_fname = f"{ts}_{safe_text}_w{writer_str}_diffusion.gif"
        gif_path = os.path.join(IMG_LOG_DIR, gif_fname)
        _make_diffusion_gif(gif_frames, gif_path)

    device_label = "cuda" if torch.cuda.is_available() else "cpu"
    info = f"Writer: {writer_str}  ·  Seed: {seed}  ·  {duration:.1f}s  ·  {device_label}"
    return paragraph, info, _load_recent_table(), gif_path


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
_CSS = """
#lucky-btn { background: #f97316 !important; border-color: #f97316 !important; color: #fff !important; }
#lucky-btn:hover { background: #ea6c08 !important; border-color: #ea6c08 !important; }
#btn-row { flex-wrap: nowrap !important; gap: 8px !important; padding: 0 !important; box-sizing: border-box !important; overflow: hidden !important; }
#btn-row > div { flex: 1 1 0 !important; min-width: 0 !important; box-sizing: border-box !important; overflow: hidden !important; }
"""

def build_ui():
    with gr.Blocks(title="DiffusionPen · Ukrainian Handwriting", css=_CSS) as demo:
        _not_lucky = gr.State(False)
        _is_lucky  = gr.State(True)
        gr.Markdown("# DiffusionPen · Ukrainian Handwriting Synthesis")
        gr.Markdown("308 real writers. 126 000 handwritten samples. One diffusion model. Type any Ukrainian word or sentence — a writer style is drawn at random and the text is generated word by word. [Read the paper](https://karl9doniz.github.io/ukr-diffusion-htg/)")

        output_img = gr.Image(label="Generated handwriting", type="pil", height=160)

        text_input = gr.Textbox(
            label="Ukrainian text",
            placeholder="Реве та стогне Дніпр широкий",
            lines=1,
            show_label=False,
        )

        with gr.Row(elem_id="btn-row"):
            lucky_btn = gr.Button("✦  I'm lucky", variant="primary", elem_id="lucky-btn")
            gen_btn   = gr.Button("✏  Generate",  variant="secondary")

        gif_output = gr.Image(
            label="Diffusion process — longest word",
            type="filepath",
            height=140,
        )

        info_box = gr.Textbox(label="", interactive=False, lines=1, show_label=False)

        with gr.Accordion("Advanced", open=False):
            seed_input = gr.Number(
                value=-1,
                label="Seed  (−1 = random each run · fix a value to reproduce exact output)",
                precision=0,
            )

        gr.Markdown("### Recent generations")
        log_table = gr.Dataframe(
            headers=["Timestamp", "Text", "Writer", "Seed", "Time", "Status"],
            datatype=["str", "str", "str", "number", "str", "str"],
            value=_load_recent_table(),
            interactive=False, wrap=True,
        )

        _outs = [output_img, info_box, log_table, gif_output]

        lucky_btn.click(fn=_lucky_fill, outputs=[text_input]).then(
            fn=generate, inputs=[text_input, seed_input, _is_lucky], outputs=_outs,
        )
        gen_btn.click(
            fn=generate, inputs=[text_input, seed_input, _not_lucky], outputs=_outs,
        )
        text_input.submit(
            fn=generate, inputs=[text_input, seed_input, _not_lucky], outputs=_outs,
        )

    return demo


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
os.makedirs(IMG_LOG_DIR, exist_ok=True)
load_models()
demo = build_ui()
demo.queue(max_size=20)

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