Update app.py

app.py

@@ -19,7 +19,7 @@ resolution = 512
 # checkpoint of the converted Stable Diffusion from KerasCV
 model_ckpt = "nielsgl/dreambooth-keras-pug-ace-sd2.1"
 pipeline = StableDiffusionPipeline.from_pretrained(model_ckpt)
-pipeline.to("cuda")
+pipeline = pipeline.to("cuda")
 
 unique_id = "puggieace"
 class_label = "dog"
@@ -30,12 +30,12 @@ image = pipeline(prompt, num_inference_steps=50).images[0]
 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:
+    while len(imgs) != 4:
         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:
+            if len(imgs) == 4:
                 break
             
     return imgs