diff --git a/src/diffusers/models/autoencoders/vae.py b/src/diffusers/models/autoencoders/vae.py index a65bca418175..4461d1540203 100644 --- a/src/diffusers/models/autoencoders/vae.py +++ b/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." ) diff --git a/tests/models/autoencoders/test_models_vq.py b/tests/models/autoencoders/test_models_vq.py index b6dbcfa64b19..54ea6da8456c 100644 --- a/tests/models/autoencoders/test_models_vq.py +++ b/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."""