+-------------------+         +---------------------------+
|   Diagram Image   +-------> |  Fine-Tuned VLM (Qwen2.5) |
+-------------------+         +-------------+-------------+
                                            |
                                            v
                           +------------------------------+
                           |  JSON (Nodes + Edges + Meta) |
                           +------------------------------+
                                            |
                                            v
                              +------------------------+
                              |     Neo4J Integration  |
                              +------------------------+

pip install -q transformers accelerate qwen-vl-utils[decord]==0.0.8

import torch
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor


MODEL_ID="zackriya/diagram2graph-adapters"
MAX_PIXELS = 1280 * 28 * 28
MIN_PIXELS = 256 * 28 * 28


model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    device_map="auto",
    torch_dtype=torch.bfloat16
)

processor = Qwen2_5_VLProcessor.from_pretrained(
    MODEL_ID,
    min_pixels=MIN_PIXELS,
    max_pixels=MAX_PIXELS
)

from qwen_vl_utils import process_vision_info

SYSTEM_MESSAGE = """You are a Vision Language Model specialized in extracting structured data from visual representations of process and flow diagrams.
Your task is to analyze the provided image of a diagram and extract the relevant information into a well-structured JSON format.
The diagram includes details such as nodes and edges. each of them have their own attributes.
Focus on identifying key data fields and ensuring the output adheres to the requested JSON structure.
Provide only the JSON output based on the extracted information. Avoid additional explanations or comments."""

def run_inference(image):
  """
  Inference with the Model
  """
  messages= [
      {
          "role": "system",
          "content": [{"type": "text", "text": SYSTEM_MESSAGE}],
      },
      {
          "role": "user",
          "content": [
              {
                  "type": "image",
                  # this image is handled by qwen_vl_utils's process_visio_Info so no need to worry about pil image or path
                  "image": image,
              },
              {
                  "type": "text",
                  "text": "Extract data in JSON format, Only give the JSON",
              },
          ],
      },
  ]

  text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
  image_inputs, _ = process_vision_info(messages)

  inputs = processor(
      text=[text],
      images=image_inputs,
      return_tensors="pt",
  )
  inputs = inputs.to('cuda')

  generated_ids = model.generate(**inputs, max_new_tokens=1024)
  generated_ids_trimmed = [
      out_ids[len(in_ids):]
      for in_ids, out_ids
      in zip(inputs.input_ids, generated_ids)
  ]

  output_text = processor.batch_decode(
      generated_ids_trimmed,
      skip_special_tokens=True,
      clean_up_tokenization_spaces=False
  )
  return output_text