add negative prompt

app.py

@@ -15,11 +15,18 @@ db_diffusion_model = from_pretrained_keras("kfahn/dreambooth-mandelbulb")
 sd_dreambooth_model._diffusion_model = db_diffusion_model
 
 # generate images
-def infer(prompt):
- generated_images = sd_dreambooth_model.text_to_image(
- prompt, batch_size=2
- )
- return generated_images 
+def infer(prompt, negative_prompt, guidance_scale=10, num_inference_steps=50):
+ neg = negative_prompt if negative_prompt else None
+ imgs = []
+ while len(imgs) != 2:
+ next_prompt = pipeline(prompt, negative_prompt=neg, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, num_images_per_prompt=5)
+ for img, is_neg in zip(next_prompt.images, next_prompt.nsfw_content_detected):
+ if not is_neg:
+ imgs.append(img)
+ if len(imgs) == 2:
+ break
+ 
+ return imgs
  
 output = gr.Gallery(label="Outputs").style(grid=(1,2))