text_plain_output_1.png
from PIL import Image

from pathlib import Path

from scipy import spatial

from sentence_transformers import SentenceTransformer

from torch import nn

from torch.optim import Adam

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

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

class DiffusionTestDataset(Dataset):
    def __init__(self, images, transform):
        self.images = images
        self.transform = transform
    def __len__(self):
        return len(self.images)
    def __getitem__(self, idx):
        image = Image.open(self.images[idx])
        image = self.transform(image)
        return image

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

def simularity(y_pred, y_true):
    cosine_sim = []
    for pred, true in zip(y_pred, y_true):
        cosine_sim.append(1 - spatial.distance.cosine(pred, true))

import torch
from torch import nn
from torch.optim import Adam
class SimModel(nn.Module):
    def __init__(self, hid_dim):
        super().__init__()
        self.fc1 = nn.Linear(384, hid_dim)
        self.fc2 = nn.Linear(hid_dim, 384)
        self.act = nn.ReLU()
    def forward(self, x):
        x = self.act(self.fc1(x))
        x = self.fc2(x)
        return x
input = torch.FloatTensor(raw_preds)
target = torch.FloatTensor(prompts_true)
model = SimModel(128)
criterion = nn.MSELoss()
optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.01)
for i in range(100):
    pred = model(input)
    optimizer.zero_grad()
    loss = criterion(pred, target - input)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    result = model(input) + input
    result = result.numpy()

simularity(result, prompts_true)