Update main.py

main.py

@@ -12,37 +12,39 @@ app = FastAPI()
 
 # Initialize Gradio Client
 hf_token = os.environ.get('HF_TOKEN')
-client = Client("Ashrafb/moondream_captioning", hf_token=hf_token)
-
+client = Client("Ashrafb/moondream1", hf_token=hf_token)
 # Function to generate captions for uploaded images
-def generate_caption(image_path):
-    context = "describe"  # Fixed context value
-    result = client.predict(image_path, context, api_name="/get_caption")
-    return result
+def get_caption(image, additional_context):
+    if additional_context.strip():  # Check if additional_context is not empty
+        context = additional_context
+    else:
+        context = "What is this image? Describe this image to someone who is visually impaired."
+    result = client.predict(image, context, api_name="/predict")
+    return result  
 
-# Route to handle image uploads and generate captions
 @app.post("/uploadfile/")
-async def generate_image_caption(file: UploadFile = File(...)):
+async def upload_file(image: UploadFile = File(...), additional_context: str = Form(...)):
     try:
-        # Save the uploaded image to a temporary file
-        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-            shutil.copyfileobj(file.file, temp_file)
-            temp_file_path = temp_file.name
-        
-        # Generate caption for the uploaded image
-        caption = generate_caption(temp_file_path)
-        
-        # Clean up temporary file
-        os.unlink(temp_file_path)
-        
-        return {"caption": caption}
-    
+        # Create a temporary directory to store the uploaded image
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_image_path = os.path.join(temp_dir, image.filename)
+            # Write the uploaded image data to a temporary file
+            with open(temp_image_path, "wb") as temp_image:
+                shutil.copyfileobj(image.file, temp_image)
+            # Read the image data
+            with open(temp_image_path, "rb") as image_file:
+                image_data = image_file.read()
+            # Generate caption using Gradio client
+            result = get_caption(image_data, additional_context)
+            return {"result": result}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 app.mount("/", StaticFiles(directory="static", html=True), name="static")
 
 @app.get("/")
 def index() -> FileResponse:
     return FileResponse(path="/app/static/index.html", media_type="text/html")
+