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-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 = "\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=[("", "")], 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.🌠GitHub' ), 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")