Update app.py

app.py

@@ -1,30 +1,34 @@
 import gradio as gr
-import torch
-from diffusers import StableDiffusionPipeline
+import requests
+import io
+from PIL import Image
+import os
+from dotenv import load_dotenv
+load_dotenv()
 
-# Load the pipeline outside the function to avoid reloading it for each request
-pipe = StableDiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4", 
-    revision="fp16", 
-    torch_dtype=torch.float16
-).to("cuda")
+# Define the API URL and headers
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}
 
-# Check if the model is on GPU
-print("Is the model on GPU?", next(pipe.parameters()).is_cuda)
+# Define the query function
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
 
-# Define a function to generate the image
-def generate_image(prompt):
-    image = pipe(prompt).images[0]
+# Define the function to generate the image
+def generate_image(input_text):
+    image_bytes = query({"inputs": input_text})
+    image = Image.open(io.BytesIO(image_bytes))
     return image
 
-# Gradio interface: Text input for the prompt, and image output
-interface = gr.Interface(
-    fn=generate_image, 
-    inputs="text", 
-    outputs="image",
-    title="Stable Diffusion Image Generator",
-    description="Generate an image based on your text prompt using Stable Diffusion."
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.inputs.Textbox(label="Input Text", placeholder="Enter your description here..."),
+    outputs=gr.outputs.Image(type="pil", label="Generated Image"),
+    title="Image Generation with FLUX",
+    description="Enter a description to generate an image."
 )
 
 # Launch the Gradio app
-interface.launch()
+iface.launch()