Update app.py

app.py

@@ -15,18 +15,21 @@ from PIL import Image
 
 from transformers.image_transforms import resize, to_channel_dimension_format
 
-
+# Install flash-attn without CUDA build isolation
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-DEVICE = torch.device("cuda")
+# Set the device to GPU if available, otherwise use CPU
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 PROCESSOR = AutoProcessor.from_pretrained(
     "HuggingFaceM4/VLM_WebSight_finetuned",
 )
 MODEL = AutoModelForCausalLM.from_pretrained(
     "HuggingFaceM4/VLM_WebSight_finetuned",
     trust_remote_code=True,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
 ).to(DEVICE)
+
+# Determine image sequence length
 if MODEL.config.use_resampler:
     image_seq_len = MODEL.config.perceiver_config.resampler_n_latents
 else:
@@ -36,12 +39,9 @@ else:
 BOS_TOKEN = PROCESSOR.tokenizer.bos_token
 BAD_WORDS_IDS = PROCESSOR.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
 
-
 ## Utils
 
 def convert_to_rgb(image):
-    # `image.convert("RGB")` would only work for .jpg images, as it creates a wrong background
-    # for transparent images. The call to `alpha_composite` handles this case
     if image.mode == "RGB":
         return image
 
@@ -51,8 +51,6 @@ def convert_to_rgb(image):
     alpha_composite = alpha_composite.convert("RGB")
     return alpha_composite
 
-# The processor is the same as the Idefics processor except for the BICUBIC interpolation inside siglip,
-# so this is a hack in order to redefine ONLY the transform method
 def custom_transform(x):
     x = convert_to_rgb(x)
     x = to_numpy_array(x)
@@ -69,13 +67,7 @@ def custom_transform(x):
 
 ## End of Utils
 
-
-IMAGE_GALLERY_PATHS = [
-    f"example_images/{ex_image}"
-    for ex_image in os.listdir(f"example_images")
-]
-
-
+# Install Playwright
 def install_playwright():
     try:
         subprocess.run(["playwright", "install"], check=True)
@@ -85,17 +77,15 @@ def install_playwright():
 
 install_playwright()
 
+IMAGE_GALLERY_PATHS = [
+    f"example_images/{ex_image}"
+    for ex_image in os.listdir(f"example_images")
+]
 
-def add_file_gallery(
-    selected_state: gr.SelectData,
-    gallery_list: List[str]
-):
+def add_file_gallery(selected_state: gr.SelectData, gallery_list: List[str]):
     return Image.open(gallery_list.root[selected_state.index].image.path)
 
-
-def render_webpage(
-    html_css_code,
-):
+def render_webpage(html_css_code):
     with sync_playwright() as p:
         browser = p.chromium.launch(headless=True)
         context = browser.new_context(
@@ -115,11 +105,8 @@ def render_webpage(
 
     return Image.open(output_path_screenshot)
 
-
 @spaces.GPU(duration=180)
-def model_inference(
-    image,
-):
+def model_inference(image):
     if image is None:
         raise ValueError("`image` is None. It should be a PIL image.")
 
@@ -132,10 +119,7 @@ def model_inference(
         [image],
         transform=custom_transform
     )
-    inputs = {
-        k: v.to(DEVICE)
-        for k, v in inputs.items()
-    }
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
 
     streamer = TextIteratorStreamer(
         PROCESSOR.tokenizer,
@@ -147,16 +131,6 @@ 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,
@@ -172,20 +146,8 @@ def model_inference(
         generated_text += new_text
         yield generated_text, rendered_image
 
-
-generated_html = gr.Code(
-    label="Extracted HTML",
-    elem_id="generated_html",
-)
-rendered_html = gr.Image(
-    label="Rendered HTML",
-    show_download_button=False,
-    show_share_button=False,
-)
-# rendered_html = gr.HTML(
-#     label="Rendered HTML"
-# )
-
+generated_html = gr.Code(label="Extracted HTML", elem_id="generated_html")
+rendered_html = gr.Image(label="Rendered HTML", show_download_button=False, show_share_button=False)
 
 css = """
 .gradio-container{max-width: 1000px!important}
@@ -193,7 +155,6 @@ h1{display: flex;align-items: center;justify-content: center;gap: .25em}
 *{transition: width 0.5s ease, flex-grow 0.5s ease}
 """
 
-
 with gr.Blocks(title="Screenshot to HTML", theme=gr.themes.Base(), css=css) as demo:
     gr.Markdown(
         "Since the model used for this demo *does not generate images*, it is more effective to input standalone website elements or sites with minimal image content."
@@ -208,15 +169,11 @@ with gr.Blocks(title="Screenshot to HTML", theme=gr.themes.Base(), css=css) as d
             )
             with gr.Group():
                 with gr.Row():
-                    submit_btn = gr.Button(
-                        value="▶️ Submit", visible=True, min_width=120
-                    )
+                    submit_btn = gr.Button(value="▶️ Submit", visible=True, min_width=120)
                     clear_btn = gr.ClearButton(
                         [imagebox, generated_html, rendered_html], value="🧹 Clear", min_width=120
                     )
-                    regenerate_btn = gr.Button(
-                        value="🔄 Regenerate", visible=True, min_width=120
-                    )
+                    regenerate_btn = gr.Button(value="🔄 Regenerate", visible=True, min_width=120)
         with gr.Column(scale=4):
             rendered_html.render()
 
@@ -235,11 +192,7 @@ with gr.Blocks(title="Screenshot to HTML", theme=gr.themes.Base(), css=css) as d
         )
 
     gr.on(
-        triggers=[
-            imagebox.upload,
-            submit_btn.click,
-            regenerate_btn.click,
-        ],
+        triggers=[imagebox.upload, submit_btn.click, regenerate_btn.click],
         fn=model_inference,
         inputs=[imagebox],
         outputs=[generated_html, rendered_html],