From bc0e4df12dfdb6883858da7ca3e792fec56be0e0 Mon Sep 17 00:00:00 2001
From: lehie <lehieuhuy19@gmail.com>
Date: Wed, 11 Feb 2026 01:41:56 +0700
Subject: [PATCH] MP3 embed: only take first 1s, resample to 8kHz mono 8bit

---
 mp3_embed_image.py | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/mp3_embed_image.py b/mp3_embed_image.py
index 97a3b16..807118a 100644
--- a/mp3_embed_image.py
+++ b/mp3_embed_image.py
@@ -1,10 +1,12 @@
 import os
+import io
 import numpy as np
 import torch
+import torchaudio
 
 
 class MP3EmbedInImage:
-    """Embed MP3 binary data into image alpha channel using steganography."""
+    """Embed first 1 second of MP3 audio into image alpha channel using steganography."""
     
     @classmethod
     def INPUT_TYPES(cls):
@@ -28,9 +30,22 @@ class MP3EmbedInImage:
         if not mp3_path or not os.path.isfile(mp3_path):
             raise Exception(f"File không tồn tại: {mp3_path}")
 
-        # Read MP3 binary data
-        with open(mp3_path, "rb") as f:
-            mp3_bytes = f.read()
+        # Load audio and take only first 1 second
+        waveform, sample_rate = torchaudio.load(mp3_path)
+        one_sec_samples = min(sample_rate, waveform.shape[1])
+        waveform_1s = waveform[:, :one_sec_samples]
+
+        # Resample to 8000Hz mono to minimize size
+        if sample_rate != 8000:
+            resampler = torchaudio.transforms.Resample(sample_rate, 8000)
+            waveform_1s = resampler(waveform_1s)
+        # Convert to mono
+        if waveform_1s.shape[0] > 1:
+            waveform_1s = waveform_1s.mean(dim=0, keepdim=True)
+
+        # Convert to 8-bit PCM bytes (compact)
+        audio_np = (waveform_1s.squeeze(0).clamp(-1, 1).numpy() * 127).astype(np.int8)
+        mp3_bytes = audio_np.tobytes()
 
         # Convert to bits
         data_bits = []
@@ -40,8 +55,9 @@ class MP3EmbedInImage:
             data_bits.append((data_len >> (31 - i)) & 1)
         # Data bits
         for b in mp3_bytes:
+            b_unsigned = b & 0xFF
             for i in range(8):
-                data_bits.append((b >> (7 - i)) & 1)
+                data_bits.append((b_unsigned >> (7 - i)) & 1)
 
         # Process image
         np_img = np.asarray(image, dtype=np.float32)
@@ -62,11 +78,9 @@ class MP3EmbedInImage:
 
         total_pixels = H * W
         if len(data_bits) > total_pixels:
-            mp3_size_kb = data_len / 1024
-            max_size_kb = (total_pixels - 32) / 8 / 1024
             raise Exception(
-                f"MP3 quá lớn ({mp3_size_kb:.1f}KB) cho ảnh {W}x{H}. "
-                f"Tối đa: {max_size_kb:.1f}KB"
+                f"Ảnh {W}x{H} quá nhỏ. Cần ít nhất {len(data_bits)} pixels, "
+                f"ảnh chỉ có {total_pixels} pixels."
             )
 
         # Embed LSB into alpha channel