Spaces:

MohsinElis
/

mfft-api

Paused

Deploy production MFFT detection API with verified 20-epoch weights

37b0e8b verified 26 days ago

19.1 kB

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Tuple, Optional


BAND_CONFIGS = {
    3: [(0.0, 0.15), (0.15, 0.45), (0.45, 1.0)],
    4: [(0.0, 0.1), (0.1, 0.3), (0.3, 0.6), (0.6, 1.0)],
}

class FrequencyDecomposition(nn.Module):
    """Radial band-pass decomposition via FFT.

    soft_masks=False reproduces the pilot behavior (binary masks).
    soft_masks=True replaces the hard cutoff with a sigmoid transition of
    width `soft_tau * r_max`, removing the ringing artifacts that hard
    masks imprint on the band images.

    Masks are cached per (H, W, device) instead of being rebuilt for every
    band of every forward pass.
    """
    def __init__(self, num_bands: int = 3, soft_masks: bool = False,
                 soft_tau: float = 0.02):
        super().__init__()
        self.bands = BAND_CONFIGS.get(num_bands, BAND_CONFIGS[3])
        self.num_bands = len(self.bands)
        self.soft_masks = soft_masks
        self.soft_tau = soft_tau
        self._mask_cache = {}

    def _get_masks(self, H: int, W: int, device) -> torch.Tensor:
        key = (H, W, str(device), self.soft_masks)
        cached = self._mask_cache.get(key)
        if cached is not None:
            return cached

        ny, nx = H // 2, W // 2
        r_max = min(ny, nx)
        y_grid, x_grid = torch.meshgrid(
            torch.arange(H, device=device),
            torch.arange(W, device=device),
            indexing='ij'
        )
        dist = torch.sqrt((y_grid - ny).float() ** 2 + (x_grid - nx).float() ** 2)

        masks = []
        for lo, hi in self.bands:
            r_low, r_high = r_max * lo, r_max * hi
            if self.soft_masks:
                tau = max(self.soft_tau * r_max, 1e-3)
                lower = torch.sigmoid((dist - r_low) / tau) if lo > 0 \
                    else torch.ones_like(dist)
                upper = torch.sigmoid((r_high - dist) / tau) if hi < 1.0 \
                    else torch.ones_like(dist)
                masks.append(lower * upper)
            else:
                masks.append(((dist >= int(r_low)) & (dist < int(r_high))).float())
        stacked = torch.stack(masks).unsqueeze(1).unsqueeze(1)  # (B_bands,1,1,H,W)
        self._mask_cache[key] = stacked
        return stacked

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        B, C, H, W = x.shape
        x_fft = torch.fft.fft2(x, norm='ortho')
        x_shifted = torch.fft.fftshift(x_fft)
        masks = self._get_masks(H, W, x.device)

        outputs = []
        for i in range(self.num_bands):
            filtered = x_shifted * masks[i]
            band = torch.fft.ifft2(torch.fft.ifftshift(filtered), norm='ortho').real
            outputs.append(band)
        return outputs


def npr_residual(x: torch.Tensor, factor: int = 2) -> torch.Tensor:
    """Neighboring-pixel-relation residual (after Tan et al., CVPR 2024):
    the difference between the image and its down-up resampled version,
    isolating the local interpolation traces that generator upsampling
    layers imprint. Returns a tensor shaped like x."""
    down = F.interpolate(x, scale_factor=1.0 / factor,
                         mode='bilinear', align_corners=False)
    up = F.interpolate(down, size=x.shape[2:],
                       mode='bilinear', align_corners=False)
    return x - up


class FrequencyFeatureExtractor(nn.Module):
    """
    CNN backbone for each frequency band.
    Uses depthwise separable convolutions for efficiency.
    """
    def __init__(self, in_channels: int = 3, feat_dim: int = 256):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.GELU(),
        )
        self.blocks = nn.ModuleList([
            self._make_block(32, 64, 3, 2),
            self._make_block(64, 128, 3, 2),
            self._make_block(128, 256, 3, 2),
        ])
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, feat_dim),
            nn.LayerNorm(feat_dim),
        )

    def _make_block(self, cin: int, cout: int, k: int, s: int) -> nn.Module:
        return nn.Sequential(
            nn.Conv2d(cin, cin, k, stride=s, padding=k//2, groups=cin, bias=False),
            nn.Conv2d(cin, cout, 1, bias=False),
            nn.BatchNorm2d(cout),
            nn.GELU(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        for block in self.blocks:
            x = block(x)
        return self.head(x)


class PretrainedFrequencyFeatureExtractor(nn.Module):
    """
    Drop-in replacement for FrequencyFeatureExtractor backed by an
    ImageNet-pretrained MobileNetV3-Small feature trunk. Used for the
    matched-pretraining MFFT variant that decouples architecture from
    initialization in comparisons against pretrained baselines.

    Band images are 3-channel (per-RGB-channel FFT filtering), so the
    pretrained stem is used unchanged.
    """
    def __init__(self, in_channels: int = 3, feat_dim: int = 256,
                 weights: str = "DEFAULT"):
        super().__init__()
        from torchvision.models import mobilenet_v3_small
        try:
            backbone = mobilenet_v3_small(weights=weights)
        except Exception as e:  # offline or weights unavailable
            import warnings
            warnings.warn(f"Pretrained weights unavailable ({e}); "
                          "falling back to random init.")
            backbone = mobilenet_v3_small(weights=None)
        if in_channels != 3:
            # inflate the stem conv: keep RGB filters, tile them (scaled)
            # over the extra channels so pretrained features are preserved
            old = backbone.features[0][0]
            new = nn.Conv2d(in_channels, old.out_channels,
                            old.kernel_size, old.stride, old.padding,
                            bias=old.bias is not None)
            with torch.no_grad():
                reps = -(-in_channels // 3)  # ceil division
                w = old.weight.repeat(1, reps, 1, 1)[:, :in_channels]
                new.weight.copy_(w * (3.0 / in_channels))
            backbone.features[0][0] = new
        self.features = backbone.features  # output: (B, 576, H/32, W/32)
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(576, feat_dim),
            nn.LayerNorm(feat_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.head(self.features(x))


class CrossAttentionFusion(nn.Module):
    """
    Cross-attention fusion across frequency bands.
    Each band attends to all others to produce a fused representation.
    """
    def __init__(self, dim: int = 256, num_heads: int = 8):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(0.1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, D = x.shape
        qkv = self.to_qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)

        out = (attn @ v).transpose(1, 2).reshape(B, N, D)
        out = self.proj(out)
        out = self.proj_drop(out)
        return out


class FrequencyGuidedAttention(nn.Module):
    """
    Novel module: uses frequency information to guide spatial attention.
    High-freq regions (edges, textures) get higher attention weights.
    """
    def __init__(self, dim: int = 256):
        super().__init__()
        self.freq_gate = nn.Sequential(
            nn.Linear(dim, dim // 4),
            nn.ReLU(),
            nn.Linear(dim // 4, dim),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor, freq_weights: torch.Tensor) -> torch.Tensor:
        gate = self.freq_gate(x)
        return x * gate * freq_weights


class MFFT(nn.Module):
    """
    Multi-Frequency Fusion Transformer (MFFT)
    ===========================================
    A novel architecture for AI-generated image detection.

    Key innovations:
    1. Frequency decomposition into Low/Mid/High bands via DCT
    2. Per-band CNN feature extraction
    3. Cross-attention fusion across frequency bands
    4. Frequency-guided spatial attention
    5. Multi-scale feature aggregation

    Input:  (B, 3, H, W)  RGB image
    Output: (B, 2)         logits [real, ai_generated]

    Ablation config (dict):
        spatial_only : bool  — skip frequency decomposition, use raw image
        skip_bands   : list  — band names to exclude: ["low", "mid", "high"]
        fusion_mode  : str   — "attention" | "concat" | "avg" | "max"
        use_fga      : bool  — enable/disable FrequencyGuidedAttention
    """
    def __init__(
        self,
        in_channels: int = 3,
        feat_dim: int = 256,
        num_bands: int = 3,
        num_heads: int = 8,
        num_classes: int = 2,
        ablation: Optional[dict] = None,
        pretrained_extractors: bool = False,
        use_npr: bool = False,
        soft_masks: bool = False,
    ):
        super().__init__()
        self.num_bands = num_bands
        self.feat_dim = feat_dim
        self.ablation = ablation or {}
        self.pretrained_extractors = pretrained_extractors
        self.use_npr = use_npr

        self.decomposer = FrequencyDecomposition(
            num_bands=num_bands, soft_masks=soft_masks)
        extractor_cls = (PretrainedFrequencyFeatureExtractor
                         if pretrained_extractors else FrequencyFeatureExtractor)
        # with NPR enabled, the highest band additionally receives the
        # 3-channel neighboring-pixel-relation residual
        self.extractors = nn.ModuleList([
            extractor_cls(
                in_channels + (3 if (use_npr and i == num_bands - 1) else 0),
                feat_dim)
            for i in range(num_bands)
        ])

        if self.ablation.get("fusion_mode", "attention") == "attention":
            self.fusion = CrossAttentionFusion(feat_dim, num_heads)
        else:
            self.fusion = None

        self.use_fga = self.ablation.get("use_fga", True)
        if self.use_fga:
            self.freq_guided_attn = FrequencyGuidedAttention(feat_dim)
        else:
            self.freq_guided_attn = None

        self.classifier = nn.Sequential(
            nn.LayerNorm(feat_dim * num_bands),
            nn.Linear(feat_dim * num_bands, feat_dim),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(feat_dim, feat_dim // 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(feat_dim // 2, num_classes),
        )

        # Ablation configs whose fused representation is narrower than the
        # classifier input (avg/max fusion, skipped bands, spatial-only)
        # get a learned projection instead of the old zero-padding, which
        # wasted classifier capacity on constant inputs.
        combined_dim = self._fused_dim(feat_dim, num_bands)
        if combined_dim != feat_dim * num_bands:
            self.input_proj = nn.Linear(combined_dim, feat_dim * num_bands)
        else:
            self.input_proj = nn.Identity()

        self._init_weights()

    def _fused_dim(self, feat_dim: int, num_bands: int) -> int:
        """Width of the flattened fused representation under the current
        ablation config (mirrors the branching in forward())."""
        if self.ablation.get("spatial_only", False):
            return feat_dim
        num_active = num_bands - len(self.ablation.get("skip_bands", []))
        mode = self.ablation.get("fusion_mode", "attention")
        if mode in ("avg", "max") and num_active > 1:
            return feat_dim
        return num_active * feat_dim  # attention / concat / single band

    def _init_weights(self):
        # pretrained extractor trunks must keep their ImageNet weights
        skip = set()
        if self.pretrained_extractors:
            for ext in self.extractors:
                for m in ext.features.modules():
                    skip.add(id(m))
        for m in self.modules():
            if id(m) in skip:
                continue
            if isinstance(m, (nn.Conv2d, nn.Linear)):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

    def forward(self, x: torch.Tensor, return_heatmap: bool = False):
        spatial_only = self.ablation.get("spatial_only", False)
        skip_bands = self.ablation.get("skip_bands", [])
        fusion_mode = self.ablation.get("fusion_mode", "attention")
        if self.num_bands == 4:
            band_names = ["low", "low_mid", "mid_high", "high"]
        else:
            band_names = ["low", "mid", "high"]

        if spatial_only:
            bands = [x]
            extractors = self.extractors[:1]
            extractor_inputs = bands
        else:
            all_bands = self.decomposer(x)
            keep_indices = [i for i, name in enumerate(band_names) if name not in skip_bands]
            bands = [all_bands[i] for i in keep_indices]
            extractors = [self.extractors[i] for i in keep_indices]
            extractor_inputs = []
            for i, band in zip(keep_indices, bands):
                if self.use_npr and i == self.num_bands - 1:
                    band = torch.cat([band, npr_residual(x)], dim=1)
                extractor_inputs.append(band)

        features = []
        for band, extractor in zip(extractor_inputs, extractors):
            feat = extractor(band)
            features.append(feat)

        B = x.shape[0]
        num_active = len(features)

        if num_active == 1 or fusion_mode == "concat":
            fused = torch.cat(features, dim=-1)
            N = num_active
            D = fused.shape[-1] // N
            fused = fused.unsqueeze(1)
        elif fusion_mode == "avg":
            fused = torch.stack(features, dim=0).mean(dim=0)
            N = 1
            D = fused.shape[-1]
            fused = fused.unsqueeze(1)
        elif fusion_mode == "max":
            fused = torch.stack(features, dim=0).max(dim=0).values
            N = 1
            D = fused.shape[-1]
            fused = fused.unsqueeze(1)
        elif fusion_mode == "attention" and self.fusion is not None and num_active > 1:
            feat_stack = torch.stack(features, dim=1)
            fused = self.fusion(feat_stack)
            N = fused.shape[1]
            D = fused.shape[2]
        else:
            fused = torch.cat(features, dim=-1)
            N = num_active
            D = fused.shape[-1] // N
            fused = fused.unsqueeze(1)

        # FGA operates on per-band tokens, so it only applies when the fused
        # representation kept one token per band (attention path). Pooled
        # representations (avg/max/concat) have no band axis to weight.
        fga_applicable = (
            self.use_fga and self.freq_guided_attn is not None
            and not spatial_only
            and N == num_active and D == fused.shape[-1]
            and fused.shape[-1] == self.feat_dim
        )
        if fga_applicable:
            all_bands_for_fga = bands
            freq_magnitudes = torch.stack([
                torch.abs(band).mean(dim=(1, 2, 3))
                for band in all_bands_for_fga
            ], dim=1)
            freq_weights = F.softmax(freq_magnitudes, dim=1).unsqueeze(-1)
            guided = self.freq_guided_attn(fused, freq_weights)
        else:
            guided = fused

        combined = guided.reshape(B, N * D)
        combined = self.input_proj(combined)
        logits = self.classifier(combined)

        if return_heatmap:
            heatmaps = []
            for band in bands:
                h = torch.abs(band).mean(dim=1, keepdim=True)
                h = F.interpolate(h, size=x.shape[2:], mode='bilinear', align_corners=False)
                heatmaps.append(h)
            return logits, torch.cat(heatmaps, dim=1)
        return logits


class MFFTWithExplainability(MFFT):
    """
    Extension of MFFT that also produces:
    - Confidence score (0-1)
    - Per-region anomaly heatmap
    - Per-frequency-band contribution scores
    """
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.register_buffer(
            'mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        )
        self.register_buffer(
            'std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
        )

    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[1] == 1:
            x = x.repeat(1, 3, 1, 1)
        x = (x / 255.0 - self.mean) / self.std
        return x

    @torch.no_grad()
    def predict(self, x: torch.Tensor) -> dict:
        x = self.preprocess(x)
        self.eval()
        logits, heatmaps = self.forward(x, return_heatmap=True)
        probs = F.softmax(logits, dim=-1)
        preds = torch.argmax(probs, dim=-1)

        return {
            "prediction": preds,
            "real_prob": probs[:, 0],
            "ai_prob": probs[:, 1],
            "confidence": probs.max(dim=-1).values,
            "heatmaps": heatmaps,
            "logits": logits,
        }


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def build_mfft(variant: str = "base", ablation: Optional[dict] = None,
               pretrained_extractors: bool = False,
               use_npr: bool = False, soft_masks: bool = False) -> MFFT:
    configs = {
        "tiny":  {"feat_dim": 128, "num_heads": 4, "num_bands": 3},
        "base":  {"feat_dim": 384, "num_heads": 6, "num_bands": 3},
        "large": {"feat_dim": 768, "num_heads": 12, "num_bands": 4},
    }
    cfg = configs.get(variant, configs["base"])
    model = MFFT(**cfg, ablation=ablation,
                 pretrained_extractors=pretrained_extractors,
                 use_npr=use_npr, soft_masks=soft_masks)
    return model

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import numpy as np
	from typing import List, Tuple, Optional


	BAND_CONFIGS = {
	3: [(0.0, 0.15), (0.15, 0.45), (0.45, 1.0)],
	4: [(0.0, 0.1), (0.1, 0.3), (0.3, 0.6), (0.6, 1.0)],
	}

	class FrequencyDecomposition(nn.Module):
	"""Radial band-pass decomposition via FFT.

	soft_masks=False reproduces the pilot behavior (binary masks).
	soft_masks=True replaces the hard cutoff with a sigmoid transition of
	width `soft_tau * r_max`, removing the ringing artifacts that hard
	masks imprint on the band images.

	Masks are cached per (H, W, device) instead of being rebuilt for every
	band of every forward pass.
	"""
	def __init__(self, num_bands: int = 3, soft_masks: bool = False,
	soft_tau: float = 0.02):
	super().__init__()
	self.bands = BAND_CONFIGS.get(num_bands, BAND_CONFIGS[3])
	self.num_bands = len(self.bands)
	self.soft_masks = soft_masks
	self.soft_tau = soft_tau
	self._mask_cache = {}

	def _get_masks(self, H: int, W: int, device) -> torch.Tensor:
	key = (H, W, str(device), self.soft_masks)
	cached = self._mask_cache.get(key)
	if cached is not None:
	return cached

	ny, nx = H // 2, W // 2
	r_max = min(ny, nx)
	y_grid, x_grid = torch.meshgrid(
	torch.arange(H, device=device),
	torch.arange(W, device=device),
	indexing='ij'
	)
	dist = torch.sqrt((y_grid - ny).float() 2 + (x_grid - nx).float() 2)

	masks = []
	for lo, hi in self.bands:
	r_low, r_high = r_max * lo, r_max * hi
	if self.soft_masks:
	tau = max(self.soft_tau * r_max, 1e-3)
	lower = torch.sigmoid((dist - r_low) / tau) if lo > 0 \
	else torch.ones_like(dist)
	upper = torch.sigmoid((r_high - dist) / tau) if hi < 1.0 \
	else torch.ones_like(dist)
	masks.append(lower * upper)
	else:
	masks.append(((dist >= int(r_low)) & (dist < int(r_high))).float())
	stacked = torch.stack(masks).unsqueeze(1).unsqueeze(1) # (B_bands,1,1,H,W)
	self._mask_cache[key] = stacked
	return stacked

	def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
	B, C, H, W = x.shape
	x_fft = torch.fft.fft2(x, norm='ortho')
	x_shifted = torch.fft.fftshift(x_fft)
	masks = self._get_masks(H, W, x.device)

	outputs = []
	for i in range(self.num_bands):
	filtered = x_shifted * masks[i]
	band = torch.fft.ifft2(torch.fft.ifftshift(filtered), norm='ortho').real
	outputs.append(band)
	return outputs


	def npr_residual(x: torch.Tensor, factor: int = 2) -> torch.Tensor:
	"""Neighboring-pixel-relation residual (after Tan et al., CVPR 2024):
	the difference between the image and its down-up resampled version,
	isolating the local interpolation traces that generator upsampling
	layers imprint. Returns a tensor shaped like x."""
	down = F.interpolate(x, scale_factor=1.0 / factor,
	mode='bilinear', align_corners=False)
	up = F.interpolate(down, size=x.shape[2:],
	mode='bilinear', align_corners=False)
	return x - up


	class FrequencyFeatureExtractor(nn.Module):
	"""
	CNN backbone for each frequency band.
	Uses depthwise separable convolutions for efficiency.
	"""
	def __init__(self, in_channels: int = 3, feat_dim: int = 256):
	super().__init__()
	self.stem = nn.Sequential(
	nn.Conv2d(in_channels, 32, 3, stride=2, padding=1, bias=False),
	nn.BatchNorm2d(32),
	nn.GELU(),
	)
	self.blocks = nn.ModuleList([
	self._make_block(32, 64, 3, 2),
	self._make_block(64, 128, 3, 2),
	self._make_block(128, 256, 3, 2),
	])
	self.head = nn.Sequential(
	nn.AdaptiveAvgPool2d(1),
	nn.Flatten(),
	nn.Linear(256, feat_dim),
	nn.LayerNorm(feat_dim),
	)

	def _make_block(self, cin: int, cout: int, k: int, s: int) -> nn.Module:
	return nn.Sequential(
	nn.Conv2d(cin, cin, k, stride=s, padding=k//2, groups=cin, bias=False),
	nn.Conv2d(cin, cout, 1, bias=False),
	nn.BatchNorm2d(cout),
	nn.GELU(),
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	x = self.stem(x)
	for block in self.blocks:
	x = block(x)
	return self.head(x)


	class PretrainedFrequencyFeatureExtractor(nn.Module):
	"""
	Drop-in replacement for FrequencyFeatureExtractor backed by an
	ImageNet-pretrained MobileNetV3-Small feature trunk. Used for the
	matched-pretraining MFFT variant that decouples architecture from
	initialization in comparisons against pretrained baselines.

	Band images are 3-channel (per-RGB-channel FFT filtering), so the
	pretrained stem is used unchanged.
	"""
	def __init__(self, in_channels: int = 3, feat_dim: int = 256,
	weights: str = "DEFAULT"):
	super().__init__()
	from torchvision.models import mobilenet_v3_small
	try:
	backbone = mobilenet_v3_small(weights=weights)
	except Exception as e: # offline or weights unavailable
	import warnings
	warnings.warn(f"Pretrained weights unavailable ({e}); "
	"falling back to random init.")
	backbone = mobilenet_v3_small(weights=None)
	if in_channels != 3:
	# inflate the stem conv: keep RGB filters, tile them (scaled)
	# over the extra channels so pretrained features are preserved
	old = backbone.features[0][0]
	new = nn.Conv2d(in_channels, old.out_channels,
	old.kernel_size, old.stride, old.padding,
	bias=old.bias is not None)
	with torch.no_grad():
	reps = -(-in_channels // 3) # ceil division
	w = old.weight.repeat(1, reps, 1, 1)[:, :in_channels]
	new.weight.copy_(w * (3.0 / in_channels))
	backbone.features[0][0] = new
	self.features = backbone.features # output: (B, 576, H/32, W/32)
	self.head = nn.Sequential(
	nn.AdaptiveAvgPool2d(1),
	nn.Flatten(),
	nn.Linear(576, feat_dim),
	nn.LayerNorm(feat_dim),
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.head(self.features(x))


	class CrossAttentionFusion(nn.Module):
	"""
	Cross-attention fusion across frequency bands.
	Each band attends to all others to produce a fused representation.
	"""
	def __init__(self, dim: int = 256, num_heads: int = 8):
	super().__init__()
	self.dim = dim
	self.num_heads = num_heads
	self.head_dim = dim // num_heads
	self.scale = self.head_dim ** -0.5

	self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
	self.proj = nn.Linear(dim, dim)
	self.proj_drop = nn.Dropout(0.1)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	B, N, D = x.shape
	qkv = self.to_qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
	qkv = qkv.permute(2, 0, 3, 1, 4)
	q, k, v = qkv[0], qkv[1], qkv[2]

	attn = (q @ k.transpose(-2, -1)) * self.scale
	attn = attn.softmax(dim=-1)

	out = (attn @ v).transpose(1, 2).reshape(B, N, D)
	out = self.proj(out)
	out = self.proj_drop(out)
	return out


	class FrequencyGuidedAttention(nn.Module):
	"""
	Novel module: uses frequency information to guide spatial attention.
	High-freq regions (edges, textures) get higher attention weights.
	"""
	def __init__(self, dim: int = 256):
	super().__init__()
	self.freq_gate = nn.Sequential(
	nn.Linear(dim, dim // 4),
	nn.ReLU(),
	nn.Linear(dim // 4, dim),
	nn.Sigmoid(),
	)

	def forward(self, x: torch.Tensor, freq_weights: torch.Tensor) -> torch.Tensor:
	gate = self.freq_gate(x)
	return x * gate * freq_weights


	class MFFT(nn.Module):
	"""
	Multi-Frequency Fusion Transformer (MFFT)
	===========================================
	A novel architecture for AI-generated image detection.

	Key innovations:
	1. Frequency decomposition into Low/Mid/High bands via DCT
	2. Per-band CNN feature extraction
	3. Cross-attention fusion across frequency bands
	4. Frequency-guided spatial attention
	5. Multi-scale feature aggregation

	Input: (B, 3, H, W) RGB image
	Output: (B, 2) logits [real, ai_generated]

	Ablation config (dict):
	spatial_only : bool — skip frequency decomposition, use raw image
	skip_bands : list — band names to exclude: ["low", "mid", "high"]
	fusion_mode : str — "attention" \| "concat" \| "avg" \| "max"
	use_fga : bool — enable/disable FrequencyGuidedAttention
	"""
	def __init__(
	self,
	in_channels: int = 3,
	feat_dim: int = 256,
	num_bands: int = 3,
	num_heads: int = 8,
	num_classes: int = 2,
	ablation: Optional[dict] = None,
	pretrained_extractors: bool = False,
	use_npr: bool = False,
	soft_masks: bool = False,
	):
	super().__init__()
	self.num_bands = num_bands
	self.feat_dim = feat_dim
	self.ablation = ablation or {}
	self.pretrained_extractors = pretrained_extractors
	self.use_npr = use_npr

	self.decomposer = FrequencyDecomposition(
	num_bands=num_bands, soft_masks=soft_masks)
	extractor_cls = (PretrainedFrequencyFeatureExtractor
	if pretrained_extractors else FrequencyFeatureExtractor)
	# with NPR enabled, the highest band additionally receives the
	# 3-channel neighboring-pixel-relation residual
	self.extractors = nn.ModuleList([
	extractor_cls(
	in_channels + (3 if (use_npr and i == num_bands - 1) else 0),
	feat_dim)
	for i in range(num_bands)
	])

	if self.ablation.get("fusion_mode", "attention") == "attention":
	self.fusion = CrossAttentionFusion(feat_dim, num_heads)
	else:
	self.fusion = None

	self.use_fga = self.ablation.get("use_fga", True)
	if self.use_fga:
	self.freq_guided_attn = FrequencyGuidedAttention(feat_dim)
	else:
	self.freq_guided_attn = None

	self.classifier = nn.Sequential(
	nn.LayerNorm(feat_dim * num_bands),
	nn.Linear(feat_dim * num_bands, feat_dim),
	nn.GELU(),
	nn.Dropout(0.2),
	nn.Linear(feat_dim, feat_dim // 2),
	nn.GELU(),
	nn.Dropout(0.1),
	nn.Linear(feat_dim // 2, num_classes),
	)

	# Ablation configs whose fused representation is narrower than the
	# classifier input (avg/max fusion, skipped bands, spatial-only)
	# get a learned projection instead of the old zero-padding, which
	# wasted classifier capacity on constant inputs.
	combined_dim = self._fused_dim(feat_dim, num_bands)
	if combined_dim != feat_dim * num_bands:
	self.input_proj = nn.Linear(combined_dim, feat_dim * num_bands)
	else:
	self.input_proj = nn.Identity()

	self._init_weights()

	def _fused_dim(self, feat_dim: int, num_bands: int) -> int:
	"""Width of the flattened fused representation under the current
	ablation config (mirrors the branching in forward())."""
	if self.ablation.get("spatial_only", False):
	return feat_dim
	num_active = num_bands - len(self.ablation.get("skip_bands", []))
	mode = self.ablation.get("fusion_mode", "attention")
	if mode in ("avg", "max") and num_active > 1:
	return feat_dim
	return num_active * feat_dim # attention / concat / single band

	def _init_weights(self):
	# pretrained extractor trunks must keep their ImageNet weights
	skip = set()
	if self.pretrained_extractors:
	for ext in self.extractors:
	for m in ext.features.modules():
	skip.add(id(m))
	for m in self.modules():
	if id(m) in skip:
	continue
	if isinstance(m, (nn.Conv2d, nn.Linear)):
	nn.init.trunc_normal_(m.weight, std=0.02)
	if m.bias is not None:
	nn.init.constant_(m.bias, 0)
	elif isinstance(m, nn.LayerNorm):
	nn.init.constant_(m.bias, 0)
	nn.init.constant_(m.weight, 1.0)

	def forward(self, x: torch.Tensor, return_heatmap: bool = False):
	spatial_only = self.ablation.get("spatial_only", False)
	skip_bands = self.ablation.get("skip_bands", [])
	fusion_mode = self.ablation.get("fusion_mode", "attention")
	if self.num_bands == 4:
	band_names = ["low", "low_mid", "mid_high", "high"]
	else:
	band_names = ["low", "mid", "high"]

	if spatial_only:
	bands = [x]
	extractors = self.extractors[:1]
	extractor_inputs = bands
	else:
	all_bands = self.decomposer(x)
	keep_indices = [i for i, name in enumerate(band_names) if name not in skip_bands]
	bands = [all_bands[i] for i in keep_indices]
	extractors = [self.extractors[i] for i in keep_indices]
	extractor_inputs = []
	for i, band in zip(keep_indices, bands):
	if self.use_npr and i == self.num_bands - 1:
	band = torch.cat([band, npr_residual(x)], dim=1)
	extractor_inputs.append(band)

	features = []
	for band, extractor in zip(extractor_inputs, extractors):
	feat = extractor(band)
	features.append(feat)

	B = x.shape[0]
	num_active = len(features)

	if num_active == 1 or fusion_mode == "concat":
	fused = torch.cat(features, dim=-1)
	N = num_active
	D = fused.shape[-1] // N
	fused = fused.unsqueeze(1)
	elif fusion_mode == "avg":
	fused = torch.stack(features, dim=0).mean(dim=0)
	N = 1
	D = fused.shape[-1]
	fused = fused.unsqueeze(1)
	elif fusion_mode == "max":
	fused = torch.stack(features, dim=0).max(dim=0).values
	N = 1
	D = fused.shape[-1]
	fused = fused.unsqueeze(1)
	elif fusion_mode == "attention" and self.fusion is not None and num_active > 1:
	feat_stack = torch.stack(features, dim=1)
	fused = self.fusion(feat_stack)
	N = fused.shape[1]
	D = fused.shape[2]
	else:
	fused = torch.cat(features, dim=-1)
	N = num_active
	D = fused.shape[-1] // N
	fused = fused.unsqueeze(1)

	# FGA operates on per-band tokens, so it only applies when the fused
	# representation kept one token per band (attention path). Pooled
	# representations (avg/max/concat) have no band axis to weight.
	fga_applicable = (
	self.use_fga and self.freq_guided_attn is not None
	and not spatial_only
	and N == num_active and D == fused.shape[-1]
	and fused.shape[-1] == self.feat_dim
	)
	if fga_applicable:
	all_bands_for_fga = bands
	freq_magnitudes = torch.stack([
	torch.abs(band).mean(dim=(1, 2, 3))
	for band in all_bands_for_fga
	], dim=1)
	freq_weights = F.softmax(freq_magnitudes, dim=1).unsqueeze(-1)
	guided = self.freq_guided_attn(fused, freq_weights)
	else:
	guided = fused

	combined = guided.reshape(B, N * D)
	combined = self.input_proj(combined)
	logits = self.classifier(combined)

	if return_heatmap:
	heatmaps = []
	for band in bands:
	h = torch.abs(band).mean(dim=1, keepdim=True)
	h = F.interpolate(h, size=x.shape[2:], mode='bilinear', align_corners=False)
	heatmaps.append(h)
	return logits, torch.cat(heatmaps, dim=1)
	return logits


	class MFFTWithExplainability(MFFT):
	"""
	Extension of MFFT that also produces:
	- Confidence score (0-1)
	- Per-region anomaly heatmap
	- Per-frequency-band contribution scores
	"""
	def __init__(self, **kwargs):
	super().__init__(**kwargs)
	self.register_buffer(
	'mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
	)
	self.register_buffer(
	'std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
	)

	def preprocess(self, x: torch.Tensor) -> torch.Tensor:
	if x.shape[1] == 1:
	x = x.repeat(1, 3, 1, 1)
	x = (x / 255.0 - self.mean) / self.std
	return x

	@torch.no_grad()
	def predict(self, x: torch.Tensor) -> dict:
	x = self.preprocess(x)
	self.eval()
	logits, heatmaps = self.forward(x, return_heatmap=True)
	probs = F.softmax(logits, dim=-1)
	preds = torch.argmax(probs, dim=-1)

	return {
	"prediction": preds,
	"real_prob": probs[:, 0],
	"ai_prob": probs[:, 1],
	"confidence": probs.max(dim=-1).values,
	"heatmaps": heatmaps,
	"logits": logits,
	}


	def count_parameters(model: nn.Module) -> int:
	return sum(p.numel() for p in model.parameters() if p.requires_grad)


	def build_mfft(variant: str = "base", ablation: Optional[dict] = None,
	pretrained_extractors: bool = False,
	use_npr: bool = False, soft_masks: bool = False) -> MFFT:
	configs = {
	"tiny": {"feat_dim": 128, "num_heads": 4, "num_bands": 3},
	"base": {"feat_dim": 384, "num_heads": 6, "num_bands": 3},
	"large": {"feat_dim": 768, "num_heads": 12, "num_bands": 4},
	}
	cfg = configs.get(variant, configs["base"])
	model = MFFT(**cfg, ablation=ablation,
	pretrained_extractors=pretrained_extractors,
	use_npr=use_npr, soft_masks=soft_masks)
	return model