Update main.py

main.py

@@ -1,37 +1,34 @@
 import os
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import JSONResponse
-
 from gradio_client import Client, handle_file
 
 app = FastAPI()
 
-
-
-
 HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is not set.")
 
-# Initialize the Gradio client with the token
-client = Client("Luisgust/moondream1", hf_token=HF_TOKEN)
+try:
+    client = Client("Luisgust/moondream1", hf_token=HF_TOKEN)
+except Exception as e:
+    print(f"Failed to initialize Gradio client: {e}")
+    raise
 
 @app.post("/get_caption")
 async def get_caption(image: UploadFile = File(...), context: str = Form(...)):
     try:
-        # Read the uploaded image file
         contents = await image.read()
-
-        # Use handle_file to prepare the image for the Gradio client
         image_data = handle_file(contents, orig_name=image.filename)
 
-        # Call the Gradio API
         result = client.predict(
             image=image_data,
             question=context,
             api_name="/answer_question"
         )
 
-        # Return the result as a JSON response
         return JSONResponse(content={"caption": result})
     except Exception as e:
+        print(f"Error during prediction: {e}")
         return JSONResponse(content={"error": str(e)}, status_code=500)