Update medgemma_api.py

medgemma_api.py

@@ -2,19 +2,31 @@ import os
 import requests
 import base64
 
-from transformers import pipeline
+# 获取 Hugging Face Token（已通过 Secrets 设置）
+HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+API_URL = "https://api-inference.huggingface.co/models/google/medgemma-4b-it"
+HEADERS = {"Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}"}
 
-pipe = pipeline("image-text-to-text", model="google/medgemma-4b-it")
+def query_medgemma(image_path, question):
+    with open(image_path, "rb") as f:
+        image_bytes = f.read()
+    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
 
-def query_medgemma_local(image_path, question):
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "source": image_path},
-                {"type": "text", "text": question}
-            ]
-        }
-    ]
-    return pipe(messages)[0]["generated_text"]
+    payload = {
+        "inputs": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": encoded_image},
+                    {"type": "text", "text": question}
+                ]
+            }
+        ]
+    }
+
+    response = requests.post(API_URL, headers=HEADERS, json=payload)
+    if response.ok:
+        return response.json()[0]["generated_text"]
+    else:
+        return f"Error: {response.text}"