update app

app.py

@@ -28,9 +28,14 @@ prompt = f"A photo of {unique_id} {class_label} in a bucket"
 
 # generate images
 def infer(prompt, negative_prompt, guidance_scale=10, num_inference_steps=50):
-    generated_images = pipeline(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, num_images_per_prompt=2)
-    
-    return generated_images.images
+    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=2)
+        if not next_prompt.nsfw_content_detected[0]:
+            imgs.append(next_prompt.images[0])
+            
+    return imgs
     
 output = gr.Gallery(label="Outputs").style(grid=(1,2))