app.py

from transformers import MllamaForConditionalGeneration, AutoProcessor, 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

ckpt = "Daemontatox/DocumentCogito"
model = MllamaForConditionalGeneration.from_pretrained(ckpt,
    torch_dtype=torch.bfloat16).to("cuda")
processor = AutoProcessor.from_pretrained(ckpt)

SYSTEM_PROMPT = """You are a helpful AI assistant specialized in analyzing documents, images, and visual content. 
Your responses should be clear, accurate, and focused on the specific details present in the provided materials. 
When analyzing documents, pay attention to key information, formatting, and context. 
For images, consider both obvious and subtle details that might be relevant to the user's query."""

@spaces.GPU()
def bot_streaming(message, history, max_new_tokens=2048, temperature=0.7):
    txt = message["text"]
    ext_buffer = f"{txt}"
    
    messages = [{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}]
    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]}]})
            images.append(Image.open(msg[0][0]).convert("RGB"))
        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
            pass
        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str):
            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})

    if len(message["files"]) == 1:
        if isinstance(message["files"][0], str):
            image = Image.open(message["files"][0]).convert("RGB")
        else:
            image = Image.open(message["files"][0]["path"]).convert("RGB")
        images.append(image)
        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
    else:
        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})

    texts = processor.apply_chat_template(messages, add_generation_prompt=True)

    if images == []:
        inputs = processor(text=texts, return_tensors="pt").to("cuda")
    else:
        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)

    generation_kwargs = dict(
        inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        temperature=temperature,  # Add temperature parameter
        do_sample=True,  # Enable sampling for temperature to take effect
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    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="Document Analyzer", 
    examples=[
        [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]}, 200, 0.7],
        [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]}, 250, 0.7],
        [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]}, 250, 0.7],
        [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]}, 250, 0.7],
        [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 250, 0.7],
    ],
    textbox=gr.MultimodalTextbox(), 
    additional_inputs=[
        gr.Slider(
            minimum=10,
            maximum=500,
            value=2048,
            step=10,
            label="Maximum number of new tokens to generate",
        ),
        gr.Slider(  # Add temperature slider
            minimum=0.1,
            maximum=2.0,
            value=0.2,
            step=0.1,
            label="Temperature (0.1 = focused, 2.0 = creative)",
        )
    ],
    cache_examples=False,
    description="MllM with Temperature Control",
    stop_btn="Stop Generation", 
    fill_height=True,
    multimodal=True
)

demo.launch(debug=True)