Update pipeline.py

pipeline.py

@@ -1,5 +1,6 @@
 import torch
 import nltk
+import io
 
 from torchvision import transforms
 from pytorch_pretrained_biggan import BigGAN, one_hot_from_names, truncated_noise_sample
@@ -34,4 +35,8 @@ class PreTrainedPipeline():
         with torch.no_grad():
             output = self.model(noise_vector, class_vector, self.truncation)
 
-        return transforms.ToPILImage()(output[0]).tobytes().decode('utf-8')
+        img = transforms.ToPILImage()(output[0])
+        buf = io.BytesIO()
+        img.save(buf, format="JPEG")
+
+        return buf.getvalue()return transforms.ToPILImage()(output[0]).tobytes().decode('utf-8')