examples/vision_example.py

#!/usr/bin/env python3
"""
Sixpert K1 - Vision/Multimodal Example
=======================================
Demonstrates how to use Sixpert K1's vision capabilities with
image inputs using llama-cpp-python.

Usage:
    python vision_example.py --image ./photo.jpg --prompt "Describe this image in detail"
"""

import argparse
import base64
import sys

try:
    from llama_cpp import Llama
except ImportError:
    print("Installing llama-cpp-python...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
    from llama_cpp import Llama


def image_to_base64(image_path: str) -> str:
    """Convert an image file to base64 data URL."""
    with open(image_path, "rb") as f:
        data = base64.b64encode(f.read()).decode("utf-8")

    # Detect MIME type from file extension
    if image_path.lower().endswith(".png"):
        mime = "image/png"
    elif image_path.lower().endswith(".jpg") or image_path.lower().endswith(".jpeg"):
        mime = "image/jpeg"
    elif image_path.lower().endswith(".webp"):
        mime = "image/webp"
    elif image_path.lower().endswith(".gif"):
        mime = "image/gif"
    else:
        mime = "image/png"

    return f"data:{mime};base64,{data}"


def analyze_image(
    model_path: str,
    image_path: str,
    prompt: str,
    max_tokens: int = 2048,
):
    """Analyze an image using Sixpert K1."""
    print(f"Loading Sixpert K1...")
    llm = Llama(
        model_path=model_path,
        n_ctx=131072,
        n_gpu_layers=-1,
        verbose=False,
    )

    image_data = image_to_base64(image_path)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": prompt,
                },
                {
                    "type": "image_url",
                    "image_url": {"url": image_data},
                },
            ],
        },
    ]

    print(f"\nPrompt: {prompt}")
    print(f"Image: {image_path}")
    print("-" * 40)
    print("Analyzing image...")
    print("-" * 40)

    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.7,
        stream=True,
    )

    for chunk in response:
        delta = chunk["choices"][0]["delta"].get("content", "")
        if delta:
            print(delta, end="", flush=True)

    print("\n")


def main():
    parser = argparse.ArgumentParser(description="Sixpert K1 Vision Example")
    parser.add_argument("--model", type=str, default="SixpertK1.gguf", help="Path to GGUF model")
    parser.add_argument("--image", type=str, required=True, help="Path to input image")
    parser.add_argument("--prompt", type=str, default="Describe this image in detail.", help="Prompt")
    parser.add_argument("--max-tokens", type=int, default=2048, help="Max tokens")

    args = parser.parse_args()
    analyze_image(args.model, args.image, args.prompt, args.max_tokens)


if __name__ == "__main__":
    main()