Update app.py

app.py

@@ -1,7 +1,7 @@
-import spaces
 import argparse
 import json
 import numpy as np
+import os
 
 import gradio as gr
 import requests
@@ -13,12 +13,12 @@ HUGGINGFACE=True
 MOCK = False
 TEST_FOLDER = "c4f5"
 
-MODEL_NAME="xu3kev/deepseekcoder-7b-logo-pbe"
-# MODEL_NAME="openlm-research/open_llama_3b" 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-hug_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map='auto', load_in_8bit=True)
-hug_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+if HUGGINGFACE:
+    MODEL_NAME="xu3kev/deepseekcoder-7b-logo-pbe"
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    hug_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, attn_implementation="flash_attention_2",).to('cuda')
+    hug_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 INPUT_STRUCTION_TEMPLATE = """Here is a gray scale images representing with integer values 0-9.
 {image_str}
@@ -37,7 +37,7 @@ MOCK_RESPONSE = [
 """
 ] * 16
 
-LOGO_HEADER = """from myturtle import Turtle
+LOGO_HEADER = """from myturtle_cv import Turtle
 from myturtle import HALF_INF, INF, EPS_DIST, EPS_ANGLE
 
 turtle = Turtle()
@@ -169,7 +169,8 @@ PORT = 8008
 MODEL_NAME="./axolotl/lora-logo_fix_full_deepseek33b_ds33i_epoch3_lr_0.0002_alpha_512_r_512_merged"
 MODEL_NAME="./axolotl/lora-logo_fix_full_deepseek7b_ds33i_lr_0.0002_alpha_512_r_512_merged"
 
-def generate_grid_images(folder):
+
+def generate_grid_images(gif_results):
     import matplotlib.patches as patches
     import matplotlib.pyplot as plt
     num_rows, num_cols = 8,8
@@ -217,13 +218,12 @@ def generate_grid_images(folder):
     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=32, do_sample=True)
+        generated_ids = hug_model.generate(**model_inputs, max_length=1400, temperature=1, num_return_sequences=8, do_sample=True)
         responses = hug_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
         codes = []
         for response in responses:
@@ -296,12 +296,15 @@ def run_code(new_folder, counter, code):
     code_with_header_and_save= f"""
 {LOGO_HEADER}
 {code}
-turtle.save('{fname}')
+# turtle.save('{fname}')
+gif = turtle.save_gif('')
 """
     try:
-        func_timeout(3, exec, args=(code_with_header_and_save, {}))
-        matplotlib.pyplot.close()
+        results = {}
+        func_timeout(3, exec, args=(code_with_header_and_save, results))
         # exec(code_with_header_and_save, globals())
+        if 'gif' in results:
+            return results['gif']
     except FunctionTimedOut:
         print("Timeout")
     except Exception as e:
@@ -317,8 +320,8 @@ def run(img_str):
         codes = MOCK_RESPONSE
     
     gradio_test_images_folder = "gradio_test_images"
-    import os
-    os.makedirs(gradio_test_images_folder, exist_ok=True)
+    # import os
+    # os.makedirs(gradio_test_images_folder, exist_ok=True)
 
     counter = 0
     # generate a random hash id
@@ -326,16 +329,21 @@ def run(img_str):
     import random
     random_id = hashlib.md5(str(random.random()).encode()).hexdigest()[0:4]
     new_folder = os.path.join(gradio_test_images_folder, random_id)
-    os.makedirs(new_folder, exist_ok=True)
+    # os.makedirs(new_folder, exist_ok=True)
+
+
+
+    for code in tqdm(codes):
+        pass
 
-    print('about to execute')
     from concurrent.futures import ProcessPoolExecutor
     from concurrent.futures import as_completed
+    gif_results = []
     with ProcessPoolExecutor() as executor:
         futures = [executor.submit(run_code, new_folder, i, code) for i, code in enumerate(codes)]
         for future in as_completed(futures):
             try:
-                future.result()
+                gif_results.append(future.result())
             except Exception as exc:
                 print(f'Generated an exception: {exc}')
 
@@ -346,10 +354,9 @@ def run(img_str):
         # out, errs = p.communicate()
         # out, errs, = out.decode(), errs.decode()
     # render
-    print('finish execute')
     print(random_id)
     folder_path = f"gradio_test_images/{random_id}"
-    return folder_path, codes
+    return gif_results, codes
 
 
 def test_gen_img_wrapper(_):
@@ -363,6 +370,16 @@ def int_img_to_str(integer_img):
     image_str = "\n".join(lines)
     return image_str
 
+
+def create_tmp_folder():
+    import hashlib
+    import random
+    random_id = hashlib.md5(str(random.random()).encode()).hexdigest()[0:4]
+    folder_name = f"generated_images_{random_id}"
+    os.makedirs(folder_name, exist_ok=True)
+    return folder_name
+
+
 def img_to_code_img(sketchpad_img):
     img = sketchpad_img['layers'][0]
     image_array = np.array(img)
@@ -392,11 +409,18 @@ def img_to_code_img(sketchpad_img):
         img_str = int_img_to_str(int_img)
         print(img_str)
 
-    folder_path, codes = run(img_str)
+    gif_results, codes = run(img_str)
 
-    generated_grid_img = generate_grid_images(folder_path)
+    # generated_grid_img = generate_grid_images()
 
-    return generated_grid_img
+    # return generated_grid_img[0] 
+    
+    folder = create_tmp_folder()
+    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])
+    return [f"{folder}/img{i}.gif" for i in range(len(gif_results))]
 
 
 def main():
@@ -411,26 +435,29 @@ def main():
         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 checkout our [paper](https://arxiv.org/abs/2406.08316) for more details!")
         gr.Markdown("## Draw logo")
-        with gr.Column():
-            canvas = gr.Sketchpad(canvas_size=(512,512), brush=Brush(colors=["black"], default_size=2, color_mode='fixed'))
-            submit_button = gr.Button("Submit")
-            output_image = gr.Image(label="output")
+        with gr.Row():
+            with gr.Column():
+                canvas = gr.Sketchpad(canvas_size=(512,512), brush=Brush(colors=["black"], default_size=2, color_mode='fixed'))
+                submit_button = gr.Button("Submit")
+            output_gallery = gr.Gallery(
+        label="Generated images", show_label=False, elem_id="gallery"
+    , columns=[3], rows=[3], object_fit="contain", height="auto")
+            # output_image = gr.Image(label="output")
         
-        submit_button.click(img_to_code_img, inputs=canvas, outputs=output_image)
-#         demo.load(
-#         None,
-#         None,
-#         js="""
-#   () => {
-#   const params = new URLSearchParams(window.location.search);
-#   if (!params.has('__theme')) {
-#     params.set('__theme', 'light');
-#     window.location.search = params.toString();
-#   }
-#   }""",
-#         )
+        submit_button.click(img_to_code_img, inputs=canvas, outputs=output_gallery)
+        demo.load(
+        None,
+        None,
+        js="""
+  () => {
+  const params = new URLSearchParams(window.location.search);
+  if (!params.has('__theme')) {
+    params.set('__theme', 'light');
+    window.location.search = params.toString();
+  }
+  }""",
+        )
 
     demo.launch(share=True)