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))