Blame for trellis2_quantized_inference.py

"""
TRELLIS 2 Quantized Inference — Run on RTX 3060 12GB (or any ≥8GB GPU)

Uses v3 GPTQ quantized weights (5.4GB) instead of FP16 (16.2GB).
Hooks into TRELLIS 2's existing low_vram sequential offloading —
zero changes to pipeline source code.

Two modes:
  1. Standalone (recommended): Uses pre-packaged quantized weights (~7.5GB).
     No FP16 weights needed. Peak CPU RAM ~8GB.
       python trellis2_quantized_inference.py --standalone-dir ./trellis2-quantized --image photo.png

  2. Legacy: Patches FP16 pipeline on-the-fly with quantized weights.
     Requires both FP16 (~16GB) and quantized (~5.4GB) weights on disk.
       python trellis2_quantized_inference.py --weights-dir /path/to/fp16 --quant-dir /path/to/quant --image photo.png

Requirements:
    - TRELLIS 2 installed (pip install trellis2 or from source)
    - Standalone package OR both FP16 + quantized weights
"""

import os, sys, gc, time, argparse, ctypes
from pathlib import Path

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

import torch
import torch.nn as nn
import torch.nn.functional as F


def _release_memory():
    """Force glibc to return freed pages to the OS (Linux only).

    Without this, Python/glibc hold freed heap pages in the process RSS
    even after gc.collect(), inflating apparent RAM usage by gigabytes.
    """
    try:
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except (OSError, AttributeError):
        pass  # macOS / non-glibc — no-op


# ── Config ───────────────────────────────────────────────────────────────────

FLOW_MODEL_NAMES = [
    'sparse_structure_flow_model',
    'shape_slat_flow_model_512',
    'shape_slat_flow_model_1024',
    'tex_slat_flow_model_512',
    'tex_slat_flow_model_1024',
]

# Map: pipeline model name → checkpoint basename (under ckpts/)
FLOW_MODEL_CKPTS = {
    'sparse_structure_flow_model': 'ss_flow_img_dit_1_3B_64_bf16',
    'shape_slat_flow_model_512': 'slat_flow_img2shape_dit_1_3B_512_bf16',
    'shape_slat_flow_model_1024': 'slat_flow_img2shape_dit_1_3B_1024_bf16',
    'tex_slat_flow_model_512': 'slat_flow_imgshape2tex_dit_1_3B_512_bf16',
    'tex_slat_flow_model_1024': 'slat_flow_imgshape2tex_dit_1_3B_1024_bf16',
}

DECODER_NAMES = [
    'sparse_structure_decoder',
    'shape_slat_decoder',
    'tex_slat_decoder',
]


# ── Vectorized GPU Dequantization ────────────────────────────────────────────

def dequantize_vectorized(entry: dict, device: torch.device,
                          out: torch.Tensor = None) -> torch.Tensor:
    """
    Dequantize v3 GPTQ weights on GPU using vectorized ops (no Python loops).

    ~30x faster than the reference loop-based implementation.

    If `out` is provided, writes into it (must be [rows, cols] FP16 on device).
    This avoids CUDA allocation churn in eager_free mode.
    """
    bits = entry['bits']
    rows, cols = entry['shape']
    group_size = entry['group_size']

    if bits == 16:
        W = entry['data'].to(device, non_blocking=True).float()
        awq = entry.get('awq_scales')
        if awq is not None:
            W.div_(awq.to(device, non_blocking=True).float().unsqueeze(0))
        result = W.half()
        if out is not None:
            out.copy_(result)
            del result
            return out
        return result

    # Transfer quantized data to GPU (pinned memory → fast)
    data_gpu = entry['data'].to(device, non_blocking=True)
    scales_gpu = entry['scales'].to(device, non_blocking=True).float()
    zeros_gpu = entry['zeros'].to(device, non_blocking=True).float()

    # Unpack integers → in-place sub/mul to reduce peak intermediates
    if bits == 4:
        flat_lo = (data_gpu & 0xF).to(torch.float32)
        flat_hi = ((data_gpu >> 4) & 0xF).to(torch.float32)
        W = torch.stack([flat_lo, flat_hi], dim=1).reshape(-1)[:rows * cols].reshape(rows, cols)
        del flat_lo, flat_hi
    elif bits == 8:
        W = data_gpu.float()
    else:
        raise ValueError(f"Unknown bits: {bits}")
    del data_gpu

    # Vectorized group dequant — in-place ops to reduce peak memory
    n_groups = scales_gpu.shape[1]
    col_idx = torch.arange(cols, device=device, dtype=torch.long) // group_size
    col_idx.clamp_(max=n_groups - 1)

    z = zeros_gpu[:, col_idx]     # [rows, cols]
    W.sub_(z)                     # in-place: W = W - z
    del z, zeros_gpu

    s = scales_gpu[:, col_idx]    # [rows, cols]
    W.mul_(s)                     # in-place: W = W * s
    del s, scales_gpu, col_idx

    # AWQ channel descaling (in-place)
    awq = entry.get('awq_scales')
    if awq is not None:
        W.div_(awq.to(device, non_blocking=True).float().unsqueeze(0))

    result = W.half()
    del W
    if out is not None:
        out.copy_(result)
        del result
        return out
    return result


# ── Weight Buffer Pool ──────────────────────────────────────────────────────

class WeightBufferPool:
    """
    Pre-allocated GPU buffer reused across eager_free dequantizations.

    Instead of 215 alloc/free cycles per forward step (CUDA fragmentation),
    each QuantizedLinear writes into a slice of this pool. Only the largest
    layer's worth of GPU memory is ever allocated.

    Usage: one pool per model, created when model moves to GPU, freed on CPU.
    """

    def __init__(self, device: torch.device, max_rows: int, max_cols: int):
        self.device = device
        self.buf = torch.empty(max_rows, max_cols, dtype=torch.float16, device=device)
        self._size_mb = self.buf.nelement() * 2 / 1e6

    def get_slice(self, rows: int, cols: int) -> torch.Tensor:
        """Return a view into the pool buffer for a [rows, cols] weight matrix."""
        return self.buf[:rows, :cols]

    def free(self):
        """Release the pool buffer."""
        del self.buf
        self.buf = None


# ── QuantizedLinear ──────────────────────────────────────────────────────────

class QuantizedLinear(nn.Module):
    """
    Drop-in nn.Linear replacement that stores weights in v3 GPTQ format.

    Weights are kept compressed on CPU. When activate() is called, they're
    dequantized on GPU and cached as FP16 for fast inference. When
    deactivate() is called, the FP16 cache is freed.

    This module is invisible to model.to() — quantized data stays on CPU,
    only bias (tiny) moves with the model. The QuantizedModelManager handles
    activation lifecycle.
    """

    def __init__(self, in_features: int, out_features: int,
                 bias_data: torch.Tensor, entry: dict):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features

        # Store quantized data as plain attributes (NOT buffers/parameters)
        # so model.to() doesn't try to move them
        self.q_entry = {}
        self.q_entry['bits'] = entry['bits']
        self.q_entry['shape'] = entry['shape']
        self.q_entry['group_size'] = entry['group_size']

        # Keep CPU tensors as-is (no pin_memory — saves ~5GB peak RAM).
        # For per-layer eager_free transfers (~12MB each), the speed
        # difference vs pinned is negligible (<2s total over full inference).
        self.q_entry['data'] = entry['data']
        self.q_entry['scales'] = entry['scales']
        self.q_entry['zeros'] = entry['zeros']
        self.q_entry['awq_scales'] = entry.get('awq_scales')

        # Bias is small — register as parameter so it moves with .to()
        if bias_data is not None:
            self.bias = nn.Parameter(bias_data.clone().half(), requires_grad=False)
        else:
            self.bias = None

        # FP16 weight cache (populated by activate, freed by deactivate)
        self._weight_cache = None
        self._using_pool = False
        self._pool_ref = None  # Set by manager when eager_free + pool is active
        # When True, dequant per-forward and free immediately (saves ~2.6GB VRAM)
        self.eager_free = False

    def activate(self, device: torch.device, pool: 'WeightBufferPool' = None):
        """Dequantize weights on GPU, cache as FP16."""
        if self._weight_cache is not None:
            return
        rows, cols = self.q_entry['shape']
        if pool is not None:
            # Write into pre-allocated pool buffer (zero CUDA alloc)
            out = pool.get_slice(rows, cols)
            self._weight_cache = dequantize_vectorized(self.q_entry, device, out=out)
            self._using_pool = True
        else:
            self._weight_cache = dequantize_vectorized(self.q_entry, device)
            self._using_pool = False

    def deactivate(self):
        """Free the FP16 cache (or just clear reference if using pool)."""
        self._weight_cache = None
        self._using_pool = False

    def forward(self, x) -> torch.Tensor:
        if self._weight_cache is None:
            # Auto-activate on the input's device (eager_free mode)
            device = x.feats.device if hasattr(x, 'feats') else x.device
            pool = getattr(self, '_pool_ref', None)
            self.activate(device, pool=pool)

        if hasattr(x, 'feats'):
            # SparseTensor: extract feats, compute, rewrap
            feats = x.feats
            w = self._weight_cache.to(feats.dtype)
            b = self.bias.to(feats.dtype) if self.bias is not None else None
            out_feats = F.linear(feats, w, b)
            result = type(x)(feats=out_feats, coords=x.coords)
        else:
            w = self._weight_cache.to(x.dtype)
            b = self.bias.to(x.dtype) if self.bias is not None else None
            result = F.linear(x, w, b)

        # In eager_free mode, release weight cache after each forward.
        # If using pool, just drop the reference (pool buffer stays allocated).
        # If not using pool, this frees the tensor entirely.
        if self.eager_free:
            self._weight_cache = None

        return result

    def extra_repr(self) -> str:
        bits = self.q_entry['bits']
        active = 'active' if self._weight_cache is not None else 'dormant'
        return f'in={self.in_features}, out={self.out_features}, bits={bits}, {active}'


# ── Module tree helpers ──────────────────────────────────────────────────────

def _get_parent_and_attr(model: nn.Module, layer_path: str):
    """Navigate module tree to find parent module and attribute name."""
    parts = layer_path.split('.')
    parent = model
    for part in parts[:-1]:
        if part.isdigit():
            parent = parent[int(part)]
        else:
            parent = getattr(parent, part)
    return parent, parts[-1]


def _parse_device(args, kwargs):
    """Extract device from .to() call arguments."""
    if 'device' in kwargs:
        return kwargs['device']
    for arg in args:
        if isinstance(arg, (str, torch.device)):
            return str(arg)
        if isinstance(arg, torch.dtype):
            continue
    return None


# ── QuantizedModelManager ────────────────────────────────────────────────────

class QuantizedModelManager:
    """
    Manages the quantized weight lifecycle for a single flow model.

    Hooks model.to() so that the pipeline's existing low_vram offloading
    automatically triggers dequant (to CUDA) and cache cleanup (to CPU).

    When activation_checkpointing=True, enables gradient checkpointing on
    transformer blocks to halve peak activation VRAM (~30% slower).
    """

    def __init__(self, model: nn.Module, model_name: str,
                 activation_checkpointing: bool = False,
                 eager_free: bool = False):
        self.model = model
        self.model_name = model_name
        self.quantized_layers = {}
        self._original_to = None
        self.activation_checkpointing = activation_checkpointing
        self.eager_free = eager_free
        self._weight_pool = None  # Created on-demand when model moves to GPU

    def patch(self, quant_dict: dict):
        """Replace nn.Linear layers with QuantizedLinear, hook .to()."""
        replaced = 0
        skipped = 0

        for layer_path, module in list(self.model.named_modules()):
            if not isinstance(module, nn.Linear):
                continue
            if layer_path not in quant_dict:
                skipped += 1
                continue

            entry = quant_dict[layer_path]
            bias_data = module.bias.data if module.bias is not None else None

            q_linear = QuantizedLinear(
                module.in_features, module.out_features,
                bias_data, entry,
            )

            # Explicitly free the original FP16 weight tensor
            module.weight = None
            if module.bias is not None:
                module.bias = None

            parent, attr = _get_parent_and_attr(self.model, layer_path)
            setattr(parent, attr, q_linear)
            self.quantized_layers[layer_path] = q_linear
            replaced += 1

        # Force garbage collection to free old FP16 weight memory
        gc.collect()

        # Hook .to() AND .cpu() for transparent activate/deactivate
        # Pipeline's low_vram code calls model.to(device) and model.cpu()
        self._original_to = self.model.to
        self._original_cpu = self.model.cpu
        manager = self

        def hooked_to(*args, **kwargs):
            device = _parse_device(args, kwargs)

            if device is not None and 'cuda' in str(device):
                extras = []
                gpu_device = torch.device(device)

                if manager.eager_free:
                    # Create shared weight buffer pool — sized for the largest layer.
                    # All layers dequant into this one buffer instead of allocating their own.
                    max_rows = max(ql.q_entry['shape'][0] for ql in manager.quantized_layers.values())
                    max_cols = max(ql.q_entry['shape'][1] for ql in manager.quantized_layers.values())
                    manager._weight_pool = WeightBufferPool(gpu_device, max_rows, max_cols)
                    for ql in manager.quantized_layers.values():
                        ql.eager_free = True
                        ql._pool_ref = manager._weight_pool
                    extras.append(f"eager_free + pool ({manager._weight_pool._size_mb:.0f}MB, {len(manager.quantized_layers)} layers)")
                else:
                    # Standard mode: pre-dequant all layers onto GPU
                    t0 = time.time()
                    for ql in manager.quantized_layers.values():
                        ql.activate(gpu_device)
                    dt = time.time() - t0
                    extras.append(f"dequantized {len(manager.quantized_layers)} layers in {dt:.2f}s")

                # Enable activation checkpointing on transformer blocks
                if manager.activation_checkpointing and hasattr(manager.model, 'blocks'):
                    n_blocks = 0
                    for block in manager.model.blocks:
                        if hasattr(block, 'use_checkpoint'):
                            block.use_checkpoint = True
                            n_blocks += 1
                    if n_blocks:
                        extras.append(f"checkpointing {n_blocks} blocks")

                print(f"    [{manager.model_name}] {', '.join(extras)}")

            # Move non-quantized params (bias, norms, embeddings)
            result = manager._original_to(*args, **kwargs)

            if device is not None and 'cpu' in str(device):
                # Free FP16 caches, pool, and disable eager_free
                for ql in manager.quantized_layers.values():
                    ql.eager_free = False
                    ql._pool_ref = None
                    ql.deactivate()
                if manager._weight_pool is not None:
                    manager._weight_pool.free()
                    manager._weight_pool = None
                # Disable checkpointing when offloading
                if hasattr(manager.model, 'blocks'):
                    for block in manager.model.blocks:
                        if hasattr(block, 'use_checkpoint'):
                            block.use_checkpoint = False
                torch.cuda.empty_cache()

            return result

        def hooked_cpu(*args, **kwargs):
            # Free FP16 caches, pool, and disable eager_free
            for ql in manager.quantized_layers.values():
                ql.eager_free = False
                ql._pool_ref = None
                ql.deactivate()
            if manager._weight_pool is not None:
                manager._weight_pool.free()
                manager._weight_pool = None
            # Disable checkpointing
            if hasattr(manager.model, 'blocks'):
                for block in manager.model.blocks:
                    if hasattr(block, 'use_checkpoint'):
                        block.use_checkpoint = False
            # Move non-quantized params to CPU
            result = manager._original_cpu(*args, **kwargs)
            torch.cuda.empty_cache()
            return result

        self.model.to = hooked_to
        self.model.cpu = hooked_cpu

        return replaced, skipped

    def memory_footprint_mb(self) -> float:
        """Estimate CPU memory used by quantized data."""
        total = 0
        for ql in self.quantized_layers.values():
            e = ql.q_entry
            total += e['data'].nelement() * e['data'].element_size()
            total += e['scales'].nelement() * e['scales'].element_size()
            total += e['zeros'].nelement() * e['zeros'].element_size()
            if e['awq_scales'] is not None:
                total += e['awq_scales'].nelement() * e['awq_scales'].element_size()
        return total / 1e6


# ── Pipeline loading ─────────────────────────────────────────────────────────

def load_quantized_pipeline(weights_dir: str, quant_dir: str, device: str = 'cuda:0',
                            activation_checkpointing: bool = False):
    """
    Load TRELLIS 2 pipeline with v3 quantized flow models.

    1. Loads standard FP16 pipeline (low_vram=True, models on CPU)
    2. Patches each flow model with QuantizedLinear (frees FP16 weights)
    3. Returns pipeline ready for inference
    """
    # Add trellis2 source to path
    # Try common install locations
    for candidate in [
        '/home/ubuntu/trellis2',
        str(Path(weights_dir).parent / 'trellis2'),
        str(Path(weights_dir).parent),
    ]:
        if Path(candidate).is_dir() and candidate not in sys.path:
            sys.path.insert(0, candidate)
    from trellis2.pipelines import Trellis2ImageTo3DPipeline

    print("=" * 60)
    print("TRELLIS 2 — Quantized Inference (v3 GPTQ)")
    print(f"  weights : {weights_dir}")
    print(f"  quant   : {quant_dir}")
    print(f"  device  : {device}")
    print("=" * 60)

    # Load pipeline — low_vram=True by default, all models stay on CPU
    print("\nLoading pipeline (FP16, low_vram=True)...")
    t0 = time.time()
    pipeline = Trellis2ImageTo3DPipeline.from_pretrained(weights_dir)
    pipeline.to(device)  # sets _device but doesn't move models (low_vram)
    print(f"  Pipeline loaded in {time.time()-t0:.1f}s")

    # Patch flow models one at a time, freeing FP16 weights incrementally
    print("\nPatching flow models with quantized weights...")
    managers = {}
    total_quant_mb = 0

    for model_name in FLOW_MODEL_NAMES:
        quant_path = Path(quant_dir) / f'{model_name}.pt'

        if not quant_path.exists():
            print(f"  {model_name}: no quantized weights, keeping FP16")
            continue
        if model_name not in pipeline.models:
            print(f"  {model_name}: not in pipeline, skipping")
            continue

        t0 = time.time()
        quant_dict = torch.load(quant_path, map_location='cpu', weights_only=False)

        mgr = QuantizedModelManager(pipeline.models[model_name], model_name,
                                     activation_checkpointing=activation_checkpointing,
                                     eager_free=activation_checkpointing)
        replaced, skipped = mgr.patch(quant_dict)
        managers[model_name] = mgr

        quant_mb = mgr.memory_footprint_mb()
        total_quant_mb += quant_mb

        del quant_dict
        gc.collect()
        _release_memory()

        dt = time.time() - t0
        print(f"  {model_name}: {replaced} quantized, {skipped} FP16 — {quant_mb:.0f} MB — {dt:.1f}s")

    print(f"\nTotal quantized weight memory: {total_quant_mb:.0f} MB")
    print(f"Models patched: {len(managers)}/{len(FLOW_MODEL_NAMES)}")
    if activation_checkpointing:
        print(f"Low-VRAM mode: eager_free + activation checkpointing (saves ~3GB VRAM, ~30% slower)")
    print()

    return pipeline, managers


# ── Standalone pipeline loading ─────────────────────────────────────────

def _load_flow_model_quantized(standalone_dir: str, model_name: str,
                                trellis2_models,
                                activation_checkpointing: bool = False) -> tuple:
    """
    Load a single flow model from the standalone package.

    Memory-optimized flow (peak ~1.2GB instead of ~3.6GB per model):
    1. Create model skeleton from JSON config (~2.5GB random weights)
    2. Immediately free all nn.Linear .weight tensors (~2.4GB freed)
    3. Load non-linear params (norms + biases) into the skeleton (~2.5MB)
    4. Load quant_dict (~1.1GB) — only now, after random weights are freed
    5. Patch with QuantizedLinear (stores refs to quant_dict tensors)
    """
    ckpt_basename = FLOW_MODEL_CKPTS[model_name]
    config_path = Path(standalone_dir) / 'ckpts' / f'{ckpt_basename}.json'
    quant_path = Path(standalone_dir) / 'quantized' / f'{model_name}.pt'

    # Read architecture config
    import json
    with open(config_path) as f:
        config = json.load(f)

    # Create model architecture (~2.5GB of random BF16 weights)
    model = trellis2_models.__getattr__(config['name'])(**config['args'])
    model.eval()

    # Immediately free nn.Linear weight tensors — we'll replace them with
    # QuantizedLinear anyway. Keeps only bias + non-Linear params (~100MB).
    # This must happen BEFORE loading quant_dict to avoid the 2.5GB+1.1GB peak.
    freed_bytes = 0
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.weight is not None:
            freed_bytes += module.weight.nelement() * module.weight.element_size()
            module.weight = None
    gc.collect()
    _release_memory()

    # Load non-linear params (norms + biases, ~2.5MB) into the now-lightweight skeleton
    quant_dict = torch.load(quant_path, map_location='cpu', weights_only=False)
    non_linear_params = quant_dict.pop('_non_linear_params', {})

    if non_linear_params:
        missing, unexpected = model.load_state_dict(non_linear_params, strict=False)
        if unexpected:
            print(f"    WARNING: unexpected keys in non-linear params: {unexpected[:5]}")
    del non_linear_params

    # Patch nn.Linear with QuantizedLinear (copies correct biases, stores quant refs)
    mgr = QuantizedModelManager(model, model_name,
                                 activation_checkpointing=activation_checkpointing,
                                 eager_free=activation_checkpointing)
    replaced, skipped = mgr.patch(quant_dict)

    del quant_dict
    gc.collect()
    _release_memory()

    return model, mgr, replaced, skipped


def load_standalone_pipeline(standalone_dir: str, device: str = 'cuda:0',
                             activation_checkpointing: bool = False):
    """
    Load TRELLIS 2 pipeline from standalone quantized package.

    Unlike load_quantized_pipeline(), this NEVER loads FP16 flow model
    safetensors. Peak CPU RAM is ~8GB instead of ~16GB.

    1. Loads pipeline.json for sampler/normalization config
    2. Loads decoders from FP16 safetensors (~2GB total)
    3. Creates flow model architectures, patches with quantized weights
    4. Assembles pipeline with all config
    """
    import json

    # Add trellis2 source to path
    for candidate in [
        '/home/ubuntu/trellis2',
        str(Path(standalone_dir).parent / 'trellis2'),
    ]:
        if Path(candidate).is_dir() and candidate not in sys.path:
            sys.path.insert(0, candidate)

    from trellis2.pipelines import Trellis2ImageTo3DPipeline
    from trellis2.pipelines import samplers, rembg
    from trellis2.modules import image_feature_extractor
    from trellis2 import models as trellis2_models

    print("=" * 60)
    print("TRELLIS 2 — Standalone Quantized Inference (v3 GPTQ)")
    print(f"  package : {standalone_dir}")
    print(f"  device  : {device}")
    print("=" * 60)

    # Read pipeline config
    with open(Path(standalone_dir) / 'pipeline.json') as f:
        args = json.load(f)['args']

    _models = {}
    managers = {}
    total_quant_mb = 0

    # ── Load flow models FIRST (quantized, peak RAM happens here) ──────
    # Decoders are deferred so they don't inflate peak RAM during the
    # flow model loading phase (~2GB overhead avoided at peak).
    print("\nLoading flow models (quantized)...")
    for model_name in FLOW_MODEL_NAMES:
        quant_path = Path(standalone_dir) / 'quantized' / f'{model_name}.pt'
        if not quant_path.exists():
            print(f"  {model_name}: no quantized weights, skipping")
            continue

        t0 = time.time()
        model, mgr, replaced, skipped = _load_flow_model_quantized(
            standalone_dir, model_name, trellis2_models,
            activation_checkpointing=activation_checkpointing,
        )
        _models[model_name] = model
        managers[model_name] = mgr

        quant_mb = mgr.memory_footprint_mb()
        total_quant_mb += quant_mb
        dt = time.time() - t0
        ram_gb = _get_process_rss_gb()
        print(f"  {model_name}: {replaced} quantized, {skipped} FP16 — {quant_mb:.0f} MB — {dt:.1f}s — RSS {ram_gb:.1f}GB")

    print(f"\nTotal quantized weight memory: {total_quant_mb:.0f} MB")
    if activation_checkpointing:
        print(f"Low-VRAM mode: eager_free + activation checkpointing (saves ~3GB VRAM, ~5% slower)")

    # ── Load decoders (FP16, ~2GB total) ───────────────────────────────
    print("\nLoading decoders (FP16)...")
    for model_name in DECODER_NAMES:
        ckpt_rel = args['models'][model_name]
        ckpt_path = str(Path(standalone_dir) / ckpt_rel)
        t0 = time.time()
        _models[model_name] = trellis2_models.from_pretrained(ckpt_path)
        _models[model_name].eval()
        dt = time.time() - t0
        print(f"  {model_name}: loaded in {dt:.1f}s")

    # ── Assemble pipeline ───────────────────────────────────────────────
    print("\nAssembling pipeline...")

    # Create samplers
    ss_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(
        **args['sparse_structure_sampler']['args']
    )
    shape_sampler = getattr(samplers, args['shape_slat_sampler']['name'])(
        **args['shape_slat_sampler']['args']
    )
    tex_sampler = getattr(samplers, args['tex_slat_sampler']['name'])(
        **args['tex_slat_sampler']['args']
    )

    # Create image conditioning and rembg
    image_cond = getattr(image_feature_extractor, args['image_cond_model']['name'])(
        **args['image_cond_model']['args']
    )
    rembg_model = getattr(rembg, args['rembg_model']['name'])(
        **args['rembg_model']['args']
    )

    # Construct pipeline
    pipeline = Trellis2ImageTo3DPipeline(
        models=_models,
        sparse_structure_sampler=ss_sampler,
        shape_slat_sampler=shape_sampler,
        tex_slat_sampler=tex_sampler,
        sparse_structure_sampler_params=args['sparse_structure_sampler']['params'],
        shape_slat_sampler_params=args['shape_slat_sampler']['params'],
        tex_slat_sampler_params=args['tex_slat_sampler']['params'],
        shape_slat_normalization=args['shape_slat_normalization'],
        tex_slat_normalization=args['tex_slat_normalization'],
        image_cond_model=image_cond,
        rembg_model=rembg_model,
        low_vram=True,
        default_pipeline_type=args.get('default_pipeline_type', '1024_cascade'),
    )
    pipeline._device = 'cpu'
    pipeline.to(device)  # sets _device, doesn't move models (low_vram)

    print(f"Pipeline assembled. Models: {len(_models)}")
    print()

    return pipeline, managers


# ── Inference ────────────────────────────────────────────────────────────────

def run_inference(pipeline, image_path: str, output_path: str,
                  pipeline_type: str = '1024_cascade', steps: int = 20,
                  seed: int = 42, render_preview: bool = False):
    """Run TRELLIS 2 inference and export mesh."""
    from PIL import Image

    print(f"Input: {image_path}")
    print(f"Pipeline: {pipeline_type}, steps={steps}, seed={seed}")

    img = Image.open(image_path).convert('RGB')
    t0 = time.time()

    with torch.inference_mode():
        outputs = pipeline.run(
            img,
            seed=seed,
            return_latent=False,
            pipeline_type=pipeline_type,
            sparse_structure_sampler_params={"steps": steps},
            shape_slat_sampler_params={"steps": steps},
            tex_slat_sampler_params={"steps": steps},
        )

    dt = time.time() - t0
    print(f"\nInference complete in {dt:.1f}s ({dt/60:.1f}m)")

    # Export mesh (supported formats: .ply, .npz, .vxz)
    mesh = outputs[0]
    output_ext = Path(output_path).suffix.lower()

    if output_ext not in ('.ply', '.npz', '.vxz'):
        output_path = str(Path(output_path).with_suffix('.ply'))
        print(f"  (format changed to .ply — supported: .ply, .npz, .vxz)")

    mesh.save(output_path)
    print(f"Mesh saved: {output_path}")

    # Optional rendered preview
    if render_preview:
        try:
            from trellis2.utils import render_utils
            from trellis2.renderers.pbr_mesh_renderer import EnvMap
            import cv2

            preview_path = str(Path(output_path).with_suffix('.png'))
            mesh_for_render = outputs[0]

            # Build envmap for PBR rendering
            import trellis2
            trellis2_root = Path(trellis2.__file__).parent.parent
            envmap = None

            # Try loading HDRI from trellis2 install
            for hdri_name in ['studio.exr', 'courtyard.exr', 'forest.exr']:
                hdri_path = trellis2_root / 'assets' / 'hdri' / hdri_name
                if not hdri_path.exists():
                    continue
                try:
                    # Try OpenCV first (needs OpenEXR support)
                    import cv2
                    img = cv2.imread(str(hdri_path), cv2.IMREAD_UNCHANGED)
                    if img is not None:
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                        envmap = EnvMap(torch.tensor(img, dtype=torch.float32, device=pipeline.device))
                        break
                except ImportError:
                    pass
                try:
                    # Fallback: imageio can read EXR
                    import imageio
                    img = imageio.imread(str(hdri_path))
                    envmap = EnvMap(torch.tensor(img, dtype=torch.float32, device=pipeline.device))
                    break
                except (ImportError, Exception):
                    pass

            if envmap is None:
                # Synthetic neutral studio lighting (no file dependency)
                # Gradient from warm top to cool bottom simulates sky+ground
                h, w = 64, 128
                env = torch.zeros(h, w, 3, dtype=torch.float32, device=pipeline.device)
                for i in range(h):
                    t = i / (h - 1)  # 0=top, 1=bottom
                    # Top: warm sky light, bottom: cool ground bounce
                    env[i, :, 0] = 0.9 * (1 - t) + 0.3 * t  # R
                    env[i, :, 1] = 0.85 * (1 - t) + 0.3 * t  # G
                    env[i, :, 2] = 0.8 * (1 - t) + 0.35 * t  # B
                envmap = EnvMap(env)

            snap = render_utils.render_snapshot(
                mesh_for_render, resolution=512, nviews=4,
                bg_color=(1, 1, 1), envmap=envmap,
            )
            from PIL import Image as PILImage
            import numpy as np
            renders = snap['shaded']
            grid = np.concatenate(renders[:4], axis=1)
            PILImage.fromarray(grid).save(preview_path)
            print(f"Preview saved: {preview_path}")
        except Exception as e:
            print(f"Preview rendering failed: {e}")

    return outputs


# ── Memory monitoring ────────────────────────────────────────────────────────

def _get_process_rss_gb() -> float:
    """Get current process RSS in GB (what the OS actually allocated to us)."""
    try:
        import psutil
        return psutil.Process().memory_info().rss / 1e9
    except ImportError:
        try:
            # Fallback: read /proc on Linux
            with open('/proc/self/status') as f:
                for line in f:
                    if line.startswith('VmRSS:'):
                        return int(line.split()[1]) / 1e6  # KB → GB
        except (FileNotFoundError, ValueError):
            pass
    return 0.0


def print_memory_status():
    """Print current CPU and GPU memory usage."""
    try:
        import psutil
        ram = psutil.virtual_memory()
        rss = psutil.Process().memory_info().rss / 1e9
        print(f"  CPU RAM: {ram.used/1e9:.1f}/{ram.total/1e9:.1f} GB ({ram.percent}%) — process RSS: {rss:.1f}GB")
    except ImportError:
        pass

    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1e9
        reserved = torch.cuda.memory_reserved() / 1e9
        total = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"  GPU VRAM: {allocated:.1f}/{total:.1f} GB allocated ({reserved:.1f} reserved)")


# ── CLI ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description='TRELLIS 2 Quantized Inference — Image to 3D on consumer GPUs',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Speed modes (auto-detected from GPU, or set explicitly):
  --fast       Pre-dequant all weights onto GPU. Fastest inference.
               Requires ≥16GB VRAM. Best for RTX 4070+/4080/4090.
  --low-vram   Dequant per-layer on the fly + activation checkpointing.
               Peak ~4GB VRAM. Works on any ≥8GB GPU including RTX 3060.
  (default)    Auto-detects: ≤13GB GPU → --low-vram, else → --fast.

Examples:
  # RTX 3060 12GB (auto low-vram):
  python trellis2_quantized_inference.py --standalone-dir ./trellis2-quantized --image photo.png

  # RTX 4090 24GB (auto fast):
  python trellis2_quantized_inference.py --standalone-dir ./trellis2-quantized --image photo.png

  # Force fast on a 16GB GPU:
  python trellis2_quantized_inference.py --standalone-dir ./trellis2-quantized --image photo.png --fast

  # Lower resolution for faster results:
  python trellis2_quantized_inference.py --standalone-dir ./trellis2-quantized --image photo.png --pipeline-type 512 --steps 12
        """,
    )
    parser.add_argument('--image', required=True, help='Input image path')
    parser.add_argument('--output', default='output.ply', help='Output mesh path (.ply, .npz, .vxz)')
    parser.add_argument('--standalone-dir', default=None,
                        help='Path to standalone quantized package (recommended)')
    parser.add_argument('--weights-dir', default='/data/weights/trellis2',
                        help='Path to TRELLIS 2 FP16 weights (legacy mode)')
    parser.add_argument('--quant-dir', default='/data/quantized-blocks-v3',
                        help='Path to v3 quantized weight files (legacy mode)')
    parser.add_argument('--pipeline-type', default='1024_cascade',
                        choices=['512', '1024', '1024_cascade'],
                        help='Pipeline resolution mode')
    parser.add_argument('--steps', type=int, default=20,
                        help='Denoising steps per stage')
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed for reproducibility')
    parser.add_argument('--render-preview', action='store_true',
                        help='Save a 4-view rendered preview image')

    speed_group = parser.add_mutually_exclusive_group()
    speed_group.add_argument('--fast', action='store_true',
                             help='Force fast mode: pre-dequant all weights onto GPU '
                                  '(needs ≥16GB VRAM, fastest inference)')
    speed_group.add_argument('--low-vram', action='store_true',
                             help='Force low-VRAM mode: per-layer dequant + checkpointing '
                                  '(works on 8-12GB GPUs, ~5%% slower)')
    parser.add_argument('--device', default='cuda:0',
                        help='Target GPU device')
    args = parser.parse_args()

    if not Path(args.image).exists():
        print(f"Error: Image not found: {args.image}")
        sys.exit(1)

    # Resolve speed mode: explicit flag → auto-detect from GPU
    if args.fast:
        low_vram = False
        print("Speed mode: --fast (pre-dequant all weights onto GPU)")
    elif args.low_vram:
        low_vram = True
        print("Speed mode: --low-vram (per-layer dequant + checkpointing)")
    elif torch.cuda.is_available():
        gpu_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
        gpu_name = torch.cuda.get_device_name(0)
        if gpu_mem_gb <= 13:
            low_vram = True
            print(f"GPU: {gpu_name} ({gpu_mem_gb:.0f}GB) → auto --low-vram")
        else:
            low_vram = False
            print(f"GPU: {gpu_name} ({gpu_mem_gb:.0f}GB) → auto --fast")
    else:
        low_vram = True
        print("No CUDA GPU detected — using --low-vram mode")

    print_memory_status()

    if args.standalone_dir is not None:
        pipeline, managers = load_standalone_pipeline(
            args.standalone_dir, args.device,
            activation_checkpointing=low_vram,
        )
    else:
        pipeline, managers = load_quantized_pipeline(
            args.weights_dir, args.quant_dir, args.device,
            activation_checkpointing=low_vram,
        )

    print("After loading:")
    print_memory_status()

    run_inference(
        pipeline, args.image, args.output,
        pipeline_type=args.pipeline_type,
        steps=args.steps,
        seed=args.seed,
        render_preview=args.render_preview,
    )

    print("\nFinal memory:")
    print_memory_status()
    print("\nDone.")


if __name__ == '__main__':
    main()