mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-04 06:37:06 +08:00
Add Google Gemini API node (#191)
This commit is contained in:
Christian Byrne
Robin Huang
parent
b9ff57aabd
commit
258b7ae7c9
@ -215,6 +215,7 @@ def download_url_to_image_tensor(url: str, timeout: int = None) -> torch.Tensor:
|
||||
image_bytesio = download_url_to_bytesio(url, timeout)
|
||||
return bytesio_to_image_tensor(image_bytesio)
|
||||
|
||||
|
||||
def process_image_response(response: requests.Response) -> torch.Tensor:
|
||||
"""Uses content from a Response object and converts it to a torch.Tensor"""
|
||||
return bytesio_to_image_tensor(BytesIO(response.content))
|
||||
@ -339,7 +340,7 @@ def upload_file_to_comfyapi(
|
||||
file_bytes_io: BytesIO,
|
||||
filename: str,
|
||||
upload_mime_type: str,
|
||||
auth_kwargs: Optional[dict[str,str]] = None,
|
||||
auth_kwargs: Optional[dict[str, str]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Uploads a single file to ComfyUI API and returns its download URL.
|
||||
@ -374,9 +375,33 @@ def upload_file_to_comfyapi(
|
||||
return response.download_url
|
||||
|
||||
|
||||
def video_to_base64_string(
|
||||
video: VideoInput,
|
||||
container_format: VideoContainer = None,
|
||||
codec: VideoCodec = None
|
||||
) -> str:
|
||||
"""
|
||||
Converts a video input to a base64 string.
|
||||
|
||||
Args:
|
||||
video: The video input to convert
|
||||
container_format: Optional container format to use (defaults to video.container if available)
|
||||
codec: Optional codec to use (defaults to video.codec if available)
|
||||
"""
|
||||
video_bytes_io = io.BytesIO()
|
||||
|
||||
# Use provided format/codec if specified, otherwise use video's own if available
|
||||
format_to_use = container_format if container_format is not None else getattr(video, 'container', VideoContainer.MP4)
|
||||
codec_to_use = codec if codec is not None else getattr(video, 'codec', VideoCodec.H264)
|
||||
|
||||
video.save_to(video_bytes_io, format=format_to_use, codec=codec_to_use)
|
||||
video_bytes_io.seek(0)
|
||||
return base64.b64encode(video_bytes_io.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def upload_video_to_comfyapi(
|
||||
video: VideoInput,
|
||||
auth_kwargs: Optional[dict[str,str]] = None,
|
||||
auth_kwargs: Optional[dict[str, str]] = None,
|
||||
container: VideoContainer = VideoContainer.MP4,
|
||||
codec: VideoCodec = VideoCodec.H264,
|
||||
max_duration: Optional[int] = None,
|
||||
@ -478,7 +503,7 @@ def audio_ndarray_to_bytesio(
|
||||
|
||||
def upload_audio_to_comfyapi(
|
||||
audio: AudioInput,
|
||||
auth_kwargs: Optional[dict[str,str]] = None,
|
||||
auth_kwargs: Optional[dict[str, str]] = None,
|
||||
container_format: str = "mp4",
|
||||
codec_name: str = "aac",
|
||||
mime_type: str = "audio/mp4",
|
||||
@ -505,8 +530,25 @@ def upload_audio_to_comfyapi(
|
||||
return upload_file_to_comfyapi(audio_bytes_io, filename, mime_type, auth_kwargs)
|
||||
|
||||
|
||||
def audio_to_base64_string(
|
||||
audio: AudioInput, container_format: str = "mp4", codec_name: str = "aac"
|
||||
) -> str:
|
||||
"""Converts an audio input to a base64 string."""
|
||||
sample_rate: int = audio["sample_rate"]
|
||||
waveform: torch.Tensor = audio["waveform"]
|
||||
audio_data_np = audio_tensor_to_contiguous_ndarray(waveform)
|
||||
audio_bytes_io = audio_ndarray_to_bytesio(
|
||||
audio_data_np, sample_rate, container_format, codec_name
|
||||
)
|
||||
audio_bytes = audio_bytes_io.getvalue()
|
||||
return base64.b64encode(audio_bytes).decode("utf-8")
|
||||
|
||||
|
||||
def upload_images_to_comfyapi(
|
||||
image: torch.Tensor, max_images=8, auth_kwargs: Optional[dict[str,str]] = None, mime_type: Optional[str] = None
|
||||
image: torch.Tensor,
|
||||
max_images=8,
|
||||
auth_kwargs: Optional[dict[str, str]] = None,
|
||||
mime_type: Optional[str] = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Uploads images to ComfyUI API and returns download URLs.
|
||||
@ -571,17 +613,24 @@ def upload_images_to_comfyapi(
|
||||
return download_urls
|
||||
|
||||
|
||||
def resize_mask_to_image(mask: torch.Tensor, image: torch.Tensor,
|
||||
upscale_method="nearest-exact", crop="disabled",
|
||||
allow_gradient=True, add_channel_dim=False):
|
||||
def resize_mask_to_image(
|
||||
mask: torch.Tensor,
|
||||
image: torch.Tensor,
|
||||
upscale_method="nearest-exact",
|
||||
crop="disabled",
|
||||
allow_gradient=True,
|
||||
add_channel_dim=False,
|
||||
):
|
||||
"""
|
||||
Resize mask to be the same dimensions as an image, while maintaining proper format for API calls.
|
||||
"""
|
||||
_, H, W, _ = image.shape
|
||||
mask = mask.unsqueeze(-1)
|
||||
mask = mask.movedim(-1,1)
|
||||
mask = common_upscale(mask, width=W, height=H, upscale_method=upscale_method, crop=crop)
|
||||
mask = mask.movedim(1,-1)
|
||||
mask = mask.movedim(-1, 1)
|
||||
mask = common_upscale(
|
||||
mask, width=W, height=H, upscale_method=upscale_method, crop=crop
|
||||
)
|
||||
mask = mask.movedim(1, -1)
|
||||
if not add_channel_dim:
|
||||
mask = mask.squeeze(-1)
|
||||
if not allow_gradient:
|
||||
@ -589,15 +638,25 @@ def resize_mask_to_image(mask: torch.Tensor, image: torch.Tensor,
|
||||
return mask
|
||||
|
||||
|
||||
def validate_string(string: str, strip_whitespace=True, field_name="prompt", min_length=None, max_length=None):
|
||||
def validate_string(
|
||||
string: str,
|
||||
strip_whitespace=True,
|
||||
field_name="prompt",
|
||||
min_length=None,
|
||||
max_length=None,
|
||||
):
|
||||
if string is None:
|
||||
raise Exception(f"Field '{field_name}' cannot be empty.")
|
||||
if strip_whitespace:
|
||||
string = string.strip()
|
||||
if min_length and len(string) < min_length:
|
||||
raise Exception(f"Field '{field_name}' cannot be shorter than {min_length} characters; was {len(string)} characters long.")
|
||||
raise Exception(
|
||||
f"Field '{field_name}' cannot be shorter than {min_length} characters; was {len(string)} characters long."
|
||||
)
|
||||
if max_length and len(string) > max_length:
|
||||
raise Exception(f" Field '{field_name} cannot be longer than {max_length} characters; was {len(string)} characters long.")
|
||||
raise Exception(
|
||||
f" Field '{field_name} cannot be longer than {max_length} characters; was {len(string)} characters long."
|
||||
)
|
||||
|
||||
|
||||
def image_tensor_pair_to_batch(
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
# generated by datamodel-codegen:
|
||||
# filename: filtered-openapi.yaml
|
||||
# timestamp: 2025-05-15T19:13:25+00:00
|
||||
# timestamp: 2025-05-17T21:08:37+00:00
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@ -446,6 +446,11 @@ class GeminiCitationMetadata(BaseModel):
|
||||
citations: Optional[List[GeminiCitation]] = None
|
||||
|
||||
|
||||
class Role1(str, Enum):
|
||||
user = 'user'
|
||||
model = 'model'
|
||||
|
||||
|
||||
class GeminiFunctionDeclaration(BaseModel):
|
||||
description: Optional[str] = None
|
||||
name: str
|
||||
@ -455,37 +460,123 @@ class GeminiFunctionDeclaration(BaseModel):
|
||||
|
||||
|
||||
class GeminiGenerationConfig(BaseModel):
|
||||
maxOutputTokens: Optional[int] = None
|
||||
maxOutputTokens: Optional[int] = Field(
|
||||
None,
|
||||
description='Maximum number of tokens that can be generated in the response. A token is approximately 4 characters. 100 tokens correspond to roughly 60-80 words.\n',
|
||||
examples=[2048],
|
||||
ge=16,
|
||||
le=8192,
|
||||
)
|
||||
seed: Optional[int] = Field(
|
||||
None,
|
||||
description="When seed is fixed to a specific value, the model makes a best effort to provide the same response for repeated requests. Deterministic output isn't guaranteed. Also, changing the model or parameter settings, such as the temperature, can cause variations in the response even when you use the same seed value. By default, a random seed value is used. Available for the following models:, gemini-2.5-flash-preview-04-1, gemini-2.5-pro-preview-05-0, gemini-2.0-flash-lite-00, gemini-2.0-flash-001\n",
|
||||
examples=[343940597],
|
||||
)
|
||||
stopSequences: Optional[List[str]] = None
|
||||
temperature: Optional[float] = None
|
||||
topK: Optional[int] = None
|
||||
topP: Optional[float] = None
|
||||
temperature: Optional[float] = Field(
|
||||
1,
|
||||
description="The temperature is used for sampling during response generation, which occurs when topP and topK are applied. Temperature controls the degree of randomness in token selection. Lower temperatures are good for prompts that require a less open-ended or creative response, while higher temperatures can lead to more diverse or creative results. A temperature of 0 means that the highest probability tokens are always selected. In this case, responses for a given prompt are mostly deterministic, but a small amount of variation is still possible. If the model returns a response that's too generic, too short, or the model gives a fallback response, try increasing the temperature\n",
|
||||
ge=0.0,
|
||||
le=2.0,
|
||||
)
|
||||
topK: Optional[int] = Field(
|
||||
40,
|
||||
description="Top-K changes how the model selects tokens for output. A top-K of 1 means the next selected token is the most probable among all tokens in the model's vocabulary. A top-K of 3 means that the next token is selected from among the 3 most probable tokens by using temperature.\n",
|
||||
examples=[40],
|
||||
ge=1,
|
||||
)
|
||||
topP: Optional[float] = Field(
|
||||
0.95,
|
||||
description='If specified, nucleus sampling is used.\nTop-P changes how the model selects tokens for output. Tokens are selected from the most (see top-K) to least probable until the sum of their probabilities equals the top-P value. For example, if tokens A, B, and C have a probability of 0.3, 0.2, and 0.1 and the top-P value is 0.5, then the model will select either A or B as the next token by using temperature and excludes C as a candidate.\nSpecify a lower value for less random responses and a higher value for more random responses.\n',
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
|
||||
|
||||
class GeminiInlineData(BaseModel):
|
||||
data: Optional[str] = None
|
||||
mimeType: Optional[str] = None
|
||||
class GeminiMimeType(str, Enum):
|
||||
application_pdf = 'application/pdf'
|
||||
audio_mpeg = 'audio/mpeg'
|
||||
audio_mp3 = 'audio/mp3'
|
||||
audio_wav = 'audio/wav'
|
||||
image_png = 'image/png'
|
||||
image_jpeg = 'image/jpeg'
|
||||
image_webp = 'image/webp'
|
||||
text_plain = 'text/plain'
|
||||
video_mov = 'video/mov'
|
||||
video_mpeg = 'video/mpeg'
|
||||
video_mp4 = 'video/mp4'
|
||||
video_mpg = 'video/mpg'
|
||||
video_avi = 'video/avi'
|
||||
video_wmv = 'video/wmv'
|
||||
video_mpegps = 'video/mpegps'
|
||||
video_flv = 'video/flv'
|
||||
|
||||
|
||||
class GeminiPart(BaseModel):
|
||||
inlineData: Optional[GeminiInlineData] = None
|
||||
text: Optional[str] = None
|
||||
class GeminiOffset(BaseModel):
|
||||
nanos: Optional[int] = Field(
|
||||
None,
|
||||
description='Signed fractions of a second at nanosecond resolution. Negative second values with fractions must still have non-negative nanos values.\n',
|
||||
examples=[0],
|
||||
ge=0,
|
||||
le=999999999,
|
||||
)
|
||||
seconds: Optional[int] = Field(
|
||||
None,
|
||||
description='Signed seconds of the span of time. Must be from -315,576,000,000 to +315,576,000,000 inclusive.\n',
|
||||
examples=[60],
|
||||
ge=-315576000000,
|
||||
le=315576000000,
|
||||
)
|
||||
|
||||
|
||||
class GeminiSafetyCategory(str, Enum):
|
||||
HARM_CATEGORY_SEXUALLY_EXPLICIT = 'HARM_CATEGORY_SEXUALLY_EXPLICIT'
|
||||
HARM_CATEGORY_HATE_SPEECH = 'HARM_CATEGORY_HATE_SPEECH'
|
||||
HARM_CATEGORY_HARASSMENT = 'HARM_CATEGORY_HARASSMENT'
|
||||
HARM_CATEGORY_DANGEROUS_CONTENT = 'HARM_CATEGORY_DANGEROUS_CONTENT'
|
||||
|
||||
|
||||
class Probability(str, Enum):
|
||||
NEGLIGIBLE = 'NEGLIGIBLE'
|
||||
LOW = 'LOW'
|
||||
MEDIUM = 'MEDIUM'
|
||||
HIGH = 'HIGH'
|
||||
UNKNOWN = 'UNKNOWN'
|
||||
|
||||
|
||||
class GeminiSafetyRating(BaseModel):
|
||||
category: Optional[str] = None
|
||||
probability: Optional[str] = None
|
||||
category: Optional[GeminiSafetyCategory] = None
|
||||
probability: Optional[Probability] = Field(
|
||||
None,
|
||||
description='The probability that the content violates the specified safety category',
|
||||
)
|
||||
|
||||
|
||||
class GeminiSafetySetting(BaseModel):
|
||||
category: str
|
||||
threshold: str
|
||||
class GeminiSafetyThreshold(str, Enum):
|
||||
OFF = 'OFF'
|
||||
BLOCK_NONE = 'BLOCK_NONE'
|
||||
BLOCK_LOW_AND_ABOVE = 'BLOCK_LOW_AND_ABOVE'
|
||||
BLOCK_MEDIUM_AND_ABOVE = 'BLOCK_MEDIUM_AND_ABOVE'
|
||||
BLOCK_ONLY_HIGH = 'BLOCK_ONLY_HIGH'
|
||||
|
||||
|
||||
class GeminiTextPart(BaseModel):
|
||||
text: Optional[str] = Field(
|
||||
None,
|
||||
description='A text prompt or code snippet.',
|
||||
examples=['Answer as concisely as possible'],
|
||||
)
|
||||
|
||||
|
||||
class GeminiTool(BaseModel):
|
||||
functionDeclarations: Optional[List[GeminiFunctionDeclaration]] = None
|
||||
|
||||
|
||||
class GeminiVideoMetadata(BaseModel):
|
||||
endOffset: Optional[GeminiOffset] = None
|
||||
startOffset: Optional[GeminiOffset] = None
|
||||
|
||||
|
||||
class IdeogramColorPalette1(BaseModel):
|
||||
name: str = Field(..., description='Name of the preset color palette')
|
||||
|
||||
@ -785,7 +876,7 @@ class InputImageContent(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class Role1(str, Enum):
|
||||
class Role3(str, Enum):
|
||||
user = 'user'
|
||||
system = 'system'
|
||||
developer = 'developer'
|
||||
@ -3419,7 +3510,7 @@ class OutputAudioContent(BaseModel):
|
||||
type: Type13 = Field(..., description='The type of output content')
|
||||
|
||||
|
||||
class Role2(str, Enum):
|
||||
class Role4(str, Enum):
|
||||
assistant = 'assistant'
|
||||
|
||||
|
||||
@ -3873,6 +3964,54 @@ class ResponseUsage(BaseModel):
|
||||
total_tokens: int = Field(..., description='The total number of tokens used.')
|
||||
|
||||
|
||||
class Rodin3DCheckStatusRequest(BaseModel):
|
||||
subscription_key: str = Field(
|
||||
..., description='subscription from generate endpoint'
|
||||
)
|
||||
|
||||
|
||||
class Rodin3DCheckStatusResponse(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class Rodin3DDownloadRequest(BaseModel):
|
||||
task_uuid: str = Field(..., description='Task UUID')
|
||||
|
||||
|
||||
class RodinGenerateJobsData(BaseModel):
|
||||
subscription_key: Optional[str] = Field(None, description='Subscription Key.')
|
||||
uuids: Optional[List[str]] = Field(None, description='subjobs uuid.')
|
||||
|
||||
|
||||
class RodinMaterialType(str, Enum):
|
||||
PBR = 'PBR'
|
||||
Shaded = 'Shaded'
|
||||
|
||||
|
||||
class RodinMeshModeType(str, Enum):
|
||||
Quad = 'Quad'
|
||||
Raw = 'Raw'
|
||||
|
||||
|
||||
class RodinQualityType(str, Enum):
|
||||
extra_low = 'extra-low'
|
||||
low = 'low'
|
||||
medium = 'medium'
|
||||
high = 'high'
|
||||
|
||||
|
||||
class RodinResourceItem(BaseModel):
|
||||
name: Optional[str] = Field(None, description='File name')
|
||||
url: Optional[str] = Field(None, description='Download url')
|
||||
|
||||
|
||||
class RodinTierType(str, Enum):
|
||||
Regular = 'Regular'
|
||||
Sketch = 'Sketch'
|
||||
Detail = 'Detail'
|
||||
Smooth = 'Smooth'
|
||||
|
||||
|
||||
class RunwayAspectRatioEnum(str, Enum):
|
||||
field_1280_720 = '1280:720'
|
||||
field_720_1280 = '720:1280'
|
||||
@ -3938,6 +4077,35 @@ class RunwayTaskStatusResponse(BaseModel):
|
||||
status: Optional[RunwayTaskStatusEnum] = None
|
||||
|
||||
|
||||
class Model4(str, Enum):
|
||||
gen4_image = 'gen4_image'
|
||||
|
||||
|
||||
class ReferenceImage(BaseModel):
|
||||
uri: Optional[str] = Field(
|
||||
None, description='A HTTPS URL or data URI containing an encoded image'
|
||||
)
|
||||
|
||||
|
||||
class RunwayTextToImageRequest(BaseModel):
|
||||
model: Model4 = Field(..., description='Model to use for generation')
|
||||
promptText: str = Field(
|
||||
..., description='Text prompt for the image generation', max_length=1000
|
||||
)
|
||||
ratio: str = Field(
|
||||
...,
|
||||
description='The resolution (aspect ratio) of the output image',
|
||||
examples=['1920:1088'],
|
||||
)
|
||||
referenceImages: Optional[List[ReferenceImage]] = Field(
|
||||
None, description='Array of reference images to guide the generation'
|
||||
)
|
||||
|
||||
|
||||
class RunwayTextToImageResponse(BaseModel):
|
||||
id: Optional[str] = Field(None, description='Task ID')
|
||||
|
||||
|
||||
class StabilityError(BaseModel):
|
||||
errors: List[str] = Field(
|
||||
...,
|
||||
@ -4025,6 +4193,219 @@ class ToolChoiceTypes(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class TripoAnimation(str, Enum):
|
||||
preset_idle = 'preset:idle'
|
||||
preset_walk = 'preset:walk'
|
||||
preset_climb = 'preset:climb'
|
||||
preset_jump = 'preset:jump'
|
||||
preset_run = 'preset:run'
|
||||
preset_slash = 'preset:slash'
|
||||
preset_shoot = 'preset:shoot'
|
||||
preset_hurt = 'preset:hurt'
|
||||
preset_fall = 'preset:fall'
|
||||
preset_turn = 'preset:turn'
|
||||
|
||||
|
||||
class TripoBalance(BaseModel):
|
||||
balance: float
|
||||
frozen: float
|
||||
|
||||
|
||||
class TripoConvertFormat(str, Enum):
|
||||
GLTF = 'GLTF'
|
||||
USDZ = 'USDZ'
|
||||
FBX = 'FBX'
|
||||
OBJ = 'OBJ'
|
||||
STL = 'STL'
|
||||
field_3MF = '3MF'
|
||||
|
||||
|
||||
class Code(int, Enum):
|
||||
integer_1001 = 1001
|
||||
integer_2000 = 2000
|
||||
integer_2001 = 2001
|
||||
integer_2002 = 2002
|
||||
integer_2003 = 2003
|
||||
integer_2004 = 2004
|
||||
integer_2006 = 2006
|
||||
integer_2007 = 2007
|
||||
integer_2008 = 2008
|
||||
integer_2010 = 2010
|
||||
|
||||
|
||||
class TripoErrorResponse(BaseModel):
|
||||
code: Code
|
||||
message: str
|
||||
suggestion: str
|
||||
|
||||
|
||||
class TripoImageToModel(str, Enum):
|
||||
image_to_model = 'image_to_model'
|
||||
|
||||
|
||||
class TripoModelStyle(str, Enum):
|
||||
person_person2cartoon = 'person:person2cartoon'
|
||||
animal_venom = 'animal:venom'
|
||||
object_clay = 'object:clay'
|
||||
object_steampunk = 'object:steampunk'
|
||||
object_christmas = 'object:christmas'
|
||||
object_barbie = 'object:barbie'
|
||||
gold = 'gold'
|
||||
ancient_bronze = 'ancient_bronze'
|
||||
|
||||
|
||||
class TripoModelVersion(str, Enum):
|
||||
v2_5_20250123 = 'v2.5-20250123'
|
||||
v2_0_20240919 = 'v2.0-20240919'
|
||||
v1_4_20240625 = 'v1.4-20240625'
|
||||
|
||||
|
||||
class TripoMultiviewMode(str, Enum):
|
||||
LEFT = 'LEFT'
|
||||
RIGHT = 'RIGHT'
|
||||
|
||||
|
||||
class TripoMultiviewToModel(str, Enum):
|
||||
multiview_to_model = 'multiview_to_model'
|
||||
|
||||
|
||||
class TripoOrientation(str, Enum):
|
||||
align_image = 'align_image'
|
||||
default = 'default'
|
||||
|
||||
|
||||
class TripoResponseSuccessCode(RootModel[int]):
|
||||
root: int = Field(
|
||||
...,
|
||||
description='Standard success code for Tripo API responses. Typically 0 for success.',
|
||||
examples=[0],
|
||||
)
|
||||
|
||||
|
||||
class TripoSpec(str, Enum):
|
||||
mixamo = 'mixamo'
|
||||
tripo = 'tripo'
|
||||
|
||||
|
||||
class TripoStandardFormat(str, Enum):
|
||||
glb = 'glb'
|
||||
fbx = 'fbx'
|
||||
|
||||
|
||||
class TripoStylizeOptions(str, Enum):
|
||||
lego = 'lego'
|
||||
voxel = 'voxel'
|
||||
voronoi = 'voronoi'
|
||||
minecraft = 'minecraft'
|
||||
|
||||
|
||||
class Code1(int, Enum):
|
||||
integer_0 = 0
|
||||
|
||||
|
||||
class Data8(BaseModel):
|
||||
task_id: str = Field(..., description='used for getTask')
|
||||
|
||||
|
||||
class TripoSuccessTask(BaseModel):
|
||||
code: Code1
|
||||
data: Data8
|
||||
|
||||
|
||||
class Topology(str, Enum):
|
||||
bip = 'bip'
|
||||
quad = 'quad'
|
||||
|
||||
|
||||
class Output(BaseModel):
|
||||
base_model: Optional[str] = None
|
||||
model: Optional[str] = None
|
||||
pbr_model: Optional[str] = None
|
||||
rendered_image: Optional[str] = None
|
||||
riggable: Optional[bool] = None
|
||||
topology: Optional[Topology] = None
|
||||
|
||||
|
||||
class Status10(str, Enum):
|
||||
queued = 'queued'
|
||||
running = 'running'
|
||||
success = 'success'
|
||||
failed = 'failed'
|
||||
cancelled = 'cancelled'
|
||||
unknown = 'unknown'
|
||||
banned = 'banned'
|
||||
expired = 'expired'
|
||||
|
||||
|
||||
class TripoTask(BaseModel):
|
||||
create_time: int
|
||||
input: Dict[str, Any]
|
||||
output: Output
|
||||
progress: int = Field(..., ge=0, le=100)
|
||||
status: Status10
|
||||
task_id: str
|
||||
type: str
|
||||
|
||||
|
||||
class TripoTextToModel(str, Enum):
|
||||
text_to_model = 'text_to_model'
|
||||
|
||||
|
||||
class TripoTextureAlignment(str, Enum):
|
||||
original_image = 'original_image'
|
||||
geometry = 'geometry'
|
||||
|
||||
|
||||
class TripoTextureFormat(str, Enum):
|
||||
BMP = 'BMP'
|
||||
DPX = 'DPX'
|
||||
HDR = 'HDR'
|
||||
JPEG = 'JPEG'
|
||||
OPEN_EXR = 'OPEN_EXR'
|
||||
PNG = 'PNG'
|
||||
TARGA = 'TARGA'
|
||||
TIFF = 'TIFF'
|
||||
WEBP = 'WEBP'
|
||||
|
||||
|
||||
class TripoTextureQuality(str, Enum):
|
||||
standard = 'standard'
|
||||
detailed = 'detailed'
|
||||
|
||||
|
||||
class TripoTopology(str, Enum):
|
||||
bip = 'bip'
|
||||
quad = 'quad'
|
||||
|
||||
|
||||
class TripoTypeAnimatePrerigcheck(str, Enum):
|
||||
animate_prerigcheck = 'animate_prerigcheck'
|
||||
|
||||
|
||||
class TripoTypeAnimateRetarget(str, Enum):
|
||||
animate_retarget = 'animate_retarget'
|
||||
|
||||
|
||||
class TripoTypeAnimateRig(str, Enum):
|
||||
animate_rig = 'animate_rig'
|
||||
|
||||
|
||||
class TripoTypeConvertModel(str, Enum):
|
||||
convert_model = 'convert_model'
|
||||
|
||||
|
||||
class TripoTypeRefineModel(str, Enum):
|
||||
refine_model = 'refine_model'
|
||||
|
||||
|
||||
class TripoTypeStylizeModel(str, Enum):
|
||||
stylize_model = 'stylize_model'
|
||||
|
||||
|
||||
class TripoTypeTextureModel(str, Enum):
|
||||
texture_model = 'texture_model'
|
||||
|
||||
|
||||
class Veo2GenVidPollRequest(BaseModel):
|
||||
operationName: str = Field(
|
||||
...,
|
||||
@ -4179,7 +4560,7 @@ class WebSearchPreviewTool(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class Status10(str, Enum):
|
||||
class Status11(str, Enum):
|
||||
in_progress = 'in_progress'
|
||||
searching = 'searching'
|
||||
completed = 'completed'
|
||||
@ -4192,7 +4573,7 @@ class Type24(str, Enum):
|
||||
|
||||
class WebSearchToolCall(BaseModel):
|
||||
id: str = Field(..., description='The unique ID of the web search tool call.\n')
|
||||
status: Status10 = Field(
|
||||
status: Status11 = Field(
|
||||
..., description='The status of the web search tool call.\n'
|
||||
)
|
||||
type: Type24 = Field(
|
||||
@ -4204,17 +4585,21 @@ class WebSearchToolCall(BaseModel):
|
||||
CreateModelResponseProperties = ModelResponseProperties
|
||||
|
||||
|
||||
class GeminiContent(BaseModel):
|
||||
parts: List[GeminiPart]
|
||||
role: str
|
||||
class GeminiInlineData(BaseModel):
|
||||
data: Optional[str] = Field(
|
||||
None,
|
||||
description='The base64 encoding of the image, PDF, or video to include inline in the prompt. When including media inline, you must also specify the media type (mimeType) of the data. Size limit: 20MB\n',
|
||||
)
|
||||
mimeType: Optional[GeminiMimeType] = None
|
||||
|
||||
|
||||
class GeminiGenerateContentRequest(BaseModel):
|
||||
contents: List[GeminiContent]
|
||||
generationConfig: Optional[GeminiGenerationConfig] = None
|
||||
safetySettings: Optional[List[GeminiSafetySetting]] = None
|
||||
systemInstruction: Optional[GeminiContent] = None
|
||||
tools: Optional[List[GeminiTool]] = None
|
||||
class GeminiPart(BaseModel):
|
||||
inlineData: Optional[GeminiInlineData] = None
|
||||
text: Optional[str] = Field(
|
||||
None,
|
||||
description='A text prompt or code snippet.',
|
||||
examples=['Write a story about a robot learning to paint'],
|
||||
)
|
||||
|
||||
|
||||
class GeminiPromptFeedback(BaseModel):
|
||||
@ -4223,6 +4608,23 @@ class GeminiPromptFeedback(BaseModel):
|
||||
safetyRatings: Optional[List[GeminiSafetyRating]] = None
|
||||
|
||||
|
||||
class GeminiSafetySetting(BaseModel):
|
||||
category: GeminiSafetyCategory
|
||||
threshold: GeminiSafetyThreshold
|
||||
|
||||
|
||||
class GeminiSystemInstructionContent(BaseModel):
|
||||
parts: List[GeminiTextPart] = Field(
|
||||
...,
|
||||
description='A list of ordered parts that make up a single message. Different parts may have different IANA MIME types. For limits on the inputs, such as the maximum number of tokens or the number of images, see the model specifications on the Google models page.\n',
|
||||
)
|
||||
role: Role1 = Field(
|
||||
...,
|
||||
description='The identity of the entity that creates the message. The following values are supported: user: This indicates that the message is sent by a real person, typically a user-generated message. model: This indicates that the message is generated by the model. The model value is used to insert messages from the model into the conversation during multi-turn conversations. For non-multi-turn conversations, this field can be left blank or unset.\n',
|
||||
examples=['user'],
|
||||
)
|
||||
|
||||
|
||||
class IdeogramV3EditRequest(BaseModel):
|
||||
color_palette: Optional[IdeogramColorPalette] = None
|
||||
image: Optional[StrictBytes] = Field(
|
||||
@ -4752,7 +5154,7 @@ class OutputContent(RootModel[Union[OutputTextContent, OutputAudioContent]]):
|
||||
|
||||
class OutputMessage(BaseModel):
|
||||
content: List[OutputContent] = Field(..., description='The content of the message')
|
||||
role: Role2 = Field(..., description='The role of the message')
|
||||
role: Role4 = Field(..., description='The role of the message')
|
||||
type: Type14 = Field(..., description='The type of output item')
|
||||
|
||||
|
||||
@ -4820,6 +5222,27 @@ class ResponseError(BaseModel):
|
||||
message: str = Field(..., description='A human-readable description of the error.')
|
||||
|
||||
|
||||
class Rodin3DDownloadResponse(BaseModel):
|
||||
list: Optional[RodinResourceItem] = None
|
||||
|
||||
|
||||
class Rodin3DGenerateRequest(BaseModel):
|
||||
images: str = Field(..., description='The reference images to generate 3D Assets.')
|
||||
material: Optional[RodinMaterialType] = None
|
||||
mesh_mode: Optional[RodinMeshModeType] = None
|
||||
quality: Optional[RodinQualityType] = None
|
||||
seed: Optional[int] = Field(None, description='Seed.')
|
||||
tier: Optional[RodinTierType] = None
|
||||
|
||||
|
||||
class Rodin3DGenerateResponse(BaseModel):
|
||||
jobs: Optional[RodinGenerateJobsData] = None
|
||||
message: Optional[str] = Field(None, description='message')
|
||||
prompt: Optional[str] = Field(None, description='prompt')
|
||||
submit_time: Optional[str] = Field(None, description='Time')
|
||||
uuid: Optional[str] = Field(None, description='Task UUID')
|
||||
|
||||
|
||||
class RunwayImageToVideoRequest(BaseModel):
|
||||
duration: RunwayDurationEnum
|
||||
model: RunwayModelEnum
|
||||
@ -4874,16 +5297,18 @@ class EasyInputMessage(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class GeminiCandidate(BaseModel):
|
||||
citationMetadata: Optional[GeminiCitationMetadata] = None
|
||||
content: Optional[GeminiContent] = None
|
||||
finishReason: Optional[str] = None
|
||||
safetyRatings: Optional[List[GeminiSafetyRating]] = None
|
||||
class GeminiContent(BaseModel):
|
||||
parts: List[GeminiPart]
|
||||
role: Role1 = Field(..., examples=['user'])
|
||||
|
||||
|
||||
class GeminiGenerateContentResponse(BaseModel):
|
||||
candidates: Optional[List[GeminiCandidate]] = None
|
||||
promptFeedback: Optional[GeminiPromptFeedback] = None
|
||||
class GeminiGenerateContentRequest(BaseModel):
|
||||
contents: List[GeminiContent]
|
||||
generationConfig: Optional[GeminiGenerationConfig] = None
|
||||
safetySettings: Optional[List[GeminiSafetySetting]] = None
|
||||
systemInstruction: Optional[GeminiSystemInstructionContent] = None
|
||||
tools: Optional[List[GeminiTool]] = None
|
||||
videoMetadata: Optional[GeminiVideoMetadata] = None
|
||||
|
||||
|
||||
class ImagenGenerateImageRequest(BaseModel):
|
||||
@ -4893,7 +5318,7 @@ class ImagenGenerateImageRequest(BaseModel):
|
||||
|
||||
class InputMessage(BaseModel):
|
||||
content: Optional[InputMessageContentList] = None
|
||||
role: Optional[Role1] = None
|
||||
role: Optional[Role3] = None
|
||||
status: Optional[Status2] = None
|
||||
type: Optional[Type9] = None
|
||||
|
||||
@ -5102,6 +5527,18 @@ class ResponseProperties(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class GeminiCandidate(BaseModel):
|
||||
citationMetadata: Optional[GeminiCitationMetadata] = None
|
||||
content: Optional[GeminiContent] = None
|
||||
finishReason: Optional[str] = None
|
||||
safetyRatings: Optional[List[GeminiSafetyRating]] = None
|
||||
|
||||
|
||||
class GeminiGenerateContentResponse(BaseModel):
|
||||
candidates: Optional[List[GeminiCandidate]] = None
|
||||
promptFeedback: Optional[GeminiPromptFeedback] = None
|
||||
|
||||
|
||||
class InputItem(RootModel[Union[EasyInputMessage, Item]]):
|
||||
root: Union[EasyInputMessage, Item]
|
||||
|
||||
|
||||
445
comfy_api_nodes/nodes_gemini.py
Normal file
445
comfy_api_nodes/nodes_gemini.py
Normal file
@ -0,0 +1,445 @@
|
||||
"""
|
||||
API Nodes for Gemini Multimodal LLM Usage via Remote API
|
||||
See: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference
|
||||
"""
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from typing import Optional, Literal
|
||||
|
||||
import torch
|
||||
|
||||
import folder_paths
|
||||
from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
|
||||
from server import PromptServer
|
||||
from comfy_api_nodes.apis import (
|
||||
GeminiContent,
|
||||
GeminiGenerateContentRequest,
|
||||
GeminiGenerateContentResponse,
|
||||
GeminiInlineData,
|
||||
GeminiPart,
|
||||
GeminiMimeType,
|
||||
)
|
||||
from comfy_api_nodes.apis.client import (
|
||||
ApiEndpoint,
|
||||
HttpMethod,
|
||||
SynchronousOperation,
|
||||
)
|
||||
from comfy_api_nodes.apinode_utils import (
|
||||
validate_string,
|
||||
audio_to_base64_string,
|
||||
video_to_base64_string,
|
||||
tensor_to_base64_string,
|
||||
)
|
||||
|
||||
|
||||
GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini"
|
||||
GEMINI_MAX_INPUT_FILE_SIZE = 20 * 1024 * 1024 # 20 MB
|
||||
|
||||
|
||||
class GeminiModel(str, Enum):
|
||||
"""
|
||||
Gemini Model Names allowed by comfy-api
|
||||
"""
|
||||
|
||||
gemini_2_5_pro_preview_05_06 = "gemini-2.5-pro-preview-05-06"
|
||||
gemini_2_5_flash_preview_04_17 = "gemini-2.5-flash-preview-04-17"
|
||||
|
||||
|
||||
def get_gemini_endpoint(
|
||||
model: GeminiModel,
|
||||
) -> ApiEndpoint[GeminiGenerateContentRequest, GeminiGenerateContentResponse]:
|
||||
"""
|
||||
Get the API endpoint for a given Gemini model.
|
||||
|
||||
Args:
|
||||
model: The Gemini model to use, either as enum or string value.
|
||||
|
||||
Returns:
|
||||
ApiEndpoint configured for the specific Gemini model.
|
||||
"""
|
||||
if isinstance(model, str):
|
||||
model = GeminiModel(model)
|
||||
return ApiEndpoint(
|
||||
path=f"{GEMINI_BASE_ENDPOINT}/{model.value}",
|
||||
method=HttpMethod.POST,
|
||||
request_model=GeminiGenerateContentRequest,
|
||||
response_model=GeminiGenerateContentResponse,
|
||||
)
|
||||
|
||||
|
||||
class GeminiNode(ComfyNodeABC):
|
||||
"""
|
||||
Node to generate text responses from a Gemini model.
|
||||
|
||||
This node allows users to interact with Google's Gemini AI models, providing
|
||||
multimodal inputs (text, images, audio, video, files) to generate coherent
|
||||
text responses. The node works with the latest Gemini models, handling the
|
||||
API communication and response parsing.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
return {
|
||||
"required": {
|
||||
"prompt": (
|
||||
IO.STRING,
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"tooltip": "Text inputs to the model, used to generate a response. You can include detailed instructions, questions, or context for the model.",
|
||||
},
|
||||
),
|
||||
"model": (
|
||||
IO.COMBO,
|
||||
{
|
||||
"tooltip": "The Gemini model to use for generating responses.",
|
||||
"options": [model.value for model in GeminiModel],
|
||||
"default": GeminiModel.gemini_2_5_pro_preview_05_06.value,
|
||||
},
|
||||
),
|
||||
"seed": (
|
||||
IO.INT,
|
||||
{
|
||||
"default": 42,
|
||||
"min": 0,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"control_after_generate": True,
|
||||
"tooltip": "When seed is fixed to a specific value, the model makes a best effort to provide the same response for repeated requests. Deterministic output isn't guaranteed. Also, changing the model or parameter settings, such as the temperature, can cause variations in the response even when you use the same seed value. By default, a random seed value is used.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"images": (
|
||||
IO.IMAGE,
|
||||
{
|
||||
"default": None,
|
||||
"tooltip": "Optional image(s) to use as context for the model. To include multiple images, you can use the Batch Images node.",
|
||||
},
|
||||
),
|
||||
"audio": (
|
||||
IO.AUDIO,
|
||||
{
|
||||
"tooltip": "Optional audio to use as context for the model.",
|
||||
"default": None,
|
||||
},
|
||||
),
|
||||
"video": (
|
||||
IO.VIDEO,
|
||||
{
|
||||
"tooltip": "Optional video to use as context for the model.",
|
||||
"default": None,
|
||||
},
|
||||
),
|
||||
"files": (
|
||||
"GEMINI_INPUT_FILES",
|
||||
{
|
||||
"default": None,
|
||||
"tooltip": "Optional file(s) to use as context for the model. Accepts inputs from the Gemini Generate Content Input Files node.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
"auth_token": "AUTH_TOKEN_COMFY_ORG",
|
||||
"comfy_api_key": "API_KEY_COMFY_ORG",
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
DESCRIPTION = "Generate text responses with Google's Gemini AI model. You can provide multiple types of inputs (text, images, audio, video) as context for generating more relevant and meaningful responses."
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "api_call"
|
||||
CATEGORY = "api node/text/Gemini"
|
||||
|
||||
def get_parts_from_response(
|
||||
self, response: GeminiGenerateContentResponse
|
||||
) -> list[GeminiPart]:
|
||||
"""
|
||||
Extract all parts from the Gemini API response.
|
||||
|
||||
Args:
|
||||
response: The API response from Gemini.
|
||||
|
||||
Returns:
|
||||
List of response parts from the first candidate.
|
||||
"""
|
||||
return response.candidates[0].content.parts
|
||||
|
||||
def get_parts_by_type(
|
||||
self, response: GeminiGenerateContentResponse, part_type: Literal["text"] | str
|
||||
) -> list[GeminiPart]:
|
||||
"""
|
||||
Filter response parts by their type.
|
||||
|
||||
Args:
|
||||
response: The API response from Gemini.
|
||||
part_type: Type of parts to extract ("text" or a MIME type).
|
||||
|
||||
Returns:
|
||||
List of response parts matching the requested type.
|
||||
"""
|
||||
parts = []
|
||||
for part in self.get_parts_from_response(response):
|
||||
if part_type == "text" and hasattr(part, "text") and part.text:
|
||||
parts.append(part)
|
||||
elif (
|
||||
hasattr(part, "inlineData")
|
||||
and part.inlineData
|
||||
and part.inlineData.mimeType == part_type
|
||||
):
|
||||
parts.append(part)
|
||||
# Skip parts that don't match the requested type
|
||||
return parts
|
||||
|
||||
def get_text_from_response(self, response: GeminiGenerateContentResponse) -> str:
|
||||
"""
|
||||
Extract and concatenate all text parts from the response.
|
||||
|
||||
Args:
|
||||
response: The API response from Gemini.
|
||||
|
||||
Returns:
|
||||
Combined text from all text parts in the response.
|
||||
"""
|
||||
parts = self.get_parts_by_type(response, "text")
|
||||
return "\n".join([part.text for part in parts])
|
||||
|
||||
def create_video_parts(self, video_input: IO.VIDEO, **kwargs) -> list[GeminiPart]:
|
||||
"""
|
||||
Convert video input to Gemini API compatible parts.
|
||||
|
||||
Args:
|
||||
video_input: Video tensor from ComfyUI.
|
||||
**kwargs: Additional arguments to pass to the conversion function.
|
||||
|
||||
Returns:
|
||||
List of GeminiPart objects containing the encoded video.
|
||||
"""
|
||||
from comfy_api.util import VideoContainer, VideoCodec
|
||||
base_64_string = video_to_base64_string(
|
||||
video_input,
|
||||
container_format=VideoContainer.MP4,
|
||||
codec=VideoCodec.H264
|
||||
)
|
||||
return [
|
||||
GeminiPart(
|
||||
inlineData=GeminiInlineData(
|
||||
mimeType=GeminiMimeType.video_mp4,
|
||||
data=base_64_string,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
def create_audio_parts(self, audio_input: IO.AUDIO) -> list[GeminiPart]:
|
||||
"""
|
||||
Convert audio input to Gemini API compatible parts.
|
||||
|
||||
Args:
|
||||
audio_input: Audio input from ComfyUI, containing waveform tensor and sample rate.
|
||||
|
||||
Returns:
|
||||
List of GeminiPart objects containing the encoded audio.
|
||||
"""
|
||||
audio_parts: list[GeminiPart] = []
|
||||
for batch_index in range(audio_input["waveform"].shape[0]):
|
||||
# Recreate an IO.AUDIO object for the given batch dimension index
|
||||
audio_at_index = {
|
||||
"waveform": audio_input["waveform"][batch_index].unsqueeze(0),
|
||||
"sample_rate": audio_input["sample_rate"],
|
||||
}
|
||||
# Convert to MP3 format for compatibility with Gemini API
|
||||
audio_bytes = audio_to_base64_string(
|
||||
audio_at_index,
|
||||
container_format="mp3",
|
||||
codec_name="libmp3lame",
|
||||
)
|
||||
audio_parts.append(
|
||||
GeminiPart(
|
||||
inlineData=GeminiInlineData(
|
||||
mimeType=GeminiMimeType.audio_mp3,
|
||||
data=audio_bytes,
|
||||
)
|
||||
)
|
||||
)
|
||||
return audio_parts
|
||||
|
||||
def create_image_parts(self, image_input: torch.Tensor) -> list[GeminiPart]:
|
||||
"""
|
||||
Convert image tensor input to Gemini API compatible parts.
|
||||
|
||||
Args:
|
||||
image_input: Batch of image tensors from ComfyUI.
|
||||
|
||||
Returns:
|
||||
List of GeminiPart objects containing the encoded images.
|
||||
"""
|
||||
image_parts: list[GeminiPart] = []
|
||||
for image_index in range(image_input.shape[0]):
|
||||
image_as_b64 = tensor_to_base64_string(
|
||||
image_input[image_index].unsqueeze(0)
|
||||
)
|
||||
image_parts.append(
|
||||
GeminiPart(
|
||||
inlineData=GeminiInlineData(
|
||||
mimeType=GeminiMimeType.image_png,
|
||||
data=image_as_b64,
|
||||
)
|
||||
)
|
||||
)
|
||||
return image_parts
|
||||
|
||||
def create_text_part(self, text: str) -> GeminiPart:
|
||||
"""
|
||||
Create a text part for the Gemini API request.
|
||||
|
||||
Args:
|
||||
text: The text content to include in the request.
|
||||
|
||||
Returns:
|
||||
A GeminiPart object with the text content.
|
||||
"""
|
||||
return GeminiPart(text=text)
|
||||
|
||||
def api_call(
|
||||
self,
|
||||
prompt: str,
|
||||
model: GeminiModel,
|
||||
images: Optional[IO.IMAGE] = None,
|
||||
audio: Optional[IO.AUDIO] = None,
|
||||
video: Optional[IO.VIDEO] = None,
|
||||
files: Optional[list[GeminiPart]] = None,
|
||||
unique_id: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> tuple[str]:
|
||||
# Validate inputs
|
||||
validate_string(prompt, strip_whitespace=False)
|
||||
|
||||
# Create parts list with text prompt as the first part
|
||||
parts: list[GeminiPart] = [self.create_text_part(prompt)]
|
||||
|
||||
# Add other modal parts
|
||||
if images is not None:
|
||||
image_parts = self.create_image_parts(images)
|
||||
parts.extend(image_parts)
|
||||
if audio is not None:
|
||||
parts.extend(self.create_audio_parts(audio))
|
||||
if video is not None:
|
||||
parts.extend(self.create_video_parts(video))
|
||||
if files is not None:
|
||||
parts.extend(files)
|
||||
|
||||
# Create response
|
||||
response = SynchronousOperation(
|
||||
endpoint=get_gemini_endpoint(model),
|
||||
request=GeminiGenerateContentRequest(
|
||||
contents=[
|
||||
GeminiContent(
|
||||
role="user",
|
||||
parts=parts,
|
||||
)
|
||||
]
|
||||
),
|
||||
auth_kwargs=kwargs,
|
||||
).execute()
|
||||
|
||||
# Get result output
|
||||
output_text = self.get_text_from_response(response)
|
||||
if unique_id and output_text:
|
||||
PromptServer.instance.send_progress_text(output_text, node_id=unique_id)
|
||||
|
||||
return (output_text or "Empty response from Gemini model...",)
|
||||
|
||||
|
||||
class GeminiInputFiles(ComfyNodeABC):
|
||||
"""
|
||||
Loads and formats input files for use with the Gemini API.
|
||||
|
||||
This node allows users to include text (.txt) and PDF (.pdf) files as input
|
||||
context for the Gemini model. Files are converted to the appropriate format
|
||||
required by the API and can be chained together to include multiple files
|
||||
in a single request.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
"""
|
||||
For details about the supported file input types, see:
|
||||
https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference
|
||||
"""
|
||||
input_dir = folder_paths.get_input_directory()
|
||||
input_files = [
|
||||
f
|
||||
for f in os.scandir(input_dir)
|
||||
if f.is_file()
|
||||
and (f.name.endswith(".txt") or f.name.endswith(".pdf"))
|
||||
and f.stat().st_size < GEMINI_MAX_INPUT_FILE_SIZE
|
||||
]
|
||||
input_files = sorted(input_files, key=lambda x: x.name)
|
||||
input_files = [f.name for f in input_files]
|
||||
return {
|
||||
"required": {
|
||||
"file": (
|
||||
IO.COMBO,
|
||||
{
|
||||
"tooltip": "Input files to include as context for the model.",
|
||||
"options": input_files,
|
||||
"default": input_files[0] if input_files else None,
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"GEMINI_INPUT_FILES": (
|
||||
"GEMINI_INPUT_FILES",
|
||||
{
|
||||
"tooltip": "An optional additional file(s) to batch together with the file loaded from this node. Allows chaining of input files so that a single message can include multiple input files.",
|
||||
"default": None,
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
DESCRIPTION = "Loads and prepares input files to include as inputs for Gemini LLM nodes. The files will be read by the Gemini model when generating a response. The contents of the text file count toward the token limit. 🛈 TIP: Can be chained together with other Gemini Input File nodes."
|
||||
RETURN_TYPES = ("GEMINI_INPUT_FILES",)
|
||||
FUNCTION = "prepare_files"
|
||||
CATEGORY = "api node/text/Gemini"
|
||||
|
||||
def create_file_part(self, file_path: str) -> GeminiPart:
|
||||
mime_type = (
|
||||
GeminiMimeType.pdf
|
||||
if file_path.endswith(".pdf")
|
||||
else GeminiMimeType.text_plain
|
||||
)
|
||||
# Use base64 string directly, not the data URI
|
||||
with open(file_path, "rb") as f:
|
||||
file_content = f.read()
|
||||
import base64
|
||||
base64_str = base64.b64encode(file_content).decode("utf-8")
|
||||
|
||||
return GeminiPart(
|
||||
inlineData=GeminiInlineData(
|
||||
mimeType=mime_type,
|
||||
data=base64_str,
|
||||
)
|
||||
)
|
||||
|
||||
def prepare_files(
|
||||
self, file: str, GEMINI_INPUT_FILES: list[GeminiPart] = []
|
||||
) -> tuple[list[GeminiPart]]:
|
||||
"""
|
||||
Loads and formats input files for Gemini API.
|
||||
"""
|
||||
file_path = folder_paths.get_annotated_filepath(file)
|
||||
input_file_content = self.create_file_part(file_path)
|
||||
files = [input_file_content] + GEMINI_INPUT_FILES
|
||||
return (files,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GeminiNode": GeminiNode,
|
||||
"GeminiInputFiles": GeminiInputFiles,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GeminiNode": "Google Gemini",
|
||||
"GeminiInputFiles": "Gemini Input Files",
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user