@@ -230,6 +230,12 @@ def clean_text(self, text: str) -> str:
230230
231231 return text .strip ()
232232
233+ def basic_clean (self , text : str ) -> str :
234+ """Light cleaning: fix mojibake and unescape HTML. Used when skip_text_cleaning=True."""
235+ text = ftfy .fix_text (text )
236+ text = html .unescape (html .unescape (text ))
237+ return text .strip ()
238+
233239
234240EXAMPLE_DOC_STRING = """
235241 Examples:
@@ -378,6 +384,8 @@ def encode_prompt(
378384 negative_prompt_embeds : torch .FloatTensor | None = None ,
379385 prompt_attention_mask : torch .BoolTensor | None = None ,
380386 negative_prompt_attention_mask : torch .BoolTensor | None = None ,
387+ tokenizer_max_length : int | None = None ,
388+ skip_text_cleaning : bool = False ,
381389 ):
382390 """Encode text prompt using standard text encoder and tokenizer, or use precomputed embeddings."""
383391 if device is None :
@@ -388,7 +396,14 @@ def encode_prompt(
388396 prompt = [prompt ]
389397 # Encode the prompts
390398 prompt_embeds , prompt_attention_mask , negative_prompt_embeds , negative_prompt_attention_mask = (
391- self ._encode_prompt_standard (prompt , device , do_classifier_free_guidance , negative_prompt )
399+ self ._encode_prompt_standard (
400+ prompt ,
401+ device ,
402+ do_classifier_free_guidance ,
403+ negative_prompt ,
404+ tokenizer_max_length = tokenizer_max_length ,
405+ skip_text_cleaning = skip_text_cleaning ,
406+ )
392407 )
393408
394409 # Duplicate embeddings for each generation per prompt
@@ -419,13 +434,21 @@ def encode_prompt(
419434 negative_prompt_attention_mask if do_classifier_free_guidance else None ,
420435 )
421436
422- def _tokenize_prompts (self , prompts : list [str ], device : torch .device ):
437+ def _tokenize_prompts (
438+ self ,
439+ prompts : list [str ],
440+ device : torch .device ,
441+ tokenizer_max_length : int | None = None ,
442+ skip_text_cleaning : bool = False ,
443+ ):
423444 """Tokenize and clean prompts."""
424- cleaned = [self .text_preprocessor .clean_text (text ) for text in prompts ]
445+ clean_fn = self .text_preprocessor .basic_clean if skip_text_cleaning else self .text_preprocessor .clean_text
446+ cleaned = [clean_fn (text ) for text in prompts ]
447+ max_length = tokenizer_max_length or self .tokenizer .model_max_length
425448 tokens = self .tokenizer (
426449 cleaned ,
427450 padding = "max_length" ,
428- max_length = self . tokenizer . model_max_length ,
451+ max_length = max_length ,
429452 truncation = True ,
430453 return_attention_mask = True ,
431454 return_tensors = "pt" ,
@@ -438,6 +461,8 @@ def _encode_prompt_standard(
438461 device : torch .device ,
439462 do_classifier_free_guidance : bool = True ,
440463 negative_prompt : str = "" ,
464+ tokenizer_max_length : int | None = None ,
465+ skip_text_cleaning : bool = False ,
441466 ):
442467 """Encode prompt using standard text encoder and tokenizer with batch processing."""
443468 batch_size = len (prompt )
@@ -450,7 +475,9 @@ def _encode_prompt_standard(
450475 else :
451476 prompts_to_encode = prompt
452477
453- input_ids , attention_mask = self ._tokenize_prompts (prompts_to_encode , device )
478+ input_ids , attention_mask = self ._tokenize_prompts (
479+ prompts_to_encode , device , tokenizer_max_length = tokenizer_max_length , skip_text_cleaning = skip_text_cleaning
480+ )
454481
455482 with torch .no_grad ():
456483 embeddings = self .text_encoder (
@@ -501,10 +528,12 @@ def check_inputs(
501528 "`negative_prompt_embeds` must also be provided for classifier-free guidance."
502529 )
503530
504- spatial_compression = self .vae_scale_factor
505- if height % spatial_compression != 0 or width % spatial_compression != 0 :
531+ # The latents must be divisible by the transformer's patch size after VAE compression.
532+ dimension_multiple = self .vae_scale_factor * self .transformer .config .patch_size
533+ if height % dimension_multiple != 0 or width % dimension_multiple != 0 :
506534 raise ValueError (
507- f"`height` and `width` have to be divisible by { spatial_compression } but are { height } and { width } ."
535+ f"`height` and `width` have to be divisible by { dimension_multiple } (vae_scale_factor *"
536+ f" transformer patch_size) but are { height } and { width } ."
508537 )
509538
510539 if guidance_scale < 1.0 :
@@ -545,6 +574,8 @@ def __call__(
545574 use_resolution_binning : bool = True ,
546575 callback_on_step_end : Callable [[int , int ], None ] | None = None ,
547576 callback_on_step_end_tensor_inputs : list [str ] = ["latents" ],
577+ tokenizer_max_length : int | None = None ,
578+ skip_text_cleaning : bool = False ,
548579 ):
549580 """
550581 Function invoked when calling the pipeline for generation.
@@ -598,6 +629,12 @@ def __call__(
598629 output_type (`str`, *optional*, defaults to `"pil"`):
599630 The output format of the generate image. Choose between
600631 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
632+ tokenizer_max_length (`int`, *optional*):
633+ Override the maximum number of tokens used when tokenizing the prompt. Defaults to the tokenizer's own
634+ ``model_max_length`` when not set.
635+ skip_text_cleaning (`bool`, *optional*, defaults to `False`):
636+ If `True`, uses only light prompt cleaning (fix encoding + unescape HTML) instead of the full DeepFloyd
637+ cleaning pipeline.
601638 return_dict (`bool`, *optional*, defaults to `True`):
602639 Whether or not to return a [`~pipelines.prx.PRXPipelineOutput`] instead of a plain tuple.
603640 use_resolution_binning (`bool`, *optional*, defaults to `True`):
@@ -627,11 +664,6 @@ def __call__(
627664 width = width or default_resolution
628665
629666 if use_resolution_binning :
630- if self .image_processor is None :
631- raise ValueError (
632- "Resolution binning requires a VAE with image_processor, but VAE is not available. "
633- "Set use_resolution_binning=False or provide a VAE."
634- )
635667 if self .default_sample_size not in ASPECT_RATIO_BINS :
636668 raise ValueError (
637669 f"Resolution binning is only supported for default_sample_size in { list (ASPECT_RATIO_BINS .keys ())} , "
@@ -684,6 +716,8 @@ def __call__(
684716 negative_prompt_embeds = negative_prompt_embeds ,
685717 prompt_attention_mask = prompt_attention_mask ,
686718 negative_prompt_attention_mask = negative_prompt_attention_mask ,
719+ tokenizer_max_length = tokenizer_max_length ,
720+ skip_text_cleaning = skip_text_cleaning ,
687721 )
688722 # Expose standard names for callbacks parity
689723 prompt_embeds = text_embeddings
0 commit comments