Update app.py

app.py

@@ -3,7 +3,7 @@ from gradio_client import Client
 from huggingface_hub import InferenceClient
 import random
 
-ss_client = Client("")
+ss_client = Client("https://xilixmeaty40-html-image-current-tabx.hf.space/")
 
 with open("models.txt", "r") as file:
     models = file.read().splitlines()
@@ -72,6 +72,14 @@ def chat_inf(system_prompt, prompt, history, memory, client_choice, seed, temp,
             print(f"Error during model inference: {e}")
             yield [("Error", "The model failed to respond, possibly due to a timeout. Please try again.")], memory
 
+def get_screenshot(chat, height=5000, width=600, chatblock=[], theme="light", wait=3000, header=True):
+    tog = 0
+    if chatblock:
+        tog = 3
+    result = ss_client.predict(str(chat), height, width, chatblock, header, theme, wait, api_name="/run_script")
+    out = f'https://xilixmeaty40-html-image-current-tabx.hf.space/file={result[tog]}'
+    return out
+
 def clear_fn():
     return None, None, None, None
 
@@ -99,7 +107,7 @@ with gr.Blocks() as app:
                     with gr.Column(scale=1):
                         stop_btn = gr.Button("Stop")
                         clear_btn = gr.Button("Clear")
-                client_choice = gr.Dropdown(label="Models", type='index', choices=[c for c in models], value=models[0], interactive=True)
+                client_choice = gr.Dropdown(label="Models", type='index', choices=[c for c in models if c in [client.model_id for client in clients]], value=models[0], interactive=True)
             with gr.Column(scale=1):
                 rand = gr.Checkbox(label="Random Seed", value=True)
                 seed = gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, step=1, value=rand_val)
@@ -108,9 +116,19 @@ with gr.Blocks() as app:
                 top_p = gr.Slider(label="Top-P", step=0.01, minimum=0.01, maximum=1.0, value=0.49)
                 rep_p = gr.Slider(label="Repetition Penalty", step=0.01, minimum=0.1, maximum=2.0, value=0.99)
                 chat_mem = gr.Number(label="Chat Memory", value=4)
+        with gr.Accordion(label="Screenshot", open=False):
+            im_btn = gr.Button("Screenshot")
+            img = gr.Image(type='filepath')
+            im_height = gr.Number(label="Height", value=5000)
+            im_width = gr.Number(label="Width", value=500)
+            wait_time = gr.Number(label="Wait Time", value=3000)
+            theme = gr.Radio(label="Theme", choices=["light", "dark"], value="light")
+            chatblock = gr.Dropdown(label="Chatblocks", choices=[c for c in range(0, 21)], value=0, type="index")
+            header = gr.Checkbox(label="Include header?", value=True)
     rand.change(check_rand, inputs=[rand, seed], outputs=seed)
     btn.click(chat_inf, inputs=[sys_inp, inp, chat_b, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, custom_prompt], outputs=[chat_b, memory])
     stop_btn.click(lambda: None, [], chat_b.stop)
     clear_btn.click(clear_fn, [], [inp, sys_inp, chat_b, memory])
+    im_btn.click(get_screenshot, inputs=[chat_b, im_height, im_width, chatblock, theme, wait_time, header], outputs=img)
 
 app.launch()