Update main.py

main.py

@@ -28,14 +28,17 @@ async def upload_file(file: UploadFile = File(...)):
         result = client.predict(temp_file_path, api_name="/predict")
         
         # Check if the result is valid
-        if result and len(result) > 0:
+        if result and len(result) == 2:
             # Convert the image data to base64 string
             with open(result[0], "rb") as image_file:
                 image_data = base64.b64encode(image_file.read()).decode("utf-8")
 
-            return {"sketch_image_base64": f"data:image/png;base64,{image_data}"}
+            return {
+                "sketch_image_base64": f"data:image/png;base64,{image_data}",
+                "result_file": result[1]
+            }
         else:
-            return {"error": "No sketch image generated."}
+            return {"error": "Invalid result from the prediction API."}
     finally:
         os.unlink(temp_file_path)