fix code

Author: junqiangwu

src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py

@@ -291,8 +291,6 @@ def rewire_prompt(self, prompt, device):
         return rewrite_prompt
 
     def _encode_prompt( self, prompt ):
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
         all_tokens = []
         for clean_prompt_sub, matched in split_quotation(prompt[0]):
             if matched:
@@ -341,23 +339,23 @@ def _encode_prompt( self, prompt ):
         prompt_embeds = text_output.hidden_states[-1].detach()
         prompt_embeds = prompt_embeds[:,self.prompt_template_encode_start_idx: -self.prompt_template_encode_end_idx ,:]
 
-        _, seq_len, _ = prompt_embeds.shape
-
-        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
         return prompt_embeds
 
     def encode_prompt(self, 
                     prompt : List[str] = None,
                     num_images_per_prompt: Optional[int] = 1,
                     prompt_embeds: Optional[torch.Tensor] = None ):
-
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
         # If prompt_embeds is provided and prompt is None, skip encoding
         if prompt_embeds is None:
-            prompt_embeds = self._encode_prompt( prompt, num_images_per_prompt )
-        
+            prompt_embeds = self._encode_prompt( prompt )
+
+        _, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
         text_ids = prepare_pos_ids(modality_id=0,
                                    type='text',
                                    start=(0, 0),

src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py

@@ -279,14 +279,10 @@ def __init__(
         self.default_sample_size = 128
         self.tokenizer_max_length = 512
 
-    def _encode_prompt( self, prompt, image, num_images_per_prompt ):
-
+    def _encode_prompt( self, prompt, image ):
         raw_vl_input = self.image_processor_vl(images=image,return_tensors="pt")
         pixel_values = raw_vl_input['pixel_values']
         image_grid_thw = raw_vl_input['image_grid_thw']
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
         all_tokens = []
         for clean_prompt_sub, matched in split_quotation(prompt[0]):
             if matched:
@@ -348,25 +344,25 @@ def _encode_prompt( self, prompt, image, num_images_per_prompt ):
         prompt_embeds = text_output.hidden_states[-1].detach()
         prompt_embeds = prompt_embeds[:,self.prompt_template_encode_start_idx: -self.prompt_template_encode_end_idx ,:]
 
-        _, seq_len, _ = prompt_embeds.shape
-
-        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
         return prompt_embeds
 
     @torch.inference_mode()
     def encode_prompt(self, 
                     prompt : List[str] = None,
                     image: Optional[torch.Tensor] = None,
                     num_images_per_prompt: Optional[int] = 1,
-                    prompt_embeds: Optional[torch.Tensor] = None,):
-
+                    prompt_embeds: Optional[torch.Tensor] = None):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
         # If prompt_embeds is provided and prompt is None, skip encoding
         if prompt_embeds is None:
-            prompt_embeds = self._encode_prompt( prompt, image, num_images_per_prompt )
+            prompt_embeds = self._encode_prompt( prompt, image )
 
+        _, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
         text_ids = prepare_pos_ids(modality_id=0,
                                    type='text',
                                    start=(0, 0),