Update app.py

app.py

@@ -1,17 +1,9 @@
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
+from transformers import pipeline
 import gradio as gr
 import timm
 import torch
 
-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"
-    ),
-)
+nsfw_tf = pipeline(model="carbon225/vit-base-patch16-224-hentai")
 
 if not os.path.exists("timm.ckpt"):
     open("timm.ckpt", "wb").write(
@@ -19,6 +11,11 @@ if not os.path.exists("timm.ckpt"):
             "https://huggingface.co/deepghs/anime_rating/resolve/main/caformer_s36_plus/model.ckpt"
         ).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")
 
@@ -26,11 +23,44 @@ else:
 nsfw_tm = timm.create_model(
     "caformer_s36.sail_in22k_ft_in1k_384",
     checkpoint_path="./timm.ckpt",
-    pretrained=True,
+    pretrained_cfg="./timmcfg.json",
+    pretrained=True
 ).eval()
-tm_config = timm.data.resolve_model_data_config(model)
-tm_trans = timm.data.create_transform(**tm_config, is_training=False)
+tm_config = timm.data.resolve_model_data_config(nsfw_tm.pretrained_cfg, model=nsfw_tm)
+tm_trans = timm.data.create_transform(**tm_config)
 
 
 def launch(img):
-    tm_output = nsfw_tm(transforms(img).unsqueeze(0))
+    weight = 0
+    img = Image.open(img).convert('RGB')
+    tm_output = model.pretrained_cfg['labels'][
+        torch.argmax(
+            torch.nn.functional.softmax(
+                nsfw_tm(transforms(img).unsqueeze(0))[0], dim=0
+            )
+        )
+    ]
+    
+    match tm_output:
+        case "safe":
+            weight -= 2
+        case "r15":
+            weight += 1
+        case "r18":
+            weight += 2
+
+
+    tf_output = nsfw_tf(img)[0]["label"]
+
+    match tf_output:
+        case "safe":
+            weight -= 2
+        case "suggestive":
+            weight += 1
+        case "r18":
+            weight += 2
+
+    return weight > 0
+
+app = gr.Interface(fn=generate, inputs="image", outputs="text")
+app.launch()