Update custom_pipeline.py

custom_pipeline.py

@@ -66,6 +66,7 @@ class FluxWithCFGPipeline(FluxPipeline):
         return_dict: bool = True,
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
         max_sequence_length: int = 300,
+        generate_with_graph = None
     ):
         """Generates images and yields intermediate results during the denoising process."""
         height = height or self.default_sample_size * self.vae_scale_factor
@@ -138,21 +139,24 @@ class FluxWithCFGPipeline(FluxPipeline):
 
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
 
-            noise_pred = self.transformer(
-                hidden_states=latents,
-                timestep=timestep / 1000,
-                guidance=guidance,
-                pooled_projections=pooled_prompt_embeds,
-                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]
-
-             # Yield intermediate result
-            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-            torch.cuda.empty_cache()
+            if generate_with_graph:
+                return generate_with_graph(latents, prompt_embeds, pooled_prompt_embeds, text_ids, latent_image_ids, timestep)
+            else:
+                noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    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]
+
+                 # Yield intermediate result
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                torch.cuda.empty_cache()
 
         # Final image
         return self._decode_latents_to_image(latents, height, width, output_type)