| from diffsynth.pipelines.z_image import ZImagePipeline, ModelConfig | |
| import torch | |
| pipe = ZImagePipeline.from_pretrained( | |
| torch_dtype=torch.bfloat16, | |
| device="cuda", | |
| model_configs=[ | |
| ModelConfig(model_id="Tongyi-MAI/Z-Image", origin_file_pattern="transformer/*.safetensors"), | |
| ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors"), | |
| ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), | |
| ], | |
| tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"), | |
| ) | |
| pipe.load_lora(pipe.dit, "./models/train/Z-Image_lora/epoch-4.safetensors") | |
| prompt = "a dog" | |
| image = pipe(prompt=prompt, seed=42, rand_device="cuda", num_inference_steps=50, cfg_scale=4) | |
| image.save("image.jpg") | |