Update app.py

app.py

@@ -7,16 +7,31 @@ from PIL import Image
 from io import BytesIO
 
 # Initialize models
-anime_model = DiffusionPipeline.from_pretrained("SmilingWolf/wd-v1-4-vit-tagger")
+anime_model = DiffusionPipeline.from_pretrained("SmilingWolf/wd-convnext-tagger-v3")
 photo_model = AutoModelForZeroShotImageClassification.from_pretrained("facebook/florence-base-in21k-retrieval")
 processor = AutoProcessor.from_pretrained("facebook/florence-base-in21k-retrieval")
 
 def get_booru_image(booru, image_id):
-    # This is a placeholder function. You'd need to implement the actual API calls for each booru.
-    url = f"https://api.{booru}.org/images/{image_id}"
+    if booru == "Gelbooru":
+        url = f"https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&id={image_id}"
+    elif booru == "Danbooru":
+        url = f"https://danbooru.donmai.us/posts/{image_id}.json"
+    elif booru == "rule34.xxx":
+        url = f"https://api.rule34.xxx/index.php?page=dapi&s=post&q=index&json=1&id={image_id}"
+    else:
+        raise ValueError("Unsupported booru")
+
     response = requests.get(url)
-    img = Image.open(BytesIO(response.content))
-    tags = ["tag1", "tag2", "tag3"]  # Placeholder
+    data = response.json()
+
+    # The exact structure of the response will vary depending on the booru
+    # You'll need to adjust this part based on each booru's API
+    image_url = data[0]['file_url'] if isinstance(data, list) else data['file_url']
+    tags = data[0]['tags'].split() if isinstance(data, list) else data['tags'].split()
+
+    img_response = requests.get(image_url)
+    img = Image.open(BytesIO(img_response.content))
+
     return img, tags
 
 def transcribe_image(image, image_type, transcriber, booru_tags=None):