27 lines
704 B
Python
27 lines
704 B
Python
from pydantic import BaseModel
|
|
from typing import Optional
|
|
from enum import Enum
|
|
|
|
class GenerationMode(str, Enum):
|
|
TEXT2IMAGE = "text2image"
|
|
IMAGE2IMAGE = "image2image"
|
|
|
|
class PromptOptimizeRequest(BaseModel):
|
|
prompt: str
|
|
mode: GenerationMode = GenerationMode.TEXT2IMAGE
|
|
original_description: Optional[str] = None
|
|
|
|
class PromptOptimizeResponse(BaseModel):
|
|
optimized_prompt: str
|
|
|
|
class ImageGenerateRequest(BaseModel):
|
|
prompt: str
|
|
optimized_prompt: str
|
|
mode: GenerationMode
|
|
model: str = "gemini-2.5-flash-image"
|
|
input_image_b64: Optional[str] = None # For I2I
|
|
|
|
class ImageGenerateResponse(BaseModel):
|
|
image_url: str
|
|
model: str
|
|
optimized_prompt: str |