Riccardo Giorato commited on
Commit
61e0f27
1 Parent(s): b5da431

update stuff

Browse files
Files changed (4) hide show
  1. .gitignore +2 -0
  2. app.py +127 -114
  3. package.json +9 -0
  4. yarn.lock +4 -0
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+
2
+ node_modules
app.py CHANGED
@@ -6,6 +6,7 @@ import utils
6
 
7
  is_colab = utils.is_google_colab()
8
 
 
9
  class Model:
10
  def __init__(self, name, path, prefix):
11
  self.name = name
@@ -14,15 +15,16 @@ class Model:
14
  self.pipe_t2i = None
15
  self.pipe_i2i = None
16
 
 
17
  models = [
18
- Model("Beeple", "riccardogiorato/beeple-diffusion", "beeple style "),
19
- Model("Avatar", "riccardogiorato/avatar-diffusion", "avatartwow style "),
20
- Model("Beksinski", "s3nh/beksinski-style-stable-diffusion", "beksinski style "),
21
- Model("Poolsuite","prompthero/poolsuite","poolsuite style "),
22
- Model("Robo Diffusion", "nousr/robo-diffusion", ""),
23
- Model("Guohua", "Langboat/Guohua-Diffusion", "guohua style ")
24
- Model("JWST", "dallinmackay/JWST-Deep-Space-diffusion", "JWST ")
25
- ]
26
 
27
  scheduler = DPMSolverMultistepScheduler(
28
  beta_start=0.00085,
@@ -37,53 +39,51 @@ scheduler = DPMSolverMultistepScheduler(
37
  lower_order_final=True,
38
  )
39
 
40
- custom_model = None
41
- if is_colab:
42
- models.insert(0, Model("Custom model", "", ""))
43
- custom_model = models[0]
44
-
45
  last_mode = "txt2img"
46
- current_model = models[1] if is_colab else models[0]
47
  current_model_path = current_model.path
48
 
49
  if is_colab:
50
- pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler)
51
-
52
- else: # download all models
53
- vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16)
54
- for model in models:
55
- try:
56
- unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16)
57
- model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
58
- model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
59
- except:
60
- models.remove(model)
61
- pipe = models[0].pipe_t2i
62
-
 
 
 
 
 
63
  if torch.cuda.is_available():
64
- pipe = pipe.to("cuda")
65
 
66
  device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
67
 
68
- def custom_model_changed(path):
69
- models[0].path = path
70
- global current_model
71
- current_model = models[0]
72
 
73
  def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
74
 
75
- global current_model
76
- for model in models:
77
- if model.name == model_name:
78
- current_model = model
79
- model_path = current_model.path
 
 
 
80
 
81
- generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
 
 
 
82
 
83
- if img is not None:
84
- return img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator)
85
- else:
86
- return txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator)
87
 
88
  def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator=None):
89
 
@@ -93,29 +93,31 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
93
  if model_path != current_model_path or last_mode != "txt2img":
94
  current_model_path = model_path
95
 
96
- if is_colab or current_model == custom_model:
97
- pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler)
 
98
  else:
99
- pipe.to("cpu")
100
- pipe = current_model.pipe_t2i
101
 
102
  if torch.cuda.is_available():
103
- pipe = pipe.to("cuda")
104
  last_mode = "txt2img"
105
 
106
- prompt = current_model.prefix + prompt
107
  result = pipe(
108
- prompt,
109
- negative_prompt = neg_prompt,
110
- # num_images_per_prompt=n_images,
111
- num_inference_steps = int(steps),
112
- guidance_scale = guidance,
113
- width = width,
114
- height = height,
115
- generator = generator)
116
-
117
  return replace_nsfw_images(result)
118
 
 
119
  def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator=None):
120
 
121
  global last_mode
@@ -124,39 +126,43 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
124
  if model_path != current_model_path or last_mode != "img2img":
125
  current_model_path = model_path
126
 
127
- if is_colab or current_model == custom_model:
128
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler)
 
129
  else:
130
- pipe.to("cpu")
131
- pipe = current_model.pipe_i2i
132
-
133
  if torch.cuda.is_available():
134
- pipe = pipe.to("cuda")
135
  last_mode = "img2img"
136
 
137
  prompt = current_model.prefix + prompt
138
  ratio = min(height / img.height, width / img.width)
139
- img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
 
140
  result = pipe(
141
  prompt,
142
- negative_prompt = neg_prompt,
143
  # num_images_per_prompt=n_images,
144
- init_image = img,
145
- num_inference_steps = int(steps),
146
- strength = strength,
147
- guidance_scale = guidance,
148
- width = width,
149
- height = height,
150
- generator = generator)
151
-
152
  return replace_nsfw_images(result)
153
 
 
154
  def replace_nsfw_images(results):
155
  for i in range(len(results.images)):
156
- if results.nsfw_content_detected[i]:
157
- results.images[i] = Image.open("nsfw.png")
158
  return results.images[0]
159
 
 
160
  css = """.playground-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.playground-diffusion-div div h1{font-weight:900;margin-bottom:7px}.playground-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
161
  """
162
  with gr.Blocks(css=css) as demo:
@@ -179,48 +185,55 @@ with gr.Blocks(css=css) as demo:
179
  """
180
  )
181
  with gr.Row():
182
-
183
- with gr.Column(scale=55):
184
- with gr.Group():
185
- model_name = gr.Dropdown(label="Model", choices=[m.name for m in models], value=current_model.name)
186
-
187
- with gr.Row():
188
- prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder="Enter prompt. Style applied automatically").style(container=False)
189
- generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))
190
 
191
-
192
- image_out = gr.Image(height=512)
193
- # gallery = gr.Gallery(
194
- # label="Generated images", show_label=False, elem_id="gallery"
195
- # ).style(grid=[1], height="auto")
196
-
197
- with gr.Column(scale=45):
198
- with gr.Tab("Options"):
199
  with gr.Group():
200
- neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
201
-
202
- # n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)
203
-
204
- with gr.Row():
205
- guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)
206
- steps = gr.Slider(label="Steps", value=25, minimum=2, maximum=75, step=1)
207
 
208
- with gr.Row():
209
- width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8)
210
- height = gr.Slider(label="Height", value=512, minimum=64, maximum=1024, step=8)
 
 
211
 
212
- seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)
 
 
 
213
 
214
- with gr.Tab("Image to image"):
215
- with gr.Group():
216
- image = gr.Image(label="Image", height=256, tool="editor", type="pil")
217
- strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)
218
-
219
- if is_colab:
220
- model_name.change(lambda x: gr.update(visible = x == models[0].name), inputs=model_name, outputs=custom_model_group)
221
- # n_images.change(lambda n: gr.Gallery().style(grid=[2 if n > 1 else 1], height="auto"), inputs=n_images, outputs=gallery)
222
-
223
- inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength, neg_prompt]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
224
  prompt.submit(inference, inputs=inputs, outputs=image_out)
225
  generate.click(inference, inputs=inputs, outputs=image_out)
226
 
@@ -234,5 +247,5 @@ with gr.Blocks(css=css) as demo:
234
  """)
235
 
236
  if not is_colab:
237
- demo.queue(concurrency_count=1)
238
- demo.launch(debug=is_colab, share=is_colab)
 
6
 
7
  is_colab = utils.is_google_colab()
8
 
9
+
10
  class Model:
11
  def __init__(self, name, path, prefix):
12
  self.name = name
 
15
  self.pipe_t2i = None
16
  self.pipe_i2i = None
17
 
18
+
19
  models = [
20
+ Model("Beeple", "riccardogiorato/beeple-diffusion", "beeple style "),
21
+ Model("Avatar", "riccardogiorato/avatar-diffusion", "avatartwow style "),
22
+ Model("Beksinski", "s3nh/beksinski-style-stable-diffusion", "beksinski style "),
23
+ Model("Poolsuite", "prompthero/poolsuite", "poolsuite style "),
24
+ Model("Robo Diffusion", "nousr/robo-diffusion", ""),
25
+ Model("Guohua", "Langboat/Guohua-Diffusion", "guohua style "),
26
+ Model("JWST", "dallinmackay/JWST-Deep-Space-diffusion", "JWST ")
27
+ ]
28
 
29
  scheduler = DPMSolverMultistepScheduler(
30
  beta_start=0.00085,
 
39
  lower_order_final=True,
40
  )
41
 
 
 
 
 
 
42
  last_mode = "txt2img"
43
+ current_model = models[0]
44
  current_model_path = current_model.path
45
 
46
  if is_colab:
47
+ pipe = StableDiffusionPipeline.from_pretrained(
48
+ current_model.path, torch_dtype=torch.float16, scheduler=scheduler)
49
+
50
+ else: # download all models
51
+ vae = AutoencoderKL.from_pretrained(
52
+ current_model.path, subfolder="vae", torch_dtype=torch.float16)
53
+ for model in models:
54
+ try:
55
+ unet = UNet2DConditionModel.from_pretrained(
56
+ model.path, subfolder="unet", torch_dtype=torch.float16)
57
+ model.pipe_t2i = StableDiffusionPipeline.from_pretrained(
58
+ model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
59
+ model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
60
+ model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
61
+ except:
62
+ models.remove(model)
63
+ pipe = models[0].pipe_t2i
64
+
65
  if torch.cuda.is_available():
66
+ pipe = pipe.to("cuda")
67
 
68
  device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
69
 
 
 
 
 
70
 
71
  def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
72
 
73
+ global current_model
74
+ for model in models:
75
+ if model.name == model_name:
76
+ current_model = model
77
+ model_path = current_model.path
78
+
79
+ generator = torch.Generator('cuda').manual_seed(
80
+ seed) if seed != 0 else None
81
 
82
+ if img is not None:
83
+ return img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator)
84
+ else:
85
+ return txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator)
86
 
 
 
 
 
87
 
88
  def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator=None):
89
 
 
93
  if model_path != current_model_path or last_mode != "txt2img":
94
  current_model_path = model_path
95
 
96
+ if is_colab:
97
+ pipe = StableDiffusionPipeline.from_pretrained(
98
+ current_model_path, torch_dtype=torch.float16, scheduler=scheduler)
99
  else:
100
+ pipe.to("cpu")
101
+ pipe = current_model.pipe_t2i
102
 
103
  if torch.cuda.is_available():
104
+ pipe = pipe.to("cuda")
105
  last_mode = "txt2img"
106
 
107
+ prompt = current_model.prefix + prompt
108
  result = pipe(
109
+ prompt,
110
+ negative_prompt=neg_prompt,
111
+ # num_images_per_prompt=n_images,
112
+ num_inference_steps=int(steps),
113
+ guidance_scale=guidance,
114
+ width=width,
115
+ height=height,
116
+ generator=generator)
117
+
118
  return replace_nsfw_images(result)
119
 
120
+
121
  def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator=None):
122
 
123
  global last_mode
 
126
  if model_path != current_model_path or last_mode != "img2img":
127
  current_model_path = model_path
128
 
129
+ if is_colab:
130
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
131
+ current_model_path, torch_dtype=torch.float16, scheduler=scheduler)
132
  else:
133
+ pipe.to("cpu")
134
+ pipe = current_model.pipe_i2i
135
+
136
  if torch.cuda.is_available():
137
+ pipe = pipe.to("cuda")
138
  last_mode = "img2img"
139
 
140
  prompt = current_model.prefix + prompt
141
  ratio = min(height / img.height, width / img.width)
142
+ img = img.resize(
143
+ (int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
144
  result = pipe(
145
  prompt,
146
+ negative_prompt=neg_prompt,
147
  # num_images_per_prompt=n_images,
148
+ init_image=img,
149
+ num_inference_steps=int(steps),
150
+ strength=strength,
151
+ guidance_scale=guidance,
152
+ width=width,
153
+ height=height,
154
+ generator=generator)
155
+
156
  return replace_nsfw_images(result)
157
 
158
+
159
  def replace_nsfw_images(results):
160
  for i in range(len(results.images)):
161
+ if results.nsfw_content_detected[i]:
162
+ results.images[i] = Image.open("nsfw.png")
163
  return results.images[0]
164
 
165
+
166
  css = """.playground-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.playground-diffusion-div div h1{font-weight:900;margin-bottom:7px}.playground-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
167
  """
168
  with gr.Blocks(css=css) as demo:
 
185
  """
186
  )
187
  with gr.Row():
 
 
 
 
 
 
 
 
188
 
189
+ with gr.Column(scale=55):
 
 
 
 
 
 
 
190
  with gr.Group():
191
+ model_name = gr.Dropdown(label="Model", choices=[
192
+ m.name for m in models], value=current_model.name)
 
 
 
 
 
193
 
194
+ with gr.Row():
195
+ prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,
196
+ placeholder="Enter prompt. Style applied automatically").style(container=False)
197
+ generate = gr.Button(value="Generate").style(
198
+ rounded=(False, True, True, False))
199
 
200
+ image_out = gr.Image(height=512)
201
+ # gallery = gr.Gallery(
202
+ # label="Generated images", show_label=False, elem_id="gallery"
203
+ # ).style(grid=[1], height="auto")
204
 
205
+ with gr.Column(scale=45):
206
+ with gr.Tab("Options"):
207
+ with gr.Group():
208
+ neg_prompt = gr.Textbox(
209
+ label="Negative prompt", placeholder="What to exclude from the image")
210
+
211
+ # n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)
212
+
213
+ with gr.Row():
214
+ guidance = gr.Slider(
215
+ label="Guidance scale", value=7.5, maximum=15)
216
+ steps = gr.Slider(
217
+ label="Steps", value=25, minimum=2, maximum=75, step=1)
218
+
219
+ with gr.Row():
220
+ width = gr.Slider(
221
+ label="Width", value=512, minimum=64, maximum=1024, step=8)
222
+ height = gr.Slider(
223
+ label="Height", value=512, minimum=64, maximum=1024, step=8)
224
+
225
+ seed = gr.Slider(
226
+ 0, 2147483647, label='Seed (0 = random)', value=0, step=1)
227
+
228
+ with gr.Tab("Image to image"):
229
+ with gr.Group():
230
+ image = gr.Image(label="Image", height=256,
231
+ tool="editor", type="pil")
232
+ strength = gr.Slider(
233
+ label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)
234
+
235
+ inputs = [model_name, prompt, guidance, steps,
236
+ width, height, seed, image, strength, neg_prompt]
237
  prompt.submit(inference, inputs=inputs, outputs=image_out)
238
  generate.click(inference, inputs=inputs, outputs=image_out)
239
 
 
247
  """)
248
 
249
  if not is_colab:
250
+ demo.queue(concurrency_count=1)
251
+ demo.launch(debug=is_colab, share=is_colab)
package.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "playground_diffusion",
3
+ "version": "1.0.0",
4
+ "repository": "https://huggingface.co/spaces/riccardogiorato/playground_diffusion",
5
+ "license": "MIT",
6
+ "scripts": {
7
+ "install": "pip install -r requirements.txt"
8
+ }
9
+ }
yarn.lock ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ # THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY.
2
+ # yarn lockfile v1
3
+
4
+