File size: 8,947 Bytes
70f55b7
 
e47ff0d
f8868b5
70f55b7
 
 
 
 
ff84eba
 
 
 
 
 
 
e47ff0d
70f55b7
 
f8868b5
e47ff0d
 
 
 
 
 
 
f8868b5
 
e47ff0d
 
 
 
 
 
a5f564c
 
e47ff0d
075e529
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74874b
075e529
 
 
 
 
bf8cd49
 
 
075e529
 
e47ff0d
 
6ea233e
e47ff0d
a4bee0b
6ea233e
e47ff0d
a4bee0b
 
 
e47ff0d
075e529
e47ff0d
ff84eba
6ea233e
 
ff84eba
6ea233e
70f55b7
ff84eba
 
6ea233e
 
e47ff0d
f8868b5
70f55b7
 
 
 
6ea233e
 
 
70f55b7
 
 
6ea233e
70f55b7
6ea233e
70f55b7
e47ff0d
 
6ea233e
e47ff0d
6ea233e
a4bee0b
e47ff0d
70f55b7
f8868b5
e47ff0d
 
 
a4bee0b
e47ff0d
 
 
 
075e529
e47ff0d
 
 
6ea233e
 
 
 
 
 
 
70f55b7
e47ff0d
 
 
6ea233e
70f55b7
 
 
6ea233e
6ba6160
a4bee0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e47ff0d
6ea233e
a4bee0b
6ea233e
70f55b7
e47ff0d
 
 
 
 
 
 
70f55b7
e47ff0d
f8868b5
 
 
e47ff0d
70f55b7
6ea233e
 
70f55b7
e47ff0d
 
a4bee0b
 
 
 
 
 
 
 
 
e47ff0d
 
 
 
075e529
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e47ff0d
f8868b5
 
 
 
 
 
 
 
c44fbfc
 
 
 
e47ff0d
f8868b5
6ea233e
 
 
 
 
e47ff0d
 
f8868b5
 
e47ff0d
075e529
 
 
e47ff0d
f8868b5
70f55b7
 
ff84eba
6ea233e
70f55b7
f8868b5
 
 
ff84eba
6ea233e
f8868b5
e47ff0d
f8868b5
 
 
 
 
 
 
e47ff0d
f8868b5
e47ff0d
f8868b5
 
 
ff84eba
f8868b5
ff84eba
 
f8868b5
e47ff0d
f8868b5
 
ff84eba
f8868b5
ff84eba
 
f8868b5
e47ff0d
f8868b5
 
 
e47ff0d
f8868b5
e47ff0d
6ea233e
f8868b5
e47ff0d
f8868b5
 
ff84eba
d54903b
f8868b5
d54903b
f8868b5
e47ff0d
f8868b5
e47ff0d
a4bee0b
e47ff0d
6ea233e
70f55b7
e47ff0d
 
 
 
 
 
 
 
 
f8868b5
075e529
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f8868b5
e47ff0d
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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
import spaces

import os
import random
import math

import torch
import numpy as np

from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
    StableDiffusionXLPipeline,
)
from diffusers.schedulers.scheduling_euler_ancestral_discrete import (
    EulerAncestralDiscreteScheduler,
)
from diffusers.models.attention_processor import AttnProcessor2_0
from transformers import AutoModelForCausalLM, AutoTokenizer

import gradio as gr

try:
    from dotenv import load_dotenv

    load_dotenv()
except:
    print("failed to import dotenv (this is not a problem on the production)")

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

HF_TOKEN = os.environ.get("HF_TOKEN")
assert HF_TOKEN is not None

IMAGE_MODEL_REPO_ID = os.environ.get(
    "IMAGE_MODEL_REPO_ID", "OnomaAIResearch/Illustrious-xl-early-release-v0"
)
DART_V3_REPO_ID = os.environ.get("DART_V3_REPO_ID", None)
assert DART_V3_REPO_ID is not None

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

TEMPLATE = (
    "<|bos|>"
    #
    "<|rating:general|>"
    "{aspect_ratio}"
    "<|length:medium|>"
    #
    "<copyright></copyright>"
    #
    "<character></character>"
    #
    "<general>{subject}"
)
QUALITY_TAGS = "masterpiece, best quality, very aesthetic, newest"
NEGATIVE_PROMPT = "(worst quality, bad quality:1.1), very displeasing, lowres, jaggy lines, 3d, watermark, signature, copyright notice, logo, blurry, scan, jpeg artifacts, chromatic aberration, white outline, film grain, artistic error, bad anatomy, bad hands, wrong hand"
BAN_TAGS = [
    "photoshop (medium)",
    "clip studio paint (medium)",
    "absurdres",
    "highres",
    "copyright request",
    "character request",
    "creature request",
]

dart = AutoModelForCausalLM.from_pretrained(
    DART_V3_REPO_ID,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
    use_cache=True,
    device_map="cpu",
)
dart = dart.eval()
dart = dart.requires_grad_(False)
dart = torch.compile(dart)
tokenizer = AutoTokenizer.from_pretrained(DART_V3_REPO_ID)
BAN_TOKENS = [tokenizer.convert_tokens_to_ids([tag]) for tag in BAN_TAGS]

pipe = StableDiffusionXLPipeline.from_pretrained(
    IMAGE_MODEL_REPO_ID,
    torch_dtype=torch.bfloat16,
    add_watermarker=False,
    custom_pipeline="lpw_stable_diffusion_xl",
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.unet.set_attn_processor(AttnProcessor2_0())
if device == "cuda":
    pipe.enable_sequential_cpu_offload(gpu_id=0, device="cuda")


def get_aspect_ratio(width: int, height: int) -> str:
    ar = width / height

    if ar <= 1 / math.sqrt(3):
        return "<|aspect_ratio:ultra_tall|>"
    elif ar <= 8 / 9:
        return "<|aspect_ratio:tall|>"
    elif ar < 9 / 8:
        return "<|aspect_ratio:square|>"
    elif ar < math.sqrt(3):
        return "<|aspect_ratio:wide|>"
    else:
        return "<|aspect_ratio:ultra_wide|>"


@torch.inference_mode
def generate_prompt(subject: str, aspect_ratio: str):
    input_ids = tokenizer.encode_plus(
        TEMPLATE.format(aspect_ratio=aspect_ratio, subject=subject),
        return_tensors="pt",
    ).input_ids
    print("input_ids:", input_ids)

    output_ids = dart.generate(
        input_ids,
        max_new_tokens=256,
        do_sample=True,
        temperature=1.0,
        top_p=1.0,
        top_k=100,
        num_beams=1,
        bad_words_ids=BAN_TOKENS,
    )[0]

    generated = output_ids[len(input_ids) :]
    decoded = ", ".join(
        [
            token
            for token in tokenizer.batch_decode(generated, skip_special_tokens=True)
            if token.strip() != ""
        ]
    )
    print("decoded:", decoded)

    return decoded


def format_prompt(prompt: str, prompt_suffix: str):
    return f"{prompt}, {prompt_suffix}"


@spaces.GPU(duration=25)
def generate_image(
    prompt: str,
    negative_prompt: str,
    generator,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image


def on_generate(
    subject: str,
    suffix: str,
    negative_prompt: str,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    ar = get_aspect_ratio(width, height)
    print("ar:", ar)
    prompt = generate_prompt(subject, ar)
    prompt = format_prompt(prompt, suffix)
    print(prompt)

    image = generate_image(
        prompt,
        negative_prompt,
        generator,
        width,
        height,
        guidance_scale,
        num_inference_steps,
    )

    return image, prompt, seed


def on_retry(
    prompt: str,
    negative_prompt: str,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    print(prompt)

    image = generate_image(
        prompt,
        negative_prompt,
        generator,
        width,
        height,
        guidance_scale,
        num_inference_steps,
    )

    return image, prompt, seed


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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
# Random IllustriousXL
Model: [IllustriousXL v0.1](https://huggingface.co/OnomaAIResearch/Illustrious-xl-early-release-v0)
""")

        with gr.Row():
            subject_radio = gr.Dropdown(
                label="Subject",
                choices=["1girl", "2girls", "1boy", "no humans"],
                value="1girl",
            )
            run_button = gr.Button("Generate random", scale=0)

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Generation details", open=False):
            with gr.Row():
                prompt_txt = gr.Textbox(label="Generated prompt", interactive=False)
                retry_button = gr.Button("πŸ”„ Retry", scale=0)

        with gr.Accordion("Advanced Settings", open=False):
            prompt_suffix = gr.Text(
                label="Prompt suffix",
                visible=True,
                value=QUALITY_TAGS,
            )
            negative_prompt = gr.Text(
                label="Negative prompt",
                placeholder="Enter a negative prompt",
                visible=True,
                value=NEGATIVE_PROMPT,
            )

            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=512,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=832,  # Replace with defaults that work for your model
                )

                height = gr.Slider(
                    label="Height",
                    minimum=512,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1152,  # Replace with defaults that work for your model
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.5,
                    value=6.5,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=20,
                    maximum=40,
                    step=1,
                    value=28,
                )

    gr.on(
        triggers=[run_button.click],
        fn=on_generate,
        inputs=[
            subject_radio,
            prompt_suffix,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, prompt_txt, seed],
    )
    gr.on(
        triggers=[retry_button.click],
        fn=on_retry,
        inputs=[
            prompt_txt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, prompt_txt, seed],
    )

demo.queue().launch()