File size: 2,431 Bytes
df8eb47 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import torch
import torch.nn.functional as F
from torchvision import transforms
import open_clip
# Set device
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, _, clip_preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
clip_model = clip_model.to(torch_device)
clip_model.eval() # model in train mode by default, impacts some models with BatchNorm or stochastic depth active
clip_tokenizer = open_clip.get_tokenizer('ViT-B-32')
def get_text_embedding(text):
text_tokens = clip_tokenizer([text]).to(torch_device)
with torch.no_grad(), torch.cuda.amp.autocast():
text_features = clip_model.encode_text(text_tokens).float()
text_features /= text_features.norm(dim=-1, keepdim=True)
return text_features
def get_image_embedding(image):
image_input = clip_preprocess(image).unsqueeze(0).to(torch_device)
with torch.no_grad(), torch.cuda.amp.autocast():
image_features = clip_model.encode_image(image_input).float()
image_features /= image_features.norm(dim=-1, keepdim=True)
return image_features
def text_image_similarity_loss(generated_images, target_text = "plain background"):
# Get text embedding
text_embedding = get_text_embedding(target_text)
# Ensure the generated_images have requires_grad=True
# generated_images.requires_grad_(True)
# Convert image tensor to the required format (normalization, resizing)
# Normalize the images (assuming they are in [0, 1])
transform = transforms.Compose([
transforms.Resize((224, 224)), # Example size, modify as needed
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# Apply the transformation
transformed_images = transform(generated_images)
# Assuming `image_encoder` is a pretrained model that returns image embeddings
# Get image embeddings
# image_embeddings = image_encoder(generated_images)
with torch.cuda.amp.autocast():
image_features = clip_model.encode_image(transformed_images).float()
norm_image_features = image_features / image_features.norm(dim=-1, keepdim=True)
# Calculate cosine similarity
cos_sim = F.cosine_similarity(norm_image_features, text_embedding, dim=-1)
# Define the loss as 1 - cosine similarity (assuming we want to maximize similarity)
loss = 1 - cos_sim.mean()
return loss
|