external/__init__.py

# CTM Wrapper Module for MorphGuard
# This module provides a bridge to Sakana AI's Continuous Thought Machine
# Until CTM is cloned, this provides the interface and mock for development

"""
Sakana AI CTM Interface for MorphGuard

This module provides compatibility layer for CTM integration.
When the full CTM is installed, import from external.sakana_ctm
For development/testing, this provides mock implementations.
"""

import torch
import torch.nn as nn
from typing import Optional, Dict, Any, List, Tuple
import numpy as np

# Check if full CTM is available
try:
    import sys
    import os
    ctm_path = os.path.join(os.path.dirname(__file__), 'continuous-thought-machines')
    if os.path.exists(ctm_path):
        sys.path.insert(0, ctm_path)
        from models.ctm import ContinuousThoughtMachine
        CTM_AVAILABLE = True
    else:
        CTM_AVAILABLE = False
except ImportError:
    CTM_AVAILABLE = False


class CTMConfig:
    """Configuration for CTM model"""
    def __init__(
        self,
        hidden_size: int = 512,
        num_neurons: int = 256,
        num_heads: int = 8,
        history_length: int = 16,
        max_thoughts: int = 16,
        sync_window: int = 8,
        dropout: float = 0.1
    ):
        self.hidden_size = hidden_size
        self.num_neurons = num_neurons
        self.num_heads = num_heads
        self.history_length = history_length
        self.max_thoughts = max_thoughts
        self.sync_window = sync_window
        self.dropout = dropout


class CTMOutput:
    """Output structure from CTM forward pass"""
    def __init__(
        self,
        prediction: int,
        confidence: float,
        synchronization: torch.Tensor,
        attention_map: Optional[torch.Tensor] = None,
        neuron_history: Optional[torch.Tensor] = None
    ):
        self.prediction = prediction
        self.confidence = confidence
        self.synchronization = synchronization
        self.attention_map = attention_map
        self.neuron_history = neuron_history


class CTMWrapper(nn.Module):
    """
    Wrapper around the real Sakana AI Continuous Thought Machine.
    Adapts the CTM API to the MorphGuard Forensic Agent interface.
    """
    def __init__(self, ctm_model):
        super().__init__()
        self.model = ctm_model
        # Map config from CTM
        self.config = CTMConfig(
            hidden_size=ctm_model.d_model,
            max_thoughts=ctm_model.iterations
        )
        
    def forward(
        self, 
        image: torch.Tensor,
        max_thoughts: Optional[int] = None,
        early_exit_threshold: float = 0.99
    ) -> Tuple[CTMOutput, List[torch.Tensor]]:
        """
        Forward pass using real CTM.
        Note: Real CTM fixed iteration count is preferred for consistency in this version.
        Early exit is technically possible but for forensic analysis we want the full trace.
        """
        # Ensure image is on correct device and type
        if image.device != next(self.model.parameters()).device:
            image = image.to(next(self.model.parameters()).device)
            
        # CTM expects images compatible with its backbone (usually 224x224)
        # Assuming input is already transformed
        
        # Run CTM with tracking enabled
        # Returns: preds, certs, (synch_out, synch_action), pre, post, attention
        with torch.no_grad():
            outputs = self.model(image, track=True)
            
        predictions, certainties, syncs, _, _, attention_stack = outputs
        
        # Parse outputs
        # predictions: (B, out_dims, T)
        # certainties: (B, 2, T)
        # attention_stack: (T, B, heads, patches) -> need to process
        
        final_idx = -1
        final_pred_logits = predictions[:, :, final_idx]
        final_certainty = certainties[:, 0, final_idx] # 0 is normalized entropy (certainty)
        
        # Sync is a tuple (out, action), we use 'out'
        sync_out = torch.from_numpy(syncs[0][final_idx]) if isinstance(syncs[0], np.ndarray) else syncs[0][..., final_idx]
        if isinstance(sync_out, np.ndarray):
             sync_out = torch.from_numpy(sync_out).to(image.device)
        
        # Process attention history
        # Stack is numpy array (T, B, heads, patches) or (T, heads, patches) depending on impl params
        # We want list of tensors for visualization
        attention_history = []
        for t in range(attention_stack.shape[0]):
            try:
                # Average over heads (dim 1 or 2 depending on batch) since we want general attention map
                # Assuming batch size 1 for forensic agent usually
                attn_step = attention_stack[t] # (B, heads, patches)
                if isinstance(attn_step, np.ndarray):
                    attn_step = torch.from_numpy(attn_step)
                
                # Average heads: (B, patches)
                attn_avg = attn_step.mean(dim=1) 
                
                # Reshape patches to spatial map
                # Patches = (H/16 * W/16)
                # For 224x224, 14x14 = 196
                side = int(np.sqrt(attn_avg.shape[-1]))
                attn_map = attn_avg.reshape(attn_avg.shape[0], side, side)
                attention_history.append(attn_map.cpu())
            except Exception as e:
                # Fallback for shape mismatch
                print(f"Warning: Attention tracking shape mismatch: {e}")
                continue

        # Create output object
        # Prediction: Argmax of logits (0 or 1 for morph/real)
        # CTM ImageNet has 1000 classes. We need to map/finetune. 
        # For now (Option 1 - Pretrained), we might be getting unrelated classes.
        # Ideally we replace the head or fine-tune.
        # But user asked to "proceed to download" and "fine-tune setup (future step)".
        # So for now we just return the raw prediction index.
        pred_idx = final_pred_logits.argmax(dim=-1).item()
        confidence = final_certainty.mean().item() # Certainty score
        
        output = CTMOutput(
            prediction=pred_idx,
            confidence=confidence,
            synchronization=sync_out,
            attention_map=attention_history[-1] if attention_history else None,
            neuron_history=None
        )
        
        return output, attention_history

    @classmethod
    def load_from_checkpoint(cls, checkpoint_path: str, device: str = 'cuda') -> 'CTMWrapper':
        """
        Load CTM model.
        If checkpoint_path looks like a repo ID, uses from_pretrained.
        """
        if "/" in checkpoint_path and not os.path.exists(checkpoint_path):
            # Assume HF Repo ID
            print(f"Loading CTM from HuggingFace Hub: {checkpoint_path}")
            model = ContinuousThoughtMachine.from_pretrained(checkpoint_path, map_location=device)
        else:
            # Load local file
            print(f"Loading local CTM checkpoint: {checkpoint_path}")
            # This path logic depends on how local loading is implemented in CTM repo
            # Typically CTM.load(path)
            # But the CTM class doesn't have a simple load method for local .pt files in the version I saw?
            # It had _from_pretrained.
            # We'll try standard torch load if it's a file
            checkpoint = torch.load(checkpoint_path, map_location=device)
            # Reconstruct args? This is hard without config.
            # Safer to fall back to 'SakanaAI/ctm-imagenet' if local fails or just use the class method if supported
             # Creating dummy CTM to load state dict
            config = CTMConfig() # Default
            # Attempt to instantiate with defaults (might fail if shape mismatch)
            # Ideally we pick a default config matching ImageNet
            model = ContinuousThoughtMachine(
                iterations=16, 
                d_model=512, 
                d_input=128, 
                heads=8, 
                n_synch_out=512, 
                n_synch_action=512, 
                synapse_depth=4, 
                memory_length=25, 
                deep_nlms=True, 
                memory_hidden_dims=4, 
                do_layernorm_nlm=False, 
                backbone_type='resnet18-4', 
                positional_embedding_type='none', 
                out_dims=2
            ) # Default params for ImageNet CTM roughly
            
            if 'model_state_dict' in checkpoint:
                 model.load_state_dict(checkpoint['model_state_dict'])
            else:
                 model.load_state_dict(checkpoint)
        
        model.to(device)
        model.eval()
        return cls(model)

# Export the wrapper as the main model class
CTMModel = CTMWrapper


def plot_attention_video(
    attention_history: List[torch.Tensor],
    original_image: torch.Tensor,
    output_path: str = 'attention_scan.gif',
    fps: int = 4
) -> str:
    """
    Generate a GIF showing the attention progression during CTM thinking.
    
    Args:
        attention_history: List of attention maps from each thinking step
        original_image: Original input image tensor (B, C, H, W)
        output_path: Path to save the output GIF
        fps: Frames per second for the GIF
    
    Returns:
        Path to the generated GIF
    """
    try:
        import imageio
        from PIL import Image
        import matplotlib.pyplot as plt
        import matplotlib.cm as cm
    except ImportError:
        print("Warning: imageio/matplotlib not available for video generation")
        return None
    
    frames = []
    
    # Convert original image to numpy
    if isinstance(original_image, torch.Tensor):
        img_np = original_image[0].permute(1, 2, 0).cpu().numpy()
        # Denormalize if needed (assuming ImageNet normalization)
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        img_np = std * img_np + mean
        img_np = np.clip(img_np, 0, 1)
    else:
        img_np = original_image
    
    for i, attention in enumerate(attention_history):
        # Create figure
        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
        
        # Show original image
        ax.imshow(img_np)
        
        # Overlay attention heatmap
        if isinstance(attention, torch.Tensor):
            attn_np = attention[0].numpy() if attention.dim() == 3 else attention.numpy()
        else:
            attn_np = attention
        
        # Resize attention to image size
        attn_resized = np.array(Image.fromarray(
            (attn_np * 255).astype(np.uint8)
        ).resize((img_np.shape[1], img_np.shape[0]), Image.BILINEAR)) / 255.0
        
        # Apply colormap
        heatmap = cm.jet(attn_resized)[:, :, :3]
        
        # Blend with original
        blended = 0.6 * img_np + 0.4 * heatmap
        
        ax.imshow(blended)
        ax.set_title(f'Thinking Step {i+1}/{len(attention_history)}')
        ax.axis('off')
        
        # Convert to image
        fig.canvas.draw()
        fig.canvas.draw()
        # buffer_rgba returns a memoryview of RGBA bytes
        frame = np.array(fig.canvas.renderer.buffer_rgba())
        # Convert RGBA to RGB
        frame = frame[:, :, :3]
        frames.append(frame)
        
        plt.close(fig)
    
    # Save as GIF
    imageio.mimsave(output_path, frames, fps=fps, loop=0)
    
    return output_path


__all__ = ['CTMModel', 'CTMConfig', 'CTMOutput', 'plot_attention_video', 'CTM_AVAILABLE']