Update custom_pipeline.py

custom_pipeline.py

@@ -48,7 +48,7 @@ class FluxWithCFGPipeline(FluxPipeline):
     with progressively increasing resolution for faster generation.
     """
     @torch.inference_mode()
-    def generate_images(
+    async def generate_images(
         self,
         prompt: Union[str, List[str]] = None,
         prompt_2: Optional[Union[str, List[str]]] = None,
@@ -64,7 +64,6 @@ class FluxWithCFGPipeline(FluxPipeline):
         pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
-        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
         max_sequence_length: int = 300,
         generate_with_graph = None
     ):
@@ -84,7 +83,6 @@ class FluxWithCFGPipeline(FluxPipeline):
         )
 
         self._guidance_scale = guidance_scale
-        self._joint_attention_kwargs = joint_attention_kwargs
         self._interrupt = False
 
         # 2. Define call parameters
@@ -92,7 +90,7 @@ class FluxWithCFGPipeline(FluxPipeline):
         device = self._execution_device
 
         # 3. Encode prompt
-        lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
+        lora_scale = None
         prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
             prompt=prompt,
             prompt_2=prompt_2,
@@ -140,7 +138,7 @@ class FluxWithCFGPipeline(FluxPipeline):
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
 
             if generate_with_graph:
-                return generate_with_graph(latents, prompt_embeds, pooled_prompt_embeds, text_ids, latent_image_ids, timestep)
+                return await generate_with_graph(latents, prompt_embeds, pooled_prompt_embeds, text_ids, latent_image_ids, timestep)
             else:
                 noise_pred = self.transformer(
                     hidden_states=latents,
@@ -150,7 +148,6 @@ class FluxWithCFGPipeline(FluxPipeline):
                     encoder_hidden_states=prompt_embeds,
                     txt_ids=text_ids,
                     img_ids=latent_image_ids,
-                    joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
                 )[0]