Update app.py

app.py

@@ -1,23 +1,15 @@
 import os
-import requests
-import io
-from PIL import Image
+import torch
+from diffusers import DiffusionPipeline
 import gradio as gr
+from PIL import Image
 
-# Define the API URL for the model
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-headers = {
-    "Authorization": f"Bearer {os.getenv('secret')}"
-}
-
-def query(prompt):
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-    response.raise_for_status()  # Raise an error for bad responses
-    return response.content
+# Load the model
+pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev").to("cuda" if torch.cuda.is_available() else "cpu")
 
 def generate_image(prompt):
-    image_bytes = query(prompt)
-    image = Image.open(io.BytesIO(image_bytes))
+    # Generate the image using the pipeline
+    image = pipeline(prompt).images[0]
     return image
 
 # Gradio interface