custom_nodes/ComfyUI-FSampler/sampling/experiment_logger.py

"""Experiment logger for FSampler research data collection."""
import os
import json
import csv
import time
from datetime import datetime
import torch
import numpy as np
from PIL import Image
import folder_paths


class ExperimentLogger:
    """Handles saving FSampler experiment runs to disk with all metadata and outputs."""

    def __init__(self, base_path=None):
        """
        Initialize the experiment logger.

        Args:
            base_path: Base directory for experiments (defaults to ComfyUI/output/experiments/)
                      If empty string, uses default. If provided, uses that exact path.
        """
        if base_path is None or base_path == "":
            # Default to ComfyUI/output/experiments/ using ComfyUI's standard path helper
            output_dir = folder_paths.get_output_directory()
            base_path = os.path.join(output_dir, "experiments")

        self.base_path = base_path
        os.makedirs(self.base_path, exist_ok=True)

    def save_run(self, image_tensor, metadata_json, baseline_image_tensor=None,
                 positive_prompt="", negative_prompt="", is_baseline=False):
        """
        Save a complete experiment run to disk.

        Args:
            image_tensor: Output image from VAE decode (torch tensor or PIL Image)
            metadata_json: JSON string containing FSampler metadata
            baseline_image_tensor: Optional baseline image for comparison
            positive_prompt: Positive prompt string
            negative_prompt: Negative prompt string
            is_baseline: Whether this run is a baseline (no skip mode)

        Returns:
            Path to the created experiment directory
        """
        # Parse metadata
        try:
            metadata = json.loads(metadata_json) if isinstance(metadata_json, str) else metadata_json
        except Exception as e:
            print(f"[ExperimentLogger] Failed to parse metadata: {e}")
            metadata = {}

        # Create timestamped directory
        timestamp_start = metadata.get("timestamp_start", time.time())
        dt = datetime.fromtimestamp(timestamp_start)
        timestamp_str = dt.strftime("%Y-%m-%d_%H-%M-%S")
        unix_timestamp = int(timestamp_start)

        # Extract model type(s)
        model_type_str = ""
        if "model_0" in metadata or "model_1" in metadata:
            # Multi-model workflow
            types = []
            if "model_0" in metadata and metadata["model_0"].get("model_type"):
                types.append(metadata["model_0"]["model_type"])
            if "model_1" in metadata and metadata["model_1"].get("model_type"):
                types.append(metadata["model_1"]["model_type"])
            if types:
                model_type_str = "_" + "+".join(types)
        elif metadata.get("model_type"):
            # Single model workflow
            model_type_str = "_" + metadata["model_type"]

        # Directory name: YYYY-MM-DD_HH-MM-SS_{unix_timestamp}_{model_type}
        dir_name = f"{timestamp_str}_{unix_timestamp}{model_type_str}"
        if is_baseline:
            dir_name += "_baseline"

        exp_dir = os.path.join(self.base_path, dir_name)
        os.makedirs(exp_dir, exist_ok=True)

        # Capture device info
        device_info = self._get_device_info()

        # Add device info and prompts to metadata
        metadata["device_info"] = device_info
        metadata["positive_prompt"] = positive_prompt
        metadata["negative_prompt"] = negative_prompt
        metadata["is_baseline"] = is_baseline

        # Save output image
        output_image_path = os.path.join(exp_dir, "output_image.png")
        self._save_image(image_tensor, output_image_path)

        # Save baseline image if provided
        if baseline_image_tensor is not None:
            baseline_image_path = os.path.join(exp_dir, "baseline.png")
            self._save_image(baseline_image_tensor, baseline_image_path)

            # Compute metrics (SSIM, RMSE, MAE) against baseline
            try:
                from .metrics import compute_metrics
                metrics = compute_metrics(image_tensor, baseline_image_tensor)
                metadata["metrics"] = metrics
                print(f"[ExperimentLogger] Computed metrics: SSIM={metrics['ssim']:.4f}, RMSE={metrics['rmse']:.4f}, MAE={metrics['mae']:.4f}")
            except Exception as e:
                print(f"[ExperimentLogger] Failed to compute metrics: {e}")
                metadata["metrics"] = None

        # Save metadata JSON
        metadata_path = os.path.join(exp_dir, "metadata.json")
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)

        # Save per-step data CSV
        per_step_data = metadata.get("per_step_data", [])
        if per_step_data:
            csv_path = os.path.join(exp_dir, "per_step_data.csv")
            self._save_per_step_csv(per_step_data, csv_path)

        # Save sigmas CSV (if available in per_step_data)
        if per_step_data:
            sigmas_csv_path = os.path.join(exp_dir, "sigmas.csv")
            self._save_sigmas_csv(per_step_data, sigmas_csv_path)

        # Save summary.txt
        summary_path = os.path.join(exp_dir, "summary.txt")
        self._save_summary(metadata, summary_path)

        print(f"[ExperimentLogger] Saved experiment to: {exp_dir}")
        return exp_dir

    def _get_device_info(self):
        """Capture device information (GPU, memory, etc.)."""
        device_info = {}

        try:
            if torch.cuda.is_available():
                device_info["device"] = f"cuda:{torch.cuda.current_device()}"
                device_info["device_name"] = torch.cuda.get_device_name(torch.cuda.current_device())
                device_info["total_memory_gb"] = torch.cuda.get_device_properties(0).total_memory / (1024**3)
                device_info["allocated_memory_gb"] = torch.cuda.memory_allocated(0) / (1024**3)
                device_info["reserved_memory_gb"] = torch.cuda.memory_reserved(0) / (1024**3)
            else:
                device_info["device"] = "cpu"
        except Exception as e:
            device_info["error"] = str(e)

        return device_info

    def _save_image(self, image_data, path):
        """
        Save image tensor or PIL Image to disk.

        Args:
            image_data: Can be torch.Tensor (from VAE) or PIL.Image
            path: Output path for PNG file
        """
        try:
            if isinstance(image_data, torch.Tensor):
                # Convert tensor to PIL Image
                # Assume tensor is in shape [B, C, H, W] or [C, H, W]
                if image_data.dim() == 4:
                    # Batch dimension, take first image
                    image_data = image_data[0]

                # Convert to numpy
                img_np = image_data.cpu().numpy()

                # If channels first, convert to channels last
                if img_np.shape[0] in [1, 3, 4]:
                    img_np = np.transpose(img_np, (1, 2, 0))

                # Clip and convert to uint8
                img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)

                # Convert to PIL
                if img_np.shape[2] == 1:
                    img_pil = Image.fromarray(img_np[:, :, 0], mode='L')
                elif img_np.shape[2] == 3:
                    img_pil = Image.fromarray(img_np, mode='RGB')
                elif img_np.shape[2] == 4:
                    img_pil = Image.fromarray(img_np, mode='RGBA')
                else:
                    raise ValueError(f"Unexpected number of channels: {img_np.shape[2]}")

                img_pil.save(path)
            elif isinstance(image_data, Image.Image):
                # Already PIL Image
                image_data.save(path)
            else:
                print(f"[ExperimentLogger] Unknown image type: {type(image_data)}")
        except Exception as e:
            print(f"[ExperimentLogger] Failed to save image to {path}: {e}")

    def _save_per_step_csv(self, per_step_data, path):
        """Save per-step data to CSV."""
        if not per_step_data:
            return

        try:
            with open(path, 'w', newline='') as f:
                # Get all unique keys from all step dicts
                all_keys = set()
                for step in per_step_data:
                    all_keys.update(step.keys())

                fieldnames = sorted(all_keys)
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(per_step_data)
        except Exception as e:
            print(f"[ExperimentLogger] Failed to save per-step CSV: {e}")

    def _save_sigmas_csv(self, per_step_data, path):
        """Save sigma schedule to CSV."""
        if not per_step_data:
            return

        try:
            with open(path, 'w', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(['step_index', 'sigma_current', 'sigma_next'])
                for step in per_step_data:
                    writer.writerow([
                        step.get('step_index', ''),
                        step.get('sigma_current', ''),
                        step.get('sigma_next', '')
                    ])
        except Exception as e:
            print(f"[ExperimentLogger] Failed to save sigmas CSV: {e}")

    def _save_summary(self, metadata, path):
        """Save human-readable summary.txt."""
        try:
            with open(path, 'w') as f:
                f.write("=" * 80 + "\n")
                f.write("FSampler Experiment Run Summary\n")
                f.write("=" * 80 + "\n\n")

                # Check if this is a multi-model workflow
                is_multi_model = "model_0" in metadata or "model_1" in metadata

                if is_multi_model:
                    # Multi-model workflow
                    f.write("Multi-Model Workflow Detected\n")
                    f.write("=" * 80 + "\n\n")

                    for model_key in ["model_0", "model_1"]:
                        if model_key in metadata:
                            model_meta = metadata[model_key]
                            f.write(f"\n{'='*80}\n")
                            f.write(f"{model_key.upper()} (Base Model)\n" if model_key == "model_0" else f"{model_key.upper()} (Refiner Model)\n")
                            f.write(f"{'='*80}\n\n")
                            self._write_single_model_summary(f, model_meta)

                    # Global info at the bottom
                    f.write("\n" + "="*80 + "\n")
                    f.write("Global Information\n")
                    f.write("="*80 + "\n")
                    f.write(f"Is Baseline: {metadata.get('is_baseline', False)}\n")

                    # Device info
                    device_info = metadata.get("device_info", {})
                    if device_info:
                        f.write("\nDevice Information:\n")
                        f.write("-" * 80 + "\n")
                        for key, value in device_info.items():
                            f.write(f"{key}: {value}\n")

                    # Prompts
                    f.write("\nPrompts:\n")
                    f.write("-" * 80 + "\n")
                    f.write(f"Positive: {metadata.get('positive_prompt', 'N/A')}\n")
                    f.write(f"Negative: {metadata.get('negative_prompt', 'N/A')}\n")

                    # Metrics
                    if "metrics" in metadata and metadata["metrics"]:
                        f.write("\nImage Quality Metrics (vs Baseline):\n")
                        f.write("-" * 80 + "\n")
                        m = metadata["metrics"]
                        f.write(f"SSIM: {m['ssim']:.6f} (structural similarity, higher=better)\n")
                        f.write(f"RMSE: {m['rmse']:.6f} (root mean square error, lower=better)\n")
                        f.write(f"MAE:  {m['mae']:.6f} (mean absolute error, lower=better)\n")

                else:
                    # Single model workflow
                    self._write_single_model_summary(f, metadata)

                    # Device info
                    device_info = metadata.get("device_info", {})
                    if device_info:
                        f.write("Device Information:\n")
                        f.write("-" * 80 + "\n")
                        for key, value in device_info.items():
                            f.write(f"{key}: {value}\n")
                        f.write("\n")

                    # Prompts
                    f.write("Prompts:\n")
                    f.write("-" * 80 + "\n")
                    f.write(f"Positive: {metadata.get('positive_prompt', 'N/A')}\n")
                    f.write(f"Negative: {metadata.get('negative_prompt', 'N/A')}\n")
                    f.write("\n")

                    # Metrics
                    if "metrics" in metadata and metadata["metrics"]:
                        f.write("Image Quality Metrics (vs Baseline):\n")
                        f.write("-" * 80 + "\n")
                        m = metadata["metrics"]
                        f.write(f"SSIM: {m['ssim']:.6f} (structural similarity, higher=better)\n")
                        f.write(f"RMSE: {m['rmse']:.6f} (root mean square error, lower=better)\n")
                        f.write(f"MAE:  {m['mae']:.6f} (mean absolute error, lower=better)\n")
                        f.write("\n")

                f.write("=" * 80 + "\n")
        except Exception as e:
            print(f"[ExperimentLogger] Failed to save summary: {e}")

    def _write_single_model_summary(self, f, metadata):
        """Write summary for a single model's metadata."""
        # Basic info
        f.write("Run Information:\n")
        f.write("-" * 80 + "\n")
        if "timestamp_start" in metadata:
            dt = datetime.fromtimestamp(metadata["timestamp_start"])
            f.write(f"Timestamp: {dt.strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Model Type: {metadata.get('model_type', 'N/A')}\n")
        f.write(f"Seed: {metadata.get('seed', 'N/A')}\n")
        if "is_baseline" in metadata:
            f.write(f"Is Baseline: {metadata.get('is_baseline', False)}\n")
        f.write("\n")

        # Sampler settings
        f.write("Sampler Settings:\n")
        f.write("-" * 80 + "\n")
        f.write(f"Sampler: {metadata.get('sampler', 'N/A')}\n")
        f.write(f"Scheduler: {metadata.get('scheduler', 'N/A')}\n")
        f.write(f"Skip Mode: {metadata.get('skip_mode', 'N/A')}\n")
        f.write(f"Adaptive Mode: {metadata.get('adaptive_mode', 'N/A')}\n")
        f.write(f"Smoothing Beta: {metadata.get('smoothing_beta', 'N/A')}\n")
        f.write(f"Protect First Steps: {metadata.get('protect_first_steps', 'N/A')}\n")
        f.write(f"Protect Last Steps: {metadata.get('protect_last_steps', 'N/A')}\n")
        if metadata.get('anchor_interval') is not None:
            f.write(f"Anchor Interval: {metadata.get('anchor_interval')}\n")
        if metadata.get('max_consecutive_skips') is not None:
            f.write(f"Max Consecutive Skips: {metadata.get('max_consecutive_skips')}\n")
        f.write("\n")

        # Performance metrics
        f.write("Performance Metrics:\n")
        f.write("-" * 80 + "\n")
        f.write(f"Total Steps: {metadata.get('total_steps', 'N/A')}\n")
        f.write(f"Model Calls: {metadata.get('model_calls', 'N/A')}\n")
        f.write(f"Skipped: {metadata.get('skipped', 'N/A')}\n")
        f.write(f"Reduction: {metadata.get('reduction_percent', 0.0):.2f}%\n")
        f.write(f"Total Time: {metadata.get('total_time_seconds', 0.0):.2f}s\n")
        f.write("\n")

        # Learning stabilizer stats
        f.write("Learning Stabilizer (L) Statistics:\n")
        f.write("-" * 80 + "\n")
        f.write(f"L Final: {metadata.get('l_final', 1.0):.4f}\n")
        f.write(f"L Mean: {metadata.get('l_mean', 1.0):.4f}\n")
        f.write(f"L Min: {metadata.get('l_min', 1.0):.4f}\n")
        f.write(f"L Max: {metadata.get('l_max', 1.0):.4f}\n")
        f.write("\n")