117 lines
3.2 KiB
Python
117 lines
3.2 KiB
Python
from datetime import datetime
|
|
|
|
from common.config_loader import loadVisionConfig
|
|
from common.runtime_log import appendRuntimeLog
|
|
from screen.screen_actions import captureOcrRegion
|
|
|
|
_vision_pipeline = None
|
|
|
|
|
|
def log(level, message):
|
|
now = datetime.now().strftime("%H:%M:%S")
|
|
print(f"[{level}] {now} {message}", flush=True)
|
|
|
|
|
|
def getVisionPipeline():
|
|
global _vision_pipeline
|
|
|
|
if _vision_pipeline is not None:
|
|
return _vision_pipeline
|
|
|
|
config = loadVisionConfig()
|
|
model_name = config["model_name"]
|
|
|
|
if not model_name:
|
|
raise RuntimeError("vision.model_name is not set")
|
|
|
|
try:
|
|
from transformers import AutoProcessor
|
|
from transformers import Gemma4ForConditionalGeneration
|
|
import torch
|
|
except Exception as e:
|
|
raise RuntimeError(f"vision dependencies are missing detail={e}")
|
|
|
|
log("INFO", f"Vision model initializing model={model_name}")
|
|
processor = AutoProcessor.from_pretrained(model_name)
|
|
model = Gemma4ForConditionalGeneration.from_pretrained(
|
|
model_name,
|
|
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
|
|
device_map="auto",
|
|
).eval()
|
|
_vision_pipeline = (processor, model)
|
|
return _vision_pipeline
|
|
|
|
|
|
def saveVisionText(image_path, text):
|
|
return appendRuntimeLog(
|
|
"VISION TEXT",
|
|
f"image={image_path}\n{text}",
|
|
)
|
|
|
|
|
|
def runVisionFromImage(image_path):
|
|
config = loadVisionConfig()
|
|
processor, model = getVisionPipeline()
|
|
|
|
prompt = config["prompt"]
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": [
|
|
{"type": "text", "text": "You are a helpful vision-language assistant."},
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image"},
|
|
{"type": "text", "text": prompt},
|
|
],
|
|
}
|
|
]
|
|
|
|
try:
|
|
import torch
|
|
from PIL import Image
|
|
except Exception as e:
|
|
log("ERROR", f"Vision runtime imports failed detail={e}")
|
|
return None
|
|
|
|
try:
|
|
image = Image.open(image_path).convert("RGB")
|
|
inputs = processor.apply_chat_template(
|
|
messages,
|
|
add_generation_prompt=True,
|
|
tokenize=True,
|
|
return_dict=True,
|
|
return_tensors="pt",
|
|
)
|
|
inputs = inputs.to(model.device, dtype=model.dtype)
|
|
with torch.no_grad():
|
|
output = model.generate(
|
|
**inputs,
|
|
max_new_tokens=config["max_new_tokens"],
|
|
do_sample=config["temperature"] > 0,
|
|
)
|
|
input_len = inputs["input_ids"].shape[-1]
|
|
generated = output[0][input_len:]
|
|
text = processor.decode(generated, skip_special_tokens=True).strip()
|
|
except Exception as e:
|
|
log("ERROR", f"Vision inference failed detail={e}")
|
|
return None
|
|
|
|
text_path = saveVisionText(image_path, text)
|
|
log("INFO", f"Vision text appended={text_path}")
|
|
print(text, flush=True)
|
|
return {
|
|
"image_path": image_path,
|
|
"text_path": text_path,
|
|
"text": text,
|
|
}
|
|
|
|
|
|
def runVisionFromScreen():
|
|
log("ACTION", "runVisionFromScreen called")
|
|
image_path = captureOcrRegion()
|
|
return runVisionFromImage(image_path)
|