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
    prompt = f'{character} {origin} {tags}'
    formatted_tags = tags.replace(" ", ", ")
    formatted_prompt = prompt.replace(" ", ", ")

    # Display the results
    return {
        "Character": character,
        "Origin": origin,
        "Tags": formatted_tags,
        "Prompt": formatted_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 Fetcher",
    description="Enter a Danbooru image URL to fetch character, origin, tags, and prompt information."
)

# Launch the interface
iface.launch()