132 lines
6.5 KiB
Python
132 lines
6.5 KiB
Python
import importlib.util
|
|
from fractions import Fraction
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
from types import ModuleType, SimpleNamespace
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import av
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from PIL import Image
|
|
from comfy_api.latest import InputImpl, Types
|
|
|
|
|
|
ROOT = Path(__file__).parents[1]
|
|
package = ModuleType("gemini_media_tests")
|
|
package.__path__ = [str(ROOT / "nodes" / "ai")]
|
|
sys.modules[package.__name__] = package
|
|
spec = importlib.util.spec_from_file_location("gemini_media_tests.FL_GeminiTextAPI", ROOT / "nodes" / "ai" / "FL_GeminiTextAPI.py")
|
|
gemini = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(gemini)
|
|
|
|
|
|
class GeminiMediaTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.node = gemini.FL_GeminiTextAPI()
|
|
self.client = mock.MagicMock()
|
|
self.client.__enter__.return_value = self.client
|
|
self.client.interactions.create.return_value = SimpleNamespace(output_text=" A caption ")
|
|
self.client.files.upload.side_effect = lambda **kw: SimpleNamespace(
|
|
name=Path(kw["file"]).name, uri="https://media/" + Path(kw["file"]).name,
|
|
state=SimpleNamespace(name="ACTIVE"))
|
|
self.factory = mock.patch.object(gemini.genai, "Client", return_value=self.client)
|
|
self.factory.start()
|
|
self.addCleanup(self.factory.stop)
|
|
|
|
def generate(self, **kwargs):
|
|
return self.node.generate_text("Describe", "test-key", "gemini-3.7-flash", 0.2, 8192, 42, **kwargs)
|
|
|
|
def test_text_only_contract(self):
|
|
self.assertEqual(self.generate(), ("A caption",))
|
|
self.assertEqual(self.client.interactions.create.call_args.kwargs["input"], "Describe")
|
|
self.client.files.upload.assert_not_called()
|
|
self.client.__exit__.assert_called_once()
|
|
|
|
def test_mixed_media_and_batches_preserve_content_and_clean_up(self):
|
|
audio = {"waveform": torch.tensor([[[0.1, -0.2], [0.3, -0.4]], [[0.2, 0.4], [0.1, 0.3]]]), "sample_rate": 48000}
|
|
images = torch.zeros(2, 8, 12, 3)
|
|
images[1, :, :, 0] = 1
|
|
video = mock.Mock()
|
|
paths = []
|
|
upload = self.client.files.upload.side_effect
|
|
|
|
def inspect_upload(**kwargs):
|
|
path = Path(kwargs["file"])
|
|
paths.append(path)
|
|
if path.suffix == ".flac":
|
|
samples, rate = sf.read(path, always_2d=True)
|
|
index = int(path.stem.split("-")[1]) - 1
|
|
np.testing.assert_allclose(samples, audio["waveform"][index].T.numpy(), atol=1e-6)
|
|
self.assertEqual(rate, 48000)
|
|
if path.name == "image-2.png":
|
|
self.assertEqual(Image.open(path).getpixel((0, 0)), (255, 0, 0))
|
|
return upload(**kwargs)
|
|
|
|
self.client.files.upload.side_effect = inspect_upload
|
|
self.assertEqual(self.generate(audio=audio, image=images, video=video), ("A caption",))
|
|
inputs = self.client.interactions.create.call_args.kwargs["input"]
|
|
self.assertEqual([part["type"] for part in inputs], ["text", "audio", "audio", "image", "image", "video"])
|
|
self.assertEqual(self.client.files.delete.call_count, 5)
|
|
self.assertFalse(any(path.exists() for path in paths))
|
|
video.save_to.assert_called_once()
|
|
|
|
def test_failed_generation_still_deletes_uploads(self):
|
|
self.client.interactions.create.side_effect = RuntimeError("Request failed")
|
|
self.assertEqual(self.generate(image=torch.zeros(1, 8, 8, 3)), ("Error: Request failed",))
|
|
self.client.files.delete.assert_called_once_with(name="image-1.png")
|
|
self.client.__exit__.assert_called_once()
|
|
|
|
def test_processing_failure_deletes_file_without_generation(self):
|
|
self.client.files.upload.side_effect = None
|
|
self.client.files.upload.return_value = SimpleNamespace(name="file", state=SimpleNamespace(name="PROCESSING"))
|
|
self.client.files.get.return_value = SimpleNamespace(name="file", state=SimpleNamespace(name="FAILED"))
|
|
with mock.patch.object(gemini.time, "sleep"):
|
|
self.assertIn("could not process", self.generate(image=torch.zeros(1, 8, 8, 3))[0])
|
|
self.client.interactions.create.assert_not_called()
|
|
self.client.files.delete.assert_called_once_with(name="file")
|
|
|
|
def test_cancellation_is_not_returned_as_caption(self):
|
|
self.client.interactions.create.side_effect = gemini.mm.InterruptProcessingException()
|
|
with self.assertRaises(gemini.mm.InterruptProcessingException):
|
|
self.generate(image=torch.zeros(1, 8, 8, 3))
|
|
self.client.files.delete.assert_called_once()
|
|
|
|
def test_processing_timeout_cleans_up(self):
|
|
self.client.files.upload.side_effect = None
|
|
self.client.files.upload.return_value = SimpleNamespace(name="file", state=SimpleNamespace(name="PROCESSING"))
|
|
with mock.patch.object(gemini.time, "monotonic", side_effect=[0, 301]):
|
|
self.assertIn("timed out", self.generate(image=torch.zeros(1, 8, 8, 3))[0])
|
|
self.client.interactions.create.assert_not_called()
|
|
self.client.files.delete.assert_called_once_with(name="file")
|
|
|
|
def test_invalid_audio_never_uploads(self):
|
|
self.assertIn("non-finite", self.generate(audio={"waveform": torch.full((1, 1, 8), float("nan")), "sample_rate": 48000})[0])
|
|
self.client.files.upload.assert_not_called()
|
|
|
|
def test_native_video_encoding_preserves_sound(self):
|
|
samples = torch.sin(torch.arange(48000) * (2 * np.pi * 440 / 48000))[None, None] * 0.1
|
|
video = InputImpl.VideoFromComponents(Types.VideoComponents(
|
|
images=torch.zeros(24, 32, 32, 3), frame_rate=Fraction(24),
|
|
audio={"waveform": samples, "sample_rate": 48000}))
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
files = self.node._media_files(Path(temporary), None, None, video)
|
|
with av.open(str(files[0][1])) as container:
|
|
self.assertEqual(len(container.streams.video), 1)
|
|
self.assertEqual(len(container.streams.audio), 1)
|
|
self.assertEqual(len(list(container.decode(video=0))), 24)
|
|
trimmed = InputImpl.VideoFromFile(str(files[0][1])).as_trimmed(start_time=0.25, duration=0.5, strict_duration=True)
|
|
trim_folder = Path(temporary) / "trim"
|
|
trim_folder.mkdir()
|
|
trimmed_files = self.node._media_files(trim_folder, None, None, trimmed)
|
|
with av.open(str(trimmed_files[0][1])) as container:
|
|
self.assertEqual(len(container.streams.audio), 1)
|
|
self.assertEqual(len(list(container.decode(video=0))), 12)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|