[Partner Nodes] fix(Google): switch to Interactions API for Omni model (#14984)
This commit is contained in:
parent
7774301a7b
commit
83082a51c4
|
|
@ -1,6 +1,6 @@
|
|||
from datetime import date
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
|
@ -242,3 +242,60 @@ class GeminiGenerateContentResponse(BaseModel):
|
|||
promptFeedback: GeminiPromptFeedback | None = Field(None)
|
||||
usageMetadata: GeminiUsageMetadata | None = Field(None)
|
||||
modelVersion: str | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionTextPart(BaseModel):
|
||||
type: Literal["text"] = "text"
|
||||
text: str = Field(...)
|
||||
|
||||
|
||||
class GeminiInteractionMediaPart(BaseModel):
|
||||
type: str = Field(..., description="One of: image, video, audio, document.")
|
||||
data: str | None = Field(None, description="Base64-encoded media bytes.")
|
||||
uri: str | None = Field(None, description="URI of the media, as an alternative to inline data.")
|
||||
mime_type: str | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionGenerationConfig(BaseModel):
|
||||
temperature: float | None = Field(None, ge=0.0, le=2.0)
|
||||
top_p: float | None = Field(None, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class GeminiInteractionRequest(BaseModel):
|
||||
model: str = Field(...)
|
||||
input: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = Field(...)
|
||||
generation_config: GeminiInteractionGenerationConfig | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionModalityTokens(BaseModel):
|
||||
modality: str | None = Field(None, description="One of: text, image, audio, video, document.")
|
||||
tokens: int | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionUsage(BaseModel):
|
||||
input_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
|
||||
output_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
|
||||
total_thought_tokens: int | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionContent(BaseModel):
|
||||
type: str | None = Field(None)
|
||||
text: str | None = Field(None)
|
||||
data: str | None = Field(None)
|
||||
uri: str | None = Field(None)
|
||||
mime_type: str | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionStep(BaseModel):
|
||||
type: str | None = Field(None)
|
||||
content: list[GeminiInteractionContent] | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteraction(BaseModel):
|
||||
id: str | None = Field(None)
|
||||
status: str | None = Field(
|
||||
None,
|
||||
description="One of: in_progress, requires_action, completed, failed, cancelled, incomplete.",
|
||||
)
|
||||
steps: list[GeminiInteractionStep] | None = Field(None)
|
||||
usage: GeminiInteractionUsage | None = Field(None)
|
||||
|
|
|
|||
|
|
@ -24,6 +24,11 @@ from comfy_api_nodes.apis.gemini import (
|
|||
GeminiImageGenerateContentRequest,
|
||||
GeminiImageGenerationConfig,
|
||||
GeminiInlineData,
|
||||
GeminiInteraction,
|
||||
GeminiInteractionGenerationConfig,
|
||||
GeminiInteractionMediaPart,
|
||||
GeminiInteractionRequest,
|
||||
GeminiInteractionTextPart,
|
||||
GeminiMimeType,
|
||||
GeminiPart,
|
||||
GeminiRole,
|
||||
|
|
@ -51,6 +56,7 @@ from comfy_api_nodes.util import (
|
|||
)
|
||||
|
||||
GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini"
|
||||
GEMINI_INTERACTIONS_ENDPOINT = "/proxy/gemini-interactions"
|
||||
GEMINI_MAX_INPUT_FILE_SIZE = 20 * 1024 * 1024 # 20 MB
|
||||
GEMINI_URL_INPUT_BUDGET = 10
|
||||
GEMINI_MAX_INLINE_BYTES = 18 * 1024 * 1024
|
||||
|
|
@ -231,29 +237,10 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
|
|||
return torch.cat(image_tensors, dim=0)
|
||||
|
||||
|
||||
async def get_video_from_response(
|
||||
response: GeminiGenerateContentResponse, cls: type[IO.ComfyNode] | None = None
|
||||
) -> InputImpl.VideoFromFile:
|
||||
parts = get_parts_by_type(response, "video/*")
|
||||
for part in parts:
|
||||
if part.inlineData and part.inlineData.data:
|
||||
return InputImpl.VideoFromFile(BytesIO(base64.b64decode(part.inlineData.data)))
|
||||
if part.fileData and part.fileData.fileUri:
|
||||
return await download_url_to_video_output(part.fileData.fileUri, cls=cls)
|
||||
model_message = get_text_from_response(response).strip()
|
||||
if model_message:
|
||||
raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
|
||||
raise ValueError(
|
||||
"Gemini did not generate a video. Try rephrasing your prompt, "
|
||||
"shortening the requested duration, or reducing the number of input images/videos."
|
||||
)
|
||||
|
||||
|
||||
def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | None:
|
||||
if not response.modelVersion:
|
||||
return None
|
||||
# Define prices (Cost per 1,000,000 tokens), see https://cloud.google.com/vertex-ai/generative-ai/pricing
|
||||
output_video_tokens_price = 0.0
|
||||
if response.modelVersion == "gemini-2.5-pro":
|
||||
input_tokens_price = 1.25
|
||||
output_text_tokens_price = 10.0
|
||||
|
|
@ -290,11 +277,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
input_tokens_price = 0.25
|
||||
output_text_tokens_price = 1.50
|
||||
output_image_tokens_price = 30.0
|
||||
elif response.modelVersion == "gemini-omni-flash-preview":
|
||||
input_tokens_price = 2.145
|
||||
output_text_tokens_price = 12.87
|
||||
output_image_tokens_price = 0.0
|
||||
output_video_tokens_price = 25.025
|
||||
else:
|
||||
return None
|
||||
final_price = response.usageMetadata.promptTokenCount * input_tokens_price
|
||||
|
|
@ -302,8 +284,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
for i in response.usageMetadata.candidatesTokensDetails:
|
||||
if i.modality == Modality.IMAGE:
|
||||
final_price += output_image_tokens_price * i.tokenCount # for Nano Banana models
|
||||
elif i.modality == Modality.VIDEO:
|
||||
final_price += output_video_tokens_price * i.tokenCount # for Omni Flash
|
||||
else:
|
||||
final_price += output_text_tokens_price * i.tokenCount
|
||||
if response.usageMetadata.thoughtsTokenCount:
|
||||
|
|
@ -311,6 +291,58 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
return final_price / 1_000_000.0
|
||||
|
||||
|
||||
def get_text_from_interaction(interaction: GeminiInteraction) -> str:
|
||||
"""Extract and concatenate all model output text from an Interactions API response."""
|
||||
texts = []
|
||||
for step in interaction.steps or []:
|
||||
if step.type != "model_output":
|
||||
continue
|
||||
for content in step.content or []:
|
||||
if content.type == "text" and content.text:
|
||||
texts.append(content.text)
|
||||
return "\n".join(texts)
|
||||
|
||||
|
||||
async def get_video_from_interaction(
|
||||
interaction: GeminiInteraction, cls: type[IO.ComfyNode] | None = None
|
||||
) -> InputImpl.VideoFromFile:
|
||||
for step in interaction.steps or []:
|
||||
if step.type != "model_output":
|
||||
continue
|
||||
for content in step.content or []:
|
||||
if content.type != "video":
|
||||
continue
|
||||
if content.data:
|
||||
return InputImpl.VideoFromFile(BytesIO(base64.b64decode(content.data)))
|
||||
if content.uri:
|
||||
return await download_url_to_video_output(content.uri, cls=cls)
|
||||
model_message = get_text_from_interaction(interaction).strip()
|
||||
if model_message:
|
||||
raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
|
||||
raise ValueError(
|
||||
"Gemini did not generate a video. Try rephrasing your prompt, "
|
||||
"shortening the requested duration, or reducing the number of input images/videos."
|
||||
)
|
||||
|
||||
|
||||
def calculate_interaction_tokens_price(interaction: GeminiInteraction) -> float | None:
|
||||
if interaction.usage is None:
|
||||
return None
|
||||
input_tokens_price = 1.5
|
||||
output_tokens_prices = {"text": 9.0, "video": 17.5}
|
||||
thoughts_tokens_price = 9.0
|
||||
final_price = 0.0
|
||||
for i in interaction.usage.input_tokens_by_modality or []:
|
||||
if i.tokens:
|
||||
final_price += input_tokens_price * i.tokens
|
||||
for i in interaction.usage.output_tokens_by_modality or []:
|
||||
if i.tokens and i.modality in output_tokens_prices:
|
||||
final_price += output_tokens_prices[i.modality] * i.tokens
|
||||
if interaction.usage.total_thought_tokens:
|
||||
final_price += thoughts_tokens_price * interaction.usage.total_thought_tokens
|
||||
return final_price / 1_000_000.0
|
||||
|
||||
|
||||
def create_video_parts(video_input: Input.Video) -> list[GeminiPart]:
|
||||
"""Convert a single video input to Gemini API compatible parts (inline MP4/H.264)."""
|
||||
base_64_string = video_to_base64_string(
|
||||
|
|
@ -445,6 +477,15 @@ async def build_gemini_media_parts(
|
|||
return parts
|
||||
|
||||
|
||||
def to_interaction_media_part(part: GeminiPart) -> GeminiInteractionMediaPart:
|
||||
"""Convert a fileData/inlineData GeminiPart into an Interactions API media part."""
|
||||
if part.fileData:
|
||||
mime = part.fileData.mimeType.value
|
||||
return GeminiInteractionMediaPart(type=mime.split("/")[0], uri=part.fileData.fileUri, mime_type=mime)
|
||||
mime = part.inlineData.mimeType.value
|
||||
return GeminiInteractionMediaPart(type=mime.split("/")[0], data=part.inlineData.data, mime_type=mime)
|
||||
|
||||
|
||||
class GeminiNode(IO.ComfyNode):
|
||||
"""
|
||||
Node to generate text responses from a Gemini model.
|
||||
|
|
@ -1676,7 +1717,7 @@ class GeminiVideoOmni(IO.ComfyNode):
|
|||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr='{"type":"usd","usd":0.146,"format":{"suffix":"/second","approximate":true}}'
|
||||
expr='{"type":"usd","usd":0.101,"format":{"suffix":"/second","approximate":true}}'
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -1695,27 +1736,34 @@ class GeminiVideoOmni(IO.ComfyNode):
|
|||
for video in videos:
|
||||
validate_video_duration(video, max_duration=10)
|
||||
|
||||
parts: list[GeminiPart] = []
|
||||
parts: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = []
|
||||
if images or videos:
|
||||
parts.extend(await build_gemini_media_parts(cls, images, [], videos))
|
||||
parts.append(GeminiPart(text=prompt))
|
||||
response = await sync_op(
|
||||
media_parts = await build_gemini_media_parts(cls, images, [], videos)
|
||||
parts.extend(to_interaction_media_part(p) for p in media_parts)
|
||||
parts.append(GeminiInteractionTextPart(text=prompt))
|
||||
interaction = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{GEMINI_BASE_ENDPOINT}/{model_id}", method="POST"),
|
||||
data=GeminiGenerateContentRequest(
|
||||
contents=[GeminiContent(role=GeminiRole.user, parts=parts)],
|
||||
generationConfig=GeminiGenerationConfig(
|
||||
responseModalities=["TEXT", "VIDEO"],
|
||||
ApiEndpoint(path=GEMINI_INTERACTIONS_ENDPOINT, method="POST"),
|
||||
data=GeminiInteractionRequest(
|
||||
model=model_id,
|
||||
input=parts,
|
||||
generation_config=GeminiInteractionGenerationConfig(
|
||||
temperature=model.get("temperature", 1.0),
|
||||
topP=model.get("top_p", 0.95),
|
||||
top_p=model.get("top_p", 0.95),
|
||||
),
|
||||
),
|
||||
response_model=GeminiGenerateContentResponse,
|
||||
price_extractor=calculate_tokens_price,
|
||||
response_model=GeminiInteraction,
|
||||
price_extractor=calculate_interaction_tokens_price,
|
||||
)
|
||||
if interaction.status != "completed":
|
||||
model_message = get_text_from_interaction(interaction).strip()
|
||||
raise ValueError(
|
||||
f"Gemini interaction did not complete (status: {interaction.status})."
|
||||
+ (f" Model response: {model_message}" if model_message else "")
|
||||
)
|
||||
return IO.NodeOutput(
|
||||
await get_video_from_response(response, cls=cls),
|
||||
get_text_from_response(response),
|
||||
await get_video_from_interaction(interaction, cls=cls),
|
||||
get_text_from_interaction(interaction),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue