[入力画像 (224×224)]
    ↓
[CLIP ViT-L/14] (凍結)
    ↓
[画像特徴 (256 tokens × 1024)]
    ↓
[線形プロジェクション] (学習)
    ↓
[画像トークン (256 tokens × 4096)]
    ↓
[画像トークン] + [テキストトークン]
    ↓
[LLaMA / Vicuna] (LoRAまたはフル学習)
    ↓
[出力テキスト]

# CLIP ViT-L/14の出力
# image: (B, 3, 224, 224)
# features: (B, 257, 1024)  # 256パッチ + 1 CLSトークン
# LLaVAでは256パッチトークンのみ使用: (B, 256, 1024)

<s> USER: <image>
この画像について説明してください。
ASSISTANT: この画像には...
</s>

A chat between a curious user and an artificial intelligence assistant.
The assistant gives helpful, detailed, and polite answers to the user's questions.

USER: <image>
{ユーザーの質問}
ASSISTANT: {モデルの応答}

USER: <image>
この画像には何が写っていますか？
ASSISTANT: この画像には赤いリンゴがテーブルの上に置かれています。

USER: そのリンゴはどのような状態ですか？
ASSISTANT: リンゴは新鮮で、艶があり、傷のない状態です。

入力: <image> この画像を簡潔に説明してください。
出力: {キャプション}

You are an AI assistant helping to create training data for a visual AI model.

Given the following image description:
Caption: "A person riding a bicycle on a sunny day"
Objects: [bicycle (person, wheels, handlebar), sun, trees]

Generate a conversation between a user and an assistant about this image.
The conversation should be natural and informative.

import torch
import torch.nn as nn
from transformers import CLIPVisionModel, LlamaForCausalLM, LlamaTokenizer


class LLaVA(nn.Module):
    """LLaVA: Large Language and Vision Assistant"""
    def __init__(
        self,
        vision_model_name="openai/clip-vit-large-patch14",
        llm_model_name="meta-llama/Llama-2-7b-hf",
        freeze_vision=True,
        freeze_llm=False,
    ):
        super().__init__()

        # 画像エンコーダ（CLIP ViT）
        self.vision_encoder = CLIPVisionModel.from_pretrained(vision_model_name)
        self.vision_hidden_size = self.vision_encoder.config.hidden_size  # 1024

        # 凍結設定
        if freeze_vision:
            for param in self.vision_encoder.parameters():
                param.requires_grad = False

        # LLM
        self.llm = LlamaForCausalLM.from_pretrained(llm_model_name)
        self.llm_hidden_size = self.llm.config.hidden_size  # 4096

        if freeze_llm:
            for param in self.llm.parameters():
                param.requires_grad = False

        # プロジェクション（2層MLP）
        self.vision_projection = nn.Sequential(
            nn.Linear(self.vision_hidden_size, self.llm_hidden_size),
            nn.GELU(),
            nn.Linear(self.llm_hidden_size, self.llm_hidden_size),
        )

        # 特殊トークンのIDを記録
        self.image_token_id = None  # <image>トークン

    def encode_image(self, images):
        """
        画像をLLMトークン空間にエンコード

        Args:
            images: (B, 3, 224, 224)
        Returns:
            (B, num_patches, llm_hidden_size)
        """
        # CLIP ViTで特徴抽出
        vision_outputs = self.vision_encoder(images)
        # last_hidden_state: (B, num_patches+1, hidden_size)
        # CLSトークンを除く
        image_features = vision_outputs.last_hidden_state[:, 1:, :]

        # プロジェクション
        image_tokens = self.vision_projection(image_features)

        return image_tokens

    def prepare_inputs(self, images, input_ids, attention_mask, image_positions):
        """
        画像トークンをテキストに挿入

        Args:
            images: (B, 3, H, W)
            input_ids: (B, seq_len)
            attention_mask: (B, seq_len)
            image_positions: (B,) 各サンプルでの<image>トークンの位置
        """
        batch_size = images.shape[0]

        # 画像エンコード
        image_tokens = self.encode_image(images)  # (B, num_img_tokens, hidden)
        num_img_tokens = image_tokens.shape[1]

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

        # 画像トークンを<image>位置に挿入
        new_embeds_list = []
        new_attention_mask_list = []

        for b in range(batch_size):
            pos = image_positions[b]

            # <image>トークンの前後で分割
            before = text_embeds[b, :pos]
            after = text_embeds[b, pos+1:]  # <image>トークンをスキップ

            # 結合: [before] + [image_tokens] + [after]
            new_embed = torch.cat([before, image_tokens[b], after], dim=0)
            new_embeds_list.append(new_embed)

            # アテンションマスクも同様に処理
            before_mask = attention_mask[b, :pos]
            after_mask = attention_mask[b, pos+1:]
            img_mask = torch.ones(num_img_tokens, device=attention_mask.device)
            new_mask = torch.cat([before_mask, img_mask, after_mask], dim=0)
            new_attention_mask_list.append(new_mask)

        # パディングして揃える
        max_len = max(e.shape[0] for e in new_embeds_list)
        new_embeds = torch.zeros(batch_size, max_len, self.llm_hidden_size, device=images.device)
        new_attention_mask = torch.zeros(batch_size, max_len, device=images.device)

        for b in range(batch_size):
            length = new_embeds_list[b].shape[0]
            new_embeds[b, :length] = new_embeds_list[b]
            new_attention_mask[b, :length] = new_attention_mask_list[b]

        return new_embeds, new_attention_mask

    def forward(
        self,
        images,
        input_ids,
        attention_mask,
        image_positions,
        labels=None,
    ):
        """
        順伝播

        Args:
            images: (B, 3, H, W)
            input_ids: (B, seq_len) <image>トークンを含む
            attention_mask: (B, seq_len)
            image_positions: (B,) <image>トークンの位置
            labels: (B, seq_len) 学習ターゲット
        """
        # 入力を準備
        inputs_embeds, new_attention_mask = self.prepare_inputs(
            images, input_ids, attention_mask, image_positions
        )

        # ラベルも調整（必要に応じて）
        if labels is not None:
            num_img_tokens = (inputs_embeds.shape[1] - input_ids.shape[1] + 1)
            # 画像トークン部分は損失計算から除外
            new_labels = self._adjust_labels(labels, image_positions, num_img_tokens)
        else:
            new_labels = None

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

        return outputs

    def _adjust_labels(self, labels, image_positions, num_img_tokens):
        """ラベルを画像トークンに合わせて調整"""
        batch_size = labels.shape[0]
        new_labels_list = []

        for b in range(batch_size):
            pos = image_positions[b]
            before = labels[b, :pos]
            after = labels[b, pos+1:]
            # 画像トークン部分は-100（無視）
            img_labels = torch.full((num_img_tokens,), -100, device=labels.device)
            new_label = torch.cat([before, img_labels, after], dim=0)
            new_labels_list.append(new_label)

        # パディング
        max_len = max(l.shape[0] for l in new_labels_list)
        new_labels = torch.full((batch_size, max_len), -100, device=labels.device)
        for b in range(batch_size):
            length = new_labels_list[b].shape[0]
            new_labels[b, :length] = new_labels_list[b]

        return new_labels

    @torch.no_grad()
    def generate(
        self,
        images,
        input_ids,
        attention_mask,
        image_positions,
        max_new_tokens=256,
        **generate_kwargs,
    ):
        """画像とプロンプトからテキストを生成"""
        inputs_embeds, new_attention_mask = self.prepare_inputs(
            images, input_ids, attention_mask, image_positions
        )

        outputs = self.llm.generate(
            inputs_embeds=inputs_embeds,
            attention_mask=new_attention_mask,
            max_new_tokens=max_new_tokens,
            **generate_kwargs,
        )

        return outputs

def train_llava(model, train_dataloader, optimizer, scheduler, num_epochs):
    """LLaVAの学習"""
    model.train()

    for epoch in range(num_epochs):
        total_loss = 0

        for batch in train_dataloader:
            images = batch['images'].cuda()
            input_ids = batch['input_ids'].cuda()
            attention_mask = batch['attention_mask'].cuda()
            image_positions = batch['image_positions'].cuda()
            labels = batch['labels'].cuda()

            optimizer.zero_grad()

            outputs = model(
                images=images,
                input_ids=input_ids,
                attention_mask=attention_mask,
                image_positions=image_positions,
                labels=labels,
            )

            loss = outputs.loss
            loss.backward()

            # 勾配クリッピング
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            optimizer.step()
            scheduler.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(train_dataloader)
        print(f"Epoch {epoch+1}: Loss = {avg_loss:.4f}")

def inference_example(model, tokenizer, image_processor, image_path, prompt):
    """推論の例"""
    from PIL import Image

    # 画像の読み込みと前処理
    image = Image.open(image_path).convert('RGB')
    pixel_values = image_processor(images=image, return_tensors='pt').pixel_values

    # プロンプトの準備
    full_prompt = f"USER: <image>\n{prompt}\nASSISTANT:"
    inputs = tokenizer(full_prompt, return_tensors='pt')

    # <image>トークンの位置を特定
    image_token_id = tokenizer.convert_tokens_to_ids('<image>')
    image_positions = (inputs.input_ids == image_token_id).nonzero(as_tuple=True)[1]

    # 生成
    model.eval()
    with torch.no_grad():
        outputs = model.generate(
            images=pixel_values.cuda(),
            input_ids=inputs.input_ids.cuda(),
            attention_mask=inputs.attention_mask.cuda(),
            image_positions=image_positions.cuda(),
            max_new_tokens=256,
            temperature=0.2,
            do_sample=True,
        )

    # デコード
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return response