import gradio as gr
from diffusers import DiffusionPipeline
import os
os.environ['HF_HOME'] = '/blabla/cache/'

# Load the diffusion model
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")

def generate_image(prompt):
    # Generate the image based on the prompt
    image = pipe(prompt).images[0]
    return image

# Create a Gradio interface using the new component style
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g., Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"),
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Image Generation with SDXL-Lightning",
    description="Enter a prompt to generate an image using the SDXL-Lightning model."
)

# Launch the Gradio interface
if __name__ == "__main__":
    iface.launch()