feat: add support of cafe_aesthetic

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Plat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,6 +1,6 @@
 ---
-title: Saltacc Beit Bestimage Salt
-emoji: 🐨
+title: Waifu Aesthetics
+emoji: 🧂☕
 colorFrom: yellow
 colorTo: indigo
 sdk: gradio
@@ -9,4 +9,9 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Original repos: 
+
+- https://huggingface.co/cafeai/cafe_aesthetic
+- https://huggingface.co/saltacc/beit-bestimage-salt
+
+About BEiT: https://huggingface.co/docs/transformers/v4.30.0/en/model_doc/beit#beit

app.py

@@ -1,3 +1,63 @@
 import gradio as gr
 
-gr.Interface.load("models/saltacc/beit-bestimage-salt").launch()
+from PIL import Image
+
+from predictor import Predictor
+
+MODEL_SALT_BESTIMAGE = "saltacc/beit-bestimage-salt"
+MODEL_CAFE_AESTHETIC = "cafeai/cafe_aesthetic"
+
+MODEL_NAMES = [
+    MODEL_SALT_BESTIMAGE,
+    MODEL_CAFE_AESTHETIC,
+]
+
+models = {
+    MODEL_SALT_BESTIMAGE: Predictor(MODEL_SALT_BESTIMAGE),
+    MODEL_CAFE_AESTHETIC: Predictor(MODEL_CAFE_AESTHETIC),
+}
+
+
+def predict(image: Image.Image) -> list[dict[str, float]]:
+    results = []
+
+    for model_name in MODEL_NAMES:
+        results.append(models[model_name].predict([image])[0])
+
+    return results
+
+
+def construct_ui():
+    with gr.Blocks() as ui:
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    label="Input Image",
+                    type="pil",
+                    source="upload",
+                    interactive=True,
+                )
+
+                submit_btn = gr.Button(
+                    value="Submit",
+                    variant="primary",
+                )
+
+            with gr.Column():
+                result_salt = gr.Label(
+                    label=MODEL_SALT_BESTIMAGE,
+                )
+
+                result_cafe = gr.Label(
+                    label=MODEL_CAFE_AESTHETIC,
+                )
+
+        submit_btn.click(
+            predict, inputs=[input_image], outputs=[result_salt, result_cafe]
+        )
+
+    return ui
+
+
+if __name__ == "__main__":
+    construct_ui().launch()

predictor.py

@@ -0,0 +1,23 @@
+from PIL import Image
+
+from transformers import BeitImageProcessor, BeitForImageClassification
+
+
+class Predictor:
+    def __init__(self, model_id: str) -> None:
+        self.processor = BeitImageProcessor.from_pretrained(model_id)
+        self.model = BeitForImageClassification.from_pretrained(model_id)
+
+    def predict(self, images: list[Image.Image]) -> list[dict[str, float]]:
+        inputs = self.processor(images, return_tensors="pt")
+        logits = self.model(**inputs).logits.softmax(1)  # 一応見た目が良いのでsoftmaxをかける
+
+        results = []
+
+        for scores in logits:
+            result = {}
+            for i, score in enumerate(scores):
+                result[self.model.config.id2label[i]] = score.item()
+            results.append(result)
+
+        return results

requirements.txt

@@ -0,0 +1,9 @@
+accelerate 
+diffusers 
+transformers 
+scipy
+omegaconf
+safetensors
+numpy
+opencv-python
+torch