src/interiorfusion/__main__.py

"""InteriorFusion CLI entry point."""

import argparse
import os
import sys
from pathlib import Path

import torch
from PIL import Image

from .pipelines import InteriorFusionPipeline


def parse_args():
    parser = argparse.ArgumentParser(
        description="InteriorFusion: Single image to 3D interior scene"
    )
    
    parser.add_argument(
        "--image", "-i",
        type=str,
        required=True,
        help="Path to input interior image",
    )
    parser.add_argument(
        "--output", "-o",
        type=str,
        default="./output",
        help="Output directory for generated 3D files",
    )
    parser.add_argument(
        "--model-size",
        type=str,
        default="L",
        choices=["S", "L", "XL"],
        help="Model size: S (fast), L (balanced), XL (quality)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for inference",
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="float16",
        choices=["float16", "bfloat16", "float32"],
        help="Data type for inference",
    )
    parser.add_argument(
        "--room-type",
        type=str,
        default=None,
        help="Optional room type hint (living_room, bedroom, kitchen, etc.)",
    )
    parser.add_argument(
        "--style",
        type=str,
        default=None,
        help="Optional style hint (modern, scandinavian, luxury, etc.)",
    )
    parser.add_argument(
        "--formats",
        type=str,
        default="glb,ply",
        help="Comma-separated export formats (glb,fbx,obj,usdz,ply)",
    )
    parser.add_argument(
        "--interactive",
        action="store_true",
        help="Launch interactive editing mode",
    )
    parser.add_argument(
        "--no-pbr",
        action="store_true",
        help="Disable PBR material generation (faster)",
    )
    parser.add_argument(
        "--no-gaussian",
        action="store_true",
        help="Disable Gaussian Splatting output",
    )
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    # Validate input
    if not os.path.exists(args.image):
        print(f"Error: Image not found: {args.image}", file=sys.stderr)
        sys.exit(1)
    
    # Parse dtype
    dtype_map = {
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
        "float32": torch.float32,
    }
    dtype = dtype_map[args.dtype]
    
    # Parse formats
    formats = [f.strip() for f in args.formats.split(",")]
    
    print(f"InteriorFusion {args.model_size}")
    print(f"  Device: {args.device}")
    print(f"  DType: {args.dtype}")
    print(f"  Input: {args.image}")
    print(f"  Output: {args.output}")
    print(f"  Formats: {formats}")
    print()
    
    # Load pipeline
    print("Loading pipeline...")
    pipeline = InteriorFusionPipeline(
        model_size=args.model_size,
        device=args.device,
        dtype=dtype,
        use_pbr=not args.no_pbr,
        use_gaussian_splatting=not args.no_gaussian,
    )
    
    # Load image
    print("Loading image...")
    image = Image.open(args.image).convert("RGB")
    print(f"  Size: {image.size}")
    
    # Generate
    print("\nGenerating 3D scene...")
    output = pipeline(
        image=image,
        room_type_hint=args.room_type,
        style_hint=args.style,
    )
    
    # Export
    print("\nExporting...")
    output.export_all(args.output)
    
    # Print results
    print(f"\n{'='*50}")
    print("Generation Complete!")
    print(f"{'='*50}")
    print(f"  Time: {output.processing_time:.1f}s")
    print(f"  Room type: {output.room_type}")
    print(f"  Style: {output.style}")
    print(f"  Objects: {len(output.object_meshes)}")
    print(f"  Materials: {len(output.pbr_materials)}")
    print()
    
    # Export paths
    if output.glb_path:
        print(f"  GLB: {output.glb_path}")
    if output.fbx_path:
        print(f"  FBX: {output.fbx_path}")
    if output.obj_path:
        print(f"  OBJ: {output.obj_path}")
    if output.usdz_path:
        print(f"  USDZ: {output.usdz_path}")
    if output.ply_path:
        print(f"  PLY (3DGS): {output.ply_path}")
    
    if args.interactive:
        print("\nInteractive editing not yet implemented in CLI.")
        print("Use the Gradio app for interactive editing.")
    
    print(f"\nDone!")


if __name__ == "__main__":
    main()