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 HighSpeedFluxPipeline
 
 # 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 = HighSpeedFluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
 ).to("cuda")
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")