mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +00:00
fix _process_image_response
This commit is contained in:
parent
f5996b2f6b
commit
f4ae14a490
1 changed files with 12 additions and 15 deletions
|
@ -23,24 +23,21 @@ def _process_image_response(response: Response, url: str) -> str:
|
|||
|
||||
image_bytes = response.content
|
||||
base64_image = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
image_type = response.headers.get("Content-Type")
|
||||
if image_type is None:
|
||||
img_type = url.split(".")[-1].lower()
|
||||
_img_type = {
|
||||
"jpg": "image/jpeg",
|
||||
"jpeg": "image/jpeg",
|
||||
"png": "image/png",
|
||||
"gif": "image/gif",
|
||||
"webp": "image/webp",
|
||||
}.get(img_type)
|
||||
if _img_type is None:
|
||||
img_type_map = {
|
||||
"jpg": "image/jpeg",
|
||||
"jpeg": "image/jpeg",
|
||||
"png": "image/png",
|
||||
"gif": "image/gif",
|
||||
"webp": "image/webp",
|
||||
}
|
||||
img_type = response.headers.get("Content-Type")
|
||||
if img_type not in img_type_map:
|
||||
_img_type = url.split('?')[0].split(".")[-1].lower()
|
||||
img_type = img_type_map.get(img_type)
|
||||
if img_type is None:
|
||||
raise Exception(
|
||||
f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']"
|
||||
)
|
||||
img_type = _img_type
|
||||
else:
|
||||
img_type = image_type
|
||||
|
||||
result = f"data:{img_type};base64,{base64_image}"
|
||||
in_memory_cache.set_cache(url, result)
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue