From 013c99e51bbb8c4b83d88f3b150a1e53251a19d2 Mon Sep 17 00:00:00 2001 From: Manan Shah <52329525+Manan17@users.noreply.github.com> Date: Fri, 17 Apr 2026 03:06:03 -0500 Subject: [PATCH] qwen3.6 patches for multi-turn chat (#5083) * qwen3.6 patches for multi-turn chat * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../patches/mlx_vlm_qwen3_5/generate.py | 1653 +++++++++++++++++ .../models/patches/mlx_vlm_qwen3_5/qwen3_5.py | 138 ++ 2 files changed, 1791 insertions(+) create mode 100644 unsloth/models/patches/mlx_vlm_qwen3_5/generate.py create mode 100644 unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py diff --git a/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py b/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py new file mode 100644 index 0000000000..01c64357c3 --- /dev/null +++ b/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py @@ -0,0 +1,1653 @@ +import argparse +import codecs +import contextlib +import functools +import json +import time +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_reduce +from mlx_lm.generate import maybe_quantize_kv_cache as mlx_maybe_quantize_kv_cache +from mlx_lm.sample_utils import make_logits_processors, make_sampler +from tqdm import tqdm +from transformers import PreTrainedTokenizer + +from .models import cache +from .prompt_utils import apply_chat_template +from .turboquant import TurboQuantKVCache, turboquant_enabled +from .utils import ( + StoppingCriteria, + ThinkingBudgetCriteria, + group_images_by_shape, + load, + prepare_inputs, +) + +DEFAULT_MODEL_PATH = "mlx-community/nanoLLaVA-1.5-8bit" +DEFAULT_IMAGE = None +DEFAULT_AUDIO = None +DEFAULT_PROMPT = "What are these?" +DEFAULT_MAX_TOKENS = 256 +DEFAULT_TEMPERATURE = 0.0 +DEFAULT_TOP_P = 1.0 +DEFAULT_SEED = 0 +DEFAULT_TOP_K = 0 +DEFAULT_MIN_P = 0.0 +DEFAULT_REPETITION_CONTEXT_SIZE = 20 +DEFAULT_KV_GROUP_SIZE = 64 +DEFAULT_KV_QUANT_SCHEME = "uniform" +DEFAULT_COMPLETION_BATCH_SIZE = 32 +DEFAULT_PREFILL_BATCH_SIZE = 8 +DEFAULT_THINKING_START_TOKEN = "" +DEFAULT_THINKING_END_TOKEN = "" +DEFAULT_QUANTIZED_KV_START = 5000 +DEFAULT_PREFILL_STEP_SIZE = 2048 + + +def parse_arguments(): + parser = argparse.ArgumentParser( + description = "Generate text from an image using a model." + ) + parser.add_argument( + "--model", + type = str, + default = DEFAULT_MODEL_PATH, + help = "The path to the local model directory or Hugging Face repo.", + ) + parser.add_argument( + "--adapter-path", + type = str, + default = None, + help = "The path to the adapter weights.", + ) + parser.add_argument( + "--image", + type = str, + nargs = "+", + default = DEFAULT_IMAGE, + help = "URL or path of the image to process.", + ) + parser.add_argument( + "--audio", + type = str, + nargs = "+", + default = DEFAULT_AUDIO, + help = "URL or path of the audio to process.", + ) + parser.add_argument( + "--resize-shape", + type = int, + nargs = "+", + default = None, + help = "Resize shape for the image.", + ) + parser.add_argument( + "--prompt", + type = str, + nargs = "+", + default = DEFAULT_PROMPT, + help = "Message to be processed by the model.", + ) + parser.add_argument( + "--system", + type = str, + default = None, + help = "System message for the model.", + ) + parser.add_argument( + "--max-tokens", + type = int, + default = DEFAULT_MAX_TOKENS, + help = "Maximum number of tokens to generate.", + ) + parser.add_argument( + "--temperature", + type = float, + default = DEFAULT_TEMPERATURE, + help = "Temperature for sampling.", + ) + parser.add_argument("--chat", action = "store_true", help = "Chat in multi-turn style.") + parser.add_argument("--verbose", action = "store_false", help = "Detailed output.") + parser.add_argument( + "--eos-tokens", + type = str, + nargs = "+", + default = None, + help = "EOS tokens to add to the tokenizer.", + ) + parser.add_argument( + "--max-kv-size", + type = int, + default = None, + help = "Maximum KV size for the prompt cache.", + ) + parser.add_argument( + "--kv-bits", + type = float, + default = None, + help = "Number of bits to quantize the KV cache to.", + ) + parser.add_argument( + "--kv-quant-scheme", + type = str, + choices = ("uniform", "turboquant"), + default = DEFAULT_KV_QUANT_SCHEME, + help = "KV cache quantization backend. Fractional --kv-bits values use " + "TurboQuant automatically.", + ) + parser.add_argument( + "--kv-group-size", + type = int, + default = DEFAULT_KV_GROUP_SIZE, + help = "Group size for uniform KV cache quantization.", + ) + parser.add_argument( + "--quantized-kv-start", + type = int, + default = DEFAULT_QUANTIZED_KV_START, + help = "Start index for the quantized KV cache.", + ) + parser.add_argument( + "--skip-special-tokens", + action = "store_true", + help = "Skip special tokens in the detokenizer.", + ) + parser.add_argument( + "--force-download", + action = "store_true", + help = "Force download the model from Hugging Face.", + ) + parser.add_argument( + "--revision", + type = str, + default = "main", + help = "The specific model version to use (branch, tag, commit).", + ) + parser.add_argument( + "--trust-remote-code", + action = "store_true", + help = "Trust remote code when loading the model.", + ) + parser.add_argument( + "--quantize-activations", + "-qa", + action = "store_true", + help = "Enable activation quantization for QQLinear layers. " + "Only supported for models quantized with 'nvfp4' or 'mxfp8' modes.", + ) + parser.add_argument( + "--processor-kwargs", + type = json.loads, + default = {}, + help = "Extra processor kwargs as JSON. " + 'Example: --processor-kwargs \'{"cropping": false, "max_patches": 3}\'', + ) + parser.add_argument( + "--prefill-step-size", + type = int, + default = DEFAULT_PREFILL_STEP_SIZE, + help = "Number of tokens to process per prefill step. " + "Lower values reduce peak memory usage but may be slower. " + "Try 512 or 256 if you hit GPU memory errors during prefill.", + ) + parser.add_argument( + "--enable-thinking", + action = "store_true", + help = "Enable thinking mode in the chat template (e.g. for Qwen3.5).", + ) + parser.add_argument( + "--thinking-budget", + type = int, + default = None, + help = "Maximum number of thinking tokens before forcing the end-of-thinking token.", + ) + parser.add_argument( + "--thinking-start-token", + type = str, + default = DEFAULT_THINKING_START_TOKEN, + help = "Token that marks the start of a thinking block (default: %(default)s).", + ) + parser.add_argument( + "--thinking-end-token", + type = str, + default = DEFAULT_THINKING_END_TOKEN, + help = "Token that marks the end of a thinking block (default: %(default)s).", + ) + + return parser.parse_args() + + +def normalize_resize_shape( + values: Optional[Sequence[int]], +) -> Optional[Tuple[int, int]]: + if values is None: + return None + if not ( + isinstance(values, Sequence) + and not isinstance(values, (str, bytes)) + and len(values) in (1, 2) + and all(type(value) is int for value in values) + ): + raise ValueError("resize_shape must contain 1 or 2 integers") + return (values[0], values[0]) if len(values) == 1 else tuple(values) + + +# A stream on the default device just for generation +generation_stream = mx.new_stream(mx.default_device()) + + +def maybe_quantize_kv_cache( + prompt_cache, + quantized_kv_start, + kv_group_size, + kv_bits, + kv_quant_scheme: str = DEFAULT_KV_QUANT_SCHEME, +): + if kv_bits is None: + return + + if turboquant_enabled(kv_bits, kv_quant_scheme): + + def quantize_entry(entry): + if isinstance(entry, TurboQuantKVCache): + return entry + if isinstance(entry, cache.RotatingKVCache): + return entry + if isinstance(entry, cache.KVCache): + if entry.offset == 0: + # Empty: replace so update_and_fetch quantizes on the fly + return TurboQuantKVCache(bits = kv_bits) + if entry.offset < quantized_kv_start: + return entry + return TurboQuantKVCache.from_cache(entry, bits = kv_bits) + if isinstance(entry, cache.CacheList): + entry.caches = [quantize_entry(sub_entry) for sub_entry in entry.caches] + return entry + if isinstance(entry, list): + for i, sub_entry in enumerate(entry): + entry[i] = quantize_entry(sub_entry) + return entry + if isinstance(entry, tuple): + return tuple(quantize_entry(sub_entry) for sub_entry in entry) + return entry + + # Skip the last layer (before final norm/LM head) — it's highly + # sensitive to quantization in deep models (e.g. gemma-4-31b). + last_idx = len(prompt_cache) - 1 if len(prompt_cache) > 2 else -1 + for index, layer_cache in enumerate(prompt_cache): + if index == last_idx: + continue + prompt_cache[index] = quantize_entry(layer_cache) + return + + mlx_maybe_quantize_kv_cache( + prompt_cache, + quantized_kv_start = quantized_kv_start, + kv_group_size = kv_group_size, + kv_bits = int(kv_bits), + ) + + +@contextlib.contextmanager +def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None): + """ + A context manager to temporarily change the wired limit. + + Note, the wired limit should not be changed during an async eval. If an + async eval could be running pass in the streams to synchronize with prior + to exiting the context manager. + """ + if not mx.metal.is_available(): + yield + return + + model_bytes = tree_reduce( + lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0 + ) + max_rec_size = mx.device_info()["max_recommended_working_set_size"] + if model_bytes > 0.9 * max_rec_size: + model_mb = model_bytes // 2**20 + max_rec_mb = max_rec_size // 2**20 + print( + f"[WARNING] Generating with a model that requires {model_mb} MB " + f"which is close to the maximum recommended size of {max_rec_mb} " + "MB. This can be slow. See the documentation for possible work-arounds: " + "https://github.com/ml-explore/mlx-lm/tree/main#large-models" + ) + old_limit = mx.set_wired_limit(max_rec_size) + try: + yield + finally: + if streams is not None: + for s in streams: + mx.synchronize(s) + else: + mx.synchronize() + mx.set_wired_limit(old_limit) + + +@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 + + +class PromptCacheState: + """Holds KV cache and token history across conversation turns. + + Pass this to stream_generate via the ``prompt_cache_state`` kwarg to + reuse the KV cache from previous turns. Only the new tokens (after + the common prefix) are processed, avoiding redundant prefill. + """ + + def __init__(self): + self.cache: Optional[List[Any]] = None + self.token_ids: Optional[List[int]] = None + + def find_prefix_length(self, new_ids: list) -> int: + """Return the number of leading tokens that match the cached ids.""" + if self.token_ids is None: + return 0 + max_len = min(len(self.token_ids), len(new_ids)) + for i in range(max_len): + if self.token_ids[i] != new_ids[i]: + return i + return max_len + + def update(self, token_ids: list, kv_cache: list): + """Store the full token sequence and corresponding KV cache.""" + self.token_ids = list(token_ids) + self.cache = kv_cache + + +def generate_step( + input_ids: mx.array, + model: nn.Module, + pixel_values, + mask, + *, + max_tokens: int = DEFAULT_MAX_TOKENS, + temperature: float = DEFAULT_TEMPERATURE, + repetition_penalty: Optional[float] = None, + repetition_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE, + top_p: float = DEFAULT_TOP_P, + min_p: float = DEFAULT_MIN_P, + top_k: int = DEFAULT_TOP_K, + logit_bias: Optional[Dict[int, float]] = None, + prompt_cache: Optional[List[Any]] = None, + max_kv_size: Optional[int] = None, + kv_bits: Optional[float] = None, + kv_group_size: int = DEFAULT_KV_GROUP_SIZE, + kv_quant_scheme: str = DEFAULT_KV_QUANT_SCHEME, + quantized_kv_start: int = DEFAULT_QUANTIZED_KV_START, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None, + prefill_step_size: Optional[int] = DEFAULT_PREFILL_STEP_SIZE, + **kwargs, +) -> Generator[Tuple[mx.array, mx.array], None, None]: + """ + A generator producing token ids based on the given prompt from the model. + + Args: + input_ids (mx.array): The input prompt token ids. + model (nn.Module): The model to use for generation. + pixel_values: The pixel values for vision models (optional). + mask: The attention mask (optional). + max_tokens (int): Maximum number of tokens to generate. + temperature (float): The temperature for sampling, if 0 the argmax is used. + repetition_penalty (float, optional): The penalty factor for repeating + tokens. + repetition_context_size (int, optional): The number of tokens to + consider for repetition penalty. + top_p (float, optional): Nucleus sampling, higher means model considers + more less likely words. + min_p (float, optional): Minimum probability threshold relative to the + highest-probability token. + top_k (int, optional): Restrict sampling to the top-k tokens. + logit_bias (dictionary, optional): Additive logit bias. + prompt_cache (list, optional): Pre-existing KV cache for the prompt. + max_kv_size (int, optional): Maximum KV cache size. + kv_bits (float, optional): Number of bits for KV cache quantization. + kv_group_size (int): Group size for uniform KV cache quantization. + kv_quant_scheme (str): KV cache quantization backend. + quantized_kv_start (int): Start index for quantized KV cache. + sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a + token from a vector of log probabilities. + logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional): + A list of functions that take tokens and logits and return the processed + logits. + prefill_step_size (int): Number of tokens to process per prefill step. + Chunked prefill processes prompts in smaller chunks to reduce peak + memory usage. + + Yields: + Generator[Tuple[mx.array, mx.array], None, None]: A generator producing + one token and a vector of log probabilities. + """ + + quantize_cache_fn = functools.partial( + maybe_quantize_kv_cache, + quantized_kv_start = quantized_kv_start, + kv_group_size = kv_group_size, + kv_bits = kv_bits, + kv_quant_scheme = kv_quant_scheme, + ) + + if sampler is None: + sampler = make_sampler( + temp = temperature, + top_p = top_p, + min_p = min_p, + top_k = top_k, + ) + + processors = make_logits_processors( + logit_bias, repetition_penalty, repetition_context_size + ) + if logits_processors is not None: + processors.extend(logits_processors) + + y = input_ids + tokens = mx.array([], dtype = input_ids.dtype) + + thinking_budget_criteria = kwargs.pop("thinking_budget_criteria", None) + + # Create the KV cache for generation + if prompt_cache is None: + prompt_cache = cache.make_prompt_cache( + model.language_model, + max_kv_size = max_kv_size, + ) + + def _step(y, inputs_embeds = None): + nonlocal tokens, kwargs + + with mx.stream(generation_stream): + if "decoder_input_ids" in kwargs: + outputs = model.language_model( + cache = prompt_cache, + **kwargs, + ) + else: + outputs = model.language_model( + y, + inputs_embeds = inputs_embeds, + cache = prompt_cache, + **kwargs, + ) + + logits = outputs.logits[:, -1, :] + + if len(processors) > 0 and len(y) > 0: + tokens = mx.concat([tokens, y.flatten()]) + + for processor in processors: + logits = processor(tokens, logits) + + quantize_cache_fn(prompt_cache) + + logprobs = logits - mx.logsumexp(logits) + y = sampler(logprobs) + + if outputs.cross_attention_states is not None: + kwargs = {"cross_attention_states": outputs.cross_attention_states} + elif outputs.encoder_outputs is not None: + kwargs = {"encoder_outputs": outputs.encoder_outputs} + else: + kwargs = {} + + return y, logprobs.squeeze(0) + + with mx.stream(generation_stream): + # Get input embeddings (handles both multimodal and text-only) + embedding_output = model.get_input_embeddings( + input_ids, pixel_values, mask = mask, **kwargs + ) + + inputs_embeds = embedding_output.inputs_embeds + + kwargs.update( + { + k: v + for k, v in embedding_output.to_dict().items() + if k != "inputs_embeds" and v is not None + } + ) + if getattr(model, "no_chunked_prefill", False): + prefill_step_size = None + if prefill_step_size is not None and inputs_embeds.shape[1] > prefill_step_size: + # Chunked prefill with embeddings + total_tokens = inputs_embeds.shape[1] + with tqdm(total = total_tokens, desc = "Prefill", unit = "tok") as pbar: + while inputs_embeds.shape[1] > 1: + n_to_process = min(prefill_step_size, inputs_embeds.shape[1] - 1) + model.language_model( + inputs = input_ids[:, :n_to_process], + inputs_embeds = inputs_embeds[:, :n_to_process], + cache = prompt_cache, + n_to_process = n_to_process, + **kwargs, + ) + quantize_cache_fn(prompt_cache) + mx.eval([c.state for c in prompt_cache]) + inputs_embeds = inputs_embeds[:, n_to_process:] + input_ids = input_ids[:, n_to_process:] + mx.clear_cache() + pbar.update(n_to_process) + + input_ids = input_ids[:, -1:] + + y, logprobs = _step(input_ids, inputs_embeds = inputs_embeds) + + mx.async_eval(y) + + n = 0 + while True: + if n != max_tokens: + next_y, next_logprobs = _step(y[None]) + mx.async_eval(next_y) + if n == 0: + mx.eval(y) + if n == max_tokens: + break + + yield y.item(), logprobs + if n % 256 == 0: + mx.clear_cache() + + if thinking_budget_criteria is not None: + next_y = thinking_budget_criteria.apply_forced_token(next_y) + y, logprobs = next_y, next_logprobs + n += 1 + + +def stream_generate( + model: nn.Module, + processor: PreTrainedTokenizer, + prompt: str, + image: Union[str, List[str]] = None, + audio: Union[str, List[str]] = None, + **kwargs, +) -> Union[str, Generator[str, None, None]]: + """ + A generator producing text based on the given prompt from the model. + + Args: + model (nn.Module): The model to use for generation. + processor (PreTrainedTokenizer): The tokenizer/processor. + prompt (str): The input prompt text. + image (Union[str, List[str]], optional): Image path(s) or URL(s). + audio (Union[str, List[str]], optional): Audio file path(s). + prefill_step_size (int, optional): Number of tokens to process per prefill + step. When set, enables chunked prefill which processes long prompts in + smaller chunks to reduce peak memory usage. + kwargs: Additional options passed to :func:`generate_step`. + See :func:`generate_step` for more details. + + Yields: + Generator[GenerationResult]: A generator producing GenerationResult objects + containing the generated text, tokens, and statistics. + """ + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Set up thinking budget criteria if requested + thinking_budget = kwargs.pop("thinking_budget", None) + thinking_end_token = kwargs.pop("thinking_end_token", DEFAULT_THINKING_END_TOKEN) + thinking_start_token = kwargs.pop( + "thinking_start_token", DEFAULT_THINKING_START_TOKEN + ) + enable_thinking = kwargs.pop("enable_thinking", False) + + # Skip special tokens + skip_special_tokens = kwargs.pop("skip_special_tokens", False) + skip_special_token_ids = ( + set(tokenizer.all_special_ids) + if skip_special_tokens and hasattr(tokenizer, "all_special_ids") + else [] + ) + + add_special_tokens = ( + getattr(processor, "chat_template", None) is None + if model.config.model_type in ["gemma3", "gemma3n", "gemma4"] + else True + ) + + resize_shape = normalize_resize_shape(kwargs.pop("resize_shape", None)) + image_token_index = getattr(model.config, "image_token_index", None) + vision_cache = kwargs.pop("vision_cache", None) + + if kwargs.get("input_ids", None) is not None: + input_ids = kwargs.pop("input_ids") + pixel_values = kwargs.pop("pixel_values", None) + mask = kwargs.pop("mask", None) + else: + inputs = prepare_inputs( + processor, + images = image, + audio = audio, + prompts = prompt, + image_token_index = image_token_index, + resize_shape = resize_shape, + add_special_tokens = add_special_tokens, + **kwargs, + ) + input_ids = inputs.get("input_ids", None) + pixel_values = inputs.get("pixel_values", None) + mask = inputs.get("attention_mask", None) + data_kwargs = { + k: v + for k, v in inputs.items() + if k not in ["input_ids", "pixel_values", "attention_mask"] + } + kwargs.update(data_kwargs) + + # Vision feature caching: reuse cached image features across turns + if vision_cache is not None and image is not None and pixel_values is not None: + cached = vision_cache.get(image) + if cached is not None: + kwargs["cached_image_features"] = cached + elif hasattr(model, "encode_image"): + features = model.encode_image(pixel_values) + mx.eval(features) + vision_cache.put(image, features) + kwargs["cached_image_features"] = features + + # Prompt cache reuse: skip common prefix from previous turn + prompt_cache_state = kwargs.pop("prompt_cache_state", None) + reused_prefix_len = 0 + full_input_ids_list = input_ids.flatten().tolist() + + if prompt_cache_state is not None and prompt_cache_state.cache is not None: + prefix_len = prompt_cache_state.find_prefix_length(full_input_ids_list) + if prefix_len > 0 and prefix_len < input_ids.shape[1]: + reused_prefix_len = prefix_len + # Trim to only new tokens + input_ids = input_ids[:, prefix_len:] + if mask is not None: + mask = mask[:, prefix_len:] + # Only skip vision if no image tokens in the new (trimmed) tokens + image_token_id = getattr(model.config, "image_token_id", None) or getattr( + model.config, "image_token_index", None + ) + new_ids = input_ids.flatten().tolist() + has_image_in_new = image_token_id is not None and image_token_id in new_ids + if not has_image_in_new: + pixel_values = None + kwargs.pop("cached_image_features", None) + # Reuse the saved KV cache (trimmed to prefix length) + kv_cache = prompt_cache_state.cache + # Trim cache to prefix_len in case it includes generated tokens + for c in kv_cache: + if hasattr(c, "keys") and c.keys is not None: + cached_len = c.keys.shape[2] + if cached_len > prefix_len: + c.keys = c.keys[:, :, :prefix_len, :] + c.values = c.values[:, :, :prefix_len, :] + if hasattr(c, "offset"): + c.offset = prefix_len + kwargs["prompt_cache"] = kv_cache + + if thinking_budget is not None: + thinking_start_token_id = tokenizer.encode( + thinking_start_token, add_special_tokens = False + )[-1] + enable_thinking = enable_thinking and ( + thinking_start_token_id in input_ids.flatten().tolist() + ) + tokenizer.thinking_budget_criteria = ThinkingBudgetCriteria( + tokenizer = tokenizer, + thinking_budget = thinking_budget, + thinking_end_token = thinking_end_token, + thinking_start_token = thinking_start_token, + enable_thinking = enable_thinking, + ) + kwargs["thinking_budget_criteria"] = tokenizer.thinking_budget_criteria + else: + tokenizer.thinking_budget_criteria = None + + # Ensure we have a prompt_cache we can track for reuse + if "prompt_cache" not in kwargs: + kwargs["prompt_cache"] = cache.make_prompt_cache( + model.language_model, + max_kv_size = kwargs.get("max_kv_size", None), + ) + tracked_cache = kwargs["prompt_cache"] + + total_prompt_tokens = reused_prefix_len + input_ids.size + + with wired_limit(model, [generation_stream]): + detokenizer = processor.detokenizer + detokenizer.reset() + thinking_criteria = getattr(tokenizer, "thinking_budget_criteria", None) + gen = generate_step(input_ids, model, pixel_values, mask, **kwargs) + tic = time.perf_counter() + + generated_tokens = [] + for n, (token, logprobs) in enumerate(gen): + if n == 0: + prompt_time = time.perf_counter() - tic + prompt_tps = total_prompt_tokens / prompt_time + tic = time.perf_counter() + + generated_tokens.append(token) + + # Check thinking budget and force token if needed + if thinking_criteria is not None: + thinking_criteria(token) + + # Stop generation if the token is in the eos_token_ids + if tokenizer.stopping_criteria(token): + break + + detokenizer.add_token(token, skip_special_token_ids = skip_special_token_ids) + + # Yield the last segment if streaming + yield GenerationResult( + text = detokenizer.last_segment, + token = token, + logprobs = logprobs, + prompt_tokens = total_prompt_tokens, + generation_tokens = n + 1, + total_tokens = total_prompt_tokens + n + 1, + prompt_tps = prompt_tps, + generation_tps = (n + 1) / (time.perf_counter() - tic), + peak_memory = mx.get_peak_memory() / 1e9, + ) + + detokenizer.finalize() + yield GenerationResult( + text = detokenizer.last_segment, + token = token, + logprobs = logprobs, + prompt_tokens = total_prompt_tokens, + generation_tokens = n + 1, + total_tokens = total_prompt_tokens + n + 1, + prompt_tps = prompt_tps, + generation_tps = (n + 1) / (time.perf_counter() - tic), + peak_memory = mx.get_peak_memory() / 1e9, + ) + + # Save cache state for potential reuse on next turn + if prompt_cache_state is not None: + all_ids = full_input_ids_list + [ + t.item() if hasattr(t, "item") else t for t in generated_tokens + ] + prompt_cache_state.update(all_ids, tracked_cache) + + # Cleanup after generation + mx.clear_cache() + + +def generate( + model: nn.Module, + processor: PreTrainedTokenizer, + prompt: str, + image: Union[str, List[str]] = None, + audio: Union[str, List[str]] = None, + verbose: bool = False, + **kwargs, +) -> GenerationResult: + """ + Generate text from the model. + + Args: + model (nn.Module): The language model. + tokenizer (PreTrainedTokenizer): The tokenizer. + prompt (str): The string prompt. + temperature (float): The temperature for sampling (default 0). + max_tokens (int): The maximum number of tokens (default 100). + verbose (bool): If ``True``, print tokens and timing information + (default ``False``). + formatter (Optional[Callable]): A function which takes a token and a + probability and displays it. + repetition_penalty (float, optional): The penalty factor for repeating tokens. + repetition_context_size (int, optional): The number of tokens to consider for repetition penalty. + """ + + if verbose: + print("=" * 10) + files = [] + if image is not None: + files.extend(image) + if audio is not None: + files.extend(audio) + if kwargs.get("video") is not None: + files.extend(kwargs.get("video")) + + print(f"Files: {files}", "\n") + + print("Prompt:", prompt) + + text = "" + last_response = None + + eos_tokens = kwargs.get("eos_tokens", None) + stopping_criteria = kwargs.get("stopping_criteria", None) + + # Get the tokenizer + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Add custom EOS tokens to the stopping criteria + if eos_tokens is not None: + tokenizer.stopping_criteria.add_eos_token_ids(eos_tokens) + + # Use custom stopping criteria + elif stopping_criteria is not None: + if isinstance(stopping_criteria, StoppingCriteria) or callable( + stopping_criteria + ): + tokenizer.stopping_criteria = stopping_criteria + else: + raise ValueError( + "stopping_criteria must be an instance of StoppingCriteria or a callable" + ) + else: + tokenizer.stopping_criteria.reset(model.config.eos_token_id) + + for response in stream_generate(model, processor, prompt, image, audio, **kwargs): + if verbose: + print(response.text, end = "", flush = True) + text += response.text + last_response = response + + if verbose: + print("\n" + "=" * 10) + if len(text) == 0: + print("No text generated for this prompt") + return GenerationResult( + text = text, + token = None, + logprobs = None, + prompt_tokens = 0, + generation_tokens = 0, + total_tokens = 0, + prompt_tps = 0.0, + generation_tps = 0.0, + peak_memory = mx.get_peak_memory() / 1e9, + ) + print( + f"Prompt: {last_response.prompt_tokens} tokens, " + f"{last_response.prompt_tps:.3f} tokens-per-sec" + ) + print( + f"Generation: {last_response.generation_tokens} tokens, " + f"{last_response.generation_tps:.3f} tokens-per-sec" + ) + print(f"Peak memory: {last_response.peak_memory:.3f} GB") + + return GenerationResult( + text = text, + token = last_response.token, + logprobs = last_response.logprobs, + prompt_tokens = last_response.prompt_tokens, + generation_tokens = last_response.generation_tokens, + total_tokens = last_response.total_tokens, + prompt_tps = last_response.prompt_tps, + generation_tps = last_response.generation_tps, + peak_memory = last_response.peak_memory, + ) + + +@dataclass +class BatchGenerationResult: + """ + Result of batch generation with optional image size tracking. + + Attributes: + texts: Generated text for each sample + tokens: Last generated token for each sample + logprobs: Log probabilities for each sample + prompt_tokens: Number of prompt tokens per sample + generation_tokens: Number of generated tokens per sample + total_tokens: Total tokens (prompt + generation) per sample + prompt_tps: Prompt tokens per second per sample + generation_tps: Generation tokens per second per sample + peak_memory: Peak memory usage in GB + image_sizes: Original (height, width) for each image (for tracking) + """ + + texts: List[str] + tokens: List[Optional[int]] + logprobs: List[Optional[List[float]]] + prompt_tokens: List[int] + generation_tokens: List[int] + total_tokens: List[int] + prompt_tps: List[float] + generation_tps: List[float] + peak_memory: float = 0.0 + image_sizes: Optional[List[Tuple[int, int]]] = None + + +def _left_pad_prompts(prompts, max_length = None): + if max_length is None: + max_length = max(len(p) for p in prompts) + + return mx.array([[0] * (max_length - len(p)) + p for p in prompts]) + + +def _make_cache(model, left_padding): + """ + Convert a list of regular caches into their corresponding + batch-aware caches. + """ + + def to_batch_cache(c): + if isinstance(c, cache.KVCache): + return cache.BatchKVCache(left_padding) + elif isinstance(c, cache.ChunkedKVCache): + return cache.BatchKVCache(left_padding) + elif isinstance(c, cache.SimpleKVCache): + return cache.BatchKVCache(left_padding) + elif isinstance(c, cache.ArraysCache): + c.left_padding = mx.array(left_padding) + return c + elif isinstance(c, cache.RotatingKVCache): + if c.keep > 0: + raise ValueError("RotatingKVCache with keep tokens is not supported.") + return cache.BatchRotatingKVCache(c.max_size, left_padding) + elif isinstance(c, cache.CacheList): + return cache.CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches)) + elif isinstance(c, tuple): + return cache.CacheList(*(to_batch_cache(sub_c) for sub_c in c)) + else: + raise ValueError(f"{type(c)} does not yet support batching") + + if hasattr(model, "make_cache"): + model_cache = model.make_cache() + return [to_batch_cache(c) for c in model_cache] + else: + return [cache.BatchKVCache(left_padding) for _ in model.layers] + + +@dataclass +class BatchStats: + """ + An data object to hold generation stats. + + Args: + prompt_tokens (int): The number of prompt tokens processed. + prompt_tps (float): The prompt processing tokens-per-second. + prompt_time (float): The time in seconds spent in prompt processing. + generation_tokens (int): The number of generated tokens. + generation_tps (float): The tokens-per-second for generation. + generation_time (float): The time in seconds spent in generation . + peak_memory (float): The peak memory used so far in GB. + """ + + prompt_tokens: int = 0 + prompt_tps: float = 0 + prompt_time: float = 0 + generation_tokens: int = 0 + generation_tps: float = 0 + generation_time: float = 0 + peak_memory: float = 0 + + +@dataclass +class BatchResponse: + """ + An data object to hold a batch generation response. + + Args: + texts: (List[str]): The generated text for each prompt. + stats (BatchStats): Statistics about the generation. + image_sizes: (Optional[List[Tuple[int, int]]]): Original (height, width) + for each image. Useful for tracking which images produced which responses + and for debugging padding/batching behavior. + """ + + texts: List[str] + stats: BatchStats + image_sizes: Optional[List[Tuple[int, int]]] = None + + +@dataclass +class Batch: + uids: List[int] + y: mx.array + logprobs: mx.array + max_tokens: List[int] + num_tokens: List[int] + cache: List[Any] + + def __len__(self): + return len(self.uids) + + def filter(self, keep_idx: List[int]): + self.uids = [self.uids[k] for k in keep_idx] + self.max_tokens = [self.max_tokens[k] for k in keep_idx] + self.num_tokens = [self.num_tokens[k] for k in keep_idx] + keep_idx = mx.array(keep_idx, mx.int32) + self.y = self.y[keep_idx] + self.logprobs = self.logprobs[keep_idx] + for c in self.cache: + c.filter(keep_idx) + + def extend(self, other): + self.uids.extend(other.uids) + self.y = mx.concatenate([self.y, other.y]) + self.logprobs = mx.concatenate([self.logprobs, other.logprobs]) + self.num_tokens.extend(other.num_tokens) + self.max_tokens.extend(other.max_tokens) + for c, o in zip(self.cache, other.cache): + c.extend(o) + + +class BatchGenerator: + @dataclass + class Response: + uid: int + token: int + logprobs: mx.array + finish_reason: Optional[str] + + def __init__( + self, + model, + processor, + max_tokens: int = DEFAULT_MAX_TOKENS, + stop_tokens: Optional[set] = None, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + completion_batch_size: int = DEFAULT_COMPLETION_BATCH_SIZE, + prefill_batch_size: int = DEFAULT_PREFILL_BATCH_SIZE, + prefill_step_size: Optional[int] = DEFAULT_PREFILL_STEP_SIZE, + prompt_cache = None, + ): + self.model = model + self.unprocessed_prompts = [] + self.max_tokens = max_tokens + self.processor = processor + self.tokenizer = ( + processor.tokenizer if hasattr(processor, "tokenizer") else processor + ) + self.sampler = sampler or (lambda x: mx.argmax(x, axis = -1)) + self.uid_count = 0 + self.prefill_step_size = prefill_step_size + self.prefill_batch_size = prefill_batch_size + self.completion_batch_size = completion_batch_size + self.prompt_cache = prompt_cache + self._stats = BatchStats() + + self.tokenizer.stopping_criteria.add_eos_token_ids(stop_tokens) + + self.active_batch = None + + def insert(self, prompts, max_tokens: Union[List[int], int, None] = None): + uids = [] + + if max_tokens is None or isinstance(max_tokens, int): + max_tokens = [max_tokens or self.max_tokens] * len(prompts) + + for p, m in zip(prompts, max_tokens): + self.unprocessed_prompts.append((self.uid_count, p, m)) + uids.append(self.uid_count) + self.uid_count += 1 + # Sort in ascending order of length + self.unprocessed_prompts = sorted( + self.unprocessed_prompts, key = lambda x: len(x[1]) + ) + return uids + + def _process_prompts(self, prompts, **kwargs) -> Batch: + uids, inputs, max_tokens = zip(*prompts) + lengths = [len(p) for p in inputs] + max_length = max(lengths) + + self._stats.prompt_tokens += sum(lengths) + left_padding = [max_length - l for l in lengths] + inputs = _left_pad_prompts(inputs, max_length = max_length) + + if self.prompt_cache is not None: + prompt_cache = self.prompt_cache + elif len(uids) == 1 and max(left_padding) == 0: + # Single prompt with no padding: use standard caches to avoid + # numerical divergence from batch cache wrappers. + prompt_cache = cache.make_prompt_cache(self.model) + else: + prompt_cache = _make_cache(self.model, left_padding) + + # Slice batch data in kwargs to match current batch size + batch_size = len(uids) + for key, value in kwargs.items(): + if isinstance(value, mx.array) and value.ndim > 0: + kwargs[key] = value[:batch_size] + + inputs_embeds = kwargs.pop("inputs_embeds", None) + if inputs_embeds is None: + raise ValueError("inputs_embeds is required") + + if ( + self.prefill_step_size is not None + and inputs_embeds.shape[1] > self.prefill_step_size + ): + # Chunked prefill with embeddings + while inputs_embeds.shape[1] > 1: + n_to_process = min(self.prefill_step_size, inputs_embeds.shape[1] - 1) + self.model( + inputs[:, :n_to_process], + cache = prompt_cache, + inputs_embeds = inputs_embeds[:, :n_to_process], + n_to_process = n_to_process, + **kwargs, + ) + mx.eval([c.state for c in prompt_cache]) + inputs_embeds = inputs_embeds[:, n_to_process:] + inputs = inputs[:, n_to_process:] + mx.clear_cache() + + y, logprobs = self._step( + inputs, prompt_cache, inputs_embeds = inputs_embeds, **kwargs + ) + + mx.async_eval(y, logprobs) + mx.clear_cache() + return Batch( + list(uids), y, logprobs, list(max_tokens), [0] * len(uids), prompt_cache + ) + + def _step(self, input_tokens: mx.array, prompt_cache: List[Any], **kwargs): + output = self.model(input_tokens, cache = prompt_cache, **kwargs) + logits = output.logits[:, -1, :] + logprobs = logits - mx.logsumexp(logits, axis = -1, keepdims = True) + sampled = self.sampler(logprobs) + + # TODO: Add KV cache quantization if specified + return sampled, logprobs + + def stats(self): + self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time + self._stats.generation_tps = ( + self._stats.generation_tokens / self._stats.generation_time + ) + self._stats.peak_memory = mx.get_peak_memory() / 1e9 + return self._stats + + def _next(self, **kwargs): + tic = time.perf_counter() + + prompt_processing = False + batch = self.active_batch + num_active = len(batch) if batch else 0 + num_to_add = self.completion_batch_size - num_active + while num_to_add >= self.prefill_batch_size: + prompts = self.unprocessed_prompts[: self.prefill_batch_size] + # Finish processing the last examples of the last batch + if len(prompts) == 0 and num_active > 0: + break + # No more prompts and no more completions, all done + elif len(prompts) == 0: + self.active_batch = None + return [] + # Process prompts + if batch is not None and not prompt_processing: + # Finish any active completion tokens + mx.eval(batch.y, batch.logprobs) + self._stats.generation_time += time.perf_counter() - tic + tic = time.perf_counter() + + batch = self._process_prompts(prompts, **kwargs) + self.unprocessed_prompts = self.unprocessed_prompts[ + self.prefill_batch_size : + ] + prompt_processing = True + # If there was no active batch, set it + if self.active_batch is None: + self.active_batch = batch + else: + self.active_batch.extend(batch) + + num_active = len(self.active_batch) + num_to_add -= len(batch) + + batch = self.active_batch + y, logprobs = batch.y, batch.logprobs + batch.y, batch.logprobs = self._step(y[:, None], batch.cache) + mx.async_eval(batch.y, batch.logprobs) + + y = y.tolist() + toc = time.perf_counter() + if prompt_processing: + self._stats.prompt_time += toc - tic + else: + self._stats.generation_time += toc - tic + keep_idx = [] + end_idx = [] + responses = [] + + for e, (t, uid, num_tok, max_tok) in enumerate( + zip(y, batch.uids, batch.num_tokens, batch.max_tokens) + ): + num_tok += 1 + batch.num_tokens[e] = num_tok + if self.tokenizer.stopping_criteria(t): + finish_reason = "stop" + end_idx.append(e) + elif num_tok >= max_tok: + finish_reason = "length" + end_idx.append(e) + else: + finish_reason = None + keep_idx.append(e) + responses.append(self.Response(uid, t, logprobs[e], finish_reason)) + + # Remove any finished completions + if len(end_idx): + if len(keep_idx) > 0: + batch.filter(keep_idx) + else: + self.active_batch = None + + self._stats.generation_tokens += len(responses) + + if len(responses) > 0 and self._stats.generation_tokens % 100 == 0: + mx.clear_cache() + + return responses + + def next(self, **kwargs): + return self._next(**kwargs) + + +def batch_generate( + model, + processor, + images: Union[str, List[str]] = None, + audios: Union[str, List[str]] = None, + prompts: List[str] = None, + max_tokens: Union[int, List[int]] = 128, + verbose: bool = False, + group_by_shape: bool = True, + track_image_sizes: bool = True, + **kwargs, +): + """ + Generate responses for the given batch of prompts with variable-sized images. + + This function implements the transformers-style approach to batching: + 1. Group images with the same shape for efficient batch processing + 2. Process each group as a batch (no padding waste within groups) + 3. Track original image sizes for proper attention masking + 4. Restore results to original batch order + + Key insight: Instead of padding all images to the same spatial dimensions + (which wastes computation and may hurt accuracy), we group same-sized + images together so there's zero padding within each group. + + Args: + model (nn.Module): The language model. + processor (PreTrainedTokenizer): The tokenizer/processor. + images (Union[str, List[str]]): Images (paths, URLs, or PIL images). + audios (Union[str, List[str]]): Audio files (not yet supported for batching). + prompts (List[str]): The input prompts. + max_tokens (Union[int, List[int]]): Maximum number of output tokens. This + can be per prompt if a list is provided. + verbose (bool): If ``True``, print tokens and timing information. + group_by_shape (bool): If ``True``, group same-shaped images for efficient + batch processing. + track_image_sizes (bool): If ``True``, track and return original image sizes. + kwargs: The remaining options get passed to :obj:`BatchGenerator`. + See :obj:`BatchGenerator` for more details. + + Returns: + BatchResponse with generated texts, statistics, and optionally image_sizes. + """ + from PIL import Image + + from .utils import process_image + + processor.detokenizer.reset() + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Handle single image case + if isinstance(images, str): + images = [images] + + # Handle no images case + if images is None: + texts, stats = _generate_batch( + model, processor, prompts, None, max_tokens, verbose, **kwargs + ) + return BatchResponse(texts, stats) + + # Load and preprocess images + image_processor = ( + processor.image_processor if hasattr(processor, "image_processor") else None + ) + + processed_images = [] + image_sizes_original = [] + for img in images: + if isinstance(img, str): + pil_img = process_image(img, None, image_processor) + elif isinstance(img, Image.Image): + pil_img = img + else: + pil_img = img + processed_images.append(pil_img) + # Track original size + if hasattr(pil_img, "height"): + image_sizes_original.append((pil_img.height, pil_img.width)) + else: + image_sizes_original.append((0, 0)) + + # Group images by shape for efficient processing (no padding within groups) + if group_by_shape and len(processed_images) > 1: + grouped_images, grouped_indices = group_images_by_shape(processed_images) + + if verbose: + print(f"[batch_generate] Found {len(grouped_images)} unique image shapes") + else: + # Single image or grouping disabled - treat as one group + shape = ( + (processed_images[0].height, processed_images[0].width) + if processed_images + else (0, 0) + ) + grouped_images = {shape: processed_images} + grouped_indices = {shape: list(range(len(processed_images)))} + + # Process each shape group + all_texts = [None] * len(prompts) + all_image_sizes = [None] * len(prompts) + total_stats = BatchStats() + + for shape, indices in grouped_indices.items(): + # Get images and prompts for this shape group + group_images = [processed_images[i] for i in indices] + group_prompts = [prompts[i] for i in indices] + group_sizes = [image_sizes_original[i] for i in indices] + + # Handle per-sample max_tokens + if isinstance(max_tokens, list): + group_max_tokens = [max_tokens[i] for i in indices] + else: + group_max_tokens = max_tokens + + # Process the entire group at once (same shape = no padding needed) + chunk_texts, chunk_stats = _generate_batch( + model, + processor, + group_prompts, + group_images, + group_max_tokens, + **kwargs, + ) + + # Store results in original order + for j, orig_idx in enumerate(indices): + all_texts[orig_idx] = chunk_texts[j] + all_image_sizes[orig_idx] = group_sizes[j] + + # Accumulate stats + total_stats.prompt_tokens += chunk_stats.prompt_tokens + total_stats.prompt_time += chunk_stats.prompt_time + total_stats.generation_tokens += chunk_stats.generation_tokens + total_stats.generation_time += chunk_stats.generation_time + + mx.clear_cache() + + # Compute final stats + if total_stats.prompt_time > 0: + total_stats.prompt_tps = total_stats.prompt_tokens / total_stats.prompt_time + if total_stats.generation_time > 0: + total_stats.generation_tps = ( + total_stats.generation_tokens / total_stats.generation_time + ) + total_stats.peak_memory = mx.get_peak_memory() / 1e9 + + if verbose: + print(f"[batch_generate] Finished processing {len(prompts)} samples") + print( + f"[batch_generate] Prompt: {total_stats.prompt_tokens} tokens, {total_stats.prompt_tps:.3f} tokens-per-sec" + ) + print( + f"[batch_generate] Generation: {total_stats.generation_tokens} tokens, " + f"{total_stats.generation_tps:.3f} tokens-per-sec" + ) + print(f"[batch_generate] Peak memory: {total_stats.peak_memory:.3f} GB") + + response = BatchResponse(all_texts, total_stats) + if track_image_sizes: + response.image_sizes = all_image_sizes + return response + + +def _generate_batch( + model, + processor, + prompts: List[str], + images: List = None, + max_tokens: Union[int, List[int]] = 100, + verbose: bool = False, + **kwargs, +) -> Tuple[List[str], BatchStats]: + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + batch_size = len(prompts) + + num_images_list = [ + 1 if i < (len(images) if images is not None else 0) else 0 + for i in range(len(prompts)) + ] + formatted_prompts = [ + apply_chat_template( + processor, + model.config, + p, + num_images = num_images_list[i], + ) + for i, p in enumerate(prompts) + ] + + add_special_tokens = ( + getattr(processor, "chat_template", None) is None + if model.config.model_type in ["gemma3", "gemma3n", "gemma4"] + else True + ) + + resize_shape = normalize_resize_shape(kwargs.pop("resize_shape", None)) + image_token_index = getattr(model.config, "image_token_index", None) + + inputs = prepare_inputs( + processor, + images = images, + audio = None, + prompts = formatted_prompts, + image_token_index = image_token_index, + resize_shape = resize_shape, + add_special_tokens = add_special_tokens, + pad_to_uniform_size = False, # Since images are pre-grouped by shape, they're already uniform size + ) + input_ids = inputs.get("input_ids", None) + pixel_values = inputs.get("pixel_values", None) + mask = inputs.get("attention_mask", None) + + data_kwargs = { + k: v + for k, v in inputs.items() + if k not in ["input_ids", "pixel_values", "attention_mask"] + } + + if getattr(model, "no_chunked_prefill", False): + kwargs.pop("prefill_step_size", None) + kwargs["prefill_step_size"] = None + + # Use batch_size for prefill and completion to ensure consistent processing + gen = BatchGenerator( + model.language_model, + processor, + prefill_batch_size = batch_size, + completion_batch_size = batch_size, + **kwargs, + ) + + with wired_limit(model, [generation_stream]): + embedding_output = model.get_input_embeddings( + input_ids, pixel_values, mask = mask, **data_kwargs + ) + + gen_kwargs = {**data_kwargs, **embedding_output.to_dict()} + + uids = gen.insert(input_ids.tolist(), max_tokens) + results = {uid: [] for uid in uids} + while responses := gen.next(**gen_kwargs): + for r in responses: + if r.finish_reason != "stop": + results[r.uid].append(r.token) + + detokenizer = processor.detokenizer + texts = [] + for uid in uids: + detokenizer.reset() + for t in results[uid]: + detokenizer.add_token(t) + detokenizer.finalize() + texts.append(detokenizer.text) + return texts, gen.stats() + + +def main(): + args = parse_arguments() + if isinstance(args.image, str): + args.image = [args.image] + + model, processor = load( + args.model, + args.adapter_path, + revision = args.revision, + trust_remote_code = args.trust_remote_code, + quantize_activations = args.quantize_activations, + ) + config = model.config + + prompt = args.prompt + + num_images = len(args.image) if args.image is not None else 0 + num_audios = ( + 1 if args.audio is not None else 0 + ) # TODO: Support multiple audio files + + chat_template_kwargs = {"enable_thinking": args.enable_thinking} + + prompt = apply_chat_template( + processor, + config, + prompt, + num_images = num_images, + num_audios = num_audios, + **chat_template_kwargs, + ) + + kwargs = {} + + if args.eos_tokens is not None: + eos_tokens = [] + for token in args.eos_tokens: + try: + decoded_token = codecs.decode(token, "unicode_escape") + eos_tokens.append(decoded_token) + except (UnicodeDecodeError, UnicodeError): + eos_tokens.append(token) + kwargs["eos_tokens"] = eos_tokens + + if args.skip_special_tokens: + kwargs["skip_special_tokens"] = args.skip_special_tokens + + # Add processor kwargs from JSON + if args.processor_kwargs: + kwargs.update(args.processor_kwargs) + + # Add thinking kwargs + kwargs["enable_thinking"] = args.enable_thinking + if args.thinking_budget is not None: + kwargs["thinking_budget"] = args.thinking_budget + kwargs["thinking_end_token"] = args.thinking_end_token + if args.thinking_start_token is not None: + kwargs["thinking_start_token"] = args.thinking_start_token + + if args.chat: + from .vision_cache import VisionFeatureCache + + vision_cache = VisionFeatureCache() + chat = [] + if args.system: + chat.append({"role": "system", "content": args.system}) + while user := input("User:"): + chat.append({"role": "user", "content": user}) + prompt = apply_chat_template( + processor, config, chat, num_images = num_images, **chat_template_kwargs + ) + response = "" + print("Assistant:", end = "") + stream_kwargs = { + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "vision_cache": vision_cache, + **kwargs, + } + if args.resize_shape is not None: + stream_kwargs["resize_shape"] = args.resize_shape + if args.prefill_step_size is not None: + stream_kwargs["prefill_step_size"] = args.prefill_step_size + + for chunk in stream_generate( + model, + processor, + prompt, + args.image, + args.audio, + **stream_kwargs, + ): + response += chunk.text + print(chunk.text, end = "") + + chat.append({"role": "assistant", "content": response}) + print() + + else: + gen_kwargs = { + "image": args.image, + "audio": args.audio, + "temperature": args.temperature, + "max_tokens": args.max_tokens, + "verbose": args.verbose, + "max_kv_size": args.max_kv_size, + "kv_bits": args.kv_bits, + "kv_group_size": args.kv_group_size, + "kv_quant_scheme": getattr( + args, "kv_quant_scheme", DEFAULT_KV_QUANT_SCHEME + ), + "quantized_kv_start": args.quantized_kv_start, + **kwargs, + } + if args.resize_shape is not None: + gen_kwargs["resize_shape"] = args.resize_shape + if args.prefill_step_size is not None: + gen_kwargs["prefill_step_size"] = args.prefill_step_size + + result = generate( + model, + processor, + prompt, + **gen_kwargs, + ) + if not args.verbose: + print(result.text) + + +if __name__ == "__main__": + print( + "Calling `python -m mlx_vlm.generate ...` directly is deprecated." + " Use `mlx_vlm generate` or `python -m mlx_vlm generate` instead." + ) + main() diff --git a/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py b/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py new file mode 100644 index 0000000000..dcc85e0b54 --- /dev/null +++ b/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py @@ -0,0 +1,138 @@ +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn + +from ..base import InputEmbeddingsFeatures +from ..qwen3_vl import Model as Qwen3VLModel +from ..qwen3_vl import processing_qwen3_vl # noqa: F401 +from ..qwen3_vl.qwen3_vl import masked_scatter +from .config import ModelConfig +from .language import LanguageModel +from .vision import VisionModel + + +class Model(Qwen3VLModel): + def __init__(self, config: ModelConfig): + # only initialize nn.Module, skip the initialization of vision_tower and language_model in the parent class + nn.Module.__init__(self) + self.config = config + self.vision_tower = VisionModel(config.vision_config) + self.language_model = LanguageModel(config.text_config, config) + + def get_input_embeddings( + self, + input_ids: Optional[mx.array] = None, + pixel_values: Optional[mx.array] = None, + **kwargs, + ): + image_grid_thw = kwargs.get("image_grid_thw", None) + video_grid_thw = kwargs.get("video_grid_thw", None) + mask = kwargs.get("mask", None) + grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw + + if pixel_values is None: + return InputEmbeddingsFeatures( + inputs_embeds = self.language_model.model.embed_tokens(input_ids) + ) + + dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.astype(dtype) + + # Get the input embeddings from the language model + inputs_embeds = self.language_model.model.embed_tokens(input_ids) + + cached = kwargs.get("cached_image_features", None) + if cached is not None: + hidden_states = cached + else: + # Get the ouptut hidden states from the vision model + hidden_states, _ = self.vision_tower(pixel_values, grid_thw) + + # Insert special image tokens in the input_ids + inputs_embeds, _ = self.merge_input_ids_with_image_features( + hidden_states, + inputs_embeds, + input_ids, + self.config.image_token_index, + self.config.video_token_index, + ) + + # Pre-calculate position_ids for chunked prefill + if image_grid_thw is not None or video_grid_thw is not None: + position_ids, rope_deltas = self.language_model.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, mask + ) + self.language_model._position_ids = position_ids + self.language_model._rope_deltas = rope_deltas + + return InputEmbeddingsFeatures( + inputs_embeds = inputs_embeds, + ) + + @staticmethod + def merge_input_ids_with_image_features( + image_features, inputs_embeds, input_ids, image_token_index, video_token_index + ): + special_image_mask = input_ids == image_token_index + special_video_mask = input_ids == video_token_index + special_image_mask = special_image_mask | special_video_mask + n_image_tokens = special_image_mask.sum() + special_image_mask = special_image_mask[..., None] + special_image_mask = mx.broadcast_to(special_image_mask, inputs_embeds.shape) + + n_image_features = image_features.shape[0] + n_image_mask_elements = special_image_mask.sum() + if n_image_mask_elements != image_features.size: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + inputs_embeds = masked_scatter( + inputs_embeds, special_image_mask, image_features + ) + + return inputs_embeds, special_image_mask + + def sanitize(self, weights): + # ignore mtp weights + weights = {key: value for key, value in weights.items() if "mtp." not in key} + + if self.config.text_config.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + norm_keys = ( + ".input_layernorm.weight", + ".post_attention_layernorm.weight", + "model.norm.weight", + ".q_norm.weight", + ".k_norm.weight", + ) + + sanitized_weights = {} + for key, value in weights.items(): + if "model" in key: + if "model.language_model" in key: + key = key.replace("model.language_model", "language_model.model") + elif "model.visual" in key: + key = key.replace("model.visual", "vision_tower") + elif "lm_head" in key: + key = key.replace("lm_head", "language_model.lm_head") + + if "conv1d.weight" in key and value.shape[-1] != 1: + value = value.moveaxis(2, 1) + if any(key.endswith(sfx) for sfx in norm_keys): + if value.ndim == 1: + value += 1.0 + + sanitized_weights[key] = value + + return sanitized_weights + + @property + def quant_predicate(self): + return self.language_model.quant_predicate + + @property + def cast_predicate(self): + return self.language_model.cast_predicate