Refactor: Separate core generation logic into a dedicated function

generate_equirect.py

@@ -34,23 +34,13 @@ def unshift_image(img: Image.Image, shift: int) -> Image.Image:
     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()
-
+def generate_equirect(prompt: str, output: str, work_dir: str | None = None) -> str | None:
     # Output-Ordner (bleibt wie gehabt)
-    output_abs = os.path.abspath(args.output)
+    output_abs = os.path.abspath(output)
+    work_dir = work_dir or os.path.dirname(os.path.abspath(__file__))
 
     # Zwischenschritte landen im eigenem temp-Ordner:
-    with tempfile.TemporaryDirectory(dir=args.work_dir) as tempdir:
+    with tempfile.TemporaryDirectory(dir=work_dir) as tempdir:
         print(f"→ Using tempdir: {tempdir}")
 
         model_path = "/Volumes/SD/ML-Models/diffusers/hdri-panorama-v1-diffusers"
@@ -76,7 +66,7 @@ def main():
 
         print("→ Generating equirectangular HDRI…")
         image = gen_pipe(
-            prompt=args.prompt,
+            prompt=prompt,
             num_inference_steps=steps,
             guidance_scale=scale-1.5,
             width=width,
@@ -102,7 +92,7 @@ def main():
 
         print("→ Inpainting seam for seamless tiling…")
         inpainted = inpaint_pipe(
-            prompt=args.prompt,
+            prompt=prompt,
             image=shifted,
             mask_image=mask,
             num_inference_steps=steps,
@@ -132,11 +122,26 @@ def main():
         )
         if not upscaled_files:
             print("→ No PNG output found in tempdir after Topaz run!")
-            return
+            return None
 
         upscaled = upscaled_files[0]
         shutil.move(upscaled, output_abs)
         print(f"→ Upscaled image moved to {output_abs}")
+        return output_abs
+
+
+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()
+    generate_equirect(args.prompt, args.output, args.work_dir)
 
 if __name__ == "__main__":
-    main()
+    main()