fix(upscale): inline SeedVR2 step after a latent step decoded with swapped args and mishandled the PIL result
decode_latent_with_vae(vae, samples) returns a PIL image; the step passed (pipe_latent, vae) and then treated it as a tensor. Also: undefined `meta` (use create_image_metadata like the Florence2 branch), manifest insert(0) + upscale_produced (no duplicate base image), and chained re-encode no longer movedims (VAE.encode takes BHWC). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Jason Hoku
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
525733d70b
commit
85378c4e78
+18
-11
@@ -27,7 +27,7 @@ from .model_loader import (
|
||||
from .lora_utils import expand_lora_folder
|
||||
from .image_generation import (
|
||||
generate_image, flush_batch_with_vae, flush_batch_with_remote_vae,
|
||||
create_image_metadata, decode_latent_with_vae, calculate_eta, print_generation_progress
|
||||
create_image_metadata, decode_latent_with_vae, pil_to_image_array, calculate_eta, print_generation_progress
|
||||
)
|
||||
from .config_utils import sanitize_session_name
|
||||
from .html_generator import get_html_template
|
||||
@@ -2201,8 +2201,8 @@ def run_generation_loop(
|
||||
from PIL import Image as PILImage
|
||||
sv_config = ucfg.get("seedvr2", {})
|
||||
# Decode current latent to PIL for SeedVR2 input
|
||||
decoded = decode_latent_with_vae(pipe_latent, loaded_vae)
|
||||
pil_img = PILImage.fromarray((decoded[0].cpu().numpy() * 255).clip(0, 255).astype("uint8"))
|
||||
# decode_latent_with_vae(vae, latent_samples) already returns a PIL image
|
||||
pil_img = decode_latent_with_vae(loaded_vae, pipe_latent["samples"])
|
||||
result_pil, up_w, up_h, sv_duration = seedvr2_upscale(pil_img, sv_config)
|
||||
total_upscale_duration += sv_duration
|
||||
is_last_step = step_idx == len(expanded_steps) - 1
|
||||
@@ -2213,12 +2213,18 @@ def run_generation_loop(
|
||||
upscaled_filename = f"img_{upscale_id}.webp"
|
||||
filepath = os.path.join(paths["images"], upscaled_filename)
|
||||
result_pil.save(filepath, quality=80)
|
||||
upscaled_meta = dict(meta)
|
||||
# `meta` is not yet assigned at this point in the function
|
||||
# (same as the Florence2 branch): build via create_image_metadata.
|
||||
upscaled_meta = create_image_metadata(
|
||||
conf, up_w, up_h, sv_duration + total_upscale_duration,
|
||||
current_seed, batch_idx,
|
||||
actual_positive_prompt, actual_negative_prompt,
|
||||
gen_index=gen_index_offset + total_generated,
|
||||
)
|
||||
upscaled_meta.update({
|
||||
"id": upscale_id,
|
||||
"base_fp": job_fp,
|
||||
"config_fp": config_fingerprint(conf, w, h, conf["seed"], pipeline_name),
|
||||
"gen_index": len(existing_data["items"]),
|
||||
"file": f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images",
|
||||
"filename": upscaled_filename,
|
||||
"width": up_w, "height": up_h,
|
||||
@@ -2230,7 +2236,10 @@ def run_generation_loop(
|
||||
"upscale_model": sv_config.get("dit_model", ""),
|
||||
})
|
||||
image_metrics.finalize(upscaled_meta, result_pil)
|
||||
existing_data["items"].append(upscaled_meta)
|
||||
# insert(0) + upscale_produced: newest-first manifest order, and
|
||||
# without the flag the base image would ALSO be saved (duplicate)
|
||||
existing_data["items"].insert(0, upscaled_meta)
|
||||
upscale_produced = True
|
||||
save_manifest(paths["manifest"], existing_data)
|
||||
if PromptServer is not None:
|
||||
try:
|
||||
@@ -2242,12 +2251,10 @@ def run_generation_loop(
|
||||
pass
|
||||
upscale_combo_idx += 1
|
||||
else:
|
||||
# Re-encode for next step in chain
|
||||
# Re-encode for next step in chain (VAE.encode takes BHWC, no movedim)
|
||||
import torch as seedvr_torch
|
||||
import numpy as seedvr_np
|
||||
result_np = seedvr_np.array(result_pil.convert("RGB")).astype(seedvr_np.float32) / 255.0
|
||||
result_tensor = seedvr_torch.from_numpy(result_np).unsqueeze(0)
|
||||
pipe_latent = loaded_vae.encode(result_tensor[:, :, :, :3].movedim(-1, 1))
|
||||
result_tensor = seedvr_torch.from_numpy(pil_to_image_array(result_pil))
|
||||
pipe_latent = loaded_vae.encode(result_tensor[:, :, :, :3])
|
||||
pipe_latent = {"samples": pipe_latent}
|
||||
pipe_w_current = up_w
|
||||
pipe_h_current = up_h
|
||||
|
||||
@@ -1207,6 +1207,17 @@ def flush_batch_with_remote_vae(pending_batch, remote_vae_worker, existing_data,
|
||||
# to SDXL KSampler. The CurrentNodeContext below sets up V3-API execution but
|
||||
# does NOT include inference_mode — that's our job at the function boundary.
|
||||
@torch.inference_mode()
|
||||
def pil_to_image_array(pil_image):
|
||||
"""PIL image -> float32 numpy [1, H, W, 3] in [0, 1] (ComfyUI IMAGE layout, BHWC).
|
||||
|
||||
Wrap in torch.from_numpy() for VAE.encode() / node inputs. Same conversion as
|
||||
seedvr2_upscale() and the Florence2 inline branch use.
|
||||
"""
|
||||
import numpy as np
|
||||
arr = np.array(pil_image.convert("RGB")).astype(np.float32) / 255.0
|
||||
return arr[None, ...]
|
||||
|
||||
|
||||
def seedvr2_upscale(pil_image, seedvr2_config):
|
||||
"""
|
||||
Upscale an image using SeedVR2 diffusion-based upscaler.
|
||||
|
||||
@@ -131,6 +131,7 @@ _GEN_ORCH_DEPS = {
|
||||
"flush_batch_with_remote_vae": lambda *a, **kw: None,
|
||||
"create_image_metadata": lambda *a, **kw: {},
|
||||
"decode_latent_with_vae": lambda *a, **kw: None,
|
||||
"pil_to_image_array": lambda *a, **kw: None,
|
||||
"calculate_eta": lambda *a, **kw: 0.0,
|
||||
"print_generation_progress": lambda *a, **kw: None,
|
||||
},
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Inline SeedVR2 step: decode_latent_with_vae(vae, samples) returns a PIL image and
|
||||
the chained re-encode needs a BHWC float tensor (VAE.encode layout, no movedim)."""
|
||||
import importlib.util
|
||||
import inspect
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, _ROOT)
|
||||
|
||||
|
||||
def _load_helper():
|
||||
src = open(os.path.join(_ROOT, "image_generation.py"), encoding="utf-8").read()
|
||||
m = re.search(r"^def pil_to_image_array\(.*?(?=^def )", src, re.S | re.M)
|
||||
ns = {}
|
||||
exec(m.group(0), ns)
|
||||
return ns["pil_to_image_array"]
|
||||
|
||||
|
||||
def test_pil_to_image_array_is_bhwc_float01():
|
||||
f = _load_helper()
|
||||
img = Image.new("RGB", (5, 3), (255, 0, 51))
|
||||
arr = f(img)
|
||||
assert arr.shape == (1, 3, 5, 3) and arr.dtype == np.float32
|
||||
assert np.allclose(arr[0, 0, 0], [1.0, 0.0, 0.2])
|
||||
|
||||
|
||||
def test_pil_to_image_array_converts_non_rgb():
|
||||
arr = _load_helper()(Image.new("RGBA", (4, 2), (0, 255, 0, 128)))
|
||||
assert arr.shape == (1, 2, 4, 3)
|
||||
|
||||
|
||||
def test_orchestrator_seedvr2_branch_call_shape():
|
||||
src = open(os.path.join(_ROOT, "generation_orchestrator.py"), encoding="utf-8").read()
|
||||
i = src.index('if mode == "seedvr2":')
|
||||
branch = src[i:i + 6000]
|
||||
assert 'decode_latent_with_vae(loaded_vae, pipe_latent["samples"])' in branch
|
||||
assert "movedim" not in branch
|
||||
assert "dict(meta)" not in branch
|
||||
Reference in New Issue
Block a user