some cleaning and on the path to having token streamign

app.py

@@ -6,7 +6,8 @@ import gradio as gr
 
 from gradio_client.client import DEFAULT_TEMP_DIR
 from playwright.sync_api import sync_playwright
-from transformers import AutoProcessor, AutoModelForCausalLM
+from threading import Thread
+from transformers import AutoProcessor, AutoModelForCausalLM, TextIteratorStreamer
 from transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension
 from typing import List
 from PIL import Image
@@ -14,15 +15,12 @@ from PIL import Image
 from transformers.image_transforms import resize, to_channel_dimension_format
 
 
-API_TOKEN = os.getenv("HF_AUTH_TOKEN")
 DEVICE = torch.device("cuda")
 PROCESSOR = AutoProcessor.from_pretrained(
     "HuggingFaceM4/VLM_WebSight_finetuned",
-    token=API_TOKEN,
 )
 MODEL = AutoModelForCausalLM.from_pretrained(
     "HuggingFaceM4/VLM_WebSight_finetuned",
-    token=API_TOKEN,
     trust_remote_code=True,
     torch_dtype=torch.bfloat16,
 ).to(DEVICE)
@@ -134,20 +132,35 @@ def model_inference(
         k: v.to(DEVICE)
         for k, v in inputs.items()
     }
-    generated_ids = MODEL.generate(
-        **inputs,
+
+    streamer = TextIteratorStreamer(
+        PROCESSOR.tokenizer,
+        decode_kwargs=dict(
+            skip_special_tokens=True
+        ),
+        skip_prompt=True,
+    )
+    generation_kwargs = dict(
+        inputs,
         bad_words_ids=BAD_WORDS_IDS,
-        max_length=4096
+        max_length=4096,
+        streamer=streamer,
     )
-    generated_text = PROCESSOR.batch_decode(
-        generated_ids,
-        skip_special_tokens=True
-    )[0]
+    thread = Thread(
+        target=MODEL.generate,
+        kwargs=generation_kwargs,
+    )
+    thread.start()
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        print("before yield")
+        # yield generated_text, image
+        print("after yield")
 
     rendered_page = render_webpage(generated_text)
     return generated_text, rendered_page
 
-
 generated_html = gr.Code(
     label="Extracted HTML",
     elem_id="generated_html",
@@ -189,7 +202,7 @@ with gr.Blocks(title="Screenshot to HTML", theme=gr.themes.Base(), css=css) as d
                     regenerate_btn = gr.Button(
                         value="🔄 Regenerate", visible=True, min_width=120
                     )
-        with gr.Column(scale=4) as result_area:
+        with gr.Column(scale=4):
             rendered_html.render()
 
     with gr.Row():