Fix naming

Author: patrickvonplaten

src/compel/compel.py

@@ -30,7 +30,7 @@ def __init__(self,
                  downweight_mode: DownweightMode = DownweightMode.MASK,
                  use_penultimate_clip_layer: bool=False,
                  device: Optional[str] = None):
-                 return_pooled: Union[str, List[bool]] = False,
+                 requires_pooled: Union[str, List[bool]] = False,
         """
         Initialize Compel. The tokenizer and text_encoder can be lifted directly from any DiffusionPipeline.
 
@@ -64,7 +64,7 @@ def __init__(self,
                                                             padding_attention_mask_value = padding_attention_mask_value,
                                                             downweight_mode=downweight_mode,
                                                             use_penultimate_clip_layer=use_penultimate_clip_layer,
-                                                            return_pooled=return_pooled,
+                                                            requires_pooled=requires_pooled,
                                                             )
         else:
             self.conditioning_provider = EmbeddingsProvider(tokenizer=tokenizer,
@@ -75,7 +75,7 @@ def __init__(self,
                                                             padding_attention_mask_value = padding_attention_mask_value,
                                                             downweight_mode=downweight_mode,
                                                             use_penultimate_clip_layer=use_penultimate_clip_layer,
-                                                            return_pooled=return_pooled,
+                                                            requires_pooled=requires_pooled,
                                                             )
             self._device = device
 

src/compel/embeddings_provider.py

@@ -312,7 +312,7 @@ def build_weighted_embedding_tensor(self,
                                         token_ids: torch.Tensor,
                                         per_token_weights: torch.Tensor,
                                         attention_mask: Optional[torch.Tensor] = None,
-                                        return_pooled: bool = False,
+                                        should_return_pooled: bool = False,
                                         device: Optional[str] = None) -> torch.Tensor:
         """
         :param token_ids: A tensor of shape `n*[self.max_length]` containing token IDs (ints) where n is some arbitrary
@@ -373,7 +373,7 @@ def build_weighted_embedding_tensor(self,
 
             chunk_start_index += chunk_size
 
-        if self.requires_pooled:
+        if should_return_pooled:
             return weighted_z, pooled
 
         return weighted_z
@@ -473,15 +473,15 @@ def __init__(self,
                 textual_inversion_manager: BaseTextualInversionManager = None,
                 dtype_for_device_getter: Callable[[torch.device], torch.dtype] = lambda device: torch.float32,
                 hidden_states_types: Union[str, List[str]] = "final",
-                return_pooled: Union[str, List[bool]] = False,
+                requires_pooled: Union[str, List[bool]] = False,
                 ):
 
         hidden_states_types = len(text_encoders) * [hidden_states_types] if not isinstance(hidden_states_types, (list, tuple)) else hidden_states_types
-        return_pooled = len(text_encoders) * [return_pooled] if not isinstance(return_pooled, (list, tuple)) else return_pooled
+        requires_pooled = len(text_encoders) * [requires_pooled] if not isinstance(requires_pooled, (list, tuple)) else requires_pooled
 
         self.embedding_providers = [
             EmbeddingsProvider(tokenizer, text_encoder, textual_inversion_manager, dtype_for_device_getter, hidden_states_type, pooled)
-            for tokenizer, text_encoder, hidden_states_type, pooled in zip(tokenizers, text_encoders, hidden_states_types, return_pooled)
+            for tokenizer, text_encoder, hidden_states_type, pooled in zip(tokenizers, text_encoders, hidden_states_types, requires_pooled)
         ]
 
     @property

test/test_compel.py

@@ -80,7 +80,7 @@ def test_basic_prompt_multi_text_encoder(self):
         tokenizer_2 = DummyTokenizer()
         text_encoder_2 = DummyTransformer()
 
-        compel = Compel(tokenizer=[tokenizer_1, tokenizer_2], text_encoder=[text_encoder_1, text_encoder_2], hidden_states_type="penultimate", return_pooled=[False, True])
+        compel = Compel(tokenizer=[tokenizer_1, tokenizer_2], text_encoder=[text_encoder_1, text_encoder_2], hidden_states_type="penultimate", requires_pooled=[False, True])
 
         # test "a b c" makes it to the Conditioning intact for t=0, t=0.5, t=1
         prompt = " ".join(KNOWN_WORDS[:3])