comfyui_nodes/interiorfusion_nodes.py

"""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",
}