text_plain_output_1.png
from pathlib import Path

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')