app.py

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 wandb 
import torchvision.utils as vutils # To save image grids
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)  # Combine embeddings into a batch

class ConditionalDiffusionModel(nn.Module):
    def __init__(self):
        super(ConditionalDiffusionModel, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(512, 768),  # Adjusted from 512 to 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),  # 3 is the number of color channels
            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):  # Upsample the image 6 times
            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):  # New class for HuggingFace 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(),
        ])
        # Apply limit if specified
        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)
        # Optionally push dataset to Hugging Face Hub

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:
            # Initialize wandb
            #self.wandb = wandb.login(key=WANDB_API_KEY) # Assuming you have already logged in.  If not, use: wandb.login(key='YOUR_WANDB_API_KEY')
            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 # Increase image size 
                #  Generate a starting image with the correct shape
                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:
                    # Log loss and learning rate
                    self.wandb.log({"train_loss": loss.item(), "lr": optimizer.param_groups[0]['lr']})

                if i % 25 == 0:
                    # Save a grid of real and generated images for monitoring
                    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()

            # Validation step
            val_loss = self.evaluate(model, dataloader, criterion)
            scheduler.step(val_loss)  # Update scheduler with validation 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())  # Normalize to [0, 1]
        image_tensor = image_tensor.permute(1, 2, 0)  # Change from (C, H, W) to (H, W, C)
        plt.imshow(image_tensor)
        plt.savefig(filename)
        #plt.show()

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  # <-- Toggle between datasets
    limit_huggingface_dataset = 1000  # <-- Set the limit for HuggingFace dataset
    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)