LuChengTHU commited on
Commit
794bf46
β€’
1 Parent(s): 08b6795

add dpmsolver

Browse files
Files changed (3) hide show
  1. app.py +286 -0
  2. nsfw.png +0 -0
  3. requirements.txt +6 -0
app.py ADDED
@@ -0,0 +1,286 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler
2
+ import gradio as gr
3
+ import torch
4
+ from PIL import Image
5
+
6
+ scheduler = DPMSolverMultistepScheduler(
7
+ beta_start=0.00085,
8
+ beta_end=0.012,
9
+ beta_schedule="scaled_linear",
10
+ num_train_timesteps=1000,
11
+ trained_betas=None,
12
+ predict_epsilon=True,
13
+ thresholding=False,
14
+ algorithm_type="dpmsolver++",
15
+ solver_type="midpoint",
16
+ lower_order_final=True,
17
+ )
18
+
19
+ def is_google_colab():
20
+ try:
21
+ import google.colab
22
+ return True
23
+ except:
24
+ return False
25
+
26
+ is_colab = is_google_colab()
27
+
28
+
29
+ class Model:
30
+ def __init__(self, name, path, prefix):
31
+ self.name = name
32
+ self.path = path
33
+ self.prefix = prefix
34
+ self.pipe_t2i = None
35
+ self.pipe_i2i = None
36
+
37
+ models = [
38
+ Model("Custom model", "", ""),
39
+ Model("Stable-Diffusion-v1.4", "runwayml/stable-diffusion-v1-4", "The 1.4 version of official stable-diffusion"),
40
+ Model("Stable-Diffusion-v1.5", "runwayml/stable-diffusion-v1-5", "The 1.5 version of official stable-diffusion"),
41
+ Model("Arcane", "nitrosocke/Arcane-Diffusion", "arcane style "),
42
+ Model("Archer", "nitrosocke/archer-diffusion", "archer style "),
43
+ Model("Elden Ring", "nitrosocke/elden-ring-diffusion", "elden ring style "),
44
+ Model("Spider-Verse", "nitrosocke/spider-verse-diffusion", "spiderverse style "),
45
+ Model("Modern Disney", "nitrosocke/mo-di-diffusion", "modern disney style "),
46
+ Model("Classic Disney", "nitrosocke/classic-anim-diffusion", "classic disney style "),
47
+ Model("Waifu", "hakurei/waifu-diffusion", ""),
48
+ Model("PokΓ©mon", "lambdalabs/sd-pokemon-diffusers", ""),
49
+ Model("Pony Diffusion", "AstraliteHeart/pony-diffusion", ""),
50
+ Model("Robo Diffusion", "nousr/robo-diffusion", ""),
51
+ Model("Cyberpunk Anime", "DGSpitzer/Cyberpunk-Anime-Diffusion", "dgs illustration style "),
52
+ Model("Tron Legacy", "dallinmackay/Tron-Legacy-diffusion", "trnlgcy ")
53
+ ]
54
+
55
+ last_mode = "txt2img"
56
+ current_model = models[1]
57
+ current_model_path = current_model.path
58
+
59
+ if is_colab:
60
+ pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler)
61
+
62
+ else: # download all models
63
+ vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16)
64
+ for model in models[1:]:
65
+ try:
66
+ unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16)
67
+ model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
68
+ model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
69
+ except:
70
+ models.remove(model)
71
+ pipe = models[1].pipe_t2i
72
+
73
+ if torch.cuda.is_available():
74
+ pipe = pipe.to("cuda")
75
+
76
+ device = "GPU πŸ”₯" if torch.cuda.is_available() else "CPU πŸ₯Ά"
77
+
78
+ def custom_model_changed(path):
79
+ models[0].path = path
80
+ global current_model
81
+ current_model = models[0]
82
+
83
+ def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
84
+
85
+ global current_model
86
+ for model in models:
87
+ if model.name == model_name:
88
+ current_model = model
89
+ model_path = current_model.path
90
+
91
+ generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
92
+
93
+ if img is not None:
94
+ return img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator)
95
+ else:
96
+ return txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator)
97
+
98
+ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator=None):
99
+
100
+ global last_mode
101
+ global pipe
102
+ global current_model_path
103
+ if model_path != current_model_path or last_mode != "txt2img":
104
+ current_model_path = model_path
105
+
106
+ if is_colab or current_model == models[0]:
107
+ pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16)
108
+ else:
109
+ pipe.to("cpu")
110
+ pipe = current_model.pipe_t2i
111
+
112
+ if torch.cuda.is_available():
113
+ pipe = pipe.to("cuda")
114
+ last_mode = "txt2img"
115
+
116
+ prompt = current_model.prefix + prompt
117
+ result = pipe(
118
+ prompt,
119
+ negative_prompt = neg_prompt,
120
+ # num_images_per_prompt=n_images,
121
+ num_inference_steps = int(steps),
122
+ guidance_scale = guidance,
123
+ width = width,
124
+ height = height,
125
+ generator = generator)
126
+
127
+ return replace_nsfw_images(result)
128
+
129
+ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator=None):
130
+
131
+ global last_mode
132
+ global pipe
133
+ global current_model_path
134
+ if model_path != current_model_path or last_mode != "img2img":
135
+ current_model_path = model_path
136
+
137
+ if is_colab or current_model == models[0]:
138
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16)
139
+ else:
140
+ pipe.to("cpu")
141
+ pipe = current_model.pipe_i2i
142
+
143
+ if torch.cuda.is_available():
144
+ pipe = pipe.to("cuda")
145
+ last_mode = "img2img"
146
+
147
+ prompt = current_model.prefix + prompt
148
+ ratio = min(height / img.height, width / img.width)
149
+ img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
150
+ result = pipe(
151
+ prompt,
152
+ negative_prompt = neg_prompt,
153
+ # num_images_per_prompt=n_images,
154
+ init_image = img,
155
+ num_inference_steps = int(steps),
156
+ strength = strength,
157
+ guidance_scale = guidance,
158
+ width = width,
159
+ height = height,
160
+ generator = generator)
161
+
162
+ return replace_nsfw_images(result)
163
+
164
+ def replace_nsfw_images(results):
165
+ for i in range(len(results.images)):
166
+ if results.nsfw_content_detected[i]:
167
+ results.images[i] = Image.open("nsfw.png")
168
+ return results.images[0]
169
+
170
+ css = """
171
+ <style>
172
+ .finetuned-diffusion-div {
173
+ text-align: center;
174
+ max-width: 700px;
175
+ margin: 0 auto;
176
+ }
177
+ .finetuned-diffusion-div div {
178
+ display: inline-flex;
179
+ align-items: center;
180
+ gap: 0.8rem;
181
+ font-size: 1.75rem;
182
+ }
183
+ .finetuned-diffusion-div div h1 {
184
+ font-weight: 900;
185
+ margin-bottom: 7px;
186
+ }
187
+ .finetuned-diffusion-div p {
188
+ margin-bottom: 10px;
189
+ font-size: 94%;
190
+ }
191
+ .finetuned-diffusion-div p a {
192
+ text-decoration: underline;
193
+ }
194
+ .tabs {
195
+ margin-top: 0px;
196
+ margin-bottom: 0px;
197
+ }
198
+ #gallery {
199
+ min-height: 20rem;
200
+ }
201
+ </style>
202
+ """
203
+ with gr.Blocks(css=css) as demo:
204
+ gr.HTML(
205
+ f"""
206
+ <div class="finetuned-diffusion-div">
207
+ <div>
208
+ <h1>Finetuned Diffusion</h1>
209
+ </div>
210
+ <p>
211
+ Demo for multiple fine-tuned Stable Diffusion models, trained on different styles: <br>
212
+ <a href="https://huggingface.co/runwayml/stable-diffusion-v1-4">Stable-Diffusion-v1.4</a>, <a href="https://huggingface.co/runwayml/stable-diffusion-v1-5">Stable-Diffusion-v1.5</a>, <a href="https://huggingface.co/nitrosocke/Arcane-Diffusion">Arcane</a>, <a href="https://huggingface.co/nitrosocke/archer-diffusion">Archer</a>, <a href="https://huggingface.co/nitrosocke/elden-ring-diffusion">Elden Ring</a>, <a href="https://huggingface.co/nitrosocke/spider-verse-diffusion">Spider-Verse</a>, <a href="https://huggingface.co/nitrosocke/modern-disney-diffusion">Modern Disney</a>, <a href="https://huggingface.co/nitrosocke/classic-anim-diffusion">Classic Disney</a>, <a href="https://huggingface.co/hakurei/waifu-diffusion">Waifu</a>, <a href="https://huggingface.co/lambdalabs/sd-pokemon-diffusers">PokΓ©mon</a>, <a href="https://huggingface.co/AstraliteHeart/pony-diffusion">Pony Diffusion</a>, <a href="https://huggingface.co/nousr/robo-diffusion">Robo Diffusion</a>, <a href="https://huggingface.co/DGSpitzer/Cyberpunk-Anime-Diffusion">Cyberpunk Anime</a>, <a href="https://huggingface.co/dallinmackay/Tron-Legacy-diffusion">Tron Legacy</a> + any other custom Diffusers 🧨 SD model hosted on HuggingFace πŸ€—.
213
+ </p>
214
+ <p>Don't want to wait in queue? <a href="https://colab.research.google.com/gist/qunash/42112fb104509c24fd3aa6d1c11dd6e0/copy-of-fine-tuned-diffusion-gradio.ipynb"><img data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" src="https://camo.githubusercontent.com/84f0493939e0c4de4e6dbe113251b4bfb5353e57134ffd9fcab6b8714514d4d1/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667"></a></p>
215
+ Running on <b>{device}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
216
+ </p>
217
+ </div>
218
+ """
219
+ )
220
+ with gr.Row():
221
+
222
+ with gr.Column(scale=55):
223
+ with gr.Group():
224
+ model_name = gr.Dropdown(label="Model", choices=[m.name for m in models], value=current_model.name)
225
+ with gr.Box(visible=False) as custom_model_group:
226
+ custom_model_path = gr.Textbox(label="Custom model path", placeholder="Path to model, e.g. nitrosocke/Arcane-Diffusion", interactive=True)
227
+ gr.HTML("<div><font size='2'>Custom models have to be downloaded first, so give it some time.</font></div>")
228
+
229
+ with gr.Row():
230
+ prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder="Enter prompt. Style applied automatically").style(container=False)
231
+ generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))
232
+
233
+
234
+ image_out = gr.Image(height=512)
235
+ # gallery = gr.Gallery(
236
+ # label="Generated images", show_label=False, elem_id="gallery"
237
+ # ).style(grid=[1], height="auto")
238
+
239
+ with gr.Column(scale=45):
240
+ with gr.Tab("Options"):
241
+ with gr.Group():
242
+ neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
243
+
244
+ # n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)
245
+
246
+ with gr.Row():
247
+ guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)
248
+ steps = gr.Slider(label="Steps", value=50, minimum=2, maximum=100, step=1)
249
+
250
+ with gr.Row():
251
+ width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8)
252
+ height = gr.Slider(label="Height", value=512, minimum=64, maximum=1024, step=8)
253
+
254
+ seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)
255
+
256
+ with gr.Tab("Image to image"):
257
+ with gr.Group():
258
+ image = gr.Image(label="Image", height=256, tool="editor", type="pil")
259
+ strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)
260
+
261
+ model_name.change(lambda x: gr.update(visible = x == models[0].name), inputs=model_name, outputs=custom_model_group)
262
+ custom_model_path.change(custom_model_changed, inputs=custom_model_path, outputs=None)
263
+ # n_images.change(lambda n: gr.Gallery().style(grid=[2 if n > 1 else 1], height="auto"), inputs=n_images, outputs=gallery)
264
+
265
+ inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength, neg_prompt]
266
+ prompt.submit(inference, inputs=inputs, outputs=image_out)
267
+ generate.click(inference, inputs=inputs, outputs=image_out)
268
+
269
+ ex = gr.Examples([
270
+ [models[1].name, "jason bateman disassembling the demon core", 7.5, 50],
271
+ [models[4].name, "portrait of dwayne johnson", 7.0, 75],
272
+ [models[5].name, "portrait of a beautiful alyx vance half life", 10, 50],
273
+ [models[6].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 45],
274
+ [models[5].name, "fantasy portrait painting, digital art", 4.0, 30],
275
+ ], [model_name, prompt, guidance, steps, seed], image_out, inference, cache_examples=False)
276
+
277
+ gr.Markdown('''
278
+ Models by [@nitrosocke](https://huggingface.co/nitrosocke), [@haruu1367](https://twitter.com/haruu1367), [@Helixngc7293](https://twitter.com/DGSpitzer) and others. ❀️<br>
279
+ Space by: [![Twitter Follow](https://img.shields.io/twitter/follow/hahahahohohe?label=%40anzorq&style=social)](https://twitter.com/hahahahohohe)
280
+
281
+ ![visitors](https://visitor-badge.glitch.me/badge?page_id=anzorq.finetuned_diffusion)
282
+ ''')
283
+
284
+ if not is_colab:
285
+ demo.queue(concurrency_count=1)
286
+ demo.launch(debug=is_colab, share=is_colab)
nsfw.png ADDED
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
+ torch
3
+ diffusers
4
+ transformers
5
+ scipy
6
+ ftfy