diff --git a/generate_equirect.py b/generate_equirect.py index 2b27c87..30b8b4f 100644 --- a/generate_equirect.py +++ b/generate_equirect.py @@ -122,10 +122,12 @@ def clear_torch_cache(device: str | None = None) -> None: empty_cache() -def configure_pipeline_memory(pipe) -> None: +def configure_pipeline_memory(pipe, *, vae_tiling: bool = True) -> None: pipe.enable_attention_slicing() - if hasattr(pipe, "enable_vae_tiling"): + if vae_tiling and hasattr(pipe, "enable_vae_tiling"): pipe.enable_vae_tiling() + elif not vae_tiling and hasattr(pipe, "disable_vae_tiling"): + pipe.disable_vae_tiling() if hasattr(pipe, "enable_vae_slicing"): pipe.enable_vae_slicing() @@ -357,7 +359,7 @@ def postprocess_image( safety_checker=None, requires_safety_checker=False, ).to(device) - configure_pipeline_memory(inpaint_pipe) + configure_pipeline_memory(inpaint_pipe, vae_tiling=False) print("→ Inpainting seam for seamless tiling…") progress_cb("inpaint", 0, steps)