nielsgl commited on
Commit
bd83a96
β€’
1 Parent(s): 5055c26

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -19,7 +19,7 @@ resolution = 512
19
  # checkpoint of the converted Stable Diffusion from KerasCV
20
  model_ckpt = "nielsgl/dreambooth-keras-pug-ace-sd2.1"
21
  pipeline = StableDiffusionPipeline.from_pretrained(model_ckpt)
22
- pipeline.to("cuda")
23
 
24
  unique_id = "puggieace"
25
  class_label = "dog"
@@ -30,12 +30,12 @@ image = pipeline(prompt, num_inference_steps=50).images[0]
30
  def infer(prompt, negative_prompt, guidance_scale=10, num_inference_steps=50):
31
  neg = negative_prompt if negative_prompt else None
32
  imgs = []
33
- while len(imgs) != 2:
34
  next_prompt = pipeline(prompt, negative_prompt=neg, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, num_images_per_prompt=5)
35
  for img, is_neg in zip(next_prompt.images, next_prompt.nsfw_content_detected):
36
  if not is_neg:
37
  imgs.append(img)
38
- if len(imgs) == 2:
39
  break
40
 
41
  return imgs
19
  # checkpoint of the converted Stable Diffusion from KerasCV
20
  model_ckpt = "nielsgl/dreambooth-keras-pug-ace-sd2.1"
21
  pipeline = StableDiffusionPipeline.from_pretrained(model_ckpt)
22
+ pipeline = pipeline.to("cuda")
23
 
24
  unique_id = "puggieace"
25
  class_label = "dog"
30
  def infer(prompt, negative_prompt, guidance_scale=10, num_inference_steps=50):
31
  neg = negative_prompt if negative_prompt else None
32
  imgs = []
33
+ while len(imgs) != 4:
34
  next_prompt = pipeline(prompt, negative_prompt=neg, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, num_images_per_prompt=5)
35
  for img, is_neg in zip(next_prompt.images, next_prompt.nsfw_content_detected):
36
  if not is_neg:
37
  imgs.append(img)
38
+ if len(imgs) == 4:
39
  break
40
 
41
  return imgs