Log VQ remapping configuration through the model logger

Author: wunianze666-netizen

src/diffusers/models/autoencoders/vae.py

@@ -17,7 +17,7 @@
 import torch
 import torch.nn as nn
 
-from ...utils import BaseOutput
+from ...utils import BaseOutput, logging
 from ...utils.torch_utils import randn_tensor
 from ..activations import get_activation
 from ..attention_processor import SpatialNorm
@@ -29,6 +29,9 @@
 )
 
 
+logger = logging.get_logger(__name__)
+
+
 @dataclass
 class EncoderOutput(BaseOutput):
     r"""
@@ -599,7 +602,7 @@ def __init__(
             if self.unknown_index == "extra":
                 self.unknown_index = self.re_embed
                 self.re_embed = self.re_embed + 1
-            print(
+            logger.info(
                 f"Remapping {self.n_e} indices to {self.re_embed} indices. "
                 f"Using {self.unknown_index} for unknown indices."
             )

tests/models/autoencoders/test_models_vq.py

@@ -13,10 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+
+import numpy as np
 import pytest
 import torch
 
 from diffusers import VQModel
+from diffusers.models.autoencoders.vae import VectorQuantizer
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import backend_manual_seed, enable_full_determinism, torch_device
@@ -119,6 +123,18 @@ def test_loss_pretrained(self):
         # fmt: on
         assert torch.allclose(output, expected_output, atol=1e-3)
 
+    def test_vector_quantizer_logs_remap_configuration(self, caplog, tmp_path):
+        remap_path = tmp_path / "used.npy"
+        np.save(remap_path, np.array([0, 2, 4], dtype=np.int64))
+
+        with caplog.at_level(logging.INFO, logger="diffusers.models.autoencoders.vae"):
+            VectorQuantizer(n_e=8, vq_embed_dim=4, beta=0.25, remap=str(remap_path), unknown_index="extra")
+
+        assert any(
+            "Remapping 8 indices to 4 indices. Using 3 for unknown indices." in record.message
+            for record in caplog.records
+        )
+
 
 class TestVQModelTraining(VQModelTesterConfig, TrainingTesterMixin):
     """Training tests for VQModel."""