unfinished app.py

app.py

@@ -1,75 +1,17 @@
-from transformers import pipeline
-from imgutils.data import rgb_encode, load_image
-from onnx_ import _open_onnx_model
-from PIL import Image
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
 import gradio as gr
-import numpy as np
-import os
-import requests
+import timm
 import torch
-import json
 
-def _img_encode(image, size=(384,384), normalize=(0.5,0.5)):
-    image = image.resize(size, Image.BILINEAR)
-    data = rgb_encode(image, order_='CHW')
+nsfw_tf = pipeline("image-classification",
+                     model=AutoModelForImageClassification.from_pretrained(
+                         "carbon225/vit-base-patch16-224-hentai"),
+                     feature_extractor=AutoFeatureExtractor.from_pretrained(
+                         "carbon225/vit-base-patch16-224-hentai"))
 
-    if normalize is not None:
-        mean_, std_ = normalize
-        mean = np.asarray([mean_]).reshape((-1, 1, 1))
-        std = np.asarray([std_]).reshape((-1, 1, 1))
-        data = (data - mean) / std
-
-    return data.astype(np.float32)
-    
-nsfw_tf = pipeline(model="carbon225/vit-base-patch16-224-hentai")
-
-if not os.path.exists("timm.onnx"):
-    open("timm.onnx", "wb").write(
-        requests.get(
-            "https://huggingface.co/deepghs/anime_rating/resolve/main/caformer_s36_plus/model.onnx"
-        ).content
-    )
-    open("timmcfg.json", "wb").write(
-        requests.get(
-            "https://huggingface.co/deepghs/anime_rating/resolve/main/caformer_s36_plus/meta.json"
-        ).content
-    )
-else:
-    print("Model already exists, skipping redownload")
-
-with open("timmcfg.json") as file:
-    tm_cfg = json.load(file)
-
-nsfw_tm = _open_onnx_model("timm.onnx")
+nsfw_tm = timm.create_model('deepghs/anime_rating', pretrained=True).eval()
+tm_config = timm.data.resolve_model_data_config(model)
+tm_trans = timm.data.create_transform(**tm_config, is_training=False)
 
 def launch(img):
-    weight = 0
-    img = img.convert('RGB')
-    tm_image = load_image(img, mode='RGB')
-    tm_input_ = _img_encode(tm_image, size=(256, 256))[None, ...]
-    tm_items, = nsfw_tm.run(['output'], {'input': tm_input_})
-    tm_output = sorted(list(zip(tm_cfg["labels"], map(lambda x: x.item(), tm_items[0]))), key=lambda x: x[1], reverse=True)[0][0]
-    
-    match tm_output:
-        case "safe":
-            weight -= 1
-        case "r15":
-            weight += 2
-        case "r18":
-            weight += 2
-
-    tf_output = nsfw_tf(img)[0]["label"]
-
-    match tf_output:
-        case "safe":
-            weight -= 1
-        case "suggestive":
-            weight += 1
-        case "r18":
-            weight += 2
-
-    print(sorted(list(zip(tm_cfg["labels"], map(lambda x: x.item(), tm_items[0]))), key=lambda x: x[1], reverse=True), tf_output)
-    return weight > 0
-
-app = gr.Interface(fn=launch, inputs="pil", outputs="text")
-app.launch()
+    tm_output = nsfw_tm(transforms(img).unsqueeze(0))

onnx_.py

@@ -1,59 +0,0 @@
-import logging
-import os
-import shutil
-from functools import lru_cache
-from typing import Optional
-
-from hbutils.system import pip_install
-
-
-def _ensure_onnxruntime():
-    try:
-        import onnxruntime
-    except (ImportError, ModuleNotFoundError):
-        logging.warning('Onnx runtime not installed, preparing to install ...')
-        if shutil.which('nvidia-smi'):
-            logging.info('Installing onnxruntime-gpu ...')
-            pip_install(['onnxruntime-gpu'], silent=True)
-        else:
-            logging.info('Installing onnxruntime (cpu) ...')
-            pip_install(['onnxruntime'], silent=True)
-
-
-_ensure_onnxruntime()
-from onnxruntime import get_available_providers, get_all_providers, InferenceSession, SessionOptions, \
-    GraphOptimizationLevel
-
-alias = {
-    'gpu': "CUDAExecutionProvider",
-    "trt": "TensorrtExecutionProvider",
-}
-
-
-def get_onnx_provider(provider: Optional[str] = None):
-    if not provider:
-        if "CUDAExecutionProvider" in get_available_providers():
-            return "CUDAExecutionProvider"
-        else:
-            return "CPUExecutionProvider"
-    elif provider.lower() in alias:
-        return alias[provider.lower()]
-    else:
-        for p in get_all_providers():
-            if provider.lower() == p.lower() or f'{provider}ExecutionProvider'.lower() == p.lower():
-                return p
-
-        raise ValueError(f'One of the {get_all_providers()!r} expected, '
-                         f'but unsupported provider {provider!r} found.')
-
-
-@lru_cache()
-def _open_onnx_model(ckpt: str, provider: str = None) -> InferenceSession:
-    options = SessionOptions()
-    options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
-    provider = provider or get_onnx_provider()
-    if provider == "CPUExecutionProvider":
-        options.intra_op_num_threads = os.cpu_count()
-
-    logging.info(f'Model {ckpt!r} loaded with provider {provider!r}')
-    return InferenceSession(ckpt, options, [provider])

requirements.txt

@@ -1,5 +1,3 @@
 torch
 transformers
-onnxruntime
-dghs-imgutils
-hbutils
+timm