anzorq commited on
Commit
b7dbd4c
1 Parent(s): d9193ea

scheduler for sd v2

Browse files
Files changed (1) hide show
  1. app.py +59 -57
app.py CHANGED
@@ -37,19 +37,6 @@ models = [
37
  Model("Pony Diffusion", "AstraliteHeart/pony-diffusion"),
38
  Model("Robo Diffusion", "nousr/robo-diffusion"),
39
  ]
40
-
41
- scheduler = DPMSolverMultistepScheduler(
42
- beta_start=0.00085,
43
- beta_end=0.012,
44
- beta_schedule="scaled_linear",
45
- num_train_timesteps=1000,
46
- trained_betas=None,
47
- predict_epsilon=True,
48
- thresholding=False,
49
- algorithm_type="dpmsolver++",
50
- solver_type="midpoint",
51
- lower_order_final=True,
52
- )
53
 
54
  custom_model = None
55
  if is_colab:
@@ -61,23 +48,20 @@ current_model = models[1] if is_colab else models[0]
61
  current_model_path = current_model.path
62
 
63
  if is_colab:
64
- pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
65
-
66
- else: # download all models
67
- pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler)
68
- # print(f"{datetime.datetime.now()} Downloading vae...")
69
- # vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16)
70
- # for model in models:
71
- # try:
72
- # print(f"{datetime.datetime.now()} Downloading {model.name} model...")
73
- # unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16)
74
- # model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
75
- # model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
76
- # except Exception as e:
77
- # print(f"{datetime.datetime.now()} Failed to load model " + model.name + ": " + str(e))
78
- # models.remove(model)
79
- # pipe = models[0].pipe_t2i
80
-
81
  if torch.cuda.is_available():
82
  pipe = pipe.to("cuda")
83
 
@@ -98,7 +82,7 @@ def on_model_change(model_name):
98
 
99
  return gr.update(visible = model_name == models[0].name), gr.update(placeholder=prefix)
100
 
101
- def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
102
 
103
  print(psutil.virtual_memory()) # print memory usage
104
 
@@ -112,13 +96,13 @@ def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0
112
 
113
  try:
114
  if img is not None:
115
- return img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator), None
116
  else:
117
- return txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator), None
118
  except Exception as e:
119
  return None, error_str(e)
120
 
121
- def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator):
122
 
123
  print(f"{datetime.datetime.now()} txt_to_img, model: {current_model.name}")
124
 
@@ -129,9 +113,18 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
129
  current_model_path = model_path
130
 
131
  if is_colab or current_model == custom_model:
132
- pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
 
 
 
 
 
133
  else:
134
- pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler)
 
 
 
 
135
  # pipe = pipe.to("cpu")
136
  # pipe = current_model.pipe_t2i
137
 
@@ -143,7 +136,7 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
143
  result = pipe(
144
  prompt,
145
  negative_prompt = neg_prompt,
146
- # num_images_per_prompt=n_images,
147
  num_inference_steps = int(steps),
148
  guidance_scale = guidance,
149
  width = width,
@@ -152,7 +145,7 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
152
 
153
  return replace_nsfw_images(result)
154
 
155
- def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):
156
 
157
  print(f"{datetime.datetime.now()} img_to_img, model: {model_path}")
158
 
@@ -163,9 +156,18 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
163
  current_model_path = model_path
164
 
165
  if is_colab or current_model == custom_model:
166
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
 
 
 
 
 
167
  else:
168
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler)
 
 
 
 
169
  # pipe = pipe.to("cpu")
170
  # pipe = current_model.pipe_i2i
171
 
@@ -179,7 +181,7 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
179
  result = pipe(
180
  prompt,
181
  negative_prompt = neg_prompt,
182
- # num_images_per_prompt=n_images,
183
  init_image = img,
184
  num_inference_steps = int(steps),
185
  strength = strength,
@@ -193,12 +195,12 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
193
  def replace_nsfw_images(results):
194
 
195
  if is_colab:
196
- return results.images[0]
197
 
198
  for i in range(len(results.images)):
199
  if results.nsfw_content_detected[i]:
200
  results.images[i] = Image.open("nsfw.png")
201
- return results.images[0]
202
 
203
  css = """.finetuned-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.finetuned-diffusion-div div h1{font-weight:900;margin-bottom:7px}.finetuned-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
204
  """
@@ -216,7 +218,8 @@ with gr.Blocks(css=css) as demo:
216
  <p>You can skip the queue and load custom models in the colab: <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>
217
  Running on <b>{device}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
218
  </p>
219
- <p>You can also duplicate this space and upgrade to gpu by going to settings: <a style="display:inline-block" href="https://huggingface.co/spaces/anzorq/finetuned_diffusion?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></p>
 
220
  </div>
221
  """
222
  )
@@ -234,10 +237,9 @@ with gr.Blocks(css=css) as demo:
234
  generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))
235
 
236
 
237
- image_out = gr.Image(height=512)
238
- # gallery = gr.Gallery(
239
- # label="Generated images", show_label=False, elem_id="gallery"
240
- # ).style(grid=[1], height="auto")
241
  error_output = gr.Markdown()
242
 
243
  with gr.Column(scale=45):
@@ -245,7 +247,7 @@ with gr.Blocks(css=css) as demo:
245
  with gr.Group():
246
  neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
247
 
248
- # n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)
249
 
250
  with gr.Row():
251
  guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)
@@ -267,18 +269,18 @@ with gr.Blocks(css=css) as demo:
267
  custom_model_path.change(custom_model_changed, inputs=custom_model_path, outputs=None)
268
  # n_images.change(lambda n: gr.Gallery().style(grid=[2 if n > 1 else 1], height="auto"), inputs=n_images, outputs=gallery)
269
 
270
- inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength, neg_prompt]
271
- outputs = [image_out, error_output]
272
  prompt.submit(inference, inputs=inputs, outputs=outputs)
273
  generate.click(inference, inputs=inputs, outputs=outputs)
274
 
275
  ex = gr.Examples([
276
- [models[7].name, "tiny cute and adorable kitten adventurer dressed in a warm overcoat with survival gear on a winters day", 7.5, 50],
277
- [models[4].name, "portrait of dwayne johnson", 7.0, 75],
278
- [models[5].name, "portrait of a beautiful alyx vance half life", 10, 50],
279
- [models[6].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 45],
280
- [models[5].name, "fantasy portrait painting, digital art", 4.0, 30],
281
- ], inputs=[model_name, prompt, guidance, steps, seed], outputs=outputs, fn=inference, cache_examples=False)
282
 
283
  gr.HTML("""
284
  <div style="border-top: 1px solid #303030;">
 
37
  Model("Pony Diffusion", "AstraliteHeart/pony-diffusion"),
38
  Model("Robo Diffusion", "nousr/robo-diffusion"),
39
  ]
 
 
 
 
 
 
 
 
 
 
 
 
 
40
 
41
  custom_model = None
42
  if is_colab:
 
48
  current_model_path = current_model.path
49
 
50
  if is_colab:
51
+ pipe = StableDiffusionPipeline.from_pretrained(
52
+ current_model.path,
53
+ torch_dtype=torch.float16,
54
+ scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
55
+ safety_checker=lambda images, clip_input: (images, False)
56
+ )
57
+
58
+ else:
59
+ pipe = StableDiffusionPipeline.from_pretrained(
60
+ current_model.path,
61
+ torch_dtype=torch.float16,
62
+ scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
63
+ )
64
+
 
 
 
65
  if torch.cuda.is_available():
66
  pipe = pipe.to("cuda")
67
 
 
82
 
83
  return gr.update(visible = model_name == models[0].name), gr.update(placeholder=prefix)
84
 
85
+ def inference(model_name, prompt, guidance, steps, n_images=1, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
86
 
87
  print(psutil.virtual_memory()) # print memory usage
88
 
 
96
 
97
  try:
98
  if img is not None:
99
+ return img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator), None
100
  else:
101
+ return txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator), None
102
  except Exception as e:
103
  return None, error_str(e)
104
 
105
+ def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator):
106
 
107
  print(f"{datetime.datetime.now()} txt_to_img, model: {current_model.name}")
108
 
 
113
  current_model_path = model_path
114
 
115
  if is_colab or current_model == custom_model:
116
+ pipe = StableDiffusionPipeline.from_pretrained(
117
+ current_model_path,
118
+ torch_dtype=torch.float16,
119
+ scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
120
+ safety_checker=lambda images, clip_input: (images, False)
121
+ )
122
  else:
123
+ pipe = StableDiffusionPipeline.from_pretrained(
124
+ current_model_path,
125
+ torch_dtype=torch.float16,
126
+ scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
127
+ )
128
  # pipe = pipe.to("cpu")
129
  # pipe = current_model.pipe_t2i
130
 
 
136
  result = pipe(
137
  prompt,
138
  negative_prompt = neg_prompt,
139
+ num_images_per_prompt=n_images,
140
  num_inference_steps = int(steps),
141
  guidance_scale = guidance,
142
  width = width,
 
145
 
146
  return replace_nsfw_images(result)
147
 
148
+ def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator):
149
 
150
  print(f"{datetime.datetime.now()} img_to_img, model: {model_path}")
151
 
 
156
  current_model_path = model_path
157
 
158
  if is_colab or current_model == custom_model:
159
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
160
+ current_model_path,
161
+ torch_dtype=torch.float16,
162
+ scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
163
+ safety_checker=lambda images, clip_input: (images, False)
164
+ )
165
  else:
166
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
167
+ current_model_path,
168
+ torch_dtype=torch.float16,
169
+ scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
170
+ )
171
  # pipe = pipe.to("cpu")
172
  # pipe = current_model.pipe_i2i
173
 
 
181
  result = pipe(
182
  prompt,
183
  negative_prompt = neg_prompt,
184
+ num_images_per_prompt=n_images,
185
  init_image = img,
186
  num_inference_steps = int(steps),
187
  strength = strength,
 
195
  def replace_nsfw_images(results):
196
 
197
  if is_colab:
198
+ return results.images
199
 
200
  for i in range(len(results.images)):
201
  if results.nsfw_content_detected[i]:
202
  results.images[i] = Image.open("nsfw.png")
203
+ return results.images
204
 
205
  css = """.finetuned-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.finetuned-diffusion-div div h1{font-weight:900;margin-bottom:7px}.finetuned-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
206
  """
 
218
  <p>You can skip the queue and load custom models in the colab: <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>
219
  Running on <b>{device}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
220
  </p>
221
+ <p>You can also duplicate this space and upgrade to gpu by going to settings:<br>
222
+ <a style="display:inline-block" href="https://huggingface.co/spaces/anzorq/finetuned_diffusion?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></p>
223
  </div>
224
  """
225
  )
 
237
  generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))
238
 
239
 
240
+ # image_out = gr.Image(height=512)
241
+ gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
242
+
 
243
  error_output = gr.Markdown()
244
 
245
  with gr.Column(scale=45):
 
247
  with gr.Group():
248
  neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
249
 
250
+ n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)
251
 
252
  with gr.Row():
253
  guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)
 
269
  custom_model_path.change(custom_model_changed, inputs=custom_model_path, outputs=None)
270
  # n_images.change(lambda n: gr.Gallery().style(grid=[2 if n > 1 else 1], height="auto"), inputs=n_images, outputs=gallery)
271
 
272
+ inputs = [model_name, prompt, guidance, steps, n_images, width, height, seed, image, strength, neg_prompt]
273
+ outputs = [gallery, error_output]
274
  prompt.submit(inference, inputs=inputs, outputs=outputs)
275
  generate.click(inference, inputs=inputs, outputs=outputs)
276
 
277
  ex = gr.Examples([
278
+ [models[7].name, "tiny cute and adorable kitten adventurer dressed in a warm overcoat with survival gear on a winters day", 7.5, 25],
279
+ [models[4].name, "portrait of dwayne johnson", 7.0, 35],
280
+ [models[5].name, "portrait of a beautiful alyx vance half life", 10, 25],
281
+ [models[6].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 30],
282
+ [models[5].name, "fantasy portrait painting, digital art", 4.0, 20],
283
+ ], inputs=[model_name, prompt, guidance, steps], outputs=outputs, fn=inference, cache_examples=False)
284
 
285
  gr.HTML("""
286
  <div style="border-top: 1px solid #303030;">