convnext v2

app.py

@@ -6,7 +6,7 @@ import os, glob
 
 
 model, transforms = get_watermarks_detection_model(
-    'convnext-wm_1102', 
+    'convnext-wm_1102_v2', 
     fp16=False, 
     cache_dir='model_files'
 )

wmdetection/models/__init__.py

@@ -9,7 +9,7 @@ from wmdetection.utils import FP16Module
 
 
 def get_convnext_model(name):
-    if name == 'convnext-tiny' or name == 'convnext-wm_1102':
+    if name == 'convnext-tiny' or name == 'convnext-wm_1102' or name == 'convnext-wm_1102_v2':
         model_ft = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768])
         model_ft.head = nn.Sequential( 
             nn.Linear(in_features=768, out_features=512),
@@ -78,6 +78,11 @@ MODELS = {
         repo_id='Inf009/wm_1102',
         filename='convnext_v1_9.pth',
     ),
+    'convnext-wm_1102_v2': dict(
+        constructor=get_convnext_model,
+        repo_id='Inf009/wm_1102',
+        filename='convnext_v2.pth',
+    ),
     'resnext101_32x8d-large': dict(
         constructor=get_resnext_model,
         repo_id='boomb0om/watermark-detectors',