Upload interiorgen3d/config/pipeline_config.py

interiorgen3d/config/pipeline_config.py

@@ -0,0 +1,106 @@
+"""
+Configuration for InteriorGen3D pipeline.
+"""
+from dataclasses import dataclass, field
+from typing import Optional, List
+
+
+@dataclass
+class DepthConfig:
+    model_name: str = "depth-anything/Depth-Anything-V2-Large-hf"
+    input_size: int = 518
+    metric_depth: bool = True
+    device: str = "cuda"
+
+
+@dataclass
+class SegmentationConfig:
+    sam_model: str = "facebook/sam2-hiera-large"
+    num_categories: int = 59
+    min_object_area: float = 0.01
+    device: str = "cuda"
+
+
+@dataclass
+class LayoutConfig:
+    model_name: str = "manycore-research/SpatialLM1.1-Qwen-0.5B"
+    quantization_bins: int = 1280
+    resolution_cm: float = 2.5
+    device: str = "cuda"
+
+
+@dataclass
+class ShapeConfig:
+    model_name: str = "microsoft/TRELLIS-image-large"
+    voxel_resolution: int = 64
+    latent_channels: int = 8
+    num_active_voxels: int = 20000
+    flow_steps: int = 50
+    cfg_scale: float = 7.5
+    device: str = "cuda"
+
+
+@dataclass
+class TextureConfig:
+    texture_resolution: int = 2048
+    num_views: int = 6
+    diffusion_resolution: int = 512
+    use_pbr: bool = True
+    device: str = "cuda"
+
+
+@dataclass
+class ExportConfig:
+    formats: List[str] = field(default_factory=lambda: ["glb", "fbx", "obj", "usdz"])
+    texture_resolution: int = 2048
+    max_triangles: int = 500000
+
+
+@dataclass
+class PipelineConfig:
+    depth: DepthConfig = field(default_factory=DepthConfig)
+    segmentation: SegmentationConfig = field(default_factory=SegmentationConfig)
+    layout: LayoutConfig = field(default_factory=LayoutConfig)
+    shape: ShapeConfig = field(default_factory=ShapeConfig)
+    texture: TextureConfig = field(default_factory=TextureConfig)
+    export: ExportConfig = field(default_factory=ExportConfig)
+    device: str = "cuda"
+    dtype: str = "float16"
+    seed: int = 42
+    output_dir: str = "./output"
+    max_parallel_objects: int = 4
+
+    @classmethod
+    def for_rtx4090(cls):
+        config = cls()
+        config.dtype = "float16"
+        config.shape.voxel_resolution = 64
+        config.texture.texture_resolution = 1024
+        config.max_parallel_objects = 2
+        return config
+
+    @classmethod
+    def for_a100(cls):
+        config = cls()
+        config.dtype = "bfloat16"
+        config.shape.voxel_resolution = 128
+        config.texture.texture_resolution = 2048
+        config.max_parallel_objects = 4
+        return config
+
+    @classmethod
+    def for_h100(cls):
+        config = cls()
+        config.dtype = "bfloat16"
+        config.shape.voxel_resolution = 128
+        config.texture.texture_resolution = 4096
+        config.max_parallel_objects = 8
+        return config
+
+    @classmethod
+    def lightweight(cls):
+        config = cls()
+        config.shape.voxel_resolution = 32
+        config.texture.texture_resolution = 512
+        config.max_parallel_objects = 1
+        return config