Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import random | |
import spaces | |
import torch | |
from diffusers import DiffusionPipeline | |
dtype = torch.bfloat16 | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device) | |
MAX_SEED = np.iinfo(np.int32).max | |
MAX_IMAGE_SIZE = 1024 | |
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)): | |
lora = "kratadata/luisa-doodle" | |
prefix = "A black and white pen doodle in the style of BOODLE of " | |
suffix = "Completly white background. Doodle in the center, and it is not touching the edges." | |
if randomize_seed: | |
seed = random.randint(0, MAX_SEED) | |
generator = torch.Generator().manual_seed(seed) | |
full_prompt = prefix + " " + prompt + " " + suffix | |
pipe.load_lora_weights(lora) | |
image = pipe( | |
prompt = full_prompt, | |
width = width, | |
height = height, | |
num_inference_steps = num_inference_steps, | |
generator = generator, | |
guidance_scale=0.0 | |
).images[0] | |
return image, seed | |
with gr.Blocks() as demo: | |
with gr.Row(): | |
with gr.Column(): | |
prompt = gr.Textbox( | |
label="Image Prompt", | |
show_label = "True", | |
info="Your image prompt", | |
max_lines=4, | |
placeholder="Enter your prompt", | |
container=True, | |
) | |
with gr.Accordion("Advanced Settings", open=False): | |
width = gr.Slider( | |
label="Width", | |
minimum=256, | |
maximum=MAX_IMAGE_SIZE, | |
step=32, | |
value=1024, | |
info = "Keep at 1024 for best results" | |
) | |
height = gr.Slider( | |
label="Height", | |
minimum=256, | |
maximum=MAX_IMAGE_SIZE, | |
step=32, | |
value=1024, | |
info = "Keep at 1024 for best results" | |
) | |
num_inference_steps = gr.Slider( | |
label="Number of inference steps", | |
minimum=1, | |
maximum=8, | |
step=1, | |
value=2, | |
info = "Increase to 4 for better results" | |
) | |
randomize_seed = gr.Checkbox(label="Randomize seed", value=True, info = "Keep true to generate a new image each time") | |
seed = gr.Slider( | |
label="Seed", | |
minimum=0, | |
maximum=MAX_SEED, | |
step=1, | |
value=0, | |
info = "Fix seed if you want to keep generating the same image" | |
) | |
run_button = gr.Button("Run", scale=0) | |
with gr.Column(): | |
result = gr.Image(label="Result", show_label=False, format="jpeg") | |
gr.on( | |
triggers=[run_button.click, prompt.submit], | |
fn = infer, | |
inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps], | |
outputs = [result, seed] | |
) | |
demo.launch() |