Update app.py

app.py

@@ -1,25 +1,31 @@
+import gradio as gr
 import os
 import torch
 from diffusers import DiffusionPipeline
-import gradio as gr
-from PIL import Image
 
-# Load the model
-pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev").to("cuda" if torch.cuda.is_available() else "cpu")
+# Load Hugging Face access token from secrets
+hf_token = os.getenv("secret")  # Ensure your secret is named "secret"
+
+# Set up the pipeline with the access token
+pipeline = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    use_auth_token=hf_token
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
+# Inference function
 def generate_image(prompt):
-    # Generate the image using the pipeline
-    image = pipeline(prompt).images[0]
+    with torch.no_grad():
+        image = pipeline(prompt).images[0]
     return image
 
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# FLUX Image Generator")
-    prompt = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want to generate...")
-    generate_button = gr.Button("Generate Image")
+    prompt = gr.Textbox(label="Enter your prompt", placeholder="e.g. Astronaut riding a horse")
+    generate_btn = gr.Button("Generate Image")
     output_image = gr.Image(label="Generated Image")
 
-    generate_button.click(fn=generate_image, inputs=prompt, outputs=output_image)
+    generate_btn.click(fn=generate_image, inputs=prompt, outputs=output_image)
 
 # Launch the app
 demo.launch()