Update main.py

main.py

@@ -14,25 +14,11 @@ app = FastAPI()
 hf_token = os.environ.get('HF_TOKEN')
 client = Client("https://ashrafb-image-to-sketch.hf.space/", hf_token=hf_token)
 
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    try:
-        # Save the uploaded image to a temporary file
-        with NamedTemporaryFile(delete=False) as tmp:
-            shutil.copyfileobj(file.file, tmp)
-            tmp_path = tmp.name
-
-        # Call the Gradio client to predict
-        result = client.predict(tmp_path, api_name="/predict")
-
-        # Delete the temporary file
-        os.unlink(tmp_path)
-
-        # Parse and return the result
-        sketch_image_url, result_file_url = result
-        return JSONResponse(content={"sketch_image_url": sketch_image_url, "result_file_url": result_file_url}, status_code=200)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+@app.post("/upload/")
+async def upload_file(file: UploadFile = File(...)):
+    contents = await file.read()
+    result = client.predict(contents, api_name="/predict")
+    return {"sketch_image": result[0]}
 
 app.mount("/", StaticFiles(directory="static", html=True), name="static")