|
from lcm_pipeline import LatentConsistencyModelPipeline |
|
from lcm_scheduler import LCMScheduler |
|
|
|
from diffusers import AutoencoderKL, UNet2DConditionModel |
|
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker |
|
from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor |
|
|
|
import os |
|
import torch |
|
from tqdm import tqdm |
|
from safetensors.torch import load_file |
|
|
|
|
|
prompt = "Self-portrait oil painting, a beautiful cyborg with golden hair" |
|
|
|
|
|
save_path = "./lcm_images" |
|
os.makedirs(save_path, exist_ok=True) |
|
|
|
|
|
|
|
model_id = "digiplay/DreamShaper_7" |
|
|
|
|
|
|
|
vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae") |
|
text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder") |
|
tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") |
|
unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", device_map=None, low_cpu_mem_usage=False, local_files_only=True) |
|
safety_checker = StableDiffusionSafetyChecker.from_pretrained(model_id, subfolder="safety_checker") |
|
feature_extractor = CLIPImageProcessor.from_pretrained(model_id, subfolder="feature_extractor") |
|
|
|
|
|
|
|
scheduler = LCMScheduler(beta_start=0.00085, beta_end=0.0120, beta_schedule="scaled_linear", prediction_type="epsilon") |
|
|
|
|
|
|
|
lcm_unet_ckpt = "./LCM_Dreamshaper_v7_4k.safetensors" |
|
ckpt = load_file(lcm_unet_ckpt) |
|
m, u = unet.load_state_dict(ckpt, strict=False) |
|
if len(m) > 0: |
|
print("missing keys:") |
|
print(m) |
|
if len(u) > 0: |
|
print("unexpected keys:") |
|
print(u) |
|
|
|
|
|
|
|
pipe = LatentConsistencyModelPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler, safety_checker=safety_checker, feature_extractor=feature_extractor) |
|
pipe = pipe.to("cuda") |
|
|
|
|
|
|
|
images = pipe(prompt=prompt, num_images_per_prompt=4, num_inference_steps=4, guidance_scale=8.0, lcm_origin_steps=50).images |
|
|
|
|
|
for i in tqdm(range(len(images))): |
|
output_path = os.path.join(save_path, "{}.png".format(i)) |
|
image = images[i] |
|
image.save(output_path) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
