[画像入力]        [テキスト入力]
    ↓                  ↓
[画像エンコーダ]   [テキストトークナイザ]
    ↓                  ↓
[画像特徴]         [テキスト埋め込み]
    ↓                  ↓
    └──────┬───────────┘
           ↓
    [融合モジュール]
           ↓
    [LLMバックボーン]
           ↓
    [テキスト出力]

[画像トークン] + [テキストトークン]
            ↓
      [Transformer層 ×N]
            ↓
      [出力トークン]

[画像] → [画像エンコーダ] → [画像特徴]
                                ↓
[テキスト] → [テキスト埋め込み] → [Cross-Attention] → [LLM]

# 2×2プーリングで256→64トークン
pooled = F.avg_pool2d(tokens.view(B, H, W, D).permute(0, 3, 1, 2), 2)

import torch
import torch.nn as nn


class SimpleMultimodalLLM(nn.Module):
    """シンプルなマルチモーダルLLM"""
    def __init__(
        self,
        vision_encoder,
        llm,
        vision_dim=1024,
        llm_dim=4096,
        num_vision_tokens=256,
    ):
        super().__init__()

        # 画像エンコーダ（事前学習済み、凍結）
        self.vision_encoder = vision_encoder
        for param in self.vision_encoder.parameters():
            param.requires_grad = False

        # プロジェクション（学習対象）
        self.vision_projection = nn.Sequential(
            nn.Linear(vision_dim, llm_dim),
            nn.GELU(),
            nn.Linear(llm_dim, llm_dim),
        )

        # LLM（LoRAなどで効率的に学習）
        self.llm = llm

    def encode_image(self, image):
        """画像をLLMトークン空間にエンコード"""
        with torch.no_grad():
            # (B, num_patches, vision_dim)
            vision_features = self.vision_encoder(image)

        # (B, num_patches, llm_dim)
        vision_tokens = self.vision_projection(vision_features)

        return vision_tokens

    def forward(self, image, input_ids, attention_mask=None, labels=None):
        """
        Args:
            image: (B, 3, H, W) 入力画像
            input_ids: (B, seq_len) テキストトークンID
            attention_mask: (B, seq_len) アテンションマスク
            labels: (B, seq_len) 学習ターゲット
        """
        # 画像をエンコード
        vision_tokens = self.encode_image(image)  # (B, num_img_tokens, llm_dim)

        # テキスト埋め込み
        text_embeds = self.llm.get_input_embeddings()(input_ids)  # (B, seq_len, llm_dim)

        # 画像トークンとテキストトークンを結合
        # [IMG_1, IMG_2, ..., IMG_N, TEXT_1, TEXT_2, ...]
        inputs_embeds = torch.cat([vision_tokens, text_embeds], dim=1)

        # アテンションマスクも調整
        num_img_tokens = vision_tokens.shape[1]
        img_attention_mask = torch.ones(
            vision_tokens.shape[0], num_img_tokens,
            dtype=attention_mask.dtype, device=attention_mask.device
        )
        full_attention_mask = torch.cat([img_attention_mask, attention_mask], dim=1)

        # ラベルも調整（画像トークン部分は-100で無視）
        if labels is not None:
            img_labels = torch.full(
                (labels.shape[0], num_img_tokens),
                -100,  # 損失計算で無視
                dtype=labels.dtype, device=labels.device
            )
            full_labels = torch.cat([img_labels, labels], dim=1)
        else:
            full_labels = None

        # LLMに入力
        outputs = self.llm(
            inputs_embeds=inputs_embeds,
            attention_mask=full_attention_mask,
            labels=full_labels,
        )

        return outputs

    @torch.no_grad()
    def generate(self, image, prompt_ids, max_new_tokens=100, **kwargs):
        """画像とプロンプトから生成"""
        vision_tokens = self.encode_image(image)
        prompt_embeds = self.llm.get_input_embeddings()(prompt_ids)
        inputs_embeds = torch.cat([vision_tokens, prompt_embeds], dim=1)

        # 生成
        outputs = self.llm.generate(
            inputs_embeds=inputs_embeds,
            max_new_tokens=max_new_tokens,
            **kwargs
        )

        return outputs

# 仮想的な使用例

# モデル作成
from transformers import CLIPVisionModel, AutoModelForCausalLM

vision_encoder = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14")
llm = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

model = SimpleMultimodalLLM(
    vision_encoder=vision_encoder,
    llm=llm,
    vision_dim=1024,
    llm_dim=4096,
)

# 推論
image = load_image("example.jpg")  # (1, 3, 224, 224)
prompt = "この画像について説明してください。"
prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids

output_ids = model.generate(image, prompt_ids, max_new_tokens=100)
response = tokenizer.decode(output_ids[0])
print(response)