Update app.py

app.py

@@ -5,7 +5,7 @@ import spaces
 import torch
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
-from custom_pipeline import FLUXPipelineWithIntermediateOutputs
+from custom_pipeline import FluxWithCFGPipeline
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -16,7 +16,7 @@ DEFAULT_INFERENCE_STEPS = 1
 
 # Device and model setup
 dtype = torch.float16
-pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
+pipe = FluxWithCFGPipeline.from_pretrained(
     "ostris/OpenFLUX.1", torch_dtype=dtype
 )
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16)