File size: 3,210 Bytes
1887c49
a3e9651
 
9b8f82d
bf459df
c3a473c
a3e9651
 
 
 
 
 
d1118c6
 
 
 
 
 
 
 
eaf8b26
ff4f40d
 
 
 
 
 
 
1fe5ec5
d1118c6
c8cac13
435c235
003a39a
 
 
 
 
 
 
 
 
 
 
 
 
 
b6d3179
 
 
 
 
67df9db
fad334a
6631477
97606cf
 
b6d3179
 
6631477
 
ad96459
b6d3179
b1db469
 
fad334a
 
 
 
 
 
 
b1db469
 
 
 
 
 
fad334a
 
6631477
 
b6d3179
6631477
 
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
import gradio as gr
from diffusers import DiffusionPipeline
import spaces
import torch
from concurrent.futures import ProcessPoolExecutor
from huggingface_hub import hf_hub_download

dev_model = "black-forest-labs/FLUX.1-dev"
schnell_model = "black-forest-labs/FLUX.1-schnell"

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

repo_name = "ByteDance/Hyper-SD"
ckpt_name = "Hyper-FLUX.1-dev-8steps-lora.safetensors"
hyper_lora = hf_hub_download(repo_name, ckpt_name)

repo_name = "alimama-creative/FLUX.1-Turbo-Alpha"
ckpt_name = "diffusion_pytorch_model.safetensors"
turbo_lora = hf_hub_download(repo_name, ckpt_name)

pipe_dev = DiffusionPipeline.from_pretrained(dev_model, torch_dtype=torch.bfloat16).to("cuda")
pipe_schnell = DiffusionPipeline.from_pretrained(
    schnell_model,
    text_encoder=pipe_dev.text_encoder,
    text_encoder_2=pipe_dev.text_encoder_2,
    tokenizer=pipe_dev.tokenizer,
    tokenizer_2=pipe_dev.tokenizer_2,
    torch_dtype=torch.bfloat16
)

@spaces.GPU(duration=75)
def run_parallel_models(prompt, progress=gr.Progress(track_tqdm=True)):
    pipe_dev.load_lora_weights(hyper_lora)
    image = pipe_dev(prompt, num_inference_steps=8, joint_attention_kwargs={"scale": 0.125}).images[0]
    pipe_dev.unload_lora_weights()
    yield image, gr.update(), gr.update()
    pipe_dev.load_lora_weights(turbo_lora)
    image = pipe_dev(prompt, num_inference_steps=8).images[0]
    yield gr.update(), image, gr.update()
    pipe_dev.unload_lora_weights()
    pipe_dev.to("cpu")
    pipe_schnell.to("cuda")
    image = pipe_schnell(prompt, num_inference_steps=4).images[0]
    yield gr.update(), gr.update(), image
   
#run_parallel_models.zerogpu = True
css = '''
#gen_btn{height: 100%}
#gen_column{align-self: stretch}
'''
with gr.Blocks(css=css) as demo:
    gr.Markdown("# Low Step Flux Comparison")
    gr.Markdown("Compare the quality (not the speed) of FLUX Schnell (4 steps), FLUX.1[dev] HyperFLUX (8 steps), FLUX.1[dev]-Turbo-Alpha (8 steps). It runs a bit slow as it's inferencing the three models.")
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Prompt")
        with gr.Column(scale=1, min_width=120, elem_id="gen_column"):
            submit = gr.Button("Run", elem_id="gen_btn")
    with gr.Row():
        hyper = gr.Image(label="FLUX.1[dev] HyperFLUX (8 steps)")
        turbo = gr.Image(label="FLUX.1[dev]-Turbo-Alpha (8 steps)")
        schnell = gr.Image(label="FLUX Schnell (4 steps)")
    
    gr.Examples(
        examples=[
            ["the spirit of a Tamagotchi wandering in the city of Vienna"],
            ["a photo of a lavender cat"],
            ["a tiny astronaut hatching from an egg on the moon"],
            ["a delicious ceviche cheesecake slice"],
            ["an insect robot preparing a delicious meal"],
            ["a Charmander fine dining with a view to la Sagrada Família"]],
        fn=run_parallel_models,
        inputs=[prompt],
        outputs=[schnell, hyper, turbo],
        cache_examples="lazy"
    )
    
    gr.on(
        triggers=[submit.click, prompt.submit],
        fn=run_parallel_models,
        inputs=[prompt],
        outputs=[hyper, turbo, schnell]
    )
demo.launch()