| """ |
| 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") |
|
|
| |
| 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 — 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 |
| |
| 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 |
| else: |
| map_idx = 2 |
|
|
| 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)" |
|
|
|
|
| |
|
|
| |
| |
| |
| 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) |
|
|
| |
| 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(): |
| |
| |
| 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) |
| |
| 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: |
| |
| 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, |
| ) |
|
|
