app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

@spaces.GPU(duration=120)
def infer(prompt, seed=47622, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, 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=guidance_scale
    ).images[0] 
    return image, seed
 
examples = [
    "a panda holding a bamboo sign that says welcome",
    "Beautiful scene with a bright sunlit valley between a large cliffside mountain and a waterfall flowing down the side into a small stream",
    "A scene full of classic video game characters as stickers on a black water bottle",
    "4k photo of a Dark Forest filled with neon lights, where the air is charged with electromagnetic fields and solar radiation from distant stars. It's a beautiful collision of waves, light, particles and complexity",
    "Mathematically perfect architecture including contrast in both materials and shapes. Utilize smooth gradients for all color transitions",
    "Colorful, whimsical, imaginative 10 course meal spread out on a beautiful hand carved wooden table in a stark modern futuristic city center restaurant",
    "The scene is set in the heart of a bustling metropolis, amidst towering skyscrapers that pierce the night sky like colossal arrows. Neon lights of all hues flicker incessantly, casting an ethereal glow",
    "In a high-tech laboratory, scientists study the effects of electromagnetism on high-energy particles. The room is bathed in an ethereal blue light emanating from advanced equipment humming with power",
    "A futuristic biocity that is located in the former site of Portsmouth, New Hampshire. It has a mix of old and new buildings, green spaces, and water features. It also has six large artificial floating islands off of its coastline,(zenithal angle), coastal city,blue sky and white clouds,the sun is shining brightly,ultra-wide angle,",
    "Depict a breathtaking scene of a meteor rain showering down from a starry night sky. The meteors should vary in size and brightness, streaking across the sky with vibrant tails of light, creating a dazzling display. Below, a serene landscape—perhaps a tranquil lake reflecting the celestial spectacle, or a rugged mountain range—should enhance the sense of wonder. The foreground can include silhouettes of trees or figures gazing up in awe at the cosmic event. The overall atmosphere should evoke feelings of magic and inspiration, capturing the beauty and mystery of the universe."
]


css="""
#body {
    font-family: Arial, sans-serif;
    background-color: #f2f2f2;
}

#gr-col-container {
    margin: 0 auto;
    max-width: 520px;
    padding: 20px;
    background-color: #fff;
    border: 1px solid #ddd;
    border-radius: 10px;
    box-shadow: 0 0 10px rgba(0,0,0,0.1);
}

.gr-markdown {
    font-size: 16px;
    color: #333;
}

.gr-input {
    padding: 10px;
    border: 1px solid #ccc;
    border-radius: 5px;
}

.gr-slider {
    width: 100%;
    padding: 10px;
}

.gr-button {
    background-color: #4CAF50;
    color: #fff;
    padding: 10px 20px;
    border: none;
    border-radius: 5px;
    cursor: pointer;
}

.gr-button:hover {
    background-color: #3e8e41;
}

.gr-accordion {
    background-color: #f2f2f2;
    border: 1px solid #ddd;
    border-radius: 10px;
    padding: 10px;
}

.gr-accordion label {
    font-weight: bold;
    margin-bottom: 10px;
}

.gr-examples {
    padding: 10px;
    background-color: #f2f2f2;
    border: 1px solid #ddd;
    border-radius: 10px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
            Generate an image with Flux. Try it out and let me know what you think!
            Expect roughly 60 seconds per generation with it's current backend.
            This can be scaled up over time as needed. Thanks!
            Not for Commercial Use - Apache 2.0 License
        """)

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

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
        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, guidance_scale, num_inference_steps],
        outputs = [result, seed]
    )

demo.launch()