Got it, just change the following line: image = pipe(prompt).images[0] to: image = pipe(prompt)[“sample”][0]