Skip to content

fix: adapt image generation/editing to new Grok cardAttachment API #392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JinchengGao-Infty wants to merge 1 commit into
base: main
Choose a base branch
from
+102 −23
Open

fix: adapt image generation/editing to new Grok cardAttachment API #392

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 14 additions & 11 deletions app/services/grok/services/image.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,10 @@ async def generate(
enable_nsfw: Optional[bool] = None,
chat_format: bool = False,
) -> ImageGenerationResult:
max_token_retries = int(get_config("retry.max_retry") or 3)
# Image generation is much more bursty than text chat.
# Keep a wider retry budget so temporary per-token image 429s
# do not fail fast when there are still many healthy tokens.
max_token_retries = max(int(get_config("retry.max_retry") or 3), 8)
tried_tokens: set[str] = set()
last_error: Optional[Exception] = None

Expand Down Expand Up @@ -285,14 +288,15 @@ async def _stream_app_chat(
enable_nsfw: Optional[bool] = None,
chat_format: bool = False,
) -> ImageGenerationResult:
overrides = self._app_chat_request_overrides(n, enable_nsfw)
overrides["modeId"] = "auto"
response = await GrokChatService().chat(
token=token,
message=prompt,
model=model_info.grok_model,
mode=model_info.model_mode,
model=None,
mode=None,
stream=True,
tool_overrides={"imageGen": True},
request_overrides=self._app_chat_request_overrides(n, enable_nsfw),
request_overrides=overrides,
)
processor = AppChatImageStreamProcessor(
model_info.model_id,
Expand Down Expand Up @@ -324,16 +328,15 @@ async def _collect_app_chat(
calls_needed = max(1, int(math.ceil(n / per_call)))

async def _call_generate(call_target: int) -> List[str]:
overrides = self._app_chat_request_overrides(call_target, enable_nsfw)
overrides["modeId"] = "auto"
response = await GrokChatService().chat(
token=token,
message=prompt,
model=model_info.grok_model,
mode=model_info.model_mode,
model=None,
mode=None,
stream=True,
tool_overrides={"imageGen": True},
request_overrides=self._app_chat_request_overrides(
call_target, enable_nsfw
),
request_overrides=overrides,
)
processor = AppChatImageCollectProcessor(
model_info.model_id,
Expand Down
84 changes: 77 additions & 7 deletions app/services/grok/services/image_edit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,15 +94,24 @@ async def edit(
tried_tokens.add(current_token)
try:
file_attachments = await self._upload_images(images, current_token)
tool_overrides: Dict[str, Any] | None = None
tool_overrides: Dict[str, Any] = {
"gmailSearch": False,
"googleCalendarSearch": False,
"outlookSearch": False,
"outlookCalendarSearch": False,
"googleDriveSearch": False,
}
request_overrides = self._build_request_overrides(n)
request_overrides["modeId"] = "auto"
request_overrides["disableMemory"] = False
request_overrides["temporary"] = False

if stream:
response = await GrokChatService().chat(
token=current_token,
message=prompt,
model=_EDIT_UPSTREAM_MODEL,
mode=_EDIT_UPSTREAM_MODE,
model=None,
mode=None,
stream=True,
file_attachments=file_attachments,
tool_overrides=tool_overrides,
Expand Down Expand Up @@ -203,15 +212,19 @@ async def _collect_images(
calls_needed = max(1, (n + per_call - 1) // per_call)

async def _call_edit():
edit_overrides = self._build_request_overrides(per_call)
edit_overrides["modeId"] = "auto"
edit_overrides["disableMemory"] = False
edit_overrides["temporary"] = False
response = await GrokChatService().chat(
token=token,
message=prompt,
model=_EDIT_UPSTREAM_MODEL,
mode=_EDIT_UPSTREAM_MODE,
model=None,
mode=None,
stream=True,
file_attachments=file_attachments,
tool_overrides=tool_overrides,
request_overrides=self._build_request_overrides(per_call),
request_overrides=edit_overrides,
)
processor = ImageCollectProcessor(
"grok-imagine-1.0-edit", token, response_format=response_format
Expand Down Expand Up @@ -331,7 +344,36 @@ async def process(
)
continue

# modelResponse
# Handle cardAttachment-based image generation (new Grok format)
if ca := resp.get("cardAttachment"):
try:
jd = orjson.loads(ca.get("jsonData", b"{}"))
if jd.get("type") in ("render_generated_image", "render_edited_image"):
chunk = jd.get("image_chunk", {})
if chunk.get("progress", 0) >= 100 and chunk.get("imageUrl"):
url = f"https://assets.grok.com/{chunk['imageUrl']}"
if self.response_format == "url":
processed = await self.process_url(url, "image")
if processed:
final_images.append(processed)
else:
try:
dl_service = self._get_dl()
base64_data = await dl_service.parse_b64(
url, self.token, "image"
)
if base64_data:
b64 = base64_data.split(",", 1)[1] if "," in base64_data else base64_data
final_images.append(b64)
except Exception as e:
logger.warning(f"Failed to convert stream card image to base64: {e}")
processed = await self.process_url(url, "image")
if processed:
final_images.append(processed)
except Exception:
pass

# modelResponse (legacy format)
if mr := resp.get("modelResponse"):
if urls := _collect_images(mr):
for url in urls:
Expand Down Expand Up @@ -490,6 +532,34 @@ async def process(self, response: AsyncIterable[bytes]) -> List[str]:
continue

resp = data.get("result", {}).get("response", {})
# Handle cardAttachment-based image generation/edit (new Grok format)
if ca := resp.get("cardAttachment"):
try:
jd = orjson.loads(ca.get("jsonData", b"{}"))
if jd.get("type") in ("render_generated_image", "render_edited_image"):
chunk = jd.get("image_chunk", {})
if chunk.get("progress", 0) >= 100 and chunk.get("imageUrl"):
url = f"https://assets.grok.com/{chunk['imageUrl']}"
if self.response_format == "url":
processed = await self.process_url(url, "image")
if processed:
images.append(processed)
else:
try:
dl_service = self._get_dl()
base64_data = await dl_service.parse_b64(
url, self.token, "image"
)
if base64_data:
b64 = base64_data.split(",", 1)[1] if "," in base64_data else base64_data
images.append(b64)
except Exception as e:
logger.warning(f"Failed to convert card image to base64: {e}")
processed = await self.process_url(url, "image")
if processed:
images.append(processed)
except Exception:
pass

if mr := resp.get("modelResponse"):
if urls := _collect_images(mr):
Expand Down
16 changes: 11 additions & 5 deletions app/services/reverse/app_chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,18 +140,24 @@ def build_payload(
"isAsyncChat": False,
"isReasoning": False,
"message": message,
"modelMode": mode,
"modelName": model,
"responseMetadata": {
"requestModelDetails": {"modelId": model},
},
"returnImageBytes": False,
"returnRawGrokInXaiRequest": False,
"sendFinalMetadata": True,
"temporary": get_config("app.temporary"),
"toolOverrides": tool_overrides or {},
}

# When model is None, use modeId-based routing (e.g. "auto")
# instead of explicit modelName/modelMode.
if model is not None:
payload["modelName"] = model
payload["modelMode"] = mode
payload["responseMetadata"] = {
"requestModelDetails": {"modelId": model},
}
else:
payload["responseMetadata"] = {}

if model == "grok-420":
payload["enable420"] = True

Expand Down
Loading