app.py

import requests
import json
import gradio as gr


def fetch_tags(image_url):
    # Extract image ID from the provided URL
    try:
        image_id = image_url.split('/')[-1].split('.')[0]
        if not image_id.isdigit():
            return '[ERROR]: Invalid image URL format.', '', '', ''
    except IndexError:
        return '[ERROR]: Invalid image URL format.', '', '', ''

    base_url = 'https://danbooru.donmai.us/posts'
    response = requests.get(f'{base_url}/{image_id}.json')

    if response.status_code != 200:
        return f'[ERROR]: {response.status_code} - Failed to retrieve data.', '', '', ''

    data = json.loads(response.text)

    # Extract required fields
    character = data.get('tag_string_character', 'N/A')
    origin = data.get('tag_string_copyright', 'N/A')
    tags = data.get('tag_string_general', '')

    # Prepare prompt and formatted tags
    formatted_tags = tags.replace(" ", ", ")
    prompt = f'{character}, {origin}, {formatted_tags}'

    return character, origin, formatted_tags, prompt


# Create a Gradio interface
iface = gr.Interface(
    fn=fetch_tags,
    inputs=gr.Textbox(label="Danbooru Image URL"),
    outputs=[
        gr.Textbox(label="Character"),
        gr.Textbox(label="Origin"),
        gr.Textbox(label="Tags (comma-separated)"),
        gr.Textbox(label="Prompt (comma-separated)")
    ],
    title="Danbooru Tag Extractor",
    description="Enter a Danbooru image URL to fetch and format character, origin, tags, and prompt information."
)

# Launch the interface
iface.launch()