diff --git a/batch_encoding.py b/batch_encoding.py index 8df7991..5554244 100644 --- a/batch_encoding.py +++ b/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) diff --git a/generation_orchestrator.py b/generation_orchestrator.py index c823b6e..102d685 100644 --- a/generation_orchestrator.py +++ b/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})