From 3a27bba9d17464f9cc4b69d22a0f12507ba7995b Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Wed, 4 Nov 2026 03:02:58 -0800 Subject: [PATCH 1/2] fix(distribution): master post-phase upscales this run's images per model group The distributed post-phase called run_deferred_upscales with the wrong arguments and raised on every run. Build the queue from this run's manifest items matched to their job configs, and upscale one model/LoRA/VAE group at a time with a fresh conditioning cache. HiRes prompt is now encoded when not cached. Co-Authored-By: Claude Fable 5.1 --- generation_orchestrator.py | 242 +++++++++++++++++++++++++---- tests/test_distributed_upscales.py | 184 ++++++++++++++++++++++ 2 files changed, 393 insertions(+), 33 deletions(-) create mode 100644 tests/test_distributed_upscales.py diff --git a/generation_orchestrator.py b/generation_orchestrator.py index 997f2ed..8ce165a 100644 --- a/generation_orchestrator.py +++ b/generation_orchestrator.py @@ -393,6 +393,11 @@ def run_deferred_upscales( Process deferred upscale jobs after generation completes. Runs all pipeline chains for each queued image. Called BEFORE finalization cleanup so VAE/model/CLIP are still loaded. + + Queue entries: config, width, height, seed, batch_idx, actual_positive_prompt, + actual_negative_prompt, gen_index, config_fp and optionally base_item (the exact + manifest item to upscale; skips the identity search). The caller supplies a + model/VAE/CLIP/conditioning_cache that matches every job in the queue. """ if not deferred_queue: return @@ -427,16 +432,18 @@ def run_deferred_upscales( actual_positive_prompt = job["actual_positive_prompt"] actual_negative_prompt = job["actual_negative_prompt"] - # Find the base image in the manifest by matching identity - base_item = None - for item in existing_data.get("items", []): - if (item.get("seed") == current_seed and - item.get("width") == pipe_w and - item.get("height") == pipe_h and - not item.get("upscaled") and - item.get("positive") == actual_positive_prompt): - base_item = item - break + # Distributed callers hand over the exact manifest item; local entries search by identity + base_item = job.get("base_item") + if base_item is None: + # Find the base image in the manifest by matching identity + for item in existing_data.get("items", []): + if (item.get("seed") == current_seed and + item.get("width") == pipe_w and + item.get("height") == pipe_h and + not item.get("upscaled") and + item.get("positive") == actual_positive_prompt): + base_item = item + break if not base_item: print(f"[GridTester] ⚠️ Deferred upscale #{job_idx+1}: base image not found in manifest, skipping") @@ -513,6 +520,15 @@ def run_deferred_upscales( else: adjusted_prompt = actual_positive_prompt hires_cond = conditioning_cache["positive"].get(adjusted_prompt) + if hires_cond is None: + # not cached (fresh per-model cache): encode it, keep the base conditioning on failure + try: + with pinned_clip(patched_clip): + hires_cond = encode_prompt_with_combinators(patched_clip, adjusted_prompt, conf.get("clip_skip", 0)) + conditioning_cache["positive"][adjusted_prompt] = hires_cond + except Exception as e: + print(f"[GridTester] ⚠️ Deferred upscale #{job_idx+1}: HiRes prompt encode failed: {e}, using base prompt") + hires_cond = None if hires_cond is not None: hires_positive_cond = hires_cond hires_prompt_active = True @@ -695,6 +711,150 @@ def run_deferred_upscales( print(f"[GridTester] ✅ Deferred upscale phase complete: {len(upscale_durations)}/{total_jobs} images upscaled") +# Item fields that describe THIS image rather than its config (stripped when a config is rebuilt from an item). +_ITEM_ONLY_KEYS = ( + "id", "file", "filename", "favorited", "rejected", "note", "duration", "gen_index", + "config_fp", "base_fp", "metrics", "upscaled", "base_width", "base_height", +) + + +def _conf_from_item(item): + """Best-effort job config rebuilt from a manifest item (fallback when no distributed job matches).""" + conf = {k: v for k, v in item.items() if k not in _ITEM_ONLY_KEYS} + conf["positive"] = item.get("config_positive", item.get("positive", "")) + conf["negative"] = item.get("config_negative", item.get("negative", "")) + conf.setdefault("lora_expanded", item.get("lora", "None")) + return conf + + +def _build_distributed_upscale_queue(items, pre_run_ids, jobs): + """ + Build the deferred-upscale queue for a distributed run. + + items: existing_data["items"] (newest first); pre_run_ids: ids present before the run + (only this run's images are upscaled); jobs: [(config, input_job)] of COMPLETED jobs. + Worker-inserted items carry only a /view URL (no filename, no config_fp), so each item is + matched back to its job config with the same rule resume uses (check_if_job_completed). + Returns entries shaped like the local deferred queue plus ``base_item``. + """ + from .config_utils import config_fingerprint + + by_key = {} + for cfg, ij in jobs: + by_key.setdefault((cfg.get("model"), cfg.get("sampler"), cfg.get("scheduler")), []).append((cfg, ij)) + + queue = [] + for item in reversed(items): # oldest first = generation order + if item.get("id") in pre_run_ids: + continue + if item.get("upscaled") or item.get("preview") or item.get("media_type") == "video": + continue + if item.get("model_type") == "ltx_video": + continue + if not item_filename(item): + print(f"[Distribution] ⚠️ Upscale: item {item.get('id')} has no filename, skipping") + continue + + w, h = item.get("width"), item.get("height") + seed = item.get("seed") + bidx = item.get("batch_idx", 0) + pos, neg = item.get("positive", ""), item.get("negative", "") + matches = [] + for cfg, _ij in by_key.get((item.get("model"), item.get("sampler"), item.get("scheduler")), []): + try: + if check_if_job_completed([item], cfg, seed, w, h, bidx, pos, neg) == 0: + matches.append(cfg) + except Exception: + continue + if matches: + cfg = next((c for c in matches if c.get("seed") == seed), matches[0]) + else: + cfg = _conf_from_item(item) + + queue.append({ + "config": cfg, "width": w, "height": h, "seed": seed, + "batch_idx": bidx, "actual_positive_prompt": pos, + "actual_negative_prompt": neg, "gen_index": item.get("gen_index", 0), + "config_fp": config_fingerprint(cfg, w, h, cfg.get("seed")), + "base_item": item, + }) + return queue + + +def _run_distributed_upscales( + queue, upscale_settings, session_settings, paths, existing_data, session_name, unique_id, + ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, + optional_positive, optional_negative, model_cache, incompatible_loras +): + """ + Run a distributed run's upscales on the master, one (model, LoRA, VAE) group at a time. + Each group loads its own model and gets a fresh conditioning cache (conditioning from another + model's CLIP is invalid), then everything is released before the next group. + """ + import comfy.model_management as mm + + def _norm_vae(c): + v = c.get("vae", "Default") + return "Default" if (not v or str(v) in ("None", "none")) else v + + groups = {} + for entry in queue: + c = entry["config"] + key = (get_model_cache_key(c), c.get("lora_expanded", c.get("lora", "None")), _norm_vae(c)) + groups.setdefault(key, []).append(entry) + + if optional_positive is not None or optional_negative is not None: + print("[Distribution] ℹ️ Positive/negative inputs are connected: upscale conditioning comes from the prompts, not those inputs") + + PromptServer = None + try: + from server import PromptServer as _PS + PromptServer = _PS + except Exception: + pass + + for (model_key, lora_str, vae_name), group in groups.items(): + if mm.processing_interrupted(): + print("[Distribution] 🛑 INTERRUPTED during distributed upscales") + break + cfg0 = group[0]["config"] + loaded_model = loaded_clip = loaded_vae = patched_model = patched_clip = None + cache = {"positive": {}, "negative": {}} + try: + loaded_model, loaded_clip, loaded_vae = load_model_by_type( + cfg0, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, + optional_positive, optional_negative, None, None, model_cache=model_cache + ) + if loaded_vae is None and optional_vae is not None: + loaded_vae = optional_vae + if vae_name != "Default": + if is_remote_vae(vae_name): + print(f"[Distribution] ℹ️ Remote VAE '{vae_name}': upscale encode/decode stays local with the model's VAE") + else: + loaded_vae = load_vae_by_name(vae_name) + patched_model, patched_clip, should_skip = load_loras( + loaded_model, loaded_clip, lora_str, model_key, incompatible_loras, model_cache=model_cache + ) + if should_skip: + print(f"[Distribution] ⚠️ LoRA incompatible for {model_key}; skipping {len(group)} upscales") + continue + if loaded_vae is None: + print(f"[Distribution] ⚠️ No local VAE for {model_key}; skipping {len(group)} upscales") + continue + run_deferred_upscales( + group, upscale_settings, session_settings, loaded_vae, patched_model, patched_clip, + cache, paths, existing_data, session_name, unique_id, PromptServer + ) + except Exception as e: + if isinstance(e, InterruptProcessingException): + raise + print(f"[Distribution] ⚠️ Upscale group {model_key} failed: {e}") + finally: + patched_model, patched_clip = cleanup_model_references(patched_model, patched_clip, cache) + loaded_model = loaded_clip = loaded_vae = None + gc.collect() + + def _expand_lora_weight_arrays(expanded_configs): """ Expand LoRA weight array bracket notation into Cartesian product of configs. @@ -838,7 +998,7 @@ def run_generation_loop( from .config_utils import ( parse_json_with_error, parse_float_input, parse_prompt_input_nested, derive_extra_seeds, - expand_configs, prepare_input_jobs + expand_configs, prepare_input_jobs, latent_job_pixel_size ) from .config_utils import ( JOB_ORDERS, build_job_plan, order_job_plan, count_switches, config_fingerprint, @@ -1627,6 +1787,12 @@ def run_generation_loop( else: model_switched = False + # A connected latent's pixel size = latent cells x this model's ratio, known only now. + if job.get("latent_cells"): + w, h = latent_job_pixel_size(job, latent_ratio) + if job_fp is not None: + job_fp = config_fingerprint(conf, w, h, conf["seed"]) + # ==== PER-CONFIG VAE SWITCHING ==== # Config Builder VAE settings take priority over sampler node's remote_vae_endpoint target_vae = conf.get("vae", "Default") @@ -3073,6 +3239,11 @@ def _run_distributed_generation( if image_format != "webp": print(f"[GridTester] 🖼️ Image save format (distributed master): {image_format}") + # Ids present before this run: only images made by this run get upscaled in the post-phase + _pre_run_ids = {it.get("id") for it in existing_data.get("items", [])} + if optional_latent is not None: + print("[Distribution] ⚠️ Latent input: remote workers render empty latents at the placeholder size") + manager.populate_jobs( expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode @@ -3339,6 +3510,9 @@ def _run_distributed_generation( latent_channels = get_latent_channels(loaded_model, optional_latent) latent_ratio = get_latent_downscale_ratio(loaded_model) if optional_latent is not None: + # pixel size of a connected latent = cells x this model's ratio + w = int(optional_latent["samples"].shape[3]) * latent_ratio + h = int(optional_latent["samples"].shape[2]) * latent_ratio latent_in = {"samples": optional_latent["samples"].clone()} else: latent_in = make_empty_latent(latent_channels, w, h, latent_ratio) @@ -3616,29 +3790,31 @@ def _run_distributed_generation( if session_settings and session_settings.get("upscaling", {}).get("enabled", False): upscale_settings = session_settings.get("upscaling", {}) print(f"[Distribution] 🔍 Running upscales on master for distributed results...") - # Build deferred upscale queue from all generated items - deferred_upscale_queue = [] - for item_meta in existing_data.get("items", []): - if not item_meta.get("upscaled"): # Skip items already upscaled - deferred_upscale_queue.append({ - "meta": item_meta, - "config": { - "model": item_meta.get("model", ckpt_name), - "steps": item_meta.get("steps", 20), - "sampler": item_meta.get("sampler", "euler"), - "scheduler": item_meta.get("scheduler", "normal"), - "cfg": item_meta.get("cfg", 7.0), - } - }) - if deferred_upscale_queue: - run_deferred_upscales( - deferred_upscale_queue, upscale_settings, existing_data, - session_name, paths, unique_id, ckpt_name, - conditioning_cache, model_cache, - use_remote_vae, remote_vae_endpoint + # Release the master's model first so the per-group loads below have the VRAM + patched_model, patched_clip = cleanup_model_references(patched_model, patched_clip, conditioning_cache) + conditioning_cache = {"positive": {}, "negative": {}} + loaded_model, loaded_clip, loaded_vae = None, None, None + gc.collect() + try: + with manager._lock: + _done_jobs = [(j.config, j.input_job) for j in manager._jobs.values() if j.state == JobState.COMPLETED] + deferred_upscale_queue = _build_distributed_upscale_queue( + existing_data.get("items", []), _pre_run_ids, _done_jobs ) - else: - print(f"[Distribution] ℹ️ No items to upscale") + if deferred_upscale_queue: + print(f"[Distribution] 🔍 {len(deferred_upscale_queue)} images from this run to upscale") + _run_distributed_upscales( + deferred_upscale_queue, upscale_settings, session_settings, paths, existing_data, + session_name, unique_id, ckpt_name, use_remote_vae, + optional_model, optional_clip, optional_vae, optional_positive, optional_negative, + model_cache, incompatible_loras + ) + else: + print(f"[Distribution] ℹ️ No items to upscale") + except InterruptProcessingException: + raise + except Exception as e: + print(f"[Distribution] Upscale phase failed: {e}") existing_data["meta"] = { "positive": positive_text, "negative": negative_text, diff --git a/tests/test_distributed_upscales.py b/tests/test_distributed_upscales.py new file mode 100644 index 0000000..1be1d16 --- /dev/null +++ b/tests/test_distributed_upscales.py @@ -0,0 +1,184 @@ +"""Distributed runs: the master's post-phase upscales this run's images, grouped per model/LoRA/VAE. + +REGRESSION (2026-11): the post-phase called run_deferred_upscales() with 11 positional arguments against a +12+ signature and entries without width/height, so it raised on every distributed run with upscaling on. +""" +import importlib.util +import inspect +import os +import re +import sys + +import pytest + +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _ROOT) + +import generation_orchestrator as go # noqa: E402 (conftest preloads the real module) +from config_utils import config_fingerprint # noqa: E402 + +_spec = importlib.util.spec_from_file_location("_mu_real_du", os.path.join(_ROOT, "manifest_utils.py")) +_mu = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mu) + +_REAL_RDU = go.run_deferred_upscales +LOCAL_KEYS = {"config", "width", "height", "seed", "batch_idx", "actual_positive_prompt", + "actual_negative_prompt", "gen_index", "config_fp"} + + +def _conf(model="a.safetensors", seed=1, lora="None", vae="Default", **kw): + c = {"model": model, "sampler": "euler", "scheduler": "normal", "steps": 8, "cfg": 5.0, + "denoise": 1.0, "seed": seed, "lora": lora, "lora_expanded": lora, "vae": vae, + "positive": "p", "negative": "n", "gguf_options": object()} + c.update(kw) + return c + + +def _item(i, conf, **kw): + it = {"id": i, "file": f"/view?filename=img_{i}.webp&type=output", "seed": conf["seed"], + "width": 768, "height": 768, "batch_idx": 0, "positive": "p", "negative": "n", + "model": conf["model"], "sampler": conf["sampler"], "scheduler": conf["scheduler"], + "steps": conf["steps"], "cfg": conf["cfg"], "denoise": conf["denoise"], "lora": conf["lora"], + "vae": conf.get("vae", "Default"), "config_positive": "cp", "config_negative": "cn"} + it.update(kw) + return it + + +@pytest.fixture(autouse=True) +def _real_item_filename(monkeypatch): + monkeypatch.setattr(go, "item_filename", _mu.item_filename) + + +def test_queue_only_this_runs_plain_items_oldest_first(): + c1, c2 = _conf(seed=1), _conf(seed=2) + jobs = [(c1, {}), (c2, {})] + items = [ # newest first + _item(5, c2), + _item(4, c1), + _item(3, c1, filename="x_upscaled.webp", upscaled=True), + _item(2, c1, file=""), # no filename + _item(1, c1), # pre-run + ] + q = go._build_distributed_upscale_queue(items, {1, 3}, jobs) + assert [e["base_item"]["id"] for e in q] == [4, 5] + for e in q: + assert set(e) == LOCAL_KEYS | {"base_item"} + assert q[0]["config"] is c1 and q[1]["config"] is c2 # job config by identity (gguf_options sentinel survives) + assert q[0]["config_fp"] == config_fingerprint(c1, 768, 768, 1) + + +def test_queue_skips_preview_video_ltx(): + c = _conf() + items = [_item(9, c, preview=True), _item(8, c, media_type="video"), _item(7, c, model_type="ltx_video")] + assert go._build_distributed_upscale_queue(items, set(), [(c, {})]) == [] + + +def test_queue_prefers_seed_sibling_and_handles_randomized_seed(): + a, b = _conf(seed=10), _conf(seed=20) + q = go._build_distributed_upscale_queue([_item(2, b), _item(1, a)], set(), [(a, {}), (b, {})]) + assert [e["config"] for e in q] == [a, b] + # randomized: the item carries a seed that is not the conf's; it still matches its job + r = _conf(seed=10, seed_behavior="randomize") + it = _item(3, r, seed=987654) + q = go._build_distributed_upscale_queue([it], set(), [(r, {})]) + assert q[0]["config"] is r and q[0]["seed"] == 987654 + + +def test_queue_fallback_conf_strips_item_fields(): + c = _conf() + it = _item(4, _conf(model="other.safetensors"), favorited=True, note="n", gen_index=3, duration=1.5) + q = go._build_distributed_upscale_queue([it], set(), [(c, {})]) + cfg = q[0]["config"] + for k in ("id", "file", "favorited", "note", "gen_index", "duration"): + assert k not in cfg + assert cfg["positive"] == "cp" and cfg["negative"] == "cn" + assert cfg["lora_expanded"] == "None" and cfg["model"] == "other.safetensors" + + +# ---- driver ------------------------------------------------------------------------------------------------- + +class _Rec: + def __init__(self, mp, vae_result="vae", skip_lora=()): + self.calls, self.loads = [], [] + self.skip_lora, self.vae_result = set(skip_lora), vae_result + mp.setattr(go, "load_model_by_type", self._load) + mp.setattr(go, "load_loras", self._loras) + mp.setattr(go, "load_vae_by_name", lambda n: f"vae:{n}") + mp.setattr(go, "cleanup_model_references", lambda m, c, cache: (None, None)) + mp.setattr(go, "run_deferred_upscales", lambda *a, **kw: self.calls.append(a)) + mp.setattr(go, "is_remote_vae", lambda v: str(v).startswith("remote")) + self.mm = sys.modules["comfy.model_management"] + mp.setattr(self.mm, "processing_interrupted", lambda: False, raising=False) + + def _load(self, conf, *a, **kw): + self.loads.append(conf["model"]) + return "m", "c", (self.vae_result if self.vae_result != "none" else None) + + def _loras(self, m, c, lora, key, inc, model_cache=None): + return "pm", "pc", lora in self.skip_lora + + +def _entry(conf): + return {"config": conf, "width": 8, "height": 8, "seed": 1, "batch_idx": 0, "actual_positive_prompt": "p", + "actual_negative_prompt": "n", "gen_index": 0, "config_fp": "x", "base_item": {}} + + +def _drive(queue): + return go._run_distributed_upscales(queue, {"pipelines": [1]}, {}, {"images": "i"}, {"items": []}, "s", "u", + "ck", False, None, None, None, None, None, {}, {}) + + +def test_driver_groups_per_model_lora_vae(monkeypatch): + rec = _Rec(monkeypatch) + q = [_entry(_conf("a")), _entry(_conf("b")), _entry(_conf("a")), _entry(_conf("a", lora="L"))] + _drive(q) + assert len(rec.calls) == 3 + assert [len(c[0]) for c in rec.calls] == [2, 1, 1] # first-appearance order + assert len({id(c[6]) for c in rec.calls}) == 3 # a fresh conditioning cache per group + n = len(inspect.signature(_REAL_RDU).parameters) + assert len(rec.calls[0]) == n - 1 # config_overrides_vae stays defaulted + assert rec.calls[0][3] == "vae" and rec.calls[0][4] == "pm" and rec.calls[0][5] == "pc" + + +def test_driver_skips_missing_vae_and_incompatible_lora(monkeypatch): + rec = _Rec(monkeypatch, vae_result="none") + _drive([_entry(_conf("a"))]) + assert rec.calls == [] + rec = _Rec(monkeypatch, skip_lora={"L"}) + _drive([_entry(_conf("a", lora="L")), _entry(_conf("b"))]) + assert len(rec.calls) == 1 and rec.calls[0][0][0]["config"]["model"] == "b" + + +def test_driver_named_vae_is_loaded(monkeypatch): + rec = _Rec(monkeypatch) + _drive([_entry(_conf("a", vae="my.vae"))]) + assert rec.calls[0][3] == "vae:my.vae" + + +def test_driver_stops_on_interrupt(monkeypatch): + rec = _Rec(monkeypatch) + monkeypatch.setattr(rec.mm, "processing_interrupted", lambda: True, raising=False) + _drive([_entry(_conf("a")), _entry(_conf("b"))]) + assert rec.calls == [] + + +# ---- source guards ------------------------------------------------------------------------------------------ + +def _src(name): + with open(os.path.join(_ROOT, "generation_orchestrator.py"), encoding="utf-8") as f: + text = f.read() + m = re.search(rf"\ndef {name}\(.*?(?=\n(?:def |class )|\Z)", text, re.S) + assert m, name + return m.group(0) + + +def test_post_phase_source(): + s = _src("_run_distributed_generation") + assert '"meta": item_meta' not in s + assert "_run_distributed_upscales(" in s and "_build_distributed_upscale_queue(" in s + assert s.index("_pre_run_ids =") < s.index("manager.populate_jobs(") + assert "run_deferred_upscales(" not in s + + +def test_deferred_uses_base_item(): + assert 'job.get("base_item")' in _src("run_deferred_upscales") From 0501686340c103390e003ac1e9d7f1c545c73176 Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Wed, 4 Nov 2026 03:03:01 -0800 Subject: [PATCH 2/2] fix(latent-input): connected latent size = cells x the model's latent ratio, not x 8 Co-Authored-By: Claude Fable 5.1 --- AI-CONTEXT.md | 4 +-- ProjectStructure.md | 4 +-- README.md | 8 ++++- config_utils.py | 22 +++++++++++-- tests/test_latent_input_dims.py | 58 +++++++++++++++++++++++++++++++++ 5 files changed, 88 insertions(+), 8 deletions(-) create mode 100644 tests/test_latent_input_dims.py diff --git a/AI-CONTEXT.md b/AI-CONTEXT.md index 98d4f30..48f147c 100644 --- a/AI-CONTEXT.md +++ b/AI-CONTEXT.md @@ -28,12 +28,12 @@ Read `ProjectStructure.md` in this directory for the full reference (file struct | `__init__.py` | 1597 | API endpoints (config CRUD, presets, dashboard save/export/delete, scan, async upscale, pairwise votes, metrics jobs, standalone dashboard page, run API, workflow sweep), node mappings, path security | | `sampler_node.py` | 401 | Main node class. Unwraps configs, extracts `_distribution` + `_session_settings`, delegates to orchestrator | | `config_builder_node.py` | 1578 | Config builder node. `state_to_configs_json(state)` is the ONLY state-to-configs implementation (`generate_config()` and `POST /configbuilder/preview` both call it). API endpoints for model lists, trigger lookups | -| `generation_orchestrator.py` | 3693 | **Largest file.** Main generation loop over a flat job plan (`config_utils.build_job_plan`), smart skip, preview pass, upscaling pipelines, deferred upscales (`run_deferred_upscales()`), GPU cooldown, distribution, ETA tracking, Start At Job # | +| `generation_orchestrator.py` | 3892 | **Largest file.** Main generation loop over a flat job plan (`config_utils.build_job_plan`), smart skip, preview pass, upscaling pipelines, deferred upscales (`run_deferred_upscales()`), GPU cooldown, distribution, ETA tracking, Start At Job # | | `image_generation.py` | 1321 | KSampler wrapper, VAE decode, `flush_batch_with_vae()` (saves images + updates manifest), `upscale_image()`, `create_image_metadata()`, rolling ETA | | `batch_encoding.py` | 878 | Prompt pre-encode with combinators, `pinned_clip()` (one text-encoder load per pass), `te_batch_size` text-encoder batching | | `civitai.py`, `distribution.py` | 40, 107 | Thin facades for the companion plugins (ComfyUI-USCG-CivitAI, ComfyUI-USCG-Distributed). No outbound network code in this package. | | `upscale_runner.py` | 642 | Dashboard async upscaling. Background thread loads models, runs pipeline chains, updates manifests, sends progress events. | -| `config_utils.py` | 1108 | Config expansion (Cartesian products), nested prompt parsing, job preparation, job plan / ordering, `config_fingerprint()` | +| `config_utils.py` | 1141 | Config expansion (Cartesian products), nested prompt parsing, job preparation, job plan / ordering, `config_fingerprint()` | | `model_loader.py` | 907 | Checkpoint/LoRA/VAE/GGUF loading, latent geometry (`get_latent_downscale_ratio`), CLIP type resolution | | `model_cache.py` | 935 | 3-tier cache: LoRA files → incremental states → patched models. Async preloading | | `manifest_utils.py` | 316 | Load/save/merge manifest. `save_manifest()` reloads from disk first to preserve concurrent user edits (favorites, votes, metrics) | diff --git a/ProjectStructure.md b/ProjectStructure.md index 0ff8101..b81f7a6 100644 --- a/ProjectStructure.md +++ b/ProjectStructure.md @@ -38,7 +38,7 @@ ComfyUI-Ultimate-Sampler-Grid/ ├── config_builder_node.py # "Config Builder" Node + API endpoints + state_to_configs_json (1578 lines) ├── json_text_node.py # "Smart JSON" Node (60 lines) │ -├── generation_orchestrator.py # The "Conductor" (Loop Logic, job plan, preview pass, 3693 lines) +├── generation_orchestrator.py # The "Conductor" (Loop Logic, job plan, preview pass, 3892 lines) ├── image_generation.py # KSampler & VAE Wrappers (1321 lines) ├── image_metrics.py # Optional per-image metrics (sharpness/entropy/...), run context, scorer-pack adapters (ADAPTERS registry: node ids/signatures pinned to pack shas, loader defaults + `prefer`, `scale`, `max_batch`, `vram_mb`; `adapter_status()` reports `problem`) (668 lines) ├── metrics_runner.py # Background job: compute metrics for an existing session (CPU, manifest lock) (167 lines) @@ -52,7 +52,7 @@ ComfyUI-Ultimate-Sampler-Grid/ ├── conditioning_cache.py # Disk-based Conditioning Cache (372 lines) ├── lora_utils.py # LoRA Searching & Validation (468 lines) ├── trigger_words.py # Trigger Word Prompt Assembly (377 lines) -├── config_utils.py # Config Parsing, Expansion, Cartesian Products, job plan + fingerprints (1108 lines) +├── config_utils.py # Config Parsing, Expansion, Cartesian Products, job plan + fingerprints (1141 lines) ├── directory_scanner.py # External Directory Session Scanner (609 lines) ├── metadata_packer.py # PNG/WebP Metadata Embedding for Exports (561 lines) ├── manifest_utils.py # Manifest Read/Write/Merge Helpers, comparisons (316 lines) diff --git a/README.md b/README.md index b1dc463..7b64dfa 100644 --- a/README.md +++ b/README.md @@ -528,7 +528,7 @@ The Generator node features built-in widgets for Model Selection and Prompts, bu - Connect `optional_positive` and `optional_negative` for pre-encoded conditioning - The node automatically detects latent channel count (4 for SD1.5/SDXL, 16 for SD3/Flux/Z-Image) - Skip detection uses conditioning tensor hashing to properly detect prompt changes -* **Latent Input:** Connect a `LATENT` to use img2img or upscaling workflows. +* **Latent Input:** Connect a `LATENT` to use img2img or upscaling workflows. The image size is the latent size times the model's latent ratio (8 px per cell for SD/SDXL/Flux.1, 16 for Flux.2/MOSS-Image/Krea 2), so the manifest and upscale target sizes are right for every family. - For SD3/Flux: Use `EmptySD3LatentImage` instead of `EmptyLatentImage` - Latent dimensions are automatically preserved @@ -1634,6 +1634,8 @@ Master (your machine): 6. Master saves everything to the session manifest ``` +**Upscales:** with upscaling enabled, the master upscales after all workers have finished, only the images made by this run, one model/LoRA/VAE group at a time (each group loads its own model; worker images are matched back to their job config). **Latent input:** remote workers never receive the latent and render an empty latent at the placeholder size; only the master honours it. + ### Master Text Encoding When enabled, the master pre-encodes ALL unique positive/negative prompts for every (model, LoRA) combination before distributing jobs. Workers receive the encoded conditionings in their job claim response and skip CLIP encoding entirely. @@ -1775,6 +1777,10 @@ After modifying backend Python files, restart ComfyUI. After modifying `web/` JS ## 📝 Changelog +### Update 11/4/26 (fixes): Distributed upscales, latent-input size +* 🌐 **Distributed upscales work:** the post-run upscale phase on the master used to crash on every distributed run with upscaling on. It now upscales this run's images (master and worker made), per model/LoRA/VAE group, with HiRes-prompt adjustment encoded per model. +* 📐 **Latent input size:** a connected latent's image size was always cells x 8; it is now cells x the loaded model's latent ratio (16 for Flux.2/MOSS-Image/Krea 2). + ### Update 11/4/26: Standalone dashboard, pairwise review, image metrics, preview pass, workflow sweeps * 📱 **Standalone Dashboard + Run API:** open the dashboard as a normal web page at `/config_tester/dashboard?session=` (phone, second monitor) with live updates over its own websocket. `POST /config_tester/run` and `GET /config_tester/run_status` start and track headless runs (Discord bot, curl) without any workflow JSON. * ⚖️ **Pairwise Review:** blind A/B compare mode (`C` key) with matched, same-settings, one-axis and random pairing, Bradley-Terry ratings (Sort by Rating) and a per-setting scoreboard with 95% intervals. Votes are stored in the session manifest under `comparisons`. diff --git a/config_utils.py b/config_utils.py index 9078b7c..b135648 100644 --- a/config_utils.py +++ b/config_utils.py @@ -785,13 +785,18 @@ def expand_configs(raw_configs, pos_prompts, neg_prompts, denoise_values, seed, return expanded -def prepare_input_jobs(optional_latent, resolutions): +def prepare_input_jobs(optional_latent, resolutions, downscale_ratio=8): """ Prepare input jobs from either optional latent or resolution list. + + For a latent input, width/height are a placeholder (cells x downscale_ratio) until the + model is known; "latent_cells" carries (cells_w, cells_h) and the orchestrator resolves + the real pixel size per job with latent_job_pixel_size() once that job's model is loaded. Args: optional_latent: Optional latent tensor input resolutions: List of [width, height] pairs + downscale_ratio: placeholder pixels per latent cell (default 8) Returns: List of job dictionaries with label, width, height, latent, batch_idx @@ -802,10 +807,12 @@ def prepare_input_jobs(optional_latent, resolutions): batch_count = optional_latent["samples"].shape[0] for i in range(batch_count): single_sample = optional_latent["samples"][i].unsqueeze(0) + cells_w, cells_h = int(single_sample.shape[3]), int(single_sample.shape[2]) input_jobs.append({ "label": f"Input {i+1}", - "width": single_sample.shape[3] * 8, - "height": single_sample.shape[2] * 8, + "width": cells_w * downscale_ratio, + "height": cells_h * downscale_ratio, + "latent_cells": (cells_w, cells_h), "latent": {"samples": single_sample}, "batch_idx": i }) @@ -822,6 +829,15 @@ def prepare_input_jobs(optional_latent, resolutions): return input_jobs +def latent_job_pixel_size(job, downscale_ratio): + """(w, h) in pixels of an input job. Latent-input jobs: cells x the model's latent ratio (8 SD/SDXL/Flux.1, 16 Flux.2/MOSS/Krea 2); resolution jobs: unchanged.""" + cells = job.get("latent_cells") if isinstance(job, dict) else None + if not cells: + return job["width"], job["height"] + r = int(downscale_ratio) if downscale_ratio else 8 + return int(cells[0]) * r, int(cells[1]) * r + + def sanitize_session_name(session_name): """Sanitize session name to be filesystem-safe""" session_name = re.sub(r'[^\w\-]', '', session_name) diff --git a/tests/test_latent_input_dims.py b/tests/test_latent_input_dims.py new file mode 100644 index 0000000..b8855dd --- /dev/null +++ b/tests/test_latent_input_dims.py @@ -0,0 +1,58 @@ +"""A connected latent's pixel size is cells x the model's latent ratio, not cells * 8.""" +import os +import re +import sys + +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _ROOT) + +from config_utils import prepare_input_jobs, latent_job_pixel_size # noqa: E402 + + +class T: + def __init__(self, shape): + self.shape = tuple(shape) + + def __getitem__(self, i): + return T(self.shape[1:]) + + def unsqueeze(self, d): + return T((1,) + self.shape) + + +def test_prepare_input_jobs_carries_cells(): + jobs = prepare_input_jobs({"samples": T((2, 128, 76, 52))}, []) + assert len(jobs) == 2 + assert [j["batch_idx"] for j in jobs] == [0, 1] + assert all(j["latent_cells"] == (52, 76) for j in jobs) + assert (jobs[0]["width"], jobs[0]["height"]) == (416, 608) + assert prepare_input_jobs({"samples": T((1, 4, 8, 8))}, [], 16)[0]["width"] == 128 + + +def test_latent_job_pixel_size(): + j = prepare_input_jobs({"samples": T((1, 128, 76, 52))}, [])[0] + assert latent_job_pixel_size(j, 16) == (832, 1216) + assert latent_job_pixel_size(j, 8) == (416, 608) + assert latent_job_pixel_size(j, None) == (416, 608) + r = prepare_input_jobs(None, [[512, 768]])[0] + assert latent_job_pixel_size(r, 16) == (512, 768) + + +def _read(n): + with open(os.path.join(_ROOT, n), encoding="utf-8") as f: + return f.read() + + +def test_no_hardcoded_8_in_prepare_input_jobs(): + body = re.search(r"def prepare_input_jobs\(.*?\ndef ", _read("config_utils.py"), re.S).group(0) + assert "* 8" not in body + + +def test_orchestrator_resolves_size_after_ratio(): + s = _read("generation_orchestrator.py") + a = s.index("latent_ratio = get_latent_downscale_ratio(loaded_model)") + b = s.index("latent_job_pixel_size(job, latent_ratio)") + c = s.index("make_empty_latent(latent_channels, gen_w") + assert a < b < c + # distributed master too + assert 'optional_latent["samples"].shape[3]) * latent_ratio' in s