Create inference.py

inference.py

@@ -0,0 +1,56 @@
+import torch
+import gc
+from transformers import TextStreamer, pipeline, BitsAndBytesConfig
+
+# 1. Clear GPU memory to recover from the previous OOM
+if 'pipe' in locals(): del pipe
+if 'outputs' in locals(): del outputs
+gc.collect()
+torch.cuda.empty_cache()
+
+# 2. Configure 4-bit quantization properly for custom architectures
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True
+)
+
+model_id = "Xerv-AI/tarn"
+pipe = pipeline(
+    "image-text-to-text",
+    model=model_id,
+    device_map="auto",
+    trust_remote_code=True,
+    model_kwargs={
+        "quantization_config": bnb_config,
+        "torch_dtype": torch.float16
+    }
+)
+
+streamer = TextStreamer(pipe.tokenizer, skip_prompt=True)
+
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"
+            },
+            {
+                "type": "text",
+                "text": "Analyze the visual artifacts present in this image and define the principles of triboelectricity."
+            }
+        ]
+    },
+]
+
+print("=== Initiating Real-Time Telemetry Stream (Quantized) ===")
+outputs = pipe(
+    text=messages,
+    generate_kwargs={
+        "max_new_tokens": 512,
+        "streamer": streamer
+    }
+)