Create smolvlm_inference.py

smolvlm_inference.py

@@ -0,0 +1,23 @@
+import torch
+from transformers import AutoModelForImageTextToText, AutoProcessor
+
+
+class TransformersModel:
+    def __init__(self, model_id: str, to_device: str = "cuda"):
+        self.model_id = model_id
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        self.processor.image_processor.size = {"longest_edge": 3 * 384}
+        self.model = AutoModelForImageTextToText.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(to_device)
+
+    def generate(self, messages: list[dict], **kwargs):
+        inputs = self.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(self.model.device, dtype=torch.bfloat16)
+        generated_ids = self.model.generate(**inputs, **kwargs)
+        return self.processor.batch_decode(
+                generated_ids[:, len(inputs["input_ids"][0]) :], skip_special_tokens=True
+            )[0]