Update app.py

app.py

@@ -1,22 +1,31 @@
 import gradio as gr
-from transformers import Flux1
+import torch
+from diffusers import FluxPipeline
 
 # Load the model
-model = Flux1.from_pretrained("black-forest-labs/flux-1-dev")
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
 
 def generate_image(prompt):
-    # Generate an image using the model
-    image = model(prompt)
+    image = pipe(
+        prompt,
+        height=1024,
+        width=1024,
+        guidance_scale=3.5,
+        num_inference_steps=50,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0)
+    ).images[0]
     return image
 
-# Create the Gradio Interface
-demo = gr.Interface(
-    fn=generate_image, 
-    inputs=gr.Textbox(label="Prompt"), 
-    outputs=gr.Image(label="Generated Image"),
-    title="FLUX.1-dev Image Generation Bot",
-    description="Enter a prompt and generate an image using the FLUX.1-dev model."
+# Create the UI
+ui = gr.Interface(
+    generate_image,
+    gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
+    "image",
+    title="FLUX.1-dev Image Generator",
+    description="Generate images using the FLUX.1-dev model.",
 )
 
-# Launch the Gradio Interface
-demo.launch()
+# Launch the UI
+ui.launch()