diff --git a/generate_equirect.py b/generate_equirect.py index 2fe8ff7..17e651e 100644 --- a/generate_equirect.py +++ b/generate_equirect.py @@ -485,7 +485,6 @@ def generate( device = select_device() is_sdxl = "sdxl" in model_path.lower() - scale = guidance # keep inpaint guidance in sync with cfg guidance enable_upscale = bool(upscale and upscale != "none") os.makedirs(work_dir, exist_ok=True) @@ -613,23 +612,11 @@ def generate( f"Unsupported scheduler '{scheduler}'. " "Try one of: euler, euler_a, heun, ddim, dpmsolver, dpmsolver-sde." ) - gen_pipe.enable_attention_slicing() - if hasattr(gen_pipe, "enable_vae_tiling"): - gen_pipe.enable_vae_tiling() - if hasattr(gen_pipe, "enable_vae_slicing"): - gen_pipe.enable_vae_slicing() + configure_pipeline_memory(gen_pipe) if "pipe_kwargs" in locals(): del pipe_kwargs - def progress_cb(phase: str, current: int, total: int): - payload = { - "phase": phase, - "current": current, - "total": total, - "upscale": enable_upscale, - "seamInpaint": seam_inpaint, - } - print(f"PROGRESS {json.dumps(payload)}", flush=True) + progress_cb = make_progress_cb(enable_upscale, seam_inpaint) print("→ Generating equirectangular HDRI…") progress_cb("gen", 0, steps) @@ -659,71 +646,35 @@ def generate( image.save(gen_path) print(f"→ Saved initial image to {gen_path}") - seamless_path = os.path.join(tempdir, os.path.basename(output_path)) if seam_inpaint: - shift_amt = width // 2 - mask_w = width // 8 + del image + clear_torch_cache(device) + restart_for_postprocess( + prompt, + gen_path, + output_path, + tempdir, + work_dir, + upscale, + steps, + guidance, + width, + height, + seam_inpaint, + ) - shifted = shift_image(image, shift_amt) - mask = create_mask(width, height, mask_w) - - inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained( - INPAINT_MODEL, - torch_dtype=torch.float32 - ).to(device) - inpaint_pipe.enable_attention_slicing() - - print("→ Inpainting seam for seamless tiling…") - progress_cb("inpaint", 0, steps) - inpainted = inpaint_pipe( - prompt=prompt, - image=shifted, - mask_image=mask, - num_inference_steps=steps, - guidance_scale=scale, - width=width, - height=height, - callback_steps=1, - callback=lambda step, timestep, kwargs: progress_cb("inpaint", step + 1, steps), - ).images[0] - - inpainted = unshift_image(inpainted, shift_amt) - inpainted.save(seamless_path) - print(f"→ Crafted seamless image: {seamless_path}") - final_source = inpainted - else: - image.save(seamless_path) - print(f"→ Using raw output (seam inpaint disabled): {seamless_path}") - final_source = image - - final_path = seamless_path - - if upscale and upscale != "none": - try: - if upscale is True or upscale == "topaz": - final_path = run_topaz(seamless_path, tempdir) - elif upscale == "realesrgan": - final_path = run_realesrgan( - final_source, - tempdir, - scale=REALESRGAN_SCALE, - model_path=REALESRGAN_MODEL, - progress_cb=progress_cb - ) - else: - raise ValueError(f"Unknown upscale option '{upscale}'") - except Exception as e: # noqa: BLE001 - print(f"Upscaling failed ({upscale}); keeping seamless image: {e}") - - with Image.open(final_path) as final_img: - final_img.load() - save_png_with_prompt(final_img, output_path, prompt) - try: - with Image.open(output_path) as _im: - print(f"→ Final image written to {output_path} [{_im.size[0]}x{_im.size[1]}]") - except Exception: - print(f"→ Final image written to {output_path}") - return output_path + return postprocess_image( + prompt, + gen_path, + output_path, + tempdir, + upscale=upscale, + steps=steps, + guidance=guidance, + width=width, + height=height, + seam_inpaint=False, + ) def main():