예제 소스

local_vision_llm.py

"""
HF Router 대신 로컬에서 비전 LLM 추론 (transformers + Qwen2.5-VL 계열).

활성화: 환경변수 HF_VISION_LOCAL=1
모델: HF_LOCAL_VISION_MODEL (허브 ID 또는 로컬 디렉터리). 미설정 시 원격 HF_VISION_MODEL을 참고하되,
      397B 등 대형 모델 ID는 로컬에 맞게 Qwen/Qwen2.5-VL-3B-Instruct 로 대체합니다.

사전 다운로드(선택): 같은 환경에서 한 번 실행하면 Hugging Face 캐시(~/.cache/huggingface)에 저장됩니다.

CLI 테스트:
  python local_vision_llm.py --download-only
  python local_vision_llm.py --image path/to/chart.jpg --prompt "차트를 보고 매수/매도/유지 중 하나만 말하세요."
  # 이미지 없으면 작은 더미 JPEG로 스모크 테스트

필요 패키지: pip install transformers accelerate torch torchvision qwen-vl-utils pillow huggingface_hub
(qwen_vl_utils 가 이미지 로드 시 torchvision 을 쓰므로 torchvision 누락 시 No module named 'torchvision' 발생)
"""

from __future__ import annotations

import base64
import logging
import os
from typing import Any

logger = logging.getLogger(__name__)

# 로컬 기본: VRAM이 적은 환경에서도 동작 가능한 instruct VLM
DEFAULT_LOCAL_VISION_MODEL = "Qwen/Qwen2.5-VL-3B-Instruct"

_MODEL: Any = None
_PROCESSOR: Any = None
_LOADED_MODEL_ID: str | None = None


def is_local_vision_enabled() -> bool:
    return os.environ.get("HF_VISION_LOCAL", "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )


def resolve_local_model_id(requested: str | None) -> str:
    """호출 인자(requested)가 있으면 최우선. 없으면 HF_LOCAL_VISION_MODEL. 397B 원격 기본값은 로컬에서 3B로 대체."""

    def _maybe_warn_397b(mid: str) -> str | None:
        low = mid.lower()
        if "397b" in low or "a17b" in low:
            logger.warning(
                "원격 기본 모델(%s)은 로컬 추론에 부적합합니다. "
                "로컬 기본값 %s 을(를) 사용합니다. HF_LOCAL_VISION_MODEL 로 변경 가능합니다.",
                mid,
                DEFAULT_LOCAL_VISION_MODEL,
            )
            return DEFAULT_LOCAL_VISION_MODEL
        return None

    rid = (requested or "").strip()
    if rid:
        repl = _maybe_warn_397b(rid)
        return repl if repl is not None else rid

    override = os.environ.get("HF_LOCAL_VISION_MODEL", "").strip()
    if override:
        repl = _maybe_warn_397b(override)
        return repl if repl is not None else override

    return DEFAULT_LOCAL_VISION_MODEL


def predownload_local_model(model_id: str | None = None) -> str:
    """허브에서 가중치를 캐시로 받기만 할 때 사용 (추론 없음)."""
    from huggingface_hub import snapshot_download

    mid = resolve_local_model_id(model_id or "")
    snapshot_download(mid)
    logger.info("snapshot_download 완료: %s", mid)
    return mid


def _image_bytes_to_data_url(image_bytes: bytes) -> str:
    if image_bytes[:8] == b"\x89PNG\r\n\x1a\n":
        mime = "image/png"
    else:
        mime = "image/jpeg"
    b64 = base64.standard_b64encode(image_bytes).decode("ascii")
    return f"data:{mime};base64,{b64}"


def _load_model(model_id: str) -> tuple[Any, Any]:
    global _MODEL, _PROCESSOR, _LOADED_MODEL_ID
    if _MODEL is not None and _PROCESSOR is not None and _LOADED_MODEL_ID == model_id:
        return _MODEL, _PROCESSOR

    try:
        from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
    except ImportError as e:
        raise RuntimeError(
            "로컬 비전 모델에 필요한 패키지가 없습니다. 예: "
            "pip install transformers accelerate torch qwen-vl-utils pillow"
        ) from e

    try:
        from qwen_vl_utils import process_vision_info
    except ImportError as e:
        raise RuntimeError(
            "qwen-vl-utils 가 필요합니다. 예: pip install qwen-vl-utils"
        ) from e

    _ = process_vision_info  # noqa: F841 — import 검증용

    logger.info("로컬 VLM 로드 중: %s (최초 실행 시 허브에서 다운로드될 수 있음)", model_id)
    print(
        f"[local_vision_llm] VLM 가중치 로드 중: {model_id} (수 분 걸릴 수 있음)",
        flush=True,
    )

    dtype = "auto"
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        model_id,
        torch_dtype=dtype,
        device_map="auto",
        trust_remote_code=True,
    )
    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

    _MODEL = model
    _PROCESSOR = processor
    _LOADED_MODEL_ID = model_id
    return model, processor


def local_vision_chat(
    *,
    image_bytes: bytes,
    user_text: str,
    model_id: str | None,
    max_new_tokens: int = 1024,
) -> str:
    """차트 이미지(JPEG/PNG bytes) + 프롬프트 → 모델 텍스트 응답."""
    from qwen_vl_utils import process_vision_info

    mid = resolve_local_model_id(model_id)
    model, processor = _load_model(mid)

    data_url = _image_bytes_to_data_url(image_bytes)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": data_url},
                {"type": "text", "text": user_text},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )

    import torch

    device = next(model.parameters()).device
    inputs = inputs.to(device)

    # 과도한 토큰은 VRAM 폭주 방지
    cap = int(os.environ.get("HF_LOCAL_MAX_NEW_TOKENS_CAP", "4096"))
    ntok = max(64, min(max_new_tokens, cap))

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=ntok,
            do_sample=False,
        )

    input_ids = inputs.input_ids
    generated_ids_trimmed = [
        out_ids[len(inp) :] for inp, out_ids in zip(input_ids, generated_ids)
    ]
    out_list = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    text_out = (out_list[0] if out_list else "").strip()
    if not text_out:
        raise RuntimeError("로컬 VLM이 빈 응답을 반환했습니다.")
    return text_out


def _minimal_test_image_jpeg_bytes() -> bytes:
    """의존성 검증용 작은 JPEG (Pillow 필요)."""
    from io import BytesIO

    try:
        from PIL import Image
    except ImportError as e:
        raise RuntimeError("더미 이미지 생성에 Pillow가 필요합니다. pip install pillow") from e

    buf = BytesIO()
    Image.new("RGB", (128, 128), color=(90, 120, 180)).save(buf, format="JPEG", quality=88)
    return buf.getvalue()


def main(argv: list[str] | None = None) -> None:
    """CLI: 모델 다운로드만 하거나, 더미/파일 이미지로 한 번 추론."""
    import argparse
    import sys

    parser = argparse.ArgumentParser(
        description="로컬 비전 LLM(Qwen2.5-VL) 다운로드 또는 스모크 테스트",
    )
    parser.add_argument(
        "--download-only",
        action="store_true",
        help="가중치만 Hugging Face 캐시에 받고 종료 (추론 없음)",
    )
    parser.add_argument(
        "--image",
        default=None,
        help="테스트 이미지 경로(JPEG/PNG). 없으면 메모리에서 더미 JPEG 생성",
    )
    parser.add_argument(
        "--prompt",
        default="이 이미지를 한두 문장으로 설명하세요.",
        help="사용자 프롬프트",
    )
    parser.add_argument(
        "--model",
        default=None,
        help="모델 ID 또는 로컬 경로 (기본: HF_LOCAL_VISION_MODEL 또는 Qwen2.5-VL-3B-Instruct)",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=256,
        help="max_new_tokens (기본 256, 스모크용으로 짧게)",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="DEBUG 로그")
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s %(name)s: %(message)s",
    )

    if args.download_only:
        mid = predownload_local_model(args.model)
        print(f"[OK] snapshot_download 완료: {mid}", file=sys.stderr)
        return

    if args.image:
        with open(args.image, "rb") as f:
            img_bytes = f.read()
    else:
        img_bytes = _minimal_test_image_jpeg_bytes()
        logger.info("이미지 경로 없음 — 더미 JPEG(%d bytes)로 테스트합니다.", len(img_bytes))

    mid = resolve_local_model_id(args.model)
    logger.info("모델: %s", mid)

    text = local_vision_chat(
        image_bytes=img_bytes,
        user_text=args.prompt,
        model_id=args.model,
        max_new_tokens=args.max_tokens,
    )
    print(text)


if __name__ == "__main__":
    main()