sample.py

import os
import sys
import torch
import torch.nn.functional as F
import json
import numpy as np
import cv2
import matplotlib.pyplot as plt
from torchvision.transforms.functional import to_pil_image
from torchvision import transforms
from tqdm import tqdm

from ldm.models.diffusion.ddim import DDIMSampler
from src.utils import load_model_from_config, apply_lora_to_model, set_seed, load_config
from src.classifier import Classifier, CIFAR10_CLASSES
from src.sampling import sample_model
from src.masking import mask as generate_mask


def get_cifar10_class_index(class_name):
    """
    Map a class name to its CIFAR-10 index.
    
    Args:
        class_name: Class name to map (e.g., 'dog', 'cat')
        
    Returns:
        int: CIFAR-10 class index (0-9), or -1 if not found
    """
    class_name = class_name.lower().strip()
    try:
        return CIFAR10_CLASSES.index(class_name)
    except ValueError:
        print(f"Warning: '{class_name}' not found in CIFAR-10 classes: {CIFAR10_CLASSES}")
        return -1


def generate_image(model, classifier, prompt, config, seed=0, apply_masking=False, target_cifar10_idx=None):
    """
    Generate image using configured sampling parameters.
    
    Args:
        model: Diffusion model
        classifier: Classifier for mask generation
        prompt: Text prompt
        config: Configuration dictionary
        seed: Random seed
        apply_masking: Whether to apply masking
        target_cifar10_idx: CIFAR-10 class index to minimize (required if apply_masking=True)
        
    Returns:
        tuple: (original PIL Image, perturbed PIL Image or None, mask array or None)
    """
    set_seed(seed)
    
    with torch.no_grad():
        sampler = DDIMSampler(model)
        cond = model.get_learned_conditioning([prompt])
        
        shape = (1, 4, 64, 64)
        start_code = torch.randn(
            shape, 
            generator=torch.Generator(device=model.device).manual_seed(seed), 
            device=model.device
        )
        
        samples = sample_model(
            model,
            sampler,
            cond,
            config['sampling']['image_size'],
            config['sampling']['image_size'],
            config['sampling']['ddim_steps'],
            config['sampling']['guidance_scale'],
            config['sampling']['ddim_eta'],
            start_code=start_code,
            verbose=False,
        )
        
        decoded = model.decode_first_stage(samples)
        decoded = (decoded + 1.0) / 2.0
        decoded = torch.clamp(decoded, 0.0, 1.0)
        
        decoded_np = decoded[0].detach().cpu().permute(1, 2, 0).numpy()
    
    with torch.enable_grad():
        img_original = to_pil_image((decoded_np * 255).astype("uint8"))
        
        img_perturbed = None
        mask_array = None
        
        if apply_masking and classifier is not None:
            if target_cifar10_idx is None:
                raise ValueError("target_cifar10_idx must be provided when apply_masking=True")
            
            img_numpy = (decoded_np * 255).astype(np.uint8)
            
            mask_array, perturbed_numpy = generate_mask(
                classifier,
                img_numpy,
                target_category=target_cifar10_idx,
                tv_beta=config['mask'].get('tv_beta', 3),
                learning_rate=config['mask'].get('learning_rate', 0.1),
                max_iterations=config['mask'].get('max_iterations', 500),
                l1_coeff=config['mask'].get('l1_coeff', 0.01),
                tv_coeff=config['mask'].get('tv_coeff', 0.2)
            )
            
            # Resize mask and perturbed image back to original resolution
            mask_array = cv2.resize(mask_array, (config['sampling']['image_size'], config['sampling']['image_size']))
            perturbed_numpy = cv2.resize(perturbed_numpy, (config['sampling']['image_size'], config['sampling']['image_size']))
            
            img_perturbed = to_pil_image((perturbed_numpy * 255).astype("uint8"))
    
    return img_original, img_perturbed, mask_array


def extract_class(prompt):
    """Extract class name from prompt."""
    prompt = prompt.lower().strip()
    for prefix in ["a photo of the ", "a photo of a ", "a photo of "]:
        if prompt.startswith(prefix):
            return prompt[len(prefix):].strip()
    return prompt


def get_classifier_predictions(classifier, img_pil, device):
    """
    Get CIFAR-10 predictions from ResNet50 classifier.
    
    Args:
        classifier: ResNet50 classifier model
        img_pil: PIL Image
        device: torch device
        
    Returns:
        numpy array of softmax probabilities for all 10 CIFAR-10 classes
    """
    # Prepare image for CIFAR-10 classifier
    # CIFAR-10 uses 32x32 images, but timm models typically use larger sizes
    preprocess = transforms.Compose([
        transforms.Resize((224, 224)),  # ResNet50 expects 224x224
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    
    img_tensor = preprocess(img_pil).unsqueeze(0).to(device)
    
    with torch.no_grad():
        logits = classifier(img_tensor)
        probs = F.softmax(logits, dim=1)
        
    return probs[0].cpu().numpy()  # Returns (10,) array


def create_comparison_plot(vanilla_metrics, lora_metrics, output_path, cifar10_classes):
    """
    Create visualization comparing CIFAR-10 classifier predictions before/after LoRA.
    
    Args:
        vanilla_metrics: Dict with vanilla model metrics
        lora_metrics: Dict with LoRA model metrics
        output_path: Path to save the plot
        cifar10_classes: List of CIFAR-10 class names
    """
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    
    prompt_classes = list(vanilla_metrics.keys())
    
    # Plot 1: Heatmap of average probabilities for each CIFAR-10 class (Vanilla)
    ax1 = axes[0, 0]
    vanilla_probs_matrix = np.zeros((len(prompt_classes), 10))
    for i, class_name in enumerate(prompt_classes):
        vanilla_probs_matrix[i] = np.mean(vanilla_metrics[class_name]['cifar10_probs'], axis=0)
    
    im1 = ax1.imshow(vanilla_probs_matrix, aspect='auto', cmap='YlOrRd', vmin=0, vmax=1)
    ax1.set_xlabel('CIFAR-10 Classes')
    ax1.set_ylabel('Prompt Classes')
    ax1.set_title('Vanilla: Average CIFAR-10 Predictions')
    ax1.set_xticks(range(10))
    ax1.set_xticklabels(cifar10_classes, rotation=45, ha='right')
    ax1.set_yticks(range(len(prompt_classes)))
    ax1.set_yticklabels(prompt_classes)
    plt.colorbar(im1, ax=ax1, label='Probability')
    
    # Plot 2: Heatmap for LoRA
    ax2 = axes[0, 1]
    lora_probs_matrix = np.zeros((len(prompt_classes), 10))
    for i, class_name in enumerate(prompt_classes):
        lora_probs_matrix[i] = np.mean(lora_metrics[class_name]['cifar10_probs'], axis=0)
    
    im2 = ax2.imshow(lora_probs_matrix, aspect='auto', cmap='YlOrRd', vmin=0, vmax=1)
    ax2.set_xlabel('CIFAR-10 Classes')
    ax2.set_ylabel('Prompt Classes')
    ax2.set_title('LoRA: Average CIFAR-10 Predictions')
    ax2.set_xticks(range(10))
    ax2.set_xticklabels(cifar10_classes, rotation=45, ha='right')
    ax2.set_yticks(range(len(prompt_classes)))
    ax2.set_yticklabels(prompt_classes)
    plt.colorbar(im2, ax=ax2, label='Probability')
    
    # Plot 3: Difference heatmap (LoRA - Vanilla)
    ax3 = axes[1, 0]
    diff_matrix = lora_probs_matrix - vanilla_probs_matrix
    im3 = ax3.imshow(diff_matrix, aspect='auto', cmap='RdBu_r', vmin=-0.5, vmax=0.5)
    ax3.set_xlabel('CIFAR-10 Classes')
    ax3.set_ylabel('Prompt Classes')
    ax3.set_title('Difference: LoRA - Vanilla')
    ax3.set_xticks(range(10))
    ax3.set_xticklabels(cifar10_classes, rotation=45, ha='right')
    ax3.set_yticks(range(len(prompt_classes)))
    ax3.set_yticklabels(prompt_classes)
    plt.colorbar(im3, ax=ax3, label='Probability Change')
    
    # Plot 4: Top predicted class confidence comparison
    ax4 = axes[1, 1]
    vanilla_top_conf = [np.mean([max(probs) for probs in vanilla_metrics[c]['cifar10_probs']]) 
                        for c in prompt_classes]
    lora_top_conf = [np.mean([max(probs) for probs in lora_metrics[c]['cifar10_probs']]) 
                     for c in prompt_classes]
    
    x = np.arange(len(prompt_classes))
    width = 0.35
    ax4.bar(x - width/2, vanilla_top_conf, width, label='Vanilla', alpha=0.7)
    ax4.bar(x + width/2, lora_top_conf, width, label='LoRA', alpha=0.7)
    ax4.set_xlabel('Prompt Classes')
    ax4.set_ylabel('Top-1 Confidence')
    ax4.set_title('Top-1 Confidence Comparison')
    ax4.set_xticks(x)
    ax4.set_xticklabels(prompt_classes, rotation=45, ha='right')
    ax4.legend()
    ax4.grid(True, alpha=0.3, axis='y')
    ax4.set_ylim(0, 1)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Comparison plot saved to: {output_path}")


def create_metrics_plot(metrics_data, output_path, cifar10_classes):
    """
    Create visualization of CIFAR-10 classifier predictions per class.
    
    Args:
        metrics_data: Dict with structure {class_name: {'cifar10_probs': []}}
        output_path: Path to save the plot
        cifar10_classes: List of CIFAR-10 class names
    """
    prompt_classes = list(metrics_data.keys())
    
    fig, axes = plt.subplots(2, 1, figsize=(14, 10))
    
    # Plot 1: Heatmap of average probabilities
    ax1 = axes[0]
    probs_matrix = np.zeros((len(prompt_classes), 10))
    for i, class_name in enumerate(prompt_classes):
        probs_matrix[i] = np.mean(metrics_data[class_name]['cifar10_probs'], axis=0)
    
    im = ax1.imshow(probs_matrix, aspect='auto', cmap='YlOrRd', vmin=0, vmax=1)
    ax1.set_xlabel('CIFAR-10 Classes')
    ax1.set_ylabel('Prompt Classes')
    ax1.set_title('Average CIFAR-10 Classifier Predictions')
    ax1.set_xticks(range(10))
    ax1.set_xticklabels(cifar10_classes, rotation=45, ha='right')
    ax1.set_yticks(range(len(prompt_classes)))
    ax1.set_yticklabels(prompt_classes)
    plt.colorbar(im, ax=ax1, label='Probability')
    
    # Add text annotations for probabilities > 0.1
    for i in range(len(prompt_classes)):
        for j in range(10):
            if probs_matrix[i, j] > 0.1:
                text = ax1.text(j, i, f'{probs_matrix[i, j]:.2f}',
                               ha="center", va="center", color="black", fontsize=8)
    
    # Plot 2: Top-1 confidence per prompt class
    ax2 = axes[1]
    top1_means = []
    top1_stds = []
    
    for class_name in prompt_classes:
        probs_array = np.array(metrics_data[class_name]['cifar10_probs'])  # (num_seeds, 10)
        top1_confs = np.max(probs_array, axis=1)  # Max prob for each seed
        top1_means.append(np.mean(top1_confs))
        top1_stds.append(np.std(top1_confs))
    
    x = np.arange(len(prompt_classes))
    ax2.bar(x, top1_means, yerr=top1_stds, capsize=5, alpha=0.7, color='steelblue')
    ax2.set_xlabel('Prompt Classes')
    ax2.set_ylabel('Top-1 Confidence')
    ax2.set_title('Top-1 Prediction Confidence per Prompt Class')
    ax2.set_xticks(x)
    ax2.set_xticklabels(prompt_classes, rotation=45, ha='right')
    ax2.grid(True, alpha=0.3, axis='y')
    ax2.set_ylim(0, 1)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Metrics plot saved to: {output_path}")


def main():
    """Sampling function with classifier-based masking."""
    if len(sys.argv) != 2:
        print("Usage: python sample.py <config_path>")
        sys.exit(1)
    
    config = load_config(sys.argv[1])
    
    print("=== Sampling Configuration ===")
    print(f"Config: {sys.argv[1]}")
    print(f"Model: {config['model']['ckpt_path']}")
    print(f"Data: {config['data']['csv_path']}")
    print(f"Output: {config['output_dir']}")
    print("=" * 40)
    
    os.makedirs(config['output_dir'], exist_ok=True)
    images_dir = os.path.join(config['output_dir'], "imgs")
    os.makedirs(images_dir, exist_ok=True)
    
    lora_path = os.path.join(config['output_dir'], "lora.pth")
    if not os.path.exists(lora_path):
        print(f"Warning: LoRA weights not found at {lora_path}")
        lora_state_dict = None
    else:
        lora_state_dict = torch.load(lora_path, map_location=config['device'])
        print(f"Loaded LoRA weights from {lora_path}")
    
    with open(config['data']['csv_path'], 'r') as f:
        prompts_data = json.load(f)
    
    target_prompt = prompts_data['target']
    target_class = extract_class(target_prompt)
    synonyms = prompts_data.get('synonyms', [])
    other_prompts = [p for p in prompts_data.get('other', []) if p.strip()]
    
    # Generate combined prompts: "a picture of <target> and <other>"
    combined_prompts = []
    for other_prompt in other_prompts:
        if other_prompt.strip():
            other_class = extract_class(other_prompt)
            combined_prompt = f"a picture of {target_class} and {other_class}"
            combined_prompts.append(combined_prompt)
    
    # Get CIFAR-10 index for target class
    target_cifar10_idx = get_cifar10_class_index(target_class)
    if target_cifar10_idx == -1:
        raise ValueError(f"Target class '{target_class}' not found in CIFAR-10 classes. "
                        f"Available classes: {CIFAR10_CLASSES}")
    print(f"Target class: {target_class} (CIFAR-10 index: {target_cifar10_idx})")
    
    eval_prompts = {
        'target': [target_prompt],
        'synonyms': synonyms,
        'others': other_prompts,
        'combined': combined_prompts
    }
    
    all_eval_prompts = []
    all_eval_classes = []
    all_categories = []
    
    for category, prompts in eval_prompts.items():
        for prompt in prompts:
            all_eval_prompts.append(prompt)
            all_eval_classes.append(extract_class(prompt))
            all_categories.append(category)
    
    print(f"\nEvaluation Set:")
    print(f"  Target: {target_prompt}")
    print(f"  Synonyms: {len(synonyms)}")
    print(f"  Others: {len(other_prompts)}")
    print(f"  Combined: {len(combined_prompts)}")
    print(f"  Total prompts: {len(all_eval_prompts)}")
    
    print("\nStarting inference...")
    
    variants = [
        ("vanilla", False, False),
        ("vanilla_masked", False, True),
    ]
    if lora_state_dict is not None:
        variants.extend([
            ("lora", True, False),
            ("lora_masked", True, True),
        ])
    
    classifier = Classifier().to(config['device'])
    classifier.eval()
    for param in classifier.parameters():
        param.requires_grad = False
    print("ResNet18 CIFAR-10 classifier loaded")
    
    # Get CIFAR-10 class names
    cifar10_classes = Classifier.get_class_names()
    print(f"CIFAR-10 classes: {cifar10_classes}")
    
    num_seeds = config.get('evaluation', {}).get('num_seeds', 10)
    
    print(f"\nGenerating images for {len(all_eval_prompts)} prompts x {num_seeds} seeds x {len(variants)} variants...")
    
    all_metrics = {}
    
    for variant_name, use_lora, apply_masking in variants:
        print(f"\n{'='*60}")
        print(f"Testing variant: {variant_name}")
        print(f"{'='*60}")
        
        print(f"Loading diffusion model...")
        model = load_model_from_config(
            config['model']['config_path'], 
            config['model']['ckpt_path'], 
            config['device']
        )
        
        if use_lora and lora_state_dict is not None:
            apply_lora_to_model(
                model.model.diffusion_model, 
                lora_state_dict, 
                alpha=config['lora']['alpha']
            )
            print("Applied LoRA weights to model")
        
        model.eval()
        
        metrics_data = {class_name: {'cifar10_probs': []} 
                        for class_name in all_eval_classes}
        
        for prompt_idx, (prompt, class_name, category) in enumerate(zip(all_eval_prompts, all_eval_classes, all_categories)):
            print(f"\n[{prompt_idx+1}/{len(all_eval_prompts)}] {category}: {prompt}")
            
            for seed in tqdm(range(num_seeds), desc=f"Seeds for {class_name}"):
                img_original, img_perturbed, mask_array = generate_image(
                    model, classifier if apply_masking else None, 
                    prompt, config, seed, 
                    apply_masking=apply_masking,
                    target_cifar10_idx=target_cifar10_idx if apply_masking else None
                )
                
                # Get CIFAR-10 classifier predictions
                cifar10_probs = get_classifier_predictions(
                    classifier, img_original, config['device']
                )
                metrics_data[class_name]['cifar10_probs'].append(cifar10_probs.tolist())
                
                save_name = f"{variant_name}_{category}_{class_name}_{seed:03d}.jpg"
                img_original.save(os.path.join(images_dir, save_name))
                
                if img_perturbed is not None:
                    save_name_masked = f"{variant_name}_{category}_{class_name}_{seed:03d}_masked.jpg"
                    img_perturbed.save(os.path.join(images_dir, save_name_masked))
        
        all_metrics[variant_name] = metrics_data
        
        del model
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
    
    print("\nGenerating metrics plots and summaries...")
    
    for variant_name, metrics_data in all_metrics.items():
        print(f"\nProcessing variant: {variant_name}")
        
        metrics_plot_path = os.path.join(config['output_dir'], f"evaluation_metrics_{variant_name}.png")
        create_metrics_plot(metrics_data, metrics_plot_path, cifar10_classes)
        
        metrics_json_path = os.path.join(config['output_dir'], f"evaluation_metrics_{variant_name}.json")
        metrics_summary = {}
        for class_name, data in metrics_data.items():
            probs_array = np.array(data['cifar10_probs'])  # (num_seeds, 10)
            mean_probs = probs_array.mean(axis=0)
            top1_idx = np.argmax(mean_probs)
            
            metrics_summary[class_name] = {
                'mean_cifar10_probs': mean_probs.tolist(),
                'top1_class': cifar10_classes[top1_idx],
                'top1_confidence': float(mean_probs[top1_idx]),
                'top1_confidence_std': float(probs_array[:, top1_idx].std()),
            }
        
        with open(metrics_json_path, 'w') as f:
            json.dump(metrics_summary, f, indent=2)
        
        print("\n" + "=" * 80)
        print(f"EVALUATION SUMMARY - {variant_name.upper()}")
        print("=" * 80)
        for class_name in all_eval_classes:
            data = metrics_summary[class_name]
            print(f"\n{class_name}:")
            print(f"  Top-1 CIFAR-10 Class: {data['top1_class']}")
            print(f"  Top-1 Confidence: {data['top1_confidence']:.3f} ± {data['top1_confidence_std']:.3f}")
        print("\n" + "=" * 80)
    
    if 'vanilla' in all_metrics and 'lora' in all_metrics:
        print("\nGenerating LoRA comparison plot...")
        comparison_plot_path = os.path.join(config['output_dir'], "lora_comparison.png")
        create_comparison_plot(all_metrics['vanilla'], all_metrics['lora'], comparison_plot_path, cifar10_classes)
        
        comparison_summary = {}
        for class_name in all_eval_classes:
            vanilla_probs = np.array(all_metrics['vanilla'][class_name]['cifar10_probs'])
            lora_probs = np.array(all_metrics['lora'][class_name]['cifar10_probs'])
            
            vanilla_mean = vanilla_probs.mean(axis=0)
            lora_mean = lora_probs.mean(axis=0)
            
            vanilla_top1_idx = np.argmax(vanilla_mean)
            lora_top1_idx = np.argmax(lora_mean)
            
            comparison_summary[class_name] = {
                'vanilla_top1_class': cifar10_classes[vanilla_top1_idx],
                'vanilla_top1_confidence': float(vanilla_mean[vanilla_top1_idx]),
                'lora_top1_class': cifar10_classes[lora_top1_idx],
                'lora_top1_confidence': float(lora_mean[lora_top1_idx]),
                'confidence_change': float(lora_mean[lora_top1_idx] - vanilla_mean[vanilla_top1_idx]),
                'vanilla_mean_probs': vanilla_mean.tolist(),
                'lora_mean_probs': lora_mean.tolist(),
            }
        
        comparison_json_path = os.path.join(config['output_dir'], "lora_comparison.json")
        with open(comparison_json_path, 'w') as f:
            json.dump(comparison_summary, f, indent=2)
        
        print("\n" + "=" * 80)
        print("LORA IMPACT SUMMARY")
        print("=" * 80)
        for class_name in all_eval_classes:
            data = comparison_summary[class_name]
            print(f"\n{class_name}:")
            print(f"  Vanilla: {data['vanilla_top1_class']} ({data['vanilla_top1_confidence']:.3f})")
            print(f"  LoRA:    {data['lora_top1_class']} ({data['lora_top1_confidence']:.3f})")
            print(f"  Change:  {data['confidence_change']:+.3f}")
        print("\n" + "=" * 80)
    
    del classifier
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    
    print("\nSampling completed!")
    print(f"Images saved to: {images_dir}")
    print(f"Metrics plots and JSON files saved to: {config['output_dir']}")


if __name__ == "__main__":
    main()