File size: 4,490 Bytes
40462a0
7aafe2f
 
 
 
 
 
 
40462a0
7aafe2f
 
 
 
 
 
40462a0
7aafe2f
 
 
 
 
 
 
 
 
23fd89e
7aafe2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
97d3c4e
7aafe2f
 
 
 
 
 
 
 
 
23fd89e
7aafe2f
 
 
 
 
 
 
 
23fd89e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
from diffusers import DiffusionPipeline
from custom_pipeline import FLUXPipelineWithIntermediateOutputs

# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
DEFAULT_INFERENCE_STEPS = 1

# Device and model setup
dtype = torch.float16
pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
).to("cuda")
torch.cuda.empty_cache()

# Inference function
@spaces.GPU(duration=25)
def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=DEFAULT_INFERENCE_STEPS):
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    start_time = time.time()

    # Only generate the last image in the sequence
    for img in pipe.generate_images(  
            prompt=prompt,
            guidance_scale=0,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator
        ): 
        latency = f"Latency: {(time.time()-start_time):.2f} seconds"    
        yield img, seed, latency


# Example prompts
examples = [
    "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",
    "a futuristic cityscape with flying cars and neon lights",
    "Photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.",
    "Imagine steve jobs as Star Wars movie character"
]

# --- Gradio UI ---
with gr.Blocks() as demo:
    with gr.Column(elem_id="app-container"):
        gr.Markdown("# 🎨 Realtime FLUX Image Generator")
        gr.Markdown("Generate stunning images in real-time with advanced AI technology.")

        with gr.Row():
            with gr.Column(scale=3):
                result = gr.Image(label="Generated Image", show_label=False, interactive=False)
            with gr.Column(scale=1):
                prompt = gr.Text(
                    label="Prompt",
                    placeholder="Describe the image you want to generate...",
                    lines=3,
                    show_label=False,
                    container=False,
                )
                enhanceBtn = gr.Button("🚀 Enhance Image")

                with gr.Column("Advanced Options"):
                    with gr.Row():
                        latency = gr.Text(show_label=False)
                    with gr.Row():
                        seed = gr.Number(label="Seed", value=42, precision=0)
                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
                    with gr.Row():
                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
                        num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)

        with gr.Row():
            gr.Markdown("### 🌟 Inspiration Gallery")
        with gr.Row():
            gr.Examples(
                examples=examples,
                fn=generate_image,
                inputs=[prompt],
                outputs=[result, seed],
                cache_examples="lazy" 
            )

    # Event handling - Trigger image generation on button click or input change
    enhanceBtn.click(
        fn=generate_image,
        inputs=[prompt, seed, width, height],
        outputs=[result, seed, latency],
        show_progress="hidden",
        show_api=False,
        queue=False
    )

    gr.on(
        triggers=[prompt.input, width.input, height.input, num_inference_steps.input],
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
        outputs=[result, seed, latency],
        show_progress="hidden",
        show_api=False,
        trigger_mode="always_last",
        queue=False
    )

# Launch the app
demo.launch()