File size: 1,414 Bytes
e96853b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
from diffusers import (
StableDiffusionPipeline,
PNDMScheduler,
)
from diffusers.models import AutoencoderKL
import os
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae = AutoencoderKL.from_pretrained(
"stabilityai/sd-vae-ft-ema", torch_dtype=torch.float32
)
common_config = {"beta_start": 0.00085, "beta_end": 0.012, "beta_schedule": "scaled_linear"}
SCHEDULER = PNDMScheduler(**common_config, skip_prk_steps=True, steps_offset=1,)
HF_API_TOKEN = os.getenv("HF_API_TOKEN")
shared_pipe_kwargs = dict(
vae=vae,
torch_dtype=torch.float32,
revision="fp16",
use_auth_token=HF_API_TOKEN,
scheduler=SCHEDULER,
)
base_sd_pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", **shared_pipe_kwargs
).to(device)
pai_model_dir = '.'
playground_v1_model_dir = os.path.join(
pai_model_dir, "snapshots/36c9e19103f6b897886fb019ebc4d8e86b566032"
)
pgv1_shared_pipe_kwargs = dict(
vae=vae,
torch_dtype=torch.float32,
tokenizer=base_sd_pipe.tokenizer,
feature_extractor=base_sd_pipe.feature_extractor,
text_encoder=base_sd_pipe.text_encoder,
scheduler=SCHEDULER,
local_files_only=True,
)
pgv1_pipe = StableDiffusionPipeline.from_pretrained(
playground_v1_model_dir,
**pgv1_shared_pipe_kwargs,
).to(device)
img = pgv1_pipe("Frog", num_inference_steps=30)
img[0][0].save('frog.png') |