pipeline.py

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

class PretrainedPipeline():
    def __init__(self):
        self.device = torch.device("cpu")
        self.generator = Generator()  # Instantiate your GAN generator class
        self.generator.load_state_dict(torch.load("pytorch_model.bin", map_location=self.device))
        self.generator.eval()

    def generate_image(self):
        with torch.no_grad():
            noise = torch.randn(1, 128, 1, 1).to(self.device)  # Assuming input noise size is 100
            generated_image_tensor = self.generator(noise)
            generated_image = generated_image_tensor.cpu().detach().squeeze(0)
            # Assuming the generator output is in the range [-1, 1]
            generated_image = (generated_image + 1) / 2.0
            pil_image = transforms.ToPILImage()(generated_image)
            return pil_image