Day-0 Muse Glimmer Support on Intel Platforms with vLLM and Hugging Face

from __future__ import annotations

import argparse
import importlib.util
import os
import time

import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoConfig
from transformers.distributed import DistributedConfig

def rank0_print(*values, **kwargs):
    if int(os.environ.get("RANK", "0")) == 0:
        print(*values, **kwargs)

def build_user_content(
    prompt: str, has_image: bool, has_video: bool
) -> str | list[dict]:
    """Turn a prompt with an <img>/<video> placeholder into chat-template content.

    Text-only -> the raw string. With one media placeholder -> a list of parts
    ({"type": "text"|"image"|"video"}) in document order, so the media lands at the
    placeholder's position; the chat template renders the image/video part as the
    single sentinel that MuseGlimmerProcessor expands.
    """
    if not (has_image or has_video):
        return prompt
    placeholder, media_type = ("<img>", "image") if has_image else ("<video>", "video")
    parts: list[dict] = []
    chunks = prompt.split(placeholder)
    for i, chunk in enumerate(chunks):
        if chunk:
            parts.append({"type": "text", "text": chunk})
        if i < len(chunks) - 1:
            parts.append({"type": media_type})
    return parts

def check_video_requirements():
    if importlib.util.find_spec("torchcodec") is None:
        raise RuntimeError(
            "Video inference requires torchcodec. Install torchcodec in the "
            "active venv."
        )

def move_to_model_device(value, device):
    if isinstance(value, torch.Tensor):
        return value.to(device)
    return value

def extract_assistant_content_ids(generated_ids: torch.Tensor, tokenizer) -> torch.Tensor:
    token_ids = generated_ids.tolist()
    message_id = tokenizer.convert_tokens_to_ids("<|message|>")
    start_id = tokenizer.convert_tokens_to_ids("<|start|>")
    stop_token_ids = {
        tokenizer.convert_tokens_to_ids(token)
        for token in ("<|eom|>", "<|eot|>", "<|end_of_text|>", "<|start|>", "<|message|>")
    }

    for message_pos, token_id in enumerate(token_ids):
        if token_id != message_id:
            continue

        header_start = 0
        for i in range(message_pos - 1, -1, -1):
            if token_ids[i] == start_id:
                header_start = i
                break

        header = tokenizer.decode(token_ids[header_start:message_pos], skip_special_tokens=True)
        if "assistant to=user" not in header:
            continue

        content_start = message_pos + 1
        content_end = len(token_ids)
        for i in range(content_start, len(token_ids)):
            if token_ids[i] in stop_token_ids:
                content_end = i
                break
        return generated_ids[content_start:content_end]

    return generated_ids

def main():
    parser = argparse.ArgumentParser(description="HF inference for MuseGlimmer")
    parser.add_argument("--hf_dir", required=True, help="Path to MuseGlimmer model directory")
    parser.add_argument(
        "--prompt",
        default="The meaning of life is",
        help="Text prompt (use <img> / <video> for vision placeholders)",
    )
    parser.add_argument(
        "--system",
        default="You are a helpful assistant.",
        help="System prompt; pass '' to omit. The SFT model is trained with a "
        "system turn and degenerates on short prompts without one.",
    )
    parser.add_argument("--image", default=None, help="Path to image file")
    parser.add_argument(
        "--video", default=None, help="Path to video file (use <video> placeholder)"
    )
    parser.add_argument("--max_new_tokens", type=int, default=1024)
    parser.add_argument("--temperature", type=float, default=0.0, help="0.0 = greedy")
    args = parser.parse_args()

    tp_plan = {
        "model.language_model.embed_tokens": "embedding_rowwise",
        "model.vision_tower.patch_embedder.patch_embedding": "colwise_gather_output",
        "model.vision_adapter.fc1": "colwise",
        "model.vision_adapter.fc2": "rowwise",
        "model.vision_projection": "colwise_gather_output",
        "model.vision_tower.layers.*.attn.q_proj": "colwise",
        "model.vision_tower.layers.*.attn.k_proj": "colwise",
        "model.vision_tower.layers.*.attn.v_proj": "colwise",
        "model.vision_tower.layers.*.attn.proj": "rowwise",
        "model.vision_tower.layers.*.mlp.fc1": "colwise",
        "model.vision_tower.layers.*.mlp.fc2": "rowwise",
        "model.language_model.layers.*.self_attn.q_proj": "colwise",
        "model.language_model.layers.*.self_attn.k_proj": "colwise",
        "model.language_model.layers.*.self_attn.v_proj": "colwise",
        "model.language_model.layers.*.self_attn.o_proj": "rowwise",
        "model.language_model.layers.*.self_attn.gate_proj": "colwise",
        "model.language_model.layers.*.mlp.gate_proj": "colwise",
        "model.language_model.layers.*.mlp.up_proj": "colwise",
        "model.language_model.layers.*.mlp.down_proj": "rowwise",
        "lm_head": "colwise_gather_output",
    }

    rank0_print(f"Loading model from {args.hf_dir}")
    t0 = time.time()
    model = AutoModelForImageTextToText.from_pretrained(
        args.hf_dir,
        torch_dtype=torch.bfloat16,
        distributed_config=DistributedConfig(tp_size=2, tp_plan=tp_plan),
    )
    model.eval()
    processor = AutoProcessor.from_pretrained(args.hf_dir)
    if hasattr(processor, "video_processor") and not hasattr(
        processor.video_processor, "patch_temporal"
    ):
        processor.video_processor.patch_temporal = processor.video_processor.temporal_patch_size
    rank0_print(f"Model loaded in {time.time() - t0:.1f}s")

    has_image = args.image is not None and "<img>" in args.prompt
    has_video = args.video is not None and "<video>" in args.prompt
    if has_video:
        check_video_requirements()

    # Read media through MuseGlimmerProcessor's shipped sub-processors: images as PIL;
    # video as a path that MuseGlimmerProcessor decodes + groups via MuseGlimmerVideoProcessor
    # (torchcodec, training-faithful) -- the same path validate_nll_hf.py uses.
    images = videos = None
    if has_image:
        images = [Image.open(args.image).convert("RGB")]
        rank0_print(f"Image: {args.image} ({images[0].width}x{images[0].height})")
    elif has_video:
        videos = [args.video]
        rank0_print(f"Video: {args.video}")

    # One user turn: the chat template emits BOS + the trailing assistant
    # generation prompt and renders each media part as a sentinel, which
    # processor(...) expands into the full span and returns pixel_values for.
    messages = []
    if args.system:
        messages.append({"role": "system", "content": args.system})
    messages.append(
        {
            "role": "user",
            "content": build_user_content(args.prompt, has_image, has_video),
        }
    )
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=text,
        images=images,
        videos=videos,
        return_tensors="pt",
    )

    input_ids = inputs["input_ids"].to(model.device)
    attention_mask = inputs["attention_mask"].to(model.device)
    # Pass the media tensors and their grids through together; the vision encoder
    # needs *_grid_thw to build patch sequence lengths.
    pixel_values = move_to_model_device(inputs.get("pixel_values"), model.device)
    image_grid_thw = move_to_model_device(inputs.get("image_grid_thw"), model.device)
    pixel_values_videos = move_to_model_device(
        inputs.get("pixel_values_videos"), model.device
    )
    video_grid_thw = move_to_model_device(inputs.get("video_grid_thw"), model.device)

    rank0_print(f"Input: {input_ids.shape[1]} tokens")
    rank0_print("Generating...")
    do_sample = args.temperature > 0
    gen_kwargs = {"max_new_tokens": args.max_new_tokens, "do_sample": do_sample}
    if do_sample:
        gen_kwargs["temperature"] = args.temperature

    t0 = time.time()
    with torch.no_grad():
        output_ids = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            pixel_values_videos=pixel_values_videos,
            video_grid_thw=video_grid_thw,
            **gen_kwargs,
        )
    elapsed = time.time() - t0

    generated_ids = output_ids[0, input_ids.shape[1] :]
    content_ids = extract_assistant_content_ids(generated_ids, processor.tokenizer)
    text_out = processor.tokenizer.decode(content_ids, skip_special_tokens=True)
    n_generated_tokens = len(generated_ids)
    n_content_tokens = len(content_ids)

    rank0_print(f"\n{'=' * 60}")
    rank0_print(f"PROMPT: {args.prompt}")
    rank0_print(f"{'=' * 60}")
    rank0_print(f"OUTPUT: {text_out}")
    rank0_print(f"{'=' * 60}")
    rank0_print(
        f"Generated {n_generated_tokens} raw tokens in {elapsed:.2f}s "
        f"({n_generated_tokens / elapsed:.1f} tok/s); decoded content tokens: {n_content_tokens}"
    )

if __name__ == "__main__":
    main()

from __future__ import annotations

import argparse
import copy
import importlib.util
import time

import torch
from PIL import Image
from transformers import AutoModel, AutoModelForImageTextToText, AutoProcessor


def build_user_content(
    prompt: str, has_image: bool, has_video: bool
) -> str | list[dict]:
    """Text-only -> the raw prompt. With an <img>/<video> placeholder -> a list of
    chat-template parts in document order, so the media lands at the placeholder.
    """
    if not (has_image or has_video):
        return prompt
    placeholder, media_type = ("<img>", "image") if has_image else ("<video>", "video")
    parts: list[dict] = []
    chunks = prompt.split(placeholder)
    for i, chunk in enumerate(chunks):
        if chunk:
            parts.append({"type": "text", "text": chunk})
        if i < len(chunks) - 1:
            parts.append({"type": media_type})
    return parts


def main():
    parser = argparse.ArgumentParser(
        description="CPU HF inference with optional DFlash speculative decoding"
    )
    parser.add_argument("--hf_dir", required=True, help="Path to HF model directory")
    parser.add_argument(
        "--assistant_dir",
        default=None,
        help=(
            "Path to the DFlash assistant HF directory. Omit to run plain "
            "autoregressive generation."
        ),
    )
    parser.add_argument(
        "--prompt",
        default="The meaning of life is",
        help="Text prompt (use <img> / <video> for vision placeholders)",
    )
    parser.add_argument(
        "--system",
        default="You are a helpful assistant.",
        help="System prompt; pass '' to omit.",
    )
    parser.add_argument("--image", default=None, help="Path to image file")
    parser.add_argument(
        "--video", default=None, help="Path to video file (use <video> placeholder)"
    )

    parser.add_argument("--max_new_tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.0, help="0.0 = greedy")
    parser.add_argument(
        "--assistant_confidence_threshold",
        type=float,
        default=0.4,
        help=(
            "DFlash draft-token confidence threshold; use 0.0 to keep full blocks. "
            "Only used with --assistant_dir."
        ),
    )
    parser.add_argument(
        "--assistant_block_size",
        type=int,
        default=None,
        help=(
            "Override the DFlash assistant block size (default: from its config). "
            "Only used with --assistant_dir."
        ),
    )
    args = parser.parse_args()

    print(f"Loading model from {args.hf_dir}")
    t0 = time.time()
    model = AutoModelForImageTextToText.from_pretrained(
        args.hf_dir,
        dtype=torch.bfloat16,
    ).eval()
    print(f"Model loaded in {time.time() - t0:.1f}s")

    generate_kwargs = {}
    if args.assistant_dir:
        print(f"Loading DFlash assistant from {args.assistant_dir}")
        t0 = time.time()
        assistant_model = AutoModel.from_pretrained(
            args.assistant_dir,
            dtype=torch.bfloat16,
        )
        if args.assistant_block_size is not None:
            assistant_model.config.block_size = args.assistant_block_size
        assistant_model.eval()
        print(
            f"Assistant loaded in {time.time() - t0:.1f}s; "
            f"block_size={assistant_model.config.block_size}, "
            f"target_layer_ids={assistant_model.config.target_layer_ids}"
        )
        generate_kwargs["assistant_model"] = assistant_model

    processor = AutoProcessor.from_pretrained(args.hf_dir)
    if hasattr(processor, "video_processor") and not hasattr(
        processor.video_processor, "patch_temporal"
    ):
        processor.video_processor.patch_temporal = processor.video_processor.temporal_patch_size

    has_image = args.image is not None and "<img>" in args.prompt
    has_video = args.video is not None and "<video>" in args.prompt
    if has_video and importlib.util.find_spec("torchcodec") is None:
        raise RuntimeError(
            "Video inference requires torchcodec; install it in the active venv."
        )

    # Images are passed as PIL objects; videos are passed as file paths and
    # decoded by the processor's video processor (requires torchcodec).
    images = videos = None
    if has_image:
        images = [Image.open(args.image).convert("RGB")]
        print(f"Image: {args.image} ({images[0].width}x{images[0].height})")
    elif has_video:
        videos = [args.video]
        print(f"Video: {args.video}")

    # Build a single user turn; the chat template adds BOS and the trailing
    # assistant generation prompt, and the processor expands each media part.
    messages = []
    if args.system:
        messages.append({"role": "system", "content": args.system})

    messages.append(
        {
            "role": "user",
            "content": build_user_content(args.prompt, has_image, has_video),
        }
    )

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=text,
        images=images,
        videos=videos,
        return_tensors="pt",
    )

    print(f"Input: {inputs['input_ids'].shape[1]} tokens")

    do_sample = args.temperature > 0
    generation_config = copy.deepcopy(model.generation_config)
    generation_config.max_new_tokens = args.max_new_tokens
    generation_config.do_sample = do_sample
    if do_sample:
        generation_config.temperature = args.temperature
    if args.assistant_dir:
        generation_config.speculation_type = "dflash"
        generation_config.assistant_confidence_threshold = (
            args.assistant_confidence_threshold
        )
        print("Generating with speculation_type='dflash' on CPU...")
    else:
        print("Generating on CPU...")

    generation_started_at = time.time()
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            **generate_kwargs,
            generation_config=generation_config,
        )
    elapsed = time.time() - generation_started_at

    generated_ids = output_ids[0, inputs["input_ids"].shape[1] :]
    n_tokens = len(generated_ids)

    # The model may emit a chat channel header before the final message marker.
    # Retain only the tokens that follow the last message marker.
    message_id = processor.tokenizer.convert_tokens_to_ids("<|message|>")
    generated_id_list = generated_ids.tolist()
    if message_id in generated_id_list:
        content_start = len(generated_id_list) - generated_id_list[::-1].index(message_id)
        content_ids = generated_ids[content_start:]
    else:
        content_ids = generated_ids
    text_out = processor.tokenizer.decode(content_ids, skip_special_tokens=True)

    print(f"\n{'=' * 60}")
    print(f"PROMPT: {args.prompt}")
    print(f"{'=' * 60}")
    print(f"OUTPUT: {text_out}")
    print(f"{'=' * 60}")
    print(
        f"Generated {n_tokens} tokens in {elapsed:.2f}s "
        f"({n_tokens / elapsed:.1f} tok/s)"
    )


if __name__ == "__main__":
    main()