osanseviero commited on
Commit
7b423cf
β€’
1 Parent(s): 4c18494

Fix output of pipeline (#10)

Browse files

- Fix output of pipeline (2d5e60271d684ea4e302a41bebe3be112914aee2)

Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -120,7 +120,7 @@ def stable_diffusion_compose(prompt, scale, steps, weights, seed):
120
  generator = th.Generator("cuda").manual_seed(int(seed))
121
  with autocast('cpu' if not th.cuda.is_available() else 'cuda'):
122
  image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps,
123
- weights=weights, generator=generator)["sample"][0]
124
  image.save(f'{"_".join(prompt.split())}.png')
125
  return image
126
 
 
120
  generator = th.Generator("cuda").manual_seed(int(seed))
121
  with autocast('cpu' if not th.cuda.is_available() else 'cuda'):
122
  image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps,
123
+ weights=weights, generator=generator).images[0]
124
  image.save(f'{"_".join(prompt.split())}.png')
125
  return image
126