mkshing commited on
Commit
b338d19
1 Parent(s): 58a9ea5

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +186 -0
app.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+ import gradio as gr
4
+ import torch
5
+ from diffusers import (
6
+ DiffusionPipeline,
7
+ DDIMScheduler,
8
+ DPMSolverMultistepScheduler,
9
+ EulerAncestralDiscreteScheduler,
10
+ EulerDiscreteScheduler,
11
+ LMSDiscreteScheduler,
12
+ PNDMScheduler,
13
+ UniPCMultistepScheduler,
14
+ )
15
+ from diffusers.utils import make_image_grid
16
+
17
+ ACCESS_TOKEN = os.environ["ACCESS_TOKEN"]
18
+ pipeline = DiffusionPipeline.from_pretrained(
19
+ "stabilityai/japanese-stable-diffusion-xl",
20
+ trust_remote_code=True,
21
+ torch_dtype=torch.float16,
22
+ use_auth_token=ACCESS_TOKEN
23
+ )
24
+ device = "cuda" if torch.cuda.is_available() else "cpu"
25
+ pipeline.to(device)
26
+ SCHEDULER_MAPPING = {
27
+ "ddim": DDIMScheduler,
28
+ "plms": PNDMScheduler,
29
+ "lms": LMSDiscreteScheduler,
30
+ "euler": EulerDiscreteScheduler,
31
+ "euler_ancestral": EulerAncestralDiscreteScheduler,
32
+ "dpm_solver++": DPMSolverMultistepScheduler,
33
+ "unipc": UniPCMultistepScheduler,
34
+ }
35
+ noise_scheduler_name = "euler"
36
+ SD_XL_BASE_RATIOS = {
37
+ "0.5": (704, 1408),
38
+ "0.52": (704, 1344),
39
+ "0.57": (768, 1344),
40
+ "0.6": (768, 1280),
41
+ "0.68": (832, 1216),
42
+ "0.72": (832, 1152),
43
+ "0.78": (896, 1152),
44
+ "0.82": (896, 1088),
45
+ "0.88": (960, 1088),
46
+ "0.94": (960, 1024),
47
+ "1.0": (1024, 1024),
48
+ "1.07": (1024, 960),
49
+ "1.13": (1088, 960),
50
+ "1.21": (1088, 896),
51
+ "1.29": (1152, 896),
52
+ "1.38": (1152, 832),
53
+ "1.46": (1216, 832),
54
+ "1.67": (1280, 768),
55
+ "1.75": (1344, 768),
56
+ "1.91": (1344, 704),
57
+ "2.0": (1408, 704),
58
+ "2.09": (1472, 704),
59
+ "2.4": (1536, 640),
60
+ "2.5": (1600, 640),
61
+ "2.89": (1664, 576),
62
+ "3.0": (1728, 576),
63
+ # "small": (512, 512), # for testing
64
+ }
65
+
66
+
67
+ def set_noise_scheduler(name) -> None:
68
+ pipeline.scheduler = SCHEDULER_MAPPING[name].from_config(pipeline.scheduler.config)
69
+
70
+
71
+ def infer(
72
+ prompt,
73
+ scale=7.5,
74
+ steps=40,
75
+ ratio="1.0",
76
+ n_samples=1,
77
+ seed="random",
78
+ negative_prompt="",
79
+ scheduler_name="euler",
80
+ ):
81
+ global noise_scheduler_name
82
+ if noise_scheduler_name != scheduler_name:
83
+ set_noise_scheduler(scheduler_name)
84
+ noise_scheduler_name = scheduler_name
85
+ scale = float(scale)
86
+ steps = int(steps)
87
+ W, H = SD_XL_BASE_RATIOS[ratio]
88
+ n_samples = int(n_samples)
89
+ if seed == "random":
90
+ seed = random.randint(0, 2**32)
91
+ seed = int(seed)
92
+
93
+ images = pipeline(
94
+ prompt=prompt,
95
+ negative_prompt=negative_prompt if len(negative_prompt) > 0 else None,
96
+ guidance_scale=scale,
97
+ generator=torch.Generator(device=device).manual_seed(seed),
98
+ num_images_per_prompt=n_samples,
99
+ num_inference_steps=steps,
100
+ height=H,
101
+ width=W,
102
+ ).images
103
+ # grid = make_image_grid(images, 1, len(images))
104
+ return (
105
+ images,
106
+ {
107
+ "seed": seed,
108
+ },
109
+ )
110
+
111
+
112
+ examples = [
113
+ ["柴犬、カラフルアート"],
114
+ ["満面の笑みのお爺さん、スケッチ"],
115
+ ["星空の中の1匹の鹿、アート"],
116
+ ["ジャングルに立っている日本男性のポートレート"],
117
+ ["茶色の猫のイラスト、アニメ"],
118
+ ["舞妓さんのポートレート、デジタルアート"],
119
+ ]
120
+ with gr.Blocks() as demo:
121
+ gr.Markdown("# Japanese Stable Diffusion XL Demo")
122
+ gr.Markdown(
123
+ """[Japanese Stable Diffusion XL](https://huggingface.co/stabilityai/japanese-stable-diffusion-xl) is a Japanese-version SDXL by [Stability AI](https://ja.stability.ai/).
124
+ - Blog: https://ja.stability.ai/blog/japanese-stable-diffusion-xl
125
+ - Twitter: https://twitter.com/StabilityAI_JP
126
+ - Discord: https://discord.com/invite/StableJP"""
127
+ )
128
+ gr.Markdown(
129
+ "### You can also try JSDXL on Google Colab [here](https://colab.research.google.com/github/Stability-AI/model-demo-notebooks/blob/main/japanese_stable_diffusion_xl.ipynb). "
130
+ )
131
+ with gr.Group():
132
+ with gr.Row():
133
+ prompt = gr.Textbox(
134
+ label="prompt",
135
+ max_lines=1,
136
+ show_label=False,
137
+ placeholder="Enter your prompt",
138
+ container=False,
139
+ )
140
+ btn = gr.Button("Run", scale=0)
141
+ gallery = gr.Gallery(label="Generated images", show_label=False)
142
+ with gr.Accordion(label="sampling info", open=False):
143
+ info = gr.JSON(label="sampling_info")
144
+ with gr.Accordion(open=False, label="Advanced options"):
145
+ scale = gr.Number(value=7.5, label="cfg_scale")
146
+ steps = gr.Number(value=25, label="steps", visible=False)
147
+ size_ratio = gr.Dropdown(
148
+ choices=list(SD_XL_BASE_RATIOS.keys()),
149
+ value="1.0",
150
+ label="size ratio",
151
+ multiselect=False,
152
+ )
153
+ n_samples = gr.Slider(
154
+ minimum=1,
155
+ maximum=3,
156
+ value=2,
157
+ label="n_samples",
158
+ )
159
+ seed = gr.Text(
160
+ value="random",
161
+ label="seed (integer or 'random')",
162
+ )
163
+ negative_prompt = gr.Textbox(
164
+ label="negative prompt",
165
+ value="",
166
+ )
167
+ noise_scheduler = gr.Dropdown(
168
+ list(SCHEDULER_MAPPING.keys()), value="euler", visible=False
169
+ )
170
+
171
+ inputs = [
172
+ prompt,
173
+ scale,
174
+ steps,
175
+ size_ratio,
176
+ n_samples,
177
+ seed,
178
+ negative_prompt,
179
+ noise_scheduler,
180
+ ]
181
+ outputs = [gallery, info]
182
+ prompt.submit(infer, inputs=inputs, outputs=outputs)
183
+ btn.click(infer, inputs=inputs, outputs=outputs)
184
+ gr.Examples(examples=examples, inputs=inputs, outputs=outputs, fn=infer)
185
+
186
+ demo.queue().launch(debug=True, share=True, show_error=True)