import os, torch # from PIL import Image 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)