From bb47f1ea893bc1f12ceaa3856cc03c0b14ec559b Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 8 Jun 2025 18:00:24 +0900 Subject: [PATCH] Fix unwrap_model handling for None text_encoders in sample_images function --- library/flux_train_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 5f6867a8..8392e559 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -67,7 +67,7 @@ def sample_images( # unwrap unet and text_encoder(s) flux = accelerator.unwrap_model(flux) if text_encoders is not None: - text_encoders = [accelerator.unwrap_model(te) for te in text_encoders] + text_encoders = [(accelerator.unwrap_model(te) if te is not None else None) for te in text_encoders] if controlnet is not None: controlnet = accelerator.unwrap_model(controlnet) # print([(te.parameters().__next__().device if te is not None else None) for te in text_encoders])