Example_Cross_Attention_for_Text_Image_Integration.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyNVwi/fEG6zDKKnnmAKyfyN",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OneFineStarstuff/OneFineStardust/blob/main/Example_Cross_Attention_for_Text_Image_Integration.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from transformers import AutoModel, AutoTokenizer, CLIPProcessor, CLIPModel\n",
        "from PIL import Image\n",
        "import requests\n",
        "from io import BytesIO\n",
        "\n",
        "class CrossAttentionModel(nn.Module):\n",
        "    def __init__(self, text_model, image_model, embed_dim=768, num_heads=8):\n",
        "        super().__init__()\n",
        "        self.text_model = text_model\n",
        "        self.image_model = image_model\n",
        "        self.cross_attention = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads)\n",
        "\n",
        "    def forward(self, text_inputs, text_attention_mask, pixel_values):\n",
        "        text_outputs = self.text_model(input_ids=text_inputs, attention_mask=text_attention_mask)\n",
        "        text_feats = text_outputs.last_hidden_state.transpose(0, 1)  # (seq_len, batch, embed_dim)\n",
        "\n",
        "        image_outputs = self.image_model(pixel_values=pixel_values)\n",
        "        image_feats = image_outputs.last_hidden_state.transpose(0, 1)  # (seq_len, batch, embed_dim)\n",
        "\n",
        "        # Cross-attention between text and image features\n",
        "        attn_output, _ = self.cross_attention(text_feats, image_feats, image_feats)\n",
        "        return attn_output\n",
        "\n",
        "# Example usage\n",
        "# Load a text model and image model from Hugging Face\n",
        "text_model = AutoModel.from_pretrained(\"bert-base-uncased\")\n",
        "image_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "\n",
        "# Initialize the CrossAttentionModel with the text and image models\n",
        "cross_attention_model = CrossAttentionModel(text_model, image_model)\n",
        "\n",
        "# Example input\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
        "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "text_input = tokenizer(\"This is an example sentence.\", return_tensors=\"pt\")\n",
        "\n",
        "# Using a valid image URL\n",
        "image_url = \"https://images.unsplash.com/photo-1516637090014-3aceb6cd25b3\"  # Replace with a valid image URL\n",
        "response = requests.get(image_url)\n",
        "image = Image.open(BytesIO(response.content))\n",
        "\n",
        "image_input = processor(images=image, return_tensors=\"pt\")[\"pixel_values\"]\n",
        "\n",
        "# Forward pass through the cross-attention model\n",
        "output = cross_attention_model(text_input[\"input_ids\"], text_input[\"attention_mask\"], image_input)\n",
        "\n",
        "print(output)"
      ],
      "metadata": {
        "id": "JwvlAehLuIQF"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}