examples/custom_config.py

"""Example of using custom configurations with PANDORA."""

from PIL import Image

from src.pandora_removal import PandoraRemoval, PandoraConfig


def example_high_quality():
    """Use higher quality settings for better results (slower)."""

    config = PandoraConfig(
        model_path="stabilityai/stable-diffusion-2-1",
        device="cuda",
        max_steps=100,  # More steps for better quality
        guidance_scale_ladg=10.0,  # Higher guidance
        percentile=95.0,  # Stricter attention control
    )

    model = PandoraRemoval(config=config)
    model.load_model()

    result = model.remove_object(
        image="input.jpg",
        mask="mask.png",
        border_size=20,  # Larger border
        num_steps=100
    )

    result.save("output_high_quality.png")
    print("✓ High-quality result saved!")


def example_fast_inference():
    """Use faster settings for quick results (lower quality)."""

    config = PandoraConfig(
        model_path="stabilityai/stable-diffusion-2-1",
        device="cuda",
        max_steps=25,  # Fewer steps for speed
        guidance_scale_ladg=5.0,  # Lower guidance
    )

    model = PandoraRemoval(config=config)
    model.load_model()

    result = model.remove_object(
        image="input.jpg",
        mask="mask.png",
        border_size=10,  # Smaller border
        num_steps=25
    )

    result.save("output_fast.png")
    print("✓ Fast result saved!")


def example_different_object_types():
    """Different settings for different object types."""

    model = PandoraRemoval()
    model.load_model()

    # Small, well-defined objects (e.g., text, logos)
    result1 = model.remove_object(
        image="image_with_text.jpg",
        mask="text_mask.png",
        border_size=2,  # Small border for crisp edges
        guidance_scale=10.0
    )
    result1.save("output_text_removed.png")

    # Large, complex objects (e.g., people, buildings)
    result2 = model.remove_object(
        image="image_with_person.jpg",
        mask="person_mask.png",
        border_size=22,  # Large border for smooth blending
        guidance_scale=7.5
    )
    result2.save("output_person_removed.png")

    # Medium objects (default settings)
    result3 = model.remove_object(
        image="image_with_object.jpg",
        mask="object_mask.png",
        border_size=17,  # Standard border
        guidance_scale=7.5
    )
    result3.save("output_object_removed.png")

    print("✓ All results saved!")


def example_cpu_inference():
    """Run inference on CPU (no GPU required)."""

    config = PandoraConfig(
        model_path="stabilityai/stable-diffusion-2-1",
        device="cpu",  # Use CPU
        max_steps=50,
    )

    model = PandoraRemoval(config=config)
    model.load_model()

    result = model.remove_object(
        image="input.jpg",
        mask="mask.png",
        border_size=17
    )

    result.save("output_cpu.png")
    print("✓ CPU inference result saved!")


def main():
    """Run all examples."""
    print("Choose an example to run:")
    print("1. High-quality inference")
    print("2. Fast inference")
    print("3. Different object types")
    print("4. CPU inference")

    choice = input("\nEnter choice (1-4): ")

    examples = {
        "1": example_high_quality,
        "2": example_fast_inference,
        "3": example_different_object_types,
        "4": example_cpu_inference,
    }

    example_func = examples.get(choice)
    if example_func:
        example_func()
    else:
        print("Invalid choice!")


if __name__ == "__main__":
    main()