Initial Commit

generate_equirect.py

@@ -0,0 +1,142 @@
+#!/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()