import os import sys import math import json import time from typing import List, Dict, Tuple, Any, Literal import torch import torch.nn as nn import numpy as np from PIL import Image from easydict import EasyDict as edict # Add TRELLIS-main to python path BASE_DIR = os.path.dirname(os.path.abspath(__file__)) TRELLIS_DIR = os.path.join(BASE_DIR, "TRELLIS-main") if TRELLIS_DIR not in sys.path: sys.path.append(TRELLIS_DIR) # Force xformers attention backend for GPU stability os.environ['ATTN_BACKEND'] = 'xformers' os.environ['SPCONV_ALGO'] = 'native' try: from trellis.pipelines import TrellisImageTo3DPipeline from trellis.representations import Gaussian, MeshExtractResult from trellis.utils import render_utils, postprocessing_utils except ImportError: print("WARNING: TRELLIS packages could not be imported. Please run within venv_trellis.") class FormaAi(nn.Module): """ FormaAi: A unified hybrid 3D neural network class. Consolidates fast generation, discrete Gaussians, and continuous NeRF fields into a single model. """ def __init__(self, pretrained_path: str = "microsoft/TRELLIS-image-large", device: str = "cuda"): super().__init__() self.device = device # Load the TRELLIS pipeline internally self.pipeline = TrellisImageTo3DPipeline.from_pretrained(pretrained_path) # Override the pipeline device property to always return self.device (cuda) # to prevent model offloading from confusing the internal device detection. self.pipeline.__class__.device = property(lambda s: torch.device(self.device)) self.pipeline.to("cpu") # Expose individual sub-networks as nn.Module attributes/layers self.image_encoder = self.pipeline.models['image_cond_model'] self.sparse_flow_model = self.pipeline.models['sparse_structure_flow_model'] self.sparse_decoder = self.pipeline.models['sparse_structure_decoder'] self.slat_flow_model = self.pipeline.models['slat_flow_model'] self.decoder_mesh = self.pipeline.models['slat_decoder_mesh'] self.decoder_gs = self.pipeline.models['slat_decoder_gs'] self.decoder_rf = self.pipeline.models['slat_decoder_rf'] # Force CPU offloading state initially for model in self.pipeline.models.values(): model.to("cpu") torch.cuda.empty_cache() print("[FormaAi] All model components loaded and ready.") def preprocess(self, image: Image.Image) -> Image.Image: """ Preprocesses the input image (masking background, centering, resizing to 518x518). """ return self.pipeline.preprocess_image(image) @torch.no_grad() def forward( self, image: Image.Image, seed: int = 42, ss_steps: int = 12, ss_cfg: float = 7.5, slat_steps: int = 12, slat_cfg: float = 3.0, formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], preprocess: bool = True, refine_gs: bool = False, refine_steps: int = 100 ) -> Dict[str, Any]: """ Single unified forward pass of the hybrid model. Args: image (Image.Image): The input image prompt. seed (int): The random seed. ss_steps (int): Sampling steps for Stage 1 (Sparse Structure). ss_cfg (float): Guidance strength for Stage 1. slat_steps (int): Sampling steps for Stage 2 (Structured Latent). slat_cfg (float): Guidance strength for Stage 2. formats (List[str]): Desired output formats. preprocess (bool): Whether to run background removal. refine_gs (bool): Whether to run differentiable Gaussian optimization. refine_steps (int): Optimization steps for Gaussian refinement. Returns: Dict[str, Any]: Aligned representations (Mesh, Gaussian, Radiance Field). """ import gc gc.collect() torch.cuda.empty_cache() if preprocess: image = self.preprocess(image) with torch.no_grad(): # 1. Condition generation self.image_encoder.to(self.device) cond = self.pipeline.get_cond([image]) torch.manual_seed(seed) # 2. Sparse structure generation (Stage 1 flow matching pass) self.sparse_flow_model.to(self.device) self.sparse_decoder.to(self.device) coords = self.pipeline.sample_sparse_structure( cond, num_samples=1, sampler_params={"steps": ss_steps, "cfg_strength": ss_cfg} ) # Offload Stage 1 models to save memory self.image_encoder.to("cpu") self.sparse_flow_model.to("cpu") self.sparse_decoder.to("cpu") torch.cuda.empty_cache() # 3. Structured latent generation (Stage 2 flow matching pass) self.slat_flow_model.to(self.device) slat = self.pipeline.sample_slat( cond, coords, sampler_params={"steps": slat_steps, "cfg_strength": slat_cfg} ) # Offload Stage 2 model self.slat_flow_model.to("cpu") torch.cuda.empty_cache() # 4. Decoding structured latent into the representations raw_outputs = self.decode_slat_offloaded(slat, formats=formats) # 4b. Differentiable Gaussian Refinement if refine_gs and 'gaussian' in raw_outputs: gs = raw_outputs['gaussian'][0] refined_gs = self.refine_gaussians(gs, image, steps=refine_steps) raw_outputs['gaussian'] = [refined_gs] with torch.no_grad(): # 5. Hybrid coordinate alignment aligned_outputs = self.align_coordinates(raw_outputs) return aligned_outputs def align_coordinates(self, outputs: Dict[str, Any]) -> Dict[str, Any]: """ Keeps Mesh, Gaussian Splatting, and NeRF coordinates in their native coordinate systems to ensure perfect geometric alignment and texture projection. Calculates bounding box center and scale for JSON metadata only. """ aligned = {} center = torch.zeros(3, device=self.device) scale = 1.0 # Determine center and scale from Mesh (most stable geometric bounding box) if 'mesh' in outputs: mesh = outputs['mesh'][0] min_coords = mesh.vertices.min(dim=0).values max_coords = mesh.vertices.max(dim=0).values center = (min_coords + max_coords) / 2.0 extents = max_coords - min_coords max_extent = extents.max().item() scale = 1.0 / max_extent if max_extent > 0 else 1.0 aligned['mesh'] = edict( vertices=mesh.vertices, # Keep original coordinates faces=mesh.faces, raw=mesh ) aligned['center'] = center aligned['scale'] = scale if 'gaussian' in outputs: aligned['gaussian'] = outputs['gaussian'][0] # Keep original coordinates if 'radiance_field' in outputs: aligned['radiance_field'] = outputs['radiance_field'][0] return aligned def decode_slat_offloaded(self, slat: Any, formats: List[str]) -> Dict[str, Any]: """ Decodes the structured latent by sequentially moving each decoder to the GPU and offloading it, keeping GPU VRAM usage at a minimum. """ ret = {} # Ensure all decoders are initially on CPU self.decoder_mesh.to("cpu") self.decoder_gs.to("cpu") self.decoder_rf.to("cpu") torch.cuda.empty_cache() if 'mesh' in formats: print("[FormaAi] Offload: Decoding Mesh on GPU...") self.decoder_mesh.to(self.device) ret['mesh'] = self.decoder_mesh(slat) self.decoder_mesh.to("cpu") torch.cuda.empty_cache() if 'gaussian' in formats: print("[FormaAi] Offload: Decoding Gaussian Splatting on GPU...") self.decoder_gs.to(self.device) ret['gaussian'] = self.decoder_gs(slat) self.decoder_gs.to("cpu") torch.cuda.empty_cache() if 'radiance_field' in formats: print("[FormaAi] Offload: Decoding Radiance Field on GPU...") self.decoder_rf.to(self.device) ret['radiance_field'] = self.decoder_rf(slat) self.decoder_rf.to("cpu") torch.cuda.empty_cache() return ret def refine_gaussians(self, gs: Gaussian, target_image: Image.Image, steps: int = 100) -> Gaussian: """ Runs a quick optimization loop on the GPU to refine the generated Gaussian Splatting representation against the input image. """ with torch.enable_grad(): print(f"[FormaAi] Starting Differentiable Gaussian Refinement ({steps} steps)...") from trellis.utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics from trellis.renderers.gaussian_render import GaussianRenderer import torch.nn.functional as F start_time = time.time() # Prepare target image tensor target_img_resized = target_image.resize((512, 512), Image.Resampling.LANCZOS) target_tensor = torch.from_numpy(np.array(target_img_resized).astype(np.float32) / 255.0).permute(2, 0, 1).to(self.device)[:3] # Enable gradients on the Gaussian parameters gs._xyz = torch.nn.Parameter(gs._xyz.clone().detach().requires_grad_(True)) gs._rotation = torch.nn.Parameter(gs._rotation.clone().detach().requires_grad_(True)) gs._scaling = torch.nn.Parameter(gs._scaling.clone().detach().requires_grad_(True)) gs._opacity = torch.nn.Parameter(gs._opacity.clone().detach().requires_grad_(True)) # Setup optimizer optimizer = torch.optim.Adam([ {"params": gs._xyz, "lr": 1e-4}, {"params": gs._rotation, "lr": 1e-3}, {"params": gs._scaling, "lr": 5e-3}, {"params": gs._opacity, "lr": 0.025}, ], lr=1e-4) # Setup renderer renderer = GaussianRenderer({ "resolution": 512, "near": 0.8, "far": 1.6, "ssaa": 1, "bg_color": (0, 0, 0), }) # Camera pose corresponding to input image: yaw=0, pitch=0, radius=2.0, fov=40 extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics([0.0], [0.0], 2.0, 40.0) extr, intr = extrinsics[0], intrinsics[0] initial_loss = None for step in range(steps + 1): rendered = renderer.render(gs, extr, intr)['color'] loss = F.l1_loss(rendered, target_tensor) if step == 0: initial_loss = loss.item() if step < steps: optimizer.zero_grad() loss.backward() optimizer.step() if step % 20 == 0 and step > 0: print(f" [Refinement] Step {step}/{steps}: Loss = {loss.item():.5f}") final_loss = loss.item() elapsed = time.time() - start_time print(f"[FormaAi] Refinement finished. L1 Loss: {initial_loss:.5f} -> {final_loss:.5f} ({elapsed:.2f} seconds)") # Detach parameters to save memory/prevent graph retention gs._xyz = gs._xyz.data gs._rotation = gs._rotation.data gs._scaling = gs._scaling.data gs._opacity = gs._opacity.data return gs def upscale_texture(self, texture_image: Image.Image, scale_factor: int = 2, upscale_device: str = "auto") -> Image.Image: """ Upscales the texture image using Swin2SR. Temporarily offloads TRELLIS models to CPU to free GPU VRAM, runs upscaling, and then restores TRELLIS models to GPU. Supports GPU, CPU, or Auto-Fallback mode to prevent OOM. """ print(f"[FormaAi] Offloading TRELLIS models to CPU to free VRAM...") for model in self.pipeline.models.values(): model.to("cpu") torch.cuda.empty_cache() import gc gc.collect() from transformers import Swin2SRForImageSuperResolution, AutoImageProcessor if scale_factor == 2: model_id = "caidas/swin2SR-lightweight-x2-64" elif scale_factor == 4: model_id = "caidas/swin2SR-classical-sr-x4-64" else: raise ValueError(f"Unsupported upscale factor: {scale_factor}. Choose 2 or 4.") processor = AutoImageProcessor.from_pretrained(model_id) # Determine initial run device dev_str = upscale_device.lower() if "cpu" in dev_str: run_device = "cpu" else: run_device = self.device # default to cuda if available orig_w, orig_h = texture_image.size has_alpha = texture_image.mode == "RGBA" if has_alpha: rgb_image = texture_image.convert("RGB") alpha_image = texture_image.split()[3] else: rgb_image = texture_image print(f"[FormaAi] Upscaling mesh texture using Swin2SR x{scale_factor} on {run_device}...") try: model = Swin2SRForImageSuperResolution.from_pretrained(model_id).to(run_device) inputs = processor(rgb_image, return_tensors="pt").to(run_device) with torch.no_grad(): outputs = model(**inputs) except torch.OutOfMemoryError as e: if "auto" in dev_str or dev_str == "auto-fallback": print(f"[FormaAi WARNING] CUDA Out of Memory on {run_device}. Falling back to CPU...") # Clear GPU memory if 'model' in locals(): del model if 'inputs' in locals(): del inputs torch.cuda.empty_cache() gc.collect() run_device = "cpu" print(f"[FormaAi] Retrying Swin2SR upscaling on CPU...") model = Swin2SRForImageSuperResolution.from_pretrained(model_id).to(run_device) inputs = processor(rgb_image, return_tensors="pt").to(run_device) with torch.no_grad(): outputs = model(**inputs) else: raise e output_pixel_values = outputs.reconstruction output_pixel_values = output_pixel_values.squeeze().cpu().permute(1, 2, 0).numpy() output_pixel_values = np.clip(output_pixel_values * 255.0, 0, 255).astype(np.uint8) upscaled_rgb = Image.fromarray(output_pixel_values) target_w, target_h = orig_w * scale_factor, orig_h * scale_factor upscaled_rgb = upscaled_rgb.crop((0, 0, target_w, target_h)) if has_alpha: upscaled_alpha = alpha_image.resize((target_w, target_h), Image.Resampling.BICUBIC) upscaled_rgb.putalpha(upscaled_alpha) del model del inputs del outputs gc.collect() torch.cuda.empty_cache() print(f"[FormaAi] Restoring TRELLIS models to GPU ({self.device})...") for model in self.pipeline.models.values(): model.to(self.device) print(f"[FormaAi] Texture upscaled from {orig_w}x{orig_h} to {target_w}x{target_h}") return upscaled_rgb def save_hybrid_asset(self, aligned_outputs: Dict[str, Any], output_dir: str, prefix: str = "hybrid", upscale_factor: int = 2, upscale_device: str = "auto"): """ Saves all aligned 3D representations to a single hybrid asset directory. """ os.makedirs(output_dir, exist_ok=True) # Save Mesh as GLB if 'mesh' in aligned_outputs: print(f"[FormaAi] Saving aligned mesh to GLB...") glb = postprocessing_utils.to_glb( aligned_outputs['gaussian'], aligned_outputs['mesh']['raw'], # Use original for to_glb as it processes internally simplify=0.95, texture_size=1024 ) # AI Texture Upscaling if upscale_factor > 1 and hasattr(glb.visual, 'material') and hasattr(glb.visual.material, 'baseColorTexture'): orig_texture = glb.visual.material.baseColorTexture if orig_texture is not None: upscaled_texture = self.upscale_texture(orig_texture, scale_factor=upscale_factor, upscale_device=upscale_device) glb.visual.material.baseColorTexture = upscaled_texture glb_path = os.path.join(output_dir, f"{prefix}.glb") glb.export(glb_path) print(f"Mesh saved: {glb_path}") # Export OBJ + MTL + Textures print(f"[FormaAi] Saving aligned mesh to OBJ/MTL...") files_before = set(os.listdir(output_dir)) obj_path = os.path.join(output_dir, f"{prefix}.obj") glb.export(obj_path) files_after = set(os.listdir(output_dir)) new_files = files_after - files_before # Create a ZIP file of the OBJ model + material + textures import zipfile zip_path = os.path.join(output_dir, f"{prefix}_obj.zip") with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for file_name in new_files: file_path = os.path.join(output_dir, file_name) zipf.write(file_path, arcname=file_name) print(f"OBJ asset package saved: {zip_path}") # Save Gaussian Splatting as PLY if 'gaussian' in aligned_outputs: print(f"[FormaAi] Saving aligned 3D Gaussian splat...") ply_path = os.path.join(output_dir, f"{prefix}.ply") aligned_outputs['gaussian'].save_ply(ply_path) print(f"Gaussians saved: {ply_path}") # Save Alignment metadata JSON meta_path = os.path.join(output_dir, f"{prefix}_metadata.json") metadata = { "center": aligned_outputs['center'].cpu().numpy().tolist(), "scale": float(aligned_outputs['scale']) } with open(meta_path, "w") as f: json.dump(metadata, f, indent=4) print(f"Alignment metadata saved: {meta_path}") print("[FormaAi] All hybrid asset formats saved successfully.")