Update app.py

app.py

@@ -1,13 +1,13 @@
 import streamlit as st
-from diffusers import FluxPipeline
+from diffusers import DiffusionPipeline
 from PIL import Image
 import torch
 
-# Load the Flux model
+# Load the FLUX model
 @st.cache_resource
 def load_pipeline():
     # Using the 'black-forest-labs/FLUX.1-schnell' model
-    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
     pipe.enable_model_cpu_offload()  # Offload to CPU to save memory
     return pipe
 
@@ -28,8 +28,6 @@ if st.button("Generate Image"):
                 user_prompt,
                 guidance_scale=0.0,  # No guidance
                 num_inference_steps=4,  # Number of steps for faster generation
-                max_sequence_length=256,
-                generator=torch.Generator("cpu").manual_seed(0)  # Ensure reproducibility
             ).images[0]
             
             # Save and display the image