diff --git a/nodes.py b/nodes.py index 35a8bda..9c57cba 100644 --- a/nodes.py +++ b/nodes.py @@ -493,6 +493,11 @@ def process(self, modelname, backend, device, low_vram, keep_models_loaded, conv isPixal3D = True pipeline = Trellis2ImageTo3DPipeline.from_pretrained(model_path, keep_models_loaded = keep_models_loaded, use_fp8=use_fp8, use_reconviagen=use_reconviagen, isPixal3D = isPixal3D) + # Pixal3D NAFF model uses NATTEN for feature upsampling. + # Windows prebuilt NATTEN wheels (torch 2.7.0, cu128, cp312): + # https://hf-mirror.com/lldacing/NATTEN-windows + # For issues/questions about Pixal3D Windows setup: + # https://space.bilibili.com/37411464 pipeline.low_vram = low_vram # if naf_chunk_size == "None": @@ -3857,6 +3862,9 @@ def INPUT_TYPES(s): "image": ("IMAGE",), "max_views": ("INT", {"default": 1, "min": 1, "max": 999}), }, + "optional": { + "naf_target_size": ("INT", {"default": 256, "min": 128, "max": 1024, "step": 64}), + }, } RETURN_TYPES = ("IMAGE_COND", "IMAGE_COND", "TRELLIS2PIPELINE", "MOGE_CAM_CONFIG") @@ -3865,9 +3873,12 @@ def INPUT_TYPES(s): CATEGORY = "Trellis2Wrapper" OUTPUT_NODE = True - def process(self, pipeline, image, max_views,): + def process(self, pipeline, image, max_views, naf_target_size=256): images = tensor_batch_to_pil_list(image, max_views=max_views) image_in = images[0] if len(images) == 1 else images + + if pipeline.isPixal3D and hasattr(pipeline, 'naf_target_size_override'): + pipeline.naf_target_size_override = naf_target_size if isinstance(image_in, (list, tuple)): images = list(image_in) diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 61e9c07..7203137 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -106,6 +106,7 @@ def __init__( self.rembg_model = rembg_model self._low_vram = low_vram self.default_pipeline_type = default_pipeline_type + self.naf_target_size_override = None self.VGGT_model = None self.pbr_attr_layout = { 'base_color': slice(0, 3), @@ -143,6 +144,12 @@ def __init__( "naf_target_size": 1024, }, } + + def _get_pixal3d_config(self, key): + config = self.PIXAL3D_IMAGE_COND_CONFIGS[key].copy() + if self.naf_target_size_override is not None and 'naf_target_size' in config: + config['naf_target_size'] = self.naf_target_size_override + return config def switch_samplers(self, sampler_type: str = "euler"): """Dynamically switches the sampler instances based on user selection.""" @@ -292,7 +299,7 @@ def load_pixal3d_image_cond_ss(self): return self.pixal3d_image_cond_ss print('Loading Pixal3D Image Cond SS Model ...') - model = build_pixal3d_image_cond_model(self.PIXAL3D_IMAGE_COND_CONFIGS["ss"]) + model = build_pixal3d_image_cond_model(self._get_pixal3d_config("ss")) self.pixal3d_image_cond_ss = model return model @@ -307,7 +314,7 @@ def load_pixal3d_image_cond_shape_512(self): return self.pixal3d_image_cond_shape_512 print('Loading Pixal3D Image Cond Shape 512 Model ...') - model = build_pixal3d_image_cond_model(self.PIXAL3D_IMAGE_COND_CONFIGS["shape_512"]) + model = build_pixal3d_image_cond_model(self._get_pixal3d_config("shape_512")) self.pixal3d_image_cond_shape_512 = model return model @@ -322,7 +329,7 @@ def load_pixal3d_image_cond_shape_1024(self): return self.pixal3d_image_cond_shape_1024 print('Loading Pixal3D Image Cond Shape 1024 Model ...') - model = build_pixal3d_image_cond_model(self.PIXAL3D_IMAGE_COND_CONFIGS["shape_1024"]) + model = build_pixal3d_image_cond_model(self._get_pixal3d_config("shape_1024")) self.pixal3d_image_cond_shape_1024 = model return model @@ -337,7 +344,7 @@ def load_pixal3d_image_cond_tex_1024(self): return self.pixal3d_image_cond_tex_1024 print('Loading Pixal3D Image Cond Tex 1024 Model ...') - model = build_pixal3d_image_cond_model(self.PIXAL3D_IMAGE_COND_CONFIGS["tex_1024"]) + model = build_pixal3d_image_cond_model(self._get_pixal3d_config("tex_1024")) self.pixal3d_image_cond_tex_1024 = model return model @@ -560,9 +567,9 @@ def get_proj_cond_ss( dict with 'cond' and 'neg_cond', each containing {'global': ..., 'proj': ...} """ print('Getting Proj Image Cond ...') - device = self.device - #image_cond_model = self.image_cond_model - if self.low_vram: + device = self.device + was_on_cpu = next(image_cond_model.parameters()).device.type == 'cpu' + if was_on_cpu: image_cond_model.to(device) cam_angle = torch.tensor([camera_angle_x], device=device) dist_tensor = torch.tensor([distance], device=device) @@ -570,7 +577,7 @@ def get_proj_cond_ss( z_global, z_proj = image_cond_model( image, camera_angle_x=cam_angle, distance=dist_tensor, mesh_scale=scale_tensor, ) - if self.low_vram: + if self.low_vram and was_on_cpu: image_cond_model.cpu() return { 'cond': {'global': z_global, 'proj': z_proj}, @@ -605,7 +612,8 @@ def get_proj_cond_shape( """ print('Getting Projected Image Cond ...') device = self.device - if self.low_vram: + was_on_cpu = next(image_cond_model.parameters()).device.type == 'cpu' + if was_on_cpu: image_cond_model.to(device) orig_grid_res = image_cond_model.grid_resolution @@ -639,7 +647,7 @@ def get_proj_cond_shape( image_resolution=image_cond_model.proj_grid.image_resolution, ).to(device) - if self.low_vram: + if self.low_vram and was_on_cpu: image_cond_model.cpu() return { 'cond': {'global': z_global, 'proj': z_proj_st}, diff --git a/trellis2/trainers/flow_matching/mixins/image_conditioned_proj.py b/trellis2/trainers/flow_matching/mixins/image_conditioned_proj.py index bb7fa84..776524c 100644 --- a/trellis2/trainers/flow_matching/mixins/image_conditioned_proj.py +++ b/trellis2/trainers/flow_matching/mixins/image_conditioned_proj.py @@ -404,6 +404,7 @@ def __init__( # NAF upsampler (frozen, no trainable params) self.naf_model = None # Lazy-loaded on first use to avoid import if not needed + self.naf_download_if_missing = True # Per-axis spatial tile factor for the NAF + proj_grid streaming path. 1 = un-tiled # (current behavior, default). >1 enables a streaming wrapper that tiles NAF and @@ -433,27 +434,32 @@ def _cached_naf_repo(self) -> Optional[str]: def _load_naf(self): """Lazy-load pretrained NAF model.""" if self.naf_model is None: - import torch.hub - device = next(self.model.parameters()).device - cached_repo = self._cached_naf_repo() - - if cached_repo is not None: - self.naf_model = torch.hub.load( - cached_repo, "naf", pretrained=True, device=device, source="local", trust_repo=True - ) - elif self.naf_download_if_missing: - self.naf_model = torch.hub.load( - "valeoai/NAF", "naf", pretrained=True, device=device, trust_repo=True - ) - - self.naf_model.eval() - self.naf_model.requires_grad_(False) + try: + import torch.hub + device = next(self.model.parameters()).device + cached_repo = self._cached_naf_repo() + + if cached_repo is not None: + self.naf_model = torch.hub.load( + cached_repo, "naf", pretrained=True, device=device, source="local", trust_repo=True + ) + elif self.naf_download_if_missing: + self.naf_model = torch.hub.load( + "valeoai/NAF", "naf", pretrained=True, device=device, trust_repo=True + ) + + self.naf_model.eval() + self.naf_model.requires_grad_(False) + except Exception as e: + import warnings + warnings.warn(f"NAF model loading failed ({e}), falling back to non-NAF mode") + self.naf_model = False def to(self, device): super().to(device) self.model.to(device) self.proj_grid.to(device) - if self.naf_model is not None: + if self.naf_model is not None and self.naf_model is not False: self.naf_model.to(device) return self @@ -461,7 +467,7 @@ def cuda(self): super().cuda() self.model.cuda() self.proj_grid.cuda() - if self.naf_model is not None: + if self.naf_model is not None and self.naf_model is not False: self.naf_model.cuda() return self @@ -560,38 +566,33 @@ def forward( ) # [B, grid_res³, D] # --- High-resolution branch (NAF): upsample then sample --- - if self.use_naf_upsample: + if self.use_naf_upsample and self.naf_model is not False: self._load_naf() # NAF expects: guide [B, 3, H, W], lr_features [B, C, h, w], target_size (H', W') lr_features_bchw = z_patchtokens_spatial.permute(0, 3, 1, 2) # [B, D, h, w] K = getattr(self, 'naf_tile_factor', 1) or 1 if K <= 1: - # Un-tiled path (default, unchanged behavior). hr_features = self.naf_model( image_for_naf, lr_features_bchw, self.naf_target_size ) # [B, D, H', W'] - # Sample from high-res feature map using same projection coordinates z_proj_hr = self.proj_grid( hr_features, camera_angle_x, distance, mesh_scale, transform_matrix, - BHWC=False # hr_features is [B, C, H', W'] + BHWC=False ) # [B, grid_res³, D] del hr_features else: - # Tiled streaming path: bound the transient HR feature map by running - # NAF and the projection sample one image-space tile at a time. z_proj_hr = self._proj_naf_tiled( image_for_naf, lr_features_bchw, camera_angle_x, distance, mesh_scale, transform_matrix, tile_factor=int(K), ) - # Concatenate lr and hr: [B, grid_res³, D*2] z_proj = torch.cat([z_proj_lr, z_proj_hr], dim=-1) else: z_proj = z_proj_lr # [B, grid_res³, D]