Update app.py

app.py

@@ -8,6 +8,10 @@ import json
 import re
 import base64
 
+def decode_image(img_b64):
+    img_data = base64.b64decode(img_b64)
+    img = Image.open(io.BytesIO(img_data))
+    return img
 
 def get_image_data(fig):
     buf = io.BytesIO()
@@ -59,7 +63,8 @@ def gpt_inference(base_url, model, openai_key, prompt):
     print("Response JSON:", response.json())
     code = response.json()["choices"][0]["message"]["content"]
     print(f"code:{code}")
-    img = execute_code(code)
+    img_b64 = execute_code(code)
+    img = decode_image(img_b64)
     return img
 
 iface = gr.Interface(