44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
# services/comfy-agent/app/models.py
|
|
from pydantic import BaseModel, Field
|
|
from typing import Any, Dict, Optional, Literal
|
|
|
|
GenType = Literal["text-to-image", "text-to-video", "image-to-video"]
|
|
|
|
class GenerateImageRequest(BaseModel):
|
|
prompt: str = Field(min_length=1)
|
|
negative_prompt: Optional[str] = None
|
|
width: int = 1024
|
|
height: int = 1024
|
|
steps: int = 28
|
|
seed: Optional[int] = None
|
|
idempotency_key: Optional[str] = None
|
|
workflow: Optional[str] = None
|
|
workflow_params: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
class GenerateVideoRequest(BaseModel):
|
|
prompt: str = Field(min_length=1)
|
|
negative_prompt: Optional[str] = None
|
|
width: int = 768
|
|
height: int = 512
|
|
frames: Optional[int] = None
|
|
seconds: int = 4
|
|
fps: int = 24
|
|
steps: int = 30
|
|
cfg: float = 2.5
|
|
seed: Optional[int] = None
|
|
format: str = "mp4"
|
|
codec: str = "h264"
|
|
idempotency_key: Optional[str] = None
|
|
workflow: Optional[str] = None
|
|
workflow_params: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
class JobStatus(BaseModel):
|
|
job_id: str
|
|
type: GenType
|
|
status: Literal["queued", "running", "succeeded", "failed"]
|
|
progress: float = 0.0
|
|
message: Optional[str] = None
|
|
result_url: Optional[str] = None
|
|
error: Optional[str] = None
|
|
comfy_prompt_id: Optional[str] = None
|