Update app.py

app.py

@@ -5,12 +5,29 @@ from diffusers import DiffusionPipeline
 import requests
 from PIL import Image
 from io import BytesIO
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
 
 # Initialize models
-anime_model = DiffusionPipeline.from_pretrained("SmilingWolf/wd-convnext-tagger-v3")
+anime_model_path = hf_hub_download("SmilingWolf/wd-convnext-tagger-v3", "model.onnx")
+anime_model = ort.InferenceSession(anime_model_path)
 photo_model = AutoModelForZeroShotImageClassification.from_pretrained("facebook/florence-base-in21k-retrieval")
 processor = AutoProcessor.from_pretrained("facebook/florence-base-in21k-retrieval")
 
+# Load labels for the anime model
+labels_path = hf_hub_download("SmilingWolf/wd-convnext-tagger-v3", "selected_tags.csv")
+with open(labels_path, 'r') as f:
+    labels = [line.strip().split(',')[0] for line in f.readlines()[1:]]  # Skip header
+
+def preprocess_image(image):
+    image = image.convert('RGB')
+    image = image.resize((448, 448), Image.LANCZOS)
+    image = np.array(image).astype(np.float32)
+    image = image[:, :, ::-1]  # RGB -> BGR
+    image = np.transpose(image, (2, 0, 1))  # HWC -> CHW
+    image = image / 255.0
+    return image[np.newaxis, ...]
+
 def get_booru_image(booru, image_id):
     if booru == "Gelbooru":
         url = f"https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&id={image_id}"
@@ -36,8 +53,14 @@ def get_booru_image(booru, image_id):
 
 def transcribe_image(image, image_type, transcriber, booru_tags=None):
     if image_type == "Anime":
-        with torch.no_grad():
-            tags = anime_model(image)
+        input_image = preprocess_image(image)
+        input_name = anime_model.get_inputs()[0].name
+        output_name = anime_model.get_outputs()[0].name
+        probs = anime_model.run([output_name], {input_name: input_image})[0]
+        
+        # Get top 50 tags
+        top_indices = probs[0].argsort()[-50:][::-1]
+        tags = [labels[i] for i in top_indices]
     else:
         inputs = processor(images=image, return_tensors="pt")
         outputs = photo_model(**inputs)