mirror of
https://github.com/open-webui/pipelines
synced 2025-06-10 16:46:31 +00:00
add style option in prompt
This commit is contained in:
parent
3526fd0368
commit
5e445df828
@ -14,6 +14,8 @@ from typing import List, Union, Generator, Iterator
|
||||
from pydantic import BaseModel
|
||||
from openai import OpenAI
|
||||
import os
|
||||
import re
|
||||
from difflib import get_close_matches
|
||||
|
||||
class Pipeline:
|
||||
class Valves(BaseModel):
|
||||
@ -23,6 +25,12 @@ class Pipeline:
|
||||
self.name = "Recraft AI Pipeline"
|
||||
self.valves = self.Valves(RECRAFT_API_TOKEN=os.getenv("RECRAFT_API_TOKEN", ""))
|
||||
self.client = None
|
||||
self.available_styles = [
|
||||
"realistic_image",
|
||||
"digital_illustration",
|
||||
"vector_illustration",
|
||||
"icon"
|
||||
]
|
||||
|
||||
async def on_startup(self):
|
||||
print(f"on_startup:{__name__}")
|
||||
@ -34,16 +42,41 @@ class Pipeline:
|
||||
async def on_shutdown(self):
|
||||
print(f"on_shutdown:{__name__}")
|
||||
|
||||
def get_closest_style(self, input_style: str) -> str:
|
||||
# Convert input and available styles to lowercase for better matching
|
||||
input_style = input_style.lower()
|
||||
style_map = {s.lower(): s for s in self.available_styles}
|
||||
|
||||
# Try to find close matches
|
||||
matches = get_close_matches(input_style, style_map.keys(), n=1, cutoff=0.6)
|
||||
|
||||
if matches:
|
||||
closest = matches[0]
|
||||
print(f"Using style '{style_map[closest]}' for input '{input_style}'")
|
||||
return style_map[closest]
|
||||
return "realistic_image" # default fallback
|
||||
|
||||
def pipe(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Union[str, Generator, Iterator]:
|
||||
print(f"pipe:{__name__}")
|
||||
|
||||
try:
|
||||
# Extract style from prompt if provided in [style] format
|
||||
style_match = re.search(r'\[(.*?)\]', user_message)
|
||||
selected_style = self.get_closest_style(style_match.group(1)) if style_match else "realistic_image"
|
||||
|
||||
# Clean the prompt by removing the style specification
|
||||
clean_prompt = re.sub(r'\[.*?\]', '', user_message).strip()
|
||||
|
||||
# Select model based on style
|
||||
model = 'recraft20b' if selected_style == 'icon' else 'recraftv3'
|
||||
|
||||
response = self.client.images.generate(
|
||||
prompt=user_message,
|
||||
style='realistic_image',
|
||||
prompt=clean_prompt,
|
||||
style=selected_style,
|
||||
size='1280x1024',
|
||||
model=model,
|
||||
)
|
||||
print(response)
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user