"""
AutoModel - Main entry point for Dokodemo AI.

Supports both GPU and CPU-only inference.

GPU usage (same as before):
    model = AutoModel.from_pretrained(
        "mistralai/Mixtral-8x7B-Instruct-v0.1",
        compression="4bit",
        max_gpu_memory="4GB",
    )

CPU-only usage — full accuracy, any model, 4 GB RAM:
    model = AutoModel.from_pretrained(
        "meta-llama/Meta-Llama-3.1-70B",
        device="cpu",
        # No compression needed — mmap = full FP16 accuracy
    )

    # Or explicitly request CPU:
    model = AutoModel.from_pretrained("...", device="cpu")

When running on CPU:
- use_mmap=True is set automatically (OS page cache = weight cache)
- compression is set to None automatically (no accuracy drop)
- compute_dtype is auto-detected (bfloat16 if AVX-512 BF16, else float32)
- CPU threads are optimized automatically
- All model sizes work on 4 GB RAM with full precision
"""

from __future__ import annotations

import gc
import logging
import os
from typing import Callable, Iterator, Optional, Union

# Default HF cache location — stored on fast NVMe alongside the project.
# Can be overridden by HF_HOME env var or the cache_dir= parameter.
_DEFAULT_HF_CACHE = "/mnt/nvme1n1/DOKODEMOAI/hf_cache"

import torch
from transformers import AutoTokenizer

from dokodemo_ai.engine.cache import AdaptiveLayerCache
from dokodemo_ai.engine.inference import InferenceEngine
from dokodemo_ai.engine.scheduler import ExpertScheduler
from dokodemo_ai.engine.streaming import TensorStreamer
from dokodemo_ai.engine.vision import VisionEncoder, merge_image_tokens
from dokodemo_ai.graph.compiler import ModelGraphCompiler
from dokodemo_ai.graph.partition import GraphPartitioner
from dokodemo_ai.quantization.dynamic import DynamicQuantizer, Precision
from dokodemo_ai.utils.cpu_optimize import (
    get_compute_dtype,
    print_cpu_info,
    setup_threads,
    supports_bf16,
)
from dokodemo_ai.utils.memory import MemoryManager
from dokodemo_ai.utils.profiler import Profiler

logger = logging.getLogger(__name__)


def _parse_memory_string(s: str) -> int:
    """Parse a memory string like '4GB', '512MB', '2.5GB' into bytes."""
    s = s.strip().upper()
    if s.endswith("GB"):
        return int(float(s[:-2]) * 1024**3)
    if s.endswith("MB"):
        return int(float(s[:-2]) * 1024**2)
    if s.endswith("TB"):
        return int(float(s[:-2]) * 1024**4)
    return int(s)


class AutoModel:
    """
    Dokodemo AI model wrapper — GPU or CPU-only, any model size.

    GPU mode:
    - 3-level async prefetch (disk → CPU → GPU)
    - FP16 compute, importance-weighted layer caching
    - Up to 64x I/O reduction for MoE via sparse expert loading

    CPU mmap mode (default when no GPU):
    - Zero-copy loading: weights live in OS page cache
    - Full FP16/FP32 accuracy — zero accuracy drop, no quantization
    - Any model size fits in 4 GB RAM (OS manages paging)
    - BF16 compute on modern CPUs (AVX-512 BF16 / Apple Silicon)
    - FP32 compute on all other CPUs
    """

    def __init__(
        self,
        graph,
        engine: InferenceEngine,
        tokenizer,
        profiler: Profiler,
        vision_encoder: Optional["VisionEncoder"] = None,
    ):
        self._graph = graph
        self._engine = engine
        self.tokenizer = tokenizer
        self._profiler = profiler
        self._vision_encoder: Optional[VisionEncoder] = vision_encoder

    @classmethod
    def from_pretrained(
        cls,
        model_path: str,
        *,
        compression: Optional[str] = None,
        max_gpu_memory: Optional[Union[str, int]] = None,
        device: Optional[Union[str, torch.device]] = None,
        hf_token: Optional[str] = None,
        trust_remote_code: bool = False,
        num_io_workers: int = 2,
        enable_speculation: bool = True,
        profiling: bool = False,
        # CPU-specific options
        use_mmap: Optional[bool] = None,
        cpu_threads: Optional[int] = None,
        # Cache directory for downloaded HF models
        cache_dir: Optional[str] = None,
    ) -> "AutoModel":
        """
        Load a model for offloaded inference (GPU or CPU-only).

        Args:
            model_path: HuggingFace model ID or local path.
            compression: Weight compression.
                - None (default on CPU): Full FP16, zero accuracy drop.
                - "dynamic": Per-layer optimal quantization (GPU recommended).
                - "8bit" / "4bit" / "2bit": Uniform quantization.
                NOTE: On CPU, compression is automatically set to None to
                guarantee zero accuracy drop. Pass explicitly to override.
            max_gpu_memory: GPU VRAM limit, e.g. "4GB". GPU only.
            device: "cuda", "cuda:1", "mps", "cpu", or None (auto-detect).
            hf_token: HuggingFace token for gated models.
            trust_remote_code: Allow custom model code.
            num_io_workers: I/O prefetch threads (GPU mode only).
            enable_speculation: Speculative expert prefetching for MoE.
            profiling: Enable layer-by-layer performance profiling.
            use_mmap: Use OS memory-mapped loading (default: True on CPU,
                      False on GPU). With mmap, all model sizes fit in 4 GB
                      RAM at full accuracy — the OS manages weight paging.
            cpu_threads: Number of CPU threads for PyTorch compute.
                         Default: all available cores.

        Returns:
            An AutoModel instance ready for generation.

        CPU Example — any model, 4 GB RAM, zero accuracy drop:
            model = AutoModel.from_pretrained(
                "meta-llama/Meta-Llama-3.1-70B",
                device="cpu",
            )
            # Automatically uses mmap + no quantization + best CPU dtype
        """
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s [dokodemo] %(message)s",
            datefmt="%H:%M:%S",
        )

        # ── 0. Resolve HF cache directory ──────────────────────────────
        # Priority: explicit cache_dir arg > HF_HOME env var > project default
        _cache = cache_dir or os.environ.get("HF_HOME") or _DEFAULT_HF_CACHE
        if "HF_HOME" not in os.environ:
            os.environ["HF_HOME"] = _cache
        logger.info("HF cache: %s", _cache)

        # ── 1. Resolve device ──────────────────────────────────────────
        if device is None:
            mem_manager = MemoryManager()
            resolved_device = mem_manager.device
        else:
            resolved_device = torch.device(device)
            mem_manager = MemoryManager(device=resolved_device)

        is_cpu = resolved_device.type == "cpu"
        is_gpu = resolved_device.type in ("cuda", "mps")

        logger.info("Using device: %s", resolved_device)

        # ── 2. CPU-specific setup ──────────────────────────────────────
        if is_cpu:
            # Optimize thread counts
            setup_threads(n_compute=cpu_threads)
            print_cpu_info()

            # Default: mmap on CPU (zero-copy, works with 4 GB RAM)
            if use_mmap is None:
                use_mmap = True

            # Default: no compression on CPU (full accuracy)
            if compression is None:
                logger.info(
                    "CPU mode: compression=None (full FP16 accuracy). "
                    "Mmap enabled — any model fits in 4 GB RAM."
                )
            else:
                logger.warning(
                    "CPU mode: compression='%s' requested. "
                    "This will reduce accuracy. "
                    "Omit compression for full FP16 accuracy on CPU.",
                    compression,
                )

            # Pick best compute dtype for this CPU
            compute_dtype = get_compute_dtype()
            logger.info(
                "CPU compute dtype: %s%s",
                compute_dtype,
                " (native BF16 via AVX-512)" if supports_bf16() else " (float32 emulation)",
            )
        else:
            if use_mmap is None:
                use_mmap = False

            # BF16 on Ampere+ (sm80+): RTX 3000 series, A100, H100, ...
            # BF16 has FP32 dynamic range — prevents NaN/Inf on long sequences.
            # Pre-Ampere (sm70 and below): BF16 is emulated and slower; use FP16.
            if resolved_device.type == "cuda":
                idx = resolved_device.index if resolved_device.index is not None else 0
                major, _ = torch.cuda.get_device_capability(idx)
                if major >= 8:
                    compute_dtype = torch.bfloat16
                    logger.info(
                        "Ampere+ GPU (sm%d0) detected — using BF16 for "
                        "wider dynamic range and NaN resistance.",
                        major,
                    )
                else:
                    compute_dtype = torch.float16
                    logger.info(
                        "Pre-Ampere GPU (sm%d0) — using FP16 compute.", major
                    )
            else:
                compute_dtype = torch.float16  # MPS

        # ── 3. Compile model graph ─────────────────────────────────────
        logger.info("Compiling model graph for: %s", model_path)
        compiler = ModelGraphCompiler()
        graph = compiler.compile(
            model_path,
            hf_token=hf_token,
            trust_remote_code=trust_remote_code,
            cache_dir=_cache,
        )
        logger.info("Model graph compiled.\n%s", graph.summary())

        # ── 4. Feasibility check ───────────────────────────────────────
        # For MoE, use the sparse active-expert layer size (only active
        # experts are streamed per token, not all experts).
        feasibility_layer_bytes = graph.max_layer_size_bytes
        if graph.is_moe:
            moe_layer = next((l for l in graph.layers if l.is_moe), None)
            if moe_layer:
                feasibility_layer_bytes = (
                    moe_layer.non_expert_size_bytes
                    + moe_layer.num_active_experts * moe_layer.expert_size_bytes
                )
        feasible, msg = mem_manager.check_feasibility(
            max_layer_size_bytes=feasibility_layer_bytes,
            embedding_size_bytes=graph.embedding.size_bytes,
            head_size_bytes=graph.head.size_bytes,
            use_mmap=use_mmap,
        )
        if not feasible:
            raise RuntimeError(msg)
        logger.info(msg)

        # ── 5. Memory budget ───────────────────────────────────────────
        if max_gpu_memory is not None and not is_cpu:
            budget = (
                _parse_memory_string(max_gpu_memory)
                if isinstance(max_gpu_memory, str)
                else int(max_gpu_memory)
            )
        else:
            budget = mem_manager.get_weight_budget_bytes(use_mmap=use_mmap)

        if use_mmap:
            logger.info(
                "Mmap mode: no Python-level weight cache (OS handles it). "
                "Budget set to 0 for the layer cache."
            )
        else:
            logger.info("Weight cache budget: %.2f GB", budget / 1e9)

        # ── 6. Compression ─────────────────────────────────────────────
        _quantizer = None
        if compression is not None:
            _map = {
                "8bit": Precision.INT8,
                "4bit": Precision.INT4,
                "2bit": Precision.INT2,
            }
            if compression == "dynamic":
                _quantizer = DynamicQuantizer(default_precision=Precision.INT4)
                logger.info("Dynamic per-layer quantization enabled")
            elif compression in _map:
                prec = _map[compression]
                _quantizer = DynamicQuantizer(default_precision=prec)
                logger.info("Uniform %s quantization enabled", compression)
            else:
                raise ValueError(
                    f"Unknown compression '{compression}'. "
                    "Use 'dynamic', '8bit', '4bit', '2bit', or None."
                )

        # ── 7. Partition the graph ─────────────────────────────────────
        partitioner = GraphPartitioner(memory_budget_bytes=max(budget, 1))
        plan = partitioner.partition(graph)
        logger.info("Execution plan:\n%s", plan.summary())

        # ── 8. Profiler ────────────────────────────────────────────────
        profiler = Profiler(enabled=profiling)

        # ── 9. Tensor streamer ─────────────────────────────────────────
        if use_mmap:
            # CPU mmap: OS manages RAM, no explicit CPU buffer needed
            streamer = TensorStreamer(
                model_graph=graph,
                device=resolved_device,
                num_workers=1,          # minimal threading on CPU
                cpu_buffer_bytes=0,
                pin_memory=False,
                use_mmap=True,
            )
        else:
            cpu_buffer = mem_manager.get_cpu_buffer_budget_bytes()
            if graph.is_vlm:
                # VLM inference reads 14+ GB from NVMe during forward passes;
                # cap the CPU prefetch buffer to leave RAM headroom for the OS
                # page cache, preventing SWAP exhaustion.
                try:
                    import psutil as _psutil
                    _avail = _psutil.virtual_memory().available
                    # 1 GB max: smaller than before (was 1.5 GB) — every MB counts
                    # when SWAP headroom is tight.
                    vlm_cap = min(_avail // 8, 1024 * 1024 * 1024)  # max 1 GB
                except ImportError:
                    vlm_cap = 512 * 1024 * 1024  # 512 MB fallback
                cpu_buffer = min(cpu_buffer, vlm_cap)
                logger.info("VLM: CPU prefetch buffer capped at %.1f GB", cpu_buffer / 1e9)
            # VLMs: disable pin_memory to avoid locking CPU buffer in physical RAM.
            # Pinned memory cannot be swapped, creating unswappable pressure when
            # SWAP headroom is limited.  The GPU transfer speed penalty is acceptable
            # given that NVMe I/O is the actual bottleneck.
            use_pin_memory = (resolved_device.type == "cuda") and not graph.is_vlm
            streamer = TensorStreamer(
                model_graph=graph,
                device=resolved_device,
                num_workers=num_io_workers,
                cpu_buffer_bytes=cpu_buffer,
                pin_memory=use_pin_memory,
                use_mmap=False,
            )

        # ── 10. Adaptive layer cache ───────────────────────────────────
        # In CPU mmap mode: budget=0, so the cache is empty.
        # The OS page cache replaces it transparently.
        cache = AdaptiveLayerCache(
            memory_budget_bytes=budget,
            model_graph=graph,
            device=resolved_device,
        )

        # ── 11. Expert scheduler ───────────────────────────────────────
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=resolved_device,
            enable_speculation=enable_speculation,
        )

        # ── 12. Inference engine ───────────────────────────────────────
        engine = InferenceEngine(
            graph=graph,
            streamer=streamer,
            cache=cache,
            scheduler=scheduler,
            device=resolved_device,
            compute_dtype=compute_dtype,
            profiler=profiler,
        )

        # Pre-load embedding + LM head (CPU-resident; only final_norm lands on GPU)
        logger.info("Loading persistent weights (embedding + head)...")
        engine.load_persistent_weights()

        # Subtract any GPU memory already used by persistent weights from the
        # layer-cache budget so the total stays within max_gpu_memory.
        if resolved_device.type == "cuda" and cache.budget > 0:
            persistent_vram = torch.cuda.memory_allocated(resolved_device)
            if persistent_vram > 0:
                new_budget = max(cache.budget - persistent_vram, 0)
                logger.info(
                    "Persistent GPU weights: %.2f MB; layer cache %.2f GB → %.2f GB",
                    persistent_vram / 1e6, cache.budget / 1e9, new_budget / 1e9,
                )
                cache.budget = new_budget

        # ── 13. Tokenizer ──────────────────────────────────────────────
        logger.info("Loading tokenizer...")
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                str(graph.model_path),
                token=hf_token,
                trust_remote_code=trust_remote_code,
            )
        except ValueError as _e:
            if "trust_remote_code" in str(_e):
                logger.warning(
                    "Tokenizer requires trust_remote_code=True; retrying. (%s)", _e
                )
                tokenizer = AutoTokenizer.from_pretrained(
                    str(graph.model_path),
                    token=hf_token,
                    trust_remote_code=True,
                )
            else:
                raise

        # ── 14. Vision encoder (VLMs only) ─────────────────────────────
        vision_encoder: Optional[VisionEncoder] = None
        if graph.is_vlm and graph.vision_params:
            logger.info(
                "VLM detected (%s). Loading vision encoder (%d params) into GPU...",
                graph.vision_arch, len(graph.vision_params),
            )
            try:
                # Load all vision + projector tensors at once (they stay resident)
                all_vis_params = graph.vision_params + (graph.projector_params or [])
                vis_tensors = streamer.load_tensors_by_name(
                    all_vis_params, to_device=resolved_device
                )
                vision_encoder = VisionEncoder.from_tensors(
                    tensors=vis_tensors,
                    config=graph.config,
                    device=resolved_device,
                    compute_dtype=compute_dtype,
                    arch=graph.vision_arch,
                )
                # Free the raw tensor dict — weights are now owned by VisionEncoder's
                # nn.Parameters. This frees ~0.9 GB of GPU/CPU memory.
                del vis_tensors
                gc.collect()
                if resolved_device.type == "cuda":
                    torch.cuda.empty_cache()
                # Move vision encoder to CPU; it will be restored to GPU on-demand
                # before each encode() call and moved back after.  This saves
                # 300–600 MB of GPU VRAM during the LLM forward pass.
                vision_encoder = vision_encoder.cpu()
                if resolved_device.type == "cuda":
                    torch.cuda.empty_cache()
                logger.info("Vision encoder ready (CPU-resident; GPU-restored per encode).")
            except Exception as exc:
                logger.warning(
                    "Vision encoder load failed (%s); VLM image input disabled.", exc
                )

        # ── 15. Summary ────────────────────────────────────────────────
        mode = "CPU mmap (full accuracy)" if use_mmap else (
            "GPU" if is_gpu else "CPU"
        )
        vlm_tag = f" + VLM/{graph.vision_arch}" if graph.is_vlm else ""
        logger.info(
            "Model ready! Mode=%s%s, dtype=%s, device=%s",
            mode, vlm_tag, compute_dtype, resolved_device,
        )

        return cls(
            graph=graph,
            engine=engine,
            tokenizer=tokenizer,
            profiler=profiler,
            vision_encoder=vision_encoder,
        )

    # ── Public API ──────────────────────────────────────────────────────

    def generate(
        self,
        input_ids: torch.Tensor,
        images=None,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 0.9,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        do_sample: bool = True,
        prefill_chunk_size: int = 512,
        kv_cache_bits: int = 16,
        on_token: Optional[Callable[[int], None]] = None,
    ) -> torch.Tensor:
        """
        Generate tokens autoregressively (blocking).

        Works on GPU and CPU. On CPU with mmap, achieves full FP16 accuracy
        with any model size on 4 GB RAM.

        Args:
            input_ids: Input token IDs [batch, seq_len].
            images: Optional PIL image(s) or pixel_values tensor for VLMs.
                    Pass a single PIL Image, a list of PIL Images, or a
                    pre-processed [batch, 3, H, W] float tensor.
                    Ignored for non-VLM models.
            max_new_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_k: Top-k sampling (0 = disabled).
            top_p: Nucleus sampling threshold (1.0 = disabled).
            min_p: Min-P — removes tokens with prob < min_p × max_prob.
                   Better than top-p for creative/open-ended generation.
            repetition_penalty: Penalise repeated tokens (1.0 = off,
                   1.1–1.3 typical for repetition prevention).
            do_sample: False = greedy, True = sampling.
            prefill_chunk_size: Max tokens per prefill chunk (prevents OOM
                   on long prompts).
            kv_cache_bits: KV cache precision — 16 (FP16) or 8 (INT8).
                   INT8 halves KV VRAM and doubles max context length.
            on_token: Optional callback called with each new token ID for
                   sync streaming without switching to stream_generate().

        Example::

            inputs = model.tokenizer("Hello!", return_tensors="pt")
            out = model.generate(inputs["input_ids"], max_new_tokens=100,
                                 repetition_penalty=1.1, min_p=0.05)
            print(model.tokenizer.decode(out[0], skip_special_tokens=True))

        VLM example::

            from PIL import Image
            img = Image.open("photo.jpg")
            prompt = "<image>\\nDescribe this image."
            inputs = model.tokenizer(prompt, return_tensors="pt")
            out = model.generate(inputs["input_ids"], images=img, max_new_tokens=200)
            print(model.tokenizer.decode(out[0], skip_special_tokens=True))
        """
        input_embeds = None
        if images is not None and self._vision_encoder is not None:
            # Flush pending garbage and GPU fragmentation before image encoding
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            input_embeds = self._prepare_vlm_embeds(input_ids, images)

        return self._engine.generate(
            input_ids=input_ids,
            input_embeds=input_embeds,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            min_p=min_p,
            repetition_penalty=repetition_penalty,
            do_sample=do_sample,
            prefill_chunk_size=prefill_chunk_size,
            kv_cache_bits=kv_cache_bits,
            on_token=on_token,
        )

    def stream_generate(
        self,
        input_ids: torch.Tensor,
        images=None,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 0.9,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        do_sample: bool = True,
        prefill_chunk_size: int = 512,
        kv_cache_bits: int = 16,
    ) -> Iterator[tuple[int, torch.Tensor]]:
        """
        Streaming generation — yields each token as it is produced.

        Essential for disk-offloaded inference: each token can take 1–10
        seconds and users need to see output appear progressively.

        Args:
            input_ids: Input token IDs [batch, seq_len].
            images: Optional PIL image(s) or pixel_values tensor for VLMs.

        Yields:
            (token_id: int, full_sequence: torch.Tensor)

        Example::

            inputs = model.tokenizer("Hello!", return_tensors="pt")
            for token_id, seq in model.stream_generate(inputs["input_ids"]):
                word = model.tokenizer.decode([token_id],
                                              skip_special_tokens=True)
                print(word, end="", flush=True)
        """
        input_embeds = None
        if images is not None and self._vision_encoder is not None:
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            input_embeds = self._prepare_vlm_embeds(input_ids, images)

        return self._engine.stream_generate(
            input_ids=input_ids,
            input_embeds=input_embeds,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            min_p=min_p,
            repetition_penalty=repetition_penalty,
            do_sample=do_sample,
            prefill_chunk_size=prefill_chunk_size,
            kv_cache_bits=kv_cache_bits,
        )

    def _prepare_vlm_embeds(
        self,
        input_ids: torch.Tensor,
        images,
    ) -> torch.Tensor:
        """Encode images and merge into token embeddings for VLM inference.

        Args:
            input_ids: Token IDs [batch, seq_len] — must contain image_token_id
                       placeholder(s) that will be replaced by image patches.
            images: PIL Image, list[PIL Image], or pre-processed
                    [batch, 3, H, W] float tensor.

        Returns:
            Merged embeddings [batch, seq_len + (num_patches - 1) * num_images,
            hidden_size] ready to pass as input_embeds to the engine.
        """
        # Pylance guard: this method is only called when _vision_encoder is not None.
        assert self._vision_encoder is not None
        assert self._engine._embedding_weight is not None

        from dokodemo_ai.utils.image import get_vlm_image_processor_config, preprocess_images

        # --- 1. Preprocess images → pixel_values ---
        if isinstance(images, torch.Tensor):
            pixel_values = images.to(
                device=self._engine.device, dtype=self._engine.compute_dtype
            )
        else:
            proc_cfg = get_vlm_image_processor_config(self._graph.config)
            pixel_values = preprocess_images(
                images,
                size=proc_cfg["size"],
                mean=proc_cfg["mean"],
                std=proc_cfg["std"],
                device=self._engine.device,
                dtype=self._engine.compute_dtype,
            )

        # --- 2. Vision encoder → image patch embeddings ---
        # Shrink the layer-cache budget to make room for the vision encoder so
        # total GPU VRAM stays within max_gpu_memory during the encode pass.
        _cache = self._engine.cache
        _orig_budget = _cache.budget
        vis_bytes = sum(p.numel() * p.element_size() for p in self._vision_encoder.parameters())
        _reduced = max(_orig_budget - vis_bytes, 0)
        if _cache.current_bytes > _reduced:
            _cache.budget = _reduced
            _cache._evict_until(_reduced)  # free lowest-importance layers to make room

        # Restore vision encoder to GPU for this encode pass (lives on CPU otherwise)
        if next(self._vision_encoder.parameters()).device.type != self._engine.device.type:
            self._vision_encoder = self._vision_encoder.to(self._engine.device)
            if self._engine.device.type == "cuda":
                torch.cuda.empty_cache()

        image_embeds = self._vision_encoder.encode(pixel_values)
        del pixel_values  # free GPU memory immediately after encoding

        # Move vision encoder back to CPU; restore full cache budget for LLM forward pass
        self._vision_encoder = self._vision_encoder.cpu()
        if self._engine.device.type == "cuda":
            torch.cuda.empty_cache()
        _cache.budget = _orig_budget

        # --- 3. Text token → embeddings (CPU indexing — no full weight on GPU) ---
        text_embeds = self._engine._embedding_weight[input_ids.cpu()].to(
            self._engine.device, self._engine.compute_dtype
        )

        # --- 4. Merge image patches at <image> token positions ---
        merged = merge_image_tokens(
            input_ids=input_ids,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            image_token_id=self._vision_encoder.image_token_id,
        )
        del text_embeds, image_embeds  # release before the expensive LLM forward pass
        return merged

    def get_profiling_report(self) -> str:
        """Return profiling report (requires profiling=True)."""
        return self._profiler.report()

    def get_expert_stats(self) -> dict:
        """Expert usage statistics for MoE models."""
        return self._engine.scheduler.get_stats()

    @property
    def graph_summary(self) -> str:
        return self._graph.summary()

    def __repr__(self) -> str:
        return (
            f"AutoModel(\n"
            f"  path={self._graph.model_path},\n"
            f"  arch={'MoE' if self._graph.is_moe else 'Dense'},\n"
            f"  layers={self._graph.num_layers},\n"
            f"  device={self._engine.device},\n"
            f"  dtype={self._engine.compute_dtype},\n"
            f")"
        )
