Update app.py

app.py

@@ -1,23 +1,48 @@
 import os
-import time
+from collections.abc import Iterator
 from threading import Thread
 import gradio as gr
+import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from transformers.image_utils import load_image
 import edge_tts
 import asyncio
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
+from transformers.image_utils import load_image
+import time
+
+DESCRIPTION = """
+# QwQ Edge 💬
+"""
 
-# Load models
-MODEL_ID = "prithivMLmods/FastThink-0.5B-Tiny"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16).eval()
+css = '''
+h1 {
+  text-align: center;
+  display: block;
+}
 
-# For multimodal OCR processing
-OCR_MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
-ocr_processor = AutoProcessor.from_pretrained(OCR_MODEL_ID, trust_remote_code=True)
-ocr_model = Qwen2VLForConditionalGeneration.from_pretrained(OCR_MODEL_ID, trust_remote_code=True, torch_dtype=torch.float16).to("cuda").eval()
+#duplicate-button {
+  margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+'''
+
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+model_id = "prithivMLmods/FastThink-0.5B-Tiny"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+model.eval()
 
 TTS_VOICES = [
     "en-US-JennyNeural",  # @tts1
@@ -30,7 +55,14 @@ TTS_VOICES = [
     "en-US-TonyNeural",   # @tts8
 ]
 
-# Handle text-to-speech conversion
+MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" 
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+model_m = Qwen2VLForConditionalGeneration.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to("auto").eval()
+
 async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
     """Convert text to speech using Edge TTS and save as MP3"""
     communicate = edge_tts.Communicate(text, voice)
@@ -39,25 +71,27 @@ async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
 
 @spaces.GPU
 def generate(
-    input_dict, 
-    history, 
-    max_new_tokens: int = 1024, 
-    temperature: float = 0.6, 
-    top_p: float = 0.9, 
-    top_k: int = 50, 
-    repetition_penalty: float = 1.2
+    input_dict: dict,
+    chat_history: list[dict],
+    max_new_tokens: int = 1024,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
 ):
-    """Generates chatbot response and handles TTS requests with multimodal support"""
-    text = input_dict.get("text", "")
+    """Generates chatbot response and handles TTS requests with multimodal input support"""
+    text = input_dict["text"]
     files = input_dict.get("files", [])
-    
-    # Handle multimodal OCR processing
-    if files:
+
+    # Check if input includes image(s)
+    if len(files) > 1:
         images = [load_image(image) for image in files]
+    elif len(files) == 1:
+        images = [load_image(files[0])]
     else:
         images = []
-    
-    # Check if the message is TTS request
+
+    # Check if message is for TTS
     tts_prefix = "@tts"
     is_tts = any(text.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 9))
     voice_index = next((i for i in range(1, 9) if text.strip().lower().startswith(f"{tts_prefix}{i}")), None)
@@ -69,29 +103,36 @@ def generate(
         voice = None
         text = text.replace(tts_prefix, "").strip()
 
-    # If images are provided, combine image and text for the prompt
+    conversation = [*chat_history, {"role": "user", "content": text}]
+
     if images:
-        # Prepare images as part of the conversation
+        # Process multimodal input
         messages = [
-            {
-                "role": "user",
-                "content": [
-                    *[{"type": "image", "image": image} for image in images],
-                    {"type": "text", "text": text},
-                ],
-            }
+            {"role": "user", "content": [
+                *[{"type": "image", "image": image} for image in images],
+                {"type": "text", "text": text},
+            ]}
         ]
-        prompt = ocr_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = ocr_processor(
-            text=[prompt],
-            images=images,
-            return_tensors="pt",
-            padding=True,
-        ).to("cuda")
+        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text=[prompt], images=images, return_tensors="pt", padding=True).to("cuda")
+
+        # Handle generation for multimodal input
+        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+
+        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        buffer = ""
+        yield "Thinking..."
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer
 
     else:
-        # Normal text-only input
-        conversation = [*history, {"role": "user", "content": text}]
+        # Process text-only input
         input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
         if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
             input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -110,38 +151,32 @@ def generate(
             num_beams=1,
             repetition_penalty=repetition_penalty,
         )
-
-        # Start generation in a separate thread
         t = Thread(target=model.generate, kwargs=generate_kwargs)
         t.start()
 
-        # Collect generated text
         outputs = []
         for text in streamer:
             outputs.append(text)
             yield "".join(outputs)
+
         final_response = "".join(outputs)
 
-    # Handle text-to-speech
-    if is_tts and voice:
-        output_file = asyncio.run(text_to_speech(final_response, voice))
-        yield gr.Audio(output_file, autoplay=True)  # Return playable audio
-    else:
-        yield final_response  # Return text response
+        if is_tts and voice:
+            output_file = asyncio.run(text_to_speech(final_response, voice))
+            yield gr.Audio(output_file, autoplay=True)  # Return playable audio
+        else:
+            yield final_response  # Return text response
 
-# Gradio Interface
-demo = gr.Interface(
+
+demo = gr.ChatInterface(
     fn=generate,
-    inputs=[
-        gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"),  # Multimodal input
-        gr.Textbox(label="Chat History", value="", placeholder="Previous conversation history"),
+    additional_inputs=[
         gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
         gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
         gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
         gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
         gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
     ],
-    outputs=["text", "audio"],
     examples=[
         ["@tts1 Who is Nikola Tesla, and why did he die?"],
         ["A train travels 60 kilometers per hour. If it travels for 5 hours, how far will it travel in total?"],
@@ -150,11 +185,15 @@ demo = gr.Interface(
         ["Rewrite the following sentence in passive voice: 'The dog chased the cat.'"],
         ["@tts5 What is the capital of France?"],
     ],
-    stop_btn="Stop Generation",
-    description="QwQ Edge: A Chatbot with Text-to-Speech and Multimodal Support",
+    cache_examples=False,
+    type="messages",
+    description=DESCRIPTION,
     css=css,
     fill_height=True,
+    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"),
+    stop_btn="Stop Generation",
+    multimodal=True,
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(max_size=20).launch()