roi/EditP23

1import argparse2import sys3from pathlib import Path4from typing import Optional5 6# --- Start of the "Messy" but Effective Path Setup ---7# This block ensures that imports work correctly without modifying the src directory.8# It adds both the project root and the src directory to the Python path.9try:10    # Get the project root directory (which is the parent of the 'scripts' directory)11    project_root = Path(__file__).resolve().parent.parent12    # Get the source code directory13    src_dir = project_root / "src"14 15    # Add both directories to the system path16    sys.path.insert(0, str(project_root))17    sys.path.insert(0, str(src_dir))18except IndexError:19    # Fallback for when the script is run in a way that __file__ is not defined20    print("Could not determine project root. Please run from the 'scripts' directory.")21    sys.exit(1)22# --- End of Path Setup ---23 24import torch25from PIL import Image26 27from pipeline import Zero123PlusPipeline # This now works because src/ is on the path28from utils import add_white_bg, load_z123_pipe29 30 31def generate_from_single_view(32    input_path: Path,33    output_path: Path,34    device_number: int = 0,35    pipeline: Optional[Zero123PlusPipeline] = None,36) -> None:37    """38    Generates a multi-view image grid from a single input image.39 40    Args:41        input_path: Path to the single input image.42        output_path: Path to save the generated multi-view .png file.43        device_number: The GPU device number to use.44        pipeline: An optional pre-loaded pipeline instance.45    """46    if not input_path.is_file():47        raise FileNotFoundError(f"Input image not found at: {input_path}")48 49    print(f"Loading pipeline on device {device_number}...")50    if pipeline is None:51        pipeline = load_z123_pipe(device_number)52 53    print(f"Processing input image: {input_path}")54    cond_image = Image.open(input_path)55    cond_image = add_white_bg(cond_image)56 57    print("Generating multi-view grid...")58    result = pipeline(cond_image, num_inference_steps=75).images[0]59 60    output_path.parent.mkdir(parents=True, exist_ok=True)61    result.save(output_path)62    print(f"Successfully saved multi-view grid to: {output_path}")63 64 65if __name__ == '__main__':66    parser = argparse.ArgumentParser(67        description="Generate a multi-view image grid from a single input view using Zero123++."68    )69    parser.add_argument(70        "--input_image",71        type=Path,72        required=True,73        help="Path to the single input image file (e.g., examples/robot_sunglasses/src.png)."74    )75    parser.add_argument(76        "--output_path",77        type=Path,78        required=True,79        help="Path to save the output multi-view grid (e.g., examples/robot_sunglasses/src_mv.png)."80    )81    parser.add_argument(82        "--device_number",83        type=int,84        default=0,85        help="GPU device number to use for generation."86    )87    args = parser.parse_args()88 89    try:90        generate_from_single_view(91            input_path=args.input_image,92            output_path=args.output_path,93            device_number=args.device_number94        )95    except Exception as e:96        print(f"An error occurred: {e}")97        sys.exit(1)98