Update main.py

main.py

@@ -17,18 +17,43 @@ exstyle, message = model.load_model("cartoon1")
 class ImageRequest(BaseModel):
     image_file: UploadFile = File(...)
 
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+import numpy as np
+import cv2
+import torch
+from vtoonify_model import Model
+
+app = FastAPI()
+
+# Load the model
+model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
+exstyle, message = model.load_model("cartoon1")
+
 @app.post("/upload/")
 async def toonify_image(image_file: UploadFile = File(...)):
     try:
+        # Read the contents of the uploaded file
         contents = await image_file.read()
-        filename = image_file.filename
+        # Determine the file extension from the content type
+        ext = image_file.content_type.split("/")[-1]
+        # Construct a unique filename with the extension
+        filename = f"uploaded_image.{ext}"
+        
+        # Convert the file contents to a numpy array
         nparr = np.frombuffer(contents, np.uint8)
+        # Decode the image using OpenCV
         img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        
+        # Perform image processing with the model
         aligned_face, instyle, message = model.detect_and_align_image(img, 200, 200, 200, 200)  # Hardcoded values
         toonified_img, message = model.image_toonify(aligned_face, instyle, exstyle, style_degree=0.5, style_type="cartoon1")
+        
+        # Return the processed image and filename
         return {"toonified_image": toonified_img, "message": message, "filename": filename}
     except Exception as e:
-        return {"error": str(e)}
+        # Return an error response if an exception occurs
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
 app.mount("/", StaticFiles(directory="AB", html=True), name="static")