Update app.py

app.py

@@ -68,12 +68,19 @@ def make_filename(inp):
     filename=first_name+str(last_name)
     return filename
     
-def get_screenshot(chat,css_in,w,h,user_c,bot_c,background_c,font_c):
-    print(len(chat))
+def get_screenshot(chat,block,css_in,w,h,user_c,bot_c,background_c,font_c):
+    chat_len=len(chat)
+    print(chat_len)
+    block_box=[]
     if user_c:
         user_c=user_c.strip('""').strip("''")
         css_in=css_in.replace("$user_c",user_c).replace("$bot_c",bot_c).replace("$background_c",background_c).replace("$font_c",font_c)
     html_body=""
+    if block:
+        for i in range(chat_len):
+            if i+1 in block:
+                block_box.append(chat[i])
+        chat=block_box
     for user,bot in chat:
         html_body += html_user.replace("$chat",user)
         html_body += html_bot.replace("$chat",bot)
@@ -159,7 +166,7 @@ with gr.Blocks(css=css) as app:
                 img=gr.Image(type='filepath')
     html_view=gr.HTML()
     btn.click(chat_inf,[sys_inp,inp,chat_b],chat_b)
-    im_btn.click(get_screenshot,[chat_b,css_box,im_width,im_height,user_c,bot_c,background_c,font_c],[img,html_view])
+    im_btn.click(get_screenshot,[chat_b,chatblock,css_box,im_width,im_height,user_c,bot_c,background_c,font_c],[img,html_view])
     #im_btn.click(get_screenshot,[chat_b,im_height,im_width,chatblock,theme,wait_time],img)
     #app.load(get_screenshot,inp,img)
 app.launch()