58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
import asyncio
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from comfy.cli_args import args
|
|
|
|
if not torch.cuda.is_available():
|
|
args.cpu = True
|
|
|
|
from comfy_api_nodes.apis.gemini import ( # noqa: E402
|
|
GeminiCandidate,
|
|
GeminiContent,
|
|
GeminiGenerateContentResponse,
|
|
GeminiInlineData,
|
|
GeminiPart,
|
|
)
|
|
from comfy_api_nodes.nodes_gemini import get_image_from_response # noqa: E402
|
|
|
|
|
|
def image_part(mode, color):
|
|
buffer = BytesIO()
|
|
Image.new(mode, (4, 4), color).save(buffer, format="PNG")
|
|
return GeminiPart(
|
|
inlineData=GeminiInlineData(
|
|
data=base64.b64encode(buffer.getvalue()).decode(),
|
|
mimeType="image/png",
|
|
)
|
|
)
|
|
|
|
|
|
def response(*parts):
|
|
return GeminiGenerateContentResponse(
|
|
candidates=[GeminiCandidate(content=GeminiContent(parts=list(parts), role="model"))]
|
|
)
|
|
|
|
|
|
def test_rgb_only_response_stays_three_channels():
|
|
out = asyncio.run(get_image_from_response(response(image_part("RGB", (10, 20, 30)))))
|
|
assert out.shape == (1, 4, 4, 3)
|
|
|
|
|
|
def test_mixed_rgb_and_rgba_parts_are_padded_to_the_same_width():
|
|
out = asyncio.run(
|
|
get_image_from_response(
|
|
response(
|
|
image_part("RGB", (10, 20, 30)),
|
|
image_part("RGBA", (10, 20, 30, 0)),
|
|
)
|
|
)
|
|
)
|
|
assert out.shape == (2, 4, 4, 4)
|
|
# the part that had no alpha is padded opaque, the transparent one is preserved
|
|
assert out[0, ..., 3].min() == 1.0
|
|
assert out[1, ..., 3].max() == 0.0
|