app.py

import gradio as gr

gr.load("models/black-forest-labs/FLUX.1-schnell").launch()

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 = 2048

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
            prompt = prompt, 
            width = width,
            height = height,
            num_inference_steps = num_inference_steps, 
            generator = generator,
            guidance_scale=0.0
    ).images[0] 
    return image, seed
 
examples = [
    "Gambarkan suasana perayaan Hari Ulang Tahun Republik Indonesia yang ke-79 di sebuah alun-alun kota. Tampak bendera merah putih berkibar di tengah lapangan, di sekelilingnya ribuan orang mengenakan pakaian tradisional dari berbagai daerah Indonesia, semua bersatu dalam semangat kebangsaan. Di latar belakang, terdapat panggung besar yang dihiasi dengan dekorasi merah dan putih, di mana seorang pemimpin upacara sedang memberikan pidato. Langit cerah biru dengan beberapa awan putih lembut, dan di udara, terlihat pesawat tempur TNI AU yang terbang membentuk formasi melintas di atas area perayaan. Di sekeliling lapangan, terlihat stand-stand kuliner khas Indonesia yang dipenuhi oleh warga yang sedang menikmati makanan bersama."
    "Depict a scene of a beautiful 20-year-old girl named Dewi, with a sweet and gentle face, shopping at the bustling Gombong morning market. She is dressed in stylish, revealing attire that accentuates her allure. Dewi moves gracefully among the stalls, attracting the admiring gazes of all the men around her. The market is filled with vibrant activity, with vendors selling fresh produce, local delicacies, and various goods. The atmosphere is lively, with bright sunlight casting soft shadows, and Dewi stands out as the captivating focal point of the scene."
    "a tiny astronaut hatching from an egg on the moon",
    "a cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 [schnell]
12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation
[[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-schnell)]
        """)
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples="lazy"
        )

    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()