src/models/multimodalIntraInterModal.py

import torch
import torch.nn as nn
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from gatedResidualBlock import GatedAlteredResidualBlock, StackedGatedResidualBlock
from loadImageModelClassifier import loadModels
from metablock import MetaBlock
from metanet import MetaNet
from multimodalMDNet import MDNet

class MultimodalModel(nn.Module):
    def __init__(
        self,
        num_classes,
        num_heads,
        device,
        cnn_model_name,
        text_model_name,
        batch_size=32,
        common_dim=512,
        text_encoder_dim_output=512,
        vocab_size=91,
        unfreeze_weights="frozen_weights",
        attention_mecanism="concatenation",
        n=2
    ):
        super().__init__()

        self.device = device
        self.common_dim = common_dim
        self.num_heads = num_heads
        self.attention_mecanism = attention_mecanism
        self.n = n
        self.vocab_size = vocab_size 
        self.num_classes = num_classes
        self.cnn_model_name = cnn_model_name
        self.text_model_name = text_model_name
        self.unfreeze_weights = unfreeze_weights
        # =====================================================
        # Text Encoder
        # =====================================================
        self.text_encoder_dim_output = text_encoder_dim_output

        # =====================================================
        # Image Encoder
        # =====================================================
        self.image_encoder, self.cnn_dim_output = loadModels.loadModelImageEncoder(
            cnn_model_name=self.cnn_model_name,
            backbone_train_mode=self.unfreeze_weights
        )

        self.image_projector = nn.Linear(self.cnn_dim_output, self.common_dim)

        if text_model_name == "one-hot-encoder":
            self.text_fc = nn.Sequential(
                nn.Linear(self.vocab_size, 256),
                nn.ReLU(),
                nn.Linear(256, 512),
                nn.ReLU(),
                nn.Linear(512, self.text_encoder_dim_output)
            )
            self.text_encoder = None
        else:
            self.text_encoder, self.text_encoder_dim_output, _ = loadModels.loadTextModelEncoder(
                text_model_encoder=self.text_model_name,
                train_mode=self.unfreeze_weights
            )
            self.text_fc = None

        self.text_projector = nn.Linear(self.text_encoder_dim_output, self.common_dim)

        # =====================================================
        # Attention blocks
        # =====================================================
        self.image_self_attention = nn.MultiheadAttention(
            embed_dim=self.common_dim,
            num_heads=self.num_heads,
            batch_first=False
        )

        self.text_self_attention = nn.MultiheadAttention(
            embed_dim=self.common_dim,
            num_heads=self.num_heads,
            batch_first=False
        )

        self.image_cross_attention = nn.MultiheadAttention(
            embed_dim=self.common_dim,
            num_heads=self.num_heads,
            batch_first=False
        )

        self.text_cross_attention = nn.MultiheadAttention(
            embed_dim=self.common_dim,
            num_heads=self.num_heads,
            batch_first=False
        )

        # =====================================================
        # Gating
        # =====================================================
        self.img_gate = nn.Linear(self.common_dim, self.common_dim)
        self.txt_gate = nn.Linear(self.common_dim, self.common_dim)

        # =====================================================
        # MetaBlock (vetorial)
        # =====================================================
        # Bloco do Metablock, caso queira usar 
        self.meta_block = MetaBlock(
            V_dim=self.common_dim if self.attention_mecanism in ["att-intramodal+residual+cross-attention-metadados+metablock"] else self.cnn_dim_output,
            U_dim=self.common_dim if self.attention_mecanism in ["att-intramodal+residual+cross-attention-metadados+metablock", "metablock-se"] else self.text_encoder_dim_output 
        )
        # Residual Blocks
        # -------------------------
        self.image_residual = GatedAlteredResidualBlock(dim=self.common_dim)
        self.text_residual = GatedAlteredResidualBlock(dim=self.common_dim)

        # =====================================================
        # Fusion MLP
        # =====================================================
        self.fc_fusion = self.fc_mlp_module(
            n=1 if attention_mecanism == "no-metadata" else self.n
        )

        # Usado no Metablock
        self.fc_visual_only = nn.Linear(self.cnn_dim_output, self.num_classes)

        self.fc_mlp_module_after_metablock_fusion_module = self.fc_mlp_module_after_metablock()

    def fc_mlp_module(self, n=1):
        fc_fusion = nn.Sequential(
            nn.Linear(self.common_dim * n, self.common_dim),
            nn.LayerNorm(self.common_dim),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(self.common_dim, self.common_dim // 2),
            nn.LayerNorm(self.common_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(self.common_dim // 2, self.num_classes)
        )
        return fc_fusion
    
    def fc_mlp_module_after_metablock(self):
        fc_fusion = nn.Sequential(
            nn.Linear(self.cnn_dim_output, self.common_dim),
            nn.LayerNorm(self.common_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(self.common_dim, self.common_dim // 2),
            nn.LayerNorm(self.common_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(self.common_dim // 2, self.num_classes)
        )
        return fc_fusion

    def forward(self, image, text_metadata):
        # =====================================================
        # Image encoding → (B, D)
        # =====================================================
        image = image.to(self.device)
        img_feat = self.image_encoder(image)

        if img_feat.dim() == 4:
            img_feat = img_feat.mean(dim=(-2, -1))

        proj_img_feat = self.image_projector(img_feat)

        # =====================================================
        # Text encoding → (B, D)
        # =====================================================
        if self.text_model_name == "one-hot-encoder":
            txt_feat = self.text_fc(text_metadata.to(self.device))
        else:
            input_ids = text_metadata["input_ids"].squeeze(1).to(self.device)
            attention_mask = text_metadata["attention_mask"].squeeze(1).to(self.device)
            outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)
            txt_feat = outputs.last_hidden_state[:, 0, :]

        proj_txt_feat = self.text_projector(txt_feat)

        # =====================================================
        # Attention (seq_len = 1)
        # =====================================================
        img_seq = proj_img_feat.unsqueeze(0)
        txt_seq = proj_txt_feat.unsqueeze(0)

        img_att, _ = self.image_self_attention(img_seq, img_seq, img_seq)
        txt_att, _ = self.text_self_attention(txt_seq, txt_seq, txt_seq)

        img_cross, _ = self.image_cross_attention(img_att, txt_att, txt_att)
        txt_cross, _ = self.text_cross_attention(txt_att, img_att, img_att)

        img_pooled = img_cross.squeeze(0)
        txt_pooled = txt_cross.squeeze(0)

        # =====================================================
        # Fusion strategies
        # =====================================================
        if self.attention_mecanism == "no-metadata":
            return self.fc_fusion(proj_img_feat)

        elif self.attention_mecanism == "no-metadata-without-mlp":
            return self.fc_visual_only(img_seq)

        elif self.attention_mecanism == "concatenation":
            fused = torch.cat([proj_img_feat, proj_txt_feat], dim=1)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "crossattention":
            fused = torch.cat([img_pooled, txt_pooled], dim=1)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "weighted":
            alpha_img = torch.sigmoid(self.img_gate(proj_img_feat))
            alpha_txt = torch.sigmoid(self.txt_gate(proj_txt_feat))
            fused = torch.cat([alpha_img * proj_img_feat, alpha_txt * proj_txt_feat], dim=1)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "gfcam":
            alpha_img = torch.sigmoid(self.img_gate(img_pooled))
            alpha_txt = torch.sigmoid(self.txt_gate(txt_pooled))
            fused = torch.cat([alpha_img * img_pooled, alpha_txt * txt_pooled], dim=1)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "cross-weights-after-crossattention":
            alpha_img = torch.sigmoid(self.img_gate(img_pooled))
            alpha_txt = torch.sigmoid(self.txt_gate(txt_pooled))
            fused = torch.cat([alpha_txt * img_pooled, alpha_img * txt_pooled], dim=1)
            return self.fc_fusion(fused)
        
        elif self.attention_mecanism == "metablock":
            meta_features = self.meta_block(
                img_feat,      
                txt_feat
            )
            return self.fc_mlp_module_after_metablock_fusion_module(meta_features)
        
        # =====================================================
        # Residual-based multimodal fusion
        # =====================================================

        # =====================================================
        # Residual-based multimodal fusion (corrigidos)
        # =====================================================

        elif self.attention_mecanism == "only-with-att-intramodal+residual":
            # Residual direto (sem self-att explícito)
            # No nosso forward atual, a forma padrão é (seq_len=1, B, D) => img_seq/txt_seq
            img_res = self.image_residual(img_seq, txt_seq, txt_seq)  # (1, B, D)
            txt_res = self.text_residual(txt_seq, img_seq, img_seq)   # (1, B, D)

            img_res = img_res.squeeze(0)  # (B, D)
            txt_res = txt_res.squeeze(0)  # (B, D)

            fused = torch.cat([img_res, txt_res], dim=1)  # (B, 2D)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "att-intramodal+residual":
            # Self-att já calculado acima: img_att, txt_att
            # Residual usa (base, att, att)
            img_res = self.image_residual(img_seq, img_att, img_att)  # (1, B, D)
            txt_res = self.text_residual(txt_seq, txt_att, txt_att)   # (1, B, D)

            img_res = img_res.squeeze(0)  # (B, D)
            txt_res = txt_res.squeeze(0)  # (B, D)

            fused = torch.cat([img_res, txt_res], dim=1)  # (B, 2D)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "att-intramodal+residual+cross-attention-metadados":
            # Self-att já calculado: img_att, txt_att
            # Residual antes do cross
            img_res = self.image_residual(img_seq, img_att, img_att)  # (1, B, D)
            txt_res = self.text_residual(txt_seq, txt_att, txt_att)   # (1, B, D)

            # Cross-attention entre os residuais
            img_cross2, _ = self.image_cross_attention(
                query=img_res, key=txt_res, value=txt_res
            )  # (1, B, D)

            txt_cross2, _ = self.text_cross_attention(
                query=txt_res, key=img_res, value=img_res
            )  # (1, B, D)

            img_pooled2 = img_cross2.squeeze(0)  # (B, D)
            txt_pooled2 = txt_cross2.squeeze(0)  # (B, D)

            fused = torch.cat([img_pooled2, txt_pooled2], dim=1)  # (B, 2D)
            return self.fc_fusion(fused)

        elif self.attention_mecanism == "att-intramodal+residual+cross-attention-metadados+metablock":
            # Self-att já calculado: img_att, txt_att
            # Residual antes do cross
            img_res = self.image_residual(img_seq, img_att, img_att)  # (1, B, D)
            txt_res = self.text_residual(txt_seq, txt_att, txt_att)   # (1, B, D)

            # Cross-attention
            img_cross2, _ = self.image_cross_attention(
                query=img_res, key=txt_res, value=txt_res
            )  # (1, B, D)

            txt_cross2, _ = self.text_cross_attention(
                query=txt_res, key=img_res, value=img_res
            )  # (1, B, D)

            img_pooled2 = img_cross2.squeeze(0)  # (B, D)
            txt_pooled2 = txt_cross2.squeeze(0)  # (B, D)

            # MetaBlock vetorial (B, D) + (B, D) -> (B, D)
            fused_meta = self.meta_block(img_pooled2, txt_pooled2)  # (B, D)

            # Classificador no espaço comum
            return self.fc_visual_only(fused_meta)

        elif self.attention_mecanism == "att-intramodal+residual+cross-attention-metadados+att-intramodal+residual":
            # Self-att inicial: img_att, txt_att
            # Residual antes do cross
            img_res1 = self.image_residual(img_seq, img_att, img_att)  # (1, B, D)
            txt_res1 = self.text_residual(txt_seq, txt_att, txt_att)   # (1, B, D)

            # Cross-attention
            img_cross2, _ = self.image_cross_attention(
                query=img_res1, key=txt_res1, value=txt_res1
            )  # (1, B, D)

            txt_cross2, _ = self.text_cross_attention(
                query=txt_res1, key=img_res1, value=img_res1
            )  # (1, B, D)

            # Self-att depois do cross
            img_att2, _ = self.image_self_attention(img_cross2, img_cross2, img_cross2)
            txt_att2, _ = self.text_self_attention(txt_cross2, txt_cross2, txt_cross2)

            # Residual final (base = cross2)
            img_res2 = self.image_residual(img_cross2, img_att2, img_att2).squeeze(0)  # (B, D)
            txt_res2 = self.text_residual(txt_cross2, txt_att2, txt_att2).squeeze(0)   # (B, D)

            fused = torch.cat([img_res2, txt_res2], dim=1)  # (B, 2D)
            return self.fc_fusion(fused)
        else:
            raise ValueError(
                f"Attention mechanism '{self.attention_mecanism}' not implemented."
            )