Update app.py

app.py

@@ -1,107 +1,243 @@
-import argparse
+import os
+import random
+import uuid
+import json
+import time
+import asyncio
+from threading import Thread
+
+import gradio as gr
 import spaces
 import torch
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import numpy as np
+from PIL import Image
+import cv2
+
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+    Qwen2VLForConditionalGeneration,
+    AutoProcessor,
+)
+from transformers.image_utils import load_image
 
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str, default="prithivMLmods/Pocket-Llama-3.2-3B-Instruct")
-    parser.add_argument("--max_length", type=int, default=512)
-    parser.add_argument("--do_sample", action="store_true")
-    # This allows ignoring unrecognized arguments, e.g., from Jupyter
-    return parser.parse_known_args()
+# Constants for text generation
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-def load_model(model_name):
-    """Load model and tokenizer from Hugging Face."""
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        torch_dtype=torch.bfloat16,
-        device_map="auto"
-    )
-    return model, tokenizer
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-def generate_reply(model, tokenizer, prompt, max_length, do_sample):
-    """Generate text from the model given a prompt."""
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    # We’re returning just the final string; no streaming here
-    output_tokens = model.generate(
-        **inputs,
-        max_length=max_length,
-        do_sample=do_sample
-    )
-    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+# Load text-only model and tokenizer
+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()
 
+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("cuda").eval()
 
-def main():
-    args, _ = get_args()
-    model, tokenizer = load_model(args.model)
-    
-    @spaces.GPU
-    def respond(user_message, chat_history):
-        """
-        Gradio expects a function that takes the last user message and the 
-        conversation history, then returns the updated history.
-        
-        chat_history is a list of (user_message, bot_reply) pairs.
-        """
-        # Build a single text prompt from the conversation so far
-        prompt = ""
-        for (old_user_msg, old_bot_msg) in chat_history:
-            prompt += f"User: {old_user_msg}\nBot: {old_bot_msg}\n"
-        # Add the new user query
-        prompt += f"User: {user_message}\nBot:"
+def clean_chat_history(chat_history):
+    """
+    Filter out any chat entries whose "content" is not a string.
+    This helps prevent errors when concatenating previous messages.
+    """
+    cleaned = []
+    for msg in chat_history:
+        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
+            cleaned.append(msg)
+    return cleaned
 
-        # Generate the response
-        bot_message = generate_reply(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=prompt,
-            max_length=args.max_length,
-            do_sample=args.do_sample
-        )
+def downsample_video(video_path):
+    """
+    Downsamples the video to 10 evenly spaced frames.
+    Each frame is returned as a PIL image along with its timestamp.
+    """
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    # Sample 10 evenly spaced frames.
+    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
+    for i in frame_indices:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, image = vidcap.read()
+        if success:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
+            pil_image = Image.fromarray(image)
+            timestamp = round(i / fps, 2)
+            frames.append((pil_image, timestamp))
+    vidcap.release()
+    return frames
 
-        # In many cases, the model output will contain the entire prompt again, 
-        # so we can strip that off or just let it show. If you see repeated 
-        # text, you can try to remove the prompt prefix from bot_message.
-        if bot_message.startswith(prompt):
-            bot_message = bot_message[len(prompt):]
+def progress_bar_html(label: str) -> str:
+    """
+    Returns an HTML snippet for a thin progress bar with a label.
+    The progress bar is styled as a dark red animated bar.
+    """
+    return f'''
+<div style="display: flex; align-items: center;">
+    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
+    <div style="width: 110px; height: 5px; background-color: #FFF0F5; border-radius: 2px; overflow: hidden;">
+        <div style="width: 100%; height: 100%; background-color: #FF69B4; animation: loading 1.5s linear infinite;"></div>
+    </div>
+</div>
+<style>
+@keyframes loading {{
+    0% {{ transform: translateX(-100%); }}
+    100% {{ transform: translateX(100%); }}
+}}
+</style>
+    '''
 
-        # Append the new user-message and bot-response to the history
-        chat_history.append((user_message, bot_message))
-        return chat_history, chat_history
+@spaces.GPU(duration=60, enable_queue=True)
+def generate(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 responses with support for multimodal input and video processing.
+    Special command:
+      - "@video-infer": triggers video processing using Qwen2VL.
+    """
+    text = input_dict["text"]
+    files = input_dict.get("files", [])
+    lower_text = text.strip().lower()
 
-    # Define the Gradio interface
-    with gr.Blocks() as demo:
-        gr.Markdown("<h2 style='text-align: center;'>Chat with Your Model</h2>")
-        
-        # A Chatbot component that will display the conversation
-        chatbot = gr.Chatbot(label="Chat")
-        
-        # A text box for user input
-        user_input = gr.Textbox(
-            show_label=False, 
-            placeholder="Type your message here and press Enter"
-        )
-        
-        # A button to clear the conversation
-        clear_button = gr.Button("Clear")
+    # Branch for video processing with Qwen2VL.
+    if lower_text.startswith("@video-infer"):
+        prompt = text[len("@video-infer"):].strip()
+        if files:
+            # Assume the first file is a video.
+            video_path = files[0]
+            frames = downsample_video(video_path)
+            messages = [
+                {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
+                {"role": "user", "content": [{"type": "text", "text": prompt}]}
+            ]
+            # Append each frame with its timestamp.
+            for frame in frames:
+                image, timestamp = frame
+                image_path = f"video_frame_{uuid.uuid4().hex}.png"
+                image.save(image_path)
+                messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
+                messages[1]["content"].append({"type": "image", "url": image_path})
+        else:
+            messages = [
+                {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
+                {"role": "user", "content": [{"type": "text", "text": prompt}]}
+            ]
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
+        ).to("cuda")
+        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            **inputs,
+            "streamer": streamer,
+            "max_new_tokens": max_new_tokens,
+            "do_sample": True,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repetition_penalty": repetition_penalty,
+        }
+        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        yield progress_bar_html("Processing video with Qwen2VL")
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer
+        return
 
-        # When the user hits Enter in the textbox, call 'respond'
-        #   - Inputs: [user_input, chatbot]  (the last user message and history)
-        #   - Outputs: [chatbot, chatbot]    (updates the chatbot display and history)
-        user_input.submit(respond, [user_input, chatbot], [chatbot, chatbot])
-        
-        # Define a helper function for clearing
-        def clear_conversation():
-            return [], []
-        
-        # When "Clear" is clicked, reset the conversation
-        clear_button.click(fn=clear_conversation, outputs=[chatbot, chatbot])
+    # Normal text or multimodal conversation processing.
+    if files:
+        if len(files) > 1:
+            images = [load_image(image) for image in files]
+        elif len(files) == 1:
+            images = [load_image(files[0])]
+        else:
+            images = []
+        messages = [{
+            "role": "user",
+            "content": [
+                *[{"type": "image", "image": image} for image in images],
+                {"type": "text", "text": text},
+            ]
+        }]
+        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text=[prompt_full], images=images, return_tensors="pt", padding=True).to("cuda")
+        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
+        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        yield progress_bar_html("Thinking...")
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer
+    else:
+        conversation = clean_chat_history(chat_history)
+        conversation.append({"role": "user", "content": text})
+        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:]
+            gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+        input_ids = input_ids.to(model.device)
+        streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            "input_ids": input_ids,
+            "streamer": streamer,
+            "max_new_tokens": max_new_tokens,
+            "do_sample": True,
+            "top_p": top_p,
+            "top_k": top_k,
+            "temperature": temperature,
+            "num_beams": 1,
+            "repetition_penalty": repetition_penalty,
+        }
+        t = Thread(target=model.generate, kwargs=generation_kwargs)
+        t.start()
+        outputs = []
+        yield progress_bar_html("Processing...")
+        for new_text in streamer:
+            outputs.append(new_text)
+            yield "".join(outputs)
+        final_response = "".join(outputs)
+        yield final_response
 
-    # Launch the Gradio app
-    demo.launch()
+demo = gr.ChatInterface(
+    fn=generate,
+    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),
+    ],
+    cache_examples=False,
+    type="messages",
+    fill_height=True,
+    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", "video"], file_count="multiple"),
+    stop_btn="Stop Generation",
+    multimodal=True,
+)
 
 if __name__ == "__main__":
-    main()
+    demo.queue(max_size=20).launch(share=True)