nexaai 1.0.16rc13__cp310-cp310-macosx_15_0_universal2.whl

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

@@ -0,0 +1,415 @@
+from __future__ import annotations
+
+import json
+import os
+import time
+from typing import Any, List, Optional, Sequence, Tuple, Union
+import mlx.core as mx
+import codecs
+from dataclasses import dataclass
+
+# Import configs and callback types from ml.py for API alignment
+from ml import (
+    VLM as BaseVLM,
+    SamplerConfig,
+    GenerationConfig,
+    ChatMessage,
+    EmbeddingConfig,
+    TokenCallback,
+    Path,
+    Tool,  # Add Path alias for type hints
+)
+
+# Import profiling module
+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
+
+# --------------------------------------------------------------------------------------
+# Updated GenerationResult to match the new structure
+# --------------------------------------------------------------------------------------
+
+@dataclass
+class GenerationResult:
+    text: str = ""
+    token: Optional[int] = None
+    logprobs: Optional[List[float]] = None
+    prompt_tokens: int = 0
+    generation_tokens: int = 0
+    total_tokens: int = 0
+    prompt_tps: float = 0.0
+    generation_tps: float = 0.0
+    peak_memory: float = 0.0
+# --------------------------------------------------------------------------------------
+# VLM (Vision-Language Model)
+# --------------------------------------------------------------------------------------
+
+class VLM(ProfilingMixin):
+    """
+    Vision-Language Models for mlx-vlm
+    API aligned with ml.py VLM abstract base class.
+    """
+
+    def __init__(
+        self,
+        model_name: Optional[str],
+        model_path: Path,
+        mmproj_path: Path,
+        context_length: int,
+        device: Optional[str] = None,
+    ) -> None:
+        # Initialize profiling mixin
+        ProfilingMixin.__init__(self)
+
+        # Check if model_path is a file, if so use its parent directory
+        if os.path.isfile(model_path):
+            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
+
+        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()
+
+    def destroy(self) -> None:
+        """Destroy the model and free resources."""
+        self.model = None
+        self.processor = None
+
+    def reset(self) -> None:
+        """Reset the model state."""
+        self._reset_cache()
+
+    def _reset_cache(self) -> None:
+        """Reset the KV cache."""
+        # If the model has a cache, reset it
+        if hasattr(self.model, "cache"):
+            self.model.cache = None
+
+    # Tokenization
+    def encode(self, text: str) -> List[int]:
+        """Encode text to token IDs."""
+        return self.processor.encode(text)
+
+    def decode(self, token_ids: Sequence[int]) -> str:
+        """Decode token IDs to text."""
+        return self.processor.decode(token_ids)
+
+    # Sampler
+    def set_sampler(self, config: SamplerConfig) -> None:
+        """Set sampler configuration."""
+        self.sampler_config = config
+
+    def reset_sampler(self) -> None:
+        """Reset sampler to default configuration."""
+        self.sampler_config = None
+
+    # Generation
+    def generate(
+        self,
+        prompt: str,
+        config: Optional[GenerationConfig] = None,
+    ) -> GenerationResult:
+        """Generate text from prompt."""
+        # Start profiling
+        self._start_profiling()
+
+        gen_kwargs = {}
+        if config is not None:
+            gen_kwargs = config.__dict__.copy()
+            # Remove image_paths and audio_paths from config as they'll be handled separately
+            gen_kwargs.pop('image_paths', None)
+            gen_kwargs.pop('audio_paths', None)
+        if self.sampler_config is not None:
+            gen_kwargs.update(self.sampler_config.__dict__)
+
+        # Get image and audio paths from config
+        image_paths = config.image_paths if config else None
+        audio_paths = config.audio_paths if config else None
+
+        # Convert paths to strings for generate function
+        image_list = [str(path) for path in image_paths] if image_paths else None
+        audio_list = [str(path) for path in audio_paths] if audio_paths else None
+
+        # End prompt processing, start decode
+        self._prompt_end()
+        self._decode_start()
+
+        try:
+            # Start timing for generation
+            generation_start_time = time.perf_counter()
+
+            text, stats = generate(
+                self.model,
+                self.processor,
+                prompt,
+                image=image_list,
+                audio=audio_list,
+                **gen_kwargs,
+            )
+
+            # End timing for generation
+            generation_end_time = time.perf_counter()
+
+            # Calculate average time per token and estimate TTFT
+            generated_tokens = stats.get("output_tokens", 0)
+            if generated_tokens > 0:
+                total_generation_time = generation_end_time - generation_start_time
+                avg_time_per_token = total_generation_time / generated_tokens
+                # TTFT = prompt processing time + first token generation time
+                # This provides a more accurate estimate than the previous approximation
+                estimated_ttft = (self._profiling_context.prompt_end_time - self._profiling_context.prompt_start_time) + avg_time_per_token
+                # Update the profiling context with estimated TTFT
+                self._profiling_context.first_token_time = self._profiling_context.prompt_start_time + estimated_ttft
+                self._profiling_context.ttft_recorded = True
+            else:
+                # If no tokens generated, use total generation time as TTFT
+                self._record_ttft()
+
+            # Update profiling data
+            prompt_tokens = stats.get("input_tokens", 0)
+            self._update_prompt_tokens(prompt_tokens)
+            self._update_generated_tokens(generated_tokens)
+            self._set_stop_reason(StopReason.ML_STOP_REASON_COMPLETED)
+            self._decode_end()
+            self._end_profiling()
+
+            return GenerationResult(
+                text=text,
+                prompt_tokens=prompt_tokens,
+                generation_tokens=generated_tokens,
+                total_tokens=stats.get("total_tokens", 0),
+                prompt_tps=stats.get("prompt_tps", 0.0),
+                generation_tps=stats.get("generation_tps", 0.0),
+                peak_memory=stats.get("peak_memory", 0.0),
+            )
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Generation error: {str(e)}")
+
+    def generate_stream(
+        self,
+        prompt: str,
+        config: Optional[GenerationConfig],
+        on_token: Optional[TokenCallback],
+    ) -> GenerationResult:
+        """Generate text with streaming callback. Unified method for both text and multimodal generation."""
+        # Start profiling
+        self._start_profiling()
+
+        gen_kwargs = {}
+        if config is not None:
+            gen_kwargs = config.__dict__.copy()
+            # Remove image_paths and audio_paths from config as they'll be handled separately
+            gen_kwargs.pop('image_paths', None)
+            gen_kwargs.pop('audio_paths', None)
+        if self.sampler_config is not None:
+            gen_kwargs.update(self.sampler_config.__dict__)
+
+        # Get image and audio paths from config
+        image_paths = config.image_paths if config else None
+        audio_paths = config.audio_paths if config else None
+
+        # Convert paths to strings for stream_generate function
+        image_list = [str(path) for path in image_paths] if image_paths else None
+        audio_list = [str(path) for path in audio_paths] if audio_paths else None
+
+        # End prompt processing, start decode
+        self._prompt_end()
+        self._decode_start()
+
+        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_impl(
+                self.model,
+                self.processor,
+                prompt,
+                image=image_list,
+                audio=audio_list,
+                **gen_kwargs,
+            ):
+                                # Record TTFT on first token
+                if first_token:
+                    self._record_ttft()
+                    first_token = False
+
+                # Call the token callback if provided
+                if on_token is not None:
+                    if not on_token(result.text):
+                        self._set_stop_reason(StopReason.ML_STOP_REASON_USER)
+                        break
+                text += result.text
+                last_result = result
+
+            # Set stop reason if not user stop
+            if self._profiling_context.stop_reason != StopReason.ML_STOP_REASON_USER:
+                self._set_stop_reason(StopReason.ML_STOP_REASON_EOS)
+
+            # Update profiling data
+            if last_result:
+                self._update_prompt_tokens(last_result.prompt_tokens)
+                self._update_generated_tokens(last_result.generation_tokens)
+
+            self._decode_end()
+            self._end_profiling()
+
+            return GenerationResult(
+                text=text,
+                token=last_result.token if last_result else None,
+                logprobs=last_result.logprobs if last_result else None,
+                prompt_tokens=last_result.prompt_tokens if last_result else 0,
+                generation_tokens=last_result.generation_tokens if last_result else 0,
+                total_tokens=(last_result.prompt_tokens + last_result.generation_tokens) if last_result else 0,
+                prompt_tps=last_result.prompt_tps if last_result else 0.0,
+                generation_tps=last_result.generation_tps if last_result else 0.0,
+                peak_memory=last_result.peak_memory if last_result else 0.0,
+            )
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Streaming generation error: {str(e)}")
+
+    # Legacy multimodal methods - kept for backward compatibility but delegate to unified method
+    def generate_multimodal(
+        self,
+        prompt: str,
+        image_paths: Optional[Sequence[Path]] = None,
+        audio_paths: Optional[Sequence[Path]] = None,
+        config: Optional[GenerationConfig] = None,
+    ) -> str:
+        """Generate text from prompt with multiple images and audio."""
+        # Create config with media paths if not provided
+        if config is None:
+            config = GenerationConfig()
+
+        # Update config with provided paths
+        if image_paths is not None:
+            config.image_paths = image_paths
+        if audio_paths is not None:
+            config.audio_paths = audio_paths
+
+        # Delegate to unified generate method and extract text
+        result = self.generate(prompt, config)
+        return result.text
+
+    def generate_stream_multimodal(
+        self,
+        prompt: str,
+        image_paths: Optional[Sequence[Path]] = None,
+        audio_paths: Optional[Sequence[Path]] = None,
+        config: Optional[GenerationConfig] = None,
+        on_token: Optional[TokenCallback] = None,
+    ) -> str:
+        """Generate text from prompt with multiple images and audio using streaming callback."""
+        # Create config with media paths if not provided
+        if config is None:
+            config = GenerationConfig()
+
+        # Update config with provided paths
+        if image_paths is not None:
+            config.image_paths = image_paths
+        if audio_paths is not None:
+            config.audio_paths = audio_paths
+
+        # Delegate to unified generate_stream method and extract text
+        result = self.generate_stream(prompt, config, on_token)
+        return result.text
+
+    def get_chat_template(self, template_name: str) -> str:
+        """Get chat template by name."""
+        # This is a stub; actual implementation depends on processor internals
+        if hasattr(self.processor, "get_chat_template"):
+            return self.processor.get_chat_template(template_name)
+        return ""
+
+    def apply_chat_template(self, messages: Sequence[ChatMessage], tools: Optional[str] = None, enable_thinking: bool = True) -> str:
+        """Apply chat template to messages with optional tools support."""
+        if hasattr(self.processor, "apply_chat_template"):
+            # Convert ChatMessage objects to dictionaries for the processor
+            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+            
+            parsed_tools = None
+            if tools is not None and tools.strip():
+                parsed_tools = json.loads(tools)
+            
+            result = apply_chat_template(self.processor, self.model.config, messages_dict, add_generation_prompt=True, enable_thinking=enable_thinking, tools=parsed_tools)
+            return result
+        # Fallback: join messages
+        return "\n".join([f"{m.role}: {m.content}" for m in messages])
+
+    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]
+
+        parsed_tools = None
+        if tools is not None and tools.strip():
+            parsed_tools = json.loads(tools)
+
+        # Use the same logic as generate.py
+        return apply_chat_template(
+            self.processor,
+            self.model.config,
+            messages_dict,
+            num_images=num_images,
+            num_audios=num_audios,
+            enable_thinking=enable_thinking,
+            tools=parsed_tools
+        )
+
+    # Embeddings
+    def embed(
+        self,
+        texts: Sequence[str],
+        config: Optional[EmbeddingConfig] = None,
+    ) -> List[List[float]]:
+        """Generate embeddings for texts with profiling."""
+        # Start profiling
+        self._start_profiling()
+
+        try:
+            # If processor/model supports embeddings, use it; otherwise, stub
+            if hasattr(self.model, "embed"):
+                embed_kwargs = config.__dict__ if config else {}
+
+                # End prompt processing, start decode
+                self._prompt_end()
+                self._decode_start()
+
+                result = self.model.embed(texts, **embed_kwargs)
+
+                # End timing and finalize profiling data
+                self._update_generated_tokens(0)  # No generation in embedding
+                self._set_stop_reason(StopReason.ML_STOP_REASON_COMPLETED)
+                self._decode_end()
+                self._end_profiling()
+
+                return result
+            else:
+                raise NotImplementedError("Embedding not supported for this model.")
+
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Error generating embeddings: {str(e)}")