Update app.py

app.py

@@ -5,7 +5,7 @@ import spaces
 import torch
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
-from custom_pipeline import HighSpeedFluxPipeline
+from custom_pipeline import FLUXPipelineWithIntermediateOutputs
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -16,7 +16,7 @@ DEFAULT_INFERENCE_STEPS = 1
 
 # Device and model setup
 dtype = torch.float16
-pipe = HighSpeedFluxPipeline.from_pretrained(
+pipe = FLUXPipelineWithIntermediateOutputs.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")
@@ -24,7 +24,7 @@ torch.cuda.empty_cache()
 
 # Inference function
 @spaces.GPU(duration=25)
-def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2, progress=gr.Progress(track_tqdm=True)):
+def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(int(float(seed)))