Create pipeline.py

pipeline.py

@@ -0,0 +1,21 @@
+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