nexaai 1.0.14__cp310-cp310-macosx_14_0_universal2.whl → 1.0.15__cp310-cp310-macosx_14_0_universal2.whl

nexaai/_version.py

@@ -1,4 +1,4 @@
 # This file is generated by CMake from _version.py.in
 # Do not modify this file manually - it will be overwritten
 
-__version__ = "1.0.14" 
+__version__ = "1.0.15" 

nexaai/mlx_backend/vlm/generate_qwen3_vl.py

@@ -0,0 +1,261 @@
+import argparse
+import json
+import sys
+import os
+import mlx.core as mx
+import mlx.nn as nn
+import time
+from PIL import Image
+import requests
+import numpy as np
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+# Add current directory to path for imports
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(curr_dir)
+sys.path.append(os.path.dirname(curr_dir))
+
+# Add the qwen3vl model directory to path
+qwen3vl_dir = os.path.join(curr_dir, "modeling", "models", "qwen3_vl")
+sys.path.append(qwen3vl_dir)
+
+# Import required modules for quantized loading
+from transformers import AutoTokenizer
+
+# Try relative imports first, fallback to sys.path approach for Nuitka compatibility
+try:
+    from .modeling.models.qwen3_vl.llm_common.generate import nexa_generate_step
+    from .modeling.models.qwen3_vl.llm_common.cache import make_prompt_cache
+    from .modeling.models.qwen3_vl.qwen3vl import (
+        VEGModel, LLMModel, ModelArgs, VisionConfig, TextConfig, handle_multimodal_embeds
+    )
+    from .modeling.models.qwen3_vl.processor import Qwen3VLProcessor
+except ImportError:
+    # Fallback for Nuitka compiled environment - use sys.path approach
+    from llm_common.generate import nexa_generate_step
+    from llm_common.cache import make_prompt_cache
+    from qwen3vl import VEGModel, LLMModel, ModelArgs, VisionConfig, TextConfig, handle_multimodal_embeds
+    from processor import Qwen3VLProcessor
+
+from ml import ChatMessage
+from dataclasses import dataclass
+from typing import Any, Generator, List, Optional, Sequence, Tuple, Union
+from .generate import GenerationResult 
+
+@dataclass
+class Qwen3VLBundledModel:
+    """Container for Qwen3-VL vision and language models."""
+    vision_model: VEGModel
+    llm_model: LLMModel
+
+
+def _ensure_list(x: Union[str, List[str], None]) -> Optional[List[str]]:
+    if x is None:
+        return None
+    return x if isinstance(x, list) else [x]
+
+
+def load_qwen3_vl(
+    path_or_repo: str,
+    adapter_path: Optional[str] = None,
+    lazy: bool = False,
+    revision: Optional[str] = None,
+    **kwargs,
+) -> Tuple[Qwen3VLBundledModel, Qwen3VLProcessor]:
+    """Load Qwen3-VL quantized models and processor.
+
+    Parameters are aligned with .generate.load for compatibility.
+    """
+    model_path = Path(path_or_repo)
+    if not model_path.exists():
+        if "/" in path_or_repo:
+            model_path = Path(snapshot_download(
+                repo_id=path_or_repo, repo_type="model", revision=revision))
+        else:
+            # Fallback to local modelfiles directory
+            model_path = Path(qwen3vl_dir) / "modelfiles"
+            if not model_path.exists():
+                model_path = Path(curr_dir) / "modelfiles"
+
+    # Model configs (kept identical to main)
+    vision_config = VisionConfig(
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_heads=16,
+        num_hidden_layers=24,
+        patch_size=16,
+        temporal_patch_size=2,
+        in_channels=3,
+        hidden_act="gelu",
+        spatial_merge_size=2,
+        out_hidden_size=2560,
+        num_position_embeddings=2304,
+        deepstack_visual_indexes=[5, 11, 17],
+    )
+
+    text_config = TextConfig(
+        model_type="qwen3vl",
+        hidden_size=2560,
+        num_hidden_layers=36,
+        intermediate_size=9728,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        rms_norm_eps=1e-6,
+        vocab_size=151936,
+        max_position_embeddings=32768,
+        rope_theta=5000000.0,
+        head_dim=128,
+        tie_word_embeddings=True,
+        attention_bias=False,
+        attention_dropout=0.0,
+        rope_scaling={"mrope_section": [24, 20, 20],
+                      "rope_type": "default", "type": "default"},
+    )
+
+    vision_model = VEGModel(vision_config)
+    llm_model = LLMModel(text_config)
+
+    # Try to load LLM model from available files in order of preference
+    preferred_order = [
+        ("qwen3vl-llm-4B-q4_0.safetensors", 4),
+        ("qwen3vl-llm-4B-q8_0.safetensors", 8),
+        ("qwen3vl-llm-4B-f32.safetensors", 32)
+    ]
+
+    llm_weights_path = None
+    quantization_bits = None
+    
+    # Try loading in order of preference
+    for filename, bits in preferred_order:
+        candidate_path = model_path / filename
+        if candidate_path.exists():
+            llm_weights_path = candidate_path
+            quantization_bits = bits
+            break
+    
+    if llm_weights_path is None:
+        # Fallback to original hardcoded path for backward compatibility
+        llm_weights_path = model_path / "qwen3vl-llm-4B-q4_0.safetensors"
+        quantization_bits = 4
+
+    vision_weights_path = model_path / "qwen3vl-vision-4B-f16.safetensors"
+
+    if not vision_weights_path.exists() or not llm_weights_path.exists():
+        raise FileNotFoundError(
+            f"Missing safetensors. Vision: {vision_weights_path}, LLM: {llm_weights_path}"
+        )
+
+    # Load weights (vision fp16, llm with detected quantization)
+    vision_model.set_dtype(mx.float16)
+    vision_model.load_weights(str(vision_weights_path), strict=True)
+
+    # Apply quantization if needed and load LLM weights
+    if quantization_bits in [4, 8]:
+        nn.quantize(llm_model, bits=quantization_bits, group_size=64,
+                    class_predicate=quant_predicate)
+    # For f32 (32-bit), no quantization needed
+    
+    llm_model.load_weights(str(llm_weights_path), strict=True)
+
+    # Tokenizer and processor
+    tokenizer = AutoTokenizer.from_pretrained(path_or_repo)
+    processor = Qwen3VLProcessor(tokenizer=tokenizer)
+
+    return Qwen3VLBundledModel(vision_model=vision_model, llm_model=llm_model), processor
+
+def apply_chat_template_qwen3_vl(messages: Sequence[ChatMessage], num_images: int = 0, num_audios: int = 0, tools: Optional[str] = None, enable_thinking: bool = False) -> str:
+    """Apply chat template: serialize messages with content as a list of typed items."""
+    messages_dict = []
+    for msg in messages:
+        content_items = [{"type": "text", "text": msg.content}]
+        messages_dict.append({"role": msg.role, "content": content_items})
+    return json.dumps(messages_dict)
+
+
+def stream_generate_qwen3_vl(
+    model: Qwen3VLBundledModel,
+    processor: Qwen3VLProcessor,
+    prompt: str,
+    image: Union[str, List[str]] = None,
+    audio: Union[str, List[str]] = None,
+    max_tokens: int = 512,
+    **kwargs,
+
+) -> Generator[Any, None, None]:
+    """Stream generation yielding .generate.GenerationResult-compatible chunks."""
+    messages = json.loads(prompt)
+    if image is not None:
+        image_list = image if isinstance(image, list) else [image]
+        pil_images = []
+        for p in image_list:
+            try:
+                pil_images.append(Image.open(p))
+            except Exception:
+                continue
+        contents = [{"type": "image", "image": img} for img in pil_images]
+        if messages:
+            if "content" not in messages[-1] or not isinstance(messages[-1]["content"], list):
+                messages[-1]["content"] = []
+            messages[-1]["content"].extend(contents)
+
+    raw_text, processed_images = processor.messages_to_text(
+        messages, add_generation_prompt=True)
+
+    inputs = processor.text_to_input_ids(
+        raw_text, images=processed_images, return_tensors="mlx")
+
+    input_ids = inputs["input_ids"]
+    pixel_values = inputs.get("pixel_values")
+    image_grid_thw = inputs.get("image_grid_thw")
+
+    inputs_embeds, deepstack_visual_embeds, visual_pos_masks, cos, sin, rope_deltas = handle_multimodal_embeds(
+        model.vision_model, model.llm_model, input_ids, pixel_values, image_grid_thw
+    )
+
+    prompt_cache = make_prompt_cache(model.llm_model, max_kv_size=4096)
+    tokenizer = processor.tokenizer
+
+    # Rough prompt TPS estimation based on input size
+    prompt_start = time.perf_counter()
+    prompt_tps = input_ids.size / max(1e-6, (time.perf_counter() - prompt_start))
+
+    gen_count = 0
+    tic = time.perf_counter()
+
+    for token, logprobs in nexa_generate_step(
+        model=model.llm_model,
+        prompt=None,
+        input_embeddings=inputs_embeds,
+        max_tokens=max_tokens,
+        max_kv_size=4096,
+        prompt_cache=prompt_cache,
+        visual_pos_masks=visual_pos_masks,
+        deepstack_visual_embeds=deepstack_visual_embeds,
+        cos=cos,
+        sin=sin,
+        rope_deltas=rope_deltas,
+    ):
+        if token == tokenizer.eos_token_id:
+            break
+
+        text_piece = tokenizer.decode([token])
+        gen_count += 1
+
+        yield GenerationResult(
+            text=text_piece,
+            token=token,
+            logprobs=logprobs,
+            prompt_tokens=int(input_ids.size),
+            generation_tokens=gen_count,
+            prompt_tps=float(prompt_tps),
+            generation_tps=float(
+                gen_count / max(1e-6, (time.perf_counter() - tic))),
+            peak_memory=float(mx.get_peak_memory() / 1e9),
+        )
+
+def quant_predicate(path: str, mod: nn.Module) -> bool:
+    """Quantization predicate to exclude certain layers from quantization."""
+    if path.endswith("lm_head") or "norm" in path.lower() or "embed" in path.lower():
+        return False
+    return isinstance(mod, (nn.Linear, nn.Embedding))

nexaai/mlx_backend/vlm/interface.py

@@ -25,6 +25,8 @@ from profiling import ProfilingMixin, ProfilingData, StopReason
 
 # Import from the actual mlx_vlm structure
 from .generate import generate, stream_generate, load
+from .generate_qwen3_vl import apply_chat_template_qwen3_vl, stream_generate_qwen3_vl, load_qwen3_vl
+
 from .modeling.prompt_utils import apply_chat_template
 
 # --------------------------------------------------------------------------------------
@@ -54,6 +56,7 @@ class VLM(ProfilingMixin):
 
     def __init__(
         self,
+        model_name: Optional[str],
         model_path: Path,
         mmproj_path: Path,
         context_length: int,
@@ -67,11 +70,13 @@ class VLM(ProfilingMixin):
             model_path = os.path.dirname(model_path)
 
         self.model_path = model_path
+        self.model_name = model_name
         self.mmproj_path = mmproj_path
         self.context_length = context_length
         self.device = device
 
-        self.model, self.processor = load(str(model_path))
+        load_impl = load_qwen3_vl if model_name == "qwen3vl" else load
+        self.model, self.processor = load_impl(str(model_path))
 
         # Init deafutl sampler config with defualt.
         self.sampler_config = SamplerConfig()
@@ -228,9 +233,10 @@ class VLM(ProfilingMixin):
         text = ""
         last_result = None
         first_token = True
+        stream_generate_impl = stream_generate_qwen3_vl if self.model_name == "qwen3vl" else stream_generate
 
         try:
-            for result in stream_generate(
+            for result in stream_generate_impl(
                 self.model,
                 self.processor,
                 prompt,
@@ -350,6 +356,9 @@ class VLM(ProfilingMixin):
 
     def apply_chat_template_with_media(self, messages: Sequence[ChatMessage], num_images: int = 0, num_audios: int = 0, tools: Optional[str] = None, enable_thinking: bool = True) -> str:
         """Apply chat template to messages with proper image/audio token insertion and optional tools support."""
+        if self.model_name == "qwen3vl":
+            return apply_chat_template_qwen3_vl(messages, num_images=num_images, num_audios=num_audios, tools=tools, enable_thinking=enable_thinking)
+        
         # Convert ChatMessage objects to dictionaries for the processor
         messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
 

nexaai/mlx_backend/vlm/main.py

@@ -3,6 +3,7 @@ from ml import GenerationConfig, SamplerConfig, ChatMessage
 import re
 import os
 import codecs
+import argparse
 
 def parse_media_from_input(user_input):
     """Parse quoted media files from user input and return prompt and media paths"""
@@ -39,12 +40,177 @@ def parse_media_from_input(user_input):
     
     return prompt, image_paths if image_paths else None, audio_paths if audio_paths else None
 
-def test_vlm_generate_stream(model_path):
+def parse_arguments():
+    """Parse command line arguments for the VLM main function."""
+    parser = argparse.ArgumentParser(
+        description="Interactive VLM (Vision-Language Model) conversation interface."
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="mlx-community/gemma-3-4b-it-8bit",
+        help="The path to the local model directory or Hugging Face repo."
+    )
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        default="",
+        help="Specific model name/type (e.g., 'qwen3vl', 'gemma3'). If empty, auto-detect from model_path."
+    )
+    parser.add_argument(
+        "--context_length",
+        type=int,
+        default=2048,
+        help="Context length for the model (default: 2048)."
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.7,
+        help="Sampling temperature (default: 0.7)."
+    )
+    parser.add_argument(
+        "--top_p",
+        type=float,
+        default=0.9,
+        help="Top-p sampling parameter (default: 0.9)."
+    )
+    parser.add_argument(
+        "--max_tokens",
+        type=int,
+        default=512,
+        help="Maximum tokens to generate (default: 512)."
+    )
+    return parser.parse_args()
+
+def main():
+    """Main function for interactive VLM conversation."""
+    args = parse_arguments()
+    
+    # Auto-detect model name if not provided
+    model_name = args.model_name
+    if not model_name:
+        if "qwen" in args.model_path.lower():
+            model_name = "qwen3vl"
+        elif "gemma" in args.model_path.lower():
+            model_name = "gemma3"
+        else:
+            model_name = ""
+    
+    # Load the VLM instance
+    vlm = VLM(
+        model_name=model_name,
+        model_path=args.model_path,
+        mmproj_path=None,  # Not needed for this model
+        context_length=args.context_length,
+        device=None
+    )
+
+    # Configure sampler
+    sampler_config = SamplerConfig(
+        temperature=args.temperature,
+        top_p=args.top_p
+    )
+    vlm.set_sampler(sampler_config)
+
+    # Chat history using ChatMessage objects
+    chat = []
+
+    print("VLM Multi-round conversation started. Type 'quit' or 'exit' to end.")
+    print("Include images/audios in quotes, e.g.: 'describe \"image1.jpg\" \"image2.png\"'")
+    print("You can also use single quotes: 'describe '/path/to/image.jpg''")
+    print("=" * 50)
+
+    def on_token(text_chunk):
+        """Token callback for streaming"""
+        print(text_chunk, end="", flush=True)
+        return True
+
+    while True:
+        # Get user input
+        user_input = input("\nUser: ").strip()
+
+        # Check for exit commands
+        if user_input.lower() in ["quit", "exit", "q"]:
+            print("Goodbye!")
+            break
+
+        if not user_input:
+            continue
+
+        # Parse media files and prompt from user input
+        prompt_text, image_paths, audio_paths = parse_media_from_input(user_input)
+        
+        # If no text prompt after parsing, use the original input
+        if not prompt_text.strip():
+            prompt_text = user_input
+            image_paths = None
+            audio_paths = None
+
+        # Add user message to chat history using ChatMessage
+        chat.append(ChatMessage(role="user", content=prompt_text))
+
+        # Calculate number of images and audios for chat template
+        num_images = len(image_paths) if image_paths else 0
+        num_audios = len(audio_paths) if audio_paths else 0
+
+        # Apply chat template with image/audio token insertion
+        try:
+            formatted_prompt = vlm.apply_chat_template_with_media(chat, num_images=num_images, num_audios=num_audios)
+        except (NotImplementedError, AttributeError):
+            # Fallback to manual formatting if chat template is not implemented
+            formatted_prompt = ""
+            for msg in chat:
+                formatted_prompt += f"{msg.role}: {msg.content}\n"
+            formatted_prompt += "Assistant: "
+
+        # Generation config with media paths
+        generation_config = GenerationConfig(
+            max_tokens=args.max_tokens,
+            sampler_config=sampler_config,
+            image_paths=image_paths,
+            audio_paths=audio_paths
+        )
+
+        # Generate response
+        print("Assistant: ", end="", flush=True)
+        
+        try:
+            # Use streaming generation with callback
+            response_text = ""
+            
+            def token_callback(text_chunk):
+                nonlocal response_text
+                print(text_chunk, end="", flush=True)
+                response_text += text_chunk
+                return True
+            
+            # Use generate_stream method for streaming generation
+            response = vlm.generate_stream(
+                prompt=formatted_prompt,
+                config=generation_config,
+                on_token=token_callback
+            )
+            
+            print()  # New line after streaming
+            
+            # Add assistant response to chat history using ChatMessage
+            chat.append(ChatMessage(role="assistant", content=response_text))
+            
+        except Exception as e:
+            print(f"Error generating response: {e}")
+            print()
+
+    # Clean up
+    vlm.destroy()
+
+def test_vlm_generate_stream(model_path, model_name):
     # Specify the checkpoint
     context_length = 2048
 
     # Load the corresponding model and VLM instance
     vlm = VLM(
+        model_name=model_name,
         model_path=model_path,
         mmproj_path=None,  # Not needed for this model
         context_length=context_length,
@@ -88,9 +254,6 @@ def test_vlm_generate_stream(model_path):
         # Parse media files and prompt from user input
         prompt_text, image_paths, audio_paths = parse_media_from_input(user_input)
         
-        print(f"image_paths: {image_paths}")
-        print(f"audio_paths: {audio_paths}")
-        
         # If no text prompt after parsing, use the original input
         if not prompt_text.strip():
             prompt_text = user_input
@@ -150,8 +313,4 @@ def test_vlm_generate_stream(model_path):
     vlm.destroy()
 
 if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default="mlx-community/gemma-3-4b-it-8bit")
-    args = parser.parse_args()
-    test_vlm_generate_stream(args.model_path)
+    main()

nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/base.py

@@ -0,0 +1,117 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+from mlx.utils import tree_map
+
+from .cache import QuantizedKVCache
+
+
+@dataclass
+class BaseModelArgs:
+    @classmethod
+    def from_dict(cls, params):
+        return cls(**{k: v for k, v in params.items() if k in inspect.signature(cls).parameters})
+
+
+def create_causal_mask(
+    N: int,
+    offset: int = 0,
+    window_size: Optional[int] = None,
+    lengths: Optional[mx.array] = None,
+):
+    rinds = mx.arange(offset + N)
+    linds = mx.arange(offset, offset + N) if offset else rinds
+    linds = linds[:, None]
+    rinds = rinds[None]
+    mask = linds >= rinds
+    if window_size is not None:
+        mask = mask & (linds <= rinds + window_size)
+    if lengths is not None:
+        lengths = lengths[:, None, None, None]
+        mask = mask & (rinds < lengths)
+    return mask
+
+
+def create_attention_mask(h: mx.array, cache: Optional[Any] = None, return_array: bool = False):
+    T = h.shape[1]
+    if T > 1:
+        offset = 0
+        window_size = None
+        if cache is not None and cache[0] is not None:
+            c = cache[0]
+            offset = c.offset
+            if hasattr(c, "max_size"):
+                window_size = c.max_size
+                offset = min(window_size, offset)
+                return_array = return_array or offset + T > window_size
+        if return_array:
+            return create_causal_mask(T, offset, window_size=window_size)
+        else:
+            return "causal"
+    else:
+        mask = None
+    return mask
+
+
+def quantized_scaled_dot_product_attention(
+    queries: mx.array,
+    q_keys: tuple[mx.array, mx.array, mx.array],
+    q_values: tuple[mx.array, mx.array, mx.array],
+    scale: float,
+    mask: Optional[mx.array],
+    group_size: int = 64,
+    bits: int = 8,
+) -> mx.array:
+    B, n_q_heads, L, D = queries.shape
+    n_kv_heads = q_keys[0].shape[-3]
+    n_repeats = n_q_heads // n_kv_heads
+
+    queries *= scale
+
+    if n_repeats > 1:
+        queries = mx.reshape(queries, (B, n_kv_heads, n_repeats, L, D))
+        q_keys = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_keys)
+        q_values = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_values)
+
+    scores = mx.quantized_matmul(queries, *q_keys, transpose=True, group_size=group_size, bits=bits)
+    if mask is not None:
+        if isinstance(mask, str):
+            qL, kL = scores.shape[-2:]
+            q_indices = mx.arange(kL - qL, kL)
+            k_indices = mx.arange(kL)
+            mask = q_indices[:, None] >= k_indices[None]
+        if mask.dtype == mx.bool_:
+            scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
+        else:
+            scores += mask
+    scores = mx.softmax(scores, axis=-1, precise=True)
+    out = mx.quantized_matmul(scores, *q_values, transpose=False, group_size=group_size, bits=bits)
+
+    if n_repeats > 1:
+        out = mx.reshape(out, (B, n_q_heads, L, D))
+
+    return out
+
+
+def scaled_dot_product_attention(
+    queries,
+    keys,
+    values,
+    cache,
+    scale: float,
+    mask: Optional[mx.array],
+) -> mx.array:
+    if isinstance(cache, QuantizedKVCache):
+        return quantized_scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            scale=scale,
+            mask=mask,
+            group_size=cache.group_size,
+            bits=cache.bits,
+        )
+    else:
+        return mx.fast.scaled_dot_product_attention(queries, keys, values, scale=scale, mask=mask)