# /// 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=''' You are an expert prompt engineer for text-to-image models. Your task is to expand the image caption below into a highly effective image-generation prompt. The attached image shows what the caption describes. Think step by step about the request before writing the answer: - What is the subject and mood? - What visual styles, mediums, and lighting options would fit? Consider two or three alternatives and pick the one that best serves the caption. - What composition, framing, and grounded details will help the text-to-image model? Then output an expanded prompt in 1-2 paragraphs. Avoid run-on sentences. Follow these rules strictly: 1. **Faithfulness First:** Preserve all original subjects, actions, colors, and spatial relationships. Do not add new objects, props, characters, or animals unless the caption clearly implies them. 2. **Practical T2I Structure:** Write a prompt that a text-to-image model can parse cleanly. Group subjects with their own attributes and actions. Use grounded phrasing for poses, interactions, and spatial layout. 3. **Style Planning Stays Internal:** Use your internal reasoning to choose style, medium, framing, and lighting. Do not emit planning tags or wrappers in the visible answer body. 4. **Text Rendering:** If the caption requests visible text, quotes, labels, or typography, specify the exact text clearly and wrap requested words in quotes. 5. **Avoid Over-Specification:** Do not invent highly specific clothing, colors, materials, or scene details unless the input supports them. 6. **Structure:** Write one cohesive paragraph after the thinking block. No bullets, JSON, or markdown. 7. **Preserve User Medium:** When the caption explicitly requests a medium (e.g. "photo of", "photograph of", "illustration of", "painting of", "sketch of", "3D render of"), honor it. Do not pivot to a different medium to avoid difficulty — match the caption's stated intent. 8. **Respect Existing Detail:** If the caption is already detailed, polish and finalize it rather than heavily expanding — preserve its phrasing and direction. Notes about the content: The dataset contains body inflation images. The captions should be rewritten to normalize the descriptions according to a few definitions. These definitions are relative to a hypothetical uninflated version of the character. Dedicate a paragraph to this. Of torso sections (above the legs, below the head, between the arms): "expanded (x)" - Only slightly larger than normal; noticeably prominent. "inflated (x)" - Significantly larger than normal; composed of round arcs. "balloon (x)" - Extremely large. Similar volume to the character's uninflated body. "blimp(x)" - Much larger than the uninflated character's entire volume. Fits several of the character inside of it. Using this keyword should mean the subject's body diverges heavily from a humanoid form. "hyper (x)" - Similar to "blimp" but can substitute when the rest of the character looks normal. "balloon" and "hyper" can be used in parallel. Of limbs and extremities (arms/legs, hands/feet): "thickened" - Noticeably swollen. Use cues about clothing or skeletal proportions to discern. "puffed" - Partly or completely melded into the torso; puffed limbs make an inflated character's silhouette diverge from a perfect sphere. "sunken" - Usually of extremities; means the relevant limbs are partly or completely absorbed into an inflated character's form *without* affecting its silhouette, leaving the extremity to protrude. User's Input: ''') 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.") parser.add_argument('--test', action='store_true', help='Process only the first image, print the result, and exit without writing files.') args = parser.parse_args() if args.test: args.dry_run = True 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) prompt = args.prompt sidecar_file = file.with_suffix('.txt') if sidecar_file.exists(): sidecar_prompt = sidecar_file.read_text(encoding='utf-8') prompt = f"{prompt}\n{sidecar_prompt}" 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": 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__': all_files = [f for f in args.folder.glob('*') if f.is_file() and f.suffix != '.txt'] if args.test: if not all_files: print('No images found to test.') exit(0) test_file = all_files[0] print(f'Test image: {test_file}') file, text = caption_single(test_file) print('-'*80) print(file) print(text + '\n') print('-'*80) exit(0) image_files = [] skipped = 0 for file in all_files: 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)