142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import os
|
|
import subprocess
|
|
import torch
|
|
from PIL import Image, ImageDraw
|
|
import shutil
|
|
import tempfile
|
|
|
|
from diffusers import (
|
|
StableDiffusionPipeline,
|
|
DPMSolverMultistepScheduler,
|
|
StableDiffusionInpaintPipeline
|
|
)
|
|
|
|
def shift_image(img: Image.Image, shift: int) -> Image.Image:
|
|
w, h = img.size
|
|
out = Image.new("RGB", (w, h))
|
|
out.paste(img.crop((shift, 0, w, h)), (0, 0))
|
|
out.paste(img.crop((0, 0, shift, h)), (w - shift, 0))
|
|
return out
|
|
|
|
def create_mask(width: int, height: int, mask_w: int) -> Image.Image:
|
|
mask = Image.new("L", (width, height), 0)
|
|
draw = ImageDraw.Draw(mask)
|
|
left = (width - mask_w) // 2
|
|
draw.rectangle([left, 0, left + mask_w, height], fill=255)
|
|
return mask
|
|
|
|
def unshift_image(img: Image.Image, shift: int) -> Image.Image:
|
|
w, h = img.size
|
|
out = Image.new("RGB", (w, h))
|
|
out.paste(img.crop((w - shift, 0, w, h)), (0, 0))
|
|
out.paste(img.crop((0, 0, w - shift, h)), (shift, 0))
|
|
return out
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate an equirectangular HDRI, make it seamless, and upscale it with Topaz Photo AI CLI."
|
|
)
|
|
parser.add_argument("--prompt", required=True,
|
|
help="Text prompt for generation and inpainting")
|
|
parser.add_argument("--output", required=True,
|
|
help="Filename for the final upscaled image (e.g. seamless.png)")
|
|
parser.add_argument("--work-dir", default=os.path.dirname(os.path.abspath(__file__)),
|
|
help="Working directory for intermediates and final outputs")
|
|
args = parser.parse_args()
|
|
|
|
# Output-Ordner (bleibt wie gehabt)
|
|
output_abs = os.path.abspath(args.output)
|
|
|
|
# Zwischenschritte landen im eigenem temp-Ordner:
|
|
with tempfile.TemporaryDirectory(dir=args.work_dir) as tempdir:
|
|
print(f"→ Using tempdir: {tempdir}")
|
|
|
|
model_path = "/Volumes/SD/ML-Models/diffusers/hdri-panorama-v1-diffusers"
|
|
topaz_cli = "/Applications/Topaz Photo AI.app/Contents/MacOS/Topaz Photo AI"
|
|
steps = 20
|
|
scale = 7.0
|
|
width, height = 1024, 512
|
|
|
|
if torch.backends.mps.is_available():
|
|
device = "mps"
|
|
elif torch.cuda.is_available():
|
|
device = "cuda"
|
|
else:
|
|
device = "cpu"
|
|
|
|
# 1) Generate base HDRI
|
|
gen_pipe = StableDiffusionPipeline.from_pretrained(
|
|
model_path,
|
|
torch_dtype=torch.float32
|
|
).to(device)
|
|
gen_pipe.scheduler = DPMSolverMultistepScheduler.from_config(gen_pipe.scheduler.config)
|
|
gen_pipe.enable_attention_slicing()
|
|
|
|
print("→ Generating equirectangular HDRI…")
|
|
image = gen_pipe(
|
|
prompt=args.prompt,
|
|
num_inference_steps=steps,
|
|
guidance_scale=scale-1.5,
|
|
width=width,
|
|
height=height
|
|
).images[0]
|
|
|
|
gen_path = os.path.join(tempdir, f"base_{width}x{height}.png")
|
|
image.save(gen_path)
|
|
print(f"→ Saved initial image to {gen_path}")
|
|
|
|
# 2) Make it seamless
|
|
shift_amt = width // 2
|
|
mask_w = width // 8
|
|
|
|
shifted = shift_image(image, shift_amt)
|
|
mask = create_mask(width, height, mask_w)
|
|
|
|
inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
|
|
"Lykon/dreamshaper-8-inpainting",
|
|
torch_dtype=torch.float32
|
|
).to(device)
|
|
inpaint_pipe.enable_attention_slicing()
|
|
|
|
print("→ Inpainting seam for seamless tiling…")
|
|
inpainted = inpaint_pipe(
|
|
prompt=args.prompt,
|
|
image=shifted,
|
|
mask_image=mask,
|
|
num_inference_steps=steps,
|
|
guidance_scale=scale,
|
|
width=width,
|
|
height=height
|
|
).images[0]
|
|
|
|
seamless_path = os.path.join(tempdir, os.path.basename(args.output))
|
|
inpainted = unshift_image(inpainted, shift_amt)
|
|
inpainted.save(seamless_path)
|
|
print(f"→ Crafted seamless image: {seamless_path}")
|
|
|
|
# 3) Upscale with Topaz Photo AI CLI
|
|
print("→ Upscaling with Topaz Photo AI CLI…")
|
|
result = subprocess.run(
|
|
[topaz_cli, "--cli", seamless_path, "-o", tempdir],
|
|
check=True
|
|
)
|
|
|
|
# Finde das letzte erstellte PNG im tempdir (das ist das hochskalierte!)
|
|
# Topaz kann einen Suffix anhängen, falls der Name schon existiert.
|
|
upscaled_files = sorted(
|
|
[os.path.join(tempdir, f) for f in os.listdir(tempdir) if f.lower().endswith(".png")],
|
|
key=os.path.getmtime,
|
|
reverse=True
|
|
)
|
|
if not upscaled_files:
|
|
print("→ No PNG output found in tempdir after Topaz run!")
|
|
return
|
|
|
|
upscaled = upscaled_files[0]
|
|
shutil.move(upscaled, output_abs)
|
|
print(f"→ Upscaled image moved to {output_abs}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |