Hugging Face's logo

Spaces:
dbaranchuk
/
SwD-FLUX
Running on Zero

App Files Community
SwD-FLUX
File size: 3,828 Bytes 
0816484
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ccc66a4
 
0816484
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
 
import spaces
import gradio as gr
import numpy as np
import random
import functools
import os
import torch
from diffusers import FluxPipeline
from peft import LoraConfig, get_peft_model, PeftModel

huggingface_token = os.getenv("HF_TOKEN")

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", 
                                    torch_dtype=torch.float16,
                                    token=huggingface_token,
                                    custom_pipeline='quickjkee/swd_pipeline_flux').to('cuda')
distill_check = 'yresearch/swd_flux'
pipe.transformer = PeftModel.from_pretrained(
    pipe.transformer,
    distill_check,
)


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


@spaces.GPU()
def infer(prompt, seed, randomize_seed):

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)
    sigmas = [1.0000, 0.8956, 0.7363, 0.6007, 0.0000]
    scales = [64, 80, 96, 128]
    
    image = pipe(
        prompt=prompt,
        height=int(scales[0] * 8),
        width=int(scales[0] * 8),
        scales=scales,
        sigmas=sigmas,
        timesteps=torch.tensor(sigmas[:-1]).to('cuda') * 1000,
        guidance_scale=4.5,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(0)
    ).images[0]

    return image


examples = [
    "3d digital art of an adorable ghost, holding a heart shaped pumpkin, Halloween, super cute, spooky haunted house background",
    'Long-exposure night photography of a starry sky over a mountain range, with light trails.',
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "A gold astronaut meditating in a lush green forest by a lake",
    "A group of friends sitting around a campfire."
]

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

if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            f"""
        # ⚡ Scale-wise Distillation | FLUX.1 [dev] ⚡ 
        # ⚡ Image Generation with 4-step SwD ⚡
        This is a demo of [Scale-wise Distillation](https://yandex-research.github.io/swd/), 
        a diffusion distillation method proposed in [Scale-wise Distillation of Diffusion Models](https://arxiv.org/abs/2503.16397)
        by [Yandex Research](https://github.com/yandex-research).
        Currently running on {power_device}.
        """
        )
        gr.Markdown(
            "If you enjoy the space, feel free to give a ⭐ to the <a href='https://github.com/yandex-research/swd' target='_blank'>Github Repo</a>. [![GitHub Stars](https://img.shields.io/github/stars/yandex-research/invertible-cd?style=social)](https://github.com/yandex-research/invertible-cd)"
        )

        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=False)


        gr.Examples(
            examples=examples,
            inputs=[prompt],
            cache_examples=False
        )
    run_button.click(
        fn=infer,
        inputs=[prompt, seed, randomize_seed],
        outputs=[result]
    )

demo.queue().launch(share=False)