Update main.py

main.py

@@ -5,13 +5,13 @@ from fastapi import FastAPI, File, UploadFile
 from pydantic import BaseModel
 import numpy as np
 import cv2
+import torch
 from vtoonify_model import Model
 
 app = FastAPI()
 
 # Load the model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = Model(device)
+model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
 exstyle, message = model.load_model("cartoon1")
 
 class ImageRequest(BaseModel):