|
import os |
|
import torch |
|
import torch.nn as nn |
|
import numpy as np |
|
import torch.optim as optim |
|
import torchvision.transforms as transforms |
|
from torch.utils.data import DataLoader, Dataset, Subset |
|
from PIL import Image |
|
from datasets import load_dataset |
|
from transformers import T5EncoderModel, T5Tokenizer, DistilBertModel, DistilBertTokenizer |
|
import matplotlib.pyplot as plt |
|
from rich import print as rp |
|
from credits import HUGGINGFACE_TOKEN, HUGGINGFACE_W_TOKEN, WANDB_API_KEY |
|
import wandb |
|
import torchvision.utils as vutils |
|
write_token = HUGGINGFACE_W_TOKEN |
|
read_token = HUGGINGFACE_TOKEN |
|
|
|
class TextEncoder(nn.Module): |
|
def __init__(self, encoder_model_name, encoder_type="t5", device='cpu'): |
|
super(TextEncoder, self).__init__() |
|
self.device = device |
|
self.encoder_type = encoder_type |
|
if encoder_type == "t5": |
|
self.tokenizer = T5Tokenizer.from_pretrained(encoder_model_name, cache_dir='./models', token=read_token) |
|
self.encoder = T5EncoderModel.from_pretrained(encoder_model_name, cache_dir='./models', token=read_token) |
|
elif encoder_type == "distilbert": |
|
self.tokenizer = DistilBertTokenizer.from_pretrained(encoder_model_name, cache_dir='./models', token=read_token) |
|
self.encoder = DistilBertModel.from_pretrained(encoder_model_name, cache_dir='./models', token=read_token) |
|
else: |
|
raise ValueError(f"Invalid encoder_type: {encoder_type}. Choose from 't5' or 'distilbert'.") |
|
self.encoder.to(self.device) |
|
|
|
def encode_text(self, text): |
|
if isinstance(text, str): |
|
text = [text] |
|
|
|
embeddings = [] |
|
for t in text: |
|
inputs = self.tokenizer(t, return_tensors="pt", padding=True, truncation=True).to(self.device) |
|
outputs = self.encoder(**inputs) |
|
embeddings.append(outputs.last_hidden_state[:, 0, :]) |
|
|
|
return torch.stack(embeddings) |
|
|
|
class ConditionalDiffusionModel(nn.Module): |
|
def __init__(self): |
|
super(ConditionalDiffusionModel, self).__init__() |
|
self.model = nn.Sequential( |
|
nn.Linear(512, 768), |
|
nn.ReLU(), |
|
nn.Linear(768, 64), |
|
nn.ReLU(), |
|
nn.Linear(64, 64) |
|
) |
|
|
|
def forward(self, text_embeddings): |
|
return self.model(text_embeddings) |
|
|
|
class SuperResolutionDiffusionModel(nn.Module): |
|
def __init__(self): |
|
super(SuperResolutionDiffusionModel, self).__init__() |
|
self.model = nn.Sequential( |
|
nn.Conv2d(3, 64, kernel_size=3, padding=1), |
|
nn.ReLU(), |
|
nn.Conv2d(64, 64, kernel_size=3, padding=1), |
|
nn.ReLU(), |
|
nn.Conv2d(64, 3, kernel_size=3, padding=1) |
|
) |
|
|
|
def forward(self, input_image): |
|
return self.model(input_image) |
|
|
|
class TextToImageModel(nn.Module): |
|
def __init__(self, text_encoder, conditional_diffusion_model, super_resolution_diffusion_model): |
|
super(TextToImageModel, self).__init__() |
|
self.text_encoder = text_encoder |
|
self.conditional_diffusion_model = conditional_diffusion_model |
|
self.super_resolution_diffusion_model = super_resolution_diffusion_model |
|
|
|
def forward(self, text): |
|
text_embeddings = self.text_encoder.encode_text(text) |
|
image_embeddings = self.conditional_diffusion_model(text_embeddings) |
|
input_image = torch.rand((image_embeddings.shape[0], 3, 128, 128)).to(text_embeddings.device) |
|
for _ in range(6): |
|
input_image = self.super_resolution_diffusion_model(input_image) |
|
return input_image |
|
|
|
class CustomDataset(Dataset): |
|
def __init__(self, annotations_file, train_img_dir, size_sqr=128): |
|
with open(annotations_file, 'r') as f: |
|
lines = f.readlines() |
|
self.transform = transforms.Compose([ |
|
transforms.Resize((size_sqr, size_sqr)), |
|
transforms.ToTensor(), |
|
]) |
|
self.img_labels = [line.strip().split(' ', 1) for line in lines] |
|
self.train_img_dir = train_img_dir |
|
|
|
def __len__(self): |
|
return len(self.img_labels) |
|
|
|
def __getitem__(self, idx): |
|
img_name, text = self.img_labels[idx] |
|
img_path = os.path.join(self.train_img_dir, img_name) |
|
image = Image.open(img_path).convert("RGB") |
|
if self.transform: |
|
image = self.transform(image) |
|
return text, image |
|
|
|
|
|
class HuggingDataset(Dataset): |
|
def __init__(self, dataset_name="vera365/lexica_dataset", size_sqr=128, limit=None): |
|
self.dataset = load_dataset(dataset_name, token=read_token, cache_dir='./datasets') |
|
self.transform = transforms.Compose([ |
|
transforms.Resize((size_sqr, size_sqr)), |
|
transforms.ToTensor(), |
|
]) |
|
|
|
if limit is not None: |
|
self.dataset = Subset(self.dataset["train"], range(limit)) |
|
|
|
def __len__(self): |
|
return len(self.dataset["train"]) |
|
|
|
def __getitem__(self, idx): |
|
item = self.dataset["train"][idx] |
|
|
|
image =item["image"].convert("RGB") |
|
text = item["subject"] |
|
if self.transform: |
|
image = self.transform(image) |
|
return text, image |
|
|
|
|
|
|
|
class StorageHandler: |
|
def __init__(self, storage_dir="./image_gen_storage", hub_model_name="K00B404/tiny_image_gen", push_dataset=False, dataset_name="K00B404/custom_image_descriptions_dataset"): |
|
self.model_name = hub_model_name |
|
self.dataset_name = dataset_name |
|
self.push_dataset = push_dataset |
|
self.storage_dir = storage_dir |
|
|
|
def save_checkpoint(self, model, optimizer, scheduler, epoch, checkpoint_path): |
|
checkpoint = { |
|
'model_state_dict': model.state_dict(), |
|
'optimizer_state_dict': optimizer.state_dict(), |
|
'scheduler_state_dict': scheduler.state_dict(), |
|
'epoch': epoch |
|
} |
|
torch.save(checkpoint, checkpoint_path) |
|
|
|
def load_checkpoint(self, model, optimizer, scheduler, checkpoint_path): |
|
if os.path.isfile(checkpoint_path): |
|
checkpoint = torch.load(checkpoint_path) |
|
model.load_state_dict(checkpoint['model_state_dict']) |
|
optimizer.load_state_dict(checkpoint['optimizer_state_dict']) |
|
scheduler.load_state_dict(checkpoint['scheduler_state_dict']) |
|
epoch = checkpoint['epoch'] |
|
return epoch, scheduler |
|
else: |
|
return 0, scheduler |
|
|
|
def push_dataset(self, dataset): |
|
if dataset: |
|
dataset.push_to_hub(self.dataset_name, token=write_token) |
|
|
|
def push(self, model, tokenizer): |
|
model.push_to_hub(self.model_name, token=write_token) |
|
tokenizer.push_to_hub(self.model_name, token=write_token) |
|
|
|
|
|
class Common: |
|
def __init__(self, device='cpu', wandb_log=False): |
|
self.wandb_log = wandb_log |
|
self.device = device |
|
self.terminal_log = rp |
|
|
|
if self.wandb_log: |
|
|
|
|
|
self.wandb = wandb.init(project="my-image-generation-project", |
|
config={ |
|
"learning_rate": learning_rate, |
|
"batch_size": batch_size, |
|
"num_epochs": num_epochs, |
|
"encoder_model": encoder |
|
}) |
|
|
|
|
|
def train(self, model, optimizer, scheduler, dataloader, criterion, num_epochs, start_epoch, checkpoint_path): |
|
for epoch in range(start_epoch, num_epochs): |
|
model.train() |
|
for i, (text_batch, image_batch) in enumerate(dataloader): |
|
image_batch = image_batch.to(self.device) |
|
image_size = 128 |
|
|
|
input_image = torch.randn((image_batch.shape[0], 3, image_size//8, image_size//8)).to(device) |
|
|
|
optimizer.zero_grad() |
|
images = model(text_batch) |
|
loss = criterion(images, image_batch) |
|
loss.backward() |
|
optimizer.step() |
|
|
|
if self.wandb_log: |
|
|
|
self.wandb.log({"train_loss": loss.item(), "lr": optimizer.param_groups[0]['lr']}) |
|
|
|
if i % 25 == 0: |
|
|
|
img_grid_real = vutils.make_grid(image_batch[:4], padding=2, normalize=True) |
|
img_grid_fake = vutils.make_grid(input_image[:4], padding=2, normalize=True) |
|
|
|
plt.figure(figsize=(15,15)) |
|
plt.subplot(1,2,1) |
|
plt.axis("off") |
|
plt.title("Real Images") |
|
plt.imshow(np.transpose(img_grid_real.cpu(),(1,2,0))) |
|
|
|
plt.subplot(1,2,2) |
|
plt.axis("off") |
|
plt.title("Generated Images") |
|
plt.imshow(np.transpose(img_grid_fake.cpu(),(1,2,0))) |
|
plt.savefig(f'generated_images_epoch_{epoch+1}_batch_{i}.png') |
|
plt.close() |
|
|
|
|
|
val_loss = self.evaluate(model, dataloader, criterion) |
|
scheduler.step(val_loss) |
|
|
|
image = self.test_inference(model, "A house next to a river.") |
|
self.visualize_image(image, f'generated_image_epoch_{epoch + 1}.png') |
|
self.terminal_log(f'Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item()}, Validation Loss: {val_loss}') |
|
StorageHandler().save_checkpoint(model, optimizer, scheduler, epoch + 1, checkpoint_path) |
|
|
|
self.terminal_log("Training completed.") |
|
|
|
def evaluate(self, model, dataloader, criterion): |
|
model.eval() |
|
total_loss = 0 |
|
with torch.no_grad(): |
|
for i, (text_batch, image_batch) in enumerate(dataloader): |
|
image_batch = image_batch.to(self.device) |
|
images = model(text_batch) |
|
loss = criterion(images, image_batch) |
|
total_loss += loss.item() |
|
|
|
avg_loss = total_loss / len(dataloader) |
|
self.terminal_log(f'Validation Loss: {avg_loss}') |
|
return avg_loss |
|
|
|
def test_inference(self, model, text): |
|
model.eval() |
|
with torch.no_grad(): |
|
if isinstance(text, str): |
|
generated_image = model(text) |
|
else: |
|
generated_image = [model(t) for t in text] |
|
return generated_image |
|
|
|
def visualize_image(self, image_tensor, filename='generated_image.png'): |
|
image_tensor = image_tensor.squeeze(0).cpu().detach() |
|
image_tensor = (image_tensor - image_tensor.min()) / (image_tensor.max() - image_tensor.min()) |
|
image_tensor = image_tensor.permute(1, 2, 0) |
|
plt.imshow(image_tensor) |
|
plt.savefig(filename) |
|
|
|
|
|
if __name__ == "__main__": |
|
batch_size = 1 |
|
learning_rate = 1e-4 |
|
num_epochs = 500 |
|
encoder = "google-t5/t5-small" |
|
checkpoint_path = './models/image_gen' |
|
os.makedirs(checkpoint_path, exist_ok=True) |
|
checkpoint_file = f"{checkpoint_path}/checkpoint_backup.pth" |
|
use_huggingface_dataset = False |
|
limit_huggingface_dataset = 1000 |
|
train_img_dir = './train_images' |
|
annotations_file = f'{train_img_dir}/annotations.txt' |
|
storage_dir = "./image_gen_storage" |
|
os.makedirs(storage_dir, exist_ok=True) |
|
hub_model_name = "K00B404/tiny_image_gen" |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
text_encoder = TextEncoder(encoder, encoder_type="t5", device=device) |
|
conditional_diffusion_model = ConditionalDiffusionModel() |
|
super_resolution_diffusion_model = SuperResolutionDiffusionModel() |
|
text_to_image_model = TextToImageModel(text_encoder, conditional_diffusion_model, super_resolution_diffusion_model) |
|
text_to_image_model.to(device) |
|
|
|
optimizer = optim.AdamW(text_to_image_model.parameters(), lr=learning_rate) |
|
criterion = nn.MSELoss() |
|
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.001, patience=2) |
|
|
|
common = Common(device, True) |
|
storage = StorageHandler(checkpoint_path, hub_model_name) |
|
|
|
start_epoch, scheduler = storage.load_checkpoint(text_to_image_model, optimizer, scheduler, checkpoint_file) |
|
|
|
if use_huggingface_dataset: |
|
dataset = HuggingDataset(size_sqr=128, limit=limit_huggingface_dataset) |
|
else: |
|
dataset = CustomDataset(annotations_file, train_img_dir, size_sqr=128) |
|
|
|
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) |
|
|
|
common.train(text_to_image_model, optimizer, scheduler, dataloader, criterion, num_epochs, start_epoch, checkpoint_file) |
|
|
|
sample_texts = ["A big ape.", "A yellow banana."] |
|
for sample_text in sample_texts: |
|
generated_image = common.test_inference(text_to_image_model, sample_text) |
|
common.visualize_image(generated_image) |