backend/app/schemas/image.py

from pydantic import BaseModel, Field
from typing import Optional, List


class ImageGenerationRequest(BaseModel):
    """Schema for image generation request"""

    prompt: str = Field(
        ...,
        description="Text prompt for image generation",
        min_length=1,
        max_length=1000,
    )
    size: Optional[str] = Field(
        default="256x256", description="Image size (256x256, 512x512, 1024x1024)"
    )
    n: Optional[int] = Field(
        default=1, ge=1, le=10, description="Number of images to generate"
    )
    model: Optional[str] = Field(
        default="dall-e-3", description="Model to use for image generation"
    )
    reference_image: Optional[str] = Field(
        default=None, description="Base64 encoded reference image (optional)"
    )


class ImageGenerationResponse(BaseModel):
    """Schema for image generation response"""

    success: bool
    message: str
    images: Optional[List[str]] = Field(
        default=None, description="List of base64 encoded generated images"
    )
    count: Optional[int] = None