nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/vlm/main.py

@@ -0,0 +1,365 @@
+from .interface import VLM
+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"""
+    # Find all quoted strings (both single and double quotes)
+    quoted_pattern = r'["\']([^"\']*)["\']'
+    quoted_matches = re.findall(quoted_pattern, user_input)
+    
+    # Remove quoted strings from the input to get the actual prompt
+    prompt = re.sub(quoted_pattern, '', user_input).strip()
+    
+    # Separate image and audio files based on extensions
+    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'}
+    audio_extensions = {'.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a'}
+    
+    image_paths = []
+    audio_paths = []
+    
+    for quoted_file in quoted_matches:
+        if quoted_file:  # Skip empty quotes
+            # Expand user path if it starts with ~
+            if quoted_file.startswith('~'):
+                quoted_file = os.path.expanduser(quoted_file)
+            
+            # Check if file exists
+            if not os.path.exists(quoted_file):
+                print(f"Warning: File '{quoted_file}' not found")
+                continue
+                
+            file_ext = os.path.splitext(quoted_file.lower())[1]
+            if file_ext in image_extensions:
+                image_paths.append(quoted_file)
+            elif file_ext in audio_extensions:
+                audio_paths.append(quoted_file)
+    
+    return prompt, image_paths if image_paths else None, audio_paths if audio_paths else None
+
+def detect_model_name_and_repo(model_path):
+    """Detect model name and corresponding HuggingFace repo based on model path or name"""
+    model_path_lower = model_path.lower()
+    
+    # Handle HuggingFace repo format
+    if "/" in model_path:
+        repo_name = model_path.split("/")[-1] if model_path.endswith("/") else model_path.split("/")[-1]
+        repo_name_lower = repo_name.lower()
+    else:
+        repo_name_lower = model_path_lower
+    
+    # Model name mapping based on the provided examples
+    model_mappings = {
+        # 4B models
+        "qwen3vl-4b-4bit-mlx": ("qwen3vl-4b", "NexaAI/qwen3vl-4B-4bit-mlx"),
+        "qwen3vl-4b-fp16-mlx": ("qwen3vl-4b", "NexaAI/qwen3vl-4B-fp16-mlx"),
+        "qwen3vl-4b-thinking-4bit-mlx": ("qwen3vl-4b-thinking", "NexaAI/qwen3vl-4B-thinking-4bit-mlx"),
+        "qwen3vl-4b-thinking-fp16-mlx": ("qwen3vl-4b-thinking", "NexaAI/qwen3vl-4B-thinking-fp16-mlx"),
+        
+        # 8B models
+        "qwen3vl-8b-4bit-mlx": ("qwen3vl-8b", "NexaAI/qwen3vl-8B-4bit-mlx"),
+        "qwen3vl-8b-fp16-mlx": ("qwen3vl-8b", "NexaAI/qwen3vl-8B-fp16-mlx"),
+        "qwen3vl-8b-thinking-4bit-mlx": ("qwen3vl-8b-thinking", "NexaAI/qwen3vl-8B-thinking-4bit-mlx"),
+        "qwen3vl-8b-thinking-fp16-mlx": ("qwen3vl-8b-thinking", "NexaAI/qwen3vl-8B-thinking-fp16-mlx"),
+    }
+    
+    # Check exact matches first
+    for key, (model_name, repo) in model_mappings.items():
+        if key in repo_name_lower:
+            return model_name, repo if "/" not in model_path else model_path
+    
+    # Fallback detection based on patterns
+    if "qwen3vl" in repo_name_lower:
+        if "8b" in repo_name_lower:
+            if "thinking" in repo_name_lower:
+                return "qwen3vl-8b-thinking", model_path
+            else:
+                return "qwen3vl-8b", model_path
+        elif "4b" in repo_name_lower:
+            if "thinking" in repo_name_lower:
+                return "qwen3vl-4b-thinking", model_path
+            else:
+                return "qwen3vl-4b", model_path
+        else:
+            # Default to 4B if size not specified
+            return "qwen3vl-4b", model_path
+    elif "gemma" in repo_name_lower:
+        return "gemma3", model_path
+    
+    return "", 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="NexaAI/qwen3vl-4B-4bit-mlx",
+        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-4b', 'qwen3vl-4b-thinking', 'qwen3vl-8b', 'qwen3vl-8b-thinking'). 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
+    model_path = args.model_path
+    
+    if not model_name:
+        model_name, model_path = detect_model_name_and_repo(args.model_path)
+        print(f"Auto-detected model: {model_name} from path: {model_path}")
+    
+    # Load the VLM instance
+    vlm = VLM(
+        model_name=model_name,
+        model_path=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,
+        device=None
+    )
+
+    # Configure sampler
+    sampler_config = SamplerConfig(
+        temperature=0.7,
+        top_p=0.9
+    )
+    vlm.set_sampler(sampler_config)
+
+    # Chat history using ChatMessage objects (following ml.py API)
+    chat = []
+
+    print("Multi-round VLM 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, user_data):
+        """Token callback for streaming"""
+        print(text_chunk, end="", flush=True)
+        if user_data is not None:
+            user_data["response"] += text_chunk
+        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 (following ml.py API)
+        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=512,
+            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 - single method handles all cases
+            user_data = {"response": ""}
+            
+            # Always use the unified generate_stream method
+            response = vlm.generate_stream(
+                prompt=formatted_prompt,
+                config=generation_config,
+                on_token=on_token,
+                user_data=user_data
+            )
+            
+            print()  # New line after streaming
+            
+            # Add assistant response to chat history using ChatMessage
+            chat.append(ChatMessage(role="assistant", content=user_data["response"]))
+            
+        except Exception as e:
+            print(f"Error generating response: {e}")
+            print()
+
+    # Clean up
+    vlm.destroy()
+
+if __name__ == "__main__":
+    main()

nexaai/mlx_backend/vlm/modeling/convert.py

@@ -0,0 +1,68 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+
+from .utils import MODEL_CONVERSION_DTYPES, convert
+
+
+def configure_parser() -> argparse.ArgumentParser:
+    """
+    Configures and returns the argument parser for the script.
+
+    Returns:
+        argparse.ArgumentParser: Configured argument parser.
+    """
+    parser = argparse.ArgumentParser(
+        description="Convert Hugging Face model to MLX format"
+    )
+
+    parser.add_argument("--hf-path", type=str, help="Path to the Hugging Face model.")
+    parser.add_argument(
+        "--mlx-path", type=str, default="mlx_model", help="Path to save the MLX model."
+    )
+    parser.add_argument(
+        "-q", "--quantize", help="Generate a quantized model.", action="store_true"
+    )
+    parser.add_argument(
+        "--q-group-size", help="Group size for quantization.", type=int, default=64
+    )
+    parser.add_argument(
+        "--q-bits", help="Bits per weight for quantization.", type=int, default=4
+    )
+    parser.add_argument(
+        "--dtype",
+        help="Type to save the parameter. Defaults to config.json's `torch_dtype` or the current model weights dtype",
+        type=str,
+        choices=MODEL_CONVERSION_DTYPES,
+        default=None,
+    )
+    parser.add_argument(
+        "--upload-repo",
+        help="The Hugging Face repo to upload the model to.",
+        type=str,
+        default=None,
+    )
+    parser.add_argument(
+        "-d",
+        "--dequantize",
+        help="Dequantize a quantized model.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--skip-vision",
+        help="Skip vision module quantization.",
+        action="store_true",
+        default=False,
+    )
+    return parser
+
+
+def main():
+    parser = configure_parser()
+    args = parser.parse_args()
+    convert(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

nexaai/mlx_backend/vlm/modeling/models/aya_vision/__init__.py

@@ -0,0 +1,8 @@
+from .aya_vision import (
+    LanguageModel,
+    Model,
+    ModelConfig,
+    TextConfig,
+    VisionConfig,
+    VisionModel,
+)

nexaai/mlx_backend/vlm/modeling/models/aya_vision/aya_vision.py

@@ -0,0 +1,193 @@
+import glob
+import inspect
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mlx_lm.utils import snapshot_download
+
+from .language import LanguageModel, TextConfig
+from .vision import VisionConfig, VisionModel
+
+
+@dataclass
+class ModelConfig:
+    text_config: TextConfig
+    vision_config: VisionConfig
+    model_type: str
+    image_token_index: int = 255036
+    max_splits_per_img: int = 12
+    downsample_factor: int = 2
+    alignment_intermediate_size: int = 28672
+    adapter_layer_norm_eps: float = 1e-06
+    vision_feature_layer: int = -1
+    vision_feature_select_strategy: str = "full"
+    eos_token_id: Optional[List[int]] = None
+
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+class AyaVisionMultiModalProjector(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.downsample_factor = config.downsample_factor
+        self.alignment_intermediate_size = getattr(
+            config, "alignment_intermediate_size", config.text_config.hidden_size
+        )
+        self.layernorm = nn.LayerNorm(
+            config.vision_config.hidden_size * (config.downsample_factor**2),
+            eps=config.adapter_layer_norm_eps,
+        )
+
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size * (config.downsample_factor**2),
+            self.alignment_intermediate_size,
+            bias=True,
+        )
+
+        self.act = nn.SiLU()  # SwiGLU uses SiLU activation
+
+        # For SwiGLU, project down to half size since we split intermediate dim
+        self.linear_2 = nn.Linear(
+            self.alignment_intermediate_size // 2,
+            config.text_config.hidden_size,
+            bias=True,
+        )
+
+    def __call__(self, image_features):
+        image_features = self.pixel_shuffle(image_features)
+        image_features = self.layernorm(image_features)
+        hidden_states = self.linear_1(image_features)
+
+        # Split along last dimension and apply SwiGLU
+        x, gate = mx.split(hidden_states, 2, axis=-1)
+        hidden_states = self.act(gate) * x
+
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+    def pixel_shuffle(self, image_features):  # B, S, D
+        batch_size, seq_length, feature_dim = image_features.shape
+        height = width = int(seq_length**0.5)
+        image_features = image_features.reshape(
+            image_features.shape[0], width, height, -1
+        )
+        channels = image_features.shape[-1]
+        image_features = image_features.reshape(
+            batch_size,
+            width,
+            int(height / self.downsample_factor),
+            int(channels * self.downsample_factor),
+        )
+        image_features = image_features.transpose(0, 2, 1, 3)
+        image_features = image_features.reshape(
+            batch_size,
+            int(height / self.downsample_factor),
+            int(width / self.downsample_factor),
+            -1,
+        )
+        image_features = image_features.transpose(0, 2, 1, 3)
+        return image_features
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.vision_tower = VisionModel(config.vision_config)
+        self.language_model = LanguageModel(config.text_config)
+        self.multi_modal_projector = AyaVisionMultiModalProjector(config)
+        self.vision_feature_layer = config.vision_feature_layer
+        self.vision_feature_select_strategy = config.vision_feature_select_strategy
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        if pixel_values is None:
+            return self.language_model.model.embed_tokens(input_ids)
+
+        # Get the input embeddings from the language model
+        inputs_embeds = self.language_model.model.embed_tokens(input_ids)
+
+        spatial_shapes = kwargs.get("spatial_shapes", None)
+        # Get the ouptut hidden states from the vision model
+        *_, hidden_states = self.vision_tower(
+            pixel_values.transpose(0, 2, 3, 1),
+            spatial_shapes=spatial_shapes,
+            output_hidden_states=True,
+        )
+
+        # Select the hidden states from the desired layer
+        selected_image_feature = hidden_states[self.vision_feature_layer]
+
+        if self.vision_feature_select_strategy == "default":
+            selected_image_feature = selected_image_feature[:, 1:]
+        elif self.vision_feature_select_strategy == "full":
+            selected_image_feature = selected_image_feature
+        else:
+            raise ValueError(
+                "Unexpected feature selection strategy: "
+                f"{self.vision_feature_select_strategy}"
+            )
+
+        # Pass image features through the multi-modal projector
+        image_features = self.multi_modal_projector(selected_image_feature)
+
+        # Insert special image tokens in the input_ids
+        final_inputs_embeds = self._merge_input_ids_with_image_features(
+            image_features, inputs_embeds, input_ids
+        )
+        return final_inputs_embeds
+
+    def _merge_input_ids_with_image_features(
+        self, image_features, inputs_embeds, input_ids
+    ):
+        image_token_index = self.config.image_token_index
+
+        # Positions of <image> tokens in input_ids, assuming batch size is 1
+        image_positions = np.where(input_ids[0] == image_token_index)[0].tolist()
+        num_images, _, _, vision_hidden_size = image_features.shape
+
+        reshaped_image_hidden_states = image_features.reshape(-1, vision_hidden_size)
+
+        # cast to the dtype of the input_embeds to support quantized models
+        reshaped_image_hidden_states = reshaped_image_hidden_states.astype(
+            inputs_embeds.dtype
+        )
+        inputs_embeds[:, image_positions, :] = reshaped_image_hidden_states
+        return inputs_embeds
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array,
+        mask: mx.array,
+        cache=None,
+        **kwargs,
+    ):
+
+        input_embddings = self.get_input_embeddings(input_ids, pixel_values, **kwargs)
+        logits = self.language_model(
+            input_ids, cache=cache, inputs_embeds=input_embddings
+        )
+        return logits