test.py

import argparse
import logging
import os

import numpy as np
import torch
from PIL import Image

from src.synclight.inference import evaluate, get_model

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def load_images(image_paths: list[str]) -> list[Image.Image]:
    """Load images from file paths and convert to RGB.
    
    Args:
        image_paths: List of paths to image files
        
    Returns:
        List of PIL Image objects in RGB format
    """
    images = []
    for path in image_paths:
        try:
            img = Image.open(path).convert("RGB")
            logger.info(f"Loaded image: {path}")
            images.append(img)
        except Exception as e:
            logger.error(f"Failed to load image {path}: {e}")
            raise
    return images


def load_lightmap(lightmap_path: str) -> np.ndarray:
    """Load lightmap from numpy file.
    
    Args:
        lightmap_path: Path to numpy file containing 4-channel lightmap
        
    Returns:
        Numpy array with shape (H, W, 4)
    """
    try:
        lightmap = np.load(lightmap_path)
        
        # Validate shape
        if lightmap.ndim != 3 or lightmap.shape[2] != 4:
            raise ValueError(
                f"Lightmap must be a 4-channel numpy array with shape (H, W, 4), "
                f"got shape {lightmap.shape}"
            )
        
        # Ensure dtype is float32
        if lightmap.dtype != np.float32:
            lightmap = lightmap.astype(np.float32)
        
        logger.info(f"Loaded lightmap from {lightmap_path} with shape {lightmap.shape}")
        return lightmap
    except Exception as e:
        logger.error(f"Failed to load lightmap {lightmap_path}: {e}")
        raise


def save_results(
    output_images: list[Image.Image],
    output_dir: str,
    prefix: str = "output"
) -> None:
    """Save output images to the specified directory.
    
    Args:
        output_images: List of PIL Image objects to save
        output_dir: Directory to save images to
        prefix: Prefix for output filenames
    """
    os.makedirs(output_dir, exist_ok=True)
    
    for i, img in enumerate(output_images):
        output_path = os.path.join(output_dir, f"{prefix}_{i}.jpg")
        img.save(output_path)
        logger.info(f"Saved output image to {output_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Run SyncLight inference on multiple images with a lightmap."
    )
    parser.add_argument(
        "--image_paths",
        type=str,
        nargs="+",
        required=True,
        help="Paths to input images (at least one required, space-separated)",
    )
    parser.add_argument(
        "--lightmap_path",
        type=str,
        required=True,
        help="Path to 4-channel numpy lightmap file (.npy)",
    )
    parser.add_argument(
        "--model_weights",
        type=str,
        required=True,
        help="Path to model weights directory (local path or HuggingFace Hub model name)",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./outputs",
        help="Directory to save output images (default: ./outputs)",
    )
    parser.add_argument(
        "--num_inference_steps",
        type=int,
        default=1,
        help="Number of inference steps (default: 1)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device to run inference on (default: cuda)",
    )
    parser.add_argument(
        "--torch_dtype",
        type=str,
        default="bfloat16",
        choices=["bfloat16", "float32"],
        help="PyTorch data type (default: bfloat16)",
    )

    args = parser.parse_args()

    # Validate inputs
    if not args.image_paths:
        logger.error("At least one image path is required")
        parser.print_help()
        return

    for image_path in args.image_paths:
        if not os.path.exists(image_path):
            logger.error(f"Image file not found: {image_path}")
            return

    if not os.path.exists(args.lightmap_path):
        logger.error(f"Lightmap file not found: {args.lightmap_path}")
        return

    # Set torch dtype
    torch_dtype = torch.bfloat16 if args.torch_dtype == "bfloat16" else torch.float32

    logger.info("=" * 80)
    logger.info("SyncLight Inference")
    logger.info("=" * 80)
    logger.info(f"Input images: {args.image_paths}")
    logger.info(f"Lightmap: {args.lightmap_path}")
    logger.info(f"Model weights: {args.model_weights}")
    logger.info(f"Output directory: {args.output_dir}")
    logger.info(f"Inference steps: {args.num_inference_steps}")
    logger.info(f"Device: {args.device}")
    logger.info(f"Torch dtype: {args.torch_dtype}")
    logger.info("=" * 80)

    # Load model
    logger.info("Loading model...")
    try:
        model = get_model(
            args.model_weights,
            torch_dtype=torch_dtype,
            device=args.device,
        )
        logger.info("Model loaded successfully")
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        return

    # Load inputs
    logger.info("Loading input images...")
    try:
        source_images = load_images(args.image_paths)
    except Exception as e:
        logger.error(f"Failed to load images: {e}")
        return

    logger.info("Loading lightmap...")
    try:
        lightmap = load_lightmap(args.lightmap_path)
    except Exception as e:
        logger.error(f"Failed to load lightmap: {e}")
        return

    # Run inference
    logger.info("Running inference...")
    try:
        output_images = evaluate(
            model=model,
            source_images=source_images,
            lightmap_image=lightmap,
            num_sampling_steps=args.num_inference_steps,
        )
        logger.info(f"Inference completed. Generated {len(output_images)} output image(s)")
    except Exception as e:
        logger.error(f"Inference failed: {e}")
        raise

    # Save results
    logger.info("Saving results...")
    try:
        save_results(output_images, args.output_dir, prefix="output")
        logger.info(f"Results saved to {args.output_dir}")
    except Exception as e:
        logger.error(f"Failed to save results: {e}")
        return

    logger.info("=" * 80)
    logger.info("Inference completed successfully!")
    logger.info("=" * 80)


if __name__ == "__main__":
    main()