Model loading is slow / Time before generating

[Feniksss]

Describe the bug

I have used a simple realization on T4 (Google Cloud)

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

access_token = ""

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=access_token,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]

image.save("astronaut_rides_horse_1.png")

And generating can take 10-15 sec per image - it's good result.
But we have a problem with time before start of generating. Pre-loading models and files can take ~ 30 sec.
For example:

How we can speed up the preloading?

Reproduction

No response

Logs

No response

System Info

Google Cloud - T4

[patrickvonplaten]

Hey @Feniksss,

That's a good question, this should be solved when using:

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=access_token,
    device_map="auto",
)

instead of:

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=access_token,
)
pipe = pipe.to("cuda")

Could you try this?

[patrickvonplaten]

Hey @Feniksss,

That's a good question, this should be solved when using:

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=access_token,
    device_map="auto",
)

instead of:

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=access_token,
)
pipe = pipe.to("cuda")

Could you try this?

[patrickvonplaten]

When using device_map="auto" the weights aren't initialized which should significantly speed up the init

[patrickvonplaten]

In my tests, the time goes down from 22 seconds to 3 seconds. Will open a PR to improve the docs here.

[Feniksss]

In my tests, the time goes down from 22 seconds to 3 seconds. Will open a PR to improve the docs here.

hm... thank you! It improved the time. A delay stay about 10-11 sec now. + 11 sec per image generating.

[piEsposito]

@patrickvonplaten nice! This was a sife effect I didn't think of when I was implementing it.