import torch
from diffusers import FluxPipeline, DPMSolverMultistepScheduler

# Load the model
model_id = "black-forest-labs/FLUX.1-dev"
pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()

if torch.cuda.is_available():
    pipe = pipe.to("cuda")

# Define the function to generate images
def generate_image(prompt):
  image = pipe(prompt, num_inference_steps=25).images[0]
  return image

# Title for your Space
title = "Kailouis AI Image Generator"

# Description for your Space
description = "Generate amazing images with Kailouis AI, powered by the FLUX.1-dev model."

# Create the UI
import gradio as gr

demo = gr.Interface(
    fn=generate_image,
    inputs=[gr.Textbox(lines=1, placeholder="Enter your prompt")],
    outputs=[gr.Image(type="pil")],
    title=title,
    description=description,
)

demo.launch(share=True)