Model Card for EdiTikZ-4B
EdiTikZ-4B is a multimodal language model for instruction-guided editing and reconstruction of scientific figures represented as TikZ/LaTeX code.
It is based on
Qwen3.5-4B
and was trained with multi-task
supervised fine-tuning (SFT)
on
DaEdiTikZ
and
DaTikZ-V4
for scientific figure editing and reconstruction.
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"
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"
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())
