Fix: Pixal3D Windows compatibility + configurable NAF target size

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)

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,17 +567,17 @@ 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)
         scale_tensor = torch.tensor([mesh_scale], device=device)
         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},

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,35 +434,40 @@ 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
 
     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]