auto-git:

[change] generate_equirect.py
This commit is contained in:
2026-05-07 12:03:22 +02:00
parent c5e8a74815
commit 52802ba7e7

View File

@@ -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,
return postprocess_image(
prompt,
gen_path,
output_path,
tempdir,
upscale=upscale,
steps=steps,
guidance=guidance,
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
seam_inpaint=False,
)
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
def main():