stable-diffusion icon indicating copy to clipboard operation
stable-diffusion copied to clipboard

Diffusers Integration in README not work

Open GreatV opened this issue 2 years ago • 0 comments

# make sure you're logged in with `huggingface-cli login`
from torch import autocast
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
	"CompVis/stable-diffusion-v1-4", 
	use_auth_token=True
).to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt)["sample"][0]  
    
image.save("astronaut_rides_horse.png")

Running the above code produces the following output:

Fetching 16 files: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:00<00:00, 10215.99it/s]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:07<00:00,  7.12it/s]
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/greatx/repos/stable-diffusion/test.py:12 in <module>                                       │
│                                                                                                  │
│    9                                                                                             │
│   10 prompt = "a photo of an astronaut riding a horse on mars"                                   │
│   11 with autocast("cuda"):                                                                      │
│ ❱ 12 │   image = pipe(prompt)["sample"][0]                                                       │
│   13                                                                                             │
│   14 image.save("astronaut_rides_horse.png")                                                     │
│   15                                                                                             │
│                                                                                                  │
│ /home/greatx/repos/stable-diffusion/venv/lib/python3.10/site-packages/diffusers/utils/outputs.py │
│ :88 in __getitem__                                                                               │
│                                                                                                  │
│    85 │   def __getitem__(self, k):                                                              │
│    86 │   │   if isinstance(k, str):                                                             │
│    87 │   │   │   inner_dict = {k: v for (k, v) in self.items()}                                 │
│ ❱  88 │   │   │   return inner_dict[k]                                                           │
│    89 │   │   else:                                                                              │
│    90 │   │   │   return self.to_tuple()[k]                                                      │
│    91                                                                                            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
KeyError: 'sample'

Replace sample with images and the above code will work.

# make sure you're logged in with `huggingface-cli login`
from torch import autocast
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
	"CompVis/stable-diffusion-v1-4", 
	use_auth_token=True
).to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt)["images"][0]  
    
image.save("astronaut_rides_horse.png")

GreatV avatar Jan 16 '23 02:01 GreatV