Update main.py

main.py

@@ -1,9 +1,7 @@
 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
+from fastapi.responses import StreamingResponse
+from fastapi.staticfiles import StaticFiles
 import torch
 import shutil
 import cv2
@@ -12,33 +10,18 @@ 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
-
-import gradio as gr
-import pathlib
-import sys
-sys.path.insert(0, 'vtoonify')
-
-from util import load_psp_standalone, get_video_crop_parameter, tensor2cv2
-import torch
-import torch.nn as nn
-import numpy as np
-import dlib
-import cv2
+from util import load_psp_standalone, get_video_crop_parameter, tensor2cv2, align_face
 from model.vtoonify import VToonify
 from model.bisenet.model import BiSeNet
-import torch.nn.functional as F
-from torchvision import transforms
-from model.encoder.align_all_parallel import align_face
-import gc
 import huggingface_hub
 import os
+from io import BytesIO
 
 app = FastAPI()
-model = None
 
 MODEL_REPO = 'PKUWilliamYang/VToonify'
 
-class Model():
+class Model:
     def __init__(self, device):
         super().__init__()
         
@@ -53,19 +36,17 @@ class Model():
         self.pspencoder = self._load_encoder()    
         self.transform = transforms.Compose([
             transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5,0.5,0.5]),
-            ])
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
         
         self.vtoonify, self.exstyle = self._load_default_model()
         self.color_transfer = False
         self.style_name = 'cartoon1'
         self.video_limit_cpu = 100
         self.video_limit_gpu = 300
-        
-    @staticmethod
-    def _create_dlib_landmark_model():
-        return dlib.shape_predictor(huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                    'models/shape_predictor_68_face_landmarks.dat'))
+
+    def _create_dlib_landmark_model(self):
+        return dlib.shape_predictor(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/shape_predictor_68_face_landmarks.dat'))
     
     def _create_parsing_model(self):
         parsingpredictor = BiSeNet(n_classes=19)
@@ -75,16 +56,16 @@ class Model():
         return parsingpredictor
     
     def _load_encoder(self) -> nn.Module:
-        style_encoder_path = huggingface_hub.hf_hub_download(MODEL_REPO,'models/encoder.pt')
+        style_encoder_path = huggingface_hub.hf_hub_download(MODEL_REPO, 'models/encoder.pt')
         return load_psp_standalone(style_encoder_path, self.device)
     
     def _load_default_model(self) -> tuple[torch.Tensor, str]:
-        vtoonify = VToonify(backbone = 'dualstylegan')
+        vtoonify = VToonify(backbone='dualstylegan')
         vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,
                                             'models/vtoonify_d_cartoon/vtoonify_s026_d0.5.pt'), 
                                             map_location=lambda storage, loc: storage)['g_ema'])
         vtoonify.to(self.device)
-        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO,'models/vtoonify_d_cartoon/exstyle_code.npy'), allow_pickle=True).item()
+        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/vtoonify_d_cartoon/exstyle_code.npy'), allow_pickle=True).item()
         exstyle = torch.tensor(tmp[list(tmp.keys())[26]]).to(self.device)
         with torch.no_grad():  
             exstyle = vtoonify.zplus2wplus(exstyle)
@@ -99,14 +80,14 @@ class Model():
             return None, 'Oops, wrong Style Type. Please select a valid model.'
         self.style_name = style_type
         model_path, ind = self.style_types[style_type]
-        style_path = os.path.join('models',os.path.dirname(model_path),'exstyle_code.npy')
-        self.vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,'models/'+model_path), 
+        style_path = os.path.join('models', os.path.dirname(model_path), 'exstyle_code.npy')
+        self.vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/' + model_path), 
                                             map_location=lambda storage, loc: storage)['g_ema'])
         tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, style_path), allow_pickle=True).item()
         exstyle = torch.tensor(tmp[list(tmp.keys())[ind]]).to(self.device)
         with torch.no_grad():  
             exstyle = self.vtoonify.zplus2wplus(exstyle)
-        return exstyle, 'Model of %s loaded.'%(style_type)
+        return exstyle, 'Model of %s loaded.' % (style_type)
     
     def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
         message = 'Error: no face detected! Please retry or change the photo.'
@@ -114,7 +95,7 @@ class Model():
         instyle = None
         h, w, scale = 0, 0, 0
         if paras is not None:
-            h,w,top,bottom,left,right,scale = paras
+            h, w, top, bottom, left, right, scale = paras
             H, W = int(bottom-top), int(right-left)
             # for HR image, we apply gaussian blur to it to avoid over-sharp stylization results
             kernel_1d = np.array([[0.125],[0.375],[0.375],[0.125]])
@@ -129,11 +110,11 @@ class Model():
                     I = self.transform(I).unsqueeze(dim=0).to(self.device)
                     instyle = self.pspencoder(I)
                     instyle = self.vtoonify.zplus2wplus(instyle)
-                    message = 'Successfully rescale the frame to (%d, %d)'%(bottom-top, right-left)
+                    message = 'Successfully rescale the frame to (%d, %d)' % (bottom-top, right-left)
                 else:
-                    frame = np.zeros((256,256,3), np.uint8)
+                    frame = np.zeros((256, 256, 3), np.uint8)
         else:
-            frame = np.zeros((256,256,3), np.uint8)
+            frame = np.zeros((256, 256, 3), np.uint8)
         if return_para:
             return frame, instyle, message, w, h, top, bottom, left, right, scale
         return frame, instyle, message
@@ -142,21 +123,21 @@ class Model():
     def detect_and_align_image(self, image: str, top: int, bottom: int, left: int, right: int
                               ) -> tuple[np.ndarray, torch.Tensor, str]:
         if image is None:
-            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
+            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load empty file.'
         frame = cv2.imread(image)
         if frame is None:
-            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load the image.'       
+            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'       
         frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
         return self.detect_and_align(frame, top, bottom, left, right)
     
     def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int
                               ) -> tuple[np.ndarray, torch.Tensor, str]:
         if video is None:
-            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
+            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load empty file.'
         video_cap = cv2.VideoCapture(video)
         if video_cap.get(7) == 0:
             video_cap.release()
-            return np.zeros((256,256,3), np.uint8), torch.zeros(1,18,512).to(self.device), 'Error: fail to load the video.'
+            return np.zeros((256, 256, 3), np.uint8), torch.zeros(1, 18, 512).to(self.device), 'Error: fail to load the video.'
         success, frame = video_cap.read()
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         video_cap.release()
@@ -166,11 +147,11 @@ class Model():
     def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[np.ndarray, str]:
         #print(style_type + ' ' + self.style_name)
         if instyle is None or aligned_face is None:
-            return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
+            return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
         if self.style_name != style_type:
             exstyle, _  = self.load_model(style_type)
         if exstyle is None:
-            return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
+            return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
         with torch.no_grad():
             if self.color_transfer:
                 s_w = exstyle
@@ -182,17 +163,13 @@ class Model():
             x_p = F.interpolate(self.parsingpredictor(2*(F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0], 
                                 scale_factor=0.5, recompute_scale_factor=False).detach()
             inputs = torch.cat((x, x_p/16.), dim=1)
-            y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s = style_degree)        
+            y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s=style_degree)        
             y_tilde = torch.clamp(y_tilde, -1, 1)
-        print('*** Toonify %dx%d image with style of %s'%(y_tilde.shape[2], y_tilde.shape[3], style_type))
-        return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s'%(self.style_name)
-    
+        print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
+        return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s' % (self.style_name)
 
+model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
 
-    model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
-
-from fastapi.responses import StreamingResponse
-from io import BytesIO
 
 @app.post("/upload/")
 async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
@@ -216,6 +193,7 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
 
 app.mount("/", StaticFiles(directory="AB", html=True), name="static")
 
+
 @app.get("/")
 def index() -> FileResponse:
     return FileResponse(path="/app/AB/index.html", media_type="text/html")