Initial release: JoyCaption Beta One + JoyTag batch captioning WebUI

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+* text=auto
+*.bat text eol=crlf
+*.sh text eol=lf

.gitignore

@@ -0,0 +1,5 @@
+venv/
+models/
+__pycache__/
+*.pyc
+.claude/

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

app.py

@@ -0,0 +1,571 @@
+"""
+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,
+		)

install_linux.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+# One-time setup: creates a venv and installs dependencies (CUDA 12.4 wheels).
+set -e
+cd "$(dirname "$0")"
+python3 -m venv venv
+source venv/bin/activate
+python -m pip install --upgrade pip
+pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
+pip install -r requirements.txt
+echo
+echo "Done. Run: ./start_linux.sh   (add --nf4 for ~7GB VRAM mode)"

install_windows.bat

@@ -0,0 +1,12 @@
+@echo off
+REM One-time setup: creates a venv and installs dependencies (CUDA 12.4 wheels).
+cd /d %~dp0
+where python >nul 2>nul || (echo Python not found on PATH. Install Python 3.10+ first. & pause & exit /b 1)
+python -m venv venv
+call venv\Scripts\activate
+python -m pip install --upgrade pip
+pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
+pip install -r requirements.txt
+echo.
+echo Done. Use start_windows.bat (24GB GPU) or start_windows_lowvram.bat (8GB+ GPU).
+pause

joytag_models.py

@@ -0,0 +1,1163 @@
+# JoyTag model code by fancyfeast (fpgaminer), Apache-2.0.
+# Source: Models.py from https://huggingface.co/spaces/fancyfeast/joytag
+# Project: https://github.com/fpgaminer/joytag
+
+import json
+from pathlib import Path
+from typing import Optional
+import torch
+import torch.backends.cuda
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+from transformers.activations import QuickGELUActivation
+import math
+from einops.layers.torch import Rearrange
+import einops
+
+
+MODEL_CONFIGS = {
+	# Custom models trained from scratch
+	# "Standard" definitions:
+	# name | layers | width | heads
+	#  B   |   12   |  768  |   12
+	#  L   |   24   | 1024  |   16
+	#  H   |   32   | 1280  |   16
+	#  G   |   48   | 1664  |   16
+	#  e   |   56   | 1792  |   16
+	#  22  |   48   | 6144  |   48
+
+	# B/16, 224, PaLM, GELU
+	'CustomTest6': {
+		'class': 'CLIPLikeModel',
+		'embedding_dim': 768,
+		'num_attention_heads': 12,
+		'activation_cls': nn.GELU,
+		'num_channels': 3,
+		'patch_size': 16,
+		'use_palm_alt': True,
+		'num_layers': 12,
+		'use_mha_alt': False,
+		'good_dropout': False,
+	},
+
+	# GAP head + Sinusoidal positional embeddings + 448 image size
+	'CustomTest18': {
+		'class': 'CLIPLikeModel',
+		'embedding_dim': 768,
+		'num_attention_heads': 12,
+		'activation_cls': nn.GELU,
+		'num_channels': 3,
+		'patch_size': 16,
+		'use_palm_alt': True,
+		'num_layers': 12,
+		'use_mha_alt': False,
+		'good_dropout': False,
+		'use_gap_head': True,
+		'sine_positional_embeddings': True,
+	},
+
+	# SW Model + B/16 + ASL + 448 image size
+	# cutout_max_pct = 0
+	# mixup_alpha = 0.8
+	# noise_level = 2
+	# random_resize_method = true
+	# total_labels = 6549
+	'SWModel1': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': False},
+
+	# Sinusoidal positional embeddings 
+	'SWModel2': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	# Sinusoidal positional embeddings + 224 image size + L/14
+	'SWModel3': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+
+	# Sinusoidal positional embeddings + 224 image size + G/14
+	'SWModel4': {'class': 'ViT', 'num_blocks': 48, 'patch_size': 14, 'd_model': 1664, 'mlp_dim': 1664*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+
+	# Sinusoidal positional embeddings + focal loss
+	'SWModel5': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	'SWModel6': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	'SWModel7': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel8': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel9': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel10': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel11': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0, 'use_sine': True},
+
+	# Trying head_mean_after
+	'SWModel12': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'head_mean_after': True},
+
+	# Fat boy
+	'SWModel13': {'class': 'ViT', 'num_blocks': 6, 'patch_size': 16, 'd_model': 1536, 'mlp_dim': 1536*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	# L/14
+	'SWModel14': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel15': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-5, 'use_sine': True},
+	'SWModel16': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.10, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel16f': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.10, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel22': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.20, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel25': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 16, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.15, 'layerscale_init': 1e-1, 'use_sine': True, 'cnn_stem': 'conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=1024;ln;relu;conv:c=1024,s=1,k=1,p=0'},
+
+	# CNN stem
+	'SWModel18': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;bn;relu;conv:c=128;bn;relu;conv:c=256;bn;relu;conv:c=512;bn;relu;conv:c=768,s=1,k=1'},
+	'SWModel19': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;bn;relu;conv:c=128;bn;relu;conv:c=128,s=1;bn;relu;conv:c=256;bn;relu;conv:c=256,s=1;bn;relu;conv:c=512;bn;relu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel20': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;relu;conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel21': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;gelu;conv:c=128;ln;gelu;conv:c=256;ln;gelu;conv:c=512;ln;gelu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel23': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;relu;conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel24': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;relu;conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=768,s=1,k=1,p=0'},
+
+	# H/14
+	'SWModel17': {'class': 'ViT', 'num_blocks': 32, 'patch_size': 14, 'd_model': 1280, 'mlp_dim': 1280*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel26': {'class': 'ViT', 'num_blocks': 32, 'patch_size': 14, 'd_model': 1280, 'mlp_dim': 1280*4, 'num_heads': 16, 'stochdepth_rate': 0.15, 'layerscale_init': 1e-1, 'use_sine': True},
+}
+
+
+class VisionModel(nn.Module):
+	image_size: int
+	n_tags: int
+
+	def __init__(self, image_size: int, n_tags: int):
+		super().__init__()
+
+		self.image_size = image_size
+		self.n_tags = n_tags
+	
+	@staticmethod
+	def load_model(path: Path | str, device: str | None = None) -> 'VisionModel':
+		"""
+		Load a model from a directory.
+		:param path: The directory containing the model.
+		:return: The model, the image size, and the number of tags.
+		"""
+		with open(Path(path) / 'config.json', 'r') as f:
+			config = json.load(f)
+		
+		if (Path(path) / 'model.safetensors').exists():
+			from safetensors.torch import load_file
+			resume = load_file(Path(path) / 'model.safetensors', device='cpu')
+		else:
+			resume = torch.load(Path(path) / 'model.pt', map_location=torch.device('cpu'))['model']
+
+		model_classes = VisionModel.__subclasses__()
+		model_cls = next(cls for cls in model_classes if cls.__name__ == config['class'])
+
+		model = model_cls(**{k: v for k, v in config.items() if k != 'class'})
+		model.load(resume)
+		if device is not None:
+			model = model.to(device)
+
+		return model
+	
+	@staticmethod
+	def from_config(config: dict) -> 'VisionModel':
+		model_classes = VisionModel.__subclasses__()
+		model_cls = next(cls for cls in model_classes if cls.__name__ == config['class'])
+		return model_cls(**{k: v for k, v in config.items() if k != 'class'})
+	
+	def get_optimized_parameters(self, lr: float):
+		raise NotImplementedError
+	
+	def save(self):
+		raise NotImplementedError
+	
+	def load(self, state_dict):
+		raise NotImplementedError
+
+
+def basic_calculate_loss(preds: dict[str, torch.Tensor], batch: dict, pos_weight: torch.Tensor | None, loss_type: str):
+	def asl_helper(preds, target):
+		p = F.softmax(preds, dim=1)
+		xs_pos = p.clamp(min=1e-6)
+		xs_neg = (1 - p).clamp(min=1e-6)
+
+		los_pos = torch.log(torch.gather(xs_pos, 1, target.unsqueeze(1))).sum()
+		los_neg = torch.log(xs_neg)
+		los_neg = los_neg.sum() - torch.gather(los_neg, 1, target.unsqueeze(1)).sum()
+		loss = los_pos + los_neg
+
+		return -loss
+
+	if loss_type == "ce":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'])
+	elif loss_type == "weighted":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight)
+	elif loss_type == "focal":
+		gamma = 2
+		p = torch.sigmoid(preds['tags'])
+		ce_loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], reduction='none')
+		p_t = p * batch['tags'] + (1 - p) * (1 - batch['tags'])
+		loss = ce_loss * ((1 - p_t) ** gamma)
+		loss = loss.mean()
+	elif loss_type == "focal2":
+		gamma = 2
+		p = torch.sigmoid(preds['tags'])
+		ce_loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], reduction='none')
+		p_t = p * batch['tags'] + (1 - p) * (1 - batch['tags'])
+		loss = ce_loss * ((1 - p_t) ** gamma) * 256
+		loss = loss.mean()
+	elif loss_type == "asl":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = los_pos + los_neg
+		loss = -loss.sum()
+
+		# Rating
+		loss = loss + asl_helper(preds['rating'], batch['rating'])
+
+		# Score
+		loss = loss + asl_helper(preds['score'], batch['score'])
+	elif loss_type == "asl2":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = -los_pos - los_neg
+		loss = loss.sum()
+	elif loss_type == "asl3":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = -los_pos - los_neg
+		loss = loss.mean()
+	elif loss_type == "asl4":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = -los_pos - los_neg
+		loss = loss.mean() * 128
+	elif loss_type == "asl5":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight) * 128
+	elif loss_type == "asl6":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight) * 256
+	elif loss_type == "asl7":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight) * 2
+	else:
+		raise ValueError(f"Invalid loss type: {loss_type}")
+	
+	return loss
+
+
+class CLIPMlp(nn.Module):
+	def __init__(self, hidden_size: int, intermediate_size: int, activation_cls):
+		super().__init__()
+		self.activation_fn = activation_cls()
+		self.fc1 = nn.Linear(hidden_size, intermediate_size)
+		self.fc2 = nn.Linear(intermediate_size, hidden_size)
+
+	def forward(self, hidden_states: torch.Tensor):
+		hidden_states = self.fc1(hidden_states)
+		hidden_states = self.activation_fn(hidden_states)
+		hidden_states = self.fc2(hidden_states)
+		return hidden_states
+
+
+class FastCLIPAttention2(nn.Module):
+	"""Fast Attention module for CLIP-like. This is NOT a drop-in replacement for CLIPAttention, since it adds additional flexibility.  Mainly uses xformers."""
+	def __init__(self, hidden_size: int, out_dim: int, num_attention_heads: int, out_seq_len: Optional[int] = None, norm_qk: bool = False):
+		super().__init__()
+		self.out_seq_len = out_seq_len
+		self.embed_dim = hidden_size
+		self.out_dim = out_dim
+		self.norm_qk = norm_qk
+		self.num_heads = num_attention_heads
+		self.head_dim = hidden_size // num_attention_heads
+		assert self.head_dim * num_attention_heads == self.embed_dim, "embed_dim must be divisible by num_attention_heads"
+
+		self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+		self.kv_proj = nn.Linear(self.embed_dim, self.embed_dim * 2)
+		self.out_proj = nn.Linear(self.embed_dim, self.out_dim)
+
+		if self.norm_qk:
+			self.query_norm = nn.LayerNorm(self.embed_dim)
+			self.key_norm = nn.LayerNorm(self.embed_dim)
+	
+	#def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+	#	return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).contiguous()
+	
+	def forward(self, query_states: torch.Tensor, kv_states: torch.Tensor) -> torch.Tensor:
+		bsz, src_len, embed_dim = kv_states.size()
+		if self.out_seq_len is not None:
+			tgt_len = self.out_seq_len
+		else:
+			tgt_len = src_len
+		
+		kv_states = self.kv_proj(kv_states)  # (bsz, src_len, embed_dim * 2)
+		q_states = self.q_proj(query_states[:, :tgt_len])   # (bsz, tgt_len, embed_dim)
+
+		# NOTE: It is not clear if LayerNorm should be applied to the embed_dim, or to the head_dim
+		if self.norm_qk:
+			q_states = self.query_norm(q_states).type(q_states.dtype)
+			k_states = self.key_norm(kv_states[:, :, :embed_dim]).type(kv_states.dtype)
+			v_states = kv_states[:, :, embed_dim:]
+		else:
+			k_states = kv_states[:, :, :embed_dim]
+			v_states = kv_states[:, :, embed_dim:]
+		
+		q_states = q_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)  # (bsz, num_heads, tgt_len, head_dim)
+		k_states = k_states.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2)  # (bsz, num_heads, src_len, head_dim)
+		v_states = v_states.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2)  # (bsz, num_heads, src_len, head_dim)
+
+		# Performs scale of query_states, attention, and softmax
+		with torch.backends.cuda.sdp_kernel(enable_math=False):
+			x = F.scaled_dot_product_attention(q_states, k_states, v_states)   # (bsz, num_heads, tgt_len, head_dim)
+			x = x.transpose(1, 2).contiguous().view(bsz, tgt_len, embed_dim)   # (bsz, tgt_len, embed_dim)
+		
+		# Projection
+		x = self.out_proj(x)  # (bsz, tgt_len, out_dim)
+
+		return x
+
+
+class SkipInit(nn.Module):
+	def __init__(self, hidden_size: int, channel_wise: bool, init_scale: float):
+		super().__init__()
+		self.hidden_size = hidden_size
+		self.channel_wise = channel_wise
+		self.init_scale = init_scale
+
+		if self.channel_wise:
+			self.scale = nn.Parameter(torch.ones(hidden_size) * init_scale)
+		else:
+			self.scale = nn.Parameter(torch.tensor(init_scale))
+
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		return x * self.scale
+
+
+class FastCLIPEncoderLayer(nn.Module):
+	def __init__(
+		self,
+		hidden_size: int,
+		num_attention_heads: int,
+		out_seq_len: Optional[int],
+		activation_cls = QuickGELUActivation,
+		use_palm_alt: bool = False,
+		norm_qk: bool = False,
+		skip_init: Optional[float] = None,
+		stochastic_depth: Optional[float] = None,
+	):
+		super().__init__()
+
+		self.use_palm_alt = use_palm_alt
+		self.stochastic_depth = stochastic_depth
+		
+		self.self_attn = FastCLIPAttention2(
+			hidden_size=hidden_size,
+			out_dim=hidden_size,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=out_seq_len,
+			norm_qk=norm_qk,
+		)
+		self.mlp = CLIPMlp(hidden_size, 4 * hidden_size, activation_cls)
+		self.layer_norm1 = nn.LayerNorm(hidden_size)
+		if not use_palm_alt:
+			self.layer_norm2 = nn.LayerNorm(hidden_size)
+		
+		if skip_init is not None:
+			self.attn_skip_init = SkipInit(hidden_size, channel_wise=True, init_scale=skip_init)
+			self.mlp_skip_init = SkipInit(hidden_size, channel_wise=True, init_scale=skip_init)
+		else:
+			self.attn_skip_init = nn.Identity()
+			self.mlp_skip_init = nn.Identity()
+	
+	def forward(self, hidden_states: torch.Tensor):
+		residual = hidden_states
+		hidden_states = self.layer_norm1(hidden_states)
+
+		if not self.use_palm_alt:
+			hidden_states = self.self_attn(query_states=hidden_states, kv_states=hidden_states)
+			hidden_states = self.attn_skip_init(hidden_states)
+			hidden_states = hidden_states + residual[:, :hidden_states.size(1)]
+
+			residual = hidden_states
+			hidden_states = self.layer_norm2(hidden_states)
+			hidden_states = self.mlp(hidden_states)
+			hidden_states = self.mlp_skip_init(hidden_states)
+			hidden_states = hidden_states + residual
+		else:
+			# An alternative implementation inspired by the PALM paper
+			# By performing the attention and MLP in parallel it's possible to fuse the linear projections of the attention and MLP layers
+			# We don't do that here yet, but that supposedly improves efficiency without hurting performance
+			attn = self.self_attn(query_states=hidden_states, kv_states=hidden_states)
+			attn = self.attn_skip_init(attn)
+			mlp = self.mlp(hidden_states[:, :attn.size(1)])
+			mlp = self.mlp_skip_init(mlp)
+
+			if self.stochastic_depth is not None:
+				attn = torchvision.ops.stochastic_depth(attn, self.stochastic_depth, mode='row', training=self.training)
+				mlp = torchvision.ops.stochastic_depth(mlp, self.stochastic_depth, mode='row', training=self.training)
+
+			hidden_states = residual[:, :attn.size(1)] + attn + mlp
+
+		return hidden_states
+
+
+def sinusoidal_position_embedding(width: int, height: int, depth: int, dtype, device, temperature = 10000):
+	"""
+	Sinusoidal position embedding. Returns a flat tensor of shape (h * w, d).
+	"""
+	assert depth % 4 == 0, "Embedding dimension must be divisible by 4."
+
+	y, x = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing="ij")
+	omega = torch.arange(depth // 4, device=device) / (depth // 4 - 1)
+	omega = 1. / (temperature ** omega)
+
+	y = y.flatten()[:, None] * omega[None, :]
+	x = x.flatten()[:, None] * omega[None, :]
+	embedding = torch.cat([x.sin(), x.cos(), y.sin(), y.cos()], dim=1)
+
+	return embedding.type(dtype)
+
+
+class CLIPEmbeddingLayer(nn.Module):
+	def __init__(self, hidden_size: int, num_channels: int, image_size: int, patch_size: int, patch_dropout: float = 0.0, good_dropout: bool = False, dpn: bool = False, sine_positional_embeddings: bool = False):
+		super().__init__()
+
+		assert image_size % patch_size == 0, "Image dimensions must be divisible by the patch size."
+
+		seq_len = (image_size // patch_size) ** 2
+		self.patch_dropout = patch_dropout
+		self.hidden_size = hidden_size
+		self.good_dropout = good_dropout
+		self.dpn = dpn
+		self.sine_positional_embeddings = sine_positional_embeddings
+		self.patch_size = patch_size
+
+		self.patch_embeddings = nn.Conv2d(
+			in_channels=num_channels,
+			out_channels=hidden_size,
+			kernel_size=patch_size,
+			stride=patch_size,
+			bias=False,
+		)
+		if not self.sine_positional_embeddings:
+			self.positional_embeddings = nn.Embedding(seq_len, hidden_size)
+		self.register_buffer("position_ids", torch.arange(seq_len))
+
+		if self.dpn:
+			self.to_patch_embeddings = nn.Sequential(
+				Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size),
+				nn.LayerNorm(3 * patch_size * patch_size),
+				nn.Linear(3 * patch_size * patch_size, hidden_size),
+				nn.LayerNorm(hidden_size),
+			)
+		else:
+			self.to_patch_embeddings = nn.Conv2d(
+				in_channels=num_channels,
+				out_channels=hidden_size,
+				kernel_size=patch_size,
+				stride=patch_size,
+				bias=False,
+			)
+	
+	def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+		B, C, H, W = pixel_values.shape
+		assert H % self.patch_size == 0, f"Input image height ({H}) needs to be divisible by the patch size ({self.patch_size})."
+		assert W % self.patch_size == 0, f"Input image width ({W}) needs to be divisible by the patch size ({self.patch_size})."
+
+		if self.dpn:
+			patches = self.to_patch_embeddings(pixel_values)
+		else:
+			patches = self.to_patch_embeddings(pixel_values)
+			patches = patches.flatten(2).transpose(1, 2)
+		
+		seq_len = patches.shape[1]
+		patch_dropout = int(math.ceil((1.0 - self.patch_dropout) * seq_len))
+		
+		if self.sine_positional_embeddings:
+			position_embeddings = sinusoidal_position_embedding(W // self.patch_size, H // self.patch_size, self.hidden_size, pixel_values.dtype, pixel_values.device)
+		else:
+			position_embeddings = self.positional_embeddings(self.position_ids)
+
+		if patch_dropout == seq_len or not self.training:
+			embeddings = patches + position_embeddings
+		elif self.good_dropout:
+			# Pick random patches to drop out
+			# The "good_dropout" variant uses random permutations for each batch item, but is slightly slower and involves more code
+
+			# The below method is a nice trick to generate a batch of random permutations.
+			# Torch (as of 1.13) doesn't have a built-in function to do this, and a for loop of torch.randperm is slow.
+			# Based on some benchmarks I measured the generation of the mask and the fetching to be only 50% slower than the non-"good_dropout" variant.
+			# And the time taken here is only a fraction of the time spent performing the embedding convolution.
+			# Generate a matrix of random numbers between 0 and 1 of shape (B, seq_len)
+			patch_mask = torch.rand(B, seq_len, device=patches.device)
+			# For each batch tensor, use argsort to convert the random numbers into a permutation of the patch indices
+			patch_mask = torch.argsort(patch_mask, dim=1)
+			# Truncate
+			patch_mask = patch_mask[:, :patch_dropout]
+
+			embeddings = patches.gather(1, patch_mask.unsqueeze(-1).expand(-1, -1, self.hidden_size)) + position_embeddings[patch_mask]
+		else:
+			# The non-"good_dropout" variant uses a single random permutation for all batch items, but is faster and uses less code
+			indices = torch.randperm(seq_len, device=pixel_values.device)[:patch_dropout]
+			embeddings = patches[:, indices, :] + position_embeddings[indices.expand(1, -1)]
+		
+		return embeddings
+
+
+class MHAPoolingHead(nn.Module):
+	def __init__(self, hidden_size: int, num_attention_heads: int, activation_cls, out_dim: int, alt_style: bool, norm_qk: bool):
+		super().__init__()
+
+		self.out_dim = out_dim if not alt_style else hidden_size
+
+		self.probe = nn.Parameter(torch.randn(hidden_size))
+
+		self.mlp = CLIPMlp(hidden_size, 4 * hidden_size, activation_cls)
+		self.layer_norm = nn.LayerNorm(hidden_size)
+		self.pooling_head = nn.Linear(hidden_size, 1)
+
+		self.self_attn = FastCLIPAttention2(
+			hidden_size=hidden_size,
+			out_dim=self.out_dim,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=1,
+			norm_qk=norm_qk,
+		)
+		self.mlp = CLIPMlp(self.out_dim, 4 * self.out_dim, activation_cls)
+		self.layer_norm1 = nn.LayerNorm(hidden_size)
+		self.layer_norm2 = nn.LayerNorm(self.out_dim)
+
+		if alt_style:
+			self.final_proj = nn.Linear(hidden_size, out_dim)
+		else:
+			self.final_proj = nn.Identity()
+	
+	def forward(self, hidden_states: torch.Tensor):
+		hidden_states = self.layer_norm1(hidden_states)
+		query_states = self.probe.unsqueeze(0).unsqueeze(0).expand(hidden_states.size(0), 1, -1)
+
+		hidden_states = self.self_attn(query_states=query_states, kv_states=hidden_states)
+		# We don't use a residual connection here because the out_dim is different from the hidden_size
+
+		residual = hidden_states
+		hidden_states = self.layer_norm2(hidden_states)
+		hidden_states = self.mlp(hidden_states)
+		hidden_states = hidden_states + residual
+		hidden_states = self.final_proj(hidden_states)
+
+		return hidden_states.squeeze(1)
+
+
+class GAPHead(nn.Module):
+	def __init__(self, hidden_size: int, out_dim: int):
+		super().__init__()
+
+		self.norm = nn.LayerNorm(hidden_size)
+		self.proj = nn.Linear(hidden_size, out_dim)
+
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		x = x.mean(dim=1)
+		x = self.norm(x)
+		x = self.proj(x)
+		return x
+
+
+class CLIPLikeModel(VisionModel):
+	def __init__(
+		self,
+		n_tags: int,
+		embedding_dim: int,
+		num_attention_heads: int,
+		activation_cls,
+		num_channels: int,
+		image_size: int,
+		patch_size: int,
+		patch_dropout: float,
+		use_palm_alt: bool,
+		num_layers: int,
+		use_mha_alt: bool,
+		loss_type: str,
+		good_dropout: bool=False,
+		dpn: bool=False,
+		sine_positional_embeddings: bool=False,
+		norm_qk: bool = False,
+		no_wd_bias: bool = False,
+		use_gap_head: bool = False,
+		skip_init: Optional[float] = None,
+		stochastic_depth: Optional[float] = None,
+	):
+		super().__init__(image_size, n_tags)
+
+		out_dim = n_tags
+		self.n_tags = n_tags
+		self.loss_type = loss_type
+		self.no_wd_bias = no_wd_bias
+		
+		stochastic_depth_space = torch.linspace(0, stochastic_depth, num_layers) if stochastic_depth is not None else None
+
+		self.embedding_layer = CLIPEmbeddingLayer(embedding_dim, num_channels, image_size, patch_size, patch_dropout, good_dropout, dpn, sine_positional_embeddings)
+		self.pre_layer_norm = nn.LayerNorm(embedding_dim)
+		self.encoder_layers = nn.ModuleList([FastCLIPEncoderLayer(
+			hidden_size=embedding_dim,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=None,
+			activation_cls=activation_cls,
+			use_palm_alt=use_palm_alt,
+			norm_qk=norm_qk,
+			skip_init=skip_init,
+			stochastic_depth=stochastic_depth_space[i].item() if stochastic_depth_space is not None else None,
+		) for i in range(num_layers)])
+
+		if use_gap_head:
+			self.pooling_head = GAPHead(embedding_dim, out_dim)
+		else:
+			self.pooling_head = MHAPoolingHead(embedding_dim, num_attention_heads, activation_cls, out_dim, use_mha_alt, norm_qk=norm_qk)
+	
+	def forward(self, batch):
+		hidden_states = self.embedding_layer(batch['image'])
+		hidden_states = self.pre_layer_norm(hidden_states)
+
+		for layer in self.encoder_layers:
+			hidden_states = layer(hidden_states)
+		
+		preds = self.pooling_head(hidden_states)
+
+		result = {
+			'tags': preds,
+		}
+
+		return result
+	
+	def calculate_loss(self, preds, batch, pos_weight):
+		return basic_calculate_loss(preds, batch, pos_weight, self.loss_type)
+	
+	def get_optimized_parameters(self, lr: float):
+		if self.no_wd_bias:
+			return self.get_optimized_parameters_no_wd_bias()
+		else:
+			return self.parameters()
+	
+	def get_optimized_parameters_no_wd_bias(self):
+		decay = []
+		no_decay = []
+
+		for name, param in self.named_parameters():
+			if not param.requires_grad:
+				continue
+
+			if len(param.shape) == 1 or name.endswith(".bias"):
+				no_decay.append(param)
+				print(f'No decay: {name}')
+			else:
+				decay.append(param)
+		
+		return [
+			{'params': decay},
+			{'params': no_decay, 'weight_decay': 0.},
+		]
+	
+	def save(self):
+		return self.state_dict()
+	
+	def load(self, state_dict):
+		self.load_state_dict(state_dict)
+
+
+class MaskedAutoEncoderViT(nn.Module):
+	def __init__(
+		self,
+		n_tags: int,
+
+		embedding_dim: int,
+		num_attention_heads: int,
+		activation_cls,
+		num_channels: int,
+		image_size: int,
+		patch_size: int,
+		num_layers: int,
+		loss_type: str,
+		sine_positional_embeddings: bool=False,
+
+		decoder_embedding_dim: int = 512,
+		decoder_num_attention_heads: int = 8,
+		decoder_num_layers: int = 6,
+		decoder_force_projection: bool = False,
+
+		masking_ratio: float = 0.75,
+		mae_loss_weight: float = 1.0,
+		mae_normalize_targets: bool = False,
+		mae_post_norm: bool = False,
+	):
+		super().__init__()
+
+		self.n_tags = n_tags
+		self.seq_len = (image_size // patch_size) ** 2
+		self.embedding_dim = embedding_dim
+		self.decoder_embedding_dim = decoder_embedding_dim
+		self.sine_positional_embeddings = sine_positional_embeddings
+		self.image_size = image_size
+		self.patch_size = patch_size
+		self.masking_ratio = masking_ratio
+		self.loss_type = loss_type
+		self.mae_loss_weight = mae_loss_weight
+		self.mae_normalize_targets = mae_normalize_targets
+
+		if not self.sine_positional_embeddings:
+			self.positional_embeddings = nn.Embedding(self.seq_len, embedding_dim)
+			self.decoder_positional_embeddings = nn.Embedding(self.seq_len, decoder_embedding_dim)
+		self.register_buffer("position_ids", torch.arange(self.seq_len))
+
+		self.to_patches = Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
+		self.patch_embedder = nn.Linear(num_channels * patch_size * patch_size, embedding_dim)
+
+		# Encoder
+		self.pre_layer_norm = nn.LayerNorm(embedding_dim)
+		self.encoder_layers = nn.ModuleList([FastCLIPEncoderLayer(
+			hidden_size=embedding_dim,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=None,
+			activation_cls=activation_cls,
+			use_palm_alt=True,
+			norm_qk=False,
+			skip_init=None,
+		) for _ in range(num_layers)])
+
+		# Head for classification
+		self.pooling_head = GAPHead(embedding_dim, n_tags)
+
+		# Decoder
+		if embedding_dim != decoder_embedding_dim or decoder_force_projection:
+			self.encoder_to_decoder_proj = nn.Linear(embedding_dim, decoder_embedding_dim)
+		else:
+			self.encoder_to_decoder_proj = nn.Identity()
+		self.decoder_pre_layer_norm = nn.LayerNorm(decoder_embedding_dim)
+		self.decoder_layers = nn.ModuleList([FastCLIPEncoderLayer(
+			hidden_size=decoder_embedding_dim,
+			num_attention_heads=decoder_num_attention_heads,
+			out_seq_len=None,
+			activation_cls=activation_cls,
+			use_palm_alt=True,
+			norm_qk=False,
+			skip_init=None,
+		) for _ in range(decoder_num_layers)])
+
+		if mae_post_norm:
+			self.decoder_to_pixel_values = nn.Sequential(
+				nn.LayerNorm(decoder_embedding_dim),
+				nn.Linear(decoder_embedding_dim, num_channels * patch_size * patch_size)
+			)
+		else:
+			self.decoder_to_pixel_values = nn.Linear(decoder_embedding_dim, num_channels * patch_size * patch_size)
+		self.mask_token = nn.Parameter(torch.zeros(decoder_embedding_dim))
+		torch.nn.init.normal_(self.mask_token, std=0.02)
+
+	def forward(self, batch):
+		pixel_values = batch['image']
+		device = pixel_values.device
+		B, C, H, W = pixel_values.shape
+		assert H % self.patch_size == 0, f"Input image height ({H}) needs to be divisible by the patch size ({self.patch_size})."
+		assert W % self.patch_size == 0, f"Input image width ({W}) needs to be divisible by the patch size ({self.patch_size})."
+
+		# Convert image to patches (B, seq_len, C * patch_size * patch_size)
+		patches = self.to_patches(pixel_values)
+		seq_len = patches.shape[1]
+		num_masked = int(self.masking_ratio * seq_len)
+
+		# For each batch tensor, use argsort to convert the random numbers into a permutation of the patch indices
+		# From this we can get the masked and unmasked indices
+		patch_mask = torch.rand(B, seq_len, device=device)
+		patch_mask = torch.argsort(patch_mask, dim=1)
+		masked_indices, unmasked_indices = patch_mask[:, :num_masked], patch_mask[:, num_masked:]
+		batch_range = torch.arange(B, device=device)[:, None]
+
+		# Masked and unmasked patches
+		unmasked_patches = patches[batch_range, unmasked_indices]
+		masked_patches = patches[batch_range, masked_indices]
+
+		# Embed unmasked patches for the encoder (B, seq_len, embedding_dim)
+		tokens = self.patch_embedder(unmasked_patches)
+
+		if self.sine_positional_embeddings:
+			position_embeddings = sinusoidal_position_embedding(W // self.patch_size, H // self.patch_size, self.embedding_dim, pixel_values.dtype, device)
+			decoder_position_embeddings = sinusoidal_position_embedding(W // self.patch_size, H // self.patch_size, self.decoder_embedding_dim, pixel_values.dtype, device)
+		else:
+			position_embeddings = self.positional_embeddings(self.position_ids)
+			decoder_position_embeddings = self.decoder_positional_embeddings(self.position_ids)
+		
+		# Add position embeddings
+		tokens = tokens + position_embeddings[unmasked_indices]
+
+		# Run the encoder
+		encoded_tokens = self.pre_layer_norm(tokens)
+
+		for layer in self.encoder_layers:
+			encoded_tokens = layer(encoded_tokens)
+		
+		# Label predictions
+		if self.training:
+			preds = self.pooling_head(encoded_tokens)
+		else:
+			# During inference, classify using the entire image
+			# But we'll do the usual for the MAE part, just so we can see how MAE is performing during validation
+			tokens = self.patch_embedder(patches)
+			tokens = tokens + position_embeddings
+			tokens = self.pre_layer_norm(tokens)
+			for layer in self.encoder_layers:
+				tokens = layer(tokens)
+			preds = self.pooling_head(tokens)
+		
+		# Projection for the decoder and position embeddings
+		decoder_tokens = self.encoder_to_decoder_proj(encoded_tokens)
+		decoder_tokens = decoder_tokens + decoder_position_embeddings[unmasked_indices]
+
+		# Fill in the masked patches
+		mask_tokens = einops.repeat(self.mask_token, 'd -> b n d', b = B, n = num_masked)
+		mask_tokens = mask_tokens + decoder_position_embeddings[masked_indices]
+		decoder_tokens = torch.cat([decoder_tokens, mask_tokens], dim=1)
+
+		# Run the decoder
+		decoded_tokens = self.decoder_pre_layer_norm(decoder_tokens)
+		
+		for layer in self.decoder_layers:
+			decoded_tokens = layer(decoded_tokens)
+
+		# Only predict the masked patches
+		# All the masked patches are at the end of the sequence
+		decoded_tokens = decoded_tokens[:, -num_masked:]
+		pred_pixel_values = self.decoder_to_pixel_values(decoded_tokens)
+
+		# Calculate the mae loss
+		if self.mae_normalize_targets:
+			# Normalize each patch by its mean and variance. The ViCHA paper says this provides better results
+			means = masked_patches.mean(dim=-1, keepdim=True)
+			vars = masked_patches.var(dim=-1, keepdim=True)
+			target = (masked_patches - means) / (vars + 1e-6)**0.5
+			mae_loss = F.mse_loss(pred_pixel_values, target)
+		else:
+			mae_loss = F.mse_loss(pred_pixel_values, masked_patches)
+		mae_loss = mae_loss * self.mae_loss_weight
+
+		return {
+			'tags': preds,
+			'mae_loss': mae_loss,
+		}
+
+	def calculate_loss(self, preds, batch, pos_weight):
+		return basic_calculate_loss(preds, batch, pos_weight, self.loss_type) + preds['mae_loss']
+	
+	def get_optimized_parameters(self, lr: float):
+		return self.parameters()
+	
+	def save(self):
+		return self.state_dict()
+	
+	def load(self, state_dict):
+		self.load_state_dict(state_dict)
+
+
+class StochDepth(nn.Module):
+	def __init__(self, drop_rate: float, scale_by_keep: bool = False):
+		super().__init__()
+		self.drop_rate = drop_rate
+		self.scale_by_keep = scale_by_keep
+	
+	def forward(self, x):
+		if not self.training:
+			return x
+		
+		batch_size = x.shape[0]
+		r = torch.rand((batch_size, 1, 1), device=x.device)
+		keep_prob = 1 - self.drop_rate
+		binary_tensor = torch.floor(keep_prob + r)
+		if self.scale_by_keep:
+			x = x / keep_prob
+		
+		return x * binary_tensor
+
+
+class SkipInitChannelwise(nn.Module):
+	def __init__(self, channels, init_val=1e-6):
+		super().__init__()
+		self.channels = channels
+		self.init_val = init_val
+		self.skip = nn.Parameter(torch.ones(channels) * init_val)
+	
+	def forward(self, x):
+		return x * self.skip
+
+
+class PosEmbedding(nn.Module):
+	def __init__(self, d_model: int, max_len: int, use_sine: bool, patch_size: int):
+		super().__init__()
+		self.d_model = d_model
+		self.max_len = max_len
+		self.use_sine = use_sine
+		self.patch_size = patch_size
+
+		if not self.use_sine:
+			self.embedding = nn.Embedding(max_len, d_model)
+			nn.init.trunc_normal_(self.embedding.weight, std=0.02)
+			self.register_buffer("position_ids", torch.arange(max_len))
+	
+	def forward(self, x, width: int, height: int):
+		if self.use_sine:
+			position_embeddings = sinusoidal_position_embedding(width // self.patch_size, height // self.patch_size, self.d_model, x.dtype, x.device)
+		else:
+			position_embeddings = self.embedding(self.position_ids)
+		
+		return x + position_embeddings
+
+
+class MLPBlock(nn.Module):
+	def __init__(self, d_model: int, d_ff: int, stochdepth_rate: float):
+		super().__init__()
+		self.linear1 = nn.Linear(d_model, d_ff)
+		self.linear2 = nn.Linear(d_ff, d_model)
+		self.activation = nn.GELU()
+		if stochdepth_rate > 0:
+			self.stochdepth = StochDepth(stochdepth_rate, scale_by_keep=True)
+		else:
+			self.stochdepth = None
+	
+	def forward(self, x):
+		x = self.linear1(x)
+		x = self.activation(x)
+		if self.stochdepth is not None:
+			x = self.stochdepth(x)
+		x = self.linear2(x)
+		return x
+
+
+class ViTBlock(nn.Module):
+	def __init__(self, num_heads: int, d_model: int, d_ff: int, layerscale_init: float, stochdepth_rate: float):
+		super().__init__()
+		self.num_heads = num_heads
+		self.d_model = d_model
+
+		assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+		# MHA
+		self.norm1 = nn.LayerNorm(d_model)
+		self.qkv_proj = nn.Linear(d_model, d_model * 3)
+		self.out_proj = nn.Linear(d_model, d_model)
+		self.skip_init1 = SkipInitChannelwise(channels=d_model, init_val=layerscale_init)
+		self.stochdepth1 = StochDepth(stochdepth_rate, scale_by_keep=True) if stochdepth_rate > 0 else None
+		
+		# MLP
+		self.norm2 = nn.LayerNorm(d_model)
+		self.mlp = MLPBlock(d_model, d_ff, stochdepth_rate)
+		self.skip_init2 = SkipInitChannelwise(channels=d_model, init_val=layerscale_init)
+		self.stochdepth2 = StochDepth(stochdepth_rate, scale_by_keep=True) if stochdepth_rate > 0 else None
+	
+	def forward(self, x):
+		bsz, src_len, embed_dim = x.shape
+
+		out = x
+		out = self.norm1(out)
+
+		# MHA
+		qkv_states = self.qkv_proj(out).split(self.d_model, dim=-1)
+		q_states = qkv_states[0].view(bsz, src_len, self.num_heads, embed_dim // self.num_heads).transpose(1, 2)  # (bsz, num_heads, src_len, embed_dim // num_heads)
+		k_states = qkv_states[1].view(bsz, src_len, self.num_heads, embed_dim // self.num_heads).transpose(1, 2)  # (bsz, num_heads, src_len, embed_dim // num_heads)
+		v_states = qkv_states[2].view(bsz, src_len, self.num_heads, embed_dim // self.num_heads).transpose(1, 2)  # (bsz, num_heads, src_len, embed_dim // num_heads)
+
+		with torch.backends.cuda.sdp_kernel(enable_math=False):
+			out = F.scaled_dot_product_attention(q_states, k_states, v_states)   # (bsz, num_heads, tgt_len, head_dim)
+			out = out.transpose(1, 2).contiguous().view(bsz, src_len, embed_dim)   # (bsz, tgt_len, embed_dim)
+		
+		out = self.out_proj(out)
+
+		out = self.skip_init1(out)
+		if self.stochdepth1 is not None:
+			out = self.stochdepth1(out)
+		x = out + x
+
+		out = self.norm2(x)
+		out = self.mlp(out)
+		out = self.skip_init2(out)
+		if self.stochdepth2 is not None:
+			out = self.stochdepth2(out)
+		
+		out = out + x
+
+		return out
+
+
+def CaiT_LayerScale_init(network_depth):
+	if network_depth <= 18:
+		return 1e-1
+	elif network_depth <= 24:
+		return 1e-5
+	else:
+		return 1e-6
+
+
+class CNNLayerNorm(nn.Module):
+	def __init__(self, d_model: int):
+		super().__init__()
+		self.norm = nn.LayerNorm(d_model)
+	
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		x = x.transpose(1, 3)
+		x = self.norm(x)
+		x = x.transpose(1, 3)
+		return x
+
+
+class CNNStem(nn.Module):
+	def __init__(self, config: str):
+		super().__init__()
+		self.config = config
+
+		layers = []
+		channels = 3
+
+		for line in config.split(";"):
+			ty, line = line.split(":") if ":" in line else (line, "")
+			options = line.split(",")
+			options = [o.split("=") for o in options] if line else []
+			options = {k: v for k, v in options}
+
+			if ty == 'conv':
+				layers.append(nn.Conv2d(
+					in_channels=channels,
+					out_channels=int(options['c']),
+					kernel_size=int(options['k'] if 'k' in options else 3),
+					stride=int(options['s'] if 's' in options else 2),
+					bias=True,
+					padding=int(options['p'] if 'p' in options else 1),
+				))
+				channels = int(options['c'])
+			elif ty == 'bn':
+				layers.append(nn.BatchNorm2d(channels))
+			elif ty == 'ln':
+				layers.append(CNNLayerNorm(channels))
+			elif ty == 'relu':
+				layers.append(nn.ReLU())
+			elif ty == 'gelu':
+				layers.append(nn.GELU())
+
+		self.conv = nn.Sequential(*layers)
+	
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		return self.conv(x)
+
+
+class ViT(VisionModel):
+	def __init__(self,
+		n_tags: int,
+		image_size: int,
+		num_blocks: int,
+		patch_size: int,
+		d_model: int,
+		mlp_dim: int,
+		num_heads: int,
+		stochdepth_rate: float,
+		use_sine: bool,
+		loss_type: str,
+		layerscale_init: Optional[float] = None,
+		head_mean_after: bool = False,
+		cnn_stem: str | None = None,
+		patch_dropout: float = 0.0,
+	):
+		super().__init__(image_size, n_tags)
+
+		#assert image_size % patch_size == 0, "image_size must be divisible by patch_size"
+		assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+		out_dim = n_tags
+		self.n_tags = n_tags
+		self.loss_type = loss_type
+		self.patch_size = patch_size
+		self.head_mean_after = head_mean_after
+		self.patch_dropout = patch_dropout
+
+		layerscale_init = CaiT_LayerScale_init(num_blocks) if layerscale_init is None else layerscale_init
+		self.patch_embeddings = nn.Conv2d(
+			in_channels=3,
+			out_channels=d_model,
+			kernel_size=patch_size,
+			stride=patch_size,
+			bias=True,
+		) if cnn_stem is None else CNNStem(cnn_stem)
+		self.pos_embedding = PosEmbedding(d_model, (image_size // patch_size) ** 2, use_sine=use_sine, patch_size=patch_size)
+
+		self.blocks = nn.ModuleList([
+			ViTBlock(num_heads, d_model, mlp_dim, layerscale_init, stochdepth_rate)
+			for _ in range(num_blocks)
+		])
+
+		self.norm = nn.LayerNorm(d_model)
+		self.head = nn.Linear(d_model, out_dim)
+
+	def forward(self, batch, return_embeddings=False, return_loss: bool = False, pos_weight = None):
+		B, C, H, W = batch['image'].shape
+		assert H % self.patch_size == 0, f"Input image height ({H}) needs to be divisible by the patch size ({self.patch_size})."
+		assert W % self.patch_size == 0, f"Input image width ({W}) needs to be divisible by the patch size ({self.patch_size})."
+
+		x = self.patch_embeddings(batch['image'])  # (bsz, d_model, patch_num, patch_num)
+		x = x.flatten(2).transpose(1, 2)  # (bsz, patch_num ** 2, d_model)
+		x = self.pos_embedding(x, W, H)   # (bsz, patch_num ** 2, d_model)
+
+		# Patch dropout
+		seq_len = x.shape[1]
+		patch_dropout = int(math.ceil((1.0 - self.patch_dropout) * seq_len))
+
+		if patch_dropout != seq_len:
+			# Generate a matrix of random numbers between 0 and 1 of shape (B, seq_len)
+			patch_mask = torch.rand(B, seq_len, device=x.device)
+			# For each batch tensor, use argsort to convert the random numbers into a permutation of the patch indices
+			patch_mask = torch.argsort(patch_mask, dim=1)
+			# Truncate
+			patch_mask = patch_mask[:, :patch_dropout]
+
+			x = x.gather(1, patch_mask.unsqueeze(-1).expand(-1, -1, x.shape[-1]))
+
+			#indices = torch.randperm(seq_len, device=x.device)[:patch_dropout]
+			#x = x[:, indices, :]
+
+		# Transformer
+		for block in self.blocks:
+			x = block(x)
+		
+		# Head
+		result = {}
+
+		x = self.norm(x)
+		if self.head_mean_after:
+			x = self.head(x)
+			x = x.mean(dim=1)
+		else:
+			x = x.mean(dim=1)
+			if return_embeddings:
+				result['embeddings'] = x
+			x = self.head(x)
+
+		result['tags'] = x
+
+		if return_loss:
+			result['loss'] = self.calculate_loss(result, batch, pos_weight)
+
+		return result
+	
+	def calculate_loss(self, preds, batch, pos_weight):
+		return basic_calculate_loss(preds, batch, pos_weight, self.loss_type)
+	
+	def get_optimized_parameters(self, lr: float):
+		return self.parameters()
+	
+	def save(self):
+		return self.state_dict()
+	
+	def load(self, state_dict):
+		if 'head.weight' in state_dict and 'head.bias' in state_dict and state_dict['head.weight'].shape[0] == (self.n_tags + 9):
+			# Support old models which included 3 rating and 6 score dimensions
+			state_dict['head.weight'] = state_dict['head.weight'][:self.n_tags]
+			state_dict['head.bias'] = state_dict['head.bias'][:self.n_tags]
+
+		self.load_state_dict(state_dict)

requirements.txt

@@ -0,0 +1,11 @@
+# IMPORTANT: gradio is pinned because 6.14.0 has a rendering regression that hides
+# the Caption Type / Caption Length dropdowns and breaks prompt updates. Do NOT upgrade.
+gradio==6.12.0
+transformers==4.46.3
+accelerate
+torch==2.6.0
+torchvision==0.21.0
+bitsandbytes==0.49.2; sys_platform != "darwin"
+pillow
+einops  # required by joytag_models.py (JoyTag tab)
+huggingface_hub

start_linux.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+cd "$(dirname "$0")"
+source venv/bin/activate
+python app.py "$@"

start_windows.bat

@@ -0,0 +1,7 @@
+@echo off
+REM Full-quality bf16 mode (~18GB VRAM)
+cd /d %~dp0
+call venv\Scripts\activate
+set HF_HUB_DISABLE_SYMLINKS_WARNING=1
+python app.py %*
+pause

start_windows_lowvram.bat

@@ -0,0 +1,7 @@
+@echo off
+REM 4-bit NF4 mode (~7GB VRAM) - use this on smaller GPUs or when ComfyUI is also running
+cd /d %~dp0
+call venv\Scripts\activate
+set HF_HUB_DISABLE_SYMLINKS_WARNING=1
+python app.py --nf4 %*
+pause