Updated fawkes to hopefully work on gradio zerogpu

app.py

@@ -1,76 +1,125 @@
 from fawkes.protection import Fawkes
 from fawkes.utils import Faces, reverse_process_cloaked
 from fawkes.differentiator import FawkesMaskGeneration
-import tensorflow as tf
+from keras import utils
 import numpy as np
 import gradio as gr
+import spaces
 
 IMG_SIZE = 112
-PREPROCESS = 'raw'
+PREPROCESS = "raw"
+
+
+def get_extractors():
+    hash_map = {
+        "extractor_2": "ce703d481db2b83513bbdafa27434703",
+        "extractor_0": "94854151fd9077997d69ceda107f9c6b",
+    }
+    for key, value in hash_map.items():
+        utils.get_file(
+            fname="{}.h5".format(key),
+            origin="http://mirror.cs.uchicago.edu/fawkes/files/{}.h5".format(key),
+            md5_hash=value,
+            cache_subdir="model",
+        )
 
-# To pre-emptively download the files at boot
-fwks_l = Fawkes("extractor_2", '0', 1, mode='low')
-fwks_m = Fawkes("extractor_2", '0', 1, mode='mid')
-fwks_h = Fawkes("extractor_2", '0', 1, mode='high')
 
 def generate_cloak_images(protector, image_X, target_emb=None):
     cloaked_image_X = protector.compute(image_X, target_emb)
     return cloaked_image_X
 
-def predict(img, level, th=0.04, sd=1e7, lr=10, max_step=500, batch_size=1, format='png',
-                       separate_target=True, debug=False, no_align=False, exp="", maximize=True,
-                       save_last_on_failed=True, progress=gr.Progress(track_tqdm=True)):
-  
-  img = img.convert('RGB')
-  img = tf.keras.utils.img_to_array(img)
-
-  if level == 'low':
-    fwks = fwks_l
-  elif level == 'mid':
-    fwks = fwks_m
-  elif level == 'high':
-    fwks = fwks_h
-
-  current_param = "-".join([str(x) for x in [fwks.th, sd, fwks.lr, fwks.max_step, batch_size, format,
-                                              separate_target, debug]])
-  faces = Faces(['./Current Face'], [img], fwks.aligner, verbose=0, no_align=False)
-  original_images = faces.cropped_faces
-
-  if len(original_images) == 0:
-      raise Exception("No face detected. ")
-  original_images = np.array(original_images)
-
-  if current_param != fwks.protector_param:
-      fwks.protector_param = current_param
-      if fwks.protector is not None:
-          del fwks.protector
-      if batch_size == -1:
-          batch_size = len(original_images)
-      fwks.protector = FawkesMaskGeneration(fwks.feature_extractors_ls,
-                                            batch_size=batch_size,
-                                            mimic_img=True,
-                                            intensity_range=PREPROCESS,
-                                            initial_const=sd,
-                                            learning_rate=fwks.lr,
-                                            max_iterations=fwks.max_step,
-                                            l_threshold=fwks.th,
-                                            verbose=0,
-                                            maximize=maximize,
-                                            keep_final=False,
-                                            image_shape=(IMG_SIZE, IMG_SIZE, 3),
-                                            loss_method='features',
-                                            tanh_process=True,
-                                            save_last_on_failed=save_last_on_failed,
-                                            )
-  protected_images = generate_cloak_images(fwks.protector, original_images)
-  faces.cloaked_cropped_faces = protected_images
-
-  final_images, _ = faces.merge_faces(
-      reverse_process_cloaked(protected_images, preprocess=PREPROCESS),
-      reverse_process_cloaked(original_images, preprocess=PREPROCESS))
-
-  return final_images[-1].astype(np.uint8)
-
-gr.Interface(fn=predict, inputs=[gr.components.Image(type='pil'),
-                                 gr.components.Radio(["low", "mid", "high"], label="Protection Level")],
-                                 outputs=gr.components.Image(type="numpy"), allow_flagging="never").launch(show_error=True, quiet=False)
+
+def preproc(img):
+    img = img.convert("RGB")
+    img = utils.img_to_array(img)
+    return img
+
+
+@spaces.GPU
+def predict(
+    img,
+    level,
+    sd=1e7,
+    format="png",
+    separate_target=True,
+    debug=False,
+    maximize=True,
+    save_last_on_failed=True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    img = preproc(img)
+
+    if level == "low":
+        fwks = Fawkes("extractor_2", 1, mode="low")
+    elif level == "mid":
+        fwks = Fawkes("extractor_2", 1, mode="mid")
+    elif level == "high":
+        fwks = Fawkes("extractor_2", 1, mode="high")
+
+    current_param = "-".join(
+        [
+            str(x)
+            for x in [
+                fwks.th,
+                sd,
+                fwks.lr,
+                fwks.max_step,
+                -1,
+                format,
+                separate_target,
+                debug,
+            ]
+        ]
+    )
+    faces = Faces(["./Current Face"], [img], fwks.aligner, verbose=0, no_align=False)
+    original_images = faces.cropped_faces
+
+    if len(original_images) == 0:
+        raise Exception("No face detected. ")
+    original_images = np.array(original_images)
+
+    if current_param != fwks.protector_param:
+        fwks.protector_param = current_param
+        if fwks.protector is not None:
+            del fwks.protector
+        batch_size = len(original_images)
+        fwks.protector = FawkesMaskGeneration(
+            fwks.feature_extractors_ls,
+            batch_size=batch_size,
+            mimic_img=True,
+            intensity_range=PREPROCESS,
+            initial_const=sd,
+            learning_rate=fwks.lr,
+            max_iterations=fwks.max_step,
+            l_threshold=fwks.th,
+            verbose=0,
+            maximize=maximize,
+            keep_final=False,
+            image_shape=(IMG_SIZE, IMG_SIZE, 3),
+            loss_method="features",
+            tanh_process=True,
+            save_last_on_failed=save_last_on_failed,
+        )
+    protected_images = generate_cloak_images(fwks.protector, original_images)
+    faces.cloaked_cropped_faces = protected_images
+
+    final_images, _ = faces.merge_faces(
+        reverse_process_cloaked(protected_images, preprocess=PREPROCESS),
+        reverse_process_cloaked(original_images, preprocess=PREPROCESS),
+    )
+
+    return final_images[-1].astype(np.uint8)
+
+
+# Download extractors pre-emptively
+get_extractors()
+
+gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.components.Image(type="pil"),
+        gr.components.Radio(["low", "mid", "high"], label="Protection Level"),
+    ],
+    outputs=gr.components.Image(type="pil"),
+    allow_flagging="never",
+).launch(show_error=True, quiet=False)

requirements.txt

@@ -1,95 +1,3 @@
-absl-py==2.1.0
-aiofiles==23.2.1
-altair==5.3.0
-annotated-types==0.7.0
-anyio==3.7.1
-astunparse==1.6.3
-attrs==23.2.0
-cachetools==5.3.3
-certifi==2024.2.2
-charset-normalizer==3.3.2
-click==8.1.7
-contourpy==1.2.1
-cycler==0.12.1
-fastapi==0.103.2
 fawkes @ git+https://github.com/logasja/fawkes@master
-ffmpy==0.3.2
-filelock==3.14.0
-flatbuffers==24.3.25
-fonttools==4.51.0
-fsspec==2024.5.0
-gast==0.4.0
-google-auth==2.29.0
-google-auth-oauthlib==1.0.0
-google-pasta==0.2.0
-gradio==4.31.5
-gradio_client==0.16.4
-grpcio==1.64.0
-h11==0.14.0
-h5py==3.11.0
-httpcore==1.0.5
-httpx==0.27.0
-huggingface-hub==0.23.1
-idna==3.7
-importlib_resources==6.4.0
-Jinja2==3.1.4
-jsonschema==4.22.0
-jsonschema-specifications==2023.12.1
-keras==2.13.1
-kiwisolver==1.4.5
-libclang==18.1.1
-Markdown==3.6
-markdown-it-py==3.0.0
-MarkupSafe==2.1.5
-matplotlib==3.9.0
-mdurl==0.1.2
-mtcnn==0.1.1
-numpy==1.24.3
-oauthlib==3.2.2
-opencv-python==4.9.0.80
-opt-einsum==3.3.0
-orjson==3.10.3
-packaging==24.0
-pandas==2.2.2
-Pillow==10.0.0
-protobuf==4.25.3
-pyasn1==0.6.0
-pyasn1_modules==0.4.0
-pydantic==2.7.1
-pydantic_core==2.18.2
-pydub==0.25.1
-Pygments==2.18.0
-pyparsing==3.1.2
-python-dateutil==2.9.0.post0
-python-multipart==0.0.9
-pytz==2024.1
-PyYAML==6.0.1
-referencing==0.35.1
-requests==2.32.2
-requests-oauthlib==2.0.0
-rich==13.7.1
-rpds-py==0.18.1
-rsa==4.9
-ruff==0.4.4
-semantic-version==2.10.0
-shellingham==1.5.4
-six==1.16.0
-sniffio==1.3.1
-starlette==0.27.0
-tensorboard==2.13.0
-tensorboard-data-server==0.7.2
-tensorflow==2.13.0
-tensorflow-estimator==2.13.0
-tensorflow-io-gcs-filesystem==0.37.0
-termcolor==2.4.0
-tomlkit==0.12.0
-toolz==0.12.1
-tqdm==4.66.4
-typer==0.12.3
-typing_extensions==4.11.0
-tzdata==2024.1
-urllib3==2.2.1
-uvicorn==0.29.0
-websockets==11.0.3
-Werkzeug==3.0.3
-wrapt==1.16.0
+gradio==4.38.1
+spaces>=0.28.3