Update main.py

main.py

@@ -2,28 +2,36 @@ from __future__ import annotations
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import FileResponse
 import shutil
-import torch
+import cv2
 import numpy as np
-
-from vtoonify_model import Model
+import dlib
+from torchvision import transforms
+import torch.nn.functional as F
+from vtoonify_model import Model  # Importing the Model class from vtoonify_model.py
 
 app = FastAPI()
-model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
+model = None
+
+@app.on_event("startup")
+async def load_model():
+    global model
+    model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
 
 @app.post("/upload/")
-async def process_image(file: UploadFile = File(...)):
+async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
+    if model is None:
+        return {"error": "Model not loaded."}
+
     # Save the uploaded image locally
     with open("uploaded_image.jpg", "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
 
-    # Load the model (assuming 'cartoon1' is always used)
-    exstyle, load_info = model.load_model('cartoon1')
-
     # Process the uploaded image
-    top, bottom, left, right = 200, 200, 200, 200  
-    aligned_face, _, input_info = model.detect_and_align_image("uploaded_image.jpg", top, bottom, left, right)
-    processed_image, message = model.image_toonify(aligned_face, exstyle, style_degree=0.5, style_type='cartoon1')
+    aligned_face, instyle, message = model.detect_and_align_image("uploaded_image.jpg", top, bottom, left, right)
+    processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
 
     # Save the processed image
     with open("result_image.jpg", "wb") as result_buffer:
@@ -32,6 +40,7 @@ async def process_image(file: UploadFile = File(...)):
     # Return the processed image
     return FileResponse("result_image.jpg", media_type="image/jpeg", headers={"Content-Disposition": "attachment; filename=result_image.jpg"})
 
+
 app.mount("/", StaticFiles(directory="AB", html=True), name="static")
 
 @app.get("/")