enable width param and contains nfsw response on API

#205
by radames - opened
Files changed (1) hide show
  1. app.py +14 -7
app.py CHANGED
@@ -19,7 +19,7 @@ torch.backends.cudnn.benchmark = True
19
  word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
20
  word_list = word_list_dataset["train"]['text']
21
 
22
- def infer(prompt, samples, steps, scale, seed):
23
  #When running locally you can also remove this filter
24
  for filter in word_list:
25
  if re.search(rf"\b{filter}\b", prompt):
@@ -32,6 +32,8 @@ def infer(prompt, samples, steps, scale, seed):
32
  num_inference_steps=steps,
33
  guidance_scale=scale,
34
  generator=generator,
 
 
35
  )
36
  images = []
37
  safe_image = Image.open(r"unsafe.png")
@@ -40,7 +42,7 @@ def infer(prompt, samples, steps, scale, seed):
40
  images.append(safe_image)
41
  else:
42
  images.append(image)
43
- return images
44
 
45
  css = """
46
  .gradio-container {
@@ -136,6 +138,7 @@ examples = [
136
  45,
137
  7.5,
138
  1024,
 
139
  ],
140
  [
141
  'A pikachu fine dining with a view to the Eiffel Tower',
@@ -143,6 +146,7 @@ examples = [
143
  45,
144
  7,
145
  1024,
 
146
  ],
147
  [
148
  'A mecha robot in a favela in expressionist style',
@@ -150,6 +154,7 @@ examples = [
150
  45,
151
  7,
152
  1024,
 
153
  ],
154
  [
155
  'an insect robot preparing a delicious meal',
@@ -157,6 +162,7 @@ examples = [
157
  45,
158
  7,
159
  1024,
 
160
  ],
161
  [
162
  "A small cabin on top of a snowy mountain in the style of Disney, artstation",
@@ -164,6 +170,7 @@ examples = [
164
  45,
165
  7,
166
  1024,
 
167
  ],
168
  ]
169
 
@@ -253,7 +260,6 @@ with block:
253
  ).style(grid=[2], height="auto")
254
 
255
  advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
256
-
257
  with gr.Row(elem_id="advanced-options"):
258
  samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
259
  steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
@@ -267,13 +273,14 @@ with block:
267
  step=1,
268
  randomize=True,
269
  )
270
-
271
- ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, steps, scale, seed], outputs=gallery, cache_examples=True)
 
272
  ex.dataset.headers = [""]
273
 
274
 
275
- text.submit(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)
276
- btn.click(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)
277
  advanced_button.click(
278
  None,
279
  [],
 
19
  word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
20
  word_list = word_list_dataset["train"]['text']
21
 
22
+ def infer(prompt, samples, steps, scale, seed, width):
23
  #When running locally you can also remove this filter
24
  for filter in word_list:
25
  if re.search(rf"\b{filter}\b", prompt):
 
32
  num_inference_steps=steps,
33
  guidance_scale=scale,
34
  generator=generator,
35
+ width=int(width),
36
+ height=int(width),
37
  )
38
  images = []
39
  safe_image = Image.open(r"unsafe.png")
 
42
  images.append(safe_image)
43
  else:
44
  images.append(image)
45
+ return [images, images_list["nsfw_content_detected"]]
46
 
47
  css = """
48
  .gradio-container {
 
138
  45,
139
  7.5,
140
  1024,
141
+ 512,
142
  ],
143
  [
144
  'A pikachu fine dining with a view to the Eiffel Tower',
 
146
  45,
147
  7,
148
  1024,
149
+ 512,
150
  ],
151
  [
152
  'A mecha robot in a favela in expressionist style',
 
154
  45,
155
  7,
156
  1024,
157
+ 512,
158
  ],
159
  [
160
  'an insect robot preparing a delicious meal',
 
162
  45,
163
  7,
164
  1024,
165
+ 512,
166
  ],
167
  [
168
  "A small cabin on top of a snowy mountain in the style of Disney, artstation",
 
170
  45,
171
  7,
172
  1024,
173
+ 512,
174
  ],
175
  ]
176
 
 
260
  ).style(grid=[2], height="auto")
261
 
262
  advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
 
263
  with gr.Row(elem_id="advanced-options"):
264
  samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
265
  steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
 
273
  step=1,
274
  randomize=True,
275
  )
276
+ contains_nfsw = gr.JSON(visible=False)
277
+ width = gr.Number(value=128, visible=False)
278
+ ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, steps, scale, seed, width], outputs=[gallery, contains_nfsw], cache_examples=False)
279
  ex.dataset.headers = [""]
280
 
281
 
282
+ text.submit(infer, inputs=[text, samples, steps, scale, seed, width], outputs=[gallery, contains_nfsw])
283
+ btn.click(infer, inputs=[text, samples, steps, scale, seed, width], outputs=[gallery, contains_nfsw])
284
  advanced_button.click(
285
  None,
286
  [],