| """Pluggable local captioners for LoRA dataset building. |
| |
| All captioners expose the same interface: ``caption(images: list[PIL.Image]) -> str``. |
| ``images`` is a list so video callers can pass multiple sampled frames; single-image |
| captioners just use the first one. |
| |
| - JoyCaptionCaptioner: fancyfeast/llama-joycaption-beta-one-hf-llava (AutoProcessor + |
| LlavaForConditionalGeneration, confirmed against the model card). The exact prompt |
| wording for a dedicated "training caption" mode is not documented on the model card, |
| so the instruction text here is our own default and is overridable via --caption-prompt. |
| - Florence2Captioner: microsoft/Florence-2-large, <MORE_DETAILED_CAPTION> task. |
| - WD14Captioner: SmilingWolf/wd-vit-tagger-v3 onnx tagger (booru-style tags), preprocessing |
| confirmed against the reference wd-tagger Space source (white-pad-to-square, RGB->BGR, |
| no normalization, category 9/4/0 = rating/character/general). |
| """ |
| import csv |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
|
|
| class JoyCaptionCaptioner: |
| MODEL_ID = "fancyfeast/llama-joycaption-beta-one-hf-llava" |
| DEFAULT_PROMPT = ( |
| "Write a detailed but concise caption for this image, suitable for training " |
| "an image/video LoRA. Describe the subject's appearance, pose, clothing, " |
| "setting, and lighting in plain descriptive sentences. Do not guess names." |
| ) |
|
|
| def __init__(self, device="cuda", load_in_4bit=False, prompt=None, system_prompt=None): |
| from transformers import AutoProcessor, LlavaForConditionalGeneration |
|
|
| self.processor = AutoProcessor.from_pretrained(self.MODEL_ID) |
| load_kwargs = {} |
| if load_in_4bit: |
| from transformers import BitsAndBytesConfig |
| load_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True) |
| load_kwargs["device_map"] = "auto" |
| else: |
| load_kwargs["torch_dtype"] = torch.bfloat16 |
| load_kwargs["device_map"] = device |
| self.model = LlavaForConditionalGeneration.from_pretrained(self.MODEL_ID, **load_kwargs) |
| self.model.eval() |
|
|
| self.system_prompt = system_prompt or "You are a helpful image captioner." |
| self.prompt = prompt or self.DEFAULT_PROMPT |
|
|
| def caption(self, images): |
| image = images[0].convert("RGB") |
| convo = [ |
| {"role": "system", "content": self.system_prompt}, |
| {"role": "user", "content": self.prompt}, |
| ] |
| convo_string = self.processor.apply_chat_template( |
| convo, tokenize=False, add_generation_prompt=True |
| ) |
| inputs = self.processor(text=[convo_string], images=[image], return_tensors="pt") |
| inputs = {k: v.to(self.model.device) for k, v in inputs.items()} |
| if "pixel_values" in inputs: |
| inputs["pixel_values"] = inputs["pixel_values"].to(self.model.dtype) |
|
|
| with torch.no_grad(): |
| generate_ids = self.model.generate( |
| **inputs, |
| max_new_tokens=300, |
| do_sample=True, |
| suppress_tokens=None, |
| use_cache=True, |
| temperature=0.6, |
| top_k=None, |
| top_p=0.9, |
| )[0] |
|
|
| generate_ids = generate_ids[inputs["input_ids"].shape[1]:] |
| text = self.processor.tokenizer.decode( |
| generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False |
| ) |
| return text.strip() |
|
|
|
|
| class Florence2Captioner: |
| MODEL_ID = "microsoft/Florence-2-large" |
| TASK = "<MORE_DETAILED_CAPTION>" |
|
|
| def __init__(self, device="cuda"): |
| from transformers import AutoModelForCausalLM, AutoProcessor |
|
|
| dtype = torch.float16 if device.startswith("cuda") else torch.float32 |
| self.device = device |
| self.model = ( |
| AutoModelForCausalLM.from_pretrained(self.MODEL_ID, trust_remote_code=True, torch_dtype=dtype) |
| .to(device) |
| .eval() |
| ) |
| self.processor = AutoProcessor.from_pretrained(self.MODEL_ID, trust_remote_code=True) |
|
|
| def caption(self, images): |
| image = images[0].convert("RGB") |
| inputs = self.processor(text=self.TASK, images=image, return_tensors="pt") |
| inputs = {k: v.to(self.device, self.model.dtype) if v.is_floating_point() else v.to(self.device) |
| for k, v in inputs.items()} |
|
|
| with torch.no_grad(): |
| generated_ids = self.model.generate( |
| input_ids=inputs["input_ids"], |
| pixel_values=inputs["pixel_values"], |
| max_new_tokens=200, |
| num_beams=3, |
| do_sample=False, |
| ) |
| text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0] |
| parsed = self.processor.post_process_generation( |
| text, task=self.TASK, image_size=(image.width, image.height) |
| ) |
| return parsed[self.TASK].strip() |
|
|
|
|
| class WD14Captioner: |
| REPO_ID = "SmilingWolf/wd-vit-tagger-v3" |
| MODEL_FILE = "model.onnx" |
| TAGS_FILE = "selected_tags.csv" |
|
|
| def __init__(self, general_thresh=0.35, character_thresh=0.75, use_gpu=True): |
| import onnxruntime as ort |
| from huggingface_hub import hf_hub_download |
|
|
| model_path = hf_hub_download(self.REPO_ID, self.MODEL_FILE) |
| tags_path = hf_hub_download(self.REPO_ID, self.TAGS_FILE) |
|
|
| providers = ["CPUExecutionProvider"] |
| if use_gpu: |
| providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] |
| self.session = ort.InferenceSession(model_path, providers=providers) |
| self.input_name = self.session.get_inputs()[0].name |
| self.output_name = self.session.get_outputs()[0].name |
| self.target_size = int(self.session.get_inputs()[0].shape[1]) |
|
|
| self.general_thresh = general_thresh |
| self.character_thresh = character_thresh |
| self.tag_names, self.general_idx, self.character_idx = self._load_tags(tags_path) |
|
|
| @staticmethod |
| def _load_tags(path): |
| names, general, character = [], [], [] |
| with open(path, newline="", encoding="utf-8") as f: |
| for i, row in enumerate(csv.DictReader(f)): |
| names.append(row["name"]) |
| cat = int(row["category"]) |
| if cat == 4: |
| character.append(i) |
| elif cat != 9: |
| general.append(i) |
| return names, general, character |
|
|
| def _preprocess(self, image: Image.Image): |
| image = image.convert("RGBA") |
| canvas = Image.new("RGBA", image.size, (255, 255, 255, 255)) |
| canvas.alpha_composite(image) |
| image = canvas.convert("RGB") |
|
|
| w, h = image.size |
| size = max(w, h) |
| padded = Image.new("RGB", (size, size), (255, 255, 255)) |
| padded.paste(image, ((size - w) // 2, (size - h) // 2)) |
| padded = padded.resize((self.target_size, self.target_size), Image.BICUBIC) |
|
|
| arr = np.asarray(padded, dtype=np.float32) |
| arr = arr[:, :, ::-1] |
| return np.expand_dims(arr, axis=0) |
|
|
| def caption(self, images): |
| image = images[0] |
| arr = self._preprocess(image) |
| preds = self.session.run([self.output_name], {self.input_name: arr})[0][0] |
|
|
| general = [(self.tag_names[i], preds[i]) for i in self.general_idx if preds[i] >= self.general_thresh] |
| character = [(self.tag_names[i], preds[i]) for i in self.character_idx if preds[i] >= self.character_thresh] |
| general.sort(key=lambda t: -t[1]) |
| character.sort(key=lambda t: -t[1]) |
|
|
| tags = [t.replace("_", " ") for t, _ in character] + [t.replace("_", " ") for t, _ in general] |
| return ", ".join(tags) |
|
|
|
|
| def build_captioner(args): |
| name = args.captioner |
| if name == "joycaption": |
| return JoyCaptionCaptioner( |
| device=args.device, |
| load_in_4bit=args.load_in_4bit, |
| prompt=args.caption_prompt, |
| ) |
| if name == "florence2": |
| return Florence2Captioner(device=args.device) |
| if name == "wd14": |
| return WD14Captioner( |
| general_thresh=args.wd14_general_thresh, |
| character_thresh=args.wd14_character_thresh, |
| use_gpu=args.device.startswith("cuda"), |
| ) |
| raise ValueError(f"Unknown captioner: {name}") |
|
|
|
|
| def with_trigger(caption, trigger): |
| caption = (caption or "").strip() |
| if not trigger: |
| return caption |
| return f"{trigger}, {caption}" if caption else trigger |
|
|
