Merge branch 'feat/orch-fixes' into int/wave3
This commit is contained in:
Jason Hoku
+2
-2
@@ -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) |
|
||||
|
||||
+2
-2
@@ -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)
|
||||
|
||||
@@ -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=<name>` (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`.
|
||||
|
||||
+19
-3
@@ -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)
|
||||
|
||||
+209
-33
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user