text_plain_output_1.png
from pathlib import Path

from sentence_transformers import SentenceTransformer

from torch.utils.data import Dataset, DataLoader

from torchvision import transforms

from tqdm.notebook import tqdm

import numpy as np

import pandas as pd

import timm

import torch

class CFG:
    model_path = '/kaggle/input/stable-diffusion-vit-baseline-train/vit_base_patch16_224.pth'
    model_name = 'vit_base_patch16_224'
    input_size = 224
    batch_size = 64

def predict(images, model_path, model_name, input_size, batch_size):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    transform = transforms.Compose([transforms.Resize(input_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    dataset = DiffusionTestDataset(images, transform)
    dataloader = DataLoader(dataset=dataset, shuffle=False, batch_size=batch_size, pin_memory=True, num_workers=2, drop_last=False)
    model = timm.create_model(model_name, pretrained=False, num_classes=384)
    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    preds = []
    for X in tqdm(dataloader, leave=False):
        X = X.to(device)
        with torch.no_grad():
            X_out = model(X)
            preds.append(X_out.cpu().numpy())
    return (np.vstack(preds).flatten(), np.array(preds)[0])

images = list(Path('/kaggle/input/stable-diffusion-image-to-prompts/images').glob('*.png'))
imgIds = [i.stem for i in images]
EMBEDDING_LENGTH = 384
imgId_eId = ['_'.join(map(str, i)) for i in zip(np.repeat(imgIds, EMBEDDING_LENGTH), np.tile(range(EMBEDDING_LENGTH), len(imgIds)))]
prompt_embeddings, raw_preds = predict(images, CFG.model_path, CFG.model_name, CFG.input_size, CFG.batch_size)
submission = pd.DataFrame(index=imgId_eId, data=prompt_embeddings, columns=['val']).rename_axis('imgId_eId')

prompts = pd.read_csv('/kaggle/input/stable-diffusion-image-to-prompts/prompts.csv')
prompts = prompts.loc[prompts['imgId'].isin(imgIds)].set_index('imgId').loc[imgIds].reset_index()

st_model = SentenceTransformer('/kaggle/input/sentence-transformers-222/all-MiniLM-L6-v2')

prompts_true = st_model.encode(prompts['prompt'])