Update modules/booru.py

As Gelbooru, Rule34 suddenly decided to stop working due to 'API authentication', so I decided to remove them temporarily (or permanently). Either way Xbooru is working, but I don't know for how long. If it fails, I'll try something else. Either way I recommend searching tags using their website instead of API.

modules/booru.py

@@ -1,161 +1,141 @@
-import requests,re,base64,io,numpy as np
-from PIL import Image,ImageOps
-import torch,gradio as gr
-
-# Custom CSS for gallery styling
-css = """
-#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
-#custom-gallery .thumbnail-item {height: var(--row-height);width: 100%;position: relative;overflow: hidden;border-radius: 8px;box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);transition: transform 0.2s ease,  box-shadow 0.2s ease;}
-#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
-#custom-gallery .thumbnail-item img {width: auto;height: 100%;max-width: 100%;max-height: var(--row-height);object-fit: contain;margin: 0 auto;display: block;}
-#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
-#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
-.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
-.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
-#custom-gallery .thumbnail-item img {width: auto;height: 100%;max-width: 100%;max-height: var(--row-height);object-fit: initial;width: fit-content;margin: 0px auto;display: block;}
-"""
-
-# Helper to load image from URL
-def loadImageFromUrl(url):
-    response = requests.get(url, timeout=10)
-    if response.status_code != 200:
-        raise Exception(f"Failed to load image from {url}")
-    i = Image.open(io.BytesIO(response.content))
-    i = ImageOps.exif_transpose(i)
-    if i.mode != "RGBA":
-        i = i.convert("RGBA")
-    alpha = i.split()[-1]
-    image = Image.new("RGB", i.size, (0, 0, 0))
-    image.paste(i, mask=alpha)
-    image = np.array(image).astype(np.float32) / 255.0
-    image = torch.from_numpy(image)[None,]
-    return image
-
-# Fetch data from multiple booru platforms
-def fetch_booru_images(site, Tags, exclude_tags, score, count, Safe, Questionable, Explicit):
-    # Clean and format tags
-    def clean_tag_list(tags):
-        return [item.strip().replace(' ', '_') for item in tags.split(',') if item.strip()]
-
-    Tags = '+'.join(clean_tag_list(Tags)) if Tags else ''
-    exclude_tags = '+'.join('-' + tag for tag in clean_tag_list(exclude_tags))
-
-    rating_filters = []
-    if not Safe:
-        rating_filters.extend(["rating:safe", "rating:general"])
-    if not Questionable:
-        rating_filters.extend(["rating:questionable", "rating:sensitive"])
-    if not Explicit:
-        rating_filters.append("rating:explicit")
-    rating_filters = '+'.join(f'-{r}' for r in rating_filters)
-
-    score_filter = f"score:>{score}"
-
-    # Build query
-    base_query = f"tags=sort:random+{Tags}+{exclude_tags}+{score_filter}+{rating_filters}&limit={count}&json=1"
-    base_query = re.sub(r"\++", "+", base_query)
-
-    # Fetch data based on site
-    if site == "Gelbooru":
-        url = f"https://gelbooru.com/index.php?page=dapi&s=post&q=index&{base_query}"
-        response = requests.get(url).json()
-        posts = response.get("post", [])
-    elif site == "Rule34":
-        url = f"https://api.rule34.xxx/index.php?page=dapi&s=post&q=index&{base_query}"
-        response = requests.get(url).json()
-        posts = response
-    elif site == "Xbooru":
-        url = f"https://xbooru.com/index.php?page=dapi&s=post&q=index&{base_query}"
-        response = requests.get(url).json()
-        posts = response
-    else:
-        return [], [], []
-
-    # Extract image URLs, tags, and post URLs
-    image_urls = []
-    tags_list = [post.get("tags", "").replace(" ", ", ").replace("_", " ").replace("(", "\\(").replace(")", "\\)").strip() for post in posts]
-    post_urls = []
-
-    for post in posts:
-        if site in ["Gelbooru", "Rule34", "Xbooru"]:
-            file_url = post.get("file_url")
-            tags = post.get("tags", "").replace(" ", ", ").strip()
-            post_id = post.get("id", "")
-        else:
-            continue
-
-        if file_url:
-            image_urls.append(file_url)
-            tags_list.append(tags)
-            if site == "Gelbooru":
-                post_urls.append(f"https://gelbooru.com/index.php?page=post&s=view&id={post_id}")
-            elif site == "Rule34":
-                post_urls.append(f"https://rule34.xxx/index.php?page=post&s=view&id={post_id}")
-            elif site == "Xbooru":
-                post_urls.append(f"https://xbooru.com/index.php?page=post&s=view&id={post_id}")
-
-    return image_urls, tags_list, post_urls
-
-# Main function to fetch and return processed images
-def booru_gradio(Tags, exclude_tags, score, count, Safe, Questionable, Explicit, site):
-    image_urls, tags_list, post_urls = fetch_booru_images(site, Tags, exclude_tags, score, count, Safe, Questionable, Explicit)
-
-    if not image_urls:
-        return [], [], [], []
-
-    image_data = []
-    for url in image_urls:
-        try:
-            image = loadImageFromUrl(url)
-            image = (image * 255).clamp(0, 255).cpu().numpy().astype(np.uint8)[0]
-            image = Image.fromarray(image)
-            image_data.append(image)
-        except Exception as e:
-            print(f"Error loading image from {url}: {e}")
-            continue
-
-    return image_data, tags_list, post_urls, image_urls
-
-# Update UI on image click
-def on_select(evt: gr.SelectData, tags_list, post_url_list, image_url_list):
-    idx = evt.index
-    if idx < len(tags_list):
-        return tags_list[idx], post_url_list[idx], image_url_list[idx]
-    return "No tags", "", ""
-
-def create_booru_interface():
-    with gr.Blocks(css=css, fill_width=True) as demo:
-        with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("### ⚙️ Search Parameters")
-                        site = gr.Dropdown(label="Select Source", choices=["Gelbooru (Not working)", "Rule34", "Xbooru"], value="Xbooru")
-                        Tags = gr.Textbox(label="Tags (comma-separated)", placeholder="e.g. solo,  1girl,  1boy,  artist name,  character,  black hair,  granblue fantasy,  ...", lines=3)
-                        exclude_tags = gr.Textbox(label="Exclude Tags (comma-separated)", placeholder="e.g. animated,  watermark,  username,  ...", lines=3)
-                        score = gr.Number(label="Minimum Score", value=0)
-                        count = gr.Slider(label="Number of Images", minimum=1, maximum=20, step=1, value=1)
-                        Safe = gr.Checkbox(label="Include Safe", value=True)
-                        Questionable = gr.Checkbox(label="Include Questionable", value=True)
-                        Explicit = gr.Checkbox(label="Include Explicit (18+)", value=False)
-                        submit_btn = gr.Button("Fetch Images", variant="primary")
-                    with gr.Column():
-                        gr.Markdown("### 📄 Results")
-                        images_output = gr.Gallery(
-                            columns=2,
-                            show_share_button=False,
-                            interactive=True,
-                            height='auto',
-                            label='Grid of images',
-                            preview=False,
-                            elem_id='custom-gallery'
-                        )
-                        tags_output = gr.Textbox(label="Tags", placeholder="Select an image to display tags", lines=6, show_copy_button=True)
-                        post_url_output = gr.Textbox(label="Post URL", lines=2, show_copy_button=True)
-                        image_url_output = gr.Textbox(label="Image URL", lines=2, show_copy_button=True)
-                    # State to store tags, URLs
-                    tags_state = gr.State([])
-                    post_url_state = gr.State([])
-                    image_url_state = gr.State([])
-                    submit_btn.click(fn=booru_gradio, inputs=[Tags, exclude_tags, score, count, Safe, Questionable, Explicit, site], outputs=[images_output, tags_state, post_url_state, image_url_state], )
-                    images_output.select(fn=on_select, inputs=[tags_state, post_url_state, image_url_state], outputs=[tags_output, post_url_output, image_url_output], )
-            
-    return demo
+import requests,re,base64,io,numpy as np
+from PIL import Image,ImageOps
+import torch,gradio as gr
+
+# Custom CSS for gallery styling
+css = """
+#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
+#custom-gallery .thumbnail-item {height: var(--row-height);width: 100%;position: relative;overflow: hidden;border-radius: 8px;box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);transition: transform 0.2s ease,  box-shadow 0.2s ease;}
+#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
+#custom-gallery .thumbnail-item img {width: auto;height: 100%;max-width: 100%;max-height: var(--row-height);object-fit: contain;margin: 0 auto;display: block;}
+#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
+#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
+.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
+.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
+#custom-gallery .thumbnail-item img {width: auto;height: 100%;max-width: 100%;max-height: var(--row-height);object-fit: initial;width: fit-content;margin: 0px auto;display: block;}
+"""
+
+# Helper to load image from URL
+def loadImageFromUrl(url):
+    response = requests.get(url, timeout=10)
+    if response.status_code != 200:
+        raise Exception(f"Failed to load image from {url}")
+    i = Image.open(io.BytesIO(response.content))
+    i = ImageOps.exif_transpose(i)
+    if i.mode != "RGBA":
+        i = i.convert("RGBA")
+    alpha = i.split()[-1]
+    image = Image.new("RGB", i.size, (0, 0, 0))
+    image.paste(i, mask=alpha)
+    image = np.array(image).astype(np.float32) / 255.0
+    image = torch.from_numpy(image)[None,]
+    return image
+
+# Fetch data from Xbooru platform only
+def fetch_booru_images(Tags, exclude_tags, score, count, Safe, Questionable, Explicit):
+    # Clean and format tags
+    def clean_tag_list(tags):
+        return [item.strip().replace(' ', '_') for item in tags.split(',') if item.strip()]
+
+    Tags = '+'.join(clean_tag_list(Tags)) if Tags else ''
+    exclude_tags = '+'.join('-' + tag for tag in clean_tag_list(exclude_tags))
+
+    rating_filters = []
+    if not Safe:
+        rating_filters.extend(["rating:safe", "rating:general"])
+    if not Questionable:
+        rating_filters.extend(["rating:questionable", "rating:sensitive"])
+    if not Explicit:
+        rating_filters.append("rating:explicit")
+    rating_filters = '+'.join(f'-{r}' for r in rating_filters)
+
+    score_filter = f"score:>{score}"
+
+    # Build query
+    base_query = f"tags=sort:random+{Tags}+{exclude_tags}+{score_filter}+{rating_filters}&limit={count}&json=1"
+    base_query = re.sub(r"\++", "+", base_query)
+
+    # Fetch data from Xbooru only atm
+    url = f"https://xbooru.com/index.php?page=dapi&s=post&q=index&{base_query}"
+    response = requests.get(url).json()
+    posts = response
+
+    # Extract image URLs, tags, and post URLs
+    image_urls = []
+    tags_list = [post.get("tags", "").replace(" ", ", ").replace("_", " ").replace("(", "\\(").replace(")", "\\)").strip() for post in posts]
+    post_urls = []
+
+    for post in posts:
+        file_url = post.get("file_url")
+        tags = post.get("tags", "").replace(" ", ", ").strip()
+        post_id = post.get("id", "")
+
+        if file_url:
+            image_urls.append(file_url)
+            tags_list.append(tags)
+            post_urls.append(f"https://xbooru.com/index.php?page=post&s=view&id={post_id}")
+
+    return image_urls, tags_list, post_urls
+
+# Main function to fetch and return processed images
+def booru_gradio(Tags, exclude_tags, score, count, Safe, Questionable, Explicit):
+    image_urls, tags_list, post_urls = fetch_booru_images(Tags, exclude_tags, score, count, Safe, Questionable, Explicit)
+
+    if not image_urls:
+        return [], [], [], []
+
+    image_data = []
+    for url in image_urls:
+        try:
+            image = loadImageFromUrl(url)
+            image = (image * 255).clamp(0, 255).cpu().numpy().astype(np.uint8)[0]
+            image = Image.fromarray(image)
+            image_data.append(image)
+        except Exception as e:
+            print(f"Error loading image from {url}: {e}")
+            continue
+
+    return image_data, tags_list, post_urls, image_urls
+
+# Update UI on image click
+def on_select(evt: gr.SelectData, tags_list, post_url_list, image_url_list):
+    idx = evt.index
+    if idx < len(tags_list):
+        return tags_list[idx], post_url_list[idx], image_url_list[idx]
+    return "No tags", "", ""
+
+def create_booru_interface():
+    with gr.Blocks(css=css, fill_width=True) as demo:
+        with gr.Row():
+                    with gr.Column():
+                        gr.Markdown("### ⚙️ Search Parameters")
+                        Tags = gr.Textbox(label="Tags (comma-separated)", placeholder="e.g. solo,  1girl,  1boy,  artist name,  character,  black hair,  granblue fantasy,  ...", lines=3)
+                        exclude_tags = gr.Textbox(label="Exclude Tags (comma-separated)", placeholder="e.g. animated,  watermark,  username,  ...", lines=3)
+                        score = gr.Number(label="Minimum Score", value=0)
+                        count = gr.Slider(label="Number of Images", minimum=1, maximum=10, step=1, value=1)
+                        Safe = gr.Checkbox(label="Include Safe", value=True)
+                        Questionable = gr.Checkbox(label="Include Questionable", value=True)
+                        Explicit = gr.Checkbox(label="Include Explicit (18+)", value=False)
+                        submit_btn = gr.Button("Fetch Images", variant="primary")
+                    with gr.Column():
+                        gr.Markdown("### 📄 Results")
+                        images_output = gr.Gallery(
+                            columns=2,
+                            show_share_button=False,
+                            interactive=True,
+                            height='auto',
+                            label='Grid of images',
+                            preview=False,
+                            elem_id='custom-gallery'
+                        )
+                        tags_output = gr.Textbox(label="Tags", placeholder="Select an image to display tags", lines=6, show_copy_button=True)
+                        post_url_output = gr.Textbox(label="Post URL", lines=2, show_copy_button=True)
+                        image_url_output = gr.Textbox(label="Image URL", lines=2, show_copy_button=True)
+                    # State to store tags, URLs
+                    tags_state = gr.State([])
+                    post_url_state = gr.State([])
+                    image_url_state = gr.State([])
+                    submit_btn.click(fn=booru_gradio, inputs=[Tags, exclude_tags, score, count, Safe, Questionable, Explicit], outputs=[images_output, tags_state, post_url_state, image_url_state], )
+                    images_output.select(fn=on_select, inputs=[tags_state, post_url_state, image_url_state], outputs=[tags_output, post_url_output, image_url_output], )
+
+    return demo