Update app.py

app.py

@@ -6,13 +6,16 @@ import os
 
 import gradio as gr
 import requests
-from openai import OpenAI
+# from openai import OpenAI
 from func_timeout import FunctionTimedOut, func_timeout
 from tqdm import tqdm
 
 HUGGINGFACE=True
-MOCK = False
+MOCK = not HUGGINGFACE
 TEST_FOLDER = "c4f5"
+NUM_RETURN_SEQ = 10
+
+DROPDOWN = None
 
 if HUGGINGFACE:
     MODEL_NAME="xu3kev/deepseekcoder-7b-logo-pbe"
@@ -36,7 +39,7 @@ MOCK_RESPONSE = [
             forward(2*i)
             left(90.0)
 """
-] * 16
+] * 10
 
 LOGO_HEADER = """from myturtle_cv import Turtle
 from myturtle import HALF_INF, INF, EPS_DIST, EPS_ANGLE
@@ -218,13 +221,14 @@ def generate_grid_images(gif_results):
     plt.close(fig)
     return image_array
 
+
 @spaces.GPU
 def llm_call(question_prompt, model_name, 
     temperature=1, max_tokens=320, 
     top_p=1, n_samples=64, stop=None):
     if HUGGINGFACE:
         model_inputs = hug_tokenizer([question_prompt], return_tensors="pt").to('cuda')
-        generated_ids = hug_model.generate(**model_inputs, max_length=1400, temperature=1, num_return_sequences=10, do_sample=True)
+        generated_ids = hug_model.generate(**model_inputs, max_length=1400, temperature=1, num_return_sequences=NUM_RETURN_SEQ, do_sample=True)
         responses = hug_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
         codes = []
         for response in responses:
@@ -334,8 +338,9 @@ def run(img_str):
 
 
 
-    for code in tqdm(codes):
-        pass
+    # for code in tqdm(codes):
+    #     pass
+    print(f"Running {len(codes)} codes")
 
     from concurrent.futures import ProcessPoolExecutor
     from concurrent.futures import as_completed
@@ -381,10 +386,26 @@ def create_tmp_folder():
     return folder_name
 
 
+CODES = []
 def img_to_code_img(sketchpad_img):
-    img = sketchpad_img['layers'][0]
+    
+    from PIL import Image
+    # with open("debug_background.png", "wb") as f:
+    #     # convert numpy to png
+    #     numpy_array = sketchpad_img['background']
+    #     img = Image.fromarray(numpy_array)
+    #     img.save(f)
+    # with open("debug_composite.png", "wb") as f:
+    #     # convert numpy to png
+    #     numpy_array = sketchpad_img['composite']
+    #     img = Image.fromarray(numpy_array)
+    #     img.save(f)
+
+    # img = sketchpad_img['layers'][0]
+    img = sketchpad_img['composite']
     image_array = np.array(img)
-    image_array = 255 - image_array[:,:,3]
+    image_array = image_array[:,:,0]
+    # image_array = 255 - image_array[:,:,3]
 
     # height, width = image_array.shape
     # output_size = 512
@@ -417,13 +438,16 @@ def img_to_code_img(sketchpad_img):
     # return generated_grid_img[0] 
     
     folder = create_tmp_folder()
-    img_names = [] 
+    global CODES
+    CODES = []
     for i in range(len(gif_results)):
         if gif_results[i]:
             with open(f"{folder}/img{i}.gif", "wb") as f:
                 f.write(gif_results[i])
-                img_names.append(f"img{i}.gif")
-    return [f"{folder}/{name}" for name in img_names]
+            CODES.append(f"```python\n{codes[i]}\n```")
+        else:
+            CODES.append("#### Execution Error/Timeout; Skip")
+    return [f"{folder}/img{i}.gif" for i in range(len(gif_results))]
 
 
 def main():
@@ -434,22 +458,46 @@ def main():
     from gradio import Brush
     theme = gr.themes.Default().set(
     ) 
+    import os
+    # get all png files under demo_example
+
+    example_input_images = []
+    for root, dirs, files in os.walk("demo_example"):
+        for file in files:
+            if file.endswith(".png"):
+                example_input_images.append(os.path.join(root, file))
+
+    canvas = gr.Sketchpad(canvas_size=(512,512), brush=Brush(colors=["black"], default_size=2, color_mode='fixed'))
     with gr.Blocks(theme=theme) as demo:
         gr.Markdown('# Visual Program Synthesis with LLM')
         gr.Markdown("""LOGO/Turtle graphics Programming-by-Example problems aims to synthesize a program that generates the given target image, where the program uses drawing library similar to Python Turtle.""")
         gr.Markdown("""Here we can draw a target image using the sketchpad, and see what kinds of graphics program LLM generates. To allow the LLM to visually perceive the input image, we convert the image to ASCII strings.""")
         gr.Markdown("Please check out our [paper](https://arxiv.org/abs/2406.08316) for more details!")
-        gr.Markdown("## Draw logo")
+        gr.Markdown("## Select an example logo input or draw your own logo!")
         with gr.Row():
             with gr.Column(scale=1):
-                canvas = gr.Sketchpad(canvas_size=(512,512), brush=Brush(colors=["black"], default_size=2, color_mode='fixed'))
-                submit_button = gr.Button("Submit")
+                gr.Examples(example_input_images, inputs=canvas)
+                canvas.render()
+                submit_button = gr.Button("Generate Programs")
             with gr.Column(scale=4):
                 output_gallery = gr.Gallery(
-            label="Generated images", show_label=False, elem_id="gallery"
-        , columns=[5], rows=[2], object_fit="contain", height="auto")
-                # output_image = gr.Image(label="output")
+        label="Generated Images", show_label=True, elem_id="gallery"
+    , columns=[5], rows=[2], object_fit="contain", height="auto")
+                with gr.Group():
+                    dropdown = gr.Dropdown([f"sample {i+1}" for i in range(NUM_RETURN_SEQ)], label='show generated program samples')
+                    code_block = gr.Markdown('')
+                    def update_code(sample_idx):
+                        int_idx = int(sample_idx.split(" ")[1]) - 1
+                        if int_idx < len(CODES):
+                            return CODES[int_idx]
+                        else:
+                            return "### Please submit an image to generate programs."
+                        #return gr.Markdown('333')
+                    dropdown.input(update_code, dropdown, code_block)
+            # output_image = gr.Image(label="output")
         
+        global DROPDOWN
+        DROPDOWN = dropdown
         submit_button.click(img_to_code_img, inputs=canvas, outputs=output_gallery)
         demo.load(
         None,