naotous commited on
Commit
137c815
1 Parent(s): aa27e70

Fix the error of 'sample' key not found

Browse files

The pipeline output looks like this currently
`ImagePipelineOutput(images=[<PIL.Image.Image image mode=RGB size=256x256 at 0x7F7190061A30>])`

Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -14,7 +14,7 @@ def predict(prompt, steps=100, seed=42, guidance_scale=6.0):
14
  #torch.cuda.empty_cache()
15
  print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
16
  generator = torch.manual_seed(seed)
17
- images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["sample"]
18
  print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
19
  return images[0]
20
 
 
14
  #torch.cuda.empty_cache()
15
  print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
16
  generator = torch.manual_seed(seed)
17
+ images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["images"]
18
  print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
19
  return images[0]
20