nexaai 1.0.19rc16__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc17__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.19-rc16" 
+__version__ = "1.0.19-rc17" 

nexaai/mlx_backend/vlm/generate_qwen3_vl.py

@@ -43,17 +43,156 @@ def _ensure_list(x: Union[str, List[str], None]) -> Optional[List[str]]:
     return x if isinstance(x, list) else [x]
 
 
+def get_model_configs(model_name: str):
+    """Get model configurations based on model name"""
+    
+    # 4B model configs (default)
+    if model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking"]:
+        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"},
+        )
+        
+    # 8B model configs
+    elif model_name in ["qwen3vl-8b", "qwen3vl-8b-thinking"]:
+        vision_config = VisionConfig(
+            hidden_size=1152,
+            intermediate_size=4304,
+            num_heads=16,
+            num_hidden_layers=27,
+            patch_size=16,
+            temporal_patch_size=2,
+            in_channels=3,
+            hidden_act="gelu",
+            spatial_merge_size=2,
+            out_hidden_size=4096,
+            num_position_embeddings=2304,
+            deepstack_visual_indexes=[8, 16, 24],
+        )
+
+        text_config = TextConfig(
+            model_type="qwen3vl",
+            hidden_size=4096,
+            num_hidden_layers=36,
+            intermediate_size=12288,
+            num_attention_heads=32,
+            num_key_value_heads=8,
+            rms_norm_eps=1e-6,
+            vocab_size=151936,
+            max_position_embeddings=262144,
+            rope_theta=5000000,
+            head_dim=128,
+            tie_word_embeddings=False,
+            attention_bias=False,
+            attention_dropout=0.0,
+            rope_scaling={"mrope_section": [24, 20, 20], "rope_type": "default", "mrope_interleaved": True},
+        )
+    else:
+        # Fallback to 4B config
+        return get_model_configs("qwen3vl-4b")
+        
+    return vision_config, text_config
+
+def get_weight_filenames(model_name: str, model_path: Path):
+    """Get appropriate weight filenames based on model name and available files"""
+    
+    # Determine model size and type based on the actual file structure
+    if "4b" in model_name:
+        size_prefix = "4b"
+    elif "8b" in model_name:
+        size_prefix = "8b"
+    else:
+        size_prefix = "4b"
+    
+    # Determine model type
+    if "thinking" in model_name:
+        model_type = f"{size_prefix}_thinking"
+    else:
+        model_type = f"{size_prefix}_instruct"
+    
+    # Try different weight file patterns matching the actual file structure
+    llm_patterns = [
+        # New naming convention matching actual files
+        f"qwen3vl-llm-{model_type}-q4_0.safetensors",
+        f"qwen3vl-llm-{model_type}-q8_0.safetensors", 
+        f"qwen3vl-llm-{model_type}-f16.safetensors",
+        # Legacy naming convention
+        f"qwen3vl-llm-{size_prefix.upper()}-q4_0.safetensors",
+        f"qwen3vl-llm-{size_prefix.upper()}-q8_0.safetensors",
+        f"qwen3vl-llm-{size_prefix.upper()}-f16.safetensors",
+        f"qwen3vl-llm-{size_prefix.upper()}-f32.safetensors",
+    ]
+    
+    vision_patterns = [
+        f"qwen3vl-vision-{model_type}-f16.safetensors",
+        f"qwen3vl-vision-{size_prefix.upper()}-f16.safetensors",
+    ]
+    
+    # Find LLM weights
+    llm_weights_path = None
+    quantization_bits = None
+    
+    for pattern in llm_patterns:
+        candidate_path = model_path / pattern
+        if candidate_path.exists():
+            llm_weights_path = candidate_path
+            if "q4_0" in pattern:
+                quantization_bits = 4
+            elif "q8_0" in pattern:
+                quantization_bits = 8
+            else:
+                quantization_bits = 16
+            break
+    
+    # Find vision weights
+    vision_weights_path = None
+    for pattern in vision_patterns:
+        candidate_path = model_path / pattern
+        if candidate_path.exists():
+            vision_weights_path = candidate_path
+            break
+    
+    return llm_weights_path, vision_weights_path, quantization_bits
+
+# Update the load_qwen3_vl function signature and implementation:
 def load_qwen3_vl(
     path_or_repo: str,
     adapter_path: Optional[str] = None,
     lazy: bool = False,
     revision: Optional[str] = None,
+    model_name: Optional[str] = None,
     **kwargs,
 ) -> Tuple[Qwen3VLBundledModel, Qwen3VLProcessor]:
-    """Load Qwen3-VL quantized models and processor.
-
-    Parameters are aligned with .generate.load for compatibility.
-    """
+    """Load Qwen3-VL quantized models and processor with support for different model sizes."""
     
     model_path = Path(path_or_repo)
     if not model_path.exists():
@@ -67,70 +206,22 @@ def load_qwen3_vl(
             if not model_path.exists():
                 model_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"},
-    )
+    # Get model configurations based on model name
+    if model_name:
+        vision_config, text_config = get_model_configs(model_name)
+    else:
+        # Default to 4B config
+        vision_config, text_config = get_model_configs("qwen3vl-4b")
 
     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"
+    # Get appropriate weight filenames
+    llm_weights_path, vision_weights_path, quantization_bits = get_weight_filenames(
+        model_name or "qwen3vl-4b", model_path
+    )
 
-    if not vision_weights_path.exists() or not llm_weights_path.exists():
+    if not vision_weights_path or not llm_weights_path:
         raise FileNotFoundError(
             f"Missing safetensors. Vision: {vision_weights_path}, LLM: {llm_weights_path}"
         )
@@ -146,8 +237,14 @@ def load_qwen3_vl(
     
     llm_model.load_weights(str(llm_weights_path), strict=True)
 
-    # Tokenizer and processor
-    tokenizer = AutoTokenizer.from_pretrained(path_or_repo)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(str(model_path))
+    except Exception:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(path_or_repo)
+        except Exception:
+            raise Exception("Failed to load tokenizer from the same path where model weights are loaded and original path_or_repo.")
+    
     processor = Qwen3VLProcessor(tokenizer=tokenizer)
 
     return Qwen3VLBundledModel(vision_model=vision_model, llm_model=llm_model), processor

nexaai/mlx_backend/vlm/interface.py

@@ -81,12 +81,13 @@ class VLM(ProfilingMixin):
 
         if model_name == "qwen3vl-moe":
             load_impl = load_qwen3_vl_moe
-        elif model_name == "qwen3vl":
+        elif model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking"]:
             load_impl = load_qwen3_vl
         else:
             load_impl = load
 
-        self.model, self.processor = load_impl(str(model_path))
+        # Pass model_name to the loader for proper configuration
+        self.model, self.processor = load_impl(str(model_path), model_name=model_name)
 
         # Init deafutl sampler config with defualt.
         self.sampler_config = SamplerConfig()
@@ -94,6 +95,19 @@ class VLM(ProfilingMixin):
         # Track global character position for incremental processing
         self.global_n_past_chars = 0
 
+        # Add conversation state tracking to VLM class
+        if model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking"]:
+            # Import here to avoid circular imports
+            from .modeling.models.qwen3_vl.llm_common.cache import make_prompt_cache
+            import mlx.core as mx
+            
+            # Initialize conversation state
+            self.rope_deltas_total = mx.zeros((1, 1), dtype=mx.int32)
+            self.prompt_cache = make_prompt_cache(self.model.llm_model, max_kv_size=4096)
+        else:
+            self.rope_deltas_total = None
+            self.prompt_cache = None
+
     def destroy(self) -> None:
         """Destroy the model and free resources."""
         self.model = None
@@ -103,6 +117,14 @@ class VLM(ProfilingMixin):
         """Reset the model state."""
         self._reset_cache()
         self.global_n_past_chars = 0
+        
+        # Reset conversation state for qwen3vl models
+        if self.model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking"]:
+            import mlx.core as mx
+            from .modeling.models.qwen3_vl.llm_common.cache import make_prompt_cache
+            
+            self.rope_deltas_total = mx.zeros((1, 1), dtype=mx.int32)
+            self.prompt_cache = make_prompt_cache(self.model.llm_model, max_kv_size=4096)
 
     def _reset_cache(self) -> None:
         """Reset the KV cache."""
@@ -280,7 +302,7 @@ class VLM(ProfilingMixin):
         
         # Apply incremental processing only for non-qwen3vl models
         # qwen3vl requires complete JSON conversation structure
-        if self.model_name != "qwen3vl":
+        if self.model_name not in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking", "qwen3vl-moe"]:
             if self.global_n_past_chars < full_prompt_len:
                 incremental_prompt = prompt[self.global_n_past_chars:]
             else:
@@ -297,7 +319,7 @@ class VLM(ProfilingMixin):
 
         if self.model_name == "qwen3vl-moe":
             stream_generate_impl = stream_generate_qwen3_vl_moe
-        elif self.model_name == "qwen3vl":
+        elif self.model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking"]:
             stream_generate_impl = stream_generate_qwen3_vl
         else:
             stream_generate_impl = stream_generate
@@ -305,28 +327,59 @@ class VLM(ProfilingMixin):
         try:
             token_count = 0
             
-            for result in stream_generate_impl(
-                self.model,
-                self.processor,
-                incremental_prompt,  # Use incremental prompt instead of full prompt
-                image=image_list,
-                audio=audio_list,
-                **gen_kwargs,
-            ):
-                token_count += 1
-                
-                # 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
+            # Pass conversation state for qwen3vl models
+            if self.model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking"]:
+                for result in stream_generate_impl(
+                    self.model,
+                    self.processor,
+                    incremental_prompt,
+                    image=image_list,
+                    audio=audio_list,
+                    rope_deltas_total=self.rope_deltas_total,  # Pass conversation state
+                    prompt_cache=self.prompt_cache,  # Pass KV cache
+                    **gen_kwargs,
+                ):
+                    token_count += 1
+                    
+                    # 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
+
+                    # Update conversation state after each token
+                    # Note: rope_deltas_total is updated inside stream_generate_qwen3_vl
+
+            else:
+                for result in stream_generate_impl(
+                    self.model,
+                    self.processor,
+                    incremental_prompt,
+                    image=image_list,
+                    audio=audio_list,
+                    **gen_kwargs,
+                ):
+                    token_count += 1
+                    
+                    # 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
@@ -339,7 +392,7 @@ class VLM(ProfilingMixin):
                 self._update_generated_tokens(last_result.generation_tokens)
 
             # Update global character position (not needed for qwen3vl JSON processing)
-            if self.model_name != "qwen3vl":
+            if self.model_name not in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking", "qwen3vl-moe"]:
                 old_pos = self.global_n_past_chars
                 self.global_n_past_chars = full_prompt_len + len(text)
 
@@ -444,11 +497,10 @@ 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":
+        if self.model_name in ["qwen3vl", "qwen3vl-4b", "qwen3vl-4b-thinking", "qwen3vl-8b", "qwen3vl-8b-thinking"]:
             return apply_chat_template_qwen3_vl(messages, num_images=num_images, num_audios=num_audios, tools=tools, enable_thinking=enable_thinking)
         if self.model_name == "qwen3vl-moe":
             return apply_chat_template_qwen3_vl_moe(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

nexaai/mlx_backend/vlm/main.py

@@ -40,6 +40,57 @@ def parse_media_from_input(user_input):
     
     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(
@@ -48,14 +99,14 @@ def parse_arguments():
     parser.add_argument(
         "--model_path",
         type=str,
-        default="mlx-community/gemma-3-4b-it-8bit",
+        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', 'qwen3vl-moe', 'gemma3'). If empty, auto-detect from model_path."
+        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",
@@ -89,22 +140,16 @@ def main():
     
     # Auto-detect model name if not provided
     model_name = args.model_name
-
-    # TODO: avoid such hardcoded model name detection
+    model_path = args.model_path
+    
     if not model_name:
-        if "qwen3vl-30B" in args.model_path.lower():
-            model_name = "qwen3vl-moe"
-        elif "qwen3" in args.model_path.lower():
-            model_name = "qwen3vl"
-        elif "gemma" in args.model_path.lower():
-            model_name = "gemma3"
-        else:
-            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=args.model_path,
+        model_path=model_path,
         mmproj_path=None,  # Not needed for this model
         context_length=args.context_length,
         device=None