kailouis's picture
Create app.py
cd068c6 verified
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)