AiDA_Python/app/schemas/generate_image.py

from typing import List, Optional

from pydantic import BaseModel, Field


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


class GenerateImageFlux2KleinModel(BaseModel):
    bucket_name: str = Field(..., description="OSS桶名，不传则为None")
    object_name: str = Field(..., description="OSS对象名（文件路径），不传则为None")
    # input_image_paths: Optional[List[str]] = Field(default=[], description="输入图片路径列表")
    width: Optional[int] = Field(default=1024, description="图片宽度，默认512像素")
    height: Optional[int] = Field(default=1024, description="图片高度，默认512像素")
    prompt: Optional[str] = Field(default="", description="文本提示词，用于模型推理等场景")
    steps: Optional[int] = Field(default=4, description="推理步数，控制模型生成过程的迭代次数")
    guidance: Optional[float] = Field(default=4.0, description="引导系数，调节提示词对生成结果的影响程度")


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
"""


# product任务子项
class ProductItemModel(BaseModel):
    tasks_id: str
    image_strength: float
    prompt: str
    image_url: str
    product_type: str


# product批处理 集合
class BatchGenerateProductImageModel(BaseModel):
    batch_tasks_id: str
    user_id: str
    batch_data_list: List[ProductItemModel]


# relight任务子项
class RelightItemModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    direction: str
    product_type: str


# relight批处理集合
class BatchGenerateRelightImageModel(BaseModel):
    batch_tasks_id: str
    user_id: str
    batch_data_list: List[RelightItemModel]


"""
    agent tool generate image
"""


class AgentTollGenerateImageModel(BaseModel):
    prompt: str
    category: str
    gender: str
    version: str
    size: int