nodes/JasonHoku-ComfyUI-Ultimate-Auto-Sampler-Config-Grid-Testing-Suite
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
a3ad11116aec3a70aa703edf320e455e5dee9ea7
JasonHoku-ComfyUI-Ultimate-…/upscale_runner.py
T
Jason HokuandClaude Opus 4.7 1f5adeb799 fix: extend image_format plumbing to remote VAE / upscales / distributed 
Commit 490c519 added a webp/png/jpg dropdown but only flush_batch_with_vae
honored it. Three other save sites still hardcoded WEBP:
- flush_batch_with_remote_vae (image_generation.py)
- upscale_runner.py (3 sites in inline + deferred upscale paths)
- distribution_worker.py (PIL bytes shipped back to master)

Now all save sites read image_format from session_settings (orchestrator)
or the job spec (distributed) and match the format/extension.

Changes:
- remote_vae.py: RemoteVAEDecodeWorker accepts image_format kwarg; _worker
  uses it to pick format/extension at save time (the real save site for the
  remote VAE path — flush_batch_with_remote_vae only queues to the worker)
- generation_orchestrator.py: passes image_format to all RemoteVAEDecodeWorker
  constructors (initialize_remote_vae + 3 per-config sites); adds image_format
  to all 8 existing_data["meta"] writes so upscale_runner can read it from disk
- upscale_runner.py: reads image_format from manifest meta at thread start;
  uses _fmt/_ext locals at all 3 save sites (seedvr2, latent-decode, PIL)
- distribution_manager.py: injects _session_image_format into worker job
  config alongside existing _session_cooldown passthrough
- distribution_worker.py: reads _session_image_format from config; encodes
  bytes in the chosen format; stamps _image_ext on meta so master uses it
- distribution_routes.py: pops _image_ext from meta to build filename with
  correct extension before writing bytes to disk

Default behavior unchanged: all paths fall back to webp if image_format
is absent or unrecognised.

Code review found this in 490c519.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 13:29:52 -07:00

492 lines
23 KiB
Python
Raw Blame History

"""
Dashboard Upscale Runner
Background thread for upscaling images triggered from the dashboard.
Loads models via ComfyUI APIs, runs upscale pipeline chains,
saves results and updates manifests.
"""
import os
import time
import json
import threading
import uuid
import torch
import numpy as np
from PIL import Image as PILImage
from .model_loader import load_checkpoint, load_vae_by_name
from .batch_encoding import encode_prompt_with_combinators
from .image_generation import upscale_image, decode_latent_with_vae, create_image_metadata
from .manifest_utils import save_manifest
import comfy.model_management as mm
# Active upscale jobs (module-level state)
_upscale_jobs = {}
def start_upscale_job(session_name, image_ids, upscale_config, all_favorited=False):
"""
Start an async upscale job. Returns (job_id, error) immediately.
The actual work runs in a background thread.
"""
import folder_paths
benchmarks_dir = os.path.realpath(os.path.join(folder_paths.get_output_directory(), "benchmarks"))
session_dir = os.path.join(benchmarks_dir, session_name)
manifest_path = os.path.join(session_dir, "manifest.json")
images_dir = os.path.join(session_dir, "images")
if not os.path.exists(manifest_path):
return None, "Manifest not found"
with open(manifest_path, "r", encoding="utf-8") as f:
manifest_data = json.load(f)
items = manifest_data.get("items", [])
meta = manifest_data.get("meta", {})
# Resolve image list
if all_favorited:
target_items = [item for item in items if item.get("favorited")]
else:
id_set = set(image_ids) if image_ids else set()
target_items = [item for item in items if item.get("id") in id_set]
if not target_items:
return None, "No matching images found"
job_id = f"upscale_{uuid.uuid4().hex[:12]}"
job = {
"id": job_id,
"status": "running",
"completed": 0,
"total": len(target_items),
"current_image": "",
"error": None,
"cancel": threading.Event(),
}
_upscale_jobs[job_id] = job
thread = threading.Thread(
target=_run_upscale_thread,
args=(job, target_items, upscale_config, meta, manifest_data, manifest_path, images_dir, session_name),
daemon=True
)
thread.start()
return job_id, None
def get_upscale_status(job_id):
"""Get status of an upscale job."""
job = _upscale_jobs.get(job_id)
if not job:
return {"status": "not_found"}
return {
"status": job["status"],
"completed": job["completed"],
"total": job["total"],
"current_image": job.get("current_image", ""),
"error": job.get("error"),
}
def cancel_upscale_job(job_id):
"""Signal an upscale job to stop."""
job = _upscale_jobs.get(job_id)
if job and job["status"] == "running":
job["cancel"].set()
return True
return False
def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, manifest_path, images_dir, session_name):
"""
Background thread that runs the upscale pipeline on each image.
Loads the session's model/VAE/CLIP, encodes prompts, runs pipeline chains.
"""
import itertools as up_itertools
import random as up_random
try:
pipelines = upscale_config.get("pipelines", [])
if not pipelines:
job["status"] = "error"
job["error"] = "No upscale pipelines configured"
return
# Resolve image format from manifest meta (written by orchestrator)
_raw_fmt = meta.get("image_format", "webp")
_fmt = (_raw_fmt or "webp").lower().strip()
if _fmt not in ("webp", "png", "jpg", "jpeg"):
_fmt = "webp"
_ext = "jpg" if _fmt in ("jpg", "jpeg") else _fmt
# Load model, VAE, CLIP for conditioning
model_name = meta.get("model", "")
vae_name = meta.get("vae", "")
print(f"[DashboardUpscale] 🔄 Loading model: {model_name}")
loaded_model, loaded_clip, loaded_vae = load_checkpoint(
target_model_name=model_name,
ckpt_name=model_name, # Same as target since we're loading directly
use_remote_vae=False,
)
if vae_name and vae_name != "None" and vae_name != "Default":
print(f"[DashboardUpscale] 🔄 Loading VAE: {vae_name}")
loaded_vae = load_vae_by_name(vae_name)
patched_model = loaded_model
patched_clip = loaded_clip
# Conditioning cache for this job
conditioning_cache = {"positive": {}, "negative": {}}
# Get PromptServer for sending events
PromptServer = None
try:
import server as comfy_server
PromptServer = comfy_server.PromptServer
# Ensure last_prompt_id exists — KSampler progress callbacks require it
if PromptServer.instance is not None and not hasattr(PromptServer.instance, 'last_prompt_id'):
PromptServer.instance.last_prompt_id = f"dashboard_upscale_{job['id']}"
except Exception:
pass
for img_idx, item in enumerate(target_items):
if job["cancel"].is_set():
print(f"[DashboardUpscale] 🛑 Cancelled — clearing VRAM...")
try:
mm.soft_empty_cache()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
job["status"] = "cancelled"
return
# Extract filename from manifest item — may be in "filename" or parsed from "file" URL
filename = item.get("filename", "")
if not filename and item.get("file"):
import re
fname_match = re.search(r'filename=([^&]+)', item["file"])
if fname_match:
filename = fname_match.group(1)
job["current_image"] = filename
image_path = os.path.join(images_dir, filename)
if not os.path.exists(image_path):
print(f"[DashboardUpscale] ⚠️ File not found: {filename}, skipping")
job["completed"] += 1
continue
# Load image and VAE-encode to latent
try:
pil_image = PILImage.open(image_path).convert("RGB")
img_array = np.array(pil_image).astype(np.float32) / 255.0
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
encoded = loaded_vae.encode(img_tensor[:, :, :, :3])
result_latent = {"samples": encoded}
except Exception as e:
print(f"[DashboardUpscale] ⚠️ Failed to load/encode {filename}: {e}")
job["completed"] += 1
continue
# Get conditioning
pos_prompt = item.get("positive", meta.get("positive", ""))
neg_prompt = item.get("negative", meta.get("negative", ""))
clip_skip = item.get("clip_skip", 0)
if pos_prompt not in conditioning_cache["positive"]:
conditioning_cache["positive"][pos_prompt] = encode_prompt_with_combinators(patched_clip, pos_prompt, clip_skip)
if neg_prompt not in conditioning_cache["negative"]:
conditioning_cache["negative"][neg_prompt] = encode_prompt_with_combinators(patched_clip, neg_prompt, clip_skip)
final_positive = conditioning_cache["positive"][pos_prompt]
final_negative = conditioning_cache["negative"][neg_prompt]
# HiRes prompt adjustment
hires_positive_cond = final_positive
hires_prompt_active = False
hires_prompt_behavior_rt = ""
hires_prompt_text_rt = ""
if upscale_config.get("hires_prompt_adjust") and upscale_config.get("hires_prompt_text", "").strip():
hires_prompt_behavior_rt = upscale_config.get("hires_prompt_behavior", "append_end")
hires_prompt_text_rt = upscale_config["hires_prompt_text"].strip()
if hires_prompt_behavior_rt == "prepend":
adjusted = hires_prompt_text_rt + " " + pos_prompt
elif hires_prompt_behavior_rt == "append_end":
adjusted = pos_prompt + " " + hires_prompt_text_rt
elif hires_prompt_behavior_rt == "replace":
adjusted = hires_prompt_text_rt
else:
adjusted = pos_prompt
if adjusted not in conditioning_cache["positive"]:
conditioning_cache["positive"][adjusted] = encode_prompt_with_combinators(patched_clip, adjusted, clip_skip)
hires_positive_cond = conditioning_cache["positive"][adjusted]
hires_prompt_active = True
pipe_w = item.get("width", 512)
pipe_h = item.get("height", 512)
upscale_combo_idx = 0
total_upscale_duration = 0
for pipeline_idx, pipeline in enumerate(pipelines):
if pipeline.get("active", True) is False:
continue
pipeline_name = pipeline.get("name", f"Pipeline {pipeline_idx + 1}")
pipeline_steps = pipeline.get("steps", [])
if not pipeline_steps:
continue
pipe_latent = result_latent
pipe_w_current = pipe_w
pipe_h_current = pipe_h
# Flatten steps with repeat
expanded_steps = []
for step in pipeline_steps:
if step.get("active", True) is False:
continue
repeat = max(1, int(step.get("repeat", 1)))
for _ in range(repeat):
expanded_steps.append(step)
for step_idx, ucfg in enumerate(expanded_steps):
mode = ucfg.get("mode", "hires_only")
show_hires = mode in ("hires_only", "model_then_hires")
show_model = mode in ("model_only", "model_then_hires")
# --- SeedVR2 upscale mode ---
if mode == "seedvr2":
from .image_generation import seedvr2_upscale
sv_config = ucfg.get("seedvr2", {})
# Convert source PIL image directly (no latent decoding needed)
pil_input = PILImage.open(os.path.join(images_dir, filename))
result_pil, up_w, up_h, sv_duration = seedvr2_upscale(pil_input, sv_config)
total_upscale_duration += sv_duration
is_last_step = step_idx == len(expanded_steps) - 1
if is_last_step:
# Save the upscaled image
upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000)
upscaled_filename = f"img_{upscale_id}_upscaled.{_ext}"
_up_path = os.path.join(images_dir, upscaled_filename)
if _fmt == "png":
result_pil.save(_up_path, format="PNG")
elif _fmt in ("jpg", "jpeg"):
result_pil.save(_up_path, format="JPEG", quality=95)
else:
result_pil.save(_up_path, format="WEBP", quality=95)
# Create manifest entry
upscaled_meta = {
k: v for k, v in item.items()
if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration",
"upscale_source", "upscale_pipeline", "upscale_mode",
"upscale_ratio", "upscale_denoise", "upscale_model")
}
upscaled_meta.update({
"id": upscale_id,
"gen_index": len(manifest_data["items"]),
"file": f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images",
"filename": upscaled_filename,
"width": up_w, "height": up_h,
"duration": round(sv_duration, 2),
"upscaled": True,
"upscale_source": "dashboard",
"upscale_pipeline": pipeline_name,
"upscale_mode": "seedvr2",
"upscale_model": sv_config.get("dit_model", ""),
})
manifest_data["items"].insert(0, upscaled_meta)
upscale_combo_idx += 1
else:
# For chained steps, the next step gets the PIL result
# Re-read from saved temp or keep in memory
pil_input = result_pil
pipe_w_current = up_w
pipe_h_current = up_h
continue # Skip the normal combo loop
raw_ratios = str(ucfg.get("upscale_ratios", "1.5"))
ratios = [float(r.strip()) for r in raw_ratios.split(",") if r.strip()] or [1.5]
raw_denoise = str(ucfg.get("hires_denoise", "0.3"))
denoises = [float(d.strip()) for d in raw_denoise.split(",") if d.strip()] or [0.3]
models = ucfg.get("upscale_models", []) or [""]
if show_hires and show_model:
combos = list(up_itertools.product(models, ratios, denoises))
elif show_hires:
combos = list(up_itertools.product([""], ratios, denoises))
elif show_model:
combos = list(up_itertools.product(models, [1.0], [0.0]))
else:
combos = []
for combo in combos:
up_model, up_ratio, up_denoise = combo
single_config = {
"mode": mode, "upscale_model": up_model,
"upscale_ratio": up_ratio, "hires_denoise": up_denoise,
"hires_steps": ucfg.get("hires_steps", 0),
"tiled_vae": ucfg.get("tiled_vae", False),
"tile_size": ucfg.get("tile_size", 512),
"upscale_size": ucfg.get("upscale_size", "2.0"),
"resize_method": ucfg.get("resize_method", "bilinear"),
"hires_tiled_sampling": ucfg.get("hires_tiled_sampling", False),
"hires_tile_width": ucfg.get("hires_tile_width", 512),
"hires_tile_height": ucfg.get("hires_tile_height", 512),
"hires_mask_blur": ucfg.get("hires_mask_blur", 8),
"hires_tile_padding": ucfg.get("hires_tile_padding", 32),
"hires_force_uniform_tiles": ucfg.get("hires_force_uniform_tiles", False)
}
up_positive = hires_positive_cond if (show_hires and hires_prompt_active) else final_positive
upscale_result, upscale_duration = upscale_image(
pipe_latent, loaded_vae, patched_model, single_config,
item, up_positive, final_negative, pipe_w_current, pipe_h_current
)
is_last_step = step_idx == len(expanded_steps) - 1
is_last_combo = combo == combos[-1]
is_final_output = is_last_step and is_last_combo
if isinstance(upscale_result, dict) and "samples" in upscale_result:
upscaled_pil = decode_latent_with_vae(loaded_vae, upscale_result["samples"])
up_w, up_h = upscaled_pil.size
if is_final_output:
upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000)
upscaled_filename = f"img_{upscale_id}_upscaled.{_ext}"
_up_path = os.path.join(images_dir, upscaled_filename)
if _fmt == "png":
upscaled_pil.save(_up_path, format="PNG")
elif _fmt in ("jpg", "jpeg"):
upscaled_pil.save(_up_path, format="JPEG", quality=95)
else:
upscaled_pil.save(_up_path, format="WEBP", quality=95)
elif isinstance(upscale_result, PILImage.Image):
up_w, up_h = upscale_result.size
if is_final_output:
upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000)
upscaled_filename = f"img_{upscale_id}_upscaled.{_ext}"
_up_path = os.path.join(images_dir, upscaled_filename)
if _fmt == "png":
upscale_result.save(_up_path, format="PNG")
elif _fmt in ("jpg", "jpeg"):
upscale_result.save(_up_path, format="JPEG", quality=95)
else:
upscale_result.save(_up_path, format="WEBP", quality=95)
else:
continue
if not is_final_output:
if isinstance(upscale_result, dict) and "samples" in upscale_result:
pipe_latent = upscale_result
pipe_w_current = up_w
pipe_h_current = up_h
continue
# Create manifest entry — copy source item fields, override with upscale info
upscaled_meta = {
k: v for k, v in item.items()
if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration",
"upscale_source", "upscale_pipeline", "upscale_mode",
"upscale_ratio", "upscale_denoise", "upscale_model")
}
# gen_index: place upscaled images at the end (newest) since they were just created
# Use current manifest length so each new upscale gets the next sequential position
upscale_gen_index = len(manifest_data["items"])
upscaled_meta.update({
"id": upscale_id,
"gen_index": upscale_gen_index,
"file": f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images",
"filename": upscaled_filename,
"width": up_w,
"height": up_h,
"duration": round(upscale_duration, 2),
"upscaled": True,
"upscale_source": "dashboard",
"upscale_pipeline": pipeline_name,
"upscale_mode": mode,
"upscale_ratio": up_ratio,
"upscale_denoise": up_denoise,
})
if up_model:
upscaled_meta["upscale_model"] = up_model
if hires_prompt_active:
upscaled_meta["hires_prompt_behavior"] = hires_prompt_behavior_rt
upscaled_meta["hires_prompt_text"] = hires_prompt_text_rt
manifest_data["items"].insert(0, upscaled_meta)
upscale_combo_idx += 1
if upscale_combo_idx > 0:
save_manifest(manifest_path, manifest_data)
# Send all new items to dashboard (may be multiple from different pipelines/combos)
if PromptServer is not None:
try:
new_items = manifest_data["items"][-upscale_combo_idx:]
PromptServer.instance.send_sync("ultimate_grid.update_data", {
"session_name": session_name,
"new_items": new_items
})
except Exception:
pass
job["completed"] = img_idx + 1
# Send progress
if PromptServer is not None:
try:
PromptServer.instance.send_sync("ultimate_grid.progress", {
"session_name": session_name,
"current_job": img_idx + 1,
"total_jobs": len(target_items),
"progress_pct": int(((img_idx + 1) / len(target_items)) * 100),
"eta_str": "",
"avg_duration": 0,
"last_duration": 0,
"phase": "upscaling",
})
except Exception:
pass
print(f"[DashboardUpscale] 🔄 Upscaled {img_idx+1}/{len(target_items)}: {filename}")
# Done — clear VRAM and RAM
print(f"[DashboardUpscale] 🧹 Clearing VRAM and RAM...")
try:
mm.soft_empty_cache()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception as cleanup_err:
print(f"[DashboardUpscale] ⚠️ Cleanup warning: {cleanup_err}")
job["status"] = "complete"
print(f"[DashboardUpscale] ✅ Complete: {job['completed']}/{job['total']} images upscaled")
except Exception as e:
import traceback
traceback.print_exc()
job["status"] = "error"
job["error"] = str(e)
print(f"[DashboardUpscale] ❌ Error: {e}")
finally:
# Always attempt cleanup on any exit path
try:
mm.soft_empty_cache()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
Reference in New Issue View Git Blame Copy Permalink