init

app.py

@@ -0,0 +1,25 @@
+from os import getenv
+from textwrap import dedent
+
+import gradio as gr
+from torch import cuda
+
+from webui import build_ui, remove_darkness, get_banner
+
+PUBLIC_DEMO = getenv("SPACE_ID") == "waleko/TikZ-Assistant"
+
+if PUBLIC_DEMO and not cuda.is_available():
+    center = ".gradio-container {text-align: center}"
+    with gr.Blocks(css=center, theme=remove_darkness(gr.themes.Soft()), title="AutomaTikZ") as demo:
+        badge = "https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-xl.svg"
+        link = "https://huggingface.co/spaces/nllg/AutomaTikZ?duplicate=true"
+        html = f'<a style="display:inline-block" href="{link}"> <img src="{badge}" alt="Duplicate this Space"> </a>'
+        message = dedent("""\
+        The size of our models exceeds the resource constraints offered by the
+        free tier of Hugging Face Spaces. For full functionality, we recommend
+        duplicating this space on a paid private GPU runtime.
+        """)
+        gr.Markdown(f'{get_banner()}\n{message}\n{html}')
+        demo.launch()
+else:
+    build_ui(lock=False, force_light=True).queue().launch(server_name="0.0.0.0", server_port=7860)

infer.py

@@ -0,0 +1,243 @@
+from collections import namedtuple
+from functools import cache, cached_property
+from io import BytesIO
+from os import environ
+from os.path import isfile, join
+from re import MULTILINE, escape, search, sub
+from subprocess import CalledProcessError, DEVNULL, TimeoutExpired
+from tempfile import NamedTemporaryFile, TemporaryDirectory
+from typing import Optional, Union
+import warnings
+
+from PIL import Image, ImageOps
+import requests
+import torch
+from torch.cuda import current_device, is_available as has_cuda
+from transformers import TextGenerationPipeline as TGP, TextStreamer, pipeline, ImageToTextPipeline as ITP
+from transformers.utils import logging
+from transformers.utils.hub import is_remote_url
+
+from pdf2image.pdf2image import convert_from_bytes
+from pdfCropMargins import crop
+import fitz
+
+logger = logging.get_logger("transformers")
+
+from os import killpg, getpgid
+from subprocess import Popen, TimeoutExpired, CalledProcessError, CompletedProcess, PIPE
+from signal import SIGKILL
+
+def run(*popenargs, input=None, timeout=None, check=False, **kwargs):
+    with Popen(*popenargs, start_new_session=True, **kwargs) as process:
+        try:
+            stdout, stderr = process.communicate(input, timeout=timeout)
+        except TimeoutExpired:
+            killpg(getpgid(process.pid), SIGKILL)
+            process.wait()
+            raise
+        except:
+            killpg(getpgid(process.pid), SIGKILL)
+            raise
+        retcode = process.poll()
+        if check and retcode:
+            raise CalledProcessError(retcode, process.args,
+                                     output=stdout, stderr=stderr)
+    return CompletedProcess(process.args, retcode, stdout, stderr) # type: ignore
+
+def check_output(*popenargs, timeout=None, **kwargs):
+    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True, **kwargs).stdout
+
+class PdfDocument:
+    def __init__(self, raw: bytes):
+        self.raw = raw
+
+    def save(self, filename):
+        with open(filename, "wb") as f:
+            f.write(self.raw)
+
+
+class TikzDocument:
+    """
+    Faciliate some operations with TikZ code. To compile the images a full
+    TeXLive installation is assumed to be on the PATH. Cropping additionally
+    requires Ghostscript, and rasterization needs poppler (apart from the 'pdf'
+    optional dependencies).
+    """
+     # engines to try, could also try: https://tex.stackexchange.com/a/495999
+    engines = ["pdflatex", "lualatex", "xelatex"]
+    Output = namedtuple("Output", ['pdf', 'status', 'log'], defaults=[None, -1, ""])
+
+    def __init__(self, code: str, timeout=120):
+        self.code = code
+        self.timeout = timeout
+
+    @property
+    def status(self) -> int:
+        return self.compile().status
+
+    @property
+    def pdf(self) -> Optional[PdfDocument]:
+        return self.compile().pdf
+
+    @property
+    def log(self) -> str:
+        return self.compile().log
+
+    @property
+    def compiled_with_errors(self) -> bool:
+        return self.status != 0
+
+    @cached_property
+    def has_content(self) -> bool:
+        """true if we have an image that isn't empty"""
+        return (img:=self.rasterize()) is not None and img.getcolors(1) is None
+
+    @classmethod
+    def set_engines(cls, engines: Union[str, list]):
+        cls.engines = [engines] if isinstance(engines, str) else engines
+
+    @cache
+    def compile(self) -> "Output":
+        output = dict()
+        with TemporaryDirectory() as tmpdirname:
+            with NamedTemporaryFile(dir=tmpdirname, buffering=0) as tmpfile:
+                codelines = self.code.split("\n")
+                # make sure we don't have page numbers in compiled pdf (for cropping)
+                codelines.insert(1, r"{cmd}\AtBeginDocument{{{cmd}}}".format(cmd=r"\thispagestyle{empty}\pagestyle{empty}"))
+                tmpfile.write("\n".join(codelines).encode())
+
+                try:
+                    # compile
+                    errorln, tmppdf, outpdf = 0, f"{tmpfile.name}.pdf", join(tmpdirname, "tikz.pdf")
+                    open(f"{tmpfile.name}.bbl", 'a').close() # some classes expect a bibfile
+
+                    def try_save_last_page():
+                        try:
+                            doc = fitz.open(tmppdf) # type: ignore
+                            doc.select([len(doc)-1])
+                            doc.save(outpdf)
+                        except:
+                            pass
+
+                    for engine in self.engines:
+                        try:
+                            check_output(
+                                cwd=tmpdirname,
+                                timeout=self.timeout,
+                                stderr=DEVNULL,
+                                env=environ | dict(max_print_line="1000"), # improve formatting of log
+                                args=["latexmk", "-f", "-nobibtex", "-norc", "-file-line-error", "-interaction=nonstopmode", f"-{engine}", tmpfile.name]
+                            )
+                        except (CalledProcessError, TimeoutExpired) as proc:
+                            log = getattr(proc, "output", b'').decode(errors="ignore")
+                            error = search(rf'^{escape(tmpfile.name)}:(\d+):.+$', log, MULTILINE)
+                            # only update status and log if first error occurs later than in previous engine
+                            if (linenr:=int(error.group(1)) if error else 0) > errorln:
+                                errorln = linenr
+                                output.update(status=getattr(proc, 'returncode', -1), log=log)
+                                try_save_last_page()
+                        else:
+                            output.update(status=0, log='')
+                            try_save_last_page()
+                            break
+
+                    # crop
+                    croppdf = f"{tmpfile.name}.crop"
+                    crop(["-gsf", "-c", "gb", "-p", "0", "-a", "-1", "-o", croppdf, outpdf], quiet=True)
+                    if isfile(croppdf):
+                        with open(croppdf, "rb") as pdf:
+                            output['pdf'] = PdfDocument(pdf.read())
+
+                except (FileNotFoundError, NameError) as e:
+                    logger.error("Missing dependencies: " + (
+                        "Install this project with the [pdf] feature name!" if isinstance(e, NameError)
+                        else "Did you install TeX Live?"
+                    ))
+                except RuntimeError: # pdf error during cropping
+                    pass
+
+        if output.get("status") == 0 and not output.get("pdf", None):
+            logger.warning("Could compile document but something seems to have gone wrong during cropping!")
+
+        return self.Output(**output)
+
+    def rasterize(self, size=336, expand_to_square=True) -> Optional[Image.Image]:
+        if self.pdf:
+            image = convert_from_bytes(self.pdf.raw, size=size, single_file=True)[0]
+            if expand_to_square:
+                image = ImageOps.pad(image, (size, size), color='white')
+
+            return image
+
+    def save(self, filename: str, *args, **kwargs):
+        match filename.split(".")[-1]:
+            case "tex": content = self.code.encode()
+            case "pdf": content = getattr(self.pdf, "raw", bytes())
+            case fmt if img := self.rasterize(*args, **kwargs):
+                img.save(imgByteArr:=BytesIO(), format=fmt)
+                content = imgByteArr.getvalue()
+            case fmt: raise ValueError(f"Couldn't save with format '{fmt}'!")
+
+        with open(filename, "wb") as f:
+            f.write(content)
+
+
+class TikzGenerator:
+    def __init__(
+        self,
+        pipe: ITP,
+        temperature: float = 0.8, # based on "a systematic evaluation of large language models of code"
+        top_p: float = 0.95,
+        top_k: int = 0,
+        stream: bool = False,
+        expand_to_square: bool = False,
+        clean_up_output: bool = True,
+    ):
+        self.expand_to_square = expand_to_square
+        self.clean_up_output = clean_up_output
+        self.pipeline = pipe
+        self.pipeline.model = torch.compile(model) # type: ignore
+
+        self.default_kwargs = dict(
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            num_return_sequences=1,
+            max_length=self.pipeline.tokenizer.model_max_length, # type: ignore
+            do_sample=True,
+            return_full_text=False,
+            streamer=TextStreamer(self.pipeline.tokenizer, # type: ignore
+                skip_prompt=True,
+                skip_special_tokens=True
+            ),
+        )
+
+        if not stream:
+            self.default_kwargs.pop("streamer")
+
+    def generate(self, image: Image.Image, **generate_kwargs):
+        prompt = "Assistant helps to write down the TikZ code for the user's image. USER: <image>\nWrite down the TikZ code to draw the diagram shown in the lol. ASSISTANT:"
+        tokenizer = self.pipeline.tokenizer
+        text = self.pipeline(image, prompt=prompt, generate_kwargs=(self.default_kwargs | generate_kwargs))[0]["generated_text"] # type: ignore
+
+        if self.clean_up_output:
+            for token in reversed(tokenizer.tokenize(prompt)): # type: ignore
+                # remove leading characters because skip_special_tokens in pipeline
+                # adds unwanted prefix spaces if prompt ends with a special tokens
+                if text and text[0].isspace() and token in tokenizer.all_special_tokens: # type: ignore
+                    text = text[1:]
+                else:
+                    break
+
+            # occasionally observed artifacts
+            artifacts = {
+                r'\bamsop\b': 'amsopn'
+            }
+            for artifact, replacement in artifacts.items():
+                text = sub(artifact, replacement, text) # type: ignore
+
+        return text
+
+
+    def __call__(self, *args, **kwargs):
+        return self.generate(*args, **kwargs)

packages.txt

@@ -0,0 +1,3 @@
+texlive-full
+ghostscript
+poppler-utils

requirements.txt

@@ -0,0 +1,7 @@
+torch<2.1
+pdfCropMargins~=2.0
+pdf2image~=1.16
+PyMuPDF~=1.22
+peft>=0.2.0
+transformers
+gradio

webui.py

@@ -0,0 +1,198 @@
+
+#!/usr/bin/env python
+
+from argparse import ArgumentParser
+from functools import lru_cache
+from importlib.resources import files
+from inspect import signature
+from multiprocessing.pool import ThreadPool
+from tempfile import NamedTemporaryFile
+from textwrap import dedent
+from typing import Optional
+
+from PIL import Image
+import fitz
+import gradio as gr
+from transformers import TextIteratorStreamer, pipeline, ImageToTextPipeline
+
+from infer import TikzDocument, TikzGenerator
+
+# assets = files(__package__) / "assets" if __package__ else files("assets") / "."
+models = {
+    "Fine-tuned Llava": "waleko/TikZ-llava-1.5-7b"
+}
+
+
+@lru_cache(maxsize=1)
+def cached_load(model_name, **kwargs) -> ImageToTextPipeline:
+    gr.Info("Instantiating model. Could take a while...") # type: ignore
+    # noinspection PyTypeChecker
+    return pipeline("image-to-text", model=model_name, **kwargs)
+
+
+def convert_to_svg(pdf):
+    doc = fitz.open("pdf", pdf.raw) # type: ignore
+    return doc[0].get_svg_image()
+
+
+def inference(
+    model_name: str,
+    image: Image.Image,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    expand_to_square: bool,
+):
+    generate = TikzGenerator(
+        cached_load(model_name, device_map="auto"),
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        expand_to_square=expand_to_square,
+    )
+    streamer = TextIteratorStreamer(
+        generate.pipeline.tokenizer, # type: ignore
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    thread = ThreadPool(processes=1)
+    async_result = thread.apply_async(generate, kwds=dict(image=image, streamer=streamer))
+
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text, None, False
+    yield async_result.get().code, None, True
+
+def tex_compile(
+    code: str,
+    timeout: int,
+    rasterize: bool
+):
+    tikzdoc = TikzDocument(code, timeout=timeout)
+    if not tikzdoc.has_content:
+        if tikzdoc.compiled_with_errors:
+            raise gr.Error("TikZ code did not compile!") # type: ignore
+        else:
+            gr.Warning("TikZ code compiled to an empty image!") # type: ignore
+    elif tikzdoc.compiled_with_errors:
+        gr.Warning("TikZ code compiled with errors!") # type: ignore
+
+    if rasterize:
+        yield tikzdoc.rasterize()
+    else:
+        with NamedTemporaryFile(suffix=".svg", buffering=0) as tmpfile:
+            if pdf:=tikzdoc.pdf:
+                tmpfile.write(convert_to_svg(pdf).encode())
+            yield tmpfile.name if pdf else None
+
+def check_inputs(image: Image.Image):
+    if not image:
+        raise gr.Error("Image is required")
+
+def get_banner():
+    return dedent('''\
+    # AutomaTi*k*Z: Text-Guided Synthesis of Scientific Vector Graphics with Ti*k*Z
+
+    <p>
+      <a style="display:inline-block" href="https://github.com/potamides/AutomaTikZ">
+        <img src="https://img.shields.io/badge/View%20on%20GitHub-green?logo=github&labelColor=gray" alt="View on GitHub">
+      </a>
+      <a style="display:inline-block" href="https://arxiv.org/abs/2310.00367">
+        <img src="https://img.shields.io/badge/View%20on%20arXiv-B31B1B?logo=arxiv&labelColor=gray" alt="View on arXiv">
+      </a>
+      <a style="display:inline-block" href="https://colab.research.google.com/drive/14S22x_8VohMr9pbnlkB4FqtF4n81khIh">
+        <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab">
+      </a>
+      <a style="display:inline-block" href="https://huggingface.co/spaces/nllg/AutomaTikZ">
+        <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/open-in-hf-spaces-sm.svg" alt="Open in HF Spaces">
+      </a>
+    </p>
+    ''')
+
+def remove_darkness(stylable):
+    """
+    Patch gradio to only contain light mode colors.
+    """
+    if isinstance(stylable, gr.themes.Base): # remove dark variants from the entire theme
+        params = signature(stylable.set).parameters
+        colors = {color: getattr(stylable, color.removesuffix("_dark")) for color in dir(stylable) if color in params}
+        return stylable.set(**colors)
+    elif isinstance(stylable, gr.Blocks): # also handle components which do not use the theme (e.g. modals)
+        stylable.load(_js="() => document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'))")
+        return stylable
+    else:
+        raise ValueError
+
+def build_ui(model=list(models)[0], lock=False, rasterize=False, force_light=False, lock_reason="locked", timeout=120):
+    theme = remove_darkness(gr.themes.Soft()) if force_light else gr.themes.Soft()
+    with gr.Blocks(theme=theme, title="AutomaTikZ") as demo: # type: ignore
+        if force_light: remove_darkness(demo)
+        gr.Markdown(get_banner())
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                info = (
+                    "Describe what you want to generate. "
+                    "Scientific graphics benefit from captions with at least 30 tokens (see examples below), "
+                    "while simple objects work best with shorter, 2-3 word captions."
+                )
+                # caption = gr.Textbox(label="Caption", info=info, placeholder="Type a caption...")
+                image = gr.Image(label="Image Input", type="pil", info=info)
+                label = "Model" + (f" ({lock_reason})" if lock else "")
+                model = gr.Dropdown(label=label, choices=list(models.items()), value=models[model], interactive=not lock) # type: ignore
+                with gr.Accordion(label="Advanced Options", open=False):
+                    temperature = gr.Slider(minimum=0, maximum=2, step=0.05, value=0.8, label="Temperature")
+                    top_p = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.95, label="Top-P")
+                    top_k = gr.Slider(minimum=0, maximum=100, step=10, value=0, label="Top-K")
+                    expand_to_square = gr.Checkbox(value=True, label="Expand image to square")
+                with gr.Row():
+                    run_btn = gr.Button("Run", variant="primary")
+                    stop_btn = gr.Button("Stop")
+                    clear_btn = gr.ClearButton([image])
+            with gr.Column():
+                with gr.Tabs() as tabs:
+                    with gr.TabItem(label:="TikZ Code", id=0):
+                        info = "Source code of the generated image."
+                        tikz_code = gr.Code(label=label, show_label=False, info=info, interactive=False)
+                    with gr.TabItem(label:="Compiled Image", id=1):
+                        result_image = gr.Image(label=label, show_label=False, show_share_button=rasterize)
+                    clear_btn.add([tikz_code, result_image])
+        # TODO: gr.Examples(examples=str(assets), inputs=[image, tikz_code, result_image])
+
+        events = list()
+        finished = gr.Textbox(visible=False) # hack to cancel compile on canceled inference
+        for listener in [run_btn.click]:
+            generate_event = listener(
+                check_inputs,
+                inputs=[image],
+                queue=False
+            ).success(
+                lambda: gr.Tabs(selected=0),
+                outputs=tabs, # type: ignore
+                queue=False
+            ).then(
+                inference,
+                inputs=[model, image, temperature, top_p, top_k, expand_to_square],
+                outputs=[tikz_code, result_image, finished]
+            )
+
+            def tex_compile_if_finished(finished, *args):
+                yield from (tex_compile(*args, timeout=timeout, rasterize=rasterize) if finished == "True" else [])
+
+            compile_event = generate_event.then(
+                lambda finished: gr.Tabs(selected=1) if finished == "True" else gr.Tabs(),
+                inputs=finished,
+                outputs=tabs, # type: ignore
+                queue=False
+            ).then(
+                tex_compile_if_finished,
+                inputs=[finished, tikz_code],
+                outputs=result_image
+            )
+            events.extend([generate_event, compile_event])
+
+        # model.select(lambda model_name: gr.Image(visible="clima" in model_name), inputs=model, outputs=image, queue=False)
+        for btn in [clear_btn, stop_btn]:
+            btn.click(fn=None, cancels=events, queue=False)
+        return demo