Update custom_pipeline.py

custom_pipeline.py

@@ -1,6 +1,7 @@
 import torch
 import numpy as np
 from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers import DiffusionPipeline
 from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
 from typing import Any, Callable, Dict, List, Optional, Union
 from PIL import Image
@@ -54,7 +55,7 @@ def prepare_timesteps(
     return timesteps, num_inference_steps
 
 # FLUX pipeline function
-class FluxWithCFGPipeline(FluxPipeline):
+class FluxWithCFGPipeline(DiffusionPipeline):
     
     @torch.inference_mode()
     def generate_images(