from pydantic import BaseModel


class GenerateMultiViewModel(BaseModel):
    tasks_id: str
    image_url: str


class GenerateImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    mode: str
    category: str
    gender: str
    version: str


class GenerateSingleLogoImageModel(BaseModel):
    tasks_id: str
    prompt: str
    seed: str


class GenerateProductImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    image_strength: float
    product_type: str


class GenerateRelightImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    direction: str
    product_type: str


"""
    batch generate image
"""


class BatchGenerateProductImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    image_strength: float
    product_type: str
    batch_size: int


class BatchGenerateRelightImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    direction: str
    product_type: str
    batch_size: int