download model

app.py

@@ -3,15 +3,34 @@ 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(
+    "image-classification",
+    model=AutoModelForImageClassification.from_pretrained(
+        "carbon225/vit-base-patch16-224-hentai"
+    ),
+    feature_extractor=AutoFeatureExtractor.from_pretrained(
+        "carbon225/vit-base-patch16-224-hentai"
+    ),
+)
 
-nsfw_tm = timm.create_model('deepghs/anime_rating', pretrained=True).eval()
+if not os.path.exists("timm.ckpt"):
+    open("timm.ckpt", "wb").write(
+        requests.get(
+            "https://huggingface.co/deepghs/anime_rating/resolve/main/caformer_s36_plus/model.ckpt"
+        ).content
+    )
+else:
+    print("Model already exists, skipping redownload")
+
+
+nsfw_tm = timm.create_model(
+    "caformer_s36.sail_in22k_ft_in1k_384",
+    checkpoint_path="./timm.ckpt",
+    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):
-    tm_output = nsfw_tm(transforms(img).unsqueeze(0))
+    tm_output = nsfw_tm(transforms(img).unsqueeze(0))