Upload comfyui_nodes/interiorfusion_nodes.py

comfyui_nodes/interiorfusion_nodes.py

@@ -0,0 +1,206 @@
+"""ComfyUI nodes for InteriorFusion.
+
+Nodes:
+- InteriorFusionSceneNode: Full pipeline from image to scene
+- InteriorFusionObjectNode: Generate single furniture object
+- InteriorFusionMaterialNode: Apply PBR materials to mesh
+- InteriorFusionExportNode: Export to various formats
+"""
+
+import os
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+class InteriorFusionSceneNode:
+    """Generate complete 3D interior scene from single image."""
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "model_size": (["S", "L", "XL"], {"default": "L"}),
+                "room_type": (["auto", "living_room", "bedroom", "kitchen",
+                              "dining_room", "office"], {"default": "auto"}),
+                "style": (["auto", "modern", "scandinavian", "luxury",
+                           "industrial", "minimalist"], {"default": "auto"}),
+                "use_pbr": ("BOOLEAN", {"default": True}),
+                "use_gaussian": ("BOOLEAN", {"default": True}),
+            }
+        }
+    
+    RETURN_TYPES = ("MESH", "SCENE_GRAPH", "STRING")
+    RETURN_NAMES = ("scene_mesh", "scene_graph", "metadata")
+    FUNCTION = "generate_scene"
+    CATEGORY = "InteriorFusion"
+    
+    def __init__(self):
+        self.pipeline = None
+    
+    def generate_scene(self, image, model_size, room_type, style, use_pbr, use_gaussian):
+        from interiorfusion.pipelines import InteriorFusionPipeline
+        
+        if self.pipeline is None or self.pipeline.model_size != model_size:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.pipeline = InteriorFusionPipeline(
+                model_size=model_size,
+                device=device,
+                dtype=torch.float16,
+                use_pbr=use_pbr,
+                use_gaussian_splatting=use_gaussian,
+            )
+        
+        # Convert ComfyUI image tensor to PIL
+        # ComfyUI images are [B, H, W, C] in range 0-1
+        img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
+        pil_image = Image.fromarray(img_np)
+        
+        # Generate
+        output = self.pipeline(
+            image=pil_image,
+            room_type_hint=room_type if room_type != "auto" else None,
+            style_hint=style if style != "auto" else None,
+        )
+        
+        metadata = f"""
+Room Type: {output.room_type}
+Style: {output.style}
+Objects: {len(output.object_meshes)}
+Materials: {len(output.pbr_materials)}
+Time: {output.processing_time:.1f}s
+        """.strip()
+        
+        return (output.scene_mesh, output.scene_graph, metadata)
+
+
+class InteriorFusionObjectNode:
+    """Generate a single furniture object from image."""
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "object_image": ("IMAGE",),
+                "object_mask": ("MASK",),
+                "model_size": (["S", "L", "XL"], {"default": "L"}),
+            }
+        }
+    
+    RETURN_TYPES = ("MESH",)
+    RETURN_NAMES = ("object_mesh",)
+    FUNCTION = "generate_object"
+    CATEGORY = "InteriorFusion"
+    
+    def generate_object(self, object_image, object_mask, model_size):
+        from interiorfusion.models.multiview_generation import MultiViewGenerationModule
+        from interiorfusion.models.reconstruction_3d import Reconstruction3DModule
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Convert inputs
+        img_np = (object_image[0].cpu().numpy() * 255).astype(np.uint8)
+        mask_np = object_mask[0].cpu().numpy()
+        
+        pil_image = Image.fromarray(img_np)
+        
+        # Generate multi-view
+        mv_gen = MultiViewGenerationModule(model_size=model_size, device=device)
+        multiviews = mv_gen.generate_object_views(pil_image, mask_np)
+        
+        # Reconstruct 3D
+        recon = Reconstruction3DModule(model_size=model_size, device=device)
+        mesh, _ = recon.reconstruct_object(multiviews)
+        
+        return (mesh,)
+
+
+class InteriorFusionMaterialNode:
+    """Apply PBR materials to a mesh."""
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mesh": ("MESH",),
+                "material_type": (["wood", "fabric", "metal", "glass",
+                                    "plastic", "leather", "wall", "floor"],
+                                   {"default": "wood"}),
+                "color_hex": ("STRING", {"default": "#8B4513"}),
+                "metallic": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
+                "roughness": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
+            }
+        }
+    
+    RETURN_TYPES = ("MESH",)
+    RETURN_NAMES = ("textured_mesh",)
+    FUNCTION = "apply_material"
+    CATEGORY = "InteriorFusion"
+    
+    def apply_material(self, mesh, material_type, color_hex, metallic, roughness):
+        # Parse color
+        color = tuple(int(color_hex.lstrip('#')[i:i+2], 16) / 255.0
+                     for i in (0, 2, 4))
+        
+        material = {
+            "type": material_type,
+            "albedo": list(color),
+            "metallic": metallic,
+            "roughness": roughness,
+        }
+        
+        # Apply to mesh
+        if hasattr(mesh, 'materials'):
+            mesh.materials = {"default": material}
+        
+        return (mesh,)
+
+
+class InteriorFusionExportNode:
+    """Export mesh to various 3D formats."""
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mesh": ("MESH",),
+                "format": (["glb", "fbx", "obj", "usdz", "ply"],
+                          {"default": "glb"}),
+                "filename": ("STRING", {"default": "scene"}),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("file_path",)
+    FUNCTION = "export_mesh"
+    CATEGORY = "InteriorFusion"
+    
+    def export_mesh(self, mesh, format, filename):
+        from interiorfusion.utils.mesh_utils import export_mesh as do_export
+        
+        output_dir = tempfile.gettempdir()
+        output_path = os.path.join(output_dir, f"{filename}.{format}")
+        
+        do_export(mesh, output_path, format=format)
+        
+        return (output_path,)
+
+
+# Node mappings for ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "InteriorFusionScene": InteriorFusionSceneNode,
+    "InteriorFusionObject": InteriorFusionObjectNode,
+    "InteriorFusionMaterial": InteriorFusionMaterialNode,
+    "InteriorFusionExport": InteriorFusionExportNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "InteriorFusionScene": "InteriorFusion: Generate Scene",
+    "InteriorFusionObject": "InteriorFusion: Generate Object",
+    "InteriorFusionMaterial": "InteriorFusion: Apply Material",
+    "InteriorFusionExport": "InteriorFusion: Export Mesh",
+}