grok2api / app /dataplane /reverse /protocol /xai_image_edit.py
FUCAT's picture
Deploy grok2api to HF Spaces (Docker)
7e55e53
Raw
History Blame Contribute Delete
3.59 kB
"""XAI image-edit protocol — payload builder and SSE field extractors."""
from typing import Any
from app.platform.config.snapshot import get_config
IMAGE_EDIT_MODEL_NAME = "imagine-image-edit"
IMAGE_EDIT_MODEL_KIND = "imagine"
IMAGE_POST_MEDIA_TYPE = "MEDIA_POST_TYPE_IMAGE"
IMAGE_EDIT_GENERATION_COUNT = 2
def build_image_edit_payload(
*,
prompt: str,
image_references: list[str],
parent_post_id: str,
) -> dict[str, Any]:
"""Build the JSON payload for imagine image-edit chat requests."""
cfg = get_config()
return {
"temporary": cfg.get_bool("features.temporary", True),
"modelName": IMAGE_EDIT_MODEL_NAME,
"message": prompt,
"enableImageGeneration": True,
"returnImageBytes": False,
"returnRawGrokInXaiRequest": False,
"enableImageStreaming": True,
"imageGenerationCount": IMAGE_EDIT_GENERATION_COUNT,
"forceConcise": False,
"enableSideBySide": True,
"sendFinalMetadata": True,
"isReasoning": False,
"disableTextFollowUps": True,
"responseMetadata": {
"modelConfigOverride": {
"modelMap": {
"imageEditModel": IMAGE_EDIT_MODEL_KIND,
"imageEditModelConfig": {
"imageReferences": image_references,
"parentPostId": parent_post_id,
},
}
}
},
"disableMemory": not cfg.get_bool("features.memory", False),
"forceSideBySide": False,
}
def extract_streaming_response(data: dict[str, Any]) -> dict[str, Any] | None:
"""Return ``response.streamingImageGenerationResponse`` when present."""
result = data.get("result")
if not isinstance(result, dict):
return None
response = result.get("response")
if not isinstance(response, dict):
return None
stream = response.get("streamingImageGenerationResponse")
return stream if isinstance(stream, dict) else None
def extract_model_response_urls(data: dict[str, Any]) -> list[str]:
"""Return fallback ``modelResponse.generatedImageUrls`` values."""
result = data.get("result")
if not isinstance(result, dict):
return []
response = result.get("response")
if not isinstance(response, dict):
return []
model_response = response.get("modelResponse")
if not isinstance(model_response, dict):
return []
urls = model_response.get("generatedImageUrls")
if not isinstance(urls, list):
return []
return [url for url in urls if isinstance(url, str) and url]
def extract_model_response_file_attachments(data: dict[str, Any]) -> list[str]:
"""Return fallback ``modelResponse.fileAttachments`` asset IDs."""
result = data.get("result")
if not isinstance(result, dict):
return []
response = result.get("response")
if not isinstance(response, dict):
return []
model_response = response.get("modelResponse")
if not isinstance(model_response, dict):
return []
attachments = model_response.get("fileAttachments")
if not isinstance(attachments, list):
return []
return [attachment for attachment in attachments if isinstance(attachment, str) and attachment]
__all__ = [
"IMAGE_EDIT_MODEL_NAME",
"IMAGE_EDIT_MODEL_KIND",
"IMAGE_POST_MEDIA_TYPE",
"IMAGE_EDIT_GENERATION_COUNT",
"build_image_edit_payload",
"extract_streaming_response",
"extract_model_response_urls",
"extract_model_response_file_attachments",
]