src/models/generate_mask_instruction.py

from pydantic import BaseModel, Field


class GenerateMaskInstruction(BaseModel):
    category: str = Field(
        ...,
        description="The editing category based on the instruction. Must be one of: Addition, Remove, Local, Global, Background.",
    )
    subject: str = Field(
        ...,
        description="The subject of the editing instruction. Must be a noun in no more than 5 words.",
    )
    caption: str = Field(..., description="The detailed description of the image.")
    target_caption: str = Field(
        ...,
        description="Apply the editing instruction to the image caption. The target caption should describe the image after the editing instruction is applied.",
    )