diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index fe75d6a976..0fd1905ff4 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -526,7 +526,7 @@ class InferenceBackend: def _generate_vision_response(self, messages, system_prompt, image, temperature, top_p, top_k, max_new_tokens, repetition_penalty) -> Generator[str, None, None]: - """Handle vision model generation.""" + """Handle vision model generation with true token-by-token streaming.""" model_info = self.models[self.active_model_name] model = model_info["model"] processor = model_info["processor"] @@ -565,31 +565,44 @@ class InferenceBackend: formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = processor.tokenizer(formatted_prompt, return_tensors="pt").to(self.device) - # Generate with streaming - captured_output = StringIO() - original_stdout = sys.stdout - + # Stream with TextIteratorStreamer + background thread try: - sys.stdout = captured_output + from transformers import TextIteratorStreamer + import threading - text_streamer = TextStreamer(processor.tokenizer, skip_prompt=True) - model.generate( + streamer = TextIteratorStreamer( + processor.tokenizer, skip_prompt=True, skip_special_tokens=True + ) + + generation_kwargs = dict( **inputs, - streamer=text_streamer, + streamer=streamer, max_new_tokens=max_new_tokens, use_cache=True, temperature=temperature, top_p=top_p, - top_k=top_k + top_k=top_k, ) - sys.stdout = original_stdout - generated_text = captured_output.getvalue() - cleaned = self._clean_generated_text(generated_text) - yield cleaned + def generate_fn(): + try: + model.generate(**generation_kwargs) + except Exception as e: + logger.error(f"Vision generation error in thread: {e}") + + thread = threading.Thread(target=generate_fn) + thread.start() + + output = "" + for new_token in streamer: + if new_token: + output += new_token + cleaned = self._clean_generated_text(output) + yield cleaned + + thread.join() except Exception as e: - sys.stdout = original_stdout logger.error(f"Vision generation error: {e}") yield f"Error: {str(e)}" pass diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 36347d4c13..64791b06f9 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -5,9 +5,9 @@ from __future__ import annotations import time import uuid -from typing import Literal, Optional, List +from typing import Annotated, Literal, Optional, List, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, Discriminator, Field, Tag class LoadRequest(BaseModel): @@ -25,7 +25,7 @@ class UnloadRequest(BaseModel): class GenerateRequest(BaseModel): - """Request for text generation""" + """Request for text generation (legacy /generate/stream endpoint)""" messages: List[dict] = Field(..., description="Chat messages in OpenAI format") system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt") temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature") @@ -64,10 +64,53 @@ class InferenceStatusResponse(BaseModel): # ===================================================================== +# ── Multimodal content parts (OpenAI vision format) ────────────── + +class TextContentPart(BaseModel): + """Text content part in a multimodal message.""" + type: Literal["text"] + text: str + + +class ImageUrl(BaseModel): + """Image URL object — supports data URIs and remote URLs.""" + url: str = Field(..., description="data:image/png;base64,... or https://...") + detail: Optional[Literal["auto", "low", "high"]] = "auto" + + +class ImageContentPart(BaseModel): + """Image content part in a multimodal message.""" + type: Literal["image_url"] + image_url: ImageUrl + + +def _content_part_discriminator(v): + if isinstance(v, dict): + return v.get("type") + return getattr(v, "type", None) + + +ContentPart = Annotated[ + Union[ + Annotated[TextContentPart, Tag("text")], + Annotated[ImageContentPart, Tag("image_url")], + ], + Discriminator(_content_part_discriminator), +] +"""Union type for multimodal content parts, discriminated by the 'type' field.""" + + +# ── Messages ───────────────────────────────────────────────────── + class ChatMessage(BaseModel): - """A single message in the conversation.""" + """ + A single message in the conversation. + + ``content`` may be a plain string (text-only) or a list of + content parts for multimodal messages (OpenAI vision format). + """ role: Literal["system", "user", "assistant"] = Field(..., description="Message role") - content: str = Field(..., description="Message content") + content: Union[str, list[ContentPart]] = Field(..., description="Message content (string or multimodal parts)") class ChatCompletionRequest(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e98165d83e..74cf37138f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -240,6 +240,59 @@ async def get_status(): # ===================================================================== +def _extract_content_parts( + messages: list, +) -> tuple[str, list[dict], "Optional[str]"]: + """ + Parse OpenAI-format messages into components the inference backend expects. + + Handles both plain-string ``content`` and multimodal content-part arrays + (``[{type: "text", ...}, {type: "image_url", ...}]``). + + Returns: + system_prompt: The system message text (or a default). + chat_messages: Non-system messages with content flattened to strings. + image_base64: Base64 data of the *first* image found, or ``None``. + """ + system_prompt = "You are a helpful AI assistant." + chat_messages: list[dict] = [] + first_image_b64: Optional[str] = None + + for msg in messages: + # ── System messages → extract as system_prompt ──────── + if msg.role == "system": + if isinstance(msg.content, str): + system_prompt = msg.content + elif isinstance(msg.content, list): + # Unlikely but handle: join text parts + system_prompt = "\n".join( + p.text for p in msg.content if p.type == "text" + ) + continue + + # ── User / assistant messages ───────────────────────── + if isinstance(msg.content, str): + # Plain string content — pass through + chat_messages.append({"role": msg.role, "content": msg.content}) + elif isinstance(msg.content, list): + # Multimodal content parts + text_parts: list[str] = [] + for part in msg.content: + if part.type == "text": + text_parts.append(part.text) + elif part.type == "image_url" and first_image_b64 is None: + url = part.image_url.url + if url.startswith("data:"): + # data:image/png;base64, → extract + first_image_b64 = url.split(",", 1)[1] if "," in url else None + else: + logger.warning( + f"Remote image URLs not yet supported: {url[:80]}..." + ) + combined_text = "\n".join(text_parts) if text_parts else "" + chat_messages.append({"role": msg.role, "content": combined_text}) + + return system_prompt, chat_messages, first_image_b64 @router.post("/chat/completions") @@ -247,6 +300,9 @@ async def openai_chat_completions(request: ChatCompletionRequest): """ OpenAI-compatible chat completions endpoint. + Supports multimodal messages: ``content`` may be a plain string or a + list of content parts (``text`` / ``image_url``). + Streaming (default): returns SSE chunks matching OpenAI's format. Non-streaming: returns a single ChatCompletion JSON object. """ @@ -258,15 +314,10 @@ async def openai_chat_completions(request: ChatCompletionRequest): detail="No model loaded. Call POST /inference/load first.", ) - # ── Extract system prompt from messages ─────────────────────── - system_prompt = "You are a helpful AI assistant." - chat_messages: list[dict] = [] - - for msg in request.messages: - if msg.role == "system": - system_prompt = msg.content - else: - chat_messages.append({"role": msg.role, "content": msg.content}) + # ── Parse messages (handles multimodal content parts) ───── + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( + request.messages + ) # If no non-system messages were provided, error out if not chat_messages: @@ -275,9 +326,12 @@ async def openai_chat_completions(request: ChatCompletionRequest): detail="At least one non-system message is required.", ) - # ── Decode image if provided (vision models) ────────────────── + # ── Decode image (from content parts OR legacy field) ───── + # Content-part images take priority; fall back to legacy field + image_b64 = extracted_image_b64 or request.image_base64 image = None - if request.image_base64: + + if image_b64: try: import base64 from PIL import Image @@ -287,10 +341,10 @@ async def openai_chat_completions(request: ChatCompletionRequest): if not model_info.get("is_vision"): raise HTTPException( status_code=400, - detail="Image provided but current model is text-only.", + detail="Image provided but current model is text-only. Load a vision model.", ) - image_data = base64.b64decode(request.image_base64) + image_data = base64.b64decode(image_b64) image = Image.open(BytesIO(image_data)) image = backend.resize_image(image) @@ -299,7 +353,7 @@ async def openai_chat_completions(request: ChatCompletionRequest): except Exception as e: raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}") - # ── Shared generation kwargs ────────────────────────────────── + # ── Shared generation kwargs ────────────────────────────── gen_kwargs = dict( messages=chat_messages, system_prompt=system_prompt,