refactor: use local torch generator for API seeds
This commit is contained in:
silveroxides
@@ -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.
|
||||
|
||||
+28
-59
@@ -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()
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user