make token streaming work

app.py

@@ -139,9 +139,6 @@ def model_inference(
 
     streamer = TextIteratorStreamer(
         PROCESSOR.tokenizer,
-        decode_kwargs=dict(
-            skip_special_tokens=True
-        ),
         skip_prompt=True,
     )
     generation_kwargs = dict(
@@ -150,6 +147,16 @@ def model_inference(
         max_length=4096,
         streamer=streamer,
     )
+    # Regular generation version
+    # generation_kwargs.pop("streamer")
+    # generated_ids = MODEL.generate(**generation_kwargs)
+    # generated_text = PROCESSOR.batch_decode(
+    #     generated_ids,
+    #     skip_special_tokens=True
+    # )[0]
+    # rendered_page = render_webpage(generated_text)
+    # return generated_text, rendered_page
+    # Token streaming version
     thread = Thread(
         target=MODEL.generate,
         kwargs=generation_kwargs,
@@ -157,15 +164,14 @@ def model_inference(
     thread.start()
     generated_text = ""
     for new_text in streamer:
+        if "</s>" in new_text:
+            new_text = new_text.replace("</s>", "")
+            rendered_image = render_webpage(generated_text)
+        else:
+            rendered_image = None
         generated_text += new_text
-        print("before yield")
-        # yield generated_text, image
-        print("after yield")
+        yield generated_text, rendered_image
 
-    # Sanity hack
-    generated_text = generated_text.replace("</s>", "")
-    rendered_page = render_webpage(generated_text)
-    return generated_text, rendered_page
 
 generated_html = gr.Code(
     label="Extracted HTML",
@@ -234,26 +240,22 @@ with gr.Blocks(title="Screenshot to HTML", theme=gr.themes.Base(), css=css) as d
         fn=model_inference,
         inputs=[imagebox],
         outputs=[generated_html, rendered_html],
-        queue=False,
     )
     regenerate_btn.click(
         fn=model_inference,
         inputs=[imagebox],
         outputs=[generated_html, rendered_html],
-        queue=False,
     )
     template_gallery.select(
         fn=add_file_gallery,
         inputs=[template_gallery],
         outputs=[imagebox],
-        queue=False,
     ).success(
         fn=model_inference,
         inputs=[imagebox],
         outputs=[generated_html, rendered_html],
-        queue=False,
     )
-    demo.load(queue=False)
+    demo.load()
 
 demo.queue(max_size=40, api_open=False)
 demo.launch(max_threads=400)