add max length

Author: zRzRzRzRzRzRzR

examples/cogview4-control/train_control_cogview4.py

@@ -132,6 +132,8 @@ def log_validation(cogview4_transformer, args, accelerator, weight_dtype, step,
                     control_image=validation_image,
                     num_inference_steps=50,
                     guidance_scale=args.guidance_scale,
+                    max_sequence_length=max_sequence_length, # For downstream task training usage, training can be performed on a batch basis.
+                    padding_type="max_length",
                     generator=generator,
                     height=args.resolution,
                     width=args.resolution,
@@ -267,6 +269,9 @@ def parse_args(input_args=None):
             " resolution"
         ),
     )
+    parser.add_argument(
+        "--max_sequence_length", type=int, default=128, help="The maximum sequence length for the prompt."
+    )
     parser.add_argument(
         "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
     )
@@ -1079,10 +1084,12 @@ def load_model_hook(models, input_dir):
                 text_encoding_pipeline = text_encoding_pipeline.to("cuda")
 
                 with torch.no_grad():
-                    (
-                        prompt_embeds,
-                        pooled_prompt_embeds,
-                    ) = text_encoding_pipeline.encode_prompt(captions, "")
+                    # Since the batch will be padded, max_length should be used for padding.
+                    prompt_embeds,pooled_prompt_embeds,= text_encoding_pipeline.encode_prompt(
+                    captions, "",
+                        max_sequence_length=args.max_sequence_length,
+                        padding_type="max_length"
+                    )
                 original_size = (args.resolution, args.resolution)
                 original_size = torch.tensor([original_size], dtype=prompt_embeds.dtype, device=prompt_embeds.device)
 
@@ -1099,7 +1106,7 @@ def load_model_hook(models, input_dir):
                 # this could be optimized by not having to do any text encoding and just
                 # doing zeros on specified shapes for `prompt_embeds` and `pooled_prompt_embeds`
                 if args.proportion_empty_prompts and random.random() < args.proportion_empty_prompts:
-                    # 这里，直接将 pooled_prompt_embeds 16个 pad token 提供给 prompt_embeds
+                    # Here, we directly pass 16 pad tokens from pooled_prompt_embeds to prompt_embeds.
                     prompt_embeds = pooled_prompt_embeds
                 if args.offload:
                     text_encoding_pipeline = text_encoding_pipeline.to("cpu")

src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py

@@ -182,6 +182,7 @@ def _get_glm_embeds(
         prompt: Union[str, List[str]] = None,
         num_images_per_prompt: int = 1,
         max_sequence_length: int = 1024,
+        padding_type: str = "longest",
         device: Optional[torch.device] = None,
         dtype: Optional[torch.dtype] = None,
     ):
@@ -193,7 +194,7 @@ def _get_glm_embeds(
 
         text_inputs = self.tokenizer(
             prompt,
-            padding="longest",  # not use max length
+            padding=padding_type,
             max_length=max_sequence_length,
             truncation=True,
             add_special_tokens=True,
@@ -239,6 +240,7 @@ def encode_prompt(
         device: Optional[torch.device] = None,
         dtype: Optional[torch.dtype] = None,
         max_sequence_length: int = 1024,
+        padding_type: str = "longest",
     ):
         r"""
         Encodes the prompt into text encoder hidden states.
@@ -275,9 +277,8 @@ def encode_prompt(
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
-
         if prompt_embeds is None:
-            prompt_embeds = self._get_glm_embeds(prompt, num_images_per_prompt, max_sequence_length, device, dtype)
+            prompt_embeds = self._get_glm_embeds(prompt, num_images_per_prompt, max_sequence_length, padding_type,  device, dtype)
 
         if do_classifier_free_guidance and negative_prompt_embeds is None:
             negative_prompt = negative_prompt or ""
@@ -296,7 +297,7 @@ def encode_prompt(
                 )
 
             negative_prompt_embeds = self._get_glm_embeds(
-                negative_prompt, num_images_per_prompt, max_sequence_length, device, dtype
+                negative_prompt, num_images_per_prompt, max_sequence_length, "longest", device, dtype
             )
 
         return prompt_embeds, negative_prompt_embeds
@@ -450,6 +451,7 @@ def __call__(
         ] = None,
         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         max_sequence_length: int = 1024,
+        padding_type: str = "longest", # For downstream tasks, it can be modified to use max_length for implementation.
     ) -> Union[CogView4PipelineOutput, Tuple]:
         """
         Function invoked when calling the pipeline for generation.
@@ -579,7 +581,8 @@ def __call__(
             prompt_embeds=prompt_embeds,
             negative_prompt_embeds=negative_prompt_embeds,
             max_sequence_length=max_sequence_length,
-            device=device,
+            padding_type=padding_type,
+            device=device
         )
 
         # Prepare latents