Update app.py

app.py

@@ -12,6 +12,7 @@ import datetime
 from pypdf import PdfReader
 import uuid
 from PIL import Image
+
 from agent import (
     PREFIX,
     GET_CHART,
@@ -441,7 +442,36 @@ def create_image(url):
     #output = cairosvg.svg2png(
     #    bytestring=open('tmp.svg').read().encode('utf-8'), write_to="output.png")        
     return "tmp.png"
-    
+
+def create_ss(graph):
+    uid=uuid.uuid4()
+    code_out=""
+    for ea in graph.split("\n"):
+        code=ea.strip().strip("\n")
+        code_out+=code
+    url=f"https://omnibus-text-to-chart.hf.space/file=merm.html?mermaid={code_out}"
+    is_url=True
+    wait_time=10000
+    if is_url:
+        options = webdriver.ChromeOptions()
+        options.add_argument('--headless')
+        options.add_argument('--no-sandbox')
+        options.add_argument('--disable-dev-shm-usage')
+        mes='<center>operation success'
+        try:
+            driver = webdriver.Chrome(options=options)
+            driver.get(url)
+            html=driver.page_source
+            print(html)            
+            driver.implicitly_wait(int(wait_time*30))
+            driver.set_window_size(int(width), int(height))
+            obj = driver.find_element(By.ID, "char_box")            
+            screenshot = obj.screenshot(f'{uid}-tmp.png')
+    return(f'{uid}-tmp.png')    
+
+
+
+
 score_js="""
 function(text_input) {
     console.log(text_input);
@@ -503,7 +533,8 @@ with gr.Blocks() as app:
     def return_score(text):
         print(text)
         return text
-    create_im.click(return_score,score,[score],_js=score_js).then(create_image,score,svg_img)
+    create_im.click(create_ss,e_box,svg_img)
+    #create_im.click(return_score,score,[score],_js=score_js).then(create_image,score,svg_img)
     #zoom_btn.change(zoom_update,zoom_btn,None)
     upd_button.click(mm,[e_box,zoom_btn],[m_box,url_box])
     #inp_query.change(search_models,inp_query,models_dd)