File size: 3,397 Bytes
cf72c4b
 
 
a188c05
cf72c4b
 
 
 
 
 
a188c05
cf72c4b
 
 
 
 
 
 
 
 
 
b857620
 
cf72c4b
 
fd62691
3513d23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a1db0e9
 
b857620
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf72c4b
3513d23
cf72c4b
dae5e79
3513d23
72c2560
dae5e79
2c5c599
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3513d23
2c5c599
 
 
 
 
 
 
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
from PIL import Image

import torch

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

try:
    import spaces
except ImportError:

    class spaces:
        def GPU(*args, **kwargs):
            return lambda x: x


import gradio as gr
from utils import NEGATIVE_PROMPT, IMAGE_SIZE_OPTIONS, QUALITY_TAGS, IMAGE_SIZES


device = "cuda"
model_name: str = "cagliostrolab/animagine-xl-3.1"
pipe = StableDiffusionXLPipeline.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    use_safetensors=True,
    add_watermarker=False,
    custom_pipeline="lpw_stable_diffusion_xl",
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_pretrained(
    model_name,
    subfolder="scheduler",
)

# sdpa
pipe.unet.set_attn_processor(AttnProcessor2_0())

pipe.to(device)


def image_generation_config_ui():
    with gr.Accordion(label="Image generation config", open=False) as accordion:
        image_size = gr.Radio(
            label="Image size",
            choices=list(IMAGE_SIZE_OPTIONS.keys()),
            value=list(IMAGE_SIZE_OPTIONS.keys())[3],
            interactive=True,
        )

        quality_tags = gr.Textbox(
            label="Quality tags",
            placeholder=QUALITY_TAGS["default"],
            value=QUALITY_TAGS["default"],
            interactive=True,
        )
        negative_prompt = gr.Textbox(
            label="Negative prompt",
            placeholder=NEGATIVE_PROMPT["default"],
            value=NEGATIVE_PROMPT["default"],
            interactive=True,
        )

        num_inference_steps = gr.Slider(
            label="Num inference steps",
            minimum=20,
            maximum=30,
            step=1,
            value=25,
            interactive=True,
        )
        guidance_scale = gr.Slider(
            label="Guidance scale",
            minimum=0.0,
            maximum=10.0,
            step=0.5,
            value=7.0,
            interactive=True,
        )

    return accordion, [
        image_size,
        quality_tags,
        negative_prompt,
        num_inference_steps,
        guidance_scale,
    ]


class ImageGenerator:
    # pipe: StableDiffusionXLPipeline

    def __init__(self):
        pass

    @spaces.GPU(duration=30)
    def generate(
        self,
        prompt: str,
        image_size: str = "768x1344",
        quality_tags: str = QUALITY_TAGS["default"],  # Light v3.1
        negative_prompt: str = NEGATIVE_PROMPT["default"],  # Light v3.1
        num_inference_steps: int = 25,
        guidance_scale: float = 7.0,
    ) -> Image.Image:
        width, height = IMAGE_SIZES[image_size]

        prompt = ", ".join([prompt, quality_tags])

        print("prompt", prompt)
        print("negative_prompt", negative_prompt)
        print("height", height)
        print("width", width)
        print("num_inference_steps", num_inference_steps)
        print("guidance_scale", guidance_scale)

        return pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
        ).images