File size: 1,720 Bytes
89fdd3f | 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 54 55 56 57 58 59 60 61 | # https://hf-mirror.com/stabilityai/stable-cascade
# https://hf-mirror.com/stabilityai/stable-cascade-prior
import torch
from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline, StableCascadeCombinedPipeline
cas = "stabilityai/stable-cascade"
cas_prior = "stabilityai/stable-cascade-prior"
def t2i_(prompt):
prior = StableCascadePriorPipeline.from_pretrained(cas_prior, variant="bf16", torch_dtype=torch.bfloat16)
decoder = StableCascadeDecoderPipeline.from_pretrained(cas, variant="bf16", torch_dtype=torch.float16)
prior.to("cuda")
decoder.to("cuda")
# prior.enable_model_cpu_offload()
# decoder.enable_model_cpu_offload()
prior_output = prior(
prompt=prompt,
height=1024,
width=1024,
negative_prompt="",
guidance_scale=4.0,
num_images_per_prompt=1,
num_inference_steps=20
)
image = decoder(
image_embeddings=prior_output.image_embeddings.to(torch.float16),
prompt=prompt,
negative_prompt="",
guidance_scale=0.0,
output_type="pil",
num_inference_steps=10
).images[0]
return image
def t2i(prompt):
pipe = StableCascadeCombinedPipeline.from_pretrained(cas, variant="bf16", torch_dtype=torch.bfloat16)
pipe.to("cuda")
image = pipe(
prompt=prompt,
negative_prompt="",
num_inference_steps=10,
prior_num_inference_steps=20,
prior_guidance_scale=3.0,
width=1024,
height=1024,
).images[0]
return image
if __name__ == "__main__":
prompt = "a girl in beijing"
image = t2i(prompt)
# image = t2i_(prompt)
image.save("stablecascade_output.png") |