b64 support:

app.py

@@ -10,6 +10,11 @@ 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
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 if torch.cuda.is_available():
@@ -61,6 +66,9 @@ layout_placeholder = """Caption: A realistic photo of a gray cat and an orange d
 Objects: [('a gray cat', [67, 243, 120, 126]), ('an orange dog', [265, 193, 190, 210])]
 Background prompt: A realistic photo of a grassy area."""
 
+canvasbase64 = ""
+oursimagebase64 = ""
+
 def get_lmd_prompt(prompt, template=default_template):
     if prompt == "":
         prompt = prompt_placeholder
@@ -69,6 +77,7 @@ def get_lmd_prompt(prompt, template=default_template):
     return simplified_prompt.format(template=template, prompt=prompt)
 
 def get_layout_image(response):
+    global canvasbase64
     if response == "":
         response = layout_placeholder
     gen_boxes, bg_prompt = parse_input(response)
@@ -82,13 +91,19 @@ def get_layout_image(response):
     # Now we can save it to a numpy array.
     data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
     data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    pic_IObytes = io.BytesIO()
+    plt.savefig(pic_IObytes,  format='png')
+    pic_IObytes.seek(0)
+    canvasbase64 = base64.b64encode(pic_IObytes.read()).decode()
+  
     plt.clf()
-    return data
+    return [data,canvasbase64]
 
 def get_layout_image_gallery(response):
-    return [get_layout_image(response)]
+    return get_layout_image(response)
 
 def get_ours_image(response, overall_prompt_override="", seed=0, num_inference_steps=20, dpm_scheduler=True, use_autocast=False, fg_seed_start=20, fg_blending_ratio=0.1, frozen_step_ratio=0.4, gligen_scheduled_sampling_beta=0.3, so_negative_prompt=DEFAULT_SO_NEGATIVE_PROMPT, overall_negative_prompt=DEFAULT_OVERALL_NEGATIVE_PROMPT, show_so_imgs=False, scale_boxes=False):
+    global oursimagebase64
     if response == "":
         response = layout_placeholder
     gen_boxes, bg_prompt = parse_input(response)
@@ -105,15 +120,20 @@ def get_ours_image(response, overall_prompt_override="", seed=0, num_inference_s
     else:
         scheduler_key = "scheduler"
         
-    image_np, so_img_list = run_ours(
+    image_np, so_img_list, b64 = run_ours(
         spec, bg_seed=seed, overall_prompt_override=overall_prompt_override, fg_seed_start=fg_seed_start, 
         fg_blending_ratio=fg_blending_ratio,frozen_step_ratio=frozen_step_ratio, use_autocast=use_autocast,
         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
     )
-    images = [image_np]
-    if show_so_imgs:
-        images.extend([np.asarray(so_img) for so_img in so_img_list])
+    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])
     return images
 
 def get_baseline_image(prompt, seed=0):
@@ -230,7 +250,7 @@ with gr.Blocks(
             inputs=[prompt],
             outputs=[output],
             fn=get_lmd_prompt,
-            cache_examples=True
+            # cache_examples=True
         )
     
     with gr.Tab("Stage 2 (New). Layout to Image generation"):
@@ -254,18 +274,19 @@ with gr.Blocks(
                 visualize_btn = gr.Button("Visualize Layout", elem_classes="btn")
                 generate_btn = gr.Button("Generate Image from Layout", variant='primary', elem_classes="btn")
             with gr.Column(scale=1):
-                gallery = gr.Gallery(
-                    label="Generated image", show_label=False, elem_id="gallery", 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"                    
                 )
-        visualize_btn.click(fn=get_layout_image_gallery, inputs=response, outputs=gallery, api_name="visualize-layout")
-        generate_btn.click(fn=get_ours_image, inputs=[response, overall_prompt_override, seed, num_inference_steps, dpm_scheduler, use_autocast, fg_seed_start, fg_blending_ratio, frozen_step_ratio, gligen_scheduled_sampling_beta, so_negative_prompt, overall_negative_prompt, show_so_imgs, scale_boxes], outputs=gallery, api_name="layout-to-image")
+                b64 = gr.Textbox(label="base64", placeholder="base64", lines = 2)
+        visualize_btn.click(fn=get_layout_image_gallery, inputs=response, outputs=[gallery, b64], api_name="visualize-layout")
+        generate_btn.click(fn=get_ours_image, inputs=[response, overall_prompt_override, seed, num_inference_steps, dpm_scheduler, use_autocast, fg_seed_start, fg_blending_ratio, frozen_step_ratio, gligen_scheduled_sampling_beta, so_negative_prompt, overall_negative_prompt, show_so_imgs, scale_boxes], outputs=[gallery, b64], api_name="layout-to-image")
 
         gr.Examples(
             examples=stage2_examples,
             inputs=[response, overall_prompt_override, seed],
             outputs=[gallery],
             fn=get_ours_image,
-            cache_examples=True
+            # cache_examples=True
         )
 
     with gr.Tab("Baseline: Stable Diffusion"):
@@ -287,7 +308,7 @@ with gr.Blocks(
             inputs=[sd_prompt],
             outputs=[gallery],
             fn=get_baseline_image,
-            cache_examples=True
+            # cache_examples=True
         )
 
 g.launch()

generation.py

@@ -8,6 +8,8 @@ from models import pipelines, sam
 from utils import parse, latents
 from shared import model_dict, sam_model_dict, DEFAULT_SO_NEGATIVE_PROMPT, DEFAULT_OVERALL_NEGATIVE_PROMPT
 import gc
+from io import BytesIO
+import base64
 
 verbose = False
 
@@ -209,6 +211,10 @@ def run(
         
     gc.collect()
     torch.cuda.empty_cache()
+
+    with BytesIO() as buffer:
+        np.save(buffer, images[0])
+        img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
         
-    return images[0], so_img_list
+    return images[0], so_img_list, img_str