BLIP-2とは？

import torch
import torch.nn as nn
import torch.nn.functional as F
    OPTForCausalLM,
    CLIPImageProcessor,
    AutoTokenizer
)
from PIL import Image
import requests
import matplotlib.pyplot as plt


class SimpleQFormer(nn.Module):
    """画像特徴量から必要な情報を抽出する簡易Q-Former"""
    def __init__(self, vision_hidden_size, qformer_hidden_size, llm_hidden_size, num_queries=32):
        super().__init__()
        # 32個の学習可能なクエリトークン
        self.query_tokens = nn.Parameter(torch.randn(1, num_queries, qformer_hidden_size))

        # Cross-Attentionを模倣するTransformerDecoder
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=qformer_hidden_size, nhead=8, batch_first=True
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=4)

        # LLMの入力次元に合わせるプロジェクション層
        self.llm_proj = nn.Linear(qformer_hidden_size, llm_hidden_size)

    def forward(self, vision_features):
        batch_size = vision_features.shape[0]
        queries = self.query_tokens.expand(batch_size, -1, -1)
        qformer_output = self.transformer(tgt=queries, memory=vision_features)
        visual_prompts = self.llm_proj(qformer_output)
        return visual_prompts

class TinyBLIP2(nn.Module):
    """CLIPとOPTをQ-Formerで繋ぐ簡易VLMアーキテクチャ"""
    def __init__(self):
        super().__init__()
        # 凍結する画像エンコーダ
        self.vision_model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32")
        for param in self.vision_model.parameters():
            param.requires_grad = False

        # 凍結する言語モデル
        self.llm_model = OPTForCausalLM.from_pretrained("facebook/opt-125m")
        for param in self.llm_model.parameters():
            param.requires_grad = False

        # 学習対象のQ-Former
        vision_dim = self.vision_model.config.hidden_size
        llm_dim = self.llm_model.config.hidden_size
        self.qformer = SimpleQFormer(
            vision_hidden_size=vision_dim,
            qformer_hidden_size=768,
            llm_hidden_size=llm_dim,
            num_queries=32
        )

    def forward(self, pixel_values, input_ids, attention_mask=None, labels=None):
        with torch.no_grad():
            vision_features = self.vision_model(pixel_values=pixel_values).last_hidden_state

        visual_prompts = self.qformer(vision_features) # [Batch, 32, LLM_Dim]

        with torch.no_grad():
            text_embeddings = self.llm_model.get_input_embeddings()(input_ids)

        # 視覚プロンプトとテキストを結合
        inputs_embeds = torch.cat([visual_prompts, text_embeddings], dim=1)

        if attention_mask is not None:
            visual_mask = torch.ones(
                attention_mask.shape[0], visual_prompts.shape[1],
                dtype=attention_mask.dtype,
                device=attention_mask.device
            )
            attention_mask = torch.cat([visual_mask, attention_mask], dim=1)

        # Loss計算用のラベルマスク処理
        if labels is not None:
            ignore_index = -100
            visual_labels = torch.full(
                (labels.shape[0], visual_prompts.shape[1]),
                ignore_index,
                dtype=labels.dtype,
                device=labels.device
            )
            labels = torch.cat([visual_labels, labels], dim=1)

        outputs = self.llm_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            labels=labels
        )
        return outputs

class VLMDataset(Dataset):
    """URLから画像を取得し、テキストと共にテンソルに変換するデータセット"""
    def __init__(self, hf_dataset, image_processor, tokenizer, max_length=32):
        self.dataset = hf_dataset
        self.image_processor = image_processor
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        # 最大試行回数を設定して無限ループを防ぐ
        for i in range(10):  # 最大10回まで試す
            try:
                item = self.dataset[(idx + i) % len(self.dataset)]
                url = item["url"]


                headers = {
                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
                }
                response = requests.get(url, headers=headers, stream=True, timeout=10)
                if response.status_code != 200:
                    continue # 次へ

                image = Image.open(response.raw).convert("RGB")

                # 正常に読み込めたら処理を続行
                pixel_values = self.image_processor(images=image, return_tensors="pt").pixel_values.squeeze(0)

                caption = item["caption"] + self.tokenizer.eos_token

                encoded_text = self.tokenizer(
                    caption, padding="max_length", truncation=True,
                    max_length=self.max_length, return_tensors="pt",
                    add_special_tokens=False
                )

                input_ids = encoded_text.input_ids.squeeze(0)
                attention_mask = encoded_text.attention_mask.squeeze(0)

                labels = input_ids.clone()
                labels[labels == self.tokenizer.pad_token_id] = -100

                return {
                    "pixel_values": pixel_values,
                    "input_ids": input_ids,
                    "attention_mask": attention_mask,
                    "labels": labels
                }
            except Exception as e:
                # 失敗した場合は次のデータへ
                continue

        # 10回試してもダメなら、仕方なく最後のエラーを発生させるか、ダミーを返す
        raise RuntimeError("有効な画像を読み込めませんでした")

print("モデルとプロセッサを初期化中...")
model = TinyBLIP2()
image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

hf_dataset = load_dataset("Spawning/megalith-cc0", split="train[:1]")
# データセットの構造を確認
train_dataset = VLMDataset(hf_dataset, image_processor, tokenizer)
train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True)

device = "cuda" if torch.cuda.is_available() else "cpu"

# モデル全体（CLIP, OPT, Q-Former全て）を確実に単精度（float32）に統一する
model.to(device, dtype=torch.float32)
model.train()

# 凍結されていないパラメータ（Q-Former関連）のみをオプティマイザに渡す
trainable_params = [p for p in model.parameters() if p.requires_grad]
optimizer = AdamW(trainable_params, lr=1e-4)

print("\n--- 【Stage 1】 簡易的な表現学習（アラインメント）を開始 ---")
# 1件のデータを過学習させる設定
epochs_stage1 = 20

for epoch in range(epochs_stage1):
    total_loss = 0
    for batch in train_dataloader:  # 1件しか入っていないDataLoader
        pixel_values = batch["pixel_values"].to(device, dtype=torch.float32)
        input_ids = batch["input_ids"].to(device)

        optimizer.zero_grad()

        # 1. Q-Formerで画像から「視覚プロンプト（32個）」を抽出
        with torch.no_grad():
            vision_features = model.vision_model(pixel_values=pixel_values).last_hidden_state
        visual_prompts = model.qformer(vision_features) # [Batch, 32, 768]

        # 2. 正解テキストから「テキスト埋め込み」を抽出
        with torch.no_grad():
            text_embeds = model.llm_model.get_input_embeddings()(input_ids) # [Batch, SeqLen, 768]

        # 3. 両者の特徴量を要約（平均プーリング）して1つのベクトルにする
        visual_repr = visual_prompts.mean(dim=1)
        text_repr = text_embeds.mean(dim=1)

        # 4. コサイン類似度ベースのLossを計算（1に近いほど正解とする）
        target = torch.ones(visual_repr.size(0)).to(device)
        loss = F.cosine_embedding_loss(visual_repr, text_repr, target)

        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    if (epoch + 1) % 5 == 0:
        print(f"Stage 1 - Epoch {epoch+1} | Cosine Loss: {total_loss:.4f}")

print("--- Stage 1 完了 ---\n")

print(f"--- 【Stage 2】 生成学習（キャプション生成）を開始 (Device: {device}) ---")
epochs = 200
# （バッチサイズは1件しかないので自動的に1になります）
for epoch in range(epochs):
    total_loss = 0
    for batch in train_dataloader:
        # 画像テンソルも明示的に float32 に変換してモデルに渡す
        pixel_values = batch["pixel_values"].to(device, dtype=torch.float32)

        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        optimizer.zero_grad()
        outputs = model(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    if (epoch+1) % 50 == 0:
        print(f"Epoch {epoch+1} 完了 | 平均Loss: {total_loss / len(train_dataloader):.4f}")
print("--- Stage 2 完了 ---")

def generate_caption(model, image_processor, tokenizer, image, device="cpu"):
    """学習済みモデルに未知の画像を渡し、キャプションを生成する"""
    model.eval()
    with torch.no_grad():
        pixel_values = image_processor(images=image, return_tensors="pt").pixel_values.to(device, dtype=torch.float32)
        vision_features = model.vision_model(pixel_values=pixel_values).last_hidden_state
        visual_prompts = model.qformer(vision_features)

        attention_mask = torch.ones(
            visual_prompts.shape[0], visual_prompts.shape[1], device=device
        )

        generated_ids = model.llm_model.generate(
            inputs_embeds=visual_prompts,
            attention_mask=attention_mask,
            max_new_tokens=150,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            do_sample=False
        )
        return tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()

sample = hf_dataset[0]
url = sample["url"]
overfit_image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
target_text = sample["caption"]

print(f"\n[正解のテキスト]: {target_text}")
# 推論実行
generated_text = generate_caption(model, image_processor, tokenizer, overfit_image, device)

print(f"[モデルの出力]: {generated_text}")

# --- 画像と生成テキストの可視化 ---
plt.figure(figsize=(6, 6))
plt.imshow(overfit_image)
plt.axis("off") # 軸のメモリを非表示にして綺麗に見せる

# 画像のタイトルとして生成されたキャプションを表示（長ければ自動で折り返し）
plt.title(f"Target: {target_text}\nGenerated: {generated_text}",
          fontsize=12, wrap=True, pad=20, backgroundcolor="white")

plt.show()

モデルとプロセッサを初期化中...
--- 【Stage 1】 簡易的な表現学習（アラインメント）を開始 ---
Stage 1 - Epoch 5 | Cosine Loss: 0.2278
Stage 1 - Epoch 10 | Cosine Loss: 0.0557
Stage 1 - Epoch 15 | Cosine Loss: 0.0154
Stage 1 - Epoch 20 | Cosine Loss: 0.0120
--- Stage 1 完了 ---

--- 【Stage 2】 生成学習（キャプション生成）を開始 (Device: cuda) ---
Epoch 50 完了 | 平均Loss: 1.4167
Epoch 100 完了 | 平均Loss: 1.0563
Epoch 150 完了 | 平均Loss: 0.8518
Epoch 200 完了 | 平均Loss: 0.8577
--- Stage 2 完了 ---

[正解のテキスト]: The image shows two ancient artifacts, possibly from the Roman Empire, displayed on a beige background. The artifacts appear to be made of metal and have a greenish-brown color. They have a pointed tip with a small handle on one end and a smaller handle on the other end. The handle is decorated with a floral design and has a small loop at the top for attaching to a belt or bag. The blade of the artifact is slightly curved and appears to be slightly pointed. The overall appearance of the artifacts is weathered and aged.
[モデルの出力]: The image shows two ancient artifacts, possibly from the Roman Empire, displayed on a beige background. The artifacts appear to be made of metal and have a green beige color, and are displayed on a beige background. The artifacts appear to be made of metal and have a green beige background. The artifacts appear to be made of metal and have a green beige color, and the artifacts appear to be made of metal. The artifacts appear to be made of metal and have a green beige background. The artifacts appear to be made of metal and have a green beige background. The artifacts appear to be made of metal and have a green beige background. The artifacts appear to be made of metal and have a green beige

import torch
from PIL import Image
import requests
from transformers import Blip2Processor, Blip2ForConditionalGeneration
import matplotlib.pyplot as plt

# デバイスの設定
device = "cuda" if torch.cuda.is_available() else "cpu"

# プロセッサとモデルの読み込み
model_id = "Salesforce/blip2-opt-2.7b"
processor = Blip2Processor.from_pretrained(model_id)
model = Blip2ForConditionalGeneration.from_pretrained(
    model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)

# 画像の準備(Huggingface Dataset Spawning/megalith-cc0 ID:f81bdee3f9e7da521751c86a35f1baf0)
img_url = "https://megalith-8m.s3.us-west-2.amazonaws.com/images/f10e65af-3cd6-5100-a0e7-50c2e43f4974.jpeg" 
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')

# 入力データの準備
inputs = processor(raw_image, return_tensors="pt").to(device, torch.float16 if device == "cuda" else torch.float32)

# キャプション生成
generated_ids = model.generate(**inputs, max_new_tokens=20)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()

plt.figure(figsize=(8, 6))
plt.imshow(raw_image)
plt.axis('off')  # 軸を非表示にする
plt.title(f"Caption: {generated_text}", fontsize=12)
plt.show()

print(f"Generated caption: {generated_text}")