Update custom_pipeline.py

custom_pipeline.py

@@ -1,8 +1,20 @@
 import torch
 import numpy as np
 from diffusers import FluxPipeline, FlowMatchEulerDiscreteScheduler
-from typing import Any, Dict, List, Optional, Union
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from typing import Any, Callable, Dict, List, Optional, Union
 from PIL import Image
+from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
+
+from diffusers.utils import is_torch_xla_available
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
 
 # Constants for shift calculation
 BASE_SEQ_LEN = 256
@@ -54,6 +66,8 @@ class FluxWithCFGPipeline(FluxPipeline):
         prompt_2: Optional[Union[str, List[str]]] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
         num_inference_steps: int = 4,
         timesteps: List[int] = None,
         guidance_scale: float = 3.5,
@@ -62,6 +76,8 @@ class FluxWithCFGPipeline(FluxPipeline):
         latents: Optional[torch.FloatTensor] = None,
         prompt_embeds: Optional[torch.FloatTensor] = None,
         pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -102,6 +118,21 @@ class FluxWithCFGPipeline(FluxPipeline):
             max_sequence_length=max_sequence_length,
             lora_scale=lora_scale,
         )
+        (
+            negative_prompt_embeds,
+            negative_pooled_prompt_embeds,
+            negative_text_ids,
+        ) = self.encode_prompt(
+            prompt=negative_prompt,
+            prompt_2=negative_prompt_2,
+            prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels // 4
         latents, latent_image_ids = self.prepare_latents(
@@ -114,6 +145,7 @@ class FluxWithCFGPipeline(FluxPipeline):
             generator,
             latents,
         )
+        
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
         image_seq_len = latents.shape[1]
@@ -149,9 +181,24 @@ class FluxWithCFGPipeline(FluxPipeline):
                 joint_attention_kwargs=self.joint_attention_kwargs,
                 return_dict=False,
             )[0]
+            
+            noise_pred_uncond = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                pooled_projections=negative_pooled_prompt_embeds,
+                encoder_hidden_states=negative_prompt_embeds,
+                txt_ids=negative_text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+            
+            noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
-             # Yield intermediate result
+            latents_dtype = latents.dtype
             latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+             # Yield intermediate result
             torch.cuda.empty_cache()
 
         # Final image