app.py

"""Gradio app for InteriorFusion."""

import os
import tempfile
from pathlib import Path

import gradio as gr
import torch
from PIL import Image

from interiorfusion.pipelines import InteriorFusionPipeline

# Initialize pipeline
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

_pipeline = InteriorFusionPipeline(
    model_size="L",
    device=DEVICE,
    dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
)


def generate_scene(
    image,
    room_type,
    style,
    export_glb,
    export_fbx,
    export_obj,
    export_ply,
    export_usdz,
):
    """Generate 3D scene from input image."""
    if image is None:
        return "Please upload an image first.", None, None, None, None, None, None, None
    
    formats = []
    if export_glb:
        formats.append("glb")
    if export_fbx:
        formats.append("fbx")
    if export_obj:
        formats.append("obj")
    if export_ply:
        formats.append("ply")
    if export_usdz:
        formats.append("usdz")
    
    if not formats:
        formats = ["glb", "ply"]
    
    # Run pipeline
    output = _pipeline(
        image=image,
        room_type_hint=room_type if room_type != "Auto" else None,
        style_hint=style if style != "Auto" else None,
    )
    
    # Export
    output_dir = tempfile.mkdtemp()
    output.export_all(output_dir)
    
    # Prepare file outputs
    files = {}
    for fmt in formats:
        path = Path(output_dir) / f"scene.{fmt}"
        if path.exists():
            files[fmt] = str(path)
    
    # Build summary
    summary = f"""## ✅ Generation Complete

- **Room Type**: {output.room_type}
- **Style**: {output.style}
- **Objects**: {len(output.object_meshes)}
- **Materials**: {len(output.pbr_materials)}
- **Processing Time**: {output.processing_time:.1f}s
"""
    
    glb_file = files.get("glb", None)
    fbx_file = files.get("fbx", None)
    obj_file = files.get("obj", None)
    ply_file = files.get("ply", None)
    usdz_file = files.get("usdz", None)
    
    return (
        summary,
        glb_file,
        fbx_file,
        obj_file,
        ply_file,
        usdz_file,
        files.get("glb", None),  # For 3D viewer
    )


# Create Gradio interface
with gr.Blocks(title="InteriorFusion - AI Interior Designer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏠 InteriorFusion
    ### Single Photo → Editable 3D Interior Scene
    
    Upload a photo of any room and get a complete 3D scene with:
    - ✅ Textured 3D meshes (GLB, FBX, OBJ, USDZ)
    - ✅ Gaussian Splatting (PLY)
    - ✅ PBR materials (metallic, roughness, normal maps)
    - ✅ Editable furniture objects
    - ✅ Scene graph representation
    
    *Powered by the open-source InteriorFusion model*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Upload Room Photo",
                type="pil",
                height=400,
            )
            
            with gr.Row():
                room_type = gr.Dropdown(
                    choices=["Auto", "living_room", "bedroom", "kitchen", 
                             "dining_room", "office", "bathroom", "hallway"],
                    value="Auto",
                    label="Room Type",
                )
                style = gr.Dropdown(
                    choices=["Auto", "modern", "scandinavian", "luxury",
                             "industrial", "minimalist", "bohemian", "indian",
                             "japanese", "traditional"],
                    value="Auto",
                    label="Style",
                )
            
            gr.Markdown("**Export Formats**")
            with gr.Row():
                export_glb = gr.Checkbox(value=True, label="GLB")
                export_fbx = gr.Checkbox(value=False, label="FBX")
                export_obj = gr.Checkbox(value=False, label="OBJ")
                export_ply = gr.Checkbox(value=True, label="PLY (3DGS)")
                export_usdz = gr.Checkbox(value=False, label="USDZ")
            
            generate_btn = gr.Button("Generate 3D Scene", variant="primary")
            
        with gr.Column(scale=2):
            output_summary = gr.Markdown()
            
            with gr.Row():
                viewer_3d = gr.Model3D(
                    label="3D Preview",
                    height=500,
                )
            
            gr.Markdown("**Download Files**")
            with gr.Row():
                glb_download = gr.File(label="GLB")
                fbx_download = gr.File(label="FBX")
            with gr.Row():
                obj_download = gr.File(label="OBJ")
                ply_download = gr.File(label="PLY (3DGS)")
                usdz_download = gr.File(label="USDZ")
    
    generate_btn.click(
        fn=generate_scene,
        inputs=[
            input_image,
            room_type,
            style,
            export_glb,
            export_fbx,
            export_obj,
            export_ply,
            export_usdz,
        ],
        outputs=[
            output_summary,
            glb_download,
            fbx_download,
            obj_download,
            ply_download,
            usdz_download,
            viewer_3d,
        ],
    )
    
    gr.Markdown("""
    ---
    ### Tips for Best Results
    
    1. **Photo Quality**: Use well-lit, clear photos with minimal blur
    2. **Camera Angle**: Eye-level shots work best (not floor-level or ceiling)
    3. **Room Visibility**: Try to capture the full room, not just a corner
    4. **Furniture**: Rooms with 2-8 furniture pieces produce the best results
    5. **Export**: GLB is best for Blender/Unity; PLY for Gaussian Splatting viewers
    
    ### Model Sizes
    - **S**: Fast preview (~5s on RTX 4090)
    - **L**: Balanced quality/speed (~15s)
    - **XL**: Maximum quality (~30s)
    
    *Note: This is a research prototype. The full trained model will produce much higher quality.*
    """)

if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)