inference_min.py

# -*- coding: utf-8 -*-
from __future__ import annotations
import os, json, re
from pathlib import Path
from typing import List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

try:
    import open_clip
    HAS_OPENCLIP = True
except Exception:
    HAS_OPENCLIP = False

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    CLIPImageProcessor as HFCLIPImageProcessor,
    CLIPModel as HFCLIPModel,
)

class PrefixProjector(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, tokens: int, p_drop: float = 0.05):
        super().__init__()
        hidden = max(512, out_dim * 2)
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, out_dim * tokens)
        self.ln  = nn.LayerNorm(out_dim)
        self.tokens = tokens
        self.drop = nn.Dropout(p_drop)
        self.alpha = nn.Parameter(torch.tensor(0.5))
        nn.init.xavier_uniform_(self.fc1.weight, gain=1.0)
        nn.init.zeros_(self.fc1.bias)
        nn.init.xavier_uniform_(self.fc2.weight, gain=0.5)
        nn.init.zeros_(self.fc2.bias)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.gelu(self.fc1(x))
        y = self.fc2(y).view(x.size(0), self.tokens, -1)
        y = self.ln(y)
        y = self.drop(self.alpha * y)
        return y

class CLIPBackend:
    def __init__(self, repo_or_kind: str, device: str):
        self.device = device
        self.repo_or_kind = repo_or_kind
        
        # Определяем тип модели
        if 'BiomedCLIP' in repo_or_kind or 'microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224' in repo_or_kind:
            # BiomedCLIP через open_clip
            assert HAS_OPENCLIP, "open_clip is required for BiomedCLIP"
            if not repo_or_kind.startswith('microsoft/'):
                repo_or_kind = 'microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224'
            model_name = f'hf-hub:{repo_or_kind}'
            self.model, self.preprocess, _ = open_clip.create_model_and_transforms(model_name)
            self.model = self.model.to(device).eval()
            self.kind = "open_clip"
            self.processor = None
        elif "/" in repo_or_kind and 'pubmed-clip' in repo_or_kind:
            # PubMedCLIP через HF
            self.model = HFCLIPModel.from_pretrained(repo_or_kind).to(device).eval()
            self.processor = HFCLIPImageProcessor.from_pretrained(repo_or_kind)
            self.kind = "hf_clip"
            self.preprocess = None
        elif "/" in repo_or_kind or repo_or_kind.startswith('redlessone/'):
            # DermLIP через open_clip
            assert HAS_OPENCLIP, "open_clip is required for DermLIP"
            model_name = f"hf-hub:{repo_or_kind}"
            self.model, self.preprocess, _ = open_clip.create_model_and_transforms(model_name)
            self.model = self.model.to(device).eval()
            self.kind = "open_clip"
            self.processor = None
        else:
            # Fallback для других моделей, включая случаи когда передается просто тип модели
            try:
                # Пытаемся определить по названию
                if 'biomedclip' in repo_or_kind.lower() or 'biomed' in repo_or_kind.lower():
                    assert HAS_OPENCLIP, "open_clip is required for BiomedCLIP"
                    model_name = "hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"
                    self.model, self.preprocess, _ = open_clip.create_model_and_transforms(model_name)
                    self.model = self.model.to(device).eval()
                    self.kind = "open_clip"
                    self.processor = None
                elif 'dermlip' in repo_or_kind.lower():
                    assert HAS_OPENCLIP, "open_clip is required for DermLIP"
                    model_name = "hf-hub:redlessone/DermLIP_ViT-B-16"
                    self.model, self.preprocess, _ = open_clip.create_model_and_transforms(model_name)
                    self.model = self.model.to(device).eval()
                    self.kind = "open_clip"
                    self.processor = None
                elif 'pubmed' in repo_or_kind.lower():
                    # PubMedCLIP через HF
                    repo_name = "flaviagiammarino/pubmed-clip-vit-base-patch32"
                    self.model = HFCLIPModel.from_pretrained(repo_name).to(device).eval()
                    self.processor = HFCLIPImageProcessor.from_pretrained(repo_name)
                    self.kind = "hf_clip"
                    self.preprocess = None
                else:
                    raise ValueError(f"Unknown model type: {repo_or_kind}")
            except Exception as e:
                # Последняя попытка - попробовать как HF модель
                try:
                    self.model = HFCLIPModel.from_pretrained(repo_or_kind).to(device).eval()
                    self.processor = HFCLIPImageProcessor.from_pretrained(repo_or_kind)
                    self.kind = "hf_clip"
                    self.preprocess = None
                except:
                    raise ValueError(f"Failed to load model {repo_or_kind}: {e}")
                
        # Определяем размер эмбеддинга
        if self.kind == "open_clip":
            with torch.no_grad():
                img = Image.new('RGB', (224, 224), color=0)
                x = self.preprocess(img).unsqueeze(0).to(device)
                feat = self.model.encode_image(x)
            self.embed_dim = int(feat.shape[-1])
        else:
            self.embed_dim = int(self.model.config.projection_dim)
            
    @torch.inference_mode()
    def encode_images(self, paths: List[str]) -> torch.Tensor:
        ims = []
        if self.kind == "open_clip":
            for p in paths:
                try:
                    im = Image.open(p).convert("RGB")
                except:
                    im = Image.new("RGB", (224, 224), color=0)
                ims.append(self.preprocess(im))
            x = torch.stack(ims).to(self.device)
            f = self.model.encode_image(x)
        else:
            # HF CLIP (PubMedCLIP)
            for p in paths:
                try:
                    im = Image.open(p).convert("RGB")
                except:
                    im = Image.new("RGB", (224, 224), color=0)
                ims.append(im)
            proc = self.processor(images=ims, return_tensors='pt')
            x = proc['pixel_values'].to(self.device)
            f = self.model.get_image_features(pixel_values=x)
        return F.normalize(f, dim=-1)

class Captioner(nn.Module):
    def __init__(self, gpt2_name: str, clip_repo: str, prefix_tokens: int, prompt: str, device: str):
        super().__init__()
        self.device = device
        self.prompt = prompt
        self.tok = AutoTokenizer.from_pretrained(gpt2_name)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
        self.gpt2 = AutoModelForCausalLM.from_pretrained(gpt2_name).to(device).eval()
        self.clip = CLIPBackend(clip_repo, device)
        self.prefix = PrefixProjector(self.clip.embed_dim, int(self.gpt2.config.n_embd), prefix_tokens).to(device).eval()
        
    @torch.inference_mode()
    def generate(self, img_paths: List[str], prompt: Optional[str] = None) -> List[str]:
        pr = prompt or self.prompt or ""
        f = self.clip.encode_images(img_paths)
        pref = self.prefix(f)
        ids = self.tok([pr]*pref.size(0), return_tensors='pt', padding=True, truncation=True).to(self.device)
        emb_prompt = self.gpt2.transformer.wte(ids['input_ids'])
        inputs_embeds = torch.cat([pref, emb_prompt], dim=1)
        attn = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long, device=self.device)
        gen = self.gpt2.generate(
            inputs_embeds=inputs_embeds, attention_mask=attn,
            max_new_tokens=60, min_new_tokens=24, num_beams=4,
            no_repeat_ngram_size=4, repetition_penalty=1.15, length_penalty=0.6,
            pad_token_id=self.tok.eos_token_id, eos_token_id=self.tok.eos_token_id, early_stopping=True
        )
        outs = self.tok.batch_decode(gen, skip_special_tokens=True)
        res = []
        for s in outs:
            cut = s.find(pr)
            if cut >= 0: s = s[cut+len(pr):]
            res.append(s.strip())
        return res

def load_model(repo_dir: str | os.PathLike) -> Captioner:
    repo_dir = Path(repo_dir)
    cfgs = sorted(repo_dir.glob("final_captioner_*.json"))
    if not cfgs:
        raise FileNotFoundError("final_captioner_*.json not found in repo snapshot")
    data = json.loads(cfgs[-1].read_text(encoding='utf-8'))
    gpt2 = data.get("gpt2_name", "gpt2-medium")
    
    # Определяем CLIP репозиторий с поддержкой TimmModel
    clip_repo = data.get("clip_weight_path", data.get("clip_repo", data.get("clip_backend_kind", "")))
    
    # Если информация о CLIP не найдена в JSON, пытаемся определить по имени файла
    if not clip_repo or clip_repo in ["open_clip", "hf_clip"]:
        ckpts = sorted(repo_dir.glob("final_captioner_*.pt"))
        if ckpts:
            ckpt_name = str(ckpts[-1])
            if "TimmModel" in ckpt_name:
                clip_repo = "microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"
            elif "VisionTransformer" in ckpt_name:
                clip_repo = "redlessone/DermLIP_ViT-B-16"
            elif "CLIPModel" in ckpt_name:
                clip_repo = "flaviagiammarino/pubmed-clip-vit-base-patch32"
            elif "biomedclip" in ckpt_name.lower():
                clip_repo = "microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"
    
    prefix_tokens = int(data.get("prefix_tokens", 32))
    prompt = data.get("prompt", "Describe the skin lesion.")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Captioner(gpt2, clip_repo, prefix_tokens, prompt, device).to(device).eval()
    # подгрузим state_dict
    ckpts = sorted(repo_dir.glob("final_captioner_*.pt"))
    if not ckpts:
        raise FileNotFoundError("final_captioner_*.pt not found in repo snapshot")
    state = torch.load(ckpts[-1], map_location="cpu")
    sd = state.get("model", state)
    model.load_state_dict(sd, strict=False)
    return model

def generate(model: Captioner, img_paths: List[str], prompt: Optional[str] = None) -> List[str]:
    return model.generate(img_paths, prompt=prompt)