Update main.py

main.py

@@ -35,14 +35,34 @@ app = FastAPI()
 
 MODEL_REPO = 'PKUWilliamYang/VToonify'
 
-class Model:
+class Model():
     def __init__(self, device):
         super().__init__()
         
         self.device = device
         self.style_types = {
             'cartoon1': ['vtoonify_d_cartoon/vtoonify_s026_d0.5.pt', 26],
-        
+            'cartoon1-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 26],
+            'cartoon2-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 64],
+            'cartoon3-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 153],
+            'cartoon4': ['vtoonify_d_cartoon/vtoonify_s299_d0.5.pt', 299],
+            'cartoon4-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 299],
+            'cartoon5-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 8],
+            'comic1-d': ['vtoonify_d_comic/vtoonify_s_d.pt', 28],
+            'comic2-d': ['vtoonify_d_comic/vtoonify_s_d.pt', 18],
+            'arcane1': ['vtoonify_d_arcane/vtoonify_s000_d0.5.pt', 0],
+            'arcane1-d': ['vtoonify_d_arcane/vtoonify_s_d.pt', 0],
+            'arcane2': ['vtoonify_d_arcane/vtoonify_s077_d0.5.pt', 77],
+            'arcane2-d': ['vtoonify_d_arcane/vtoonify_s_d.pt', 77],
+            'caricature1': ['vtoonify_d_caricature/vtoonify_s039_d0.5.pt', 39],
+            'caricature2': ['vtoonify_d_caricature/vtoonify_s068_d0.5.pt', 68],
+            'pixar': ['vtoonify_d_pixar/vtoonify_s052_d0.5.pt', 52],
+            'pixar-d': ['vtoonify_d_pixar/vtoonify_s_d.pt', 52],
+            'illustration1-d': ['vtoonify_d_illustration/vtoonify_s054_d_c.pt', 54],
+            'illustration2-d': ['vtoonify_d_illustration/vtoonify_s004_d_c.pt', 4],
+            'illustration3-d': ['vtoonify_d_illustration/vtoonify_s009_d_c.pt', 9],
+            'illustration4-d': ['vtoonify_d_illustration/vtoonify_s043_d_c.pt', 43],
+            'illustration5-d': ['vtoonify_d_illustration/vtoonify_s086_d_c.pt', 86],
         }
         
         self.landmarkpredictor = self._create_dlib_landmark_model()
@@ -50,15 +70,15 @@ 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
-
+        
     def _create_dlib_landmark_model(self):
         return dlib.shape_predictor(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/shape_predictor_68_face_landmarks.dat'))