Spaces:
Running
Running
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() |