Implement text encoder outputs extra conditioning data

Author: JasonHoku

batch_encoding.py

@@ -76,8 +76,11 @@ def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positi
                             original_layer = clip_model.cond_stage_model.clip_layer
                             clip_model.cond_stage_model.set_clip_options({"layer": clip_skip})
 
-                        cond, pooled = clip_model.encode_from_tokens(tokens, return_pooled=True)
-                        conditioning = [[cond, {"pooled_output": pooled}]]
+                        # Use return_dict=True to preserve all extra conditioning keys
+                        # (e.g. t5xxl_ids/t5xxl_weights for Anima, attention_mask for Lumina, etc.)
+                        pooled_dict = clip_model.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
+                        cond = pooled_dict.pop("cond")
+                        conditioning = [[cond, pooled_dict]]
                         results[prompt] = conditioning
                         cond_cache.set(prompt, conditioning, prompt_type)
 

generation_orchestrator.py

@@ -616,8 +616,11 @@ def run_generation_loop(
                                         patched_clip.cond_stage_model.set_clip_options({"layer": clip_skip})
 
                                     tokens = patched_clip.tokenize(prompt)
-                                    cond, pooled = patched_clip.encode_from_tokens(tokens, return_pooled=True)
-                                    conditioning_cache["positive"][prompt] = [[cond, {"pooled_output": pooled}]]
+                                    # Use return_dict=True to preserve all extra conditioning keys
+                                    # (e.g. t5xxl_ids/t5xxl_weights for Anima, attention_mask for Lumina, etc.)
+                                    pooled_dict = patched_clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
+                                    cond = pooled_dict.pop("cond")
+                                    conditioning_cache["positive"][prompt] = [[cond, pooled_dict]]
 
                                     if original_layer is not None:
                                         patched_clip.cond_stage_model.set_clip_options({"layer": original_layer})
@@ -635,8 +638,10 @@ def run_generation_loop(
                                         patched_clip.cond_stage_model.set_clip_options({"layer": clip_skip})
 
                                     tokens = patched_clip.tokenize(prompt)
-                                    cond, pooled = patched_clip.encode_from_tokens(tokens, return_pooled=True)
-                                    conditioning_cache["negative"][prompt] = [[cond, {"pooled_output": pooled}]]
+                                    # Use return_dict=True to preserve all extra conditioning keys
+                                    pooled_dict = patched_clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
+                                    cond = pooled_dict.pop("cond")
+                                    conditioning_cache["negative"][prompt] = [[cond, pooled_dict]]
 
                                     if original_layer is not None:
                                         patched_clip.cond_stage_model.set_clip_options({"layer": original_layer})