update generation script, add pillow

baseline.py

@@ -1,13 +1,13 @@
 # Original Stable Diffusion (1.4)
 
 import torch
-import numpy as np
 import models
 from models import pipelines
 from shared import model_dict, DEFAULT_OVERALL_NEGATIVE_PROMPT
 import gc
 from io import BytesIO
 import base64
+import PIL.Image
 
 vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype
 
@@ -42,7 +42,6 @@ def run(prompt, scheduler_key='dpm_scheduler', bg_seed=1, num_inference_steps=20
     )
 
     # Convert to PIL Image
-    import PIL.Image
     image = PIL.Image.fromarray(images[0])
     
     # Save as PNG in memory

generation.py

@@ -10,8 +10,7 @@ from shared import model_dict, sam_model_dict, DEFAULT_SO_NEGATIVE_PROMPT, DEFAU
 import gc
 from io import BytesIO
 import base64
-import pickle
-import codecs
+import PIL.Image
 
 verbose = False
 
@@ -214,12 +213,16 @@ def run(
     gc.collect()
     torch.cuda.empty_cache()
 
-    obj = np.random.normal(size=(10, 10))
-    obj_base64string = codecs.encode(pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL), "base64").decode('latin1')
-    img_str = pickle.loads(codecs.decode(obj_base64string.encode('latin1'), "base64"))
-    # with BytesIO() as buffer:
-    #     np.save(buffer, images[0])
-    #     img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    # Convert to PIL Image
+    image = PIL.Image.fromarray(images[0])
+    
+    # Save as PNG in memory
+    buffer = BytesIO()
+    image.save(buffer, format='PNG')
+    
+    # Encode PNG to base64
+    png_bytes = buffer.getvalue()
+    base64_string = base64.b64encode(png_bytes).decode('utf-8')\
         
     return images[0], so_img_list, img_str
 

requirements.txt

@@ -10,3 +10,4 @@ inflect==6.0.4
 easydict
 accelerate==0.18.0
 gradio==3.35.2
+Pillow==10.0.0