""" 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()