Update app.py

app.py

@@ -1,24 +1,26 @@
 # -*- coding: utf-8 -*-
 """
-Swin-Large AI detection
+Swin/CAFormer AI detection
 -------------------------------------------------------------------
-• V2 / V4                                 : 2-class  (AI vs. Non-AI)
-• V7 / V8 / V9                            : 4-class  (photo / anime × AI / Non-AI)
+• Swin-V2 / V4                          : 2-class  (AI vs. Non-AI)
+• Swin-V7 / V8 / V9                     : 4-class  (photo / anime × AI / Non-AI)
+• CAFormer-V10                          : 4-class  (photo / anime × AI / Non-AI)
 -------------------------------------------------------------------
 Author: telecomadm1145
 """
 
-import os, torch, timm, math, numpy as np
+import os, torch, timm, numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file  # Added for .safetensors support
 
 # --------------------------------------------------
 # 1. Model & Checkpoint Meta-data
 # --------------------------------------------------
-REPO_ID = "telecomadm1145/swin-ai-detection"     # 同一个 repo 存两种 ckpt 也 OK
+REPO_ID = "telecomadm1145/swin-ai-detection"
 HF_FILENAMES = {
     "V2":                 "swin_classifier_stage1_v2_epoch_3.pth",
     "V4":                 "swin_classifier_stage1_v4.pth",
@@ -26,25 +28,27 @@ HF_FILENAMES = {
     "V7":                 "swin_classifier_4class_fp16_v7.pth",
     "V8":                 "swin_classifier_4class_fp16_v8_epoch7_acc9740.pth",
     "V9":                 "swin_classifier_4class_fp16_v9_acc9861.pth",
+    "V1-CAFormer":        "caformer_b36_4class.safetensors",
 }
 
 CKPT_META = {
-    "V2": { "n_cls": 2, "head": "v4",
+    "V2": { "n_cls": 2, "head": "v4", "backbone": "swin_large_patch4_window12_384",
             "names": ["Non-AI Generated", "AI Generated"]},
-    "V4": { "n_cls": 2, "head": "v4",
+    "V4": { "n_cls": 2, "head": "v4", "backbone": "swin_large_patch4_window12_384",
             "names": ["Non-AI Generated", "AI Generated"]},
-    #"V5(underfitting)": { "n_cls": 2, "head": "v5",
+    #"V5(underfitting)": { "n_cls": 2, "head": "v5", "backbone": "swin_large_patch4_window12_384",
     #        "names": ["Non-AI Generated", "AI Generated"]},
-    "V7": { "n_cls": 4, "head": "v7",
+    "V7": { "n_cls": 4, "head": "v7", "backbone": "swin_large_patch4_window12_384",
             "names": ["non_ai", "ai", "ani_non_ai", "ani_ai"]},
-    "V8": { "n_cls": 4, "head": "v7",
+    "V8": { "n_cls": 4, "head": "v7", "backbone": "swin_large_patch4_window12_384",
             "names": ["non_ai", "ai", "ani_non_ai", "ani_ai"]},
-    "V9": { "n_cls": 4, "head": "v7",
+    "V9": { "n_cls": 4, "head": "v7", "backbone": "swin_large_patch4_window12_384",
+            "names": ["non_ai", "ai", "ani_non_ai", "ani_ai"]},
+    "V1-CAFormer": { "n_cls": 4, "head": "v7", "backbone": "caformer_b36.sail_in22k_ft_in1k_384",
             "names": ["non_ai", "ai", "ani_non_ai", "ani_ai"]},
 }
 
-DEFAULT_CKPT = "V9"
-MODEL_NAME  = "swin_large_patch4_window12_384"
+DEFAULT_CKPT = "V1-CAFormer"
 LOCAL_CKPT_DIR = "./checkpoints"
 SEED = 4421
 DROP_RATE = 0.1
@@ -56,6 +60,7 @@ print(f"Using device: {device}")
 model, current_ckpt = None, None
 current_meta = None
 
+# Renamed to ImageClassifier for clarity, but keeping original name to avoid breaking changes if subclassed elsewhere.
 class SwinClassifier(nn.Module):
     def __init__(self, model_name, num_classes, pretrained=True,
                  head_version="v4"):
@@ -66,7 +71,7 @@ class SwinClassifier(nn.Module):
         self.data_config = timm.data.resolve_data_config({}, model=self.backbone)
 
         # ------- 根据版本选择不同 head -------
-        if head_version == "v7":            # <-- V7: 极简 64-hidden, GELU
+        if head_version == "v7":            # <-- V7, V8, V9, V10: 极简 64-hidden, GELU
             self.classifier = nn.Sequential(
                 nn.Dropout(DROP_RATE),
                 nn.Linear(self.backbone.num_features, 64),
@@ -120,23 +125,28 @@ def load_model(ckpt_name: str):
 
     print(f"\n🔄 Switching to {ckpt_name} ...")
     meta = CKPT_META[ckpt_name]
+    ckpt_filename = HF_FILENAMES[ckpt_name]
     ckpt_file = hf_hub_download(
         repo_id=REPO_ID,
-        filename=HF_FILENAMES[ckpt_name],
+        filename=ckpt_filename,
         local_dir=LOCAL_CKPT_DIR, force_download=False
     )
     print(f"Checkpoint: {ckpt_file}")
 
     # Build model structure
     model = SwinClassifier(
-        MODEL_NAME,
+        meta["backbone"],   # Use backbone from meta
         num_classes       = meta["n_cls"],
         pretrained        = False,
         head_version      = meta["head"]
     ).to(device)
 
-    # compatible load
-    state = torch.load(ckpt_file, map_location=device, weights_only=False)
+    # Compatible load for .pth and .safetensors
+    if ckpt_filename.endswith(".safetensors"):
+        state = load_file(ckpt_file, device=device)
+    else:
+        state = torch.load(ckpt_file, map_location=device, weights_only=False)
+
     model.load_state_dict(state.get("model_state_dict", state), strict=True)
     model.eval()
 
@@ -180,7 +190,7 @@ def launch():
     load_model(DEFAULT_CKPT)      # 预加载
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🖼️ Swin-Large — AI / Non-AI  Detector (V2)")
+        gr.Markdown("# AI Detector")
         gr.Markdown(
             "Choose a model checkpoint on the left, upload an image, "
             "and click **Run** to see predictions. Checkpoint V7+ outputs 4 classes."