encode as jpeg with simplejpeg

Dockerfile

@@ -20,6 +20,7 @@ RUN git clone https://github.com/NovelAI/stable-diffusion
 RUN pip3 install -e stable-diffusion/.
 RUN pip3 install pytorch_lightning
 RUN pip3 install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+RUN pip3 install simplejpeg
 
 #Open ports
 EXPOSE 8080

hydra_node/models.py

@@ -122,9 +122,6 @@ class StableDiffusionModel(nn.Module):
         images = []
         for x_sample in x_samples_ddim:
             x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-            x_sample = x_sample.tobytes()
-            #get base64 of x_sample
-            x_sample = str(base64.b64encode(x_sample))
             images.append(x_sample) 
 
         return images
\ No newline at end of file

main.py

@@ -14,6 +14,8 @@ import time
 import gc
 import os
 import signal
+import simplejpeg
+import base64
 
 #Initialize model and config
 model, config = init_config_model()
@@ -72,6 +74,15 @@ def generate(request: GenerationRequest):
             return ErrorOutput(error=output[1])
         
         images = model.sample(request)
+        for x in range(images):
+            image = simplejpeg.encode_jpeg(images[x], quality=95) 
+            image = image.tobytes()
+            #get base64 of image
+            image = str(base64.b64encode(image))
+            #remove b' from base64
+            image = image[2:-1]
+            image[x] = image
+            
         process_time = time.perf_counter() - t
         logger.info(f"Request took {process_time:0.3f} seconds")
         return GenerationOutput(generation=images)