app.py

"""
Room Layout Generator Web Interface

A web application for generating indoor layouts from natural language descriptions
"""
import warnings
import gradio as gr
from pathlib import Path

from config.settings import settings
from core import ConfigGenerator, ModelConverter
from backend import InfinigenBackend
from utils.logger import logger

warnings.filterwarnings("ignore", category=UserWarning)


class LayoutApp:
    """Layout generation application"""
    
    def __init__(self):
        # Validate configuration
        try:
            settings.validate()
        except ValueError as e:
            logger.error(f"Configuration error: {e}")
            raise
        
        # Initialize components
        self.config_gen = ConfigGenerator()
        self.backend = InfinigenBackend()
        self.converter = ModelConverter()
        
        logger.info("Application initialized successfully")
    
    def generate_layout(self, user_input: str) -> str:
        """
        Main pipeline: from user input to 3D model
        
        Args:
            user_input: User's layout requirement description
            
        Returns:
            Path to generated .glb file
        """
        if not user_input or not user_input.strip():
            return "Please enter your layout requirements"
        
        try:
            # 1. Generate configuration
            logger.info("Step 1/3: Generating configuration...")
            yaml_content = self.config_gen.generate(user_input)
            
            # Save configuration
            self.config_gen.save(yaml_content, settings.CONFIG_OUTPUT)
            
            # 2. Generate layout
            logger.info("Step 2/3: Generating layout (this may take a while)...")
            blend_file = self.backend.generate(
                config_path=settings.CONFIG_OUTPUT,
                output_dir=settings.OUTPUT_DIR,
                seed=settings.DEFAULT_SEED
            )
            
            # 3. Convert format
            logger.info("Step 3/3: Converting to web format...")
            glb_file = self.converter.blend_to_glb(blend_file)
            
            logger.info(f"Generation complete: {glb_file}")
            return str(glb_file)
            
        except Exception as e:
            error_msg = f"Generation failed: {str(e)}"
            logger.error(error_msg)
            return error_msg


def create_interface(app: LayoutApp) -> gr.Interface:
    """Create Gradio interface"""
    
    demo = gr.Interface(
        fn=app.generate_layout,
        inputs=gr.Textbox(
            lines=6,
            placeholder=(
                "Describe your desired room layout...\n\n"
                "Example:\n"
                "I want a 2-bedroom house with private bathrooms in each bedroom, "
                "living room 1.3x larger than bedrooms, kitchen connected to dining room."
            ),
            label="Layout Requirements"
        ),
        outputs=gr.Model3D(
            scale=2,
            label="Generated 3D Layout"
        ),
        title="Room Layout Generator",
        description=(
            "Generate indoor layouts from natural language descriptions.\n"
            "Supports English/Chinese input. Max 3 bedrooms, 10 total rooms."
        ),
        examples=[
            ["2 bedrooms with private bathrooms, living room 1.5x bedroom size"],
            ["3-bedroom house, 2 bathrooms, kitchen connected to dining room"],
            ["1 bedroom, 1 bathroom, open kitchen connected to living room"]
        ],
        theme=gr.themes.Soft()
    )
    
    return demo


def main():
    """Launch application"""
    print("Room Layout Generator")
    print("=" * 50)
    print(f"Project root: {settings.INFINIGEN_ROOT}")
    print(f"Output dir: {settings.OUTPUT_DIR}")
    print(f"Blender: {settings.BLENDER_PATH}")
    print(f"AI model: {settings.OPENAI_MODEL}")
    print("=" * 50)
    
    try:
        app = LayoutApp()
        demo = create_interface(app)
        
        demo.launch(
            share=settings.SHARE_GRADIO,
            allowed_paths=[str(settings.OUTPUT_DIR)]
        )
        
    except Exception as e:
        logger.error(f"Failed to start application: {e}")
        raise


if __name__ == "__main__":
    main()