Update app.py

app.py

@@ -1,109 +1,134 @@
+import time
+from threading import Thread
+from typing import Dict, List
+
+import gradio as gr
+import spaces
+import torch
+from PIL import Image
 from transformers import (
-    MllamaForConditionalGeneration,
     AutoProcessor,
+    MllamaForConditionalGeneration,
     TextIteratorStreamer,
 )
-from PIL import Image
-import requests
-import torch
-from threading import Thread
-import gradio as gr
-from gradio import FileData
-import time
-import spaces
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-ckpt = "toandev/Viet-Receipt-Llama-3.2-11B-Vision-Instruct"
+# Constants
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CHECKPOINT = "toandev/Viet-Receipt-Llama-3.2-11B-Vision-Instruct"
+
+# Model initialization
 model = MllamaForConditionalGeneration.from_pretrained(
-    ckpt, torch_dtype=torch.bfloat16
-).to(device)
-processor = AutoProcessor.from_pretrained(ckpt)
+    CHECKPOINT, torch_dtype=torch.bfloat16
+).to(DEVICE)
+processor = AutoProcessor.from_pretrained(CHECKPOINT)
 
 
-@spaces.GPU
-def bot_streaming(message, history, max_new_tokens=250):
+def process_chat_history(history: List) -> tuple[List[Dict], List[Image.Image]]:
+    """
+    Process chat history to extract messages and images.
 
-    txt = message["text"]
-    ext_buffer = f"{txt}"
+    Args:
+        history: List of chat messages
 
+    Returns:
+        Tuple containing processed messages and images
+    """
     messages = []
     images = []
 
     for i, msg in enumerate(history):
         if isinstance(msg[0], tuple):
-            messages.append(
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": history[i + 1][0]},
-                        {"type": "image"},
-                    ],
-                }
-            )
-            messages.append(
-                {
-                    "role": "assistant",
-                    "content": [{"type": "text", "text": history[i + 1][1]}],
-                }
+            messages.extend(
+                [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": history[i + 1][0]},
+                            {"type": "image"},
+                        ],
+                    },
+                    {
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": history[i + 1][1]}],
+                    },
+                ]
             )
             images.append(Image.open(msg[0][0]).convert("RGB"))
         elif isinstance(history[i - 1], tuple) and isinstance(msg[0], str):
-            # messages are already handled
-            pass
-        elif isinstance(history[i - 1][0], str) and isinstance(
-            msg[0], str
-        ):  # text only turn
-            messages.append(
-                {"role": "user", "content": [{"type": "text", "text": msg[0]}]}
-            )
-            messages.append(
-                {"role": "assistant", "content": [{"type": "text", "text": msg[1]}]}
+            continue
+        elif isinstance(history[i - 1][0], str) and isinstance(msg[0], str):
+            messages.extend(
+                [
+                    {"role": "user", "content": [{"type": "text", "text": msg[0]}]},
+                    {
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": msg[1]}],
+                    },
+                ]
             )
 
-    # add current message
-    if len(message["files"]) == 1:
+    return messages, images
+
 
-        if isinstance(message["files"][0], str):  # examples
-            image = Image.open(message["files"][0]).convert("RGB")
-        else:  # regular input
-            image = Image.open(message["files"][0]["path"]).convert("RGB")
+@spaces.GPU
+def bot_streaming(message: Dict, history: List, max_new_tokens: int = 250) -> str:
+    """
+    Generate streaming responses for the chatbot.
+
+    Args:
+        message: Current message containing text and files
+        history: Chat history
+        max_new_tokens: Maximum number of tokens to generate
+
+    Yields:
+        Generated text buffer
+    """
+    text = message["text"]
+    messages, images = process_chat_history(history)
+
+    # Handle current message
+    if len(message["files"]) == 1:
+        image = (
+            Image.open(message["files"][0])
+            if isinstance(message["files"][0], str)
+            else Image.open(message["files"][0]["path"])
+        ).convert("RGB")
         images.append(image)
         messages.append(
             {
                 "role": "user",
-                "content": [{"type": "text", "text": txt}, {"type": "image"}],
+                "content": [{"type": "text", "text": text}, {"type": "image"}],
             }
         )
     else:
-        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
+        messages.append({"role": "user", "content": [{"type": "text", "text": text}]})
 
+    # Process inputs
     texts = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = (
+        processor(text=texts, images=images, return_tensors="pt")
+        if images
+        else processor(text=texts, return_tensors="pt")
+    ).to(DEVICE)
 
-    if images == []:
-        inputs = processor(text=texts, return_tensors="pt").to(device)
-    else:
-        inputs = processor(text=texts, images=images, return_tensors="pt").to(device)
+    # Setup streaming
     streamer = TextIteratorStreamer(
         processor, skip_special_tokens=True, skip_prompt=True
     )
-
     generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
-    generated_text = ""
 
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-    buffer = ""
 
+    buffer = ""
     for new_text in streamer:
         buffer += new_text
-        generated_text_without_prompt = buffer
         time.sleep(0.01)
         yield buffer
 
 
 demo = gr.ChatInterface(
     fn=bot_streaming,
-    title="Multimodal Llama",
     textbox=gr.MultimodalTextbox(),
     additional_inputs=[
         gr.Slider(
@@ -115,10 +140,10 @@ demo = gr.ChatInterface(
         )
     ],
     cache_examples=False,
-    description="Try Multimodal Llama by Meta with transformers in this demo. Upload an image, and start chatting about it, or simply try one of the examples below. To learn more about Llama Vision, visit [our blog post](https://huggingface.co/blog/llama32). ",
     stop_btn="Stop Generation",
     fill_height=True,
     multimodal=True,
 )
 
-demo.launch(debug=True)
+if __name__ == "__main__":
+    demo.launch(debug=True)