mirror of
https://github.com/tcsenpai/pensieve.git
synced 2025-06-06 19:25:24 +00:00
49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
from pydantic import BaseModel
|
|
from typing import List, Dict, Any, Optional
|
|
import httpx
|
|
from PIL import Image
|
|
import base64
|
|
import io
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
model: str
|
|
messages: List[Dict[str, Any]]
|
|
max_tokens: Optional[int] = None
|
|
|
|
class ChatCompletionResponse(BaseModel):
|
|
id: str
|
|
object: str
|
|
created: int
|
|
model: str
|
|
choices: List[Dict[str, Any]]
|
|
usage: Dict[str, int]
|
|
|
|
class ModelData(BaseModel):
|
|
id: str
|
|
object: str = "model"
|
|
created: int
|
|
owned_by: str = "transformers"
|
|
root: str = "models"
|
|
parent: Optional[str] = None
|
|
max_model_len: int
|
|
permission: List[Dict[str, Any]]
|
|
|
|
class ModelsResponse(BaseModel):
|
|
object: str = "list"
|
|
data: List[ModelData]
|
|
|
|
async def get_image_from_url(image_url):
|
|
if image_url.startswith("data:image/"):
|
|
image_data = base64.b64decode(image_url.split(",")[1])
|
|
return Image.open(io.BytesIO(image_data))
|
|
elif image_url.startswith("file://"):
|
|
file_path = image_url[len("file://") :]
|
|
return Image.open(file_path)
|
|
else:
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.get(image_url)
|
|
response.raise_for_status()
|
|
image_data = response.content
|
|
return Image.open(io.BytesIO(image_data))
|
|
|