wifix199's picture
Update app.py
7467e8a verified
raw
history blame
857 Bytes
import gradio as gr
import torch
from diffusers import FluxPipeline
# Load the model
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
def generate_image(prompt):
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0)
).images[0]
return image
# Create the UI
ui = gr.Interface(
generate_image,
gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
"image",
title="FLUX.1-dev Image Generator",
description="Generate images using the FLUX.1-dev model.",
)
# Launch the UI
ui.launch()