use new proxy model

app.py

@@ -7,9 +7,11 @@ import numpy as np
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
-from transformers import AutoModel, BlipImageProcessor
+from transformers import AutoModelForImageClassification, BlipImageProcessor
 from diffusers import DiffusionPipeline, AutoencoderKL
 import torchvision.transforms as transforms
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
 
 from copy import deepcopy
 from collections import OrderedDict
@@ -45,8 +47,10 @@ class BZHStableSignatureDemo(object):
 
         # load the proxy detector
         self.detector_image_processor = BlipImageProcessor.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
-        commit_hash = "584a7bc01dc0f02e53bf8b8b295717ed09ed7294"
-        self.detector_model = AutoModel.from_pretrained("imatag/stable-signature-bzh-detector-resnet18", trust_remote_code=True, revision=commit_hash)
+        self.detector_model = AutoModelForImageClassification.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
+        calibration = hf_hub_download("imatag/stable-signature-bzh-detector-resnet18", filename="calibration.safetensors")
+        with safe_open(calibration, framework="pt") as f:
+            self.calibration_logits = f.get_tensor("logits")
 
     def generate(self, mode, seed, prompt):
         generator = torch.Generator(device=device)
@@ -132,7 +136,9 @@ class BZHStableSignatureDemo(object):
         inputs = self.detector_image_processor(img, return_tensors="pt")
 
         with torch.no_grad():
-            pvalue = torch.sigmoid(self.detector_model(**inputs).logits).item()
+            logit = self.detector_model(**inputs).logits[...,0]
+            pvalue = (1 + torch.searchsorted(self.calibration_logits, logit)) / self.calibration_logits.shape[0]
+            pvalue = pvalue.item()
 
         return pvalue
  

detector_calibration.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5146e98a98a65559fbdb8ccf8cfd9a18982e09e13aa16f3005a640639ca5b298
+size 1999942