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