app.py

"""
JoyCaption Batch WebUI — local Gradio UI for JoyCaption Beta One + JoyTag,
with batch dataset captioning (writes .txt files next to images, kohya-style).

Models by fancyfeast (fpgaminer): https://github.com/fpgaminer/joycaption (Apache-2.0)
UI based on the HF spaces fancyfeast/joy-caption-beta-one and MegaTronX/JoyCaptionBetaOne.

Defaults target a 24GB GPU (bf16 + SDPA + TF32); run with --nf4 for 4-bit
(~7GB VRAM) if you need to keep a diffusion UI like ComfyUI loaded alongside.
"""

import argparse
import gc
import os
from pathlib import Path
from threading import Thread
from typing import Generator

import torch
import gradio as gr
from PIL import Image
from transformers import LlavaForConditionalGeneration, TextIteratorStreamer, AutoProcessor

parser = argparse.ArgumentParser()
parser.add_argument("--nf4", action="store_true", help="Load model in 4-bit NF4 (~7GB VRAM instead of ~18GB)")
parser.add_argument("--port", type=int, default=7861, help="Port (default 7861, avoids ComfyUI's 8188/7860)")
parser.add_argument("--listen", action="store_true", help="Listen on 0.0.0.0 (LAN access)")
parser.add_argument("--cache-dir", type=str, default=None, help="Model download/cache dir (default: ./models next to this file)")
parser.add_argument("--no-browser", action="store_true", help="Don't auto-open a browser tab on launch")
args = parser.parse_args()

ON_SPACES = os.environ.get("SPACE_ID") is not None
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
CACHE_DIR = args.cache_dir or str(Path(__file__).parent / "models")

# Ampere+ optimizations
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

TITLE = """<h1 style="text-align:center;">JoyCaption Batch WebUI</h1>
<p style="text-align:center;">JoyCaption Beta One + JoyTag &mdash; single-image and batch dataset captioning</p>
<hr>"""

CAPTION_TYPE_MAP = {
	"Descriptive": [
		"Write a detailed description for this image.",
		"Write a detailed description for this image in {word_count} words or less.",
		"Write a {length} detailed description for this image.",
	],
	"Descriptive (Casual)": [
		"Write a descriptive caption for this image in a casual tone.",
		"Write a descriptive caption for this image in a casual tone within {word_count} words.",
		"Write a {length} descriptive caption for this image in a casual tone.",
	],
	"Straightforward": [
		"Write a straightforward caption for this image. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with \"This image is…\" or similar phrasing.",
		"Write a straightforward caption for this image within {word_count} words. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with \"This image is…\" or similar phrasing.",
		"Write a {length} straightforward caption for this image. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with \"This image is…\" or similar phrasing.",
	],
	"Stable Diffusion Prompt": [
		"Output a stable diffusion prompt that is indistinguishable from a real stable diffusion prompt.",
		"Output a stable diffusion prompt that is indistinguishable from a real stable diffusion prompt. {word_count} words or less.",
		"Output a {length} stable diffusion prompt that is indistinguishable from a real stable diffusion prompt.",
	],
	"MidJourney": [
		"Write a MidJourney prompt for this image.",
		"Write a MidJourney prompt for this image within {word_count} words.",
		"Write a {length} MidJourney prompt for this image.",
	],
	"Danbooru tag list": [
		"Generate only comma-separated Danbooru tags (lowercase_underscores). Strict order: `artist:`, `copyright:`, `character:`, `meta:`, then general tags. Include counts (1girl), appearance, clothing, accessories, pose, expression, actions, background. Use precise Danbooru syntax. No extra text.",
		"Generate only comma-separated Danbooru tags (lowercase_underscores). Strict order: `artist:`, `copyright:`, `character:`, `meta:`, then general tags. Include counts (1girl), appearance, clothing, accessories, pose, expression, actions, background. Use precise Danbooru syntax. No extra text. {word_count} words or less.",
		"Generate only comma-separated Danbooru tags (lowercase_underscores). Strict order: `artist:`, `copyright:`, `character:`, `meta:`, then general tags. Include counts (1girl), appearance, clothing, accessories, pose, expression, actions, background. Use precise Danbooru syntax. No extra text. {length} length.",
	],
	"e621 tag list": [
		"Write a comma-separated list of e621 tags in alphabetical order for this image. Start with the artist, copyright, character, species, meta, and lore tags (if any), prefixed by 'artist:', 'copyright:', 'character:', 'species:', 'meta:', and 'lore:'. Then all the general tags.",
		"Write a comma-separated list of e621 tags in alphabetical order for this image. Start with the artist, copyright, character, species, meta, and lore tags (if any), prefixed by 'artist:', 'copyright:', 'character:', 'species:', 'meta:', and 'lore:'. Then all the general tags. Keep it under {word_count} words.",
		"Write a {length} comma-separated list of e621 tags in alphabetical order for this image. Start with the artist, copyright, character, species, meta, and lore tags (if any), prefixed by 'artist:', 'copyright:', 'character:', 'species:', 'meta:', and 'lore:'. Then all the general tags.",
	],
	"Rule34 tag list": [
		"Write a comma-separated list of rule34 tags in alphabetical order for this image. Start with the artist, copyright, character, and meta tags (if any), prefixed by 'artist:', 'copyright:', 'character:', and 'meta:'. Then all the general tags.",
		"Write a comma-separated list of rule34 tags in alphabetical order for this image. Start with the artist, copyright, character, and meta tags (if any), prefixed by 'artist:', 'copyright:', 'character:', and 'meta:'. Then all the general tags. Keep it under {word_count} words.",
		"Write a {length} comma-separated list of rule34 tags in alphabetical order for this image. Start with the artist, copyright, character, and meta tags (if any), prefixed by 'artist:', 'copyright:', 'character:', and 'meta:'. Then all the general tags.",
	],
	"Booru-like tag list": [
		"Write a list of Booru-like tags for this image.",
		"Write a list of Booru-like tags for this image within {word_count} words.",
		"Write a {length} list of Booru-like tags for this image.",
	],
	"Art Critic": [
		"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.",
		"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it within {word_count} words.",
		"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it {length}.",
	],
	"Product Listing": [
		"Write a caption for this image as though it were a product listing.",
		"Write a caption for this image as though it were a product listing. Keep it under {word_count} words.",
		"Write a {length} caption for this image as though it were a product listing.",
	],
	"Social Media Post": [
		"Write a caption for this image as if it were being used for a social media post.",
		"Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.",
		"Write a {length} caption for this image as if it were being used for a social media post.",
	],
}

NAME_OPTION = "If there is a person/character in the image you must refer to them as {name}."

EXTRA_OPTIONS_LIST = [
	"Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
	"Include information about lighting.",
	"Include information about camera angle.",
	"Include information about whether there is a watermark or not.",
	"Include information about whether there are JPEG artifacts or not.",
	"If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.",
	"Do NOT include anything sexual; keep it PG.",
	"Do NOT mention the image's resolution.",
	"You MUST include information about the subjective aesthetic quality of the image from low to very high.",
	"Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.",
	"Do NOT mention any text that is in the image.",
	"Specify the depth of field and whether the background is in focus or blurred.",
	"If applicable, mention the likely use of artificial or natural lighting sources.",
	"Do NOT use any ambiguous language.",
	"Include whether the image is sfw, suggestive, or nsfw.",
	"ONLY describe the most important elements of the image.",
]

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}

processor = None
model = None

JOYTAG_REPO = "fancyfeast/joytag"
joytag_model = None
joytag_tags = None


def load_model() -> str:
	global processor, model
	if model is not None:
		return "✅ Model is already loaded."
	print("Loading processor...")
	processor = AutoProcessor.from_pretrained(MODEL_PATH, cache_dir=CACHE_DIR)
	print("Loading model..." + (" (4-bit NF4)" if args.nf4 else " (bf16)"))
	load_kwargs = dict(
		torch_dtype=torch.bfloat16,
		attn_implementation="sdpa",
		cache_dir=CACHE_DIR,
	)
	if args.nf4:
		from transformers import BitsAndBytesConfig
		load_kwargs["quantization_config"] = BitsAndBytesConfig(
			load_in_4bit=True,
			bnb_4bit_quant_type="nf4",
			bnb_4bit_compute_dtype=torch.bfloat16,
			bnb_4bit_use_double_quant=True,
		)
		load_kwargs["device_map"] = DEVICE
	else:
		load_kwargs["device_map"] = DEVICE
	model = LlavaForConditionalGeneration.from_pretrained(MODEL_PATH, **load_kwargs)
	model.eval()
	print("Model loaded and ready!")
	return "✅ Model loaded and ready!" + (" (4-bit NF4)" if args.nf4 else " (bf16)")


def unload_model() -> str:
	global processor, model
	if model is None:
		return "⚠️ Model is not currently loaded."
	del model
	del processor
	model = None
	processor = None
	gc.collect()
	torch.cuda.empty_cache()
	return "✅ Model unloaded from VRAM."


def load_joytag() -> str:
	global joytag_model, joytag_tags
	if joytag_model is not None:
		return "✅ JoyTag is already loaded."
	import huggingface_hub
	from joytag_models import VisionModel
	print("Loading JoyTag...")
	path = huggingface_hub.snapshot_download(JOYTAG_REPO, cache_dir=CACHE_DIR)
	joytag_model = VisionModel.load_model(path, device="cuda" if torch.cuda.is_available() else "cpu")
	joytag_model.eval()
	joytag_tags = [line.strip() for line in (Path(path) / "top_tags.txt").read_text(encoding="utf-8").splitlines() if line.strip()]
	print("JoyTag loaded!")
	return "✅ JoyTag loaded (~0.4GB VRAM)."


def unload_joytag() -> str:
	global joytag_model, joytag_tags
	if joytag_model is None:
		return "⚠️ JoyTag is not currently loaded."
	joytag_model = None
	joytag_tags = None
	gc.collect()
	torch.cuda.empty_cache()
	return "✅ JoyTag unloaded from VRAM."


def prepare_joytag_image(image: Image.Image, target_size: int) -> torch.Tensor:
	import torchvision.transforms.functional as TVF
	# Pad to square on white, resize, CLIP-normalize (mirrors the official space)
	max_dim = max(image.size)
	padded = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
	padded.paste(image, ((max_dim - image.size[0]) // 2, (max_dim - image.size[1]) // 2))
	if max_dim != target_size:
		padded = padded.resize((target_size, target_size), Image.BICUBIC)
	tensor = TVF.pil_to_tensor(padded) / 255.0
	return TVF.normalize(tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])


@torch.no_grad()
def joytag_predict(image: Image.Image, threshold: float) -> tuple[list[tuple[str, float]], dict[str, float]]:
	"""Returns (tags over threshold sorted by confidence, all scores)."""
	if joytag_model is None:
		load_joytag()
	device = next(joytag_model.parameters()).device
	tensor = prepare_joytag_image(image.convert("RGB"), joytag_model.image_size).unsqueeze(0).to(device)
	with torch.autocast(device.type, enabled=True):
		preds = joytag_model({"image": tensor})
		tag_preds = preds["tags"].sigmoid().float().cpu()[0]
	scores = {joytag_tags[i]: tag_preds[i].item() for i in range(len(joytag_tags))}
	predicted = sorted(((t, s) for t, s in scores.items() if s > threshold), key=lambda x: -x[1])
	return predicted, scores


def format_joytag_tags(predicted: list[tuple[str, float]], spaces: bool) -> str:
	tags = [t for t, _ in predicted]
	if spaces:
		tags = [t.replace("_", " ") for t in tags]
	return ", ".join(tags)


def joytag_single(input_image: Image.Image, threshold: float, spaces: bool):
	if input_image is None:
		return "Please upload an image first.", ""
	predicted, scores = joytag_predict(input_image, threshold)
	top = sorted(scores.items(), key=lambda x: -x[1])[:50]
	score_text = "\n".join(f"{s:.3f}  {t}" for t, s in top)
	return format_joytag_tags(predicted, spaces), score_text


def joytag_batch(
	folder: str,
	threshold: float,
	spaces: bool,
	skip_existing: bool,
	prefix: str,
	progress=gr.Progress(),
) -> Generator[str, None, None]:
	folder_path = Path(folder.strip().strip('"'))
	if not folder_path.is_dir():
		yield f"❌ Not a folder: {folder_path}"
		return
	if joytag_model is None:
		yield "Loading JoyTag..."
		load_joytag()

	images = sorted(p for p in folder_path.iterdir() if p.suffix.lower() in IMAGE_EXTS)
	if not images:
		yield f"❌ No images found in {folder_path}"
		return

	log = []
	done = skipped = failed = 0
	for i, img_path in enumerate(progress.tqdm(images, desc="Tagging")):
		txt_path = img_path.with_suffix(".txt")
		if skip_existing and txt_path.exists():
			skipped += 1
			continue
		try:
			predicted, _ = joytag_predict(Image.open(img_path), threshold)
			tag_string = format_joytag_tags(predicted, spaces)
			if prefix.strip():
				tag_string = prefix.strip() + " " + tag_string
			txt_path.write_text(tag_string, encoding="utf-8")
			done += 1
			log.append(f"[{i + 1}/{len(images)}] {img_path.name} ✓ ({len(predicted)} tags)")
		except Exception as e:
			failed += 1
			log.append(f"[{i + 1}/{len(images)}] {img_path.name} ❌ {e}")
		yield "\n".join(log[-30:]) + f"\n\n— {done} tagged, {skipped} skipped, {failed} failed —"

	yield "\n".join(log[-30:]) + f"\n\n✅ Finished: {done} tagged, {skipped} skipped, {failed} failed ({len(images)} images total)"


def build_prompt(caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str) -> str:
	if caption_length == "any":
		map_idx = 0
	elif isinstance(caption_length, str) and caption_length.isdigit():
		map_idx = 1  # numeric word count template
	else:
		map_idx = 2  # length descriptor template

	prompt = CAPTION_TYPE_MAP[caption_type][map_idx]

	if extra_options:
		prompt += " " + " ".join(extra_options)

	return prompt.format(name=name_input, length=caption_length, word_count=caption_length)


def toggle_name_box(extra_options: list[str]):
	return gr.update(visible=NAME_OPTION in extra_options)


def prepare_inputs(input_image: Image.Image, prompt: str):
	convo = [
		{"role": "system", "content": "You are a helpful image captioner."},
		{"role": "user", "content": prompt},
	]
	convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
	inputs = processor(text=[convo_string], images=[input_image.convert("RGB")], return_tensors="pt").to(model.device)
	inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
	return inputs


def generation_kwargs(temperature: float, top_p: float, max_new_tokens: int) -> dict:
	return dict(
		max_new_tokens=int(max_new_tokens),
		do_sample=temperature > 0,
		suppress_tokens=None,
		use_cache=True,
		temperature=temperature if temperature > 0 else None,
		top_k=None,
		top_p=top_p if temperature > 0 else None,
	)


@torch.no_grad()
def chat_joycaption(
	input_image: Image.Image,
	prompt: str,
	temperature: float,
	top_p: float,
	max_new_tokens: int,
) -> Generator[str, None, None]:
	if input_image is None:
		yield "Please upload an image first."
		return
	if model is None:
		load_model()

	torch.cuda.empty_cache()
	inputs = prepare_inputs(input_image, prompt)

	streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
	gen_kwargs = dict(**inputs, streamer=streamer, **generation_kwargs(temperature, top_p, max_new_tokens))

	thread = Thread(target=model.generate, kwargs=gen_kwargs)
	thread.start()

	output = ""
	for new_text in streamer:
		output += new_text
		yield output.strip()

	thread.join()


@torch.no_grad()
def caption_single(input_image: Image.Image, prompt: str, temperature: float, top_p: float, max_new_tokens: int) -> str:
	inputs = prepare_inputs(input_image, prompt)
	generate_ids = model.generate(**inputs, **generation_kwargs(temperature, top_p, max_new_tokens))[0]
	generate_ids = generate_ids[inputs["input_ids"].shape[1]:]
	return processor.tokenizer.decode(generate_ids, skip_special_tokens=True).strip()


def batch_caption(
	folder: str,
	prompt: str,
	temperature: float,
	top_p: float,
	max_new_tokens: int,
	skip_existing: bool,
	prefix: str,
	progress=gr.Progress(),
) -> Generator[str, None, None]:
	folder_path = Path(folder.strip().strip('"'))
	if not folder_path.is_dir():
		yield f"❌ Not a folder: {folder_path}"
		return
	if model is None:
		yield "Loading model..."
		load_model()

	images = sorted(p for p in folder_path.iterdir() if p.suffix.lower() in IMAGE_EXTS)
	if not images:
		yield f"❌ No images found in {folder_path}"
		return

	log = []
	done = skipped = failed = 0
	for i, img_path in enumerate(progress.tqdm(images, desc="Captioning")):
		txt_path = img_path.with_suffix(".txt")
		if skip_existing and txt_path.exists():
			skipped += 1
			continue
		try:
			image = Image.open(img_path)
			caption = caption_single(image, prompt, temperature, top_p, max_new_tokens)
			if prefix.strip():
				caption = prefix.strip() + " " + caption
			txt_path.write_text(caption, encoding="utf-8")
			done += 1
			log.append(f"[{i + 1}/{len(images)}] {img_path.name} ✓")
		except Exception as e:
			failed += 1
			log.append(f"[{i + 1}/{len(images)}] {img_path.name} ❌ {e}")
		yield "\n".join(log[-30:]) + f"\n\n— {done} captioned, {skipped} skipped, {failed} failed —"

	yield "\n".join(log[-30:]) + f"\n\n✅ Finished: {done} captioned, {skipped} skipped, {failed} failed ({len(images)} images total)"


# ─── Gradio UI ───────────────────────────────────────────────────────────────

# gradio 6.12's tab bar moves tabs that don't fit into an "overflow" dropdown, and
# selecting a tab from that dropdown hangs the page in a layout loop. Letting the tab
# bar wrap keeps scrollWidth <= clientWidth so the overflow logic never engages.
TAB_OVERFLOW_FIX_CSS = """
.tab-container { flex-wrap: wrap !important; }
"""

with gr.Blocks(title="JoyCaption Batch WebUI", css=TAB_OVERFLOW_FIX_CSS) as demo:
	gr.HTML(TITLE)

	with gr.Row():
		load_btn = gr.Button("⚡ Load Model", variant="primary")
		unload_btn = gr.Button("🗑️ Unload Model from VRAM", variant="stop")
	model_status = gr.Textbox(label="Model Status", value="Model loads automatically on first caption.", interactive=False)
	load_btn.click(load_model, outputs=model_status)
	unload_btn.click(unload_model, outputs=model_status)

	# Shared prompt controls
	with gr.Row():
		with gr.Column():
			caption_type = gr.Dropdown(
				choices=list(CAPTION_TYPE_MAP.keys()),
				label="Caption Type",
				value="Descriptive",
			)
			caption_length = gr.Dropdown(
				choices=["any", "very short", "short", "medium-length", "long", "very long"]
				+ [str(i) for i in range(20, 261, 10)],
				label="Caption Length",
				value="long",
			)
			with gr.Accordion("Extra Options", open=False):
				extra_options = gr.CheckboxGroup(
					choices=[NAME_OPTION] + EXTRA_OPTIONS_LIST,
					label="Extra Options",
				)
				name_input = gr.Textbox(
					label="Person/Character Name (if applicable)",
					visible=False,
				)
			with gr.Accordion("Generation Settings", open=False):
				temperature_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.6, step=0.05, label="Temperature")
				top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.01, label="Top-p")
				max_tokens_slider = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens")

		with gr.Column():
			# Initial value set here directly: a demo.load() handler for this hangs the
			# page when switching to a third tab (gradio 6.12 bug, see repro notes).
			prompt_box = gr.Textbox(
				lines=4, label="Prompt", interactive=True,
				value=build_prompt("Descriptive", "long", [], ""),
			)

	extra_options.change(toggle_name_box, inputs=extra_options, outputs=name_input)
	for ctrl in (caption_type, caption_length, extra_options, name_input):
		ctrl.change(
			build_prompt,
			inputs=[caption_type, caption_length, extra_options, name_input],
			outputs=prompt_box,
		)

	with gr.Tab("Single Image"):
		with gr.Row():
			with gr.Column():
				input_image = gr.Image(type="pil", label="Input Image")
				run_button = gr.Button("Caption", variant="primary")
			with gr.Column():
				output_caption = gr.Textbox(label="Caption", lines=10)
				copy_btn = gr.Button("📋 Copy to Clipboard")
				copy_btn.click(None, output_caption, None, js="(text) => navigator.clipboard.writeText(text)")

		run_button.click(
			chat_joycaption,
			inputs=[input_image, prompt_box, temperature_slider, top_p_slider, max_tokens_slider],
			outputs=output_caption,
		)

	with gr.Tab("Batch Folder"):
		gr.Markdown("Dataset captioning: captions every image in a folder and writes a `.txt` with the same name next to each image (kohya/diffusion-pipe style).")
		folder_box = gr.Textbox(label="Folder Path", placeholder=r"C:\path\to\my_dataset")
		with gr.Row():
			skip_existing_cb = gr.Checkbox(value=True, label="Skip images that already have a .txt")
			prefix_box = gr.Textbox(label="Caption Prefix (optional, e.g. trigger word)", placeholder="myTriggerWord,")
		batch_button = gr.Button("Caption Folder", variant="primary")
		batch_log = gr.Textbox(label="Progress", lines=15)

		batch_button.click(
			batch_caption,
			inputs=[folder_box, prompt_box, temperature_slider, top_p_slider, max_tokens_slider, skip_existing_cb, prefix_box],
			outputs=batch_log,
		)

	with gr.Tab("JoyTag"):
		gr.Markdown(
			"[JoyTag](https://huggingface.co/fancyfeast/joytag) is a small ViT Danbooru tagger by the same author — "
			"near-instant per image, ~0.4GB VRAM, loads independently of the JoyCaption model. "
			"It loads automatically on first use."
		)
		with gr.Row():
			joytag_load_btn = gr.Button("⚡ Load JoyTag")
			joytag_unload_btn = gr.Button("🗑️ Unload JoyTag", variant="stop")
		joytag_status = gr.Textbox(label="JoyTag Status", value="JoyTag loads automatically on first use.", interactive=False)
		joytag_load_btn.click(load_joytag, outputs=joytag_status)
		joytag_unload_btn.click(unload_joytag, outputs=joytag_status)

		with gr.Row():
			joytag_threshold = gr.Slider(minimum=0.05, maximum=0.95, value=0.4, step=0.05, label="Tag Threshold (lower = more tags)")
			joytag_spaces_cb = gr.Checkbox(value=False, label="Replace underscores with spaces")

		with gr.Row():
			with gr.Column():
				joytag_image = gr.Image(type="pil", label="Input Image")
				joytag_run_btn = gr.Button("Tag", variant="primary")
			with gr.Column():
				joytag_output = gr.Textbox(label="Tags", lines=6)
				# NOTE: gr.Label hangs the page on tab switch in gradio 6.12 — use a textbox
				joytag_scores = gr.Textbox(label="Tag Confidences (top 50)", lines=10)
		joytag_run_btn.click(
			joytag_single,
			inputs=[joytag_image, joytag_threshold, joytag_spaces_cb],
			outputs=[joytag_output, joytag_scores],
		)

		with gr.Accordion("Batch Folder (dataset tagging)", open=False):
			gr.Markdown("Tags every image in a folder and writes a `.txt` with the same name next to each image.")
			joytag_folder_box = gr.Textbox(label="Folder Path", placeholder=r"C:\path\to\my_dataset")
			with gr.Row():
				joytag_skip_cb = gr.Checkbox(value=True, label="Skip images that already have a .txt")
				joytag_prefix_box = gr.Textbox(label="Tag Prefix (optional, e.g. trigger word)", placeholder="myTriggerWord,")
			joytag_batch_btn = gr.Button("Tag Folder", variant="primary")
			joytag_batch_log = gr.Textbox(label="Progress", lines=15)
			joytag_batch_btn.click(
				joytag_batch,
				inputs=[joytag_folder_box, joytag_threshold, joytag_spaces_cb, joytag_skip_cb, joytag_prefix_box],
				outputs=joytag_batch_log,
			)

if __name__ == "__main__":
	if ON_SPACES:
		# On HF Spaces let gradio pick up the Spaces defaults (0.0.0.0:7860).
		demo.launch()
	else:
		demo.launch(
			server_name="0.0.0.0" if args.listen else "127.0.0.1",
			server_port=args.port,
			inbrowser=not args.no_browser,
		)