Upload 5 files

models/__init__.py

@@ -0,0 +1 @@
+# Makes this directory a python package

models/base_encoder.py

@@ -0,0 +1,20 @@
+# models/base_encoder.py
+
+import torch
+import torch.nn as nn
+from abc import ABC, abstractmethod
+
+
+class BaseVisionEncoder(nn.Module, ABC):
+    def __init__(self, embed_dim: int):
+        super().__init__()
+        self.embed_dim = embed_dim
+
+    @abstractmethod
+    def forward(self, pixel_values: torch.Tensor):
+        pass
+
+    @abstractmethod
+    def get_output_dim(self):
+        """Return the dimensionality of the encoder output embedding."""
+        pass

models/encoder_projection_t5.py

@@ -0,0 +1,27 @@
+# models/image_projection.py
+
+import torch
+import torch.nn as nn
+
+
+class ImageProjection(nn.Module):
+    """
+    Projects encoder image embeddings into the T5 hidden size.
+
+    Example:
+        - CLIP ViT-L/14 gives 1024-d embeddings
+        - T5-small expects 512-d hidden states
+        → This linear layer maps 1024 → 512
+
+    Forward:
+        image_embeds: (B, D_enc) or (B, S, D_enc)
+        returns:
+            projected_embeds: (B, D_t5) or (B, S, D_t5)
+    """
+
+    def __init__(self, encoder_dim: int, t5_hidden_size: int):
+        super().__init__()
+        self.proj = nn.Linear(encoder_dim, t5_hidden_size)
+
+    def forward(self, image_embeds: torch.Tensor):
+        return self.proj(image_embeds)

models/encoders.py

@@ -0,0 +1,223 @@
+# models/resnet_encoder.py
+# Just for setting up the pipeline, this will be replaced
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from .base_encoder import BaseVisionEncoder
+from transformers import CLIPVisionModel
+from transformers import AutoModel
+
+
+class ResnetCNNEncoder(nn.Module): 
+    def __init__(self, model_name="resnet50", fine_tune_all_encoder_layers=False, freeze_encoder_entirely=False, freeze_until=3, vision_mode="patch"):
+        super().__init__()
+        self.embed_dim = 2048 # Fixed output dimension for ResNet-50/101
+        self.vision_mode = vision_mode
+
+        if model_name == "resnet50":
+            resnet = models.resnet50(
+                weights=models.ResNet50_Weights.IMAGENET1K_V2
+            )
+        elif model_name == "resnet101":
+            resnet = models.resnet101(
+                weights=models.ResNet101_Weights.IMAGENET1K_V2
+            )
+        else:
+            raise ValueError("model_name must be 'resnet50' or 'resnet101'")
+
+        # Keep encoder layers only (no classifier head)
+        # ResNet children indices: 0:conv1, 1:bn1, 2:relu, 3:maxpool, 4:layer1, 5:layer2, 6:layer3, 7:layer4
+        self.features = nn.Sequential(*list(resnet.children())[:-1])
+        self.model = self.features
+
+
+        # Full Fine-Tuning Mode
+        if fine_tune_all_encoder_layers:
+            print("[INFO] ResNet-50: Fine-tuning ALL layers (1-4).")
+            # PyTorch defaults to requires_grad=True, so no action is needed here.
+            return
+
+        # Full Freezing Mode
+        elif freeze_encoder_entirely:
+            print("[INFO] ResNet-50: Freezing ALL layers (1-4).")
+            for param in self.features.parameters():
+                param.requires_grad = False
+            return
+            
+        # Dynamic Partial Freezing Mode 
+        else:
+            # freeze_until=3 is the default behavior (freeze L1-L3, train L4)
+            # freeze_until=2 means freeze L1-L2, train L3-L4
+            print(f"[INFO] ResNet-50: Dynamic partial fine-tuning (Freezing Layers 1-{freeze_until}).")
+            
+            # To freeze up to layer N (L1, L2, L3, or L4), we freeze all indices from 0 up to N+3.
+            max_freeze_idx = freeze_until + 3 
+            
+            # Create a set of indices to freeze (from 0 up to max_freeze_idx, inclusive)
+            freeze_indices = set(range(max_freeze_idx + 1))
+
+            for idx, layer in enumerate(self.features):
+                # Ensure we only process layers up to the target index
+                if idx in freeze_indices:
+                    for param in layer.parameters():
+                        param.requires_grad = False
+                        
+
+    def forward(self, pixel_values):
+        x = self.features(pixel_values)      # (B, 2048, H, W)
+
+        if self.vision_mode == "cls":
+            x_flat = x.flatten(1)            # (B, 2048*H*W)
+            return {"image_embeds": x_flat}
+
+        tokens = x.flatten(2).transpose(1, 2)  # (B, S, 2048)
+        return {"image_embeds": tokens}
+
+
+    def get_output_dim(self):
+        return self.embed_dim
+
+
+
+# ViT  Encoders
+
+class ViTEncoder(BaseVisionEncoder):
+
+    def __init__(self, model_name="google/vit-base-patch16-224", train_last_n_layers=4, vision_mode="patch"):
+        super().__init__(embed_dim=None) 
+
+        self.model = AutoModel.from_pretrained(model_name)
+        self.vision_mode = vision_mode
+
+        self.embed_dim = self.model.config.hidden_size
+        if self.embed_dim is None:
+            raise ValueError("Could not determine embed_dim from model config.")
+        
+        # Partial Fine-Tuning Strategy
+        # Strategy: Freeze first 8 layers (0-7), train last 4 layers (8-11), 
+        # plus embeddings and final LayerNorm. (Total layers = 12 for ViT-Base)
+        
+        # Freeze all parameters initially
+        for param in self.model.parameters():
+            param.requires_grad = False
+            
+        # Unfreeze the final N transformer blocks
+        NUM_LAYERS_TO_TRAIN = train_last_n_layers
+        
+        try:
+            # The layers are typically stored in .encoder.layer
+            encoder_layers = self.model.encoder.layer
+            num_layers = len(encoder_layers)
+            
+            # Unfreeze the last NUM_LAYERS_TO_TRAIN blocks
+            for i in range(num_layers - NUM_LAYERS_TO_TRAIN, num_layers): 
+                layer = encoder_layers[i] 
+                for param in layer.parameters():
+                    param.requires_grad = True
+            
+            print(f"ViT Encoder: Unfrozen the final {NUM_LAYERS_TO_TRAIN} blocks ({num_layers - NUM_LAYERS_TO_TRAIN} to {num_layers - 1}).")
+            
+        except AttributeError:
+            print("Warning: Could not find standard ViT layer structure for partial fine-tuning.")
+            
+        #  Unfreeze Positional Embeddings (often gives a small boost)
+        if hasattr(self.model.embeddings, 'position_embeddings'):
+            self.model.embeddings.position_embeddings.requires_grad = True
+            print("ViT Encoder: Unfrozen positional embeddings.")
+
+        #  Unfreeze the final LayerNorm (for stabilization)
+        if hasattr(self.model.encoder, 'layernorm'):
+            for param in self.model.encoder.layernorm.parameters():
+                param.requires_grad = True
+            print("ViT Encoder: Unfrozen final LayerNorm.")
+
+
+    def forward(self, pixel_values):
+        out = self.model(pixel_values=pixel_values)
+
+        # CLS MODE
+        if self.vision_mode == "cls":
+            if hasattr(out, 'pooler_output') and out.pooler_output is not None:
+                pooled = out.pooler_output             # (B, D)
+            elif hasattr(out, 'last_hidden_state'):
+                pooled = out.last_hidden_state[:, 0, :]  # CLS token (B, D)
+            else:
+                raise RuntimeError("Model output format not recognized.")
+
+            return {"image_embeds": pooled}
+
+        # PATCH
+        seq = out.last_hidden_state       # (B, S, D)
+        return {"image_embeds": seq}
+
+
+    def get_output_dim(self):
+        return self.embed_dim
+
+
+# Clip Encoders
+
+class CLIPEncoder(BaseVisionEncoder):
+
+    def __init__(self, model_name="openai/clip-vit-base-patch32", train_last_n_layers=4, vision_mode="patch"):
+        # The output dimension (hidden size) will be set after loading the model config
+        super().__init__(embed_dim=None) 
+
+        self.model = CLIPVisionModel.from_pretrained(model_name)
+        self.vision_mode = vision_mode
+        
+        self.embed_dim = self.model.config.hidden_size
+        if self.embed_dim is None:
+            raise ValueError("Could not determine embed_dim from model config.")
+        
+        # Partial Fine-Tuning Strategy 
+        # Strategy: Freeze first 8 layers (0-7), train last 4 layers (8-11), 
+        # plus embeddings and final LayerNorm. (Total layers = 12 for ViT-Base)
+        
+        # Freeze all parameters initially
+        for param in self.model.parameters():
+            param.requires_grad = False
+            
+        #  Unfreeze the final N transformer blocks
+        NUM_LAYERS_TO_TRAIN = train_last_n_layers
+        
+        try:
+            encoder_layers = self.model.vision_model.encoder.layers
+            num_layers = len(encoder_layers)
+            
+            for i in range(num_layers - NUM_LAYERS_TO_TRAIN, num_layers): 
+                layer = encoder_layers[i] 
+                for param in layer.parameters():
+                    param.requires_grad = True
+            
+            print(f"CLIP Encoder: Unfrozen the final {NUM_LAYERS_TO_TRAIN} blocks ({num_layers - NUM_LAYERS_TO_TRAIN} to {num_layers - 1}).")
+            
+        except AttributeError:
+            print("Warning: Could not find standard CLIP layer structure for partial fine-tuning. Ensure model structure is correct.")
+            
+
+        if hasattr(self.model.vision_model.embeddings, 'position_embedding'):
+            self.model.vision_model.embeddings.position_embedding.requires_grad = True
+            print("CLIP Encoder: Unfrozen positional embeddings.")
+
+        if hasattr(self.model.vision_model, 'post_layernorm'):
+            for param in self.model.vision_model.post_layernorm.parameters():
+                param.requires_grad = True
+            print("CLIP Encoder: Unfrozen final LayerNorm.")
+
+ 
+
+    def forward(self, pixel_values):
+        out = self.model(pixel_values=pixel_values)
+        seq = out.last_hidden_state  # (B, S, D)
+
+        if self.vision_mode == "cls":
+            return {"image_embeds": seq[:, 0, :]}  # (B, D)
+
+        return {"image_embeds": seq}               # (B, S, D)
+
+
+
+    def get_output_dim(self):
+        return self.embed_dim

models/vision_t5.py

@@ -0,0 +1,187 @@
+import torch
+import torch.nn as nn
+from transformers import T5ForConditionalGeneration
+from transformers.modeling_outputs import BaseModelOutput
+from peft import LoraConfig, get_peft_model
+import os
+import requests
+from transformers import AutoConfig
+
+
+def safe_load_t5(model_name, local_path):
+    has_local = os.path.exists(local_path)
+
+    try:
+        print(f"[INFO] Trying to load {model_name} from HuggingFace…")
+        model = T5ForConditionalGeneration.from_pretrained(model_name)
+        print("[INFO] Loaded from HF successfully.")
+        return model
+    
+    except Exception as e:
+        print(f"[WARN] HF load failed: {e}")
+
+        if not has_local:
+            raise RuntimeError(
+                f"No local copy available at {local_path} and HF download failed."
+            )
+
+        print("[INFO] Falling back to local Drive copy...")
+        return T5ForConditionalGeneration.from_pretrained(local_path)
+
+
+class VisionT5(nn.Module):
+
+
+    def __init__(self, vision_encoder, projector, t5_name="t5-small", decoder_params=None):
+        super().__init__()
+
+        decoder_params = decoder_params or {}
+
+        self.vision_encoder = vision_encoder
+        self.projector = projector
+
+        # Load full T5, but we only use decoder
+        local_large = "/content/drive/MyDrive/Models/t5-large"
+
+        if t5_name == "t5-large":
+            self.t5 = safe_load_t5("t5-large", local_large)
+        else:
+            self.t5 = T5ForConditionalGeneration.from_pretrained(t5_name)
+            
+        self.apply_decoder_options(decoder_params)
+        
+        for p in self.t5.encoder.parameters():
+            p.requires_grad = False
+
+        self.hidden_size = self.t5.config.d_model
+
+
+    def apply_decoder_options(self, params):
+
+        # LoRA setup
+        if params.get("use_lora", False):
+            lora_rank = params.get("lora_rank", 8)
+            lora_alpha = params.get("lora_alpha", 16)
+            
+            print(f"[INFO] LoRA enabled for T5 decoder (Rank={lora_rank})")
+            
+            # Target query and value matrices in all T5 attention blocks
+            lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_alpha,
+                target_modules=["q", "v"],
+                lora_dropout=params.get("lora_dropout", 0.1),
+                bias="none",
+                task_type="CAUSAL_LM"
+            )
+            
+            self.t5 = get_peft_model(self.t5, lora_config)
+            
+            self.t5.print_trainable_parameters()
+            
+            # The freeze_decoder flag (if present) is ignored when using LoRA, as LoRA automatically handles freezing and only exposes the adapter weights.
+
+            return
+
+
+        num_decoder_layers = self.t5.config.num_decoder_layers # 
+        
+        trainable_layers = params.get("trainable_decoder_layers")
+
+        if trainable_layers is not None:
+            num_frozen = num_decoder_layers - trainable_layers
+            
+            if num_frozen > 0:
+                print(f"[INFO] Partial Tuning: Freezing first {num_frozen} of {num_decoder_layers} decoder blocks.")
+                
+                for i, block in enumerate(self.t5.decoder.block):
+                    if i < num_frozen:
+                        for p in block.parameters():
+                            p.requires_grad = False
+                        print(f"  > Block {i} frozen.")
+                    else:
+                        for p in block.parameters():
+                            p.requires_grad = True
+                        print(f"  > Block {i} trainable.")
+                
+                if num_frozen > 0:
+                    for p in self.t5.decoder.embed_tokens.parameters():
+                        p.requires_grad = False
+                    print("  > Decoder embeddings frozen.")
+                    
+                return
+
+        if params.get("freeze_decoder", False):
+            print("[INFO] Freezing all T5 decoder parameters.")
+            for p in self.t5.decoder.parameters():
+                p.requires_grad = False
+        
+        if params.get("dropout_override") is not None:
+            self.t5.config.dropout_rate = params["dropout_override"]
+
+
+    def forward(
+        self,
+        pixel_values=None,
+        input_ids=None,
+        attention_mask=None,
+        labels=None
+    ):
+
+        vision_out = self.vision_encoder(pixel_values)
+        image_embeds = vision_out["image_embeds"]
+
+        if image_embeds.dim() == 2:
+            image_embeds = image_embeds.unsqueeze(1)
+
+        projected = self.projector(image_embeds)
+
+        B, S, _ = projected.shape
+        encoder_attention_mask = torch.ones(B, S, dtype=torch.long, device=projected.device)
+
+        encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+        decoder_attention_mask = attention_mask
+
+        output = self.t5(
+            input_ids=input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            attention_mask=encoder_attention_mask, 
+            encoder_outputs=encoder_outputs,
+            labels=labels,
+            return_dict=True,
+        )
+
+        return output
+
+    @torch.no_grad()
+    def generate(self, pixel_values, tokenizer, max_length=32, num_beams=3):
+
+        vision_out = self.vision_encoder(pixel_values)
+        image_embeds = vision_out["image_embeds"]
+
+        if image_embeds.dim() == 2:
+            image_embeds = image_embeds.unsqueeze(1) # (B, 1, D)
+
+        projected = self.projector(image_embeds)    # (B, S, d_model)
+
+        encoder_outputs = BaseModelOutput(
+            last_hidden_state=projected
+        )
+        
+        generated_ids = self.t5.generate(
+            encoder_outputs=encoder_outputs,
+            decoder_start_token_id=self.t5.config.decoder_start_token_id,
+            input_ids=torch.tensor([[tokenizer.pad_token_id]]).to(projected.device),
+            max_length=max_length,
+            num_beams=num_beams
+        )
+
+        return tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+
+    @staticmethod
+    def get_t5_hidden_size(t5_name):
+
+        cfg = AutoConfig.from_pretrained(t5_name)
+        return cfg.d_model