run_methods_comparison.py

#!/usr/bin/env python
"""
Compare different editing methods on the Messi->Ronaldo transformation.
Tests: null-text-inversion+p2p, negative-prompt-inversion+p2p, and directinversion+p2p variants.
"""
import os
import sys
import torch
import numpy as np
from PIL import Image
from models.p2p_editor import P2PEditor


def load_and_prepare_image(image_path, target_size=512):
    """Load image and resize to 512x512."""
    img = Image.open(image_path).convert('RGB')
    img = img.resize((target_size, target_size), Image.Resampling.LANCZOS)
    return img


def main():
    # Configuration
    image_path = "data/messi.webp"
    prompt_src = "Lionel Messi in Argentina jersey lifting the FIFA world cup"
    prompt_tar = "Arturo Vidal in Chile jersey lifting the FIFA world cup"

    # Methods to compare
    methods = [
        {
            'name': 'null-text-inversion+p2p',
            'description': 'Null-Text Inversion + P2P',
            'params': {}
        },
        {
            'name': 'negative-prompt-inversion+p2p',
            'description': 'Negative Prompt Inversion + P2P',
            'params': {}
        },
        {
            'name': 'directinversion+p2p',
            'description': 'DirectInversion + P2P (α=0, β=0.9)',
            'params': {'alpha': 0.0, 'beta': 0.9, 'alpha_sch': 'exp'}
        },
        {
            'name': 'directinversion+p2p',
            'description': 'DirectInversion + P2P (α=0.5, β=0.9)',
            'params': {'alpha': 0.5, 'beta': 0.9, 'alpha_sch': 'exp'}
        },
        {
            'name': 'directinversion+p2p',
            'description': 'DirectInversion + P2P (α=1.0, β=0.9)',
            'params': {'alpha': 1.0, 'beta': 0.9, 'alpha_sch': 'exp'}
        },
    ]

    # Output directory
    base_output_dir = "outputs_methods_comparison"
    os.makedirs(base_output_dir, exist_ok=True)

    # Check if image exists
    if not os.path.exists(image_path):
        print(f"Error: Image not found at {image_path}")
        return

    print(f"Loading image from {image_path}")

    # Convert webp to jpg and save a processed version
    processed_image_path = os.path.join(base_output_dir, "messi_original.jpg")
    img = load_and_prepare_image(image_path)
    img.save(processed_image_path)
    print(f"Saved processed image to {processed_image_path}")

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

    # Get all unique method names for initialization
    unique_methods = list(set([m['name'] for m in methods]))
    p2p_editor = P2PEditor(unique_methods, device, num_ddim_steps=50)

    print("\n" + "="*100)
    print(f"SOURCE PROMPT: {prompt_src}")
    print(f"TARGET PROMPT: {prompt_tar}")
    print(f"Methods to compare: {len(methods)}")
    print("="*100 + "\n")

    # Run each method
    total_runs = len(methods)

    for idx, method_config in enumerate(methods, 1):
        method_name = method_config['name']
        description = method_config['description']
        params = method_config['params']

        print(f"\n[{idx}/{total_runs}] Running: {description}")
        print("-" * 80)

        # Create safe folder name
        folder_name = description.replace(' ', '_').replace('(', '').replace(')', '').replace(',', '').replace('=', '')
        output_dir = os.path.join(base_output_dir, folder_name)
        os.makedirs(output_dir, exist_ok=True)

        output_image_path = os.path.join(output_dir, "edited_image.jpg")

        # Skip if already exists
        if os.path.exists(output_image_path):
            print(f"Skipping (already exists): {output_image_path}")
            continue

        try:
            # Prepare parameters
            edit_params = {
                'image_path': processed_image_path,
                'prompt_src': prompt_src,
                'prompt_tar': prompt_tar,
                'guidance_scale': 7.5,
                'cross_replace_steps': 0.4,
                'self_replace_steps': 0.6,
                'blend_word': None,
                'eq_params': None,
            }

            # Add method-specific parameters
            if method_name == 'directinversion+p2p':
                edit_params.update({
                    'proximal': "l0",
                    'quantile': 0.75,
                    'use_inversion_guidance': True,
                    'recon_lr': 1,
                    'recon_t': 400,
                    'alpha': params.get('alpha', 1.0),
                    'beta': params.get('beta', 1.0),
                    'alpha_sch': params.get('alpha_sch', 'exp'),
                })
            elif method_name in ['negative-prompt-inversion+p2p']:
                edit_params.update({
                    'proximal': "l0",
                    'quantile': 0.75,
                    'use_inversion_guidance': True,
                    'recon_lr': 1,
                    'recon_t': 400,
                })

            # Run editing
            torch.cuda.empty_cache()

            result = p2p_editor(method_name, **edit_params)

            # Handle different return types
            if isinstance(result, tuple):
                if len(result) == 2:
                    edited_image, output_dicts = result
                else:
                    edited_image = result[0]
                    output_dicts = None
            else:
                edited_image = result
                output_dicts = None

            # Save edited image
            edited_image.save(output_image_path)
            print(f"✓ Saved to: {output_image_path}")

            # Save output metrics plot if available
            if output_dicts:
                import matplotlib.pyplot as plt
                for key, values in output_dicts.items():
                    plot_path = os.path.join(output_dir, f"plot_{key}.png")
                    plt.figure(figsize=(10, 6))

                    # Convert tensor values to floats
                    vals = [v.item() if isinstance(v, torch.Tensor) else v for v in values]

                    plt.plot(vals, linewidth=2, marker='o')
                    plt.title(f'{key} - {description}', fontsize=14, fontweight='bold')
                    plt.xlabel('Step', fontsize=12)
                    plt.ylabel('Value', fontsize=12)
                    plt.grid(True, alpha=0.3)
                    plt.tight_layout()
                    plt.savefig(plot_path, dpi=150)
                    plt.close()

        except Exception as e:
            print(f"✗ Error: {e}")
            import traceback
            traceback.print_exc()
            continue

    print("\n" + "="*100)
    print("METHODS COMPARISON COMPLETE!")
    print("="*100)
    print(f"\nResults saved to: {base_output_dir}/")
    print(f"Total methods tested: {total_runs}")
    print("\nTo view results, check the folders:")
    for method_config in methods:
        description = method_config['description']
        folder_name = description.replace(' ', '_').replace('(', '').replace(')', '').replace(',', '').replace('=', '')
        print(f"  - {base_output_dir}/{folder_name}/edited_image.jpg")


if __name__ == "__main__":
    main()