96 lines
2.5 KiB
Python
96 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class KieBase64UploadRequest(BaseModel):
|
|
base64Data: str = Field(..., description="Base64 data URL payload.")
|
|
uploadPath: str = Field("images/base64", description="Kie upload path bucket.")
|
|
fileName: str = Field("reference.png", description="Uploaded file name.")
|
|
|
|
|
|
class KieUploadFileData(BaseModel):
|
|
fileName: str | None = None
|
|
filePath: str | None = None
|
|
downloadUrl: str | None = None
|
|
mimeType: str | None = None
|
|
|
|
|
|
class KieBase64UploadResponse(BaseModel):
|
|
code: int
|
|
msg: str
|
|
data: KieUploadFileData
|
|
|
|
|
|
class KieNanoBanana2Input(BaseModel):
|
|
prompt: str
|
|
aspect_ratio: str = "auto"
|
|
resolution: str = "2K"
|
|
output_format: str = "png"
|
|
google_search: bool = False
|
|
image_input: list[str] = Field(default_factory=list)
|
|
|
|
|
|
class KieCreateTaskRequest(BaseModel):
|
|
model: str = "nano-banana-2"
|
|
input: KieNanoBanana2Input
|
|
callBackUrl: str | None = None
|
|
|
|
|
|
class KieTaskCreateData(BaseModel):
|
|
taskId: str
|
|
|
|
|
|
class KieCreateTaskResponse(BaseModel):
|
|
code: int
|
|
msg: str
|
|
data: KieTaskCreateData
|
|
|
|
|
|
class KieTaskRecordData(BaseModel):
|
|
taskId: str
|
|
state: str
|
|
resultJson: str | dict[str, Any] | None = None
|
|
failCode: str | None = None
|
|
failMsg: str | None = None
|
|
|
|
def parsed_result(self) -> dict[str, Any]:
|
|
if self.resultJson is None or self.resultJson == "":
|
|
return {}
|
|
if isinstance(self.resultJson, dict):
|
|
return self.resultJson
|
|
try:
|
|
parsed = json.loads(self.resultJson)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
return parsed if isinstance(parsed, dict) else {}
|
|
|
|
def result_urls(self) -> list[str]:
|
|
payload = self.parsed_result()
|
|
urls = payload.get("resultUrls", [])
|
|
if not isinstance(urls, list):
|
|
return []
|
|
return [url for url in urls if isinstance(url, str) and url]
|
|
|
|
def failure_reason(self) -> str:
|
|
fail_code = (self.failCode or "").strip()
|
|
fail_msg = (self.failMsg or "").strip()
|
|
if fail_code and fail_msg:
|
|
return f"{fail_code}: {fail_msg}"
|
|
return fail_code or fail_msg
|
|
|
|
|
|
class KieTaskRecordResponse(BaseModel):
|
|
code: int
|
|
msg: str
|
|
data: KieTaskRecordData
|
|
|
|
|
|
KieBase64UploadResponse.model_rebuild()
|
|
KieCreateTaskRequest.model_rebuild()
|
|
KieCreateTaskResponse.model_rebuild()
|
|
KieTaskRecordResponse.model_rebuild()
|