studio/contracts.py

from pathlib import Path
from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator


class Request(BaseModel):
    image: Path | None = None
    prompt: str = Field(default="", max_length=4000)
    animation_prompt: str = Field(default="", max_length=2000)
    animation_seconds: float = Field(default=6, ge=2, le=12)
    seed: int = Field(default=0, ge=0, le=2**31 - 1)
    parts: int = Field(default=8, ge=2, le=32)

    @model_validator(mode="after")
    def check_input(self):
        self.prompt = self.prompt.strip()
        self.animation_prompt = self.animation_prompt.strip()
        if not self.prompt and self.image is None:
            raise ValueError("Upload an image, enter a prompt, or supply both.")
        if self.image is not None and not self.image.is_file():
            raise ValueError("The uploaded image is no longer available.")
        return self


ColorGain = Annotated[float, Field(ge=0.75, le=1.25)]


class PartEdit(BaseModel):
    model_config = ConfigDict(extra="forbid")
    part_id: str
    name: str = Field(min_length=1, max_length=80)
    roughness: float | None = Field(default=None, ge=0, le=1)
    metallic: float | None = Field(default=None, ge=0, le=1)
    color_gain: tuple[ColorGain, ColorGain, ColorGain] = (1, 1, 1)


class Refinement(BaseModel):
    model_config = ConfigDict(extra="forbid")
    summary: str = Field(max_length=4000)
    parts: list[PartEdit] = Field(default_factory=list, max_length=32)
    project_reference_texture: bool = False
    reference_front_axis: Literal["-Y", "+Y", "-X", "+X"] = "-Y"
    projection_strength: float = Field(default=0.9, ge=0, le=1)
    remove_degenerate_faces: bool = True
    remove_duplicate_faces: bool = True
    unresolved: list[str] = Field(default_factory=list, max_length=32)


class StageEvent(BaseModel):
    stage: Literal["reference", "geometry", "components", "textures", "refinement", "rigging", "motion", "animation", "animation_review", "output"]
    status: Literal["running", "complete", "failed"]
    detail: str