How to Run Pre-trained and Optimized Models

1. Load Pre-trained Model

To run the pre-trained model, first load the model checkpoint and set the model to evaluation mode.

import torch
from transformers import BertTokenizer
from PIL import Image
import numpy as np

# Load the trained model
model_path = '/content/drive/My Drive/Transformer_Project/models/text_to_image_transformer_original.pth'
model.load_state_dict(torch.load(model_path))
model.eval()

# Set up the tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

2. Generate Image from Text

Provide the model with a text prompt to generate an image.

# Input text
input_text = "A beautiful sunset over the mountains."

# Tokenize the input text
tokens = tokenizer(input_text, padding="max_length", truncation=True, return_tensors="pt")

# Generate image
with torch.no_grad():
    generated_image = model(tokens)

# Convert image tensor to numpy array and save it
image_array = (np.clip(generated_image[0].permute(1, 2, 0).numpy(), 0, 1) * 255).astype(np.uint8)
image = Image.fromarray(image_array)
image.save('/content/drive/My Drive/Transformer_Project/results/generated_image_pretrained.png')

3. Load Optimized Model

Follow similar steps to load the optimized model and run inference.

# Load the optimized model
optimized_model_path = '/content/drive/My Drive/Transformer_Project/models/text_to_image_transformer_optimized.pth'
model.load_state_dict(torch.load(optimized_model_path))
model.eval()

4. Generate Image with Optimized Model

Use the optimized model to generate images and compare with pre-trained model outputs.

# Input text (same as used for the pre-trained model)
input_text = "A beautiful sunset over the mountains."

# Tokenize the input text
tokens = tokenizer(input_text, padding="max_length", truncation=True, return_tensors="pt")

# Generate image with optimized model
with torch.no_grad():
    generated_image_optimized = model(tokens)

# Convert image tensor to numpy array and save it
image_array_optimized = (np.clip(generated_image_optimized[0].permute(1, 2, 0).numpy(), 0, 1) * 255).astype(np.uint8)
image_optimized = Image.fromarray(image_array_optimized)
image_optimized.save('/content/drive/My Drive/Transformer_Project/results/generated_image_optimized.png')

5. Compare Generated Images

To visually compare the images generated by both models, use the following code:

import matplotlib.pyplot as plt

# Load the images
img_pretrained = Image.open('/content/drive/My Drive/Transformer_Project/results/generated_image_pretrained.png')
img_optimized = Image.open('/content/drive/My Drive/Transformer_Project/results/generated_image_optimized.png')

# Visualization
fig, ax = plt.subplots(1, 2, figsize=(12, 6))
ax[0].imshow(img_pretrained)
ax[0].set_title("Pre-trained Model")
ax[0].axis('off')

ax[1].imshow(img_optimized)
ax[1].set_title("Optimized Model")
ax[1].axis('off')

plt.show()