Update app.py

app.py

@@ -1,50 +1,69 @@
 import gradio as gr
-from huggingface_hub import InferenceApi
-import base64
+from transformers import AutoProcessor, Idefics2ForConditionalGeneration
+import re
+import time
 from PIL import Image
-import io
-
-client = InferenceApi("HuggingFaceM4/idefics2-8b")
-
-def image_to_base64(image):
-    buffered = io.BytesIO()
-    image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
-    return img_str
-
-def model_inference(image, text):
-    image_base64 = image_to_base64(image)
-    inputs = {
-        "inputs": {
-            "text": text,
-            "image": image
-        }
-    }
+import torch
+import spaces
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+
+processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b")
+
+model = Idefics2ForConditionalGeneration.from_pretrained(
+        "HuggingFaceM4/idefics2-8b",
+        torch_dtype=torch.bfloat16,
+        #_attn_implementation="flash_attention_2",
+        trust_remote_code=True).to("cuda")
+
+# import gradio as gr
+# from huggingface_hub import InferenceApi
+# import base64
+# from PIL import Image
+# import io
+
+# client = InferenceApi("HuggingFaceM4/idefics2-8b")
+
+# def image_to_base64(image):
+#     buffered = io.BytesIO()
+#     image.save(buffered, format="JPEG")
+#     img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+#     return img_str
+
+# def model_inference(image, text):
+#     image_base64 = image_to_base64(image)
+#     inputs = {
+#         "inputs": {
+#             "text": text,
+#             "image": image
+#         }
+#     }
     
-    result = client(inputs)
+#     result = client(inputs)
     
-    generated_text = result['generated_text']
-    return generated_text
-
-with gr.Blocks(css=".input_image {max-width: 100%; border: 1px solid #ccc; box-shadow: 0 0 10px #ccc; margin-bottom: 10px;} .output_textbox {min-height: 100px;}") as demo:
-    gr.Markdown("## Enhanced IDEFICS2 Demo")
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(label="Upload Image", type="pil", height=240, width=320)
-            query_input = gr.Textbox(label="Enter Prompt", placeholder="Type your prompt here...")
-        with gr.Column(scale=1):
-            output = gr.Textbox(label="Model Output", interactive=True, placeholder="Output will be displayed here...")
-
-    submit_btn = gr.Button("Generate")
-    submit_btn.click(model_inference, inputs=[image_input, query_input], outputs=output)
-
-    examples = [
-        ["american_football.png", "Explain in detail what is depicted in the picture"],
-        ["bike.png", "Explore the image closely and describe in detail what you discover."],
-        ["finance.png", "Provide a detailed description of everything you see in the image."],
-        ["science.png", "Please perform optical character recognition (OCR) on the uploaded image. Extract all text visible in the image accurately. Ensure to capture the text in its entirety and maintain the formatting as closely as possible to how it appears in the image. After extracting the text, display it in a clear and readable format, making sure that any special characters or symbols are also accurately represented. Provide the extracted text as output."],
-        ["spirituality.png", "Please perform optical character recognition (OCR) on the uploaded image. Extract all text visible in the image accurately. Ensure to capture the text in its entirety and maintain the formatting as closely as possible to how it appears in the image. After extracting the text, display it in a clear and readable format, making sure that any special characters or symbols are also accurately represented. Provide the extracted text as output."]
-    ]
-    gr.Examples(examples=examples, inputs=[image_input, query_input], outputs=output)
-
-demo.launch(debug=True)
+#     generated_text = result['generated_text']
+#     return generated_text
+
+# with gr.Blocks(css=".input_image {max-width: 100%; border: 1px solid #ccc; box-shadow: 0 0 10px #ccc; margin-bottom: 10px;} .output_textbox {min-height: 100px;}") as demo:
+#     gr.Markdown("## Enhanced IDEFICS2 Demo")
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             image_input = gr.Image(label="Upload Image", type="pil", height=240, width=320)
+#             query_input = gr.Textbox(label="Enter Prompt", placeholder="Type your prompt here...")
+#         with gr.Column(scale=1):
+#             output = gr.Textbox(label="Model Output", interactive=True, placeholder="Output will be displayed here...")
+
+#     submit_btn = gr.Button("Generate")
+#     submit_btn.click(model_inference, inputs=[image_input, query_input], outputs=output)
+
+#     examples = [
+#         ["american_football.png", "Explain in detail what is depicted in the picture"],
+#         ["bike.png", "Explore the image closely and describe in detail what you discover."],
+#         ["finance.png", "Provide a detailed description of everything you see in the image."],
+#         ["science.png", "Please perform optical character recognition (OCR) on the uploaded image. Extract all text visible in the image accurately. Ensure to capture the text in its entirety and maintain the formatting as closely as possible to how it appears in the image. After extracting the text, display it in a clear and readable format, making sure that any special characters or symbols are also accurately represented. Provide the extracted text as output."],
+#         ["spirituality.png", "Please perform optical character recognition (OCR) on the uploaded image. Extract all text visible in the image accurately. Ensure to capture the text in its entirety and maintain the formatting as closely as possible to how it appears in the image. After extracting the text, display it in a clear and readable format, making sure that any special characters or symbols are also accurately represented. Provide the extracted text as output."]
+#     ]
+#     gr.Examples(examples=examples, inputs=[image_input, query_input], outputs=output)
+
+# demo.launch(debug=True)