for upcoming security features

engines/emica_encoder/EmicaEncoder.py

@@ -16,7 +16,7 @@ def _init_models(self, ):
         _abs_path = os.path.dirname(os.path.abspath(__file__))
         _model_path = os.path.join(_abs_path, '../../assets/emica/EMICA-CVT_flame2020_notexture.pt')
         assert os.path.exists(_model_path), f"Model not found: {_model_path}."
-        ckpt = torch.load(_model_path, map_location='cpu')
+        ckpt = torch.load(_model_path, map_location='cpu', weights_only=True)
         self.mica_encoder = MicaEncoder()
         self.deca_encoder = DecaEncoder(outsize=86)
         self.expression_encoder = DecaEncoder(outsize=100)

engines/emica_encoder/MicaEncoder.py

@@ -84,7 +84,6 @@ def __init__(self,
                  block, layers, dropout=0, num_features=512, zero_init_residual=False,
                  groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):
         super(IResNet, self).__init__()
-        self.fp16 = fp16
         self.inplanes = 64
         self.dilation = 1
         self.block = block
@@ -160,18 +159,17 @@ def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
         return nn.Sequential(*layers)
 
     def forward(self, x):
-        with torch.cuda.amp.autocast(self.fp16):
-            x = self.conv1(x)
-            x = self.bn1(x)
-            x = self.prelu(x)
-            x = self.layer1(x)
-            x = self.layer2(x)
-            x = self.layer3(x)
-            x = self.layer4(x)
-            x = self.bn2(x)
-            x = torch.flatten(x, 1)
-            x = self.dropout(x)
-        x = self.fc(x.float() if self.fp16 else x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.prelu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.bn2(x)
+        x = torch.flatten(x, 1)
+        x = self.dropout(x)
+        x = self.fc(x)
         x = self.features(x)
         return x
 
@@ -180,7 +178,7 @@ class Arcface(IResNet):
     def __init__(self, pretrained_path=None, **kwargs):
         super(Arcface, self).__init__(IBasicBlock, [3, 13, 30, 3], **kwargs)
         if pretrained_path is not None and os.path.exists(pretrained_path):
-            self.load_state_dict(torch.load(pretrained_path))
+            self.load_state_dict(torch.load(pretrained_path, map_location='cpu', weights_only=True))
         self.freezer([self.layer1, self.layer2, self.layer3, self.conv1, self.bn1, self.prelu])
 
     def freezer(self, layers):
@@ -193,21 +191,20 @@ def forward(self, images):
         return x
 
     def forward_arcface(self, x):
-        with torch.cuda.amp.autocast(self.fp16):
-            ### FROZEN ###
-            with torch.no_grad():
-                x = self.conv1(x)
-                x = self.bn1(x)
-                x = self.prelu(x)
-                x = self.layer1(x)
-                x = self.layer2(x)
-                x = self.layer3(x)
-
-            x = self.layer4(x)
-            x = self.bn2(x)
-            x = torch.flatten(x, 1)
-            x = self.dropout(x)
-        x = self.fc(x.float() if self.fp16 else x)
+        ### FROZEN ###
+        with torch.no_grad():
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.prelu(x)
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+
+        x = self.layer4(x)
+        x = self.bn2(x)
+        x = torch.flatten(x, 1)
+        x = self.dropout(x)
+        x = self.fc(x)
         x = self.features(x)
         return x
 

engines/flame_model/FLAME.py

@@ -25,7 +25,9 @@ def __init__(self, n_shape, n_exp, scale=1.0, no_lmks=False):
         # print("creating the FLAME Model")
         _abs_path = os.path.dirname(os.path.abspath(__file__))
         self.flame_path = os.path.join(_abs_path, '../../assets/flame')
-        self.flame_ckpt = torch.load(os.path.join(self.flame_path, 'FLAME_with_eye.pt'))
+        self.flame_ckpt = torch.load(
+            os.path.join(self.flame_path, 'FLAME_with_eye.pt'), map_location='cpu', weights_only=True
+        )
         flame_model = self.flame_ckpt['flame_model']
         flame_lmk = self.flame_ckpt['lmk_embeddings']
 

engines/human_matting/matting_engine.py

@@ -15,7 +15,7 @@ def _init_models(self, ):
         # load dict
         _abs_script_path = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
         _ckpt_path = os.path.join(_abs_script_path, '../../assets/matting', 'stylematte_synth.pt')
-        state_dict = torch.load(_ckpt_path, map_location='cpu')
+        state_dict = torch.load(_ckpt_path, map_location='cpu', weights_only=True)
         # build model
         model = StyleMatte()
         model.load_state_dict(state_dict)

engines/utils_lmdb.py

@@ -38,7 +38,7 @@ def __getitem__(self, key_name):
             image_buf = torch.tensor(np.frombuffer(payload, dtype=np.uint8))
             data = torchvision.io.decode_image(image_buf, mode=torchvision.io.ImageReadMode.RGB)
         except:
-            data = torch.load(io.BytesIO(payload))
+            data = torch.load(io.BytesIO(payload), weights_only=True)
         return data
 
     def __del__(self,):
@@ -52,7 +52,7 @@ def load(self, key_name, type='image', **kwargs):
         if payload is None:
             raise KeyError('Key:{} Not Found!'.format(key_name))
         if type == 'torch':
-            torch_data = torch.load(io.BytesIO(payload))
+            torch_data = torch.load(io.BytesIO(payload), weights_only=True)
             return torch_data
         elif type == 'image':
             image_buf = torch.tensor(np.frombuffer(payload, dtype=np.uint8))
@@ -94,7 +94,7 @@ def dump(self, key_name, payload, type='image', encode_jpeg=True):
                     payload[key] = payload[key].detach().float().cpu()
                 torch.save(payload, torch_buf)
             payload_encoded = torch_buf.getvalue()
-            # torch_data = torch.load(io.BytesIO(payload_encoded))
+            # torch_data = torch.load(io.BytesIO(payload_encoded), weights_only=True)
             self._lmdb_txn.put(key_name.encode(), payload_encoded)
         elif type == 'image':
             assert payload.dim() == 3 and payload.shape[0] == 3

engines/vgghead_detector/utils_lmks_detector.py

@@ -568,7 +568,7 @@ def __init__(self, model_path):
 
         self.heatmap_head = HeatmapHead()
 
-        self.load_state_dict(torch.load(model_path, map_location='cpu'))
+        self.load_state_dict(torch.load(model_path, map_location='cpu', weights_only=True))
 
     def forward(self, img):
         heatmap, landmark = self.heatmap_head(self.backbone(img))