Update app.py

app.py

@@ -1,47 +1,25 @@
 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 torch.utils.data import DataLoader, Dataset
 from PIL import Image
-from datasets import load_dataset
-from transformers import T5EncoderModel, T5Tokenizer, DistilBertModel, DistilBertTokenizer
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 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
-
+device ="cpu"
 class TextEncoder(nn.Module):
-    def __init__(self, encoder_model_name, encoder_type="t5", device='cpu'):
+    def __init__(self, encoder_model_name):
         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)
+        self.tokenizer = T5Tokenizer.from_pretrained(encoder_model_name)
+        self.encoder = T5ForConditionalGeneration.from_pretrained(encoder_model_name)
+        self.encoder.to(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
+        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        inputs = {key: value.to(device) for key, value in inputs.items()}
+        outputs = self.encoder.encoder(**inputs)
+        embeddings = outputs.last_hidden_state[:, 0, :]
+        return embeddings
 
 class ConditionalDiffusionModel(nn.Module):
     def __init__(self):
@@ -81,233 +59,110 @@ class TextToImageModel(nn.Module):
     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 = torch.rand((1, 3, 64, 64))  # Initialize input image with random values
+        for i 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):
+    def __init__(self, annotations_file, img_dir, transform=None):
         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
+        self.img_dir = img_dir
+        self.transform = transform
 
     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)
+        img_path = os.path.join(self.img_dir, img_name)
         image = Image.open(img_path).convert("RGB")
         if self.transform:
             image = self.transform(image)
         return text, image
 
+def save_checkpoint(model, optimizer, epoch, checkpoint_path):
+    checkpoint = {
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'epoch': epoch
+    }
+    torch.save(checkpoint, checkpoint_path)
+
+def load_checkpoint(model, optimizer, 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'])
+        epoch = checkpoint['epoch']
+        return epoch
+    else:
+        return 0
 
-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 test_inference(model, text):
+    model.eval()
+    with torch.no_grad():
+        generated_image = model(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()
+def visualize_image(image_tensor):
+    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.show()
 
 if __name__ == "__main__":
-    batch_size = 1
+    # Define hyperparameters and paths
+    batch_size = 4
     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)
+    num_epochs = 1000
+    checkpoint_path = 'checkpoint.pth'
+    annotations_file = 'annotations.txt'
+    img_dir = 'images/'
+    
+    # Initialize models
+    text_encoder = TextEncoder("google-t5/t5-small")
     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)
+    # Define optimizer and criterion
+    optimizer = torch.optim.Adam(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)
+    # Load checkpoint if available
+    start_epoch = load_checkpoint(text_to_image_model, optimizer, checkpoint_path)
 
-    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)
+    # Define transformations for the images
+    transform = transforms.Compose([
+        transforms.Resize((64, 64)),
+        transforms.ToTensor(),
+    ])
 
+    # Initialize dataset and dataloader
+    dataset = CustomDataset(annotations_file, img_dir, transform)
     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)
+    # Training loop
+    text_to_image_model.train()
+    for epoch in range(start_epoch, num_epochs):
+        for i, (text_batch, image_batch) in enumerate(dataloader):
+            optimizer.zero_grad()
+            images = text_to_image_model(text_batch)
+            target_images = image_batch.to(device)
+            loss = criterion(images, target_images)
+            loss.backward()
+            optimizer.step()
+        
+        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')
+        save_checkpoint(text_to_image_model, optimizer, epoch+1, checkpoint_path)
+
+    print("Training completed.")
+    
+    # Test inference
+    sample_text = "A big ape."
+    generated_image = test_inference(text_to_image_model, sample_text)
+    visualize_image(generated_image)