Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
14 changes: 11 additions & 3 deletions comfy_api_nodes/nodes_gemini.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import base64
import os
from enum import Enum
from fnmatch import fnmatch
from io import BytesIO
from typing import Literal

Expand Down Expand Up @@ -119,6 +120,13 @@ async def create_image_parts(
return image_parts


def _mime_matches(mime: GeminiMimeType | None, pattern: str) -> bool:
"""Check if a MIME type matches a pattern. Supports fnmatch globs (e.g. 'image/*')."""
if mime is None:
return False
return fnmatch(mime.value, pattern)


def get_parts_by_type(response: GeminiGenerateContentResponse, part_type: Literal["text"] | str) -> list[GeminiPart]:
"""
Filter response parts by their type.
Expand Down Expand Up @@ -151,9 +159,9 @@ def get_parts_by_type(response: GeminiGenerateContentResponse, part_type: Litera
for part in candidate.content.parts:
if part_type == "text" and part.text:
parts.append(part)
elif part.inlineData and part.inlineData.mimeType == part_type:
elif part.inlineData and _mime_matches(part.inlineData.mimeType, part_type):
parts.append(part)
elif part.fileData and part.fileData.mimeType == part_type:
elif part.fileData and _mime_matches(part.fileData.mimeType, part_type):
parts.append(part)

if not parts and blocked_reasons:
Expand All @@ -178,7 +186,7 @@ def get_text_from_response(response: GeminiGenerateContentResponse) -> str:

async def get_image_from_response(response: GeminiGenerateContentResponse) -> Input.Image:
image_tensors: list[Input.Image] = []
parts = get_parts_by_type(response, "image/png")
parts = get_parts_by_type(response, "image/*")
for part in parts:
if part.inlineData:
image_data = base64.b64decode(part.inlineData.data)
Expand Down
3 changes: 3 additions & 0 deletions comfy_api_nodes/nodes_rodin.py
Original file line number Diff line number Diff line change
Expand Up @@ -505,6 +505,9 @@ def define_schema(cls) -> IO.Schema:
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.4}""",
),
)

@classmethod
Expand Down
Loading