EdiTikZ-4B-RL

Model Card for EdiTikZ-4B-RL

EdiTikZ-4B-RL is a multimodal language model for instruction-guided editing and reconstruction of scientific figures represented as TikZ/LaTeX code.

It is based on EdiTikZ-4B and was trained with multi-reward reinforcement learning (RL) on DaEdiTikZ for scientific figure editing.

Although the RL stage was applied only to the editing task, EdiTikZ-4B-RL retains the reconstruction capability learned during supervised fine-tuning.

Installation

pip install torch==2.7.1 transformers==5.3.0 accelerate==1.12.0 pillow

Usage

import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

model_id = "nllg/EdiTikZ-4B-RL"
fixed_image_size = 448
fixed_pixels = fixed_image_size * fixed_image_size

processor = AutoProcessor.from_pretrained(
    model_id,
    use_fast=False,
    min_pixels=fixed_pixels,
    max_pixels=fixed_pixels,
    trust_remote_code=True,
)

model = AutoModelForImageTextToText.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
model.eval()

tokenizer = processor.tokenizer
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

model.config.eos_token_id = tokenizer.eos_token_id
model.config.pad_token_id = tokenizer.pad_token_id
if getattr(model, "generation_config", None) is not None:
    model.generation_config.eos_token_id = tokenizer.eos_token_id
    model.generation_config.pad_token_id = tokenizer.pad_token_id

def build_editing_prompt(instruction: str) -> str:
    return (
        "<TASK_EDITING>\n"
        "This is an image of a scientific figure. Reconstruct it in TikZ and "
        "apply the following changes:\n"
        f"{instruction.strip()}\n"
        "Wrap your code using \\documentclass[tikz]{standalone}, and include "
        "\\begin{document} ... \\end{document}. Only output valid LaTeX code "
        "with no extra text."
    )

def build_reconstruction_prompt() -> str:
    return (
        "<TASK_RECONSTRUCTION>\n"
        "This is an image of a scientific figure. Reconstruct it in TikZ.\n"
        "Wrap your code using \\documentclass[tikz]{standalone}, and include "
        "\\begin{document} ... \\end{document}. Only output valid LaTeX code "
        "with no extra text."
    )

def load_image(image_path: str) -> Image.Image:
    image = Image.open(image_path).convert("RGB")
    if image.size != (fixed_image_size, fixed_image_size):
        image = image.resize(
            (fixed_image_size, fixed_image_size),
            Image.Resampling.BICUBIC,
        )
    return image

@torch.inference_mode()
def generate_tikz(image_path: str, prompt: str) -> str:
    image = load_image(image_path)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    inputs = processor(
        text=[text],
        images=[image],
        padding=False,
        return_tensors="pt",
    )
    model_device = next(model.parameters()).device
    inputs = {name: value.to(model_device) for name, value in inputs.items()}

    generated_ids = model.generate(
        **inputs,
        max_new_tokens=2048,
        do_sample=True,
        temperature=0.2,
        top_p=0.9,
        top_k=50,
        repetition_penalty=1.0,
        use_cache=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
    )
    generated_ids = generated_ids[:, inputs["input_ids"].shape[1]:]
    return processor.batch_decode(
        generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0].strip()

image_path = "path/to/source_figure.png"
editing_instruction = (
    "The thick black polyline consisting of a horizontal segment along the "
    "bottom edge and a diagonal segment connecting the bottom-middle to the "
    "center of the grid is removed. A thick black outline with rounded corners "
    "is added around the entire 3x2 grid, replacing the sharp corners of the "
    "original rectangular boundary."
)

# Instruction-guided editing
edited_tikz = generate_tikz(
    image_path,
    build_editing_prompt(editing_instruction),
)
print(edited_tikz)

# Image-to-TikZ reconstruction
# reconstructed_tikz = generate_tikz(image_path, build_reconstruction_prompt())
# print(reconstructed_tikz)

Faster batched inference with vLLM

pip install vllm==0.25.0 pillow

from PIL import Image
from transformers import AutoProcessor
from vllm import LLM, SamplingParams

model_id = "nllg/EdiTikZ-4B-RL"
fixed_image_size = 448
fixed_pixels = fixed_image_size * fixed_image_size

processor = AutoProcessor.from_pretrained(
    model_id,
    use_fast=False,
    min_pixels=fixed_pixels,
    max_pixels=fixed_pixels,
    trust_remote_code=True,
)

llm = LLM(
    model=model_id,
    dtype="bfloat16",
    trust_remote_code=True,
    limit_mm_per_prompt={"image": 1},
    # tensor_parallel_size=2,  # Uncomment for multi-GPU tensor parallelism.
)

def make_request(image_path: str, instruction: str) -> dict:
    image = Image.open(image_path).convert("RGB")
    if image.size != (fixed_image_size, fixed_image_size):
        image = image.resize(
            (fixed_image_size, fixed_image_size),
            Image.Resampling.BICUBIC,
        )

    prompt = (
        "<TASK_EDITING>\n"
        "This is an image of a scientific figure. Reconstruct it in TikZ and "
        "apply the following changes:\n"
        f"{instruction.strip()}\n"
        "Wrap your code using \\documentclass[tikz]{standalone}, and include "
        "\\begin{document} ... \\end{document}. Only output valid LaTeX code "
        "with no extra text."
    )
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    formatted_prompt = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    return {
        "prompt": formatted_prompt,
        "multi_modal_data": {"image": image},
    }

examples = [
    {
        "image_path": "path/to/source_figure_1.png",
        "instruction": "first_editing_instruction",
    },
    {
        "image_path": "path/to/source_figure_2.png",
        "instruction": "second_editing_instruction",
    },
]

requests = [
    make_request(item["image_path"], item["instruction"])
    for item in examples
]

sampling_params = SamplingParams(
    max_tokens=2048,
    temperature=0.2,
    top_p=0.9,
    top_k=50,
    repetition_penalty=1.0,
)

outputs = llm.generate(
    requests,
    sampling_params=sampling_params,
    use_tqdm=True,
)

for output in outputs:
    print(output.outputs[0].text.strip())