diff --git a/README.md b/README.md index 0cc517b..1580e13 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,8 @@ If you are in China or other regions requiring a proxy: - `cache_ttl_minutes`: Number of minutes Gemini retains cached context. - `cache_seed`: Generation seed used while `use_cache` is enabled. +With `use_seed` enabled, the selected nonzero seed is sent unchanged; zero chooses a fresh seed. Retry seeds come from a dedicated CPU `torch.Generator` for that execution, without reseeding global Python, NumPy, PyTorch, or CUDA RNGs. Fixed inputs produce a repeatable retry sequence. With `use_seed` disabled, no seed is sent, including image requests and retries. Input fingerprinting and cached-result reuse are unchanged. + ### Video Input Configuration Node - `video`: Native ComfyUI video input. diff --git a/gemini_nodes.py b/gemini_nodes.py index 4aa9deb..6d59598 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -559,26 +559,26 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): @classmethod def _handle_seed(cls, use_seed, seed): - actual_seed = None - if use_seed: - if seed == 0: - current_time = int(time.time() * 1000) - random_component = random.randint(0, 1000000) - actual_seed = (current_time + random_component) % 2147483647 - print(f"[INFO] Generated random seed: {actual_seed}") - else: - actual_seed = seed - print(f"[INFO] Using specified seed: {actual_seed}") - - random.seed(actual_seed) - np.random.seed(actual_seed) - torch.manual_seed(actual_seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(actual_seed) - else: + if not use_seed: print("[INFO] Seed not used") + return None, None - return actual_seed + generator = torch.Generator(device="cpu") + if seed == 0: + generator.seed() + actual_seed = cls._next_seed(generator) + print(f"[INFO] Generated random seed: {actual_seed}") + else: + generator.manual_seed(seed) + actual_seed = seed + print(f"[INFO] Using specified seed: {actual_seed}") + return actual_seed, generator + + @staticmethod + def _next_seed(generator): + if generator is None: + return None + return torch.randint(1, 2**31, (), generator=generator, device="cpu", dtype=torch.int64).item() @classmethod def _build_proxy_url(cls, proxy_host, proxy_port): @@ -869,7 +869,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): print(padded_system_instruction) api_seed = cache_seed if use_cache else seed - actual_seed = cls._handle_seed(use_seed, api_seed) + actual_seed, seed_generator = cls._handle_seed(use_seed, api_seed) input_seed = api_seed # Store the original API seed for retry cache key # Check if we have a cached successful gemini seed for this input seed @@ -879,6 +879,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): if cached_gemini_seed is not None: print(f"[INFO] Using cached successful gemini seed {cached_gemini_seed} for input seed {input_seed}") actual_seed = cached_gemini_seed + seed_generator.manual_seed(actual_seed) text_output = "" thoughts_output = "" @@ -1107,7 +1108,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): top_p=top_p, top_k=top_k, max_output_tokens=max_output_tokens, - seed=api_seed, + seed=actual_seed, include_images=include_images, response_modalities=response_modalities, aspect_ratio=aspect_ratio, @@ -1124,16 +1125,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): generate_content_config.response_mime_type = "application/json" generate_content_config.response_json_schema = schema_snapshot - if use_seed and actual_seed is not None: - try: - generate_content_config.seed = actual_seed - except Exception: - pass - elif use_cache: - try: - generate_content_config.seed = None - except Exception: - pass + generate_content_config.seed = actual_seed # API call in background thread start_time = time.time() @@ -1187,11 +1179,6 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): max_retries = 1 for attempt in range(max_retries): try: - if use_seed and actual_seed is not None: - try: - generate_content_config.seed = actual_seed + attempt - except Exception: - pass response = generate_once() if not (response.candidates and getattr(response.candidates[0].content, 'parts', None)): finish_reason = "UNKNOWN" @@ -1273,18 +1260,9 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): retry_attempt += 1 print(f"[INFO] Retry pattern matched in response. Retry attempt {retry_attempt}/{max_retries}") - # Generate new random gemini seed for retry - current_time = int(time.time() * 1000) - random_component = random.randint(0, 1000000) - actual_seed = (current_time + random_component) % 2147483647 - print(f"[INFO] Retrying with new gemini seed: {actual_seed}") - - # Update config seed and retry - if use_seed: - try: - generate_content_config.seed = actual_seed - except Exception: - pass + actual_seed = cls._next_seed(seed_generator) + generate_content_config.seed = actual_seed + print(f"[INFO] Retrying with generation seed: {actual_seed}") # Create new queue and thread for retry result_queue = queue.Queue() @@ -1309,18 +1287,9 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): retry_attempt += 1 print(f"[INFO] Retry pattern matched in error. Retry attempt {retry_attempt}/{max_retries}") - # Generate new random gemini seed for retry - current_time = int(time.time() * 1000) - random_component = random.randint(0, 1000000) - actual_seed = (current_time + random_component) % 2147483647 - print(f"[INFO] Retrying with new gemini seed: {actual_seed}") - - # Update config seed and retry - if use_seed: - try: - generate_content_config.seed = actual_seed - except Exception: - pass + actual_seed = cls._next_seed(seed_generator) + generate_content_config.seed = actual_seed + print(f"[INFO] Retrying with generation seed: {actual_seed}") # Create new queue and thread for retry result_queue = queue.Queue() diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index 1cf3931..7f5bb13 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -248,6 +248,56 @@ def test_cache_seed_explicitly_enables_seed_widget_control(): assert options["max"] == 2147483647 +@pytest.mark.parametrize("use_cache", [False, True]) +def test_seed_generator_is_local_and_drives_retries(structured_client, monkeypatch, use_cache): + node = gemini_nodes.SSL_GeminiTextPrompt + def forbidden(*args, **kwargs): + raise AssertionError("Seed generation must not use or reseed a global RNG") + monkeypatch.setattr(gemini_nodes.random, "randint", forbidden) + monkeypatch.setattr(gemini_nodes.random, "seed", forbidden) + monkeypatch.setattr(gemini_nodes.np.random, "seed", forbidden) + monkeypatch.setattr(torch, "manual_seed", forbidden) + monkeypatch.setattr(torch.cuda, "manual_seed_all", forbidden) + before = torch.get_rng_state().clone() + args = _execute_kwargs(_config(use_cache=use_cache, cache_seed=2**32 + 17)) + args.update(use_seed=True, seed=2**32 + 17, retry_pattern="retry", max_retries=1) + sequences = [] + for _ in range(2): + node._cache.clear() + node._seed_map_cache.clear() + structured_client.calls.clear() + structured_client.responses = [[("final", "retry")], [("final", "done")]] + result = node.execute(**args) + sequences.append([call["config"].seed for call in structured_client.calls]) + assert result[2] == sequences[-1][-1] + assert sequences[0] == sequences[1] + assert sequences[0][0] == 2**32 + 17 + assert 1 <= sequences[0][1] <= 2147483647 + assert torch.equal(before, torch.get_rng_state()) + + +def test_zero_seed_uses_independent_torch_generator(): + node = gemini_nodes.SSL_GeminiTextPrompt + before = torch.get_rng_state().clone() + seed, generator = node._handle_seed(True, 0) + assert isinstance(generator, torch.Generator) + assert generator.device.type == "cpu" + assert 1 <= seed <= 2147483647 + assert torch.equal(before, torch.get_rng_state()) + assert node._handle_seed(False, 0) == (None, None) + + +@pytest.mark.parametrize("use_cache", [False, True]) +def test_disabled_seed_stays_omitted_on_image_retry(structured_client, use_cache): + args = _execute_kwargs(_config(use_cache=use_cache, cache_seed=123)) + args.update(model="gemini-3-pro-image", include_images=True, use_seed=False, seed=42, retry_pattern="retry", max_retries=1) + structured_client.responses = [[("final", "retry")], [("final", "done")]] + result = gemini_nodes.SSL_GeminiTextPrompt.execute(**args) + assert len(structured_client.calls) == 2 + assert all(call["config"].seed is None for call in structured_client.calls) + assert result[2] == 0 + + def _execute_kwargs(config): return { "config": config, @@ -546,7 +596,7 @@ def test_context_cache_creation_failure_uses_original_request(monkeypatch): output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) assert output[0] == "ok" - assert captured["config"].seed == 2 + assert captured["config"].seed is None def test_missing_context_cache_is_recreated_once(monkeypatch):