app.py

import logging
import time
from queue import Empty
from threading import Event, Thread

import gradio as gr
import spaces
import torch
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    StoppingCriteria,
    StoppingCriteriaList,
    TextIteratorStreamer,
)

MODEL_NAME = "Azure99/Blossom-V7.1-35B-A3B"
MAX_INPUT_TOKENS = 32768
MAX_IMAGE_PIXELS = 1024 * 1024
DEFAULT_MAX_NEW_TOKENS = 4096
MAX_NEW_TOKENS = 32768
DEFAULT_GPU_DURATION_SECONDS = 60
STREAM_TIMEOUT_SECONDS = 90
STREAM_INTERVAL_SECONDS = 0.05

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

processor = AutoProcessor.from_pretrained(MODEL_NAME)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_NAME,
    dtype=torch.bfloat16,
    device_map="auto",
    use_kernels=True,
)
model.eval()


class CancelStoppingCriteria(StoppingCriteria):
    def __init__(self, cancel_event):
        self.cancel_event = cancel_event

    def __call__(self, input_ids, scores, **kwargs):
        return torch.full(
            (input_ids.shape[0],),
            self.cancel_event.is_set(),
            dtype=torch.bool,
            device=input_ids.device,
        )


def content_text(content):
    if isinstance(content, str):
        return content

    parts = []
    for item in content:
        if isinstance(item, str):
            parts.append(item)
        elif item.get("type") == "text":
            parts.append(item["text"])
    return "".join(parts)


def content_blocks(content):
    if isinstance(content, str):
        return [{"type": "text", "text": content}] if content else []

    blocks = []
    for item in content:
        if isinstance(item, str):
            if item:
                blocks.append({"type": "text", "text": item})
        elif item.get("type") == "text":
            blocks.append({"type": "text", "text": item["text"]})
        elif item.get("type") == "file":
            blocks.append({"type": "image", "path": item["file"]["path"]})
        elif "path" in item:
            blocks.append({"type": "image", "path": item["path"]})
    return blocks


def append_user_message(messages, content):
    if messages and messages[-1]["role"] == "user":
        content = messages[-1]["content"] + content
        messages.pop()

    images = [item for item in content if item["type"] == "image"]
    text = [item for item in content if item["type"] == "text"]
    messages.append({"role": "user", "content": images + text})


def get_messages(user, history):
    messages = []
    pending_reasoning = []

    for message in history or []:
        role = message["role"]
        content = content_text(message["content"])
        metadata = message.get("metadata") or {}

        if role == "assistant" and metadata.get("title") == "Reasoning":
            pending_reasoning.append(content)
            continue

        if pending_reasoning:
            messages.append(
                {
                    "role": "assistant",
                    "reasoning": "\n\n".join(pending_reasoning),
                    "content": content if role == "assistant" else "",
                }
            )
            pending_reasoning.clear()
            if role == "assistant":
                continue

        if role == "user":
            append_user_message(messages, content_blocks(message["content"]))
        elif role == "assistant":
            messages.append({"role": role, "content": content})

    if pending_reasoning:
        messages.append(
            {
                "role": "assistant",
                "reasoning": "\n\n".join(pending_reasoning),
                "content": "",
            }
        )

    content = [{"type": "image", "path": path} for path in user["files"]]
    if user["text"]:
        content.append({"type": "text", "text": user["text"]})
    append_user_message(messages, content)
    return messages


def get_gpu_duration(user, history, max_new_tokens, gpu_duration_seconds):
    return int(gpu_duration_seconds)


@spaces.GPU(size="xlarge", duration=get_gpu_duration)
def chat(user, history, max_new_tokens, gpu_duration_seconds):
    if not user["text"].strip() and not user["files"]:
        raise gr.Error("Please enter a message.")

    messages = get_messages(user, history)
    inputs = processor.apply_chat_template(
        messages,
        chat_template=processor.tokenizer.chat_template,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        processor_kwargs={"images_kwargs": {"max_pixels": MAX_IMAGE_PIXELS}},
    )
    input_tokens = inputs["input_ids"].shape[-1]
    if input_tokens > MAX_INPUT_TOKENS:
        raise gr.Error(
            f"This conversation is {input_tokens:,} tokens. "
            f"The demo limit is {MAX_INPUT_TOKENS:,} tokens."
        )
    inputs = inputs.to(model.device)

    cancel_event = Event()
    generation_errors = []
    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        timeout=STREAM_TIMEOUT_SECONDS,
    )
    generation_kwargs = {
        **inputs,
        "streamer": streamer,
        "max_new_tokens": int(max_new_tokens),
        "stopping_criteria": StoppingCriteriaList(
            [CancelStoppingCriteria(cancel_event)]
        ),
    }

    def generate():
        try:
            with torch.inference_mode():
                model.generate(**generation_kwargs)
        except Exception as error:
            generation_errors.append(error)
            logger.exception("Generation failed")
            streamer.on_finalized_text("", stream_end=True)

    thread = Thread(target=generate, daemon=True)
    thread.start()

    output = "<think>\n"
    last_yield = time.monotonic()
    timed_out = False

    try:
        yield output
        for new_text in streamer:
            output += new_text
            now = time.monotonic()
            if now - last_yield >= STREAM_INTERVAL_SECONDS:
                yield output
                last_yield = now
    except Empty:
        timed_out = True
    finally:
        cancel_event.set()
        thread.join()

    if generation_errors:
        raise gr.Error(f"Generation failed: {generation_errors[0]}")
    if timed_out:
        raise gr.Error("Generation timed out.")
    yield output


gr.ChatInterface(
    chat,
    multimodal=True,
    chatbot=gr.Chatbot(
        show_label=False,
        height=500,
        buttons=["copy"],
        render_markdown=True,
        reasoning_tags=[("<think>", "</think>")],
        latex_delimiters=[{"left": "\\[", "right": "\\]", "display": True}],
    ),
    textbox=gr.MultimodalTextbox(
        placeholder="",
        container=False,
        scale=7,
        file_types=["image"],
        file_count="single",
        submit_btn=True,
        stop_btn=True,
    ),
    title=f"{MODEL_NAME} Demo",
    description=(
        "Hello, I am Blossom, an open source conversational large language "
        'model.🌠<a href="https://github.com/Azure99/BlossomLM">GitHub</a>'
    ),
    additional_inputs=[
        gr.Slider(
            minimum=1,
            maximum=MAX_NEW_TOKENS,
            value=DEFAULT_MAX_NEW_TOKENS,
            step=1,
            label="Max Output Tokens",
        ),
        gr.Slider(
            minimum=30,
            maximum=240,
            value=DEFAULT_GPU_DURATION_SECONDS,
            step=30,
            label="ZeroGPU Duration (seconds)",
        ),
    ],
    additional_inputs_accordion=gr.Accordion(label="Config", open=True),
    examples=[
        ["Hello"],
        ["What is MBTI"],
        ["用Python实现二分查找"],
        ["为switch写一篇小红书种草文案，带上emoji"],
    ],
    cache_examples=False,
).queue().launch(theme="soft")