From 1bd351ffa362b749e195c737a7eb7a9647d8fac6 Mon Sep 17 00:00:00 2001
From: Gromit Park <bemore.one@gmail.com>
Date: Thu, 12 Feb 2026 19:46:14 +0900
Subject: [PATCH] Enhance image_spectrum to return amplitude and phase

Updated RETURN_TYPES to include amplitude and phase spectra. Modified image_spectrum function to return both amplitude and phase tensors.
---
 __init__.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/__init__.py b/__init__.py
index 7f7f0ed..b1bd70c 100644
--- a/__init__.py
+++ b/__init__.py
@@ -18,7 +18,8 @@ class Image_Spectrum:
             },
         }
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    RETURN_NAMES = ("amplitude_spectrum", "phase_spectrum")
     FUNCTION = "image_spectrum"
     CATEGORY = "image"
 
@@ -26,17 +27,22 @@ class Image_Spectrum:
         image = 255. * image[0].cpu().numpy()
         image = Image.fromarray(np.clip(image, 0, 255).astype(np.uint8))
         image = ImageOps.grayscale(image)
-        
+
         # 2D Fourier 변환 수행
         f_transform = fft2(image)
         f_shift = fftshift(f_transform)
-        spectrum = amplitude*np.log1p(np.abs(f_shift))
-        
-        image = np.array(spectrum).astype(np.float32) / 255.0
-        mask = 1. - torch.from_numpy(image)
-        image = torch.from_numpy(image)[None,]
-        
-        return (image,)
+
+        # 진폭 스펙트럼
+        spectrum = amplitude * np.log1p(np.abs(f_shift))
+        amp_img = np.array(spectrum).astype(np.float32) / 255.0
+        amp_tensor = torch.from_numpy(amp_img)[None,]
+
+        # 위상 스펙트럼 (-π ~ π → 0 ~ 1 정규화)
+        phase = np.angle(f_shift)
+        phase_img = ((phase + np.pi) / (2 * np.pi)).astype(np.float32)
+        phase_tensor = torch.from_numpy(phase_img)[None,]
+
+        return (amp_tensor, phase_tensor)
 
 NODE_CLASS_MAPPINGS = {
     "Image_Spectrum": Image_Spectrum