Migrate hf_hub to transformers fuser

vlm.py

@@ -1,6 +1,9 @@
 # vlm.py
-import os, logging, traceback, json
-from huggingface_hub import InferenceClient
+import os, logging, traceback, json, base64
+from io import BytesIO
+from PIL import Image
+# from huggingface_hub import InferenceClient # Render model on HF hub
+from transformers import pipeline             # Render model on transformers
 from translation import translate_query
 
 # Initialise once
@@ -10,6 +13,16 @@ client = InferenceClient(provider="auto", api_key=HF_TOKEN)
 logger = logging.getLogger("vlm-agent")
 logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(name)s — %(levelname)s — %(message)s", force=True) # Change INFO to DEBUG for full-ctx JSON loader
 
+# ✅ Load VLM pipeline once (lazy load allowed)
+vlm_pipe = None
+def load_vlm():
+    global vlm_pipe
+    if vlm_pipe is None:
+        logger.info("⏳ Loading MedGEMMA model via Transformers pipeline...")
+        vlm_pipe = pipeline("image-to-text", model="google/medgemma-4b", device_map="auto")
+        logger.info("✅ MedGEMMA model ready.")
+    return vlm_pipe
+
 def process_medical_image(base64_image: str, prompt: str = None, lang: str = "EN") -> str:
     """
     Send base64 image + prompt to MedGEMMA and return output.
@@ -20,16 +33,22 @@ def process_medical_image(base64_image: str, prompt: str = None, lang: str = "EN
         user_query = translate_query(user_query, lang.lower())
     # Send over API
     try:
-        response = client.chat.completions.create(
-            model="google/medgemma-4b-it",
-            messages=[{
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
-                ]
-            }]
-        )
+        # HF hub
+        # response = client.chat.completions.create(
+        #     model="google/medgemma-4b-it",
+        #     messages=[{
+        #         "role": "user",
+        #         "content": [
+        #             {"type": "text", "text": prompt},
+        #             {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+        #         ]
+        #     }]
+        # )
+        # Transformers
+        image_data = base64.b64decode(base64_image) # Decode base64 to PIL Image
+        image = Image.open(BytesIO(image_data)).convert("RGB")
+        pipe = load_vlm()
+        response = pipe(image, prompt=prompt, max_new_tokens=100)[0]["generated_text"]
         # Validate response
         if not response or not hasattr(response, "choices") or not response.choices:
             raise ValueError("Empty or malformed response from MedGEMMA.")