昨天看到一位佬友发布的函数，尝试了下不知道为什么没有成功，所以也去用ai糊了一个，可以正常使用，分享给大家，有需要的可以体验一下
佬友帖子：[ open-webui集成x.ai新出的grok-2-image模型生图，150刀有去处啦]
使用方式：
open-webui管理员面板——函数，添加以下代码，保存后填写xAI api并开启函数即可。默认模型列表会出现GROK/Image Generator，前缀可在函数参数处修改，可以用于生图
若本地未开代理图片无法显示，可在函数设置里把Response Format从url改为b64_json
欢迎大家体验，大佬们有更好的方法也请分享下
from pydantic import BaseModel, Field
import requests
import base64
from typing import Optional, List
import json
import asyncio
import time
class Pipe:
"""
A Pipe function for Open WebUI that handles image generation using the grok-2-image model.
Includes built-in prompt enhancement and streaming updates.
"""
class Valves(BaseModel):
"""Configuration options for the Grok Image Pipe"""
XAI_API_KEY: str = Field(
default="",
description="API key for authenticating requests to the xAI API"
)
XAI_API_BASE_URL: str = Field(
default="https://api.x.ai/v1",
description="Base URL for accessing xAI API endpoints"
)
NAME_PREFIX: str = Field(
default="GROK/",
description="Prefix to be added before model names"
)
DEFAULT_NUM_IMAGES: int = Field(
default=1,
description="Default number of images to generate (1-10)"
)
RESPONSE_FORMAT: str = Field(
default="url",
description="Response format for generated images ('url' or 'b64_json')"
)
DISPLAY_REVISED_PROMPT: bool = Field(
default=True,
description="Whether to display the revised prompt in the response"
)
ENHANCE_PROMPTS: bool = Field(
default=True,
description="Whether to enhance prompts before sending to the API"
)
ADD_DETAIL_LEVEL: str = Field(
default="medium",
description="Level of detail enhancement (none, low, medium, high)"
)
STYLE_PRESET: str = Field(
default="",
description="Style preset to apply to all prompts (leave empty for default)"
)
SHOW_PROGRESS_UPDATES: bool = Field(
default=True,
description="Whether to show progress updates during image generation"
)
def __init__(self):
"""Initialize the Pipe instance with default Valves settings"""
self.valves = self.Valves()
def pipes(self):
"""
Define the custom "models" this pipe provides.
Returns a list of models that will appear in the model selector.
"""
if not self.valves.XAI_API_KEY:
return [
{
"id": "grok-2-image",
"name": f"{self.valves.NAME_PREFIX}Image Generator (API Key not provided)",
}
]
# Define the image generation model
return [
{
"id": "grok-2-image",
"name": f"{self.valves.NAME_PREFIX}Image Generator",
}
]
def _enhance_prompt(self, prompt: str) -> str:
"""
Enhance the image generation prompt with additional details and context
Args:
prompt: The original user prompt
Returns:
Enhanced prompt with additional details and context
"""
# Return original if enhancement is disabled
if not self.valves.ENHANCE_PROMPTS:
return prompt
# Add style preset if specified
if self.valves.STYLE_PRESET:
if self.valves.STYLE_PRESET.lower() not in prompt.lower():
prompt = f"{prompt}, {self.valves.STYLE_PRESET} style"
# Add detail level enhancements
detail_level = self.valves.ADD_DETAIL_LEVEL.lower()
if detail_level == "low":
# Add minimal enhancements
if "resolution" not in prompt.lower() and "quality" not in prompt.lower():
prompt = f"{prompt}, high quality"
elif detail_level == "medium":
# Add moderate level of detail
details = []
if "lighting" not in prompt.lower():
details.append("detailed lighting")
if "resolution" not in prompt.lower() and "quality" not in prompt.lower():
details.append("high resolution")
if details:
prompt = f"{prompt}, {', '.join(details)}"
elif detail_level == "high":
# Add comprehensive details for better generation
details = []
if "lighting" not in prompt.lower():
details.append("professional lighting")
if "resolution" not in prompt.lower():
details.append("ultra high resolution")
if "detailed" not in prompt.lower():
details.append("highly detailed")
if "composition" not in prompt.lower():
details.append("perfect composition")
if details:
prompt = f"{prompt}, {', '.join(details)}"
return prompt
def _create_progress_update(self, message, index=0):
"""
Create a streaming progress update message for Open WebUI
Args:
message: The progress message to display
index: Choice index
Returns:
Formatted streaming event for Open WebUI
"""
return {
"id": f"progress-{time.time()}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": "grok-2-image",
"choices": [
{
"index": index,
"delta": {
"role": "assistant",
"content": f"{message}\n\n"
},
"finish_reason": None
}
]
}
async def _stream_progress(self, body, original_prompt, enhanced_prompt):
"""
Generate streaming progress updates during image generation
Args:
body: The original request body
original_prompt: The user's original prompt
enhanced_prompt: The enhanced prompt (if enhancement is enabled)
Yields:
Streaming progress update events
"""
# Initial message
yield self._create_progress_update("🎨 **Starting image generation process...**")
await asyncio.sleep(0.5)
# Show prompts
yield self._create_progress_update(f"📝 **Original prompt:** {original_prompt}")
await asyncio.sleep(0.8)
if self.valves.ENHANCE_PROMPTS and enhanced_prompt != original_prompt:
yield self._create_progress_update(f"✨ **Enhanced prompt:** {enhanced_prompt}")
await asyncio.sleep(0.8)
# Generation status
yield self._create_progress_update("🔄 **Sending request to xAI API...**")
await asyncio.sleep(1.0)
yield self._create_progress_update("⏳ **Waiting for image generation...**")
await asyncio.sleep(0.5)
num_images = body.get("n", self.valves.DEFAULT_NUM_IMAGES)
if num_images > 1:
yield self._create_progress_update(f"🖼️ **Generating {num_images} images...**")
else:
yield self._create_progress_update("🖼️ **Generating your image...**")
def pipe(self, body: dict, __user__: Optional[dict] = None):
"""
Main pipe function that processes the request and returns the generated image(s)
Args:
body: Input data containing the prompt and other parameters
__user__: User information (optional)
Returns:
A response containing the generated image(s) or an error message
"""
print(f"Processing image generation request with body: {body}")
# Check if API key is provided
if not self.valves.XAI_API_KEY:
return "Error: xAI API key not provided. Please configure the API key in the pipe settings."
# Extract parameters from the body
messages = body.get("messages", [])
# Get the last user message as the prompt
original_prompt = None
for message in reversed(messages):
if message.get("role") == "user":
original_prompt = message.get("content", "")
break
if not original_prompt:
return "Error: No prompt found in the messages. Please provide a description of the image you want to generate."
# Enhance the prompt if enabled
enhanced_prompt = self._enhance_prompt(original_prompt) if self.valves.ENHANCE_PROMPTS else original_prompt
# Extract other parameters or use defaults
num_images = body.get("n", self.valves.DEFAULT_NUM_IMAGES)
response_format = body.get("response_format", self.valves.RESPONSE_FORMAT)
# Ensure parameters are within valid ranges
num_images = max(1, min(10, num_images)) # Limit to 1-10 images
if response_format not in ["url", "b64_json"]:
response_format = "url" # Default to URL if invalid format
# If streaming is enabled, return a generator for progress updates
if body.get("stream", False) and self.valves.SHOW_PROGRESS_UPDATES:
import asyncio
async def stream_response():
# Stream progress updates first
async for update in self._stream_progress(body, original_prompt, enhanced_prompt):
yield update
# Set up the request to the xAI API
headers = {
"Authorization": f"Bearer {self.valves.XAI_API_KEY}",
"Content-Type": "application/json",
}
payload = {
"model": "grok-2-image",
"prompt": enhanced_prompt,
"n": num_images,
"response_format": response_format
}
try:
# Yield a message that we're making the API call
yield self._create_progress_update("📡 **Making API request...**")
# Make the API request to generate images (not in async context)
response = requests.post(
f"{self.valves.XAI_API_BASE_URL}/images/generations",
headers=headers,
json=payload
)
# Check for errors in the response
response.raise_for_status()
result = response.json()
# Yield the final result with images
yield self._create_progress_update("✅ **Image generation complete!**")
await asyncio.sleep(0.5)
# Format the response as a streaming completion
formatted_response = self._format_response(result, response_format)
final_content = formatted_response["choices"][0]["message"]["content"]
# Send the final content as a delta
yield {
"id": formatted_response["id"],
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": "grok-2-image",
"choices": [
{
"index": 0,
"delta": {
"content": final_content
},
"finish_reason": "stop"
}
]
}
# Signal completion
yield {
"id": formatted_response["id"],
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": "grok-2-image",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
}
except requests.exceptions.RequestException as e:
error_message = f"Error making request to xAI API: {str(e)}"
if hasattr(e, "response") and e.response is not None:
try:
error_details = e.response.json()
error_message += f"\nDetails: {json.dumps(error_details, indent=2)}"
except:
error_message += f"\nStatus code: {e.response.status_code}"
# Yield the error message
yield self._create_progress_update(f"❌ **Error:** {error_message}")
# Signal completion
yield {
"id": f"error-{time.time()}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": "grok-2-image",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
}
# Return the generator
return stream_response()
else:
# Non-streaming response
# Set up the request to the xAI API
headers = {
"Authorization": f"Bearer {self.valves.XAI_API_KEY}",
"Content-Type": "application/json",
}
payload = {
"model": "grok-2-image",
"prompt": enhanced_prompt,
"n": num_images,
"response_format": response_format
}
try:
# Make the API request to generate images
response = requests.post(
f"{self.valves.XAI_API_BASE_URL}/images/generations",
headers=headers,
json=payload
)
# Check for errors in the response
response.raise_for_status()
result = response.json()
# Format the response to display in Open WebUI
formatted_response = self._format_response(result, response_format)
return formatted_response
except requests.exceptions.RequestException as e:
error_message = f"Error making request to xAI API: {str(e)}"
if hasattr(e, "response") and e.response is not None:
try:
error_details = e.response.json()
error_message += f"\nDetails: {json.dumps(error_details, indent=2)}"
except:
error_message += f"\nStatus code: {e.response.status_code}"
return error_message
def _format_response(self, result, response_format):
"""
Format the xAI API response for display in Open WebUI
Args:
result: The raw API response
response_format: The requested response format ('url' or 'b64_json')
Returns:
Formatted response for Open WebUI
"""
images_data = result.get("data", [])
if not images_data:
return "No images were generated."
formatted_content = ""
for i, image_data in enumerate(images_data):
# Add image number if multiple images
if len(images_data) > 1:
formatted_content += f"### Image {i+1}\n\n"
# Add revised prompt if available and enabled
if self.valves.DISPLAY_REVISED_PROMPT and "revised_prompt" in image_data:
formatted_content += f"**Revised Prompt:** {image_data['revised_prompt']}\n\n"
# Add image based on response format
if response_format == "url" and "url" in image_data:
formatted_content += f"\n\n"
elif response_format == "b64_json" and "b64_json" in image_data:
# Convert base64 to a data URL for display
img_data = image_data['b64_json']
if not img_data.startswith("data:"):
img_data = f"data:image/jpeg;base64,{img_data}"
formatted_content += f"\n\n"
formatted_content += "---\n\n"
# Wrap the response in the expected format for Open WebUI
return {
"id": f"grok-image-{time.time()}",
"object": "chat.completion",
"created": int(time.time()),
"model": "grok-2-image",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": formatted_content.strip()
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0
}
}
实际体验：
USER
卡通版兔子战士穿着军装、扛着枪在战场上冒着炮火向前突进
ASSISTANT
Starting image generation process…
Original prompt: 卡通版兔子战士穿着军装、扛着枪在战场上冒着炮火向前突进
Enhanced prompt: 卡通版兔子战士穿着军装、扛着枪在战场上冒着炮火向前突进, detailed lighting, high resolution
Sending request to xAI API…
Waiting for image generation…
Generating 2 images…
Making API request…
Image generation complete!
Image 1
Revised Prompt: A high-resolution cartoon illustration of a rabbit soldier wearing detailed military camouflage gear and holding a rifle, set in a daytime battlefield. The rabbit, with a focused expression, is centrally positioned and advancing forward amidst a war-torn landscape with smoke and subtle battlefield debris. The background features a slightly overcast sky and distant, blurred battlefield elements to maintain focus on the rabbit soldier, enhancing the scene’s realism and intensity without overwhelming the main subject.
Image 2
Revised Prompt: A high-resolution cartoon illustration of a rabbit soldier wearing detailed military camouflage gear and holding a rifle, set in a daytime battlefield. The rabbit, with a focused expression, is centrally positioned and advancing forward amidst a war-torn landscape with smoke and subtle battlefield debris. The background features a slightly overcast sky and distant, blurred battlefield elements to maintain focus on the rabbit soldier, enhancing the scene’s realism and intensity without overwhelming the main subject.
