Update app.py

app.py

@@ -1,130 +1,118 @@
-import sys
-from threading import Thread
-
 import gradio as gr
 import torch
+from PIL import Image
 from transformers import AutoModel, AutoProcessor
-from transformers import StoppingCriteria, TextIteratorStreamer, StoppingCriteriaList
 
+# Set device
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-model = AutoModel.from_pretrained("unum-cloud/uform-gen2-dpo", trust_remote_code=True).to(device)
-processor = AutoProcessor.from_pretrained("unum-cloud/uform-gen2-dpo", trust_remote_code=True)
-
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = [151645]
-        for stop_id in stop_ids:
-            if input_ids[0][-1] == stop_id:
-                return True
-        return False
-
-@torch.no_grad()
-def response(message, history, image):
-    stop = StopOnTokens()
-
-    messages = [{"role": "system", "content": "You are a helpful assistant."}]
-
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
-
-    if len(messages) == 1:
-        message = f" <image>{message}"
-
-    messages.append({"role": "user", "content": message})
-
-    model_inputs = processor.tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
+# Load the model and processor
+# We use trust_remote_code=True because this model has custom code.
+model = AutoModel.from_pretrained(
+    "unum-cloud/uform-gen2-dpo",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16
+).to(device)
+processor = AutoProcessor.from_pretrained(
+    "unum-cloud/uform-gen2-dpo",
+    trust_remote_code=True
+)
+
+def transcribe_image(image):
+    """
+    Generates a caption for the given image.
+    """
+    if image is None:
+        return "Please upload an image."
+    
+    prompt = "a photo of"
+    inputs = processor(
+        text=[prompt],
+        images=[image],
         return_tensors="pt"
-    )
-
-    image = (
-            processor.feature_extractor(image)
-            .unsqueeze(0)
-    )
+    ).to(device)
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            do_sample=False,
+            use_cache=True,
+            max_new_tokens=128,
+            eos_token_id=32001,
+            pad_token_id=processor.tokenizer.pad_token_id
+        )
 
-    attention_mask = torch.ones(
-        1, model_inputs.shape[1] + processor.num_image_latents - 1
-    )
+    prompt_len = inputs["input_ids"].shape[1]
+    decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
     
-    model_inputs = {
-        "input_ids": model_inputs,
-        "images": image,
-        "attention_mask": attention_mask
-    }
-
-    model_inputs = {k: v.to(device) for k, v in model_inputs.items()}
-
-    streamer = TextIteratorStreamer(processor.tokenizer, timeout=30., skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=1024,
-        stopping_criteria=StoppingCriteriaList([stop])
+    # Remove the end-of-sequence token
+    result = decoded_text.replace("<|im_end|>", "").strip()
+    return result
+
+def visual_question_answer(image, question):
+    """
+    Answers a question about the given image.
+    """
+    if image is None:
+        return "Please upload an image."
+    if not question:
+        return "Please ask a question."
+
+    # The model expects the prompt to be in a specific format.
+    prompt = f"<|im_start|>question\n{question}<|im_end|><|im_start|>answer\n"
+    
+    inputs = processor(
+        text=[prompt],
+        images=[image],
+        return_tensors="pt"
+    ).to(device)
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            do_sample=False,
+            use_cache=True,
+            max_new_tokens=128,
+            eos_token_id=32001,
+            pad_token_id=processor.tokenizer.pad_token_id
         )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    history.append([message, ""])
-    partial_response = ""
-    for new_token in streamer:
-        partial_response += new_token
-        history[-1][1] = partial_response
-        yield history, gr.Button(visible=False), gr.Button(visible=True, interactive=True)
 
+    prompt_len = inputs["input_ids"].shape[1]
+    decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
+    
+    # Remove the end-of-sequence token
+    result = decoded_text.replace("<|im_end|>", "").strip()
+    return result
 
+# Create the Gradio interface
 with gr.Blocks() as demo:
-    with gr.Row():
-        image = gr.Image(type="pil")
-
-        with gr.Column():
-            chat = gr.Chatbot(show_label=False)
-            message = gr.Textbox(interactive=True, show_label=False, container=False)
-
-            with gr.Row():
-                gr.ClearButton([chat, message])
-                stop = gr.Button(value="Stop", variant="stop", visible=False)
-                submit = gr.Button(value="Submit", variant="primary")
-
-    with gr.Row():
-        gr.Examples(
-            [
-                ["images/interior.jpg", "Describe the image accurately."],
-                ["images/cat.jpg", "Describe the image in three sentences."],
-                ["images/child.jpg", "Describe the image in one sentence."],
-            ],
-            [image, message],
-            label="Captioning"
-        )
-        gr.Examples(
-            [
-                ["images/scream.jpg", "What is the main emotion of this image?"],
-                ["images/louvre.jpg", "Where is this landmark located?"],
-                ["images/three_people.jpg", "What are these people doing?"]
-            ],
-            [image, message],
-            label="VQA"
-        )
-
-    response_handler = (
-        response,
-        [message, chat, image],
-        [chat, submit, stop]
+    gr.Markdown("# Image Transcription and Visual Question Answering")
+    gr.Markdown("Powered by the unum-cloud/uform-gen2-dpo model.")
+
+    with gr.Tab("Image Transcription"):
+        with gr.Row():
+            transcribe_image_input = gr.Image(type="pil", label="Upload Image")
+            transcribe_output = gr.Textbox(label="Generated Caption")
+        transcribe_button = gr.Button("Generate Caption")
+
+    with gr.Tab("Visual Question Answering"):
+        with gr.Row():
+            vqa_image_input = gr.Image(type="pil", label="Upload Image")
+            with gr.Column():
+                vqa_question_input = gr.Textbox(label="Ask a question")
+                vqa_output = gr.Textbox(label="Answer")
+        vqa_button = gr.Button("Get Answer")
+
+    # Connect the functions to the Gradio components
+    transcribe_button.click(
+        fn=transcribe_image,
+        inputs=[transcribe_image_input],
+        outputs=[transcribe_output]
     )
-    postresponse_handler = (
-        lambda: (gr.Button(visible=False), gr.Button(visible=True)),
-        None,
-        [stop, submit]
+    vqa_button.click(
+        fn=visual_question_answer,
+        inputs=[vqa_image_input, vqa_question_input],
+        outputs=[vqa_output]
     )
 
-    event1 = message.submit(*response_handler)
-    event1.then(*postresponse_handler)
-    event2 = submit.click(*response_handler)
-    event2.then(*postresponse_handler)
-
-    stop.click(None, None, None, cancels=[event1, event2])
-
-demo.queue()
-demo.launch()
+demo.launch()