Add --folder option to export without USDZ packaging for large PLY files

threedgrut/export/nurec_templates.py

@@ -15,6 +15,7 @@
 
 import zipfile
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict, Union
 
 import numpy as np
@@ -38,6 +39,19 @@ def save_to_zip(self, zip_file: zipfile.ZipFile):
         """
         zip_file.writestr(self.filename, self.serialized)
 
+    def save_to_folder(self, out_dir: Path):
+        """
+        Save the serialized data to a directory.
+
+        Args:
+            out_dir: Directory to save the data to
+        """
+        out_dir.mkdir(parents=True, exist_ok=True)
+        out_path = out_dir / self.filename
+        mode = "wb" if isinstance(self.serialized, bytes) else "w"
+        with open(out_path, mode) as f:
+            f.write(self.serialized)
+
 
 def _fill_state_dict_tensors(
     template: Dict[str, Any],

threedgrut/export/scripts/ply_to_usd.py

@@ -49,10 +49,18 @@ def load_default_config(
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Convert PLY to USDZ")
+    parser = argparse.ArgumentParser(description="Convert PLY to USDZ or USD folder")
     parser.add_argument("input_file", type=str, help="Input PLY file path")
     parser.add_argument(
-        "--output_file", type=str, help="Output USDZ file path (defaults to input file path with .usdz extension)"
+        "--output_file",
+        type=str,
+        help="Output USDZ file path or folder path (defaults to input file path with .usdz extension)",
+    )
+    parser.add_argument(
+        "--folder",
+        action="store_true",
+        help="Export to a folder containing USD files instead of a USDZ archive. "
+        "This is useful for large models that may fail to load in Omniverse Kit when packaged as USDZ."
     )
 
     args = parser.parse_args()
@@ -73,9 +81,13 @@ def main():
         output_path = Path(args.output_file)
         output_path.parent.mkdir(parents=True, exist_ok=True)
     else:
-        output_path = input_path.with_suffix(".usdz")
+        if args.folder:
+            output_path = input_path
+        else:
+            output_path = input_path.with_suffix(".usdz")
 
-    logger.info(f"Converting {input_path} to {output_path}")
+    output_type = "folder" if args.folder else "USDZ"
+    logger.info(f"Converting {input_path} to {output_type}: {output_path}")
 
     try:
         # 1. Create model with default config
@@ -90,9 +102,9 @@ def main():
         # 3. Create USDZExporter
         exporter = USDZExporter()
 
-        # 4. Export to USDZ
+        # 4. Export to USDZ or folder
         logger.info(f"Exporting with USDZExporter: {output_path}")
-        exporter.export(model, output_path, dataset=None, conf=conf)
+        exporter.export(model, output_path, dataset=None, conf=conf, as_folder=args.folder)
 
         logger.info(f"Successfully exported to {output_path}")
     except Exception as e:

threedgrut/export/usd_util.py

@@ -37,6 +37,10 @@ def save(self, out_dir: Path):
         out_dir.mkdir(parents=True, exist_ok=True)
         self.stage.Export(str(out_dir / self.filename))
 
+    def save_to_folder(self, out_dir: Path):
+        out_dir.mkdir(parents=True, exist_ok=True)
+        self.stage.GetRootLayer().Export(str(out_dir / self.filename))
+
     def save_to_zip(self, zip_file: zipfile.ZipFile):
         with tempfile.NamedTemporaryFile(mode="wb", suffix=self.filename, delete=False) as temp_file:
             temp_file_path = temp_file.name
@@ -279,3 +283,29 @@ def write_to_usdz(file_path: Path, model_file, gauss_usd: NamedUSDStage, default
         gauss_usd.save_to_zip(zip_file)
 
     logger.info(f"USDZ file created successfully at {file_path}")
+
+
+def write_to_folder(folder_path: Path, model_file, gauss_usd: NamedUSDStage, default_usd: NamedUSDStage) -> None:
+    """
+    Write the USD files and model data to a folder.
+
+    This is an alternative to write_to_usdz for large models that may fail to load
+    when packaged in a USDZ archive.
+
+    Args:
+        folder_path: Path to the folder to write files to
+        model_file: The compressed model data
+        gauss_usd: The gauss USD stage
+        default_usd: The default USD stage
+    """
+    # Create the output folder
+    folder_path.mkdir(parents=True, exist_ok=True)
+
+    # Save the USD stages
+    default_usd.save_to_folder(folder_path)
+    gauss_usd.save_to_folder(folder_path)
+
+    # Save the model file
+    model_file.save_to_folder(folder_path)
+
+    logger.info(f"USD files created successfully in folder {folder_path}")

threedgrut/export/usdz_exporter.py

@@ -28,6 +28,7 @@
 from threedgrut.export.usd_util import (
     serialize_nurec_usd,
     serialize_usd_default_layer,
+    write_to_folder,
     write_to_usdz,
 )
 from threedgrut.utils.logger import logger
@@ -42,18 +43,28 @@ class USDZExporter(ModelExporter):
 
     @torch.no_grad()
     def export(
-        self, model: ExportableModel, output_path: Path, dataset=None, conf: Dict[str, Any] = None, **kwargs
+        self,
+        model: ExportableModel,
+        output_path: Path,
+        dataset=None,
+        conf: Dict[str, Any] = None,
+        as_folder: bool = False,
+        **kwargs,
     ) -> None:
-        """Export the model to a USDZ file.
+        """Export the model to a USDZ file or folder.
 
         Args:
             model: The model to export (must implement ExportableModel)
-            output_path: Path where the USDZ file will be saved
+            output_path: Path where the USDZ file or folder will be saved
             dataset: Optional dataset to get camera poses for upright transform
             conf: Configuration parameters for the renderer
+            as_folder: If True, export to a folder instead of a USDZ archive.
+                This is useful for large models that may fail to load in Omniverse Kit
+                when packaged as USDZ.
             **kwargs: Additional parameters for export
         """
-        logger.info(f"exporting usdz file to {output_path}...")
+        output_type = "folder" if as_folder else "usdz file"
+        logger.info(f"exporting {output_type} to {output_path}...")
 
         if not conf.render.method in ["3dgut", "3dgrt"]:
             raise ValueError(f"Not supported for USDZ export: {conf.render.method}")
@@ -133,5 +144,8 @@ def export(
         gauss_usd = serialize_nurec_usd(model_file, positions, normalizing_transform)
         default_usd = serialize_usd_default_layer(gauss_usd)
 
-        # Write the final USDZ file
-        write_to_usdz(output_path, model_file, gauss_usd, default_usd)
+        # Write the final USDZ file or folder
+        if as_folder:
+            write_to_folder(output_path, model_file, gauss_usd, default_usd)
+        else:
+            write_to_usdz(output_path, model_file, gauss_usd, default_usd)