|
import os, torch |
|
|
|
from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline |
|
from kolors.models.modeling_chatglm import ChatGLMModel |
|
from kolors.models.tokenization_chatglm import ChatGLMTokenizer |
|
from diffusers import UNet2DConditionModel, AutoencoderKL |
|
from diffusers import EulerDiscreteScheduler |
|
|
|
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
|
|
|
def infer(prompt): |
|
ckpt_dir = f'{root_dir}/weights/Kolors' |
|
text_encoder = ChatGLMModel.from_pretrained( |
|
f'{ckpt_dir}/text_encoder', |
|
torch_dtype=torch.float16).half() |
|
tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder') |
|
vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half() |
|
scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler") |
|
unet = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half() |
|
pipe = StableDiffusionXLPipeline( |
|
vae=vae, |
|
text_encoder=text_encoder, |
|
tokenizer=tokenizer, |
|
unet=unet, |
|
scheduler=scheduler, |
|
force_zeros_for_empty_prompt=False) |
|
pipe = pipe.to("cuda") |
|
pipe.enable_model_cpu_offload() |
|
image = pipe( |
|
prompt=prompt, |
|
height=1024, |
|
width=1024, |
|
num_inference_steps=50, |
|
guidance_scale=5.0, |
|
num_images_per_prompt=1, |
|
generator= torch.Generator(pipe.device).manual_seed(66)).images[0] |
|
image.save(f'{root_dir}/scripts/outputs/sample_test.jpg') |
|
|
|
|
|
if __name__ == '__main__': |
|
import fire |
|
fire.Fire(infer) |
|
|
