Moxin-7B VLM

use moxin_vlm_mlx::{load_model, load_tokenizer, normalize_dino, normalize_siglip, Generate, KVCache};
use mlx_rs::Array;
use image::imageops::FilterType;

// Load model and tokenizer
let mut vlm = load_model("~/.OminiX/models/moxin-vlm-7b")?;
let tokenizer = load_tokenizer("~/.OminiX/models/moxin-vlm-7b")?;

// Load and preprocess image to 224×224
let img = image::open("photo.jpg")?;
let img = img.resize_exact(224, 224, FilterType::CatmullRom);
let rgb = img.to_rgb8();

// Convert to [1, 224, 224, 3] float32 tensor in [0, 1]
let pixels: Vec<f32> = rgb
    .pixels()
    .flat_map(|p| p.0.iter().map(|&v| v as f32 / 255.0))
    .collect();
let tensor = Array::from_slice(&pixels, &[1, 224, 224, 3]);

// Normalize for each encoder
let dino_img = normalize_dino(&tensor)?;
let siglip_img = normalize_siglip(&tensor)?;

// Format prompt (Prismatic "Pure" format)
let prompt_text = format!("In: {}\nOut:", "Describe this image.");

// Tokenize with BOS token
let encoding = tokenizer.encode(prompt_text.as_str(), true)?;
let input_ids = Array::from_iter(
    encoding.get_ids().iter().map(|&id| id as i32),
    &[1, encoding.get_ids().len() as i32],
);

// Generate tokens
let mut cache: Vec<KVCache> = Vec::new();
let generator = Generate::new(
    &mut vlm,
    &mut cache,
    0.0,  // temperature (0 = greedy)
    dino_img,
    siglip_img,
    input_ids,
);

let eos_token_id = 2u32; // </s>
let mut generated = Vec::new();

for token_result in generator.take(256) {
    let token = token_result?;
    let token_id = token.item::<u32>();
    
    if token_id == eos_token_id {
        break;
    }
    
    generated.push(token_id);
    
    // Decode and print incrementally
    let text = tokenizer.decode(&generated, true)?;
    print!("{}", text);
}

let generator = Generate::new(
    &mut vlm,
    &mut cache,
    0.0,  // temperature = 0 for greedy
    dino_img,
    siglip_img,
    input_ids,
);