stream

endpoint_handler.py

@@ -22,7 +22,7 @@ class EndpointHandler:
 
         print(f"[Model Load] Loading model from: {model_path}")
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path,trust_remote_code=True)
             self.model = AutoModel.from_pretrained(
                 model_path,
                 trust_remote_code=True,
@@ -50,7 +50,6 @@ class EndpointHandler:
 
     def predict(self, request):
         print(f"[Predict] Received request: {request}")
-
         image_base64 = request.get("inputs", {}).get("image")
         question = request.get("inputs", {}).get("question")
         stream = request.get("inputs", {}).get("stream", False)

main.py

@@ -63,15 +63,32 @@ async def predict_endpoint(payload: PredictRequest):
     except Exception as e:
         return JSONResponse({"error": "Internal server error"}, status_code=500)
 
+    # ─── If it's a generator, return SSE/streaming ──────────────────────────
     if isinstance(result, types.GeneratorType):
         def event_stream():
             try:
                 for chunk in result:
+                    # Each chunk should already be a Python dict or
+                    # a string containing JSON. We wrap it in "data: …\n\n"
                     yield f"data: {json.dumps(chunk)}\n\n"
-                # Return structured JSON to indicate end of stream
+                # Finally send an end‐of‐stream marker
                 yield f"data: {json.dumps({'end': True})}\n\n"
             except Exception as e:
-                # Return structured JSON to indicate error
                 yield f"data: {json.dumps({'error': str(e)})}\n\n"
-    return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+        return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+    # ─── Otherwise (non‐streaming), return a single JSON response ──────────
+    #    result is expected to be a JSON‐string or a dict
+    try:
+        # If handler.predict returned a JSON‐encoded str, parse it to dict
+        if isinstance(result, str):
+            parsed = json.loads(result)
+        else:
+            parsed = result  # assume it's already a dict
+    except Exception:
+        # Fall back to returning the raw result
+        return JSONResponse({"error": "Invalid JSON from handler"}, status_code=500)
+
+    return JSONResponse(parsed)