Update main.py

main.py

@@ -1,46 +1,37 @@
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException
-from fastapi.responses import HTMLResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from fastapi.responses import FileResponse
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from gradio_client import Client
-import shutil
 import os
 import tempfile
-
-
+import shutil
 
 app = FastAPI()
 
+# Initialize Gradio Client
+hf_token = os.environ.get('HF_TOKEN')
+client = Client("Ashrafb/moondream_captioning", hf_token=hf_token)
 
-
-# Function to make API prediction
-def predict_image_description(file_path):
-    hf_token = os.environ.get('HF_TOKEN')
-    client = Client("Ashrafb/moondream1", hf_token=hf_token)
-    result = client.predict(file_path, api_name="/get_caption")
+# 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
 
-
-
-
-# Route to handle file upload
-@app.post("/uploadfile/")
-async def create_upload_file(file: UploadFile = File(...)):
+# Route to handle image uploads and generate captions
+@app.post("/generate_caption/")
+async def generate_image_caption(image: UploadFile = File(...)):
     try:
+        # Save the uploaded image to a temporary file
         with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-            shutil.copyfileobj(file.file, temp_file)
+            shutil.copyfileobj(image.file, temp_file)
             temp_file_path = temp_file.name
-        description = predict_image_description(temp_file_path)
+        
+        # Generate caption for the uploaded image
+        caption = generate_caption(temp_file_path)
+        
+        # Clean up temporary file
         os.unlink(temp_file_path)
-        return {"description": description}
+        
+        return {"caption": caption}
+    
     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")
-