add b64 to baseline

app.py

@@ -10,9 +10,6 @@ from baseline import run as run_baseline
 import torch
 from shared import DEFAULT_SO_NEGATIVE_PROMPT, DEFAULT_OVERALL_NEGATIVE_PROMPT
 from examples import stage1_examples, stage2_examples
-import pickle
-import codecs
-import subprocess
 import base64
 import io
 
@@ -126,11 +123,7 @@ def get_ours_image(response, overall_prompt_override="", seed=0, num_inference_s
         gligen_scheduled_sampling_beta=gligen_scheduled_sampling_beta, num_inference_steps=num_inference_steps, scheduler_key=scheduler_key,
         so_negative_prompt=so_negative_prompt, overall_negative_prompt=overall_negative_prompt, so_batch_size=2
     )
-    print(type(image_np))
-    pic_IObytes = io.BytesIO()
-    plt.savefig(pic_IObytes,  format='png')
-    pic_IObytes.seek(0)
-    canvasbase64 = base64.b64encode(pic_IObytes.read()).decode()
+    
     images = [image_np, b64]
     # if show_so_imgs:
     #     images.extend([np.asarray(so_img) for so_img in so_img_list])
@@ -143,8 +136,9 @@ def get_baseline_image(prompt, seed=0):
     scheduler_key = "dpm_scheduler"
     num_inference_steps = 20
     
-    image_np = run_baseline(prompt, bg_seed=seed, scheduler_key=scheduler_key, num_inference_steps=num_inference_steps)
-    return [image_np]
+    image_np, b64 = run_baseline(prompt, bg_seed=seed, scheduler_key=scheduler_key, num_inference_steps=num_inference_steps)
+    images = [image_np, b64]
+    return images
 
 def parse_input(text=None):
     try:
@@ -298,10 +292,11 @@ with gr.Blocks(
             # with gr.Column(scale=1):
             #     output = gr.Image(shape=(512, 512), elem_classes="img", elem_id="img")
             with gr.Column(scale=1):
-                gallery = gr.Gallery(
-                    label="Generated image", show_label=False, elem_id="gallery2", columns=[1], rows=[1], object_fit="contain", preview=True
+                gallery = gr.Image(
+                    label="Generated image", show_label=False, elem_id="gallery", columns=[1], rows=[1], object_fit="contain"                    
                 )
-        generate_btn.click(fn=get_baseline_image, inputs=[sd_prompt, seed], outputs=gallery, api_name="baseline")
+                b64 = gr.Textbox(label="base64", placeholder="base64", lines = 2)
+        generate_btn.click(fn=get_baseline_image, inputs=[sd_prompt, seed], outputs=[gallery,b64], api_name="baseline")
 
         gr.Examples(
             examples=stage1_examples,

baseline.py

@@ -5,6 +5,8 @@ import models
 from models import pipelines
 from shared import model_dict, DEFAULT_OVERALL_NEGATIVE_PROMPT
 import gc
+from io import BytesIO
+import base64
 
 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
 
@@ -41,4 +43,8 @@ def run(prompt, scheduler_key='dpm_scheduler', bg_seed=1, num_inference_steps=20
     gc.collect()
     torch.cuda.empty_cache()
 
-    return images[0]
+     with BytesIO() as buffer:
+        np.save(buffer, images[0])
+        img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+    return images[0], img_str