88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
"""
|
|
Decode saved stepwise latents into panorama image strips using StepwiseLatentPreview.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import glob
|
|
import time
|
|
from config import MODEL_ORDER
|
|
from utils import check_comfyui, queue_prompt, wait_for_completion
|
|
|
|
|
|
MODEL_PREFIXES = {tag: tag.lower().replace("_", "") for tag in MODEL_ORDER}
|
|
# Q2_K -> q2k, Q3_K_S -> q3ks, Q4_0 -> q40, etc.
|
|
# Override the ones that don't follow the pattern:
|
|
MODEL_PREFIXES["Q2_K"] = "q2k"
|
|
MODEL_PREFIXES["Q3_K_S"] = "q3ks"
|
|
MODEL_PREFIXES["Q4_0"] = "q4_0"
|
|
MODEL_PREFIXES["Q4_1"] = "q4_1"
|
|
MODEL_PREFIXES["Q5_0"] = "q5_0"
|
|
MODEL_PREFIXES["Q5_1"] = "q5_1"
|
|
MODEL_PREFIXES["Q6_K"] = "q6k"
|
|
MODEL_PREFIXES["Q8_0"] = "q8_0"
|
|
|
|
|
|
def find_stepwise_dir(base_dir, prefix, seed):
|
|
pattern = os.path.join(base_dir, f"stepwise_{prefix}_s{seed}_*")
|
|
dirs = sorted(d for d in glob.glob(pattern) if os.path.isdir(d))
|
|
return dirs[-1] if dirs else None
|
|
|
|
|
|
def build_preview_workflow(save_dir, filename_prefix, steps="1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20"):
|
|
return {
|
|
"1": {"class_type": "VAELoader",
|
|
"inputs": {"vae_name": "ae.safetensors"}},
|
|
"2": {"class_type": "StepwiseLatentPreview",
|
|
"inputs": {"vae": ["1", 0], "save_dir": save_dir,
|
|
"preview_steps": steps, "use_x0": True}},
|
|
"3": {"class_type": "SaveImage",
|
|
"inputs": {"images": ["2", 0], "filename_prefix": filename_prefix}},
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generate panorama strips from stepwise latents")
|
|
parser.add_argument("--stepwise-dir", required=True, help="ComfyUI output dir with stepwise_* folders")
|
|
parser.add_argument("--output-dir", required=True)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
check_comfyui()
|
|
|
|
available = []
|
|
for tag in MODEL_ORDER:
|
|
prefix = MODEL_PREFIXES[tag]
|
|
d = find_stepwise_dir(args.stepwise_dir, prefix, args.seed)
|
|
if d:
|
|
available.append((tag, prefix, d))
|
|
print(f" [OK] {tag}: {os.path.basename(d)}")
|
|
else:
|
|
print(f" [--] {tag}: not found")
|
|
|
|
if not available:
|
|
print("\nNo stepwise data found. Run generate_stepwise.py first.")
|
|
return
|
|
|
|
print(f"\nGenerating {len(available)} panoramas\n")
|
|
|
|
for i, (tag, prefix, save_dir) in enumerate(available):
|
|
fname = f"panorama_{prefix}_s{args.seed}"
|
|
print(f"[{i+1}/{len(available)}] {tag} -> {fname} ...", end=" ", flush=True)
|
|
|
|
wf = build_preview_workflow(save_dir, fname)
|
|
t0 = time.time()
|
|
try:
|
|
prompt_id = queue_prompt(wf)
|
|
ok = wait_for_completion(prompt_id, timeout=300)
|
|
print(f"{'ok' if ok else 'FAIL'} ({time.time()-t0:.1f}s)")
|
|
except Exception as e:
|
|
print(f"ERROR: {e}")
|
|
|
|
print("\nDone! Check output directory for panorama_* files.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|