diff --git a/generation_orchestrator.py b/generation_orchestrator.py index 1bc57e4..de42771 100644 --- a/generation_orchestrator.py +++ b/generation_orchestrator.py @@ -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 diff --git a/image_generation.py b/image_generation.py index cf25332..1f0c925 100644 --- a/image_generation.py +++ b/image_generation.py @@ -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. diff --git a/tests/conftest.py b/tests/conftest.py index d313211..bbe6367 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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, }, diff --git a/tests/test_seedvr2_inline.py b/tests/test_seedvr2_inline.py new file mode 100644 index 0000000..8735748 --- /dev/null +++ b/tests/test_seedvr2_inline.py @@ -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