TTI / models.py
Sam3838's picture
Upload 4 files
3b9f744 verified
from pydantic import BaseModel, Field
from typing import List, Optional, Literal
from enum import Enum
class ImageSize(str, Enum):
"""Supported image sizes (OpenAI compatible)"""
SMALL = "256x256"
MEDIUM = "512x512"
LARGE = "1024x1024"
WIDE = "1792x1024"
TALL = "1024x1792"
class ImageQuality(str, Enum):
"""Image quality options"""
STANDARD = "standard"
HD = "hd"
class ImageStyle(str, Enum):
"""Image style options"""
VIVID = "vivid"
NATURAL = "natural"
class ResponseFormat(str, Enum):
"""Response format options"""
URL = "url"
B64_JSON = "b64_json"
class ImageGenerationRequest(BaseModel):
"""OpenAI compatible image generation request"""
prompt: str = Field(..., description="A text description of the desired image(s)")
model: str = Field(default="dall-e-3", description="The model to use for image generation")
n: int = Field(default=1, ge=1, le=10, description="Number of images to generate")
quality: ImageQuality = Field(default=ImageQuality.STANDARD, description="Quality of the image")
response_format: ResponseFormat = Field(default=ResponseFormat.URL, description="Response format")
size: ImageSize = Field(default=ImageSize.LARGE, description="Size of the generated images")
style: ImageStyle = Field(default=ImageStyle.VIVID, description="Style of the generated images")
user: Optional[str] = Field(default=None, description="A unique identifier representing your end-user")
class ImageData(BaseModel):
"""Individual image data in response"""
url: Optional[str] = Field(default=None, description="URL of the generated image")
b64_json: Optional[str] = Field(default=None, description="Base64 encoded image data")
revised_prompt: Optional[str] = Field(default=None, description="The revised prompt used for generation")
class ImageGenerationResponse(BaseModel):
"""OpenAI compatible image generation response"""
created: int = Field(..., description="Unix timestamp of when the image was created")
data: List[ImageData] = Field(..., description="List of generated images")
class ErrorResponse(BaseModel):
"""Error response format"""
error: dict = Field(..., description="Error details")
class ModelInfo(BaseModel):
"""Model information"""
id: str
object: str = "model"
created: int
owned_by: str
class ModelsResponse(BaseModel):
"""Models list response"""
object: str = "list"
data: List[ModelInfo]