#!/usr/bin/env python3
"""Moondream 2 — пакетная генерация подписей к картинкам.

Usage:
    python caption.py <image_or_dir>...

Для каждого изображения:
  - проверка, не битое ли
  - ресайз до 768 по большей стороне (пропорции сохраняются)
  - конверт в RGB
  - генерация normal caption
  - сохранение в <имя_файла>.txt
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image, UnidentifiedImageError
import sys, os, torch
from pathlib import Path
from tqdm import tqdm

MODEL_ID = "vikhyatk/moondream2"
REVISION = "2025-06-21"
MAX_SIZE = 768


def load_image(path: str) -> Image.Image | None:
    """Загружает, конвертит в RGB, ресайзит. Возвращает None если битая."""
    try:
        img = Image.open(path)
        img.load()  # проверить, что не битая
    except (UnidentifiedImageError, OSError, Exception) as e:
        print(f"  [SKIP] не удалось прочитать: {e}")
        return None

    # RGB
    if img.mode != "RGB":
        img = img.convert("RGB")

    # Ресайз: большая сторона = MAX_SIZE
    w, h = img.size
    if max(w, h) > MAX_SIZE:
        scale = MAX_SIZE / max(w, h)
        new_w = round(w * scale)
        new_h = round(h * scale)
        img = img.resize((new_w, new_h), Image.LANCZOS)

    return img


def process_image(model, tokenizer, img_path: str):
    """Обрабатывает одно изображение."""
    path = Path(img_path)
    txt_path = path.with_suffix(".txt")

    if txt_path.exists():
        return

    img = load_image(img_path)
    if img is None:
        return

    caption = model.caption(img, length="normal")["caption"]
    txt_path.write_text(caption.strip() + "\n")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    # Собираем все файлы из аргументов
    EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif"}
    paths: list[str] = []
    for arg in sys.argv[1:]:
        p = Path(arg)
        if p.is_dir():
            paths.extend(str(f) for f in p.rglob("*") if f.suffix.lower() in EXTS)
        elif p.is_file():
            paths.append(arg)
        else:
            print(f"Не найден: {arg}")

    if not paths:
        print("Нет файлов для обработки.")
        sys.exit(1)

    if not torch.cuda.is_available():
        print("Ошибка: CUDA не доступна. Нужен GPU.")
        sys.exit(1)
    device = "cuda:0"
    print(f"Загрузка модели {MODEL_ID} (rev {REVISION}) на {device}...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        revision=REVISION,
        trust_remote_code=True,
        dtype=torch.bfloat16,
    ).to(device)
    model = torch.compile(model, mode="reduce-overhead")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION)

    print()
    for p in tqdm(paths, desc="Captioning", unit="img"):
        process_image(model, tokenizer, p)

    print("\nГотово.")


if __name__ == "__main__":
    main()