# /// script # requires-python = ">=3.13" # dependencies = [ # "pillow>=12.2.0", # "openai>=1.0.0", # ] # /// import argparse from pathlib import Path import math import base64 import io import concurrent.futures from openai import OpenAI from PIL import Image MAX_SEQ_LENGTH = 16384 MAX_IMAGE_PIXELS = 1024**2 BATCH_SIZE = 4 parser = argparse.ArgumentParser() parser.add_argument('--model', type=str, required=True, help="Model identifier for the OpenAI-compatible API.") parser.add_argument('--folder', type=Path, required=True) parser.add_argument('--api_base', type=str, default=None, help="Base URL for the OpenAI-compatible API.") parser.add_argument('--api_key', type=str, required=True, help="API key for the OpenAI-compatible API.") parser.add_argument('--system_prompt', type=str) parser.add_argument('--prompt', type=str, default=''' Write a descriptive caption for this body inflation image that would allow a modern, natural-language-based image generation model to create something like it. The subject is inflated into a compound shape; we're looking to capture it in detail to train a LoRA on multi-part body inflation as a concept. The prompt will be in four parts. The first part is an overview. For example: " from is . is . 's body is inflated to a huge round shape." The second part is the inflation section. Start from the largest/most prominent region and then move outwards. For example: "'s belly is inflated into a wide sphere that pushes out in all directions. 's arms and legs are massive and round. 's [forearms, shoulders, upper arms] are pumped full, [thighs, hips, calves] plump and bloated. 's breasts are huge and bouncy, squeezing out from between thickened arms." The third part describes the subject's pose, taking into consideration what you said before. Include the locations of the subject's extremities to help orient the MMDiT. For example: " stands upright in a wide stance, arms resting between breasts and belly. holds a small helium canister in puffy hands. 's head is on top of body, ." For the fourth part, briefly describe the art style, background and 's expression. Note: - Every part is mandatory. Don't use headers or any formatting, just separate with newlines. - Keep to 2 grammatical clauses per sentence as described. The text encoder in use is Qwen3 0.6B, so it's very small and easy to confuse. - If you can't be *certain* of the character's identity, simply describing them is enough; just use a consistent synechdoche in place of (e.g. "the woman".) For anthropomorphic characters, that synechdoche should include their species (e.g. "the dog man".) - Don't caption dialogue, text, or watermarks. Pretend they don't exist. The encoder WILL hyperfocus on them if you do. - You needn't follow the examples exactly; describe what you *see,* and write the caption in the general shape of the examples. ''') parser.add_argument('--prompt_file', type=Path, default=None, help="Path to a text file containing the prompt. Overrides --prompt if given.") parser.add_argument('--max_tokens', type=int, default=8192) parser.add_argument('--temperature', type=float, default=0.7) parser.add_argument('--overwrite', action='store_true', help='Re-caption images even if a .txt file already exists.') parser.add_argument('--dry_run', action='store_true', help="Don't write any files.") args = parser.parse_args() assert args.folder.is_dir() if args.prompt_file is not None: assert args.prompt_file.is_file(), f"Prompt file not found: {args.prompt_file}" args.prompt = args.prompt_file.read_text(encoding='utf-8') client = OpenAI(base_url=args.api_base, api_key=args.api_key) def rgba_to_rgb( image: Image.Image, background_color=(255, 255, 255)) -> Image.Image: """Convert an RGBA image to RGB with filled background color.""" assert image.mode == "RGBA" converted = Image.new("RGB", image.size, background_color) converted.paste(image, mask=image.split()[3]) # 3 is the alpha channel return converted def convert_image_mode(image: Image.Image, to_mode: str): if image.mode == to_mode: return image elif image.mode == "RGBA" and to_mode == "RGB": return rgba_to_rgb(image) else: return image.convert(to_mode) def limit_size(img): w, h = img.size num_pixels = w*h if num_pixels > MAX_IMAGE_PIXELS: scale = math.sqrt(num_pixels / MAX_IMAGE_PIXELS) img.thumbnail((int(w/scale), int(h/scale))) return img def image_to_base64(image: Image.Image) -> str: buffer = io.BytesIO() image.save(buffer, format="JPEG") return base64.b64encode(buffer.getvalue()).decode("utf-8") def caption_single(file: Path) -> tuple[Path, str]: img = Image.open(file) img = convert_image_mode(limit_size(img), 'RGB') b64 = image_to_base64(img) messages = [] if args.system_prompt: messages.append({"role": "system", "content": args.system_prompt}) messages.append({ "role": "user", "content": [ { "type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}, }, { "type": "text", "text": args.prompt, }, ], }) response = client.chat.completions.create( model=args.model, messages=messages, max_tokens=args.max_tokens, temperature=args.temperature, ) text = response.choices[0].message.content or "" txt_file = file.with_suffix('.txt') if not args.dry_run: with open(txt_file, 'w', encoding='utf-8') as f: f.write(text) return file, text if __name__ == '__main__': image_files = [] skipped = 0 for file in args.folder.glob('*'): if file.is_dir() or file.suffix == '.txt': continue if not args.overwrite and file.with_suffix('.txt').exists(): skipped += 1 continue image_files.append(file) if skipped: print(f'Skipped {skipped} image(s) that already have captions (use --overwrite to re-caption).') num_images = len(image_files) for i in range(0, num_images, BATCH_SIZE): batch = image_files[i:i+BATCH_SIZE] print(f'Progress: {i} / {num_images}') with concurrent.futures.ThreadPoolExecutor(max_workers=BATCH_SIZE) as executor: futures = [executor.submit(caption_single, f) for f in batch] for future in concurrent.futures.as_completed(futures): file, text = future.result() if file == batch[0]: print('-'*80) print(file) print(text + '\n') print('-'*80)