예제 소스

local_vision_mlx.py

"""
Apple Silicon 전용: Ollama 없이 MLX로 비전 LLM 추론 (mlx-vlm).

  pip install mlx-vlm

모델은 Hugging Face의 mlx-community 양자화 VLM(예: Qwen2-VL-2B-Instruct-4bit)을 쓰면
다운로드 후 로컬 캐시에서 로드됩니다.

환경변수:
  MLX_VISION_MODEL — 기본 모델 ID (미설정 시 아래 DEFAULT 사용)
  MLX_VISION_MAX_TOKENS — 생성 최대 토큰 (기본 512)

주의:
  - 차트 매매 신호 품질은 모델·프롬프트·후처리에 달려 있으며, MLX로 바꾼다고
    항상 정확해지지는 않습니다. temperature=0 에 가깝게 두는 것이 재현성에 유리합니다.
  - 기존 local_vision_llm.py 는 PyTorch+MPS 경로이고, 본 모듈은 MLX 경로입니다.
"""

from __future__ import annotations

import logging
import os
from io import BytesIO
from typing import Any

logger = logging.getLogger(__name__)

# mlx-community 쪽 Qwen2-VL 소형(로컬 테스트용 기본값)
DEFAULT_MLX_VISION_MODEL = "********"

_MLX_MODEL: Any = None
_MLX_PROCESSOR: Any = None
_MLX_CONFIG: Any = None
_LOADED_MODEL_ID: str | None = None


def resolve_mlx_model_id(requested: str | None) -> str:
    """CLI/호출 인자가 있으면 최우선, 없으면 MLX_VISION_MODEL, 없으면 기본값."""
    if requested and str(requested).strip():
        return str(requested).strip()
    env = os.environ.get("MLX_VISION_MODEL", "").strip()
    if env:
        return env
    return DEFAULT_MLX_VISION_MODEL


def _load_mlx(model_id: str) -> tuple[Any, Any, Any]:
    global _MLX_MODEL, _MLX_PROCESSOR, _MLX_CONFIG, _LOADED_MODEL_ID
    if (
        _MLX_MODEL is not None
        and _MLX_PROCESSOR is not None
        and _MLX_CONFIG is not None
        and _LOADED_MODEL_ID == model_id
    ):
        return _MLX_MODEL, _MLX_PROCESSOR, _MLX_CONFIG

    try:
        from mlx_vlm import load  # type: ignore
        from mlx_vlm.utils import load_config  # type: ignore
    except ImportError as e:
        raise RuntimeError(
            "mlx-vlm 이 필요합니다. Apple Silicon 맥에서: pip install mlx-vlm"
        ) from e

    logger.info("MLX-VLM 로드 중: %s (최초 실행 시 HF에서 받을 수 있음)", model_id)
    model, processor = load(model_id)
    config = load_config(model_id)

    _MLX_MODEL = model
    _MLX_PROCESSOR = processor
    _MLX_CONFIG = config
    _LOADED_MODEL_ID = model_id
    return model, processor, config


def mlx_vision_chat(
    *,
    image_bytes: bytes,
    user_text: str,
    model_id: str | None = None,
    max_tokens: int | None = None,
    temperature: float = 0.0,
) -> str:
    """차트 이미지(JPEG/PNG bytes) + 프롬프트 → 모델 텍스트 응답 (MLX)."""
    try:
        from mlx_vlm import generate  # type: ignore
        from mlx_vlm.prompt_utils import apply_chat_template  # type: ignore
    except ImportError as e:
        raise RuntimeError("mlx-vlm 패키지를 설치하세요: pip install mlx-vlm") from e

    try:
        from PIL import Image
    except ImportError as e:
        raise RuntimeError("Pillow가 필요합니다: pip install pillow") from e

    mid = resolve_mlx_model_id(model_id)
    model, processor, config = _load_mlx(mid)

    img = Image.open(BytesIO(image_bytes)).convert("RGB")
    formatted_prompt = apply_chat_template(
        processor,
        config,
        user_text,
        num_images=1,
    )

    cap = int(os.environ.get("MLX_VISION_MAX_TOKENS", "512"))
    ntok = max(32, min(max_tokens or cap, cap))

    try:
        out = generate(
            model,
            processor,
            formatted_prompt,
            [img],
            verbose=False,
            max_tokens=ntok,
            temperature=temperature,
        )
    except TypeError:
        out = generate(model, processor, formatted_prompt, [img], verbose=False)

    if isinstance(out, str):
        text_out = out.strip()
    else:
        text_out = str(out).strip()

    if not text_out:
        raise RuntimeError("MLX-VLM이 빈 응답을 반환했습니다.")
    return text_out


def _minimal_test_image_jpeg_bytes() -> bytes:
    from PIL import Image

    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:
    import argparse
    import sys

    parser = argparse.ArgumentParser(description="MLX-VLM 스모크 테스트 (이미지+프롬프트)")
    parser.add_argument("--image", default=None, help="테스트 이미지 경로. 없으면 더미 JPEG")
    parser.add_argument("--prompt", default="이 이미지를 한 문장으로 설명하세요.")
    parser.add_argument("--model", default=None, help="mlx-community 모델 ID")
    parser.add_argument("--max-tokens", type=int, default=256)
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    if args.image:
        with open(args.image, "rb") as f:
            img_bytes = f.read()
    else:
        img_bytes = _minimal_test_image_jpeg_bytes()

    text = mlx_vision_chat(
        image_bytes=img_bytes,
        user_text=args.prompt,
        model_id=args.model,
        max_tokens=args.max_tokens,
        temperature=0.0,
    )
    print(text, file=sys.stdout)


if __name__ == "__main__":
    main()