Update app.py

app.py

@@ -35,7 +35,7 @@ def get_concat_v_cut(in1, in2):
 
 
 
-def run_script(url: str, height: int, width: int, check_b):
+def run_script(url: str, height: int, width: int, check_b,check_h):
     mes_box=[]
     out_box=[]
     uid=uuid.uuid4()
@@ -64,18 +64,22 @@ def run_script(url: str, height: int, width: int, check_b):
             messages = driver.find_elements(By.CLASS_NAME, "message")
             print (len(messages))
             mes_blocks=len(messages)%2
-            try:
-                for ea in check_b:
+            for ea in check_b:
+                try:    
                     ea = int(ea)
                     ea = (ea*2)-2
                     print (ea)
                     messages[ea].screenshot(f'{ea}-{uid}-tmp.png')
                     messages[ea+1].screenshot(f'{ea+1}-{uid}-tmp.png')
-                    out_box.append(get_concat_v_cut(f'{ea}-{uid}-tmp.png',f'{ea+1}-{uid}-tmp.png'))
+                    if check_h:
+                        head_paste = get_concat_v_cut(f'head-{uid}-tmp.png',f'{ea}-{uid}-tmp.png'))
+                        out_box.append(get_concat_v_cut(head_paste,f'{ea+1}-{uid}-tmp.png'))
+                    if not check_h:    
+                        out_box.append(get_concat_v_cut(f'{ea}-{uid}-tmp.png',f'{ea+1}-{uid}-tmp.png'))
                     print(f'out_box:: {out_box}')
                     #mes_box.clear()
-            except Exception:
-                mes="<center>Selected blocks returned an error"
+                except Exception:
+                    mes="<center>Some blocks returned an error"
             '''
             for i,mes in enumerate(messages):
                 print (mes.text)
@@ -128,8 +132,9 @@ with gr.Blocks() as app:
             btn= gr.Button()
             
         with gr.Column():
+            check_h=gr.Checkbox(label="Show Header", value=True)
+            check_b=gr.CheckboxGroup(label="Chatblocks", choices=[1,2,3,4,5,6,7,8,9,10], value=[1])
             with gr.Row():
-                check_b=gr.CheckboxGroup(label="Chatblocks", choices=[1,2,3,4,5,6,7,8,9,10])
                 height=gr.Number(label="Height", value=4096)
                 width=gr.Number(label="Width",value=800)
     message=gr.HTML('<center>Enter URL')
@@ -138,5 +143,5 @@ with gr.Blocks() as app:
             outgal=gr.Gallery()
         outim = gr.Image()
 
-    btn.click(run_script,[inp,height,width,check_b],[outim,message,outgal])
+    btn.click(run_script,[inp,height,width,check_b,check_h],[outim,message,outgal])
 app.launch()