Spaces:
Running on Zero
Running on Zero
File size: 7,170 Bytes
88fe4e7 c92596a 88fe4e7 8417e5e 88fe4e7 8417e5e c92596a 4cd2f1b c92596a 88fe4e7 210dd01 8417e5e 88fe4e7 93a00cf 88fe4e7 c92596a 88fe4e7 8417e5e 88fe4e7 8417e5e 88fe4e7 | 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 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | import torch
from transformers import AutoModelForCausalLM, AutoProcessor
from PIL import Image
from typing import Optional, Dict, Any
from src.vision.vision_engine import VisionEngine
from src.config import CONFIG
from src.vision.utils import preprocess_image, auto_enhance
from src.vision.captioning import format_caption
from src.vision.ocr import format_ocr
from src.vision.detection import format_object_detection
try:
import spaces
IS_SPACES = True
gpu_decorator = spaces.GPU
except ImportError:
IS_SPACES = False
def gpu_decorator(func):
return func
def get_device() -> str:
"""Select best available device."""
if torch.cuda.is_available():
return "cuda"
elif torch.backends.mps.is_available():
return "mps"
return "cpu"
DEVICE: str = "cpu" if IS_SPACES else get_device()
DTYPE: torch.dtype = torch.float16 if IS_SPACES or DEVICE == "cuda" else torch.float32
class FlorenceVisionEngine(VisionEngine):
def __init__(self):
self.model: Optional[AutoModelForCausalLM] = None
self.processor: Optional[AutoProcessor] = None
self.paddle_ocr = None
def load(self):
"""Load the Florence-2 model."""
if self.model is not None:
return
try:
print(f"Loading Florence-2 on {DEVICE.upper()} (will move to GPU during inference if on Spaces)...")
# Hotfix for Florence-2 in newer transformers versions
import transformers
if not hasattr(transformers.PretrainedConfig, "forced_bos_token_id"):
transformers.PretrainedConfig.forced_bos_token_id = None
# Force _supports_sdpa to False on the actual base class
import transformers.modeling_utils
transformers.modeling_utils.PreTrainedModel._supports_sdpa = False
self.model = AutoModelForCausalLM.from_pretrained(
CONFIG.MODEL_NAME,
trust_remote_code=True,
torch_dtype=DTYPE,
attn_implementation="eager"
).eval()
# Hotfix for Florence-2 processor tokenizer compatibility
if not hasattr(transformers.PreTrainedTokenizerBase, "additional_special_tokens"):
transformers.PreTrainedTokenizerBase.additional_special_tokens = property(
lambda self: getattr(self, "_additional_special_tokens", [])
)
self.processor = AutoProcessor.from_pretrained(
CONFIG.MODEL_NAME,
trust_remote_code=True,
)
print("Florence-2 loaded successfully")
try:
from paddleocr import PaddleOCR
print("Loading PaddleOCR...")
self.paddle_ocr = PaddleOCR(use_angle_cls=True, lang='en', show_log=False)
print("PaddleOCR loaded successfully")
except Exception as e:
print(f"PaddleOCR load failed: {e}")
self._warmup()
except Exception as e:
print(f"Model loading failed: {e}")
raise
def _warmup(self):
"""Run a dummy inference to warm up kernels."""
if IS_SPACES:
print("Skipping warmup on ZeroGPU Spaces")
return
try:
dummy = Image.new("RGB", (224, 224), 128)
self._run_inference(dummy, "<CAPTION>")
print("Model warmed up")
except Exception as e:
print(f"Warmup warning: {e}")
@gpu_decorator
def _run_inference(self, image: Image.Image, task_token: str) -> Dict[str, Any]:
"""Core inference logic."""
if self.model is None or self.processor is None:
raise RuntimeError("Model not loaded. Call load() first.")
# Ensure model is on the right device when inference runs
target_device = "cuda" if IS_SPACES else DEVICE
if next(self.model.parameters()).device.type != target_device:
self.model.to(target_device)
image = preprocess_image(image)
image = auto_enhance(image)
max_tokens = CONFIG.MAX_NEW_TOKENS.get(task_token, 64)
inputs = self.processor(
text=task_token,
images=image,
return_tensors="pt",
).to(target_device)
if "pixel_values" in inputs:
inputs["pixel_values"] = inputs["pixel_values"].to(DTYPE)
with torch.inference_mode():
output_ids = self.model.generate(
input_ids=inputs["input_ids"],
pixel_values=inputs["pixel_values"],
max_new_tokens=max_tokens,
do_sample=False,
num_beams=1,
use_cache=True,
)
raw_text = self.processor.batch_decode(output_ids, skip_special_tokens=False)[0]
result = self.processor.post_process_generation(
raw_text,
task=task_token,
image_size=(image.width, image.height),
)
return result
def describe_scene(self, image: Image.Image, detailed: bool = False) -> str:
task = "<MORE_DETAILED_CAPTION>" if detailed else "<DETAILED_CAPTION>"
try:
result = self._run_inference(image, task)
return format_caption(result.get(task, ""))
except Exception as e:
print(f"describe_scene error: {e}")
return "I couldn't analyze the scene right now."
def read_text(self, image: Image.Image) -> str:
try:
if hasattr(self, 'paddle_ocr') and self.paddle_ocr:
import numpy as np
# Convert PIL Image to RGB Numpy array for PaddleOCR
img_array = np.array(image.convert("RGB"))
result = self.paddle_ocr.ocr(img_array, cls=True)
if not result or result[0] is None:
return "I couldn't find any clear text in the image."
lines = []
for line in result[0]:
text = line[1][0]
lines.append(text)
final_text = " ".join(lines).strip()
if not final_text:
return "I couldn't find any clear text."
return f"The text says: {final_text}"
else:
# Fallback to Florence-2 OCR
result = self._run_inference(image, "<OCR>")
return format_ocr(result.get("<OCR>", ""))
except Exception as e:
print(f"read_text error: {e}")
return "I couldn't read the text right now."
def analyze(self, image: Image.Image, task: str) -> str:
# Generic handler
try:
result = self._run_inference(image, task)
return str(result)
except Exception as e:
return f"Error: {e}"
|