Spaces:
Running on Zero
Running on Zero
File size: 5,120 Bytes
f7e32a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """Vision engine — image analysis, OCR, captioning via ZeroGPU."""
import logging
from typing import Optional
from config import config
from models.zerogpu import requires_gpu
logger = logging.getLogger("synapse.vision")
class VisionEngine:
"""Image understanding and document analysis."""
def __init__(self):
self._florence_model = None
self._qwen_model = None
def _load_florence(self):
if self._florence_model is not None:
return
try:
import torch
from transformers import AutoModelForCausalLM, AutoProcessor
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
model_id = config.vision.florence_model
self._florence_model = {
"model": AutoModelForCausalLM.from_pretrained(
model_id, torch_dtype=dtype, trust_remote_code=True,
).to(device),
"processor": AutoProcessor.from_pretrained(
model_id, trust_remote_code=True,
),
"device": device,
}
logger.info("Florence model loaded")
except Exception as e:
logger.warning(f"Failed to load Florence: {e}")
@requires_gpu("vision")
def describe_image(self, image_path: str,
task: str = "caption") -> str:
"""Analyze an image with Florence-2."""
self._load_florence()
if self._florence_model is None:
return self._describe_fallback(image_path, task)
try:
from PIL import Image
import torch
image = Image.open(image_path).convert("RGB")
model = self._florence_model["model"]
processor = self._florence_model["processor"]
device = self._florence_model["device"]
task_prompts = {
"caption": "<CAPTION>",
"detailed": "<DETAILED_CAPTION>",
"ocr": "<OCR>",
"objects": "<OD>",
"regions": "<REGION_PROPOSAL>",
}
prompt = task_prompts.get(task, "<CAPTION>")
inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
with torch.no_grad():
generated_ids = model.generate(
**inputs, max_new_tokens=256, num_beams=3,
)
result = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
return result.strip()
except Exception as e:
logger.error(f"Florence vision error: {e}")
return self._describe_fallback(image_path, task)
def _describe_fallback(self, image_path: str, task: str) -> str:
try:
from PIL import Image
img = Image.open(image_path)
size = img.size
mode = img.mode
return f"Image: {size[0]}x{size[1]}, mode={mode}. Detailed analysis unavailable — vision model not loaded."
except Exception:
return "Image analysis unavailable."
def caption(self, image_path: str) -> str:
return self.describe_image(image_path, "caption")
def detailed_caption(self, image_path: str) -> str:
return self.describe_image(image_path, "detailed")
def ocr(self, image_path: str) -> str:
return self.describe_image(image_path, "ocr")
def detect_objects(self, image_path: str) -> str:
return self.describe_image(image_path, "objects")
def analyze_for_question(self, image_path: str, question: str) -> str:
self._load_florence()
if self._florence_model is None:
return f"Cannot answer '{question}' — vision model not loaded."
try:
from PIL import Image
import torch
image = Image.open(image_path).convert("RGB")
model = self._florence_model["model"]
processor = self._florence_model["processor"]
device = self._florence_model["device"]
inputs = processor(text=question, images=image, return_tensors="pt").to(device)
with torch.no_grad():
generated_ids = model.generate(
**inputs, max_new_tokens=256, num_beams=3,
)
result = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
return result.strip()
except Exception as e:
return f"Error analyzing image: {e}"
def get_tasks(self) -> list[dict]:
return [
{"id": "caption", "name": "Caption", "description": "Brief image description"},
{"id": "detailed", "name": "Detailed Caption", "description": "Detailed description"},
{"id": "ocr", "name": "OCR", "description": "Extract text from image"},
{"id": "objects", "name": "Object Detection", "description": "Detect objects"},
{"id": "regions", "name": "Regions", "description": "Detect regions"},
]
vision_engine = VisionEngine()
|