[diffusion] fix: fix Flux2-Klein prompt tokenization length to 512 and add regression coverage (#21407)

This commit is contained in:
Aditya Sharma
2026-03-28 17:28:02 +08:00
committed by GitHub
parent edd4d54023
commit 627e162335
@@ -435,6 +435,17 @@ class Flux2PipelineConfig(FluxPipelineConfig):
default_factory=lambda: (flux2_postprocess_text,) default_factory=lambda: (flux2_postprocess_text,)
) )
vae_config: VAEConfig = field(default_factory=Flux2VAEConfig) vae_config: VAEConfig = field(default_factory=Flux2VAEConfig)
text_encoder_extra_args: list[dict] = field(
default_factory=lambda: [
dict(
max_length=512,
padding="max_length",
truncation=True,
return_overflowing_tokens=False,
return_length=False,
)
]
)
def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
# flatten to 1-d list # flatten to 1-d list
@@ -679,7 +690,9 @@ class Flux2KleinPipelineConfig(Flux2PipelineConfig):
texts = [_apply_chat_template(prompt) for prompt in prompts] texts = [_apply_chat_template(prompt) for prompt in prompts]
tok_kwargs = dict(tok_kwargs or {}) tok_kwargs = dict(tok_kwargs or {})
max_length = tok_kwargs.pop("max_length", 512) tok_kwargs.pop("max_length", None)
# Flux2 Klein uses max_length 512.
max_length = 512
padding = tok_kwargs.pop("padding", "max_length") padding = tok_kwargs.pop("padding", "max_length")
truncation = tok_kwargs.pop("truncation", True) truncation = tok_kwargs.pop("truncation", True)
return_tensors = tok_kwargs.pop("return_tensors", "pt") return_tensors = tok_kwargs.pop("return_tensors", "pt")