From f791f4fdac158377d599f2473a91467eb10279c3 Mon Sep 17 00:00:00 2001 From: Johnson Sun Date: Thu, 5 Feb 2026 15:37:14 +0800 Subject: [PATCH] Add `--folder` option to export USD files without USDZ packaging Large USDZ archives can fail to export due to size limit. This adds a --folder option to `ply_to_usd.py` that exports USD files (default.usda, gauss.usda, and .nurec) directly to a folder instead of packaging them in a USDZ archive. Reference: https://github.com/nv-tlabs/3dgrut/issues/139 --- threedgrut/export/nurec_templates.py | 14 ++++++++++++ threedgrut/export/scripts/ply_to_usd.py | 24 +++++++++++++++----- threedgrut/export/usd_util.py | 30 +++++++++++++++++++++++++ threedgrut/export/usdz_exporter.py | 26 ++++++++++++++++----- 4 files changed, 82 insertions(+), 12 deletions(-) diff --git a/threedgrut/export/nurec_templates.py b/threedgrut/export/nurec_templates.py index d1217031..0db3e412 100644 --- a/threedgrut/export/nurec_templates.py +++ b/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], diff --git a/threedgrut/export/scripts/ply_to_usd.py b/threedgrut/export/scripts/ply_to_usd.py index 4bb8934f..90788be7 100644 --- a/threedgrut/export/scripts/ply_to_usd.py +++ b/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: diff --git a/threedgrut/export/usd_util.py b/threedgrut/export/usd_util.py index e84b66f1..cdf25829 100644 --- a/threedgrut/export/usd_util.py +++ b/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}") diff --git a/threedgrut/export/usdz_exporter.py b/threedgrut/export/usdz_exporter.py index 048c5aec..5ea44cd8 100644 --- a/threedgrut/export/usdz_exporter.py +++ b/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)