File size: 6,138 Bytes
40462a0
7aafe2f
 
 
 
 
c82ffd4
2097b43
40462a0
7aafe2f
 
 
 
 
 
40462a0
7aafe2f
 
2097b43
7aafe2f
 
deb928c
7aafe2f
 
 
7d90483
f1c2277
7aafe2f
 
c82ffd4
7aafe2f
 
 
 
b7bdba8
7aafe2f
 
 
b7bdba8
7aafe2f
b7bdba8
 
 
7aafe2f
 
 
 
ffe402c
7aafe2f
ffe402c
9494d71
5d87e58
7aafe2f
 
 
 
 
 
 
a25ba7a
a8abbc8
7aafe2f
 
bf1be4a
7aafe2f
 
c9efd38
7aafe2f
 
 
 
 
 
db97287
7aafe2f
 
 
 
db97287
f1c2277
7aafe2f
1c4647b
7aafe2f
 
 
 
 
 
 
 
 
 
 
 
 
5d87e58
7aafe2f
 
 
3bad26c
 
 
 
7aafe2f
3bad26c
97d3c4e
7aafe2f
5358c2f
8374546
 
7aafe2f
 
f1c2277
7aafe2f
23fd89e
7aafe2f
f1c2277
05a0589
66892aa
f1c2277
 
 
 
 
 
 
 
 
 
 
 
8374546
 
f1c2277
 
 
 
 
 
 
3bad26c
 
f1c2277
3bad26c
8374546
 
7aafe2f
 
f1c2277
 
 
 
 
 
c9c4986
8374546
 
f1c2277
 
7aafe2f
f1c2277
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
from diffusers import DiffusionPipeline, AutoencoderTiny
from custom_pipeline import HighSpeedFluxPipeline

# 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 = HighSpeedFluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
).to("cuda")
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).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=2, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(int(float(seed)))

    start_time = time.time()

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

# Example prompts
examples = [
    "a tiny astronaut hatching from an egg on the moon",
    "a cute white cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
    "Create mage of Modern house in minecraft style",
    "Imagine steve jobs as Star Wars movie character",
    "Lion",
    "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.",
]

# --- 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 Modified Flux.Schnell pipeline.")
        gr.Markdown("<span style='color: red;'>Note: Sometimes it stucks or stops generating images (I don't know why). In that situation just refresh the site.</span>")

        with gr.Row():
            with gr.Column(scale=2.5):
                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,
                )
                generateBtn = gr.Button("🖼️ Generate Image")
                enhanceBtn = gr.Button("🚀 Enhance Image")

                with gr.Column("Advanced Options"):
                    with gr.Row():
                        realtime = gr.Checkbox(label="Realtime Toggler", info="If TRUE then uses more GPU but create image in realtime.", value=False)
                        latency = gr.Text(label="Latency")
                    with gr.Row():
                        seed = gr.Number(label="Seed", value=42)
                        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, latency],
                cache_examples="lazy" 
            )

    def enhance_image(*args):
        gr.Info("Enhancing Image") # currently just runs optimized pipeline for 2 steps. Further implementations later.
        return next(generate_image(*args))

    enhanceBtn.click(
        fn=enhance_image,
        inputs=[prompt, seed, width, height],
        outputs=[result, seed, latency],
        show_progress="full",
        queue=False,
        concurrency_limit=None
    )

    generateBtn.click(
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
        outputs=[result, seed, latency],
        show_progress="full",
        api_name="RealtimeFlux",
        queue=False
    )

    def update_ui(realtime_enabled):
        return {
            prompt: gr.update(interactive=True),
            generateBtn: gr.update(visible=not realtime_enabled)
        }

    realtime.change(
        fn=update_ui,
        inputs=[realtime],
        outputs=[prompt, generateBtn],
        queue=False,
        concurrency_limit=None
    )

    def realtime_generation(*args):
        if args[0]:  # If realtime is enabled
            return next(generate_image(*args[1:]))

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
        outputs=[result, seed, latency],
        show_progress="full",
        queue=False,
        concurrency_limit=None
    )

    for component in [prompt, width, height, num_inference_steps]:
        component.input(
            fn=realtime_generation,
            inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
            outputs=[result, seed, latency],
            show_progress="hidden",
            trigger_mode="always_last",
            queue=False,
            concurrency_limit=None
        )

# Launch the app
demo.launch()