Fix output of pipeline

#10
by osanseviero - opened
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