screenshot2html

Running on Zero

App Files Files Community

VictorSanh commited on Jan 12

Commit

844c526

•

1 Parent(s): 5c49818

very big update

Browse files

Files changed (1) hide show

app.py +85 -16

app.py CHANGED Viewed

@@ -1,24 +1,72 @@
 import os
 import subprocess
 from playwright.sync_api import sync_playwright
 from typing import List
 from PIL import Image
-import gradio as gr
-from gradio_client.client import DEFAULT_TEMP_DIR
-from transformers import AutoProcessor, AutoModelForCausalLM
 API_TOKEN = os.getenv("HF_AUTH_TOKEN")
-# PROCESSOR = AutoProcessor.from_pretrained(
-#     "HuggingFaceM4/img2html",
-#     token=API_TOKEN,
-# )
-# MODEL = AutoModelForCausalLM.from_pretrained(
-#     "HuggingFaceM4/img2html",
-#     token=API_TOKEN,
-# )
 IMAGE_GALLERY_PATHS = [
@@ -36,11 +84,13 @@ def install_playwright():
 install_playwright()
 def add_file_gallery(
     selected_state: gr.SelectData,
     gallery_list: List[str]
 ):
-    return f"example_images/{gallery_list.root[selected_state.index].image.orig_name}"
 def render_webpage(
     html_css_code,
@@ -68,6 +118,22 @@ def render_webpage(
 def model_inference(
     image,
 ):
     CAR_COMPNAY = """<!DOCTYPE html>
 <html lang="en">
 <head>
@@ -189,8 +255,8 @@ def model_inference(
 </body>
 </html>"""
-    rendered_page = render_webpage(CAR_COMPNAY)
-    return CAR_COMPNAY, rendered_page
 generated_html = gr.Code(
@@ -216,7 +282,7 @@ with gr.Blocks(title="Img2html", theme=gr.themes.Base(), css=css) as demo:
     with gr.Row(equal_height=True):
         with gr.Column(scale=4, min_width=250) as upload_area:
             imagebox = gr.Image(
-                type="filepath",
                 label="Screenshot to extract",
                 visible=True,
                 sources=["upload", "clipboard"],
@@ -253,7 +319,6 @@ with gr.Blocks(title="Img2html", theme=gr.themes.Base(), css=css) as demo:
         triggers=[
             imagebox.upload,
             submit_btn.click,
-            template_gallery.select,
             regenerate_btn.click,
         ],
         fn=model_inference,
@@ -274,6 +339,10 @@ with gr.Blocks(title="Img2html", theme=gr.themes.Base(), css=css) as demo:
         inputs=[template_gallery],
         outputs=[imagebox],
         queue=False,
     )
     demo.load(queue=False)

 import os
 import subprocess
+import torch
+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 transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension
 from typing import List
 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/img2html",
+    token=API_TOKEN,
+)
+MODEL = AutoModelForCausalLM.from_pretrained(
+    "HuggingFaceM4/img2html", #TODO
+    token=API_TOKEN,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to(DEVICE)
+if MODEL.config.use_resampler:
+    image_seq_len = MODEL.config.perceiver_config.resampler_n_latents
+else:
+    image_seq_len = (
+        MODEL.config.vision_config.image_size // MODEL.config.vision_config.patch_size
+    ) ** 2
+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
+    image_rgba = image.convert("RGBA")
+    background = Image.new("RGBA", image_rgba.size, (255, 255, 255))
+    alpha_composite = Image.alpha_composite(background, image_rgba)
+    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)
+    x = resize(x, (960, 960), resample=PILImageResampling.BILINEAR)
+    x = PROCESSOR.image_processor.rescale(x, scale=1 / 255)
+    x = PROCESSOR.image_processor.normalize(
+        x,
+        mean=PROCESSOR.image_processor.image_mean,
+        std=PROCESSOR.image_processor.image_std
+    )
+    x = to_channel_dimension_format(x, ChannelDimension.FIRST)
+    x = torch.tensor(x)
+    return x
+## End of Utils
 IMAGE_GALLERY_PATHS = [
 install_playwright()
 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 model_inference(
     image,
 ):
+    if image is None:
+        raise ValueError("`image` is None. It should be a PIL image.")
+    inputs = PROCESSOR.tokenizer(
+        f"{BOS_TOKEN}<fake_token_around_image>{'<image>' * image_seq_len}<fake_token_around_image>",
+        return_tensors="pt"
+    )
+    inputs["pixel_values"] = PROCESSOR.image_processor(
+        [image],
+        transform=custom_transform
+    )
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    generated_ids = MODEL.generate(**inputs, bad_words_ids=BAD_WORDS_IDS)
+    generated_text = PROCESSOR.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    print(generated_text)
     CAR_COMPNAY = """<!DOCTYPE html>
 <html lang="en">
 <head>
 </body>
 </html>"""
+    rendered_page = render_webpage(generated_text)
+    return generated_text, rendered_page
 generated_html = gr.Code(
     with gr.Row(equal_height=True):
         with gr.Column(scale=4, min_width=250) as upload_area:
             imagebox = gr.Image(
+                type="pil",
                 label="Screenshot to extract",
                 visible=True,
                 sources=["upload", "clipboard"],
         triggers=[
             imagebox.upload,
             submit_btn.click,
             regenerate_btn.click,
         ],
         fn=model_inference,
         inputs=[template_gallery],
         outputs=[imagebox],
         queue=False,
+    ).success(
+        fn=model_inference,
+        inputs=[imagebox],
+        outputs=[generated_html, rendered_html],
     )
     demo.load(queue=False)