auto-git:
[change] generate_equirect.py
This commit is contained in:
@@ -485,7 +485,6 @@ def generate(
|
|||||||
|
|
||||||
device = select_device()
|
device = select_device()
|
||||||
is_sdxl = "sdxl" in model_path.lower()
|
is_sdxl = "sdxl" in model_path.lower()
|
||||||
scale = guidance # keep inpaint guidance in sync with cfg guidance
|
|
||||||
enable_upscale = bool(upscale and upscale != "none")
|
enable_upscale = bool(upscale and upscale != "none")
|
||||||
|
|
||||||
os.makedirs(work_dir, exist_ok=True)
|
os.makedirs(work_dir, exist_ok=True)
|
||||||
@@ -613,23 +612,11 @@ def generate(
|
|||||||
f"Unsupported scheduler '{scheduler}'. "
|
f"Unsupported scheduler '{scheduler}'. "
|
||||||
"Try one of: euler, euler_a, heun, ddim, dpmsolver, dpmsolver-sde."
|
"Try one of: euler, euler_a, heun, ddim, dpmsolver, dpmsolver-sde."
|
||||||
)
|
)
|
||||||
gen_pipe.enable_attention_slicing()
|
configure_pipeline_memory(gen_pipe)
|
||||||
if hasattr(gen_pipe, "enable_vae_tiling"):
|
|
||||||
gen_pipe.enable_vae_tiling()
|
|
||||||
if hasattr(gen_pipe, "enable_vae_slicing"):
|
|
||||||
gen_pipe.enable_vae_slicing()
|
|
||||||
if "pipe_kwargs" in locals():
|
if "pipe_kwargs" in locals():
|
||||||
del pipe_kwargs
|
del pipe_kwargs
|
||||||
|
|
||||||
def progress_cb(phase: str, current: int, total: int):
|
progress_cb = make_progress_cb(enable_upscale, seam_inpaint)
|
||||||
payload = {
|
|
||||||
"phase": phase,
|
|
||||||
"current": current,
|
|
||||||
"total": total,
|
|
||||||
"upscale": enable_upscale,
|
|
||||||
"seamInpaint": seam_inpaint,
|
|
||||||
}
|
|
||||||
print(f"PROGRESS {json.dumps(payload)}", flush=True)
|
|
||||||
|
|
||||||
print("→ Generating equirectangular HDRI…")
|
print("→ Generating equirectangular HDRI…")
|
||||||
progress_cb("gen", 0, steps)
|
progress_cb("gen", 0, steps)
|
||||||
@@ -659,71 +646,35 @@ def generate(
|
|||||||
image.save(gen_path)
|
image.save(gen_path)
|
||||||
print(f"→ Saved initial image to {gen_path}")
|
print(f"→ Saved initial image to {gen_path}")
|
||||||
|
|
||||||
seamless_path = os.path.join(tempdir, os.path.basename(output_path))
|
|
||||||
if seam_inpaint:
|
if seam_inpaint:
|
||||||
shift_amt = width // 2
|
del image
|
||||||
mask_w = width // 8
|
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)
|
return postprocess_image(
|
||||||
mask = create_mask(width, height, mask_w)
|
prompt,
|
||||||
|
gen_path,
|
||||||
inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
|
output_path,
|
||||||
INPAINT_MODEL,
|
tempdir,
|
||||||
torch_dtype=torch.float32
|
upscale=upscale,
|
||||||
).to(device)
|
steps=steps,
|
||||||
inpaint_pipe.enable_attention_slicing()
|
guidance=guidance,
|
||||||
|
|
||||||
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,
|
width=width,
|
||||||
height=height,
|
height=height,
|
||||||
callback_steps=1,
|
seam_inpaint=False,
|
||||||
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
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
|||||||
Reference in New Issue
Block a user