Update custom_pipeline.py

custom_pipeline.py

@@ -1,7 +1,7 @@
 import torch
 import numpy as np
-from diffusers import FluxPipeline, FlowMatchEulerDiscreteScheduler
-from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.pipelines.flux.pipeline_output import FluxPipeline, 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
@@ -55,10 +55,7 @@ def prepare_timesteps(
 
 # FLUX pipeline function
 class FluxWithCFGPipeline(FluxPipeline):
-    """
-    Extends the FluxPipeline to yield intermediate images during the denoising process 
-    with progressively increasing resolution for faster generation.
-    """
+    
     @torch.inference_mode()
     def generate_images(
         self,
@@ -83,7 +80,6 @@ class FluxWithCFGPipeline(FluxPipeline):
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
         max_sequence_length: int = 300,
     ):
-        """Generates images and yields intermediate results during the denoising process."""
         height = height or self.default_sample_size * self.vae_scale_factor
         width = width or self.default_sample_size * self.vae_scale_factor