mirror of
https://github.com/open-webui/pipelines
synced 2025-05-10 23:50:45 +00:00
improve: enhance image format detection in process_image method
- Add proper MIME type detection for both data URLs and HTTP requests - Extract media type from Content-Type header or MIME type - Make format detection more robust and generic - Remove hardcoded PNG/JPEG format assumptions
This commit is contained in:
parent
327062733a
commit
488c43edd9
@ -227,17 +227,24 @@ class Pipeline:
|
|||||||
|
|
||||||
def process_image(self, image: str):
|
def process_image(self, image: str):
|
||||||
img_stream = None
|
img_stream = None
|
||||||
if image["url"].startswith("data:image"):
|
content_type = None
|
||||||
if ',' in image["url"]:
|
|
||||||
base64_string = image["url"].split(',')[1]
|
|
||||||
image_data = base64.b64decode(base64_string)
|
|
||||||
|
|
||||||
|
if image["url"].startswith("data:image"):
|
||||||
|
mime_type, base64_string = image["url"].split(",", 1)
|
||||||
|
content_type = mime_type.split(":")[1].split(";")[0]
|
||||||
|
image_data = base64.b64decode(base64_string)
|
||||||
img_stream = BytesIO(image_data)
|
img_stream = BytesIO(image_data)
|
||||||
else:
|
else:
|
||||||
img_stream = requests.get(image["url"]).content
|
response = requests.get(image["url"])
|
||||||
|
img_stream = BytesIO(response.content)
|
||||||
|
content_type = response.headers.get('Content-Type', 'image/jpeg')
|
||||||
|
|
||||||
|
media_type = content_type.split('/')[-1] if '/' in content_type else content_type
|
||||||
return {
|
return {
|
||||||
"image": {"format": "png" if image["url"].endswith(".png") else "jpeg",
|
"image": {
|
||||||
"source": {"bytes": img_stream.read()}}
|
"format": media_type,
|
||||||
|
"source": {"bytes": img_stream.read()}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
def stream_response(self, model_id: str, payload: dict) -> Generator:
|
def stream_response(self, model_id: str, payload: dict) -> Generator:
|
||||||
|
Loading…
Reference in New Issue
Block a user