Jason HokuandClaude Opus 4.6
1e574d6927
Dashboard fast-path sort assumes newest items are at the start of the items array. append() put upscaled images at the end (oldest position), making them appear as the oldest items when sorting by newest. Changed both SeedVR2 and regular upscale paths from append() to insert(0, ...) so upscaled images appear at the top (newest position). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
467 lines
21 KiB
Python
467 lines
21 KiB
Python
"""
|
|
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
|
|
|
|
# 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.webp"
|
|
result_pil.save(os.path.join(images_dir, upscaled_filename), 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.webp"
|
|
upscaled_pil.save(os.path.join(images_dir, upscaled_filename), 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.webp"
|
|
upscale_result.save(os.path.join(images_dir, upscaled_filename), 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
|